116 lines
3.5 KiB
Go
116 lines
3.5 KiB
Go
package ratelimit
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
type Policy struct {
|
|
RatePerSecond float64
|
|
Burst int
|
|
}
|
|
|
|
type Result struct {
|
|
Allowed bool
|
|
Remaining int
|
|
RetryAfter time.Duration
|
|
}
|
|
|
|
type Limiter struct {
|
|
client *redis.Client
|
|
prefix string
|
|
policies map[string]Policy
|
|
script *redis.Script
|
|
}
|
|
|
|
func New(client *redis.Client, prefix string, policy Policy) (*Limiter, error) {
|
|
if client == nil || prefix == "" || policy.RatePerSecond <= 0 || policy.Burst <= 0 {
|
|
return nil, errors.New("redis client, prefix and positive rate policy are required")
|
|
}
|
|
return NewPolicies(client, prefix, map[string]Policy{"*": policy})
|
|
}
|
|
|
|
// NewPolicies builds a limiter whose token-bucket policy is selected by scope.
|
|
// A "*" policy is used as the fallback for unrecognised scopes.
|
|
func NewPolicies(client *redis.Client, prefix string, policies map[string]Policy) (*Limiter, error) {
|
|
if client == nil || prefix == "" || len(policies) == 0 {
|
|
return nil, errors.New("redis client, prefix and rate policies are required")
|
|
}
|
|
cloned := make(map[string]Policy, len(policies))
|
|
for scope, policy := range policies {
|
|
if scope == "" || policy.RatePerSecond <= 0 || policy.Burst <= 0 {
|
|
return nil, errors.New("rate policy scopes, rates and bursts must be positive")
|
|
}
|
|
cloned[scope] = policy
|
|
}
|
|
return &Limiter{client: client, prefix: prefix, policies: cloned, script: redis.NewScript(tokenBucketScript)}, nil
|
|
}
|
|
|
|
func (l *Limiter) Allow(ctx context.Context, scope, identifier string) (Result, error) {
|
|
if scope == "" || identifier == "" {
|
|
return Result{}, errors.New("rate limit scope and identifier are required")
|
|
}
|
|
policy, ok := l.policies[scope]
|
|
if !ok {
|
|
policy, ok = l.policies["*"]
|
|
}
|
|
if !ok {
|
|
return Result{}, fmt.Errorf("rate limit policy for scope %q is not configured", scope)
|
|
}
|
|
key := fmt.Sprintf("%s:ratelimit:%s:%s", l.prefix, scope, IdentifierHash(identifier))
|
|
values, err := l.script.Run(ctx, l.client, []string{key}, policy.RatePerSecond, policy.Burst).Slice()
|
|
if err != nil {
|
|
return Result{}, fmt.Errorf("execute rate limit: %w", err)
|
|
}
|
|
if len(values) != 3 {
|
|
return Result{}, errors.New("unexpected rate limit response")
|
|
}
|
|
allowed, ok1 := values[0].(int64)
|
|
remaining, ok2 := values[1].(int64)
|
|
retryMilliseconds, ok3 := values[2].(int64)
|
|
if !ok1 || !ok2 || !ok3 {
|
|
return Result{}, errors.New("invalid rate limit response types")
|
|
}
|
|
return Result{
|
|
Allowed: allowed == 1, Remaining: int(remaining),
|
|
RetryAfter: time.Duration(retryMilliseconds) * time.Millisecond,
|
|
}, nil
|
|
}
|
|
|
|
func IdentifierHash(identifier string) string {
|
|
digest := sha256.Sum256([]byte(identifier))
|
|
return hex.EncodeToString(digest[:16])
|
|
}
|
|
|
|
const tokenBucketScript = `
|
|
local current_time = redis.call('TIME')
|
|
local now = tonumber(current_time[1]) + tonumber(current_time[2]) / 1000000
|
|
local rate = tonumber(ARGV[1])
|
|
local burst = tonumber(ARGV[2])
|
|
local values = redis.call('HMGET', KEYS[1], 'tokens', 'updated_at')
|
|
local tokens = tonumber(values[1])
|
|
local updated_at = tonumber(values[2])
|
|
if tokens == nil then tokens = burst end
|
|
if updated_at == nil then updated_at = now end
|
|
if now > updated_at then
|
|
tokens = math.min(burst, tokens + (now - updated_at) * rate)
|
|
end
|
|
local allowed = 0
|
|
local retry_ms = 0
|
|
if tokens >= 1 then
|
|
allowed = 1
|
|
tokens = tokens - 1
|
|
else
|
|
retry_ms = math.ceil(((1 - tokens) / rate) * 1000)
|
|
end
|
|
redis.call('HSET', KEYS[1], 'tokens', tokens, 'updated_at', now)
|
|
redis.call('PEXPIRE', KEYS[1], math.ceil((burst / rate) * 2000) + 1000)
|
|
return {allowed, math.floor(tokens), retry_ms}
|
|
`
|