vibedns/internal/ratelimit/ratelimit.go

// 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(),
	}
}