Files
volumen/internal/ratelimit/ratelimit.go
T
petrbalvin f8ed33df83
Test / test (push) Successful in 7m5s
Release / gates (push) Successful in 7m28s
Release / build (amd64, freebsd) (push) Successful in 2m52s
Release / build (amd64, linux) (push) Successful in 2m46s
Release / build (arm64, freebsd) (push) Successful in 2m22s
Release / build (arm64, linux) (push) Successful in 2m38s
Release / build (loong64, linux) (push) Successful in 2m7s
Release / build (riscv64, linux) (push) Successful in 2m17s
Release / release (push) Successful in 1m0s
Initial commit
Assisted-by: GLM 5.3
2026-09-29 10:03:32 +02:00

231 lines
6.5 KiB
Go

// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: PolyForm-Noncommercial-1.0.0
// Package ratelimit implements sliding-window request counters keyed by
// client IP, with LRU-bounded storage and periodic sweeps.
package ratelimit
import (
"container/list"
"sync"
"time"
)
const (
apiMaxTracked = 20_000
apiCleanupInterval = 300 * time.Second
loginWindow = 60 * time.Second
loginMaxAttempts = 10
loginMaxTracked = 10_000
loginCleanupEvery = 300 * time.Second
attemptsPerKeyLimit = 64
)
// window is a sliding-window counter over opaque string keys. Both the
// public API limiter and the login limiter are one of these with
// different limits, which is why the storage, the sweep and the
// eviction live here once.
type window struct {
limit int
duration time.Duration
maxTracked int
cleanupEvery time.Duration
now func() time.Time
mu sync.Mutex
hits map[string][]time.Time
order *list.List // front = most recently used key
positions map[string]*list.Element
lastCleanup time.Time
haveCleaned bool
}
func newWindow(limit int, duration time.Duration, maxTracked int, cleanupEvery time.Duration, now func() time.Time) *window {
if now == nil {
now = time.Now
}
return &window{
limit: limit,
duration: duration,
maxTracked: maxTracked,
cleanupEvery: cleanupEvery,
now: now,
hits: map[string][]time.Time{},
order: list.New(),
positions: map[string]*list.Element{},
}
}
// hit records one attempt for key and returns the number of attempts in
// the window, including this one.
func (w *window) hit(key string) int {
w.mu.Lock()
defer w.mu.Unlock()
now := w.now()
w.sweepLocked(now)
kept := w.pruneLocked(key, now)
kept = append(kept, now)
w.hits[key] = kept
w.touchLocked(key)
w.evictLocked()
return len(kept)
}
// count returns the number of attempts for key inside the window,
// without recording one.
func (w *window) count(key string) int {
w.mu.Lock()
defer w.mu.Unlock()
return len(w.pruneLocked(key, w.now()))
}
// retryAfter returns how long the caller should wait before the oldest
// tracked attempt for key leaves the window.
func (w *window) retryAfter(key string) time.Duration {
w.mu.Lock()
defer w.mu.Unlock()
now := w.now()
kept := w.pruneLocked(key, now)
if len(kept) == 0 {
return 0
}
if wait := kept[0].Add(w.duration).Sub(now); wait > 0 {
return wait
}
return 0
}
// pruneLocked drops the attempts for key that have left the window and
// stores what is left. It also caps how many attempts one key may keep:
// a key that is refused on every request would otherwise grow a slice of
// timestamps for the whole window, unbounded by anything. The cap is far
// above the limit, so the count a caller sees is unaffected; the retry
// delay of a hammering client becomes approximate, which is all it can
// be once it is already being refused. The caller must hold w.mu.
func (w *window) pruneLocked(key string, now time.Time) []time.Time {
cutoff := now.Add(-w.duration)
hits := w.hits[key]
kept := hits[:0]
for _, t := range hits {
if t.After(cutoff) {
kept = append(kept, t)
}
}
if len(kept) == 0 {
delete(w.hits, key)
return kept
}
if len(kept) > w.limit+attemptsPerKeyLimit {
kept = kept[len(kept)-w.limit-attemptsPerKeyLimit:]
}
w.hits[key] = kept
return kept
}
func (w *window) touchLocked(key string) {
if elem, ok := w.positions[key]; ok {
w.order.MoveToFront(elem)
return
}
w.positions[key] = w.order.PushFront(key)
}
// evictLocked drops the least recently used keys once the bound is
// exceeded. It runs on every insert, so the maps cannot grow between
// sweeps. The caller must hold w.mu.
func (w *window) evictLocked() {
for w.order.Len() > w.maxTracked {
oldest := w.order.Back()
if oldest == nil {
return
}
key := oldest.Value.(string)
w.order.Remove(oldest)
delete(w.positions, key)
delete(w.hits, key)
}
}
// sweepLocked drops every attempt that has left its window and trims
// the key set. It is rate-limited by cleanupEvery so a busy server does
// not walk the map on every request. The caller must hold w.mu.
func (w *window) sweepLocked(now time.Time) {
if w.haveCleaned && now.Sub(w.lastCleanup) < w.cleanupEvery {
return
}
cutoff := now.Add(-w.duration)
for key, hits := range w.hits {
kept := hits[:0]
for _, t := range hits {
if t.After(cutoff) {
kept = append(kept, t)
}
}
if len(kept) == 0 {
delete(w.hits, key)
if elem, ok := w.positions[key]; ok {
w.order.Remove(elem)
delete(w.positions, key)
}
continue
}
w.hits[key] = kept
}
w.evictLocked()
w.lastCleanup = now
w.haveCleaned = true
}
// Limiter is a thread-safe sliding-window request counter.
type Limiter struct {
window *window
}
// New creates a limiter allowing limit hits per window.
func New(limit int, window time.Duration) *Limiter {
return &Limiter{window: newWindow(limit, window, apiMaxTracked, apiCleanupInterval, time.Now)}
}
// Limit returns the configured request budget per window.
func (l *Limiter) Limit() int { return l.window.limit }
// Check records a hit for key and reports whether it is allowed, how
// many requests remain, and the retry-after delay in seconds.
func (l *Limiter) Check(key string) (allowed bool, remaining, retryAfter int) {
count := l.window.hit(key)
allowed = count <= l.window.limit
remaining = max(0, l.window.limit-count)
if !allowed {
retryAfter = int(l.window.retryAfter(key).Seconds()) + 1
}
return allowed, remaining, retryAfter
}
// LoginLimiter counts login attempts per IP. Record returns the number
// of attempts in the current window; Blocked reports whether the
// threshold was exceeded together with the retry-after delay.
type LoginLimiter struct {
window *window
}
// NewLoginLimiter creates the shared login-attempt tracker.
func NewLoginLimiter() *LoginLimiter {
return &LoginLimiter{window: newWindow(loginMaxAttempts, loginWindow, loginMaxTracked, loginCleanupEvery, time.Now)}
}
// Record logs an attempt for ip and returns the count in the window.
func (l *LoginLimiter) Record(ip string) int { return l.window.hit(ip) }
// Blocked reports whether ip exceeded the login budget and, when it
// did, the whole-second wait before the oldest attempt expires.
func (l *LoginLimiter) Blocked(ip string) (bool, int) {
if l.window.count(ip) <= l.window.limit {
return false, 0
}
retryAfter := l.window.retryAfter(ip)
if retryAfter <= 0 {
retryAfter = l.window.duration
}
return true, int(retryAfter.Seconds()) + 1
}