Initial commit

internal/server/ratelimit.go

@@ -0,0 +1,101 @@
+package server
+
+import (
+	"net"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"send/internal/config"
+)
+
+// limiter is a per-key token bucket. It is deliberately small: there is no
+// dependency to pull in for this, and uploads are the only thing worth limiting.
+type limiter struct {
+	rate  float64 // tokens per second
+	burst float64
+
+	mu      sync.Mutex
+	buckets map[string]*bucket
+	lastGC  time.Time
+}
+
+type bucket struct {
+	tokens float64
+	last   time.Time
+}
+
+func newLimiter(perHour, burst int) *limiter {
+	if burst < 1 {
+		burst = 1
+	}
+	return &limiter{
+		rate:    float64(perHour) / 3600,
+		burst:   float64(burst),
+		buckets: map[string]*bucket{},
+	}
+}
+
+// allow consumes a token for key, reporting whether one was available.
+func (l *limiter) allow(key string, now time.Time) bool {
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	b, ok := l.buckets[key]
+	if !ok {
+		b = &bucket{tokens: l.burst, last: now}
+		l.buckets[key] = b
+	}
+	b.tokens = min(l.burst, b.tokens+now.Sub(b.last).Seconds()*l.rate)
+	b.last = now
+
+	l.gcLocked(now)
+
+	if b.tokens < 1 {
+		return false
+	}
+	b.tokens--
+	return true
+}
+
+// gcLocked drops buckets that have refilled completely and so carry no state.
+func (l *limiter) gcLocked(now time.Time) {
+	if now.Sub(l.lastGC) < 10*time.Minute {
+		return
+	}
+	l.lastGC = now
+	for k, b := range l.buckets {
+		if b.tokens >= l.burst && now.Sub(b.last) > time.Hour {
+			delete(l.buckets, k)
+		}
+	}
+}
+
+// clientIP identifies the caller for rate-limiting and logging.
+//
+// X-Forwarded-For is trusted only when the immediate peer is a configured
+// proxy, and then only to skip past further trusted hops: the first
+// untrusted address from the right is the real client. Without this the header
+// is pure client input and rate limiting would be trivially bypassed.
+func clientIP(r *http.Request, c *config.Config) string {
+	host, _, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		host = r.RemoteAddr
+	}
+	peer := net.ParseIP(host)
+	if peer == nil || !c.TrustsProxy(peer) {
+		return host
+	}
+	hops := strings.Split(r.Header.Get("X-Forwarded-For"), ",")
+	for i := len(hops) - 1; i >= 0; i-- {
+		ip := net.ParseIP(strings.TrimSpace(hops[i]))
+		if ip == nil {
+			break // malformed: stop trusting the rest of the chain
+		}
+		if !c.TrustsProxy(ip) {
+			return ip.String()
+		}
+	}
+	return host
+}