Codevault initial commit

internal/ratelimit/limiter.go

@@ -0,0 +1,145 @@
+// Package ratelimit implements a per-IP sliding-window failure counter used to
+// hide the database behind a cap on failed lookups.
+package ratelimit
+
+import (
+	"sync"
+	"time"
+)
+
+type bucket struct {
+	index int64
+	count int
+}
+
+type ipState struct {
+	current  bucket
+	previous bucket
+}
+
+// Limiter is a two-bucket sliding-window counter per IP. Allowed is a pure read
+// (RLock) that estimates the failure count in the trailing window from whatever
+// is already stored, without mutating it. RecordFailure takes the write lock and
+// performs the actual window roll and increment. Keeping the previous window's
+// count means a client blocked near the end of one window stays throttled into
+// the next, instead of getting an instant reset at the window boundary.
+type Limiter struct {
+	mu     sync.RWMutex
+	state  map[string]ipState
+	window time.Duration
+	max    int
+}
+
+func NewLimiter(max int, window time.Duration) *Limiter {
+	return &Limiter{
+		state:  make(map[string]ipState),
+		window: window,
+		max:    max,
+	}
+}
+
+func (l *Limiter) windowIndex(t time.Time) int64 {
+	return t.UnixNano() / int64(l.window)
+}
+
+// Allowed reports whether ip is currently under the failure threshold. It never
+// mutates the limiter's state.
+func (l *Limiter) Allowed(ip string) bool {
+	now := time.Now()
+	idx := l.windowIndex(now)
+
+	l.mu.RLock()
+	state, ok := l.state[ip]
+	l.mu.RUnlock()
+	if !ok {
+		return true
+	}
+
+	current, previous := effectiveBuckets(state, idx)
+
+	overlap := overlapFraction(now, l.window, previous.index)
+	estimate := float64(current.count) + float64(previous.count)*overlap
+
+	return estimate < float64(l.max)
+}
+
+// RecordFailure increments the failure count for ip in the current window,
+// rolling the previous stored window forward if it has become stale.
+func (l *Limiter) RecordFailure(ip string) {
+	now := time.Now()
+	idx := l.windowIndex(now)
+
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	state := l.state[ip]
+	current, previous := effectiveBuckets(state, idx)
+	current.count++
+	l.state[ip] = ipState{current: current, previous: previous}
+}
+
+// Cleanup removes IPs whose state is more than 2 windows stale, bounding memory.
+func (l *Limiter) Cleanup() {
+	idx := l.windowIndex(time.Now())
+
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	for ip, state := range l.state {
+		if idx-state.current.index > 2 {
+			delete(l.state, ip)
+		}
+	}
+}
+
+// Run periodically calls Cleanup until ctx is done.
+func (l *Limiter) Run(done <-chan struct{}, interval time.Duration) {
+	ticker := time.NewTicker(interval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-done:
+			return
+		case <-ticker.C:
+			l.Cleanup()
+		}
+	}
+}
+
+// effectiveBuckets derives the current/previous buckets as seen from window idx,
+// without mutating the stored state: if the stored current bucket is stale by
+// exactly one window it becomes the effective previous bucket; if it is older
+// than that, both buckets are treated as empty.
+func effectiveBuckets(state ipState, idx int64) (current, previous bucket) {
+	switch {
+	case state.current.index == idx:
+		current = state.current
+		if state.previous.index == idx-1 {
+			previous = state.previous
+		} else {
+			previous = bucket{index: idx - 1}
+		}
+	case state.current.index == idx-1:
+		current = bucket{index: idx}
+		previous = state.current
+	default:
+		current = bucket{index: idx}
+		previous = bucket{index: idx - 1}
+	}
+	return current, previous
+}
+
+// overlapFraction returns how much of the previous window (identified by
+// previousIdx) still falls within the trailing lookback window of length
+// windowSize measured back from now.
+func overlapFraction(now time.Time, windowSize time.Duration, previousIdx int64) float64 {
+	previousWindowEnd := time.Unix(0, (previousIdx+1)*int64(windowSize))
+	elapsedSincePreviousEnd := now.Sub(previousWindowEnd)
+	if elapsedSincePreviousEnd >= windowSize {
+		return 0
+	}
+	if elapsedSincePreviousEnd <= 0 {
+		return 1
+	}
+	return 1 - float64(elapsedSincePreviousEnd)/float64(windowSize)
+}