sub2api/backend/internal/service/invalid_auth_abuse_limiter.go

package service

import (
	"sync"
	"sync/atomic"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/config"
)

const invalidAuthAbuseShardCount = 16

type invalidAuthAbuseEntry struct {
	failures     int
	windowStart  time.Time
	blockedUntil time.Time
}

type invalidAuthAbuseShard struct {
	mu      sync.Mutex
	entries map[string]*invalidAuthAbuseEntry
}

type invalidAuthOverflow struct {
	mu           sync.Mutex
	failures     int
	windowStart  time.Time
	blockedUntil time.Time
}

type invalidAuthAbuseLimiter struct {
	threshold int
	window    time.Duration
	block     time.Duration
	capacity  int64
	shards    [invalidAuthAbuseShardCount]invalidAuthAbuseShard
	overflow  invalidAuthOverflow
	now       func() time.Time

	tracked       atomic.Int64
	recorded      atomic.Uint64
	blocked       atomic.Uint64
	rejected      atomic.Uint64
	expired       atomic.Uint64
	overflowed    atomic.Uint64
	globalBlocked atomic.Uint64
	cleanupNext   atomic.Int64
	cleanupCursor atomic.Uint32
}

type InvalidAuthAbuseHealth struct {
	Enabled       bool   `json:"enabled"`
	Tracked       int64  `json:"tracked"`
	Capacity      int64  `json:"capacity"`
	Recorded      uint64 `json:"recorded"`
	Blocks        uint64 `json:"blocks"`
	Rejected      uint64 `json:"rejected"`
	Expired       uint64 `json:"expired"`
	Overflowed    uint64 `json:"overflowed"`
	GlobalBlocked uint64 `json:"global_blocked"`
}

func newInvalidAuthAbuseLimiter(cfg *config.Config) *invalidAuthAbuseLimiter {
	if cfg == nil || !cfg.APIKeyAuth.InvalidAbuse.Enabled {
		return nil
	}
	c := cfg.APIKeyAuth.InvalidAbuse
	if c.Threshold <= 0 || c.WindowSeconds <= 0 || c.BlockSeconds <= 0 || c.Capacity <= 0 {
		return nil
	}
	l := &invalidAuthAbuseLimiter{
		threshold: c.Threshold,
		window:    time.Duration(c.WindowSeconds) * time.Second,
		block:     time.Duration(c.BlockSeconds) * time.Second,
		capacity:  int64(c.Capacity),
		now:       time.Now,
	}
	for i := range l.shards {
		l.shards[i].entries = make(map[string]*invalidAuthAbuseEntry)
	}
	return l
}

func (s *APIKeyService) CheckInvalidAuthAbuse(clientKey string) (time.Duration, bool) {
	if s == nil || s.invalidAuthAbuse == nil {
		return 0, false
	}
	return s.invalidAuthAbuse.check(clientKey)
}

func (s *APIKeyService) RecordInvalidAuthFailure(clientKey string) {
	if s == nil || s.invalidAuthAbuse == nil {
		return
	}
	s.invalidAuthAbuse.record(clientKey)
}

func (s *APIKeyService) InvalidAuthAbuseHealth() InvalidAuthAbuseHealth {
	if s == nil || s.invalidAuthAbuse == nil {
		return InvalidAuthAbuseHealth{}
	}
	return s.invalidAuthAbuse.health()
}

func (l *invalidAuthAbuseLimiter) check(clientKey string) (time.Duration, bool) {
	if l == nil || clientKey == "" {
		return 0, false
	}
	now := l.now()
	l.maybeCleanupAtCapacity(now)
	shard := l.shard(clientKey)
	shard.mu.Lock()
	entry := shard.entries[clientKey]
	if entry != nil && l.entryExpired(entry, now) {
		delete(shard.entries, clientKey)
		l.tracked.Add(-1)
		l.expired.Add(1)
		entry = nil
	}
	if entry != nil && entry.blockedUntil.After(now) {
		retry := entry.blockedUntil.Sub(now)
		shard.mu.Unlock()
		l.rejected.Add(1)
		return retry, true
	}
	shard.mu.Unlock()

	if entry == nil && l.tracked.Load() >= l.capacity {
		if retry, blocked := l.checkOverflow(now); blocked {
			l.rejected.Add(1)
			l.globalBlocked.Add(1)
			return retry, true
		}
	}
	return 0, false
}

func (l *invalidAuthAbuseLimiter) record(clientKey string) {
	if l == nil || clientKey == "" {
		return
	}
	l.recorded.Add(1)
	now := l.now()
	l.maybeCleanupAtCapacity(now)
	shard := l.shard(clientKey)
	shard.mu.Lock()
	entry := shard.entries[clientKey]
	if entry != nil && l.entryExpired(entry, now) {
		delete(shard.entries, clientKey)
		l.tracked.Add(-1)
		l.expired.Add(1)
		entry = nil
	}
	if entry == nil {
		if !l.reserveEntry() {
			shard.mu.Unlock()
			l.recordOverflow(now)
			return
		}
		entry = &invalidAuthAbuseEntry{windowStart: now}
		shard.entries[clientKey] = entry
	}
	if entry.blockedUntil.After(now) {
		shard.mu.Unlock()
		return
	}
	if entry.windowStart.After(now) || !now.Before(entry.windowStart.Add(l.window)) {
		entry.windowStart = now
		entry.failures = 0
	}
	entry.failures++
	if entry.failures >= l.threshold {
		entry.failures = 0
		entry.blockedUntil = now.Add(l.block)
		entry.windowStart = entry.blockedUntil
		l.blocked.Add(1)
	}
	shard.mu.Unlock()
}

func (l *invalidAuthAbuseLimiter) reserveEntry() bool {
	for {
		current := l.tracked.Load()
		if current >= l.capacity {
			return false
		}
		if l.tracked.CompareAndSwap(current, current+1) {
			return true
		}
	}
}

func (l *invalidAuthAbuseLimiter) maybeCleanupAtCapacity(now time.Time) {
	if l.tracked.Load() < l.capacity {
		return
	}
	nowUnixNano := now.UnixNano()
	for {
		next := l.cleanupNext.Load()
		if nowUnixNano < next {
			return
		}
		if l.cleanupNext.CompareAndSwap(next, now.Add(100*time.Millisecond).UnixNano()) {
			break
		}
	}
	index := l.cleanupCursor.Add(1) - 1
	shard := &l.shards[index%invalidAuthAbuseShardCount]
	shard.mu.Lock()
	for key, entry := range shard.entries {
		if l.entryExpired(entry, now) {
			delete(shard.entries, key)
			l.tracked.Add(-1)
			l.expired.Add(1)
		}
	}
	shard.mu.Unlock()
}

func (l *invalidAuthAbuseLimiter) entryExpired(entry *invalidAuthAbuseEntry, now time.Time) bool {
	return entry != nil && !entry.blockedUntil.After(now) && !entry.windowStart.After(now) && !now.Before(entry.windowStart.Add(l.window))
}

func (l *invalidAuthAbuseLimiter) shard(clientKey string) *invalidAuthAbuseShard {
	const fnvOffset32 = uint32(2166136261)
	const fnvPrime32 = uint32(16777619)
	hash := fnvOffset32
	for i := 0; i < len(clientKey); i++ {
		hash ^= uint32(clientKey[i])
		hash *= fnvPrime32
	}
	return &l.shards[hash%invalidAuthAbuseShardCount]
}

func (l *invalidAuthAbuseLimiter) recordOverflow(now time.Time) {
	l.overflowed.Add(1)
	l.overflow.mu.Lock()
	defer l.overflow.mu.Unlock()
	if l.overflow.blockedUntil.After(now) {
		return
	}
	if l.overflow.windowStart.IsZero() || !now.Before(l.overflow.windowStart.Add(l.window)) {
		l.overflow.windowStart = now
		l.overflow.failures = 0
	}
	l.overflow.failures++
	if l.overflow.failures >= l.threshold {
		l.overflow.failures = 0
		l.overflow.blockedUntil = now.Add(l.block)
		l.overflow.windowStart = l.overflow.blockedUntil
		l.blocked.Add(1)
	}
}

func (l *invalidAuthAbuseLimiter) checkOverflow(now time.Time) (time.Duration, bool) {
	l.overflow.mu.Lock()
	defer l.overflow.mu.Unlock()
	if l.overflow.blockedUntil.After(now) {
		return l.overflow.blockedUntil.Sub(now), true
	}
	return 0, false
}

func (l *invalidAuthAbuseLimiter) health() InvalidAuthAbuseHealth {
	return InvalidAuthAbuseHealth{
		Enabled:       true,
		Tracked:       l.tracked.Load(),
		Capacity:      l.capacity,
		Recorded:      l.recorded.Load(),
		Blocks:        l.blocked.Load(),
		Rejected:      l.rejected.Load(),
		Expired:       l.expired.Load(),
		Overflowed:    l.overflowed.Load(),
		GlobalBlocked: l.globalBlocked.Load(),
	}
}