99 lines
2.6 KiB
Go
99 lines
2.6 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
)
|
|
|
|
// ScriptRunner is the slice of a Redis client this package needs. Narrow on
|
|
// purpose: it keeps go-redis out of this file's imports, which means the window
|
|
// logic below is testable without a Redis server. See redis_adapter.go for the
|
|
// ten-line bridge to *redis.Client.
|
|
type ScriptRunner interface {
|
|
Eval(ctx context.Context, script string, keys []string, args ...any) (any, error)
|
|
}
|
|
|
|
// INCR and EXPIRE must be one atomic step. Done as two commands, a crash
|
|
// between them leaves a counter with no TTL, which silently becomes a permanent
|
|
// ban for that key.
|
|
//
|
|
// Returns {count, ttl_ms}.
|
|
const allowScript = `
|
|
local current = redis.call('INCR', KEYS[1])
|
|
if current == 1 then
|
|
redis.call('PEXPIRE', KEYS[1], ARGV[1])
|
|
end
|
|
return {current, redis.call('PTTL', KEYS[1])}
|
|
`
|
|
|
|
// Redis is a fixed-window limiter.
|
|
//
|
|
// Fixed window has a known flaw: a client can send `limit` requests at the end
|
|
// of one window and `limit` more at the start of the next, briefly achieving
|
|
// twice the rate. A sliding window or token bucket avoids it at the cost of
|
|
// more state. For protecting a backend from a runaway bot instance the simple
|
|
// version is enough, and the burst is bounded at 2x rather than unbounded.
|
|
type Redis struct {
|
|
Client ScriptRunner
|
|
Prefix string // namespace, so SCAN can find these keys when debugging
|
|
}
|
|
|
|
func NewRedis(c ScriptRunner) *Redis {
|
|
return &Redis{Client: c, Prefix: "rl:"}
|
|
}
|
|
|
|
func (r *Redis) Allow(ctx context.Context, key string, limit int, dur time.Duration) (Result, error) {
|
|
if limit <= 0 {
|
|
return Result{Allowed: true}, nil
|
|
}
|
|
|
|
raw, err := r.Client.Eval(ctx, allowScript,
|
|
[]string{r.Prefix + key}, dur.Milliseconds())
|
|
if err != nil {
|
|
return Result{}, fmt.Errorf("ratelimit: eval: %w", err)
|
|
}
|
|
|
|
vals, ok := raw.([]any)
|
|
if !ok || len(vals) != 2 {
|
|
return Result{}, fmt.Errorf("ratelimit: unexpected script result %T", raw)
|
|
}
|
|
|
|
count, ok1 := toInt64(vals[0])
|
|
ttlMS, ok2 := toInt64(vals[1])
|
|
|
|
if !ok1 || !ok2 {
|
|
return Result{}, fmt.Errorf("ratelimit: non-integer script result")
|
|
}
|
|
|
|
remaining := limit - int(count)
|
|
if remaining < 0 {
|
|
remaining = 0
|
|
}
|
|
|
|
retry := time.Duration(ttlMS) * time.Millisecond
|
|
if ttlMS < 0 {
|
|
retry = dur
|
|
}
|
|
|
|
return Result{
|
|
Allowed: int(count) <= limit,
|
|
Limit: limit,
|
|
Remaining: remaining,
|
|
RetryAfter: retry,
|
|
}, nil
|
|
}
|
|
|
|
// Redis integers arrive as int64 through go-redis, but a fake or a future
|
|
// client version may hand back int. Accept both rather than panicking.
|
|
func toInt64(v any) (int64, bool) {
|
|
switch n := v.(type) {
|
|
case int64:
|
|
return n, true
|
|
case int:
|
|
return int64(n), true
|
|
default:
|
|
return 0, false
|
|
}
|
|
}
|