diff options
| author | Peter Stone <thepeterstone@gmail.com> | 2026-03-08 21:03:50 +0000 |
|---|---|---|
| committer | Peter Stone <thepeterstone@gmail.com> | 2026-03-08 21:03:50 +0000 |
| commit | 632ea5a44731af94b6238f330a3b5440906c8ae7 (patch) | |
| tree | d8c780412598d66b89ef390b5729e379fdfd9d5b /internal/api/ratelimit.go | |
| parent | 406247b14985ab57902e8e42898dc8cb8960290d (diff) | |
| parent | 93a4c852bf726b00e8014d385165f847763fa214 (diff) | |
merge: pull latest from master and resolve conflicts
- Resolve conflicts in API server, CLI, and executor.
- Maintain Gemini classification and assignment logic.
- Update UI to use generic agent config and project_dir.
- Fix ProjectDir/WorkingDir inconsistencies in Gemini runner.
- All tests passing after merge.
Diffstat (limited to 'internal/api/ratelimit.go')
| -rw-r--r-- | internal/api/ratelimit.go | 99 |
1 files changed, 99 insertions, 0 deletions
diff --git a/internal/api/ratelimit.go b/internal/api/ratelimit.go new file mode 100644 index 0000000..089354c --- /dev/null +++ b/internal/api/ratelimit.go @@ -0,0 +1,99 @@ +package api + +import ( + "net" + "net/http" + "sync" + "time" +) + +// ipRateLimiter provides per-IP token-bucket rate limiting. +type ipRateLimiter struct { + mu sync.Mutex + limiters map[string]*tokenBucket + rate float64 // tokens replenished per second + burst int // maximum token capacity +} + +// newIPRateLimiter creates a limiter with the given replenishment rate (tokens/sec) +// and burst capacity. Use rate=0 to disable replenishment (tokens never refill). +func newIPRateLimiter(rate float64, burst int) *ipRateLimiter { + return &ipRateLimiter{ + limiters: make(map[string]*tokenBucket), + rate: rate, + burst: burst, + } +} + +func (l *ipRateLimiter) allow(ip string) bool { + l.mu.Lock() + b, ok := l.limiters[ip] + if !ok { + b = &tokenBucket{ + tokens: float64(l.burst), + capacity: float64(l.burst), + rate: l.rate, + lastTime: time.Now(), + } + l.limiters[ip] = b + } + l.mu.Unlock() + return b.allow() +} + +// middleware wraps h with per-IP rate limiting, returning 429 when exceeded. +func (l *ipRateLimiter) middleware(next http.Handler) http.Handler { + return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + ip := realIP(r) + if !l.allow(ip) { + writeJSON(w, http.StatusTooManyRequests, map[string]string{"error": "rate limit exceeded"}) + return + } + next.ServeHTTP(w, r) + }) +} + +// tokenBucket is a simple token-bucket rate limiter for a single key. +type tokenBucket struct { + mu sync.Mutex + tokens float64 + capacity float64 + rate float64 // tokens per second + lastTime time.Time +} + +func (b *tokenBucket) allow() bool { + b.mu.Lock() + defer b.mu.Unlock() + now := time.Now() + if !b.lastTime.IsZero() { + elapsed := now.Sub(b.lastTime).Seconds() + b.tokens = min(b.capacity, b.tokens+elapsed*b.rate) + } + b.lastTime = now + if b.tokens >= 1.0 { + b.tokens-- + return true + } + return false +} + +// realIP extracts the client's real IP from a request. +func realIP(r *http.Request) string { + if xff := r.Header.Get("X-Forwarded-For"); xff != "" { + for i, c := range xff { + if c == ',' { + return xff[:i] + } + } + return xff + } + if xri := r.Header.Get("X-Real-IP"); xri != "" { + return xri + } + host, _, err := net.SplitHostPort(r.RemoteAddr) + if err != nil { + return r.RemoteAddr + } + return host +} |
