xtablo-source/backend/internal/auth/ratelimit.go

package auth

import (
	"sync"
	"time"

	"golang.org/x/time/rate"
)

// entry holds a per-key token-bucket limiter and the last time the key was
// accessed. lastSeen is used by the janitor to evict idle entries (D-16).
type entry struct {
	lim      *rate.Limiter
	lastSeen time.Time
}

// LimiterStore is a map of token-bucket rate limiters keyed on an arbitrary
// string (e.g. lower(email)+":"+clientIP). Each key gets an independent
// Limiter so exhausting one key never affects another (D-16).
//
// The store uses an injectable clock (now) so unit tests can control time
// deterministically without sleeps (Pattern 8). In production, now is time.Now.
//
// Memory is bounded: a janitor goroutine (StartJanitor) removes entries idle
// longer than idleTTL, preventing unlimited growth in high-cardinality key
// spaces (D-16, Pitfall 11).
type LimiterStore struct {
	mu      sync.Mutex
	limits  map[string]*entry
	r       rate.Limit
	burst   int
	idleTTL time.Duration
	now     func() time.Time
}

// NewLimiterStore returns a LimiterStore configured for 5 requests per minute
// (rate.Every(12s)), burst=5, idleTTL=10min. The clock defaults to time.Now.
//
// Start the janitor with StartJanitor before the store is used in production
// to prevent unbounded memory growth.
func NewLimiterStore() *LimiterStore {
	return &LimiterStore{
		limits:  make(map[string]*entry),
		r:       rate.Every(12 * time.Second), // 5 tokens/min
		burst:   5,
		idleTTL: 10 * time.Minute,
		now:     time.Now,
	}
}

// newLimiterStoreWithClock creates a LimiterStore with an injectable clock.
// Used in same-package tests to drive time deterministically (Pattern 8).
func newLimiterStoreWithClock(now func() time.Time) *LimiterStore {
	s := NewLimiterStore()
	s.now = now
	return s
}

// NewLimiterStoreWithClock creates a LimiterStore with an injectable clock.
// Exported for cross-package integration tests (e.g. internal/web handlers tests).
func NewLimiterStoreWithClock(now func() time.Time) *LimiterStore {
	return newLimiterStoreWithClock(now)
}

// SetLimiterClock sets the clock function on an existing LimiterStore.
// Exported for cross-package tests that need to freeze time after construction.
func SetLimiterClock(s *LimiterStore, now func() time.Time) {
	s.mu.Lock()
	defer s.mu.Unlock()
	s.now = now
}

// Allow reports whether the key has a token available. It uses AllowN(t, 1)
// with the injectable clock (not Allow() which uses wall time internally and
// is untestable — Pitfall 8 / Pattern 8).
//
// Allow is safe for concurrent use.
func (s *LimiterStore) Allow(key string) bool {
	s.mu.Lock()
	defer s.mu.Unlock()

	e, ok := s.limits[key]
	if !ok {
		e = &entry{lim: rate.NewLimiter(s.r, s.burst)}
		s.limits[key] = e
	}
	t := s.now()
	e.lastSeen = t
	return e.lim.AllowN(t, 1)
}

// StartJanitor launches a background goroutine that calls cleanupNow on each
// tick. Send to stop to shut the goroutine down cleanly (e.g. on server
// shutdown). interval is typically 1 minute in production.
func (s *LimiterStore) StartJanitor(interval time.Duration, stop <-chan struct{}) {
	go func() {
		tick := time.NewTicker(interval)
		defer tick.Stop()
		for {
			select {
			case <-stop:
				return
			case <-tick.C:
				s.cleanupNow()
			}
		}
	}()
}

// cleanupNow removes all entries whose lastSeen is older than idleTTL.
// It is called by the janitor goroutine and is also exported for direct use
// in tests (same package).
func (s *LimiterStore) cleanupNow() {
	s.mu.Lock()
	defer s.mu.Unlock()

	cutoff := s.now().Add(-s.idleTTL)
	for k, e := range s.limits {
		if e.lastSeen.Before(cutoff) {
			delete(s.limits, k)
		}
	}
}

// size returns the number of tracked keys. Used only in tests.
func (s *LimiterStore) size() int {
	s.mu.Lock()
	defer s.mu.Unlock()
	return len(s.limits)
}