Skip to content

Commit c52a01e

Browse files
0x46616c6b0x46616c6b
authored
Merge pull request #77 from systemli/configurable-rate-limit-message
✨ Make rate limit rejection message configurable
2 parents 6229500 + 03fd788 commit c52a01e

7 files changed

Lines changed: 61 additions & 20 deletions

File tree

‎.env.example‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ DOMAIN_LISTEN_ADDR=":10002"
66
MAILBOX_LISTEN_ADDR=":10003"
77
SENDERS_LISTEN_ADDR=":10004"
88
METRICS_LISTEN_ADDR=":10005"
9+
RATE_LIMIT_MESSAGE="Rate limit exceeded, please try again later"
910

1011
LOG_LEVEL=debug
1112
LOG_FORMAT=text

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ The adapter is configured via environment variables:
1515
- `SOCKETMAP_LISTEN_ADDR`: The address to listen on for socketmap requests. Default: `:10001`.
1616
- `POLICY_LISTEN_ADDR`: The address to listen on for policy requests (rate limiting). Default: `:10003`.
1717
- `METRICS_LISTEN_ADDR`: The address to listen on for metrics. Default: `:10002`.
18+
- `RATE_LIMIT_MESSAGE`: The rejection message returned when a sender exceeds their quota. Default: `Rate limit exceeded, please try again later`.
1819

1920
In Postfix, you can configure the adapter using the socketmap protocol like this:
2021

‎config.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,9 @@ type Config struct {
2424

2525
// MetricsListenAddr is the address to listen for metrics requests.
2626
MetricsListenAddr string
27+
28+
// RateLimitMessage is the message returned when rate limit is exceeded.
29+
RateLimitMessage string
2730
}
2831

2932
// NewConfig creates a new Config with default values.
@@ -55,12 +58,18 @@ func NewConfig() (*Config, error) {
5558
policyListenAddr = ":10003"
5659
}
5760

61+
rateLimitMessage := os.Getenv("RATE_LIMIT_MESSAGE")
62+
if rateLimitMessage == "" {
63+
rateLimitMessage = "Rate limit exceeded, please try again later"
64+
}
65+
5866
return &Config{
5967
UserliBaseURL: userliBaseURL,
6068
UserliToken: userliToken,
6169
PostfixRecipientDelimiter: postfixRecipientDelimiter,
6270
SocketmapListenAddr: socketmapListenAddr,
6371
PolicyListenAddr: policyListenAddr,
6472
MetricsListenAddr: metricsListenAddr,
73+
RateLimitMessage: rateLimitMessage,
6574
}, nil
6675
}

‎config_test.go‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@ func (s *ConfigTestSuite) TestNewConfig() {
3939
s.Equal("", config.PostfixRecipientDelimiter)
4040
s.Equal(":10001", config.SocketmapListenAddr)
4141
s.Equal(":10002", config.MetricsListenAddr)
42+
s.Equal("Rate limit exceeded, please try again later", config.RateLimitMessage)
4243
})
4344

4445
s.Run("custom config", func() {
@@ -47,6 +48,7 @@ func (s *ConfigTestSuite) TestNewConfig() {
4748
os.Setenv("POSTFIX_RECIPIENT_DELIMITER", "+")
4849
os.Setenv("SOCKETMAP_LISTEN_ADDR", ":20001")
4950
os.Setenv("METRICS_LISTEN_ADDR", ":20002")
51+
os.Setenv("RATE_LIMIT_MESSAGE", "Too many emails")
5052

5153
config, err := NewConfig()
5254

@@ -56,6 +58,7 @@ func (s *ConfigTestSuite) TestNewConfig() {
5658
s.Equal("+", config.PostfixRecipientDelimiter)
5759
s.Equal(":20001", config.SocketmapListenAddr)
5860
s.Equal(":20002", config.MetricsListenAddr)
61+
s.Equal("Too many emails", config.RateLimitMessage)
5962
})
6063
}
6164

‎main.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ func main() {
4848
defer stop()
4949

5050
rateLimiter := NewRateLimiter(ctx)
51-
policyServer := NewPolicyServer(userli, rateLimiter, logger.Named("policy"))
51+
policyServer := NewPolicyServer(userli, rateLimiter, config.RateLimitMessage, logger.Named("policy"))
5252

5353
var wg sync.WaitGroup
5454

‎policy.go‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -16,17 +16,19 @@ import (
1616
// PolicyServer implements a Postfix SMTP Access Policy Delegation server
1717
// for rate limiting outgoing mail based on sender quotas.
1818
type PolicyServer struct {
19-
client UserliService
20-
rateLimiter *RateLimiter
21-
logger *zap.Logger
19+
client UserliService
20+
rateLimiter *RateLimiter
21+
rateLimitMessage string
22+
logger *zap.Logger
2223
}
2324

2425
// NewPolicyServer creates a new PolicyServer with the given UserliService
25-
func NewPolicyServer(client UserliService, rateLimiter *RateLimiter, logger *zap.Logger) *PolicyServer {
26+
func NewPolicyServer(client UserliService, rateLimiter *RateLimiter, rateLimitMessage string, logger *zap.Logger) *PolicyServer {
2627
return &PolicyServer{
27-
client: client,
28-
rateLimiter: rateLimiter,
29-
logger: logger,
28+
client: client,
29+
rateLimiter: rateLimiter,
30+
rateLimitMessage: rateLimitMessage,
31+
logger: logger,
3032
}
3133
}
3234

@@ -235,7 +237,7 @@ func (p *PolicyServer) handleRequest(ctx context.Context, req *PolicyRequest) st
235237
policyRequestDuration.WithLabelValues("check", "reject").Observe(time.Since(startTime).Seconds())
236238
quotaExceededTotal.Inc()
237239

238-
return "REJECT Rate limit exceeded, please try again later"
240+
return "REJECT " + p.rateLimitMessage
239241
}
240242

241243
p.logger.Debug("Message allowed",

‎policy_test.go‎

Lines changed: 36 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -81,7 +81,7 @@ func TestPolicyServer_HandleRequest_SkipNonEndOfMessage(t *testing.T) {
8181
rateLimiter := &RateLimiter{
8282
counters: make(map[string]*senderCounter),
8383
}
84-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
84+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
8585

8686
req := &PolicyRequest{
8787
ProtocolState: "RCPT",
@@ -103,7 +103,7 @@ func TestPolicyServer_HandleRequest_NoSender(t *testing.T) {
103103
rateLimiter := &RateLimiter{
104104
counters: make(map[string]*senderCounter),
105105
}
106-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
106+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
107107

108108
req := &PolicyRequest{
109109
ProtocolState: "END-OF-MESSAGE",
@@ -125,7 +125,7 @@ func TestPolicyServer_HandleRequest_APIError(t *testing.T) {
125125
rateLimiter := &RateLimiter{
126126
counters: make(map[string]*senderCounter),
127127
}
128-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
128+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
129129

130130
req := &PolicyRequest{
131131
ProtocolState: "END-OF-MESSAGE",
@@ -148,7 +148,7 @@ func TestPolicyServer_HandleRequest_NoLimits(t *testing.T) {
148148
rateLimiter := &RateLimiter{
149149
counters: make(map[string]*senderCounter),
150150
}
151-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
151+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
152152

153153
req := &PolicyRequest{
154154
ProtocolState: "END-OF-MESSAGE",
@@ -170,7 +170,7 @@ func TestPolicyServer_HandleRequest_AllowedMessage(t *testing.T) {
170170
rateLimiter := &RateLimiter{
171171
counters: make(map[string]*senderCounter),
172172
}
173-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
173+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
174174

175175
req := &PolicyRequest{
176176
ProtocolState: "END-OF-MESSAGE",
@@ -192,7 +192,7 @@ func TestPolicyServer_HandleRequest_RateLimited(t *testing.T) {
192192
rateLimiter := &RateLimiter{
193193
counters: make(map[string]*senderCounter),
194194
}
195-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
195+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
196196

197197
req := &PolicyRequest{
198198
ProtocolState: "END-OF-MESSAGE",
@@ -218,14 +218,39 @@ func TestPolicyServer_HandleRequest_RateLimited(t *testing.T) {
218218
}
219219
}
220220

221+
func TestPolicyServer_HandleRequest_CustomRateLimitMessage(t *testing.T) {
222+
mockClient := &MockUserliServiceForPolicy{
223+
quota: &Quota{PerHour: 1, PerDay: 100},
224+
}
225+
rateLimiter := &RateLimiter{
226+
counters: make(map[string]*senderCounter),
227+
}
228+
server := NewPolicyServer(mockClient, rateLimiter, "Too many emails", zap.NewNop())
229+
230+
req := &PolicyRequest{
231+
ProtocolState: "END-OF-MESSAGE",
232+
Sender: "user@example.org",
233+
SaslUsername: "user@example.org",
234+
}
235+
236+
// First message passes
237+
server.handleRequest(context.Background(), req)
238+
239+
// Second message should be rejected with custom message
240+
response := server.handleRequest(context.Background(), req)
241+
if response != "REJECT Too many emails" {
242+
t.Errorf("Expected 'REJECT Too many emails', got %s", response)
243+
}
244+
}
245+
221246
func TestPolicyServer_HandleRequest_UsesSaslUsername(t *testing.T) {
222247
mockClient := &MockUserliServiceForPolicy{
223248
quota: &Quota{PerHour: 1, PerDay: 100},
224249
}
225250
rateLimiter := &RateLimiter{
226251
counters: make(map[string]*senderCounter),
227252
}
228-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
253+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
229254

230255
// First request uses sasl_username
231256
req1 := &PolicyRequest{
@@ -255,7 +280,7 @@ func TestPolicyServer_Integration(t *testing.T) {
255280
rateLimiter := &RateLimiter{
256281
counters: make(map[string]*senderCounter),
257282
}
258-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
283+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
259284

260285
// Start server on random port
261286
listener, err := net.Listen("tcp", "127.0.0.1:0")
@@ -317,7 +342,7 @@ func TestPolicyServer_HandleConnection(t *testing.T) {
317342
rateLimiter := &RateLimiter{
318343
counters: make(map[string]*senderCounter),
319344
}
320-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
345+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
321346

322347
serverConn, clientConn := net.Pipe()
323348
defer serverConn.Close()
@@ -367,7 +392,7 @@ func TestPolicyServer_HandleConnection_MultipleRequests(t *testing.T) {
367392
rateLimiter := &RateLimiter{
368393
counters: make(map[string]*senderCounter),
369394
}
370-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
395+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
371396

372397
serverConn, clientConn := net.Pipe()
373398
defer serverConn.Close()
@@ -412,7 +437,7 @@ func TestPolicyServer_StartPolicyServer(t *testing.T) {
412437
rateLimiter := &RateLimiter{
413438
counters: make(map[string]*senderCounter),
414439
}
415-
server := NewPolicyServer(mockClient, rateLimiter, zap.NewNop())
440+
server := NewPolicyServer(mockClient, rateLimiter, "Rate limit exceeded, please try again later", zap.NewNop())
416441

417442
ctx, cancel := context.WithCancel(context.Background())
418443
var wg sync.WaitGroup

0 commit comments

Comments
 (0)