1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
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)
}
}
}
|