initial commit

2026-08-16 21:18:45 -05:00
commit 1e05a01bcf
122 changed files with 29178 additions and 0 deletions
@@ -0,0 +1,204 @@
+// Package ratelimit provides per-client token buckets for DNS abuse
+// protection.
+//
+// The limiter is sharded by client address so that a busy resolver does not
+// serialise every query behind one mutex, and idle buckets are swept
+// periodically so a flood of unique source addresses cannot grow the map
+// without bound.
+package ratelimit
+
+import (
+	"hash/maphash"
+	"net/netip"
+	"sync"
+	"sync/atomic"
+	"time"
+
+	"github.com/owen/vibedns/internal/netutil"
+)
+
+const shardCount = 32
+
+// Config controls the limiter.
+type Config struct {
+	Enabled bool
+	QPS     int
+	Burst   int
+	Exempt  []string
+}
+
+// bucket is a token bucket that refills continuously.
+type bucket struct {
+	tokens   float64
+	lastFill time.Time
+	lastSeen time.Time
+}
+
+type shard struct {
+	mu      sync.Mutex
+	buckets map[netip.Addr]*bucket
+}
+
+// Limiter enforces a per-client query rate.
+type Limiter struct {
+	shards [shardCount]*shard
+	seed   maphash.Seed
+
+	mu      sync.RWMutex
+	enabled bool
+	qps     float64
+	burst   float64
+	exempt  *netutil.PrefixSet
+
+	allowed atomic.Int64
+	denied  atomic.Int64
+	clients atomic.Int64
+}
+
+// New creates a limiter.
+func New(cfg Config) *Limiter {
+	l := &Limiter{seed: maphash.MakeSeed()}
+	for i := range l.shards {
+		l.shards[i] = &shard{buckets: map[netip.Addr]*bucket{}}
+	}
+	l.SetConfig(cfg)
+	return l
+}
+
+// SetConfig replaces the limiter configuration.
+func (l *Limiter) SetConfig(cfg Config) {
+	if cfg.QPS < 1 {
+		cfg.QPS = 1
+	}
+	if cfg.Burst < cfg.QPS {
+		cfg.Burst = cfg.QPS
+	}
+	l.mu.Lock()
+	l.enabled = cfg.Enabled
+	l.qps = float64(cfg.QPS)
+	l.burst = float64(cfg.Burst)
+	l.exempt = netutil.NewPrefixSet(cfg.Exempt)
+	l.mu.Unlock()
+}
+
+func (l *Limiter) shardFor(addr netip.Addr) *shard {
+	b, _ := addr.MarshalBinary()
+	h := maphash.Bytes(l.seed, b)
+	return l.shards[h%shardCount]
+}
+
+// Allow reports whether a query from addr may be answered.
+//
+// Exempt networks — internal infrastructure, by default loopback — are never
+// limited, so a busy local forwarder cannot be throttled by accident.
+func (l *Limiter) Allow(addr netip.Addr) bool {
+	l.mu.RLock()
+	enabled, qps, burst, exempt := l.enabled, l.qps, l.burst, l.exempt
+	l.mu.RUnlock()
+
+	if !enabled {
+		return true
+	}
+	if !addr.IsValid() {
+		return true
+	}
+	if exempt.Contains(addr) {
+		l.allowed.Add(1)
+		return true
+	}
+
+	now := time.Now()
+	sh := l.shardFor(addr)
+
+	sh.mu.Lock()
+	b, ok := sh.buckets[addr]
+	if !ok {
+		b = &bucket{tokens: burst, lastFill: now}
+		sh.buckets[addr] = b
+		l.clients.Add(1)
+	} else {
+		elapsed := now.Sub(b.lastFill).Seconds()
+		if elapsed > 0 {
+			b.tokens += elapsed * qps
+			if b.tokens > burst {
+				b.tokens = burst
+			}
+			b.lastFill = now
+		}
+	}
+	b.lastSeen = now
+
+	if b.tokens >= 1 {
+		b.tokens--
+		sh.mu.Unlock()
+		l.allowed.Add(1)
+		return true
+	}
+	sh.mu.Unlock()
+	l.denied.Add(1)
+	return false
+}
+
+// Sweep drops buckets that have been idle for longer than maxIdle and returns
+// how many were removed.
+func (l *Limiter) Sweep(maxIdle time.Duration) int {
+	cutoff := time.Now().Add(-maxIdle)
+	removed := 0
+	for _, sh := range l.shards {
+		sh.mu.Lock()
+		for addr, b := range sh.buckets {
+			if b.lastSeen.Before(cutoff) {
+				delete(sh.buckets, addr)
+				removed++
+			}
+		}
+		sh.mu.Unlock()
+	}
+	l.clients.Add(-int64(removed))
+	return removed
+}
+
+// Run sweeps idle buckets until done is closed.
+func (l *Limiter) Run(done <-chan struct{}, interval, maxIdle time.Duration) {
+	t := time.NewTicker(interval)
+	defer t.Stop()
+	for {
+		select {
+		case <-done:
+			return
+		case <-t.C:
+			l.Sweep(maxIdle)
+		}
+	}
+}
+
+// Reset clears every bucket.
+func (l *Limiter) Reset() {
+	for _, sh := range l.shards {
+		sh.mu.Lock()
+		sh.buckets = map[netip.Addr]*bucket{}
+		sh.mu.Unlock()
+	}
+	l.clients.Store(0)
+}
+
+// Stats reports limiter activity.
+type Stats struct {
+	Enabled        bool  `json:"enabled"`
+	Allowed        int64 `json:"allowed"`
+	Denied         int64 `json:"denied"`
+	TrackedClients int64 `json:"tracked_clients"`
+}
+
+// Stats returns the limiter counters.
+func (l *Limiter) Stats() Stats {
+	l.mu.RLock()
+	enabled := l.enabled
+	l.mu.RUnlock()
+	return Stats{
+		Enabled:        enabled,
+		Allowed:        l.allowed.Load(),
+		Denied:         l.denied.Load(),
+		TrackedClients: l.clients.Load(),
+	}
+}