feat: harden HTTP server with rate limiting, request timeouts, and sanitized error messages
This commit is contained in:
@@ -0,0 +1,86 @@
|
||||
package serve
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
type visitor struct {
|
||||
tokens float64
|
||||
lastSeen time.Time
|
||||
}
|
||||
|
||||
type rateLimiter struct {
|
||||
mu sync.Mutex
|
||||
visitors map[string]*visitor
|
||||
rate float64 // tokens per second
|
||||
burst float64 // max tokens
|
||||
}
|
||||
|
||||
func newRateLimiter(rate float64, burst int) *rateLimiter {
|
||||
rl := &rateLimiter{
|
||||
visitors: make(map[string]*visitor),
|
||||
rate: rate,
|
||||
burst: float64(burst),
|
||||
}
|
||||
go rl.cleanup()
|
||||
return rl
|
||||
}
|
||||
|
||||
func (rl *rateLimiter) allow(ip string) bool {
|
||||
rl.mu.Lock()
|
||||
defer rl.mu.Unlock()
|
||||
|
||||
v, exists := rl.visitors[ip]
|
||||
now := time.Now()
|
||||
if !exists {
|
||||
rl.visitors[ip] = &visitor{tokens: rl.burst - 1, lastSeen: now}
|
||||
return true
|
||||
}
|
||||
|
||||
elapsed := now.Sub(v.lastSeen).Seconds()
|
||||
v.lastSeen = now
|
||||
v.tokens += elapsed * rl.rate
|
||||
if v.tokens > rl.burst {
|
||||
v.tokens = rl.burst
|
||||
}
|
||||
|
||||
if v.tokens < 1 {
|
||||
return false
|
||||
}
|
||||
v.tokens--
|
||||
return true
|
||||
}
|
||||
|
||||
func (rl *rateLimiter) cleanup() {
|
||||
for range time.Tick(time.Minute) {
|
||||
rl.mu.Lock()
|
||||
for ip, v := range rl.visitors {
|
||||
if time.Since(v.lastSeen) > 5*time.Minute {
|
||||
delete(rl.visitors, ip)
|
||||
}
|
||||
}
|
||||
rl.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func clientIP(r *http.Request) string {
|
||||
if ip := r.Header.Get("X-Forwarded-For"); ip != "" {
|
||||
return ip
|
||||
}
|
||||
if ip := r.Header.Get("X-Real-IP"); ip != "" {
|
||||
return ip
|
||||
}
|
||||
return r.RemoteAddr
|
||||
}
|
||||
|
||||
func withRateLimit(rl *rateLimiter, next http.Handler) http.Handler {
|
||||
return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if !rl.allow(clientIP(r)) {
|
||||
writeError(w, http.StatusTooManyRequests, "rate limit exceeded")
|
||||
return
|
||||
}
|
||||
next.ServeHTTP(w, r)
|
||||
})
|
||||
}
|
||||
Reference in New Issue
Block a user