Added avatars to make it a bit more friendly

internal/api/ratelimit.go

@@ -1,31 +1,33 @@
 package api
 
 import (
-	"net"
-	"net/http"
 	"sync"
 	"time"
 )
 
+// Simple token-bucket rate limiter used by Server.cors middleware.
+
+type tokenBucket struct {
+	tokens   float64
+	lastFill time.Time
+}
+
 type rateLimiter struct {
-	mu     sync.Mutex
-	bk     map[string]*bucket
-	rate   float64 // tokens per second
-	burst  float64
-	window time.Duration
+	rate     float64 // tokens per second
+	burst    float64
+	mu       sync.Mutex
+	bk       map[string]*tokenBucket
+	evictDur time.Duration
+	lastGC   time.Time
 }
 
-type bucket struct {
-	tokens float64
-	last   time.Time
-}
-
-func newRateLimiter(rps float64, burst int, window time.Duration) *rateLimiter {
+func newRateLimiter(rate float64, burst int, evict time.Duration) *rateLimiter {
 	return &rateLimiter{
-		bk:     make(map[string]*bucket),
-		rate:   rps,
-		burst:  float64(burst),
-		window: window,
+		rate:     rate,
+		burst:    float64(burst),
+		bk:       make(map[string]*tokenBucket),
+		evictDur: evict,
+		lastGC:   time.Now(),
 	}
 }
 
@@ -34,45 +36,37 @@ func (rl *rateLimiter) allow(key string) bool {
 	rl.mu.Lock()
 	defer rl.mu.Unlock()
 
-	b := rl.bk[key]
-	if b == nil {
-		b = &bucket{tokens: rl.burst, last: now}
+	// GC old buckets occasionally
+	if now.Sub(rl.lastGC) > rl.evictDur {
+		for k, b := range rl.bk {
+			if now.Sub(b.lastFill) > rl.evictDur {
+				delete(rl.bk, k)
+			}
+		}
+		rl.lastGC = now
+	}
+
+	b, ok := rl.bk[key]
+	if !ok {
+		b = &tokenBucket{tokens: rl.burst, lastFill: now}
 		rl.bk[key] = b
 	}
-	// refill
-	elapsed := now.Sub(b.last).Seconds()
-	b.tokens = min(rl.burst, b.tokens+elapsed*rl.rate)
-	b.last = now
 
-	if b.tokens < 1.0 {
-		return false
-	}
-	b.tokens -= 1.0
+	// Refill
+	elapsed := now.Sub(b.lastFill).Seconds()
+	b.tokens = minf(rl.burst, b.tokens+elapsed*rl.rate)
+	b.lastFill = now
 
-	// occasional cleanup
-	for k, v := range rl.bk {
-		if now.Sub(v.last) > rl.window {
-			delete(rl.bk, k)
-		}
+	if b.tokens >= 1.0 {
+		b.tokens -= 1.0
+		return true
 	}
-	return true
+	return false
 }
 
-func min(a, b float64) float64 {
+func minf(a, b float64) float64 {
 	if a < b {
 		return a
 	}
 	return b
 }
-
-func clientIP(r *http.Request) string {
-	// Prefer Cloudflare’s header if present; fall back to RemoteAddr.
-	if ip := r.Header.Get("CF-Connecting-IP"); ip != "" {
-		return ip
-	}
-	host, _, err := net.SplitHostPort(r.RemoteAddr)
-	if err != nil {
-		return r.RemoteAddr
-	}
-	return host
-}