archivdms/internal/auth/ratelimit.go

package auth

import (
	"sync"
	"time"
)

// keyedRateLimiter is a minimal in-memory token-bucket limiter keyed by an
// arbitrary string (used for per-(tenant,loginName) and per-IP LDAP login
// throttling). No external dependency, no Redis — buckets are created lazily
// and idle ones are swept periodically to bound memory.
type keyedRateLimiter struct {
	mu           sync.Mutex
	buckets      map[string]*bucket
	burst        float64
	refillPerSec float64
	lastSweep    time.Time
}

type bucket struct {
	tokens float64
	last   time.Time
}

func newKeyedRateLimiter(burst, refillPerSec float64) *keyedRateLimiter {
	return &keyedRateLimiter{
		buckets:      make(map[string]*bucket),
		burst:        burst,
		refillPerSec: refillPerSec,
		lastSweep:    time.Now(),
	}
}

// allow consumes one token for key, returning false when the bucket is empty.
func (l *keyedRateLimiter) allow(key string) bool {
	now := time.Now()
	l.mu.Lock()
	defer l.mu.Unlock()

	l.sweepLocked(now)

	b, ok := l.buckets[key]
	if !ok {
		b = &bucket{tokens: l.burst, last: now}
		l.buckets[key] = b
	}
	b.tokens += now.Sub(b.last).Seconds() * l.refillPerSec
	if b.tokens > l.burst {
		b.tokens = l.burst
	}
	b.last = now

	if b.tokens < 1 {
		return false
	}
	b.tokens--
	return true
}

func (l *keyedRateLimiter) sweepLocked(now time.Time) {
	if now.Sub(l.lastSweep) < time.Minute {
		return
	}
	l.lastSweep = now
	for k, b := range l.buckets {
		if now.Sub(b.last) > 30*time.Minute {
			delete(l.buckets, k)
		}
	}
}