104 lines
2.6 KiB
Go
104 lines
2.6 KiB
Go
package api
|
|
|
|
import (
|
|
"net/http"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
// Per-IP token bucket rate limiting (#2). Goroutine-safe via mutex.
|
|
|
|
type bucket struct {
|
|
tokens float64
|
|
last time.Time
|
|
rate float64 // tokens per second
|
|
burst float64
|
|
}
|
|
|
|
type limiter struct {
|
|
mu sync.Mutex
|
|
buckets map[string]*bucket
|
|
}
|
|
|
|
func newLimiter() *limiter {
|
|
return &limiter{buckets: make(map[string]*bucket)}
|
|
}
|
|
|
|
func (l *limiter) allow(key string, rate, burst float64) bool {
|
|
l.mu.Lock()
|
|
defer l.mu.Unlock()
|
|
now := time.Now()
|
|
b, ok := l.buckets[key]
|
|
if !ok {
|
|
b = &bucket{tokens: burst, last: now, rate: rate, burst: burst}
|
|
l.buckets[key] = b
|
|
}
|
|
elapsed := now.Sub(b.last).Seconds()
|
|
b.tokens += elapsed * b.rate
|
|
if b.tokens > b.burst {
|
|
b.tokens = b.burst
|
|
}
|
|
b.last = now
|
|
if b.tokens < 1 {
|
|
return false
|
|
}
|
|
b.tokens--
|
|
return true
|
|
}
|
|
|
|
// clientIP extracts the request IP (no reverse proxy header by default).
|
|
func clientIP(r *http.Request) string {
|
|
host := r.RemoteAddr
|
|
if i := strings.LastIndex(host, ":"); i > 0 {
|
|
host = host[:i]
|
|
}
|
|
return host
|
|
}
|
|
|
|
var globalLimiter = newLimiter()
|
|
|
|
// globalSettingsFn is set at startup; tests can point it at fixed settings.
|
|
var globalSettingsFn func() Settings
|
|
|
|
func globalSettings() Settings {
|
|
if globalSettingsFn != nil {
|
|
return globalSettingsFn()
|
|
}
|
|
return defaultSettings(Config{})
|
|
}
|
|
|
|
// rateLimitCreate uses the admin-tunable burst and per-minute refill (#40).
|
|
func rateLimitCreate(r *http.Request, s Settings) bool {
|
|
return globalLimiter.allow("create:"+clientIP(r), s.RateLimitPerMinute/60.0, s.RateLimitBurst)
|
|
}
|
|
|
|
// rateLimitGuess: 1 req/sec refill, burst 5, per IP.
|
|
func rateLimitGuess(r *http.Request) bool {
|
|
return globalLimiter.allow("guess:"+clientIP(r), 1, 5)
|
|
}
|
|
|
|
// rateLimitUnlock: 5 per minute per IP+paste.
|
|
func rateLimitUnlock(id string, r *http.Request) bool {
|
|
return globalLimiter.allow("unlock:"+id+":"+clientIP(r), 5.0/60.0, 5)
|
|
}
|
|
|
|
// rateLimitAdmin: 5 attempts per minute per IP on the admin key check (#66),
|
|
// same pattern as the unlock limiter (#34).
|
|
func rateLimitAdmin(r *http.Request) bool {
|
|
return globalLimiter.allow("admin:"+clientIP(r), 5.0/60.0, 5)
|
|
}
|
|
|
|
// writeRateLimited responds 429 with Retry-After based on refill rate.
|
|
func writeRateLimited(w http.ResponseWriter, retryAfterSecs int) {
|
|
w.Header().Set("Retry-After", strconv.Itoa(retryAfterSecs))
|
|
writeErr(w, 429, "rate limit exceeded")
|
|
}
|
|
|
|
// setRateLimitHeaders sets informational X-RateLimit headers for create/guess.
|
|
func setRateLimitHeaders(w http.ResponseWriter, limit, burst int) {
|
|
w.Header().Set("X-RateLimit-Limit", strconv.Itoa(limit))
|
|
w.Header().Set("X-RateLimit-Burst", strconv.Itoa(burst))
|
|
}
|