internal/server/api/ratelimit.go

Ref: main v0.0.5   Size: 2.9 KiB   History

package api

import (
	"net"
	"net/http"
	"net/netip"
	"sync"
	"time"
)

// enrollBurst is the per-IP bucket capacity for POST /api/v1/enroll — the
// only unauthenticated endpoint. Enrollment is operator-paced (one paste per
// host), so a small burst with a slow refill comfortably covers legitimate
// use while blunting token-guessing.
const enrollBurst = 5

// enrollRefillEvery is how often one token drips back into a bucket.
const enrollRefillEvery = 30 * time.Second

// maxLimiterEntries caps the bucket map so an address-spoofing flood cannot
// grow it without bound; when full, entries idle past a full refill are
// pruned, and if nothing is prunable the request is allowed (fail open — the
// limiter is a brake, not the auth boundary).
const maxLimiterEntries = 10_000

// ipLimiter is a per-IP token bucket. now is injectable for tests.
type ipLimiter struct {
	mu      sync.Mutex
	buckets map[string]*bucket
	now     func() time.Time
}

type bucket struct {
	tokens   float64
	lastSeen time.Time
}

func newIPLimiter(now func() time.Time) *ipLimiter {
	return &ipLimiter{buckets: map[string]*bucket{}, now: now}
}

func (l *ipLimiter) allow(ip string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()
	now := l.now()

	b := l.buckets[ip]
	if b == nil {
		if len(l.buckets) >= maxLimiterEntries && !l.prune(now) {
			return true // fail open; see maxLimiterEntries
		}
		b = &bucket{tokens: enrollBurst}
		l.buckets[ip] = b
	} else {
		refill := float64(now.Sub(b.lastSeen)) / float64(enrollRefillEvery)
		b.tokens = min(enrollBurst, b.tokens+refill)
	}
	b.lastSeen = now
	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}

func (l *ipLimiter) prune(now time.Time) bool {
	idle := time.Duration(enrollBurst) * enrollRefillEvery
	freed := false
	for ip, b := range l.buckets {
		if now.Sub(b.lastSeen) > idle {
			delete(l.buckets, ip)
			freed = true
		}
	}
	return freed
}

// bucketKey normalizes a client IP into its rate-limit bucket: IPv4 (and
// v4-mapped v6) per address; IPv6 by /64, since a single host trivially owns
// billions of v6 addresses and per-address buckets would under-limit it.
// Unparseable input is used raw (still a stable key).
func bucketKey(ip string) string {
	addr, err := netip.ParseAddr(ip)
	if err != nil {
		return ip
	}
	if addr.Is4() || addr.Is4In6() {
		return addr.Unmap().String()
	}
	p, err := addr.Prefix(64)
	if err != nil {
		// Unreachable for a parsed addr (BitLen is always 128 here); kept as
		// belt-and-braces so a refactor can't turn this into a panic.
		return ip
	}
	return p.String()
}

// clientIP extracts the bare IP from r.RemoteAddr. Deployments fronted by a
// reverse proxy see the proxy's address here — enroll rate limiting then
// applies to the proxy as a whole, which is still a meaningful brake.
func clientIP(r *http.Request) string {
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr
	}
	return host
}