NextMyDay/internal/core/middleware/ratelimit.go

package middleware

import (
	"net/http"
	"strconv"
	"sync"
	"time"

	"git.lohmar.co.uk/cclohmar/NextMyDay/internal/core/routing"
)

// RateLimiter is a per-key fixed-window limiter (A3, SPEC-003 §4/AC-9). It
// guards the auth-sensitive routes: /{GUID}, /access, /setup. State is
// in-memory and single-process, which matches the single-instance model
// (D-012).
//
// The state map is pruned: elapsed windows are swept as new keys arrive and the
// map has a hard cap, so a flood of distinct keys cannot grow it without bound
// (M3s — a memory-DoS defence).
type RateLimiter struct {
	mu      sync.Mutex
	limit   int
	window  time.Duration
	now     func() time.Time
	state   map[string]*window
	sweptAt time.Time // when the last sweep ran
}

type window struct {
	count int
	reset time.Time
}

// DefaultRateLimit / DefaultRateWindow are the code defaults for the
// auth-sensitive routes.
const (
	DefaultRateLimit  = 30
	DefaultRateWindow = time.Minute
)

// The map's bounds (M3s). sweepThreshold is the size above which an elapsed
// window triggers a sweep — below it the map is too small to be worth it — and
// sweeps run at most once per window. maxStateEntries is the hard cap, so the
// map can never grow without bound whatever the arrival pattern.
const (
	sweepThreshold  = 1 << 12
	maxStateEntries = 1 << 16
)

// NewRateLimiter builds a limiter allowing limit requests per window per key.
// Non-positive values fall back to the defaults.
func NewRateLimiter(limit int, win time.Duration) *RateLimiter {
	if limit <= 0 {
		limit = DefaultRateLimit
	}
	if win <= 0 {
		win = DefaultRateWindow
	}
	return &RateLimiter{
		limit:  limit,
		window: win,
		now:    time.Now,
		state:  make(map[string]*window),
	}
}

// Allow records a request for key and reports whether it is within the limit.
func (l *RateLimiter) Allow(key string) bool {
	l.mu.Lock()
	defer l.mu.Unlock()

	now := l.now()
	w, ok := l.state[key]
	if !ok || !now.Before(w.reset) {
		if !ok {
			l.pruneLocked(now)
		}
		l.state[key] = &window{count: 1, reset: now.Add(l.window)}
		return true
	}
	w.count++
	return w.count <= l.limit
}

// pruneLocked drops elapsed windows (M3s). It is called on insert of a NEW
// key, and is cheap: a sweep needs the map to be worth sweeping
// (sweepThreshold) and runs at most once per window. At the hard cap the map is
// reset outright — a bounded limiter that occasionally forgets beats an
// unbounded one an attacker can hold open.
func (l *RateLimiter) pruneLocked(now time.Time) {
	if len(l.state) >= maxStateEntries {
		l.state = make(map[string]*window)
		l.sweptAt = now
		return
	}
	if len(l.state) < sweepThreshold || now.Sub(l.sweptAt) < l.window {
		return
	}
	for k, w := range l.state {
		if !now.Before(w.reset) {
			delete(l.state, k)
		}
	}
	l.sweptAt = now
}

// Len reports how many distinct keys the limiter is tracking. It is for tests
// and observability; the map is bounded by maxStateEntries (M3s).
func (l *RateLimiter) Len() int {
	l.mu.Lock()
	defer l.mu.Unlock()
	return len(l.state)
}

// SetClock injects a clock (tests).
func (l *RateLimiter) SetClock(fn func() time.Time) {
	if fn == nil {
		return
	}
	l.mu.Lock()
	defer l.mu.Unlock()
	l.now = fn
}

// RetryAfter returns the whole seconds until the window resets for key (for the
// Retry-After header). It is best-effort.
func (l *RateLimiter) RetryAfter(key string) int {
	l.mu.Lock()
	defer l.mu.Unlock()
	w, ok := l.state[key]
	if !ok {
		return 0
	}
	d := time.Until(w.reset)
	if d < 0 {
		return 0
	}
	return int(d.Seconds()) + 1
}

// Class classifies a path into an auth-sensitive bucket, or "" when the path is
// not rate-limited. /access and /setup are explicit; any non-reserved first
// segment is a GUID lookup.
func Class(path string) string {
	p := routing.Parse(path)
	if len(p) == 0 || p[0] == "" {
		return ""
	}
	switch p[0] {
	case "access", "setup":
		return p[0]
	}
	if routing.IsReservedWorkflow(p[0]) {
		return ""
	}
	return "guid"
}

// AuthSensitive reports whether path is one of the rate-limited routes (A3).
func AuthSensitive(path string) bool { return Class(path) != "" }

// Middleware wraps next with a limiter keyed by keyFn. It answers 429 with a
// Retry-After header once the limit is exceeded.
func (l *RateLimiter) Middleware(keyFn func(*http.Request) string, next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		key := keyFn(r)
		if !l.Allow(key) {
			if secs := l.RetryAfter(key); secs > 0 {
				w.Header().Set("Retry-After", strconv.Itoa(secs))
			}
			http.Error(w, "too many requests", http.StatusTooManyRequests)
			return
		}
		next.ServeHTTP(w, r)
	})
}