// 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