1 files changed, 139 insertions, 0 deletions
diff --git a/internal/bot/ratelimit.go b/internal/bot/ratelimit.go
new file mode 100644
index 0000000..b5333fd
--- /dev/null
+++ b/internal/bot/ratelimit.go
@@ -0,0 +1,139 @@
+package bot
+
+import (
+	"sync"
+	"time"
+)
+
+// The public bot answers anyone, so one person with a script can spend the
+// whole service's budget against the cep.is API. These constants bound what a
+// single Telegram account may cost.
+//
+// Burst is what a human actually does: paste a handful of CEPs in a row, then
+// think. Refill is the sustained rate afterwards — 20 lookups a minute, which
+// no person reaches by hand and which a loop hits within a second.
+//
+// WarnInterval keeps the refusal from becoming its own flood: someone held at
+// the limit is told once, then ignored until the window passes, so a script
+// hammering the bot does not turn into the bot hammering Telegram.
+const (
+	Burst        = 5
+	Refill       = 3 * time.Second
+	WarnInterval = 30 * time.Second
+)
+
+// The bucket table is swept rather than kept forever: every distinct user gets
+// an entry, and a public bot meets a lot of users. An entry is dropped once it
+// has been idle long enough that it would have refilled to full anyway, at
+// which point forgetting it and recreating it are the same thing.
+const (
+	bucketIdleTTL = 15 * time.Minute
+	sweepEvery    = 5 * time.Minute
+)
+
+// Decision is the answer of [Limiter.Allow].
+type Decision struct {
+	// Allowed reports whether the request may proceed.
+	Allowed bool
+
+	// Warn reports whether the user should be told they are being limited.
+	// It is true at most once per [WarnInterval] per key.
+	Warn bool
+
+	// RetryAfter is roughly how long until the next token, meant for the
+	// message shown to the user.
+	RetryAfter time.Duration
+}
+
+type bucket struct {
+	tokens   float64
+	lastSeen time.Time
+	lastWarn time.Time
+}
+
+// Limiter is a token bucket per key, where the key is the Telegram user id — or
+// the chat id for channel posts, which carry no author.
+//
+// It is in-memory and per-process on purpose: the bot runs as a single
+// container, and a shared store would add a dependency and a failure mode to
+// protect against something a single process can already see.
+type Limiter struct {
+	burst        float64
+	refill       time.Duration
+	warnInterval time.Duration
+	now          func() time.Time
+
+	mu        sync.Mutex
+	buckets   map[int64]*bucket
+	lastSweep time.Time
+}
+
+// NewLimiter builds a limiter with the package constants.
+func NewLimiter() *Limiter {
+	return NewLimiterWith(Burst, Refill, WarnInterval, time.Now)
+}
+
+// NewLimiterWith builds a limiter with explicit parameters and an injectable
+// clock, which is how the test suite exercises refilling without waiting for it.
+func NewLimiterWith(burst int, refill, warnInterval time.Duration, now func() time.Time) *Limiter {
+	if now == nil {
+		now = time.Now
+	}
+	return &Limiter{
+		burst:        float64(burst),
+		refill:       refill,
+		warnInterval: warnInterval,
+		now:          now,
+		buckets:      make(map[int64]*bucket),
+	}
+}
+
+// Allow takes one token for key and reports what the caller should do.
+func (l *Limiter) Allow(key int64) Decision {
+	now := l.now()
+
+	l.mu.Lock()
+	defer l.mu.Unlock()
+
+	l.sweep(now)
+
+	b, ok := l.buckets[key]
+	if !ok {
+		b = &bucket{tokens: l.burst}
+		l.buckets[key] = b
+	} else if l.refill > 0 {
+		b.tokens += float64(now.Sub(b.lastSeen)) / float64(l.refill)
+		if b.tokens > l.burst {
+			b.tokens = l.burst
+		}
+	}
+	b.lastSeen = now
+
+	if b.tokens >= 1 {
+		b.tokens--
+		return Decision{Allowed: true}
+	}
+
+	decision := Decision{RetryAfter: time.Duration((1 - b.tokens) * float64(l.refill))}
+	if b.lastWarn.IsZero() || now.Sub(b.lastWarn) >= l.warnInterval {
+		b.lastWarn = now
+		decision.Warn = true
+	}
+	return decision
+}
+
+func (l *Limiter) sweep(now time.Time) {
+	if l.lastSweep.IsZero() {
+		l.lastSweep = now
+		return
+	}
+	if now.Sub(l.lastSweep) < sweepEvery {
+		return
+	}
+	l.lastSweep = now
+	for key, b := range l.buckets {
+		if now.Sub(b.lastSeen) >= bucketIdleTTL {
+			delete(l.buckets, key)
+		}
+	}
+}