veylant/internal/ratelimit/limiter.go

// Package ratelimit implements per-tenant and per-user token bucket rate limiting.
// Buckets are kept in memory for fast Allow() checks; configuration is persisted
// in PostgreSQL and refreshed on admin mutations without restart.
package ratelimit

import (
	"fmt"
	"sync"

	"go.uber.org/zap"
	"golang.org/x/time/rate"
)

// RateLimitConfig holds rate limiting parameters for one tenant.
// When IsEnabled is false, all requests are allowed regardless of token counts.
type RateLimitConfig struct {
	TenantID       string
	RequestsPerMin int  // tenant-wide requests per minute
	BurstSize      int  // tenant-wide burst capacity
	UserRPM        int  // per-user requests per minute within the tenant
	UserBurst      int  // per-user burst capacity
	IsEnabled      bool // false → bypass all limits for this tenant
}

// Limiter is a thread-safe rate limiter backed by per-tenant and per-user
// token buckets. It is the zero value of a sync.RWMutex-protected map.
type Limiter struct {
	mu         sync.RWMutex
	tenantBkts map[string]*rate.Limiter // key: tenantID
	userBkts   map[string]*rate.Limiter // key: tenantID:userID
	configs    map[string]RateLimitConfig
	defaultCfg RateLimitConfig
	logger     *zap.Logger
}

// New creates a Limiter with the given default configuration applied to every
// tenant that has no explicit per-tenant override.
func New(defaultCfg RateLimitConfig, logger *zap.Logger) *Limiter {
	return &Limiter{
		tenantBkts: make(map[string]*rate.Limiter),
		userBkts:   make(map[string]*rate.Limiter),
		configs:    make(map[string]RateLimitConfig),
		defaultCfg: defaultCfg,
		logger:     logger,
	}
}

// Allow returns true if both the tenant-wide and per-user budgets allow the
// request. Returns false (→ HTTP 429) if either limit is exceeded or if the
// tenant config has IsEnabled=false (always allowed in that case).
func (l *Limiter) Allow(tenantID, userID string) bool {
	cfg := l.configFor(tenantID)
	if !cfg.IsEnabled {
		return true
	}

	l.mu.Lock()
	defer l.mu.Unlock()

	// Tenant bucket
	tb, ok := l.tenantBkts[tenantID]
	if !ok {
		tb = newBucket(cfg.RequestsPerMin, cfg.BurstSize)
		l.tenantBkts[tenantID] = tb
	}
	if !tb.Allow() {
		l.logger.Debug("rate limit: tenant bucket exceeded",
			zap.String("tenant_id", tenantID))
		return false
	}

	// User bucket (key: tenantID:userID)
	userKey := tenantID + ":" + userID
	ub, ok := l.userBkts[userKey]
	if !ok {
		ub = newBucket(cfg.UserRPM, cfg.UserBurst)
		l.userBkts[userKey] = ub
	}
	if !ub.Allow() {
		l.logger.Debug("rate limit: user bucket exceeded",
			zap.String("tenant_id", tenantID),
			zap.String("user_id", userID))
		return false
	}

	return true
}

// SetConfig stores a per-tenant config and resets the tenant's buckets so the
// new limits take effect immediately (next request builds fresh buckets).
func (l *Limiter) SetConfig(cfg RateLimitConfig) {
	l.mu.Lock()
	defer l.mu.Unlock()

	l.configs[cfg.TenantID] = cfg

	// Invalidate old buckets so they are rebuilt with the new limits.
	delete(l.tenantBkts, cfg.TenantID)
	for k := range l.userBkts {
		if len(k) > len(cfg.TenantID)+1 && k[:len(cfg.TenantID)+1] == cfg.TenantID+":" {
			delete(l.userBkts, k)
		}
	}
}

// DeleteConfig removes a per-tenant override, reverting to default limits.
// Existing buckets are reset immediately.
func (l *Limiter) DeleteConfig(tenantID string) {
	l.mu.Lock()
	defer l.mu.Unlock()

	delete(l.configs, tenantID)
	delete(l.tenantBkts, tenantID)
	for k := range l.userBkts {
		if len(k) > len(tenantID)+1 && k[:len(tenantID)+1] == tenantID+":" {
			delete(l.userBkts, k)
		}
	}
}

// ListConfigs returns all explicit per-tenant configurations (not the default).
func (l *Limiter) ListConfigs() []RateLimitConfig {
	l.mu.RLock()
	defer l.mu.RUnlock()

	out := make([]RateLimitConfig, 0, len(l.configs))
	for _, c := range l.configs {
		out = append(out, c)
	}
	return out
}

// GetConfig returns the effective config for a tenant (explicit or default).
func (l *Limiter) GetConfig(tenantID string) RateLimitConfig {
	return l.configFor(tenantID)
}

// configFor returns the tenant-specific config if one exists, otherwise default.
func (l *Limiter) configFor(tenantID string) RateLimitConfig {
	l.mu.RLock()
	cfg, ok := l.configs[tenantID]
	l.mu.RUnlock()
	if ok {
		return cfg
	}
	c := l.defaultCfg
	c.TenantID = tenantID
	return c
}

// newBucket creates a token bucket refilling at rpm tokens per minute with the
// given burst capacity. Uses golang.org/x/time/rate (token bucket algorithm).
func newBucket(rpm, burst int) *rate.Limiter {
	if rpm <= 0 {
		rpm = 1
	}
	if burst <= 0 {
		burst = 1
	}
	// rate.Limit is tokens per second; convert RPM → RPS.
	rps := rate.Limit(float64(rpm) / 60.0)
	return rate.NewLimiter(rps, burst)
}

// userKey builds the map key for per-user buckets.
func userKey(tenantID, userID string) string {
	return fmt.Sprintf("%s:%s", tenantID, userID)
}

// Ensure userKey is used (suppress unused warning from compiler).
var _ = userKey