job-search/internal/middleware/ratelimit.go

package middleware

import (
	"context"
	"hash/maphash"
	"math"
	"net/http"
	"strconv"
	"sync"
	"time"

	"golang.org/x/time/rate"

	"git.avsdev.uk/Go/job-search/internal/httpx"
)

// Limiter is the rate limiting hook. Implementations decide whether a request
// identified by key may proceed. Swap in a Redis/memcached-backed limiter for
// multi-instance deployments by implementing this interface.
type Limiter interface {
	Allow(key string) (allowed bool, retryAfter time.Duration)
}

// LimiterFunc adapts a function to the Limiter interface.
type LimiterFunc func(key string) (bool, time.Duration)

func (f LimiterFunc) Allow(key string) (bool, time.Duration) { return f(key) }

// KeyFunc derives the rate limit key from a request.
type KeyFunc func(*http.Request) string

// KeyByIP keys on the real client IP (after proxy processing).
func KeyByIP(r *http.Request) string { return ClientIP(r) }

// RateLimit turns a Limiter into middleware. Rejected requests get 429 with
// a Retry-After header. A nil limiter yields a pass-through middleware.
func RateLimit(l Limiter, key KeyFunc) func(http.Handler) http.Handler {
	if l == nil {
		return func(next http.Handler) http.Handler { return next }
	}
	if key == nil {
		key = KeyByIP
	}
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			ok, retry := l.Allow(key(r))
			if !ok {
				secs := int(math.Ceil(retry.Seconds()))
				if secs < 1 {
					secs = 1
				}
				w.Header().Set("Retry-After", strconv.Itoa(secs))
				httpx.WriteProblem(w, r, http.StatusTooManyRequests, "rate limit exceeded")
				return
			}
			next.ServeHTTP(w, r)
		})
	}
}

// KeyedLimiter is an in-memory token bucket per key, sharded to reduce lock
// contention, with idle-entry eviction.
type KeyedLimiter struct {
	limit  rate.Limit
	burst  int
	ttl    time.Duration
	seed   maphash.Seed
	shards [32]limiterShard
}

type limiterShard struct {
	mu      sync.Mutex
	entries map[string]*limiterEntry
}

type limiterEntry struct {
	lim  *rate.Limiter
	last time.Time
}

// NewKeyedLimiter creates a limiter allowing rps sustained requests per
// second with the given burst. A janitor goroutine evicts entries idle for
// longer than ttl until ctx is cancelled.
func NewKeyedLimiter(ctx context.Context, rps float64, burst int, ttl time.Duration) *KeyedLimiter {
	k := &KeyedLimiter{limit: rate.Limit(rps), burst: burst, ttl: ttl, seed: maphash.MakeSeed()}
	for i := range k.shards {
		k.shards[i].entries = make(map[string]*limiterEntry)
	}
	if ttl > 0 {
		go k.janitor(ctx)
	}
	return k
}

func (k *KeyedLimiter) shard(key string) *limiterShard {
	return &k.shards[maphash.String(k.seed, key)%uint64(len(k.shards))]
}

// Allow implements Limiter.
func (k *KeyedLimiter) Allow(key string) (bool, time.Duration) {
	now := time.Now()
	s := k.shard(key)
	s.mu.Lock()
	e := s.entries[key]
	if e == nil {
		e = &limiterEntry{lim: rate.NewLimiter(k.limit, k.burst)}
		s.entries[key] = e
	}
	e.last = now
	lim := e.lim
	s.mu.Unlock()

	res := lim.ReserveN(now, 1)
	if !res.OK() {
		return false, time.Second
	}
	if d := res.DelayFrom(now); d > 0 {
		res.CancelAt(now)
		return false, d
	}
	return true, 0
}

func (k *KeyedLimiter) janitor(ctx context.Context) {
	t := time.NewTicker(k.ttl / 2)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case now := <-t.C:
			for i := range k.shards {
				s := &k.shards[i]
				s.mu.Lock()
				for key, e := range s.entries {
					if now.Sub(e.last) > k.ttl {
						delete(s.entries, key)
					}
				}
				s.mu.Unlock()
			}
		}
	}
}