1501

Committed 10 Feb 2025 06:27PM UTC coverage: 95.581% (-1.4%) from 97.013%

Build # 1501

Build Type

Pull #211

travis-pro

Committed by

robertodauria

Commit Message

Update comment

Pull Request Pull Request #211: Implement rate limiting based on IP+UA

Run Details

43 of 74 new or added lines in 2 files covered. (58.11%)

2 existing lines in 1 file now uncovered.

1925 of 2014 relevant lines covered (95.58%)

1.06 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

90.7

/limits/ratelimiter.go

package limits

import (
        "fmt"
        "strconv"
        "time"

        "github.com/gomodule/redigo/redis"
)

// RateLimitConfig holds the configuration for IP+UA rate limiting.
type RateLimitConfig struct {
        // Interval defines the duration of the sliding window.
        Interval time.Duration
        // MaxEvents defines the maximum number of events allowed in the interval.
        MaxEvents int
        // KeyPrefix is the prefix for Redis keys.
        KeyPrefix string
}

// RateLimiter implements a distributed rate limiter using Redis sorted sets (ZSET).
// It maintains a sliding window of events for each IP+UA combination, where:
//   - Each event is stored in a ZSET with the timestamp as score
//   - Old events (outside the window) are automatically removed
//   - Keys automatically expire after the configured interval
//
// The limiter considers a request to be rate-limited if the number of events
// in the current window exceeds MaxEvents.
type RateLimiter struct {
        pool      *redis.Pool
        interval  time.Duration
        maxEvents int
        keyPrefix string
}

// NewRateLimiter creates a new rate limiter.
func NewRateLimiter(pool *redis.Pool, config RateLimitConfig) *RateLimiter {
        return &RateLimiter{
                pool:      pool,
                interval:  config.Interval,
                maxEvents: config.MaxEvents,
                keyPrefix: config.KeyPrefix,
        }
}

// generateKey creates a Redis key from IP and User-Agent.
func (rl *RateLimiter) generateKey(ip, ua string) string {
        return fmt.Sprintf("%s%s:%s", rl.keyPrefix, ip, ua)
}

// IsLimited checks if the given IP and User-Agent combination should be rate limited.
func (rl *RateLimiter) IsLimited(ip, ua string) (bool, error) {
        conn := rl.pool.Get()
        defer conn.Close()

        now := time.Now().UnixMicro()
        windowStart := now - rl.interval.Microseconds()
        redisKey := rl.generateKey(ip, ua)

        // Send all commands in pipeline.
        // 1. Remove events outside the window
        conn.Send("ZREMRANGEBYSCORE", redisKey, "-inf", windowStart)
        // 2. Add current event
        conn.Send("ZADD", redisKey, now, strconv.FormatInt(now, 10))
        // 3. Set key expiration
        conn.Send("EXPIRE", redisKey, int64(rl.interval.Seconds()))
        // 4. Get total event count
        conn.Send("ZCARD", redisKey)

        // Flush pipeline
        if err := conn.Flush(); err != nil {
                return false, fmt.Errorf("failed to flush pipeline: %w", err)
        }

        // Receive all replies
        for i := 0; i < 3; i++ {
                // Receive replies for ZREMRANGEBYSCORE, ZADD, and EXPIRE
                if _, err := conn.Receive(); err != nil {
                        return false, fmt.Errorf("failed to receive reply %d: %w", i, err)
                }
        }

        // Receive and process ZCARD reply
        count, err := redis.Int64(conn.Receive())
        if err != nil {
                return false, fmt.Errorf("failed to receive count: %w", err)
        }

        return count > int64(rl.maxEvents), nil
}

1	package limits
2
3	import (
4	"fmt"
5	"strconv"
6	"time"
7
8	"github.com/gomodule/redigo/redis"
9	)
10
11	// RateLimitConfig holds the configuration for IP+UA rate limiting.
12	type RateLimitConfig struct {
13	// Interval defines the duration of the sliding window.
14	Interval time.Duration
15	// MaxEvents defines the maximum number of events allowed in the interval.
16	MaxEvents int
17	// KeyPrefix is the prefix for Redis keys.
18	KeyPrefix string
19	}
20
21	// 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:
23	// - Each event is stored in a ZSET with the timestamp as score
24	// - Old events (outside the window) are automatically removed
25	// - Keys automatically expire after the configured interval
26	//
27	// The limiter considers a request to be rate-limited if the number of events
28	// in the current window exceeds MaxEvents.
29	type RateLimiter struct {
30	pool *redis.Pool
31	interval time.Duration
32	maxEvents int
33	keyPrefix string
34	}
35
36	// NewRateLimiter creates a new rate limiter.
37	func NewRateLimiter(pool redis.Pool, config RateLimitConfig) RateLimiter {	1✔
38	return &RateLimiter{	1✔
39	pool: pool,	1✔
40	interval: config.Interval,	1✔
41	maxEvents: config.MaxEvents,	1✔
42	keyPrefix: config.KeyPrefix,	1✔
43	}	1✔
44	}	1✔
45
46	// generateKey creates a Redis key from IP and User-Agent.
47	func (rl *RateLimiter) generateKey(ip, ua string) string {	1✔
48	return fmt.Sprintf("%s%s:%s", rl.keyPrefix, ip, ua)	1✔
49	}	1✔
50
51	// IsLimited checks if the given IP and User-Agent combination should be rate limited.
52	func (rl *RateLimiter) IsLimited(ip, ua string) (bool, error) {	1✔
53	conn := rl.pool.Get()	1✔
54	defer conn.Close()	1✔
55		1✔
56	now := time.Now().UnixMicro()	1✔
57	windowStart := now - rl.interval.Microseconds()	1✔
58	redisKey := rl.generateKey(ip, ua)	1✔
59		1✔
60	// Send all commands in pipeline.	1✔
61	// 1. Remove events outside the window	1✔
62	conn.Send("ZREMRANGEBYSCORE", redisKey, "-inf", windowStart)	1✔
63	// 2. Add current event	1✔
64	conn.Send("ZADD", redisKey, now, strconv.FormatInt(now, 10))	1✔
65	// 3. Set key expiration	1✔
66	conn.Send("EXPIRE", redisKey, int64(rl.interval.Seconds()))	1✔
67	// 4. Get total event count	1✔
68	conn.Send("ZCARD", redisKey)	1✔
69		1✔
70	// Flush pipeline	1✔
71	if err := conn.Flush(); err != nil {	2✔
72	return false, fmt.Errorf("failed to flush pipeline: %w", err)	1✔
73	}	1✔
74
75	// Receive all replies
76	for i := 0; i < 3; i++ {	2✔
77	// Receive replies for ZREMRANGEBYSCORE, ZADD, and EXPIRE	1✔
78	if _, err := conn.Receive(); err != nil {	1✔
NEW 79	return false, fmt.Errorf("failed to receive reply %d: %w", i, err)	×
NEW 80	}	×
81	}
82
83	// Receive and process ZCARD reply
84	count, err := redis.Int64(conn.Receive())	1✔
85	if err != nil {	1✔
NEW 86	return false, fmt.Errorf("failed to receive count: %w", err)	×
NEW 87	}	×
88
89	return count > int64(rl.maxEvents), nil	1✔
90	}

m-lab / locate / 1501

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous