134 lines
3.4 KiB
Go
134 lines
3.4 KiB
Go
// Package iplimit provides a per-IP rate-limiting HTTP middleware backed by
|
|
// a token bucket from golang.org/x/time/rate. Each unique IP gets its own
|
|
// bucket; rejected requests are served a 429 response with a Retry-After
|
|
// header.
|
|
package iplimit
|
|
|
|
import (
|
|
"math"
|
|
"net"
|
|
"net/http"
|
|
"strconv"
|
|
"sync"
|
|
"time"
|
|
|
|
"git.juancwu.dev/juancwu/errx"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
const (
|
|
op = "iplimit"
|
|
|
|
// idleTTL is how long a per-IP bucket survives without activity before
|
|
// being evicted from the in-memory map.
|
|
idleTTL = 10 * time.Minute
|
|
|
|
// sweepEvery bounds how often the eviction pass runs (lazily, on the
|
|
// next request after the interval expires).
|
|
sweepEvery = idleTTL / 10
|
|
)
|
|
|
|
// timeNow is the clock the middleware reads. Tests override it to drive the
|
|
// limiter and the eviction sweep deterministically.
|
|
var timeNow = time.Now
|
|
|
|
type entry struct {
|
|
limiter *rate.Limiter
|
|
seen time.Time
|
|
}
|
|
|
|
// New returns a per-IP rate-limiting middleware. r is the steady-state token
|
|
// refill rate; burst is the maximum burst size. Use rate.Every(d) to express
|
|
// the rate as one event per duration d.
|
|
//
|
|
// Pair with realip.New() upstream so the limiter keys on the originating
|
|
// client IP rather than the proxy peer. Loopback, private, link-local, and
|
|
// unspecified addresses pass through unlimited — they are typically internal
|
|
// callers (health checks, dev) that should not be rate limited.
|
|
//
|
|
// Idle buckets are evicted after 10 minutes of inactivity. The sweep runs
|
|
// inline on the next request after the interval expires, so the only state
|
|
// retained between requests is the entries map itself — no background
|
|
// goroutines.
|
|
//
|
|
// Rejected requests receive a 429 Too Many Requests response with a
|
|
// Retry-After header (in whole seconds, ceiling).
|
|
func New(r rate.Limit, burst int) func(http.Handler) http.Handler {
|
|
if r <= 0 {
|
|
panic(errx.New(op, "rate must be > 0"))
|
|
}
|
|
if burst <= 0 {
|
|
panic(errx.New(op, "burst must be > 0"))
|
|
}
|
|
|
|
var (
|
|
mu sync.Mutex
|
|
entries = make(map[string]*entry)
|
|
nextSweep time.Time
|
|
)
|
|
|
|
return func(next http.Handler) http.Handler {
|
|
return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
|
|
ip := parseClientIP(req.RemoteAddr)
|
|
if ip == nil || isLocal(ip) {
|
|
next.ServeHTTP(w, req)
|
|
return
|
|
}
|
|
|
|
now := timeNow()
|
|
key := ip.String()
|
|
|
|
mu.Lock()
|
|
if now.After(nextSweep) {
|
|
for k, e := range entries {
|
|
if now.Sub(e.seen) > idleTTL {
|
|
delete(entries, k)
|
|
}
|
|
}
|
|
nextSweep = now.Add(sweepEvery)
|
|
}
|
|
e, ok := entries[key]
|
|
if !ok {
|
|
e = &entry{limiter: rate.NewLimiter(r, burst)}
|
|
entries[key] = e
|
|
}
|
|
e.seen = now
|
|
lim := e.limiter
|
|
mu.Unlock()
|
|
|
|
res := lim.ReserveN(now, 1)
|
|
delay := res.DelayFrom(now)
|
|
if delay > 0 {
|
|
res.CancelAt(now)
|
|
w.Header().Set("Retry-After", strconv.Itoa(retryAfterSeconds(delay)))
|
|
http.Error(w, "Too Many Requests", http.StatusTooManyRequests)
|
|
return
|
|
}
|
|
next.ServeHTTP(w, req)
|
|
})
|
|
}
|
|
}
|
|
|
|
func retryAfterSeconds(d time.Duration) int {
|
|
s := int(math.Ceil(d.Seconds()))
|
|
if s < 1 {
|
|
return 1
|
|
}
|
|
return s
|
|
}
|
|
|
|
func parseClientIP(remoteAddr string) net.IP {
|
|
host := remoteAddr
|
|
if h, _, err := net.SplitHostPort(remoteAddr); err == nil {
|
|
host = h
|
|
}
|
|
return net.ParseIP(host)
|
|
}
|
|
|
|
func isLocal(ip net.IP) bool {
|
|
return ip.IsLoopback() ||
|
|
ip.IsPrivate() ||
|
|
ip.IsLinkLocalUnicast() ||
|
|
ip.IsLinkLocalMulticast() ||
|
|
ip.IsUnspecified()
|
|
}
|