62 lines
1.4 KiB
Go
62 lines
1.4 KiB
Go
|
// Package ratelimit provides a simple in-memory fixed-window limiter
|
||
|
|
// keyed by client IP, used to protect authentication endpoints.
|
||
|
|
package ratelimit
|
||
|
|
|
||
|
|
import (
|
||
|
|
"net"
|
||
|
|
"net/http"
|
||
|
|
"strings"
|
||
|
|
"sync"
|
||
|
|
"time"
|
||
|
|
)
|
||
|
|
|
||
|
|
// Limiter is a fixed-window, per-IP request limiter.
|
||
|
|
type Limiter struct {
|
||
|
|
mu sync.Mutex
|
||
|
|
max int
|
||
|
|
window time.Duration
|
||
|
|
hits map[string]*hits
|
||
|
|
}
|
||
|
|
|
||
|
|
type hits struct {
|
||
|
|
windowStart time.Time
|
||
|
|
count int
|
||
|
|
}
|
||
|
|
|
||
|
|
// New returns a limiter allowing max requests per window per IP.
|
||
|
|
func New(max int, window time.Duration) *Limiter {
|
||
|
|
return &Limiter{max: max, window: window, hits: make(map[string]*hits)}
|
||
|
|
}
|
||
|
|
|
||
|
|
// Allow reports whether the request passes the limit, consuming one slot.
|
||
|
|
func (l *Limiter) Allow(r *http.Request) bool {
|
||
|
|
ip := ClientIP(r)
|
||
|
|
now := time.Now()
|
||
|
|
|
||
|
|
l.mu.Lock()
|
||
|
|
defer l.mu.Unlock()
|
||
|
|
h := l.hits[ip]
|
||
|
|
if h == nil || now.Sub(h.windowStart) >= l.window {
|
||
|
|
h = &hits{windowStart: now}
|
||
|
|
l.hits[ip] = h
|
||
|
|
}
|
||
|
|
h.count++
|
||
|
|
return h.count <= l.max
|
||
|
|
}
|
||
|
|
|
||
|
|
// ClientIP extracts the remote address, preferring X-Forwarded-For when
|
||
|
|
// running behind a reverse proxy.
|
||
|
|
func ClientIP(r *http.Request) string {
|
||
|
|
if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
|
||
|
|
if i := strings.IndexByte(fwd, ','); i > 0 {
|
||
|
|
return strings.TrimSpace(fwd[:i])
|
||
|
|
}
|
||
|
|
return strings.TrimSpace(fwd)
|
||
|
|
}
|
||
|
|
host, _, err := net.SplitHostPort(r.RemoteAddr)
|
||
|
|
if err != nil {
|
||
|
|
return r.RemoteAddr
|
||
|
|
}
|
||
|
|
return host
|
||
|
|
}
|