popertotsandSisyphus
6064a06b7d
Move RateLimiter, ProgressiveRateLimiter, and cleanup functions to ratelimit_shared.go. Add dev and prod rate limiter tests. Remove inline cleanup goroutines. Ultraworked with [Sisyphus](https://github.com/code-yeongyu/oh-my-openagent) Co-authored-by: Sisyphus <clio-agent@sisyphuslabs.ai>
105 lines
2.6 KiB
Go
105 lines
2.6 KiB
Go
// Shared types, registration, and cleanup for RateLimiter + ProgressiveRateLimiter.
|
|
// Used by both production (!dev) and dev (dev) builds — keep tag-free.
|
|
|
|
package mw
|
|
|
|
import (
|
|
"context"
|
|
"crussell/clock"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// RateLimiter implements a simple in-memory rate limiter
|
|
type RateLimiter struct {
|
|
requests map[string][]time.Time
|
|
mu sync.RWMutex
|
|
limit int
|
|
window time.Duration
|
|
}
|
|
|
|
// ProgressiveRateLimiter implements per-IP rate limiting with increasing backoff.
|
|
// Designed for bot-spam prevention across accounts (not account-specific lockout).
|
|
type ProgressiveRateLimiter struct {
|
|
requests map[string]*ipProgressiveState
|
|
mu sync.RWMutex
|
|
}
|
|
|
|
type ipProgressiveState struct {
|
|
// Timestamps of all requests within the tracking window
|
|
timestamps []time.Time
|
|
}
|
|
|
|
var (
|
|
registeredLimiters []*RateLimiter
|
|
registeredLimitersMu sync.Mutex
|
|
)
|
|
|
|
func registerLimiter(rl *RateLimiter) {
|
|
registeredLimitersMu.Lock()
|
|
registeredLimiters = append(registeredLimiters, rl)
|
|
registeredLimitersMu.Unlock()
|
|
}
|
|
|
|
// CleanupAllRateLimiters runs Cleanup on every registered RateLimiter.
|
|
// Called by the centralised jobs scheduler.
|
|
func CleanupAllRateLimiters(ctx context.Context) (int, error) {
|
|
registeredLimitersMu.Lock()
|
|
limiters := make([]*RateLimiter, len(registeredLimiters))
|
|
copy(limiters, registeredLimiters)
|
|
registeredLimitersMu.Unlock()
|
|
for _, l := range limiters {
|
|
l.Cleanup()
|
|
}
|
|
return 0, nil
|
|
}
|
|
|
|
// CleanupProgressiveRateLimiter runs Cleanup on the global progressive rate limiter.
|
|
// Called by the centralised jobs scheduler.
|
|
func CleanupProgressiveRateLimiter(ctx context.Context) (int, error) {
|
|
globalProgressiveLimiter.Cleanup()
|
|
return 0, nil
|
|
}
|
|
|
|
// Cleanup removes expired entries from the rate limiter map.
|
|
func (rl *RateLimiter) Cleanup() {
|
|
rl.mu.Lock()
|
|
defer rl.mu.Unlock()
|
|
now := clock.Now()
|
|
for key, times := range rl.requests {
|
|
var valid []time.Time
|
|
for _, t := range times {
|
|
if now.Sub(t) < rl.window {
|
|
valid = append(valid, t)
|
|
}
|
|
}
|
|
if len(valid) == 0 {
|
|
delete(rl.requests, key)
|
|
} else {
|
|
rl.requests[key] = valid
|
|
}
|
|
}
|
|
}
|
|
|
|
// Cleanup removes expired entries from the progressive rate limiter map.
|
|
func (prl *ProgressiveRateLimiter) Cleanup() {
|
|
prl.mu.Lock()
|
|
defer prl.mu.Unlock()
|
|
cutoff := clock.Now().Add(-60 * time.Second)
|
|
for ip, state := range prl.requests {
|
|
var valid []time.Time
|
|
for _, t := range state.timestamps {
|
|
if t.After(cutoff) {
|
|
valid = append(valid, t)
|
|
}
|
|
}
|
|
if len(valid) == 0 {
|
|
delete(prl.requests, ip)
|
|
} else {
|
|
state.timestamps = valid
|
|
}
|
|
}
|
|
}
|
|
|
|
var globalProgressiveLimiter = NewProgressiveRateLimiter()
|