@@ -8,83 +8,147 @@ import (
88 "github.com/gomodule/redigo/redis"
99)
1010
11- // RateLimitConfig holds the configuration for IP+UA rate limiting.
12- type RateLimitConfig struct {
13- // Interval defines the duration of the sliding window.
11+ // LimitStatus indicates the result of a rate limit check
12+ type LimitStatus struct {
13+ // IsLimited indicates if the request should be rate limited
14+ IsLimited bool
15+ // LimitType indicates which limit was exceeded ("ip" or "ipua" or "")
16+ LimitType string
17+ }
18+
19+ // LimitConfig holds the configuration for a single rate limit type
20+ type LimitConfig struct {
21+ // Interval defines the duration of the sliding window
1422 Interval time.Duration
15- // MaxEvents defines the maximum number of events allowed in the interval.
23+
24+ // MaxEvents defines the maximum number of events allowed in the interval
1625 MaxEvents int
17- // KeyPrefix is the prefix for Redis keys.
26+ }
27+
28+ // RateLimitConfig holds the configuration for both IP-only and IP+UA rate limiting.
29+ type RateLimitConfig struct {
30+ // IPConfig defines the rate limiting configuration for IP-only checks
31+ IPConfig LimitConfig
32+
33+ // IPUAConfig defines the rate limiting configuration for IP+UA checks
34+ IPUAConfig LimitConfig
35+
36+ // KeyPrefix is the prefix for Redis keys
1837 KeyPrefix string
1938}
2039
2140// RateLimiter implements a distributed rate limiter using Redis sorted sets (ZSET).
22- // It maintains a sliding window of events for each IP+UA combination , where:
41+ // It maintains sliding windows for both IP-only and IP+UA combinations , where:
2342// - Each event is stored in a ZSET with the timestamp as score
2443// - Old events (outside the window) are automatically removed
2544// - Keys automatically expire after the configured interval
2645//
2746// The limiter considers a request to be rate-limited if the number of events
28- // in the current window exceeds MaxEvents.
47+ // in either window exceeds their respective MaxEvents.
2948type RateLimiter struct {
30- pool * redis.Pool
31- interval time. Duration
32- maxEvents int
33- keyPrefix string
49+ pool * redis.Pool
50+ ipConfig LimitConfig
51+ ipuaConfig LimitConfig
52+ keyPrefix string
3453}
3554
3655// NewRateLimiter creates a new rate limiter.
3756func NewRateLimiter (pool * redis.Pool , config RateLimitConfig ) * RateLimiter {
3857 return & RateLimiter {
39- pool : pool ,
40- interval : config .Interval ,
41- maxEvents : config .MaxEvents ,
42- keyPrefix : config .KeyPrefix ,
58+ pool : pool ,
59+ ipConfig : config .IPConfig ,
60+ ipuaConfig : config .IPUAConfig ,
61+ keyPrefix : config .KeyPrefix ,
4362 }
4463}
4564
46- // generateKey creates a Redis key from IP and User-Agent.
47- func (rl * RateLimiter ) generateKey (ip , ua string ) string {
65+ // generateIPKey creates a Redis key from IP only.
66+ func (rl * RateLimiter ) generateIPKey (ip string ) string {
67+ return fmt .Sprintf ("%s:%s" , rl .keyPrefix , ip )
68+ }
69+
70+ // generateIPUAKey creates a Redis key from IP and User-Agent.
71+ func (rl * RateLimiter ) generateIPUAKey (ip , ua string ) string {
72+ // If User-Agent is empty, use "none" as the value. This allows to distinguish
73+ // between IP-only keys and IPUA keys with an empty User-Agent.
74+ if ua == "" {
75+ ua = "none"
76+ }
4877 return fmt .Sprintf ("%s:%s:%s" , rl .keyPrefix , ip , ua )
4978}
5079
5180// IsLimited checks if the given IP and User-Agent combination should be rate limited.
52- func (rl * RateLimiter ) IsLimited (ip , ua string ) (bool , error ) {
81+ // It first checks the IP-only limit, then the IP+UA limit if the IP-only check passes.
82+ func (rl * RateLimiter ) IsLimited (ip , ua string ) (LimitStatus , error ) {
5383 conn := rl .pool .Get ()
5484 defer conn .Close ()
5585
5686 now := time .Now ().UnixMicro ()
57- windowStart := now - rl .interval .Microseconds ()
58- redisKey := rl .generateKey (ip , ua )
59-
60- // Send all commands in pipeline.
61- // 1. Remove events outside the window
62- conn .Send ("ZREMRANGEBYSCORE" , redisKey , "-inf" , windowStart )
63- // 2. Add current event
64- conn .Send ("ZADD" , redisKey , now , strconv .FormatInt (now , 10 ))
65- // 3. Set key expiration
66- conn .Send ("EXPIRE" , redisKey , int64 (rl .interval .Seconds ()))
67- // 4. Get total event count
68- conn .Send ("ZCARD" , redisKey )
87+ ipKey := rl .generateIPKey (ip )
88+ ipuaKey := rl .generateIPUAKey (ip , ua )
89+
90+ // Start pipeline for both checks
91+ // 1. IP-only check
92+ conn .Send ("ZREMRANGEBYSCORE" , ipKey , "-inf" , now - rl .ipConfig .Interval .Microseconds ())
93+ conn .Send ("ZADD" , ipKey , now , strconv .FormatInt (now , 10 ))
94+ conn .Send ("EXPIRE" , ipKey , int64 (rl .ipConfig .Interval .Seconds ()))
95+ conn .Send ("ZCARD" , ipKey )
96+
97+ // 2. IP+UA limit check
98+ conn .Send ("ZREMRANGEBYSCORE" , ipuaKey , "-inf" , now - rl .ipuaConfig .Interval .Microseconds ())
99+ conn .Send ("ZADD" , ipuaKey , now , strconv .FormatInt (now , 10 ))
100+ conn .Send ("EXPIRE" , ipuaKey , int64 (rl .ipuaConfig .Interval .Seconds ()))
101+ conn .Send ("ZCARD" , ipuaKey )
69102
70103 // Flush pipeline
71104 if err := conn .Flush (); err != nil {
72- return false , fmt .Errorf ("failed to flush pipeline: %w" , err )
105+ return LimitStatus {} , fmt .Errorf ("failed to flush pipeline: %w" , err )
73106 }
74107
75- // Receive all replies
108+ // Receive first 3 replies for IP limit (ZREMRANGEBYSCORE, ZADD, EXPIRE)
76109 for i := 0 ; i < 3 ; i ++ {
77- // Receive replies for ZREMRANGEBYSCORE, ZADD, and EXPIRE
78110 if _ , err := conn .Receive (); err != nil {
79- return false , fmt .Errorf ("failed to receive reply %d: %w" , i , err )
111+ return LimitStatus {} , fmt .Errorf ("failed to receive IP limit reply %d: %w" , i , err )
80112 }
81113 }
82114
83- // Receive and process ZCARD reply
84- count , err := redis .Int64 (conn .Receive ())
115+ // Receive IP limit count
116+ ipCount , err := redis .Int64 (conn .Receive ())
85117 if err != nil {
86- return false , fmt .Errorf ("failed to receive count: %w" , err )
118+ return LimitStatus {}, fmt .Errorf ("failed to receive IP limit count: %w" , err )
119+ }
120+
121+ // Check IP-only limit first
122+ if ipCount > int64 (rl .ipConfig .MaxEvents ) {
123+ return LimitStatus {
124+ IsLimited : true ,
125+ LimitType : "ip" ,
126+ }, nil
127+ }
128+
129+ // Receive next 3 replies for IP+UA limit (ZREMRANGEBYSCORE, ZADD, EXPIRE)
130+ for i := 0 ; i < 3 ; i ++ {
131+ if _ , err := conn .Receive (); err != nil {
132+ return LimitStatus {}, fmt .Errorf ("failed to receive IP+UA limit reply %d: %w" , i , err )
133+ }
134+ }
135+
136+ // Receive IP+UA limit count
137+ ipuaCount , err := redis .Int64 (conn .Receive ())
138+ if err != nil {
139+ return LimitStatus {}, fmt .Errorf ("failed to receive IP+UA limit count: %w" , err )
140+ }
141+
142+ // Check IP+UA limit
143+ if ipuaCount > int64 (rl .ipuaConfig .MaxEvents ) {
144+ return LimitStatus {
145+ IsLimited : true ,
146+ LimitType : "ipua" ,
147+ }, nil
87148 }
88149
89- return count > int64 (rl .maxEvents ), nil
150+ return LimitStatus {
151+ IsLimited : false ,
152+ LimitType : "" ,
153+ }, nil
90154}
0 commit comments