summaryrefslogtreecommitdiffstats
path: root/internal/ratelimit/config.go
blob: 132c96b566b475d25b078be64ec94381ad52e888 (plain)
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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
package ratelimit

import "time"

// Config configures the rate limiter behavior.
type Config struct {
	// RequestsPerSecond is the default rate limit in requests per second.
	// This applies to authenticated users (identified by pubkey).
	// Default: 10
	RequestsPerSecond float64

	// BurstSize is the maximum burst size (token bucket capacity).
	// Allows users to make burst requests up to this limit.
	// Default: 20
	BurstSize int

	// IPRequestsPerSecond is the rate limit for unauthenticated users.
	// These are identified by IP address.
	// Typically set lower than authenticated user limits.
	// Default: 5
	IPRequestsPerSecond float64

	// IPBurstSize is the burst size for IP-based rate limiting.
	// Default: 10
	IPBurstSize int

	// MethodLimits provides per-method rate limit overrides.
	// Key is the full gRPC method name (e.g., "/nostr.v1.NostrRelay/PublishEvent")
	// If not specified, uses the default RequestsPerSecond and BurstSize.
	MethodLimits map[string]MethodLimit

	// UserLimits provides per-user custom rate limits.
	// Key is the pubkey. Useful for VIP/premium users or admins.
	// If not specified, uses the default limits.
	UserLimits map[string]UserLimit

	// SkipMethods is a list of gRPC methods that bypass rate limiting.
	// Useful for health checks or public endpoints.
	// Example: []string{"/grpc.health.v1.Health/Check"}
	SkipMethods []string

	// SkipUsers is a list of pubkeys that bypass rate limiting.
	// Useful for admins or monitoring services.
	SkipUsers []string

	// CleanupInterval is how often to remove idle rate limiters from memory.
	// Limiters that haven't been used recently are removed to save memory.
	// Default: 5 minutes
	CleanupInterval time.Duration

	// MaxIdleTime is how long a limiter can be idle before being cleaned up.
	// Default: 10 minutes
	MaxIdleTime time.Duration
}

// MethodLimit defines rate limits for a specific gRPC method.
type MethodLimit struct {
	RequestsPerSecond float64
	BurstSize         int
}

// UserLimit defines custom rate limits for a specific user (pubkey).
type UserLimit struct {
	// RequestsPerSecond is the default rate for this user.
	RequestsPerSecond float64

	// BurstSize is the burst size for this user.
	BurstSize int

	// MethodLimits provides per-method overrides for this user.
	// Allows fine-grained control like "VIP user gets 100 req/s for queries
	// but still only 5 req/s for publishes"
	MethodLimits map[string]MethodLimit
}

// DefaultConfig returns the default rate limit configuration.
func DefaultConfig() *Config {
	return &Config{
		RequestsPerSecond:   10,
		BurstSize:           20,
		IPRequestsPerSecond: 5,
		IPBurstSize:         10,
		CleanupInterval:     5 * time.Minute,
		MaxIdleTime:         10 * time.Minute,
	}
}

// Validate checks if the configuration is valid.
func (c *Config) Validate() error {
	if c.RequestsPerSecond <= 0 {
		c.RequestsPerSecond = 10
	}
	if c.BurstSize <= 0 {
		c.BurstSize = 20
	}
	if c.IPRequestsPerSecond <= 0 {
		c.IPRequestsPerSecond = 5
	}
	if c.IPBurstSize <= 0 {
		c.IPBurstSize = 10
	}
	if c.CleanupInterval <= 0 {
		c.CleanupInterval = 5 * time.Minute
	}
	if c.MaxIdleTime <= 0 {
		c.MaxIdleTime = 10 * time.Minute
	}
	return nil
}

// GetLimitForMethod returns the rate limit for a specific method and user.
// Precedence: UserLimit.MethodLimit > MethodLimit > UserLimit > Default
func (c *Config) GetLimitForMethod(pubkey, method string) (requestsPerSecond float64, burstSize int) {
	// Check user-specific method limit first (highest precedence)
	if userLimit, ok := c.UserLimits[pubkey]; ok {
		if methodLimit, ok := userLimit.MethodLimits[method]; ok {
			return methodLimit.RequestsPerSecond, methodLimit.BurstSize
		}
	}

	// Check global method limit
	if methodLimit, ok := c.MethodLimits[method]; ok {
		return methodLimit.RequestsPerSecond, methodLimit.BurstSize
	}

	// Check user-specific default limit
	if userLimit, ok := c.UserLimits[pubkey]; ok {
		return userLimit.RequestsPerSecond, userLimit.BurstSize
	}

	// Fall back to global default
	return c.RequestsPerSecond, c.BurstSize
}

// ShouldSkipMethod returns true if the method should bypass rate limiting.
func (c *Config) ShouldSkipMethod(method string) bool {
	for _, skip := range c.SkipMethods {
		if skip == method {
			return true
		}
	}
	return false
}

// ShouldSkipUser returns true if the user should bypass rate limiting.
func (c *Config) ShouldSkipUser(pubkey string) bool {
	for _, skip := range c.SkipUsers {
		if skip == pubkey {
			return true
		}
	}
	return false
}