1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3package clientlimit_test
4
5import (
6 "context"
7 "fmt"
8 "net/http"
9 "net/netip"
10 "testing"
11 "time"
12
13 "github.com/stretchr/testify/assert"
14 "github.com/stretchr/testify/require"
15
16 "bnema.dev/ermite/internal/clientlimit"
17)
18
19var bg = context.Background()
20
21type fakeClock struct{ t time.Time }
22
23func (c *fakeClock) now() time.Time { return c.t }
24
25func TestLimiter(t *testing.T) {
26 t.Parallel()
27 clk := &fakeClock{t: time.Date(2026, 10, 10, 8, 0, 0, 0, time.UTC)}
28 l := clientlimit.NewLimiter(clk.now)
29
30 for range clientlimit.MaxFailures - 1 {
31 l.Fail(bg, "1.2.3.4")
32 }
33 blocked, _ := l.Blocked(bg, "1.2.3.4")
34 assert.False(t, blocked, "9 failures do not block")
35
36 l.Fail(bg, "1.2.3.4")
37 blocked, wait := l.Blocked(bg, "1.2.3.4")
38 assert.True(t, blocked, "10 failures block")
39 assert.Equal(t, clientlimit.Window, wait)
40
41 other, _ := l.Blocked(bg, "5.6.7.8")
42 assert.False(t, other, "another address is not affected")
43
44 clk.t = clk.t.Add(clientlimit.Window - time.Second)
45 blocked, wait = l.Blocked(bg, "1.2.3.4")
46 assert.True(t, blocked)
47 assert.Equal(t, time.Second, wait)
48
49 clk.t = clk.t.Add(time.Second)
50 blocked, _ = l.Blocked(bg, "1.2.3.4")
51 assert.False(t, blocked, "free once the window passed")
52 assert.Zero(t, l.Tracked(), "expired entries are pruned")
53}
54
55func TestLimiterSlidingWindow(t *testing.T) {
56 t.Parallel()
57 clk := &fakeClock{t: time.Date(2026, 10, 10, 8, 0, 0, 0, time.UTC)}
58 l := clientlimit.NewLimiter(clk.now)
59 for range 5 {
60 l.Fail(bg, "a")
61 }
62 clk.t = clk.t.Add(clientlimit.Window - time.Minute)
63 for range 5 {
64 l.Fail(bg, "a")
65 }
66 blocked, wait := l.Blocked(bg, "a")
67 assert.True(t, blocked)
68 assert.Equal(t, time.Minute, wait, "free when the oldest of the last ten leaves")
69
70 clk.t = clk.t.Add(time.Minute)
71 blocked, _ = l.Blocked(bg, "a")
72 assert.False(t, blocked, "five old failures expired, five remain")
73}
74
75func TestLimiterOneAddressStaysSmall(t *testing.T) {
76 t.Parallel()
77 clk := &fakeClock{t: time.Date(2026, 10, 10, 8, 0, 0, 0, time.UTC)}
78 l := clientlimit.NewLimiter(clk.now)
79 for range 5000 {
80 l.Fail(bg, "x")
81 }
82 blocked, _ := l.Blocked(bg, "x")
83 assert.True(t, blocked)
84 assert.Equal(t, 1, l.Tracked())
85}
86
87func TestLimiterBoundsMemory(t *testing.T) {
88 t.Parallel()
89 clk := &fakeClock{t: time.Date(2026, 10, 10, 8, 0, 0, 0, time.UTC)}
90 l := clientlimit.NewLimiter(clk.now)
91 for i := range clientlimit.MaxTracked + 5000 {
92 clk.t = clk.t.Add(time.Microsecond) // distinct, increasing last failures
93 l.Fail(bg, fmt.Sprintf("10.%d.%d.%d", i>>16&255, i>>8&255, i&255))
94 }
95 assert.LessOrEqual(t, l.Tracked(), clientlimit.MaxTracked)
96
97 // The oldest were evicted, the newest kept.
98 last := clientlimit.MaxTracked + 5000 - 1
99 for range clientlimit.MaxFailures - 1 {
100 l.Fail(bg, fmt.Sprintf("10.%d.%d.%d", last>>16&255, last>>8&255, last&255))
101 }
102 blocked, _ := l.Blocked(bg, fmt.Sprintf("10.%d.%d.%d", last>>16&255, last>>8&255, last&255))
103 assert.True(t, blocked, "a recent address keeps its failures")
104 l.Fail(bg, "10.0.0.0")
105 blocked, _ = l.Blocked(bg, "10.0.0.0")
106 assert.False(t, blocked)
107}
108
109func TestLimiterSweepsExpiredEntriesOncePerInterval(t *testing.T) {
110 t.Parallel()
111 clk := &fakeClock{t: time.Date(2026, 10, 10, 8, 0, 0, 0, time.UTC)}
112 l := clientlimit.NewLimiter(clk.now)
113 for i := range 100 {
114 l.Fail(bg, fmt.Sprintf("192.0.2.%d", i))
115 }
116 require.Equal(t, 100, l.Tracked())
117
118 clk.t = clk.t.Add(clientlimit.Window + time.Second)
119 l.Fail(bg, "198.51.100.1") // first call past the interval: sweeps
120 assert.Equal(t, 1, l.Tracked(), "expired entries are gone")
121}
122
123// The per-failure cost stays flat with a full table: the guard against an
124// O(n) pass per Fail.
125func BenchmarkLimiterFailAtMaxTracked(b *testing.B) {
126 clk := &fakeClock{t: time.Date(2026, 10, 10, 8, 0, 0, 0, time.UTC)}
127 l := clientlimit.NewLimiter(clk.now)
128 keys := make([]string, clientlimit.MaxTracked)
129 for i := range keys {
130 keys[i] = fmt.Sprintf("10.%d.%d.%d", i>>16&255, i>>8&255, i&255)
131 clk.t = clk.t.Add(time.Microsecond)
132 l.Fail(bg, keys[i])
133 }
134 b.ResetTimer()
135 for i := 0; b.Loop(); i++ {
136 clk.t = clk.t.Add(time.Microsecond)
137 l.Fail(bg, fmt.Sprintf("172.%d.%d.%d", i>>16&255, i>>8&255, i&255))
138 l.Blocked(bg, keys[i%len(keys)])
139 }
140}
141
142func TestLimiterGroupsIPv6By64(t *testing.T) {
143 t.Parallel()
144 clk := &fakeClock{t: time.Date(2026, 10, 10, 8, 0, 0, 0, time.UTC)}
145 l := clientlimit.NewLimiter(clk.now)
146
147 // Ten different addresses of one /64 count together.
148 for i := range clientlimit.MaxFailures {
149 l.Fail(bg, fmt.Sprintf("2001:db8:1:2::%x", i+1))
150 }
151 assert.Equal(t, 1, l.Tracked())
152 for _, addr := range []string{"2001:db8:1:2::ffff", "2001:db8:1:2:aaaa:bbbb:cccc:dddd"} {
153 blocked, _ := l.Blocked(bg, addr)
154 assert.True(t, blocked, addr)
155 }
156 // Another /64, and another family, are not affected.
157 for _, addr := range []string{"2001:db8:1:3::1", "198.51.100.1"} {
158 blocked, _ := l.Blocked(bg, addr)
159 assert.False(t, blocked, addr)
160 }
161}
162
163func TestLimiterIPv4StaysPerAddress(t *testing.T) {
164 t.Parallel()
165 clk := &fakeClock{t: time.Date(2026, 10, 10, 8, 0, 0, 0, time.UTC)}
166 l := clientlimit.NewLimiter(clk.now)
167 for range clientlimit.MaxFailures {
168 l.Fail(bg, "203.0.113.5")
169 }
170 blocked, _ := l.Blocked(bg, "203.0.113.5")
171 assert.True(t, blocked)
172 blocked, _ = l.Blocked(bg, "203.0.113.6")
173 assert.False(t, blocked, "a neighbouring IPv4 address is a different client")
174 // IPv4-mapped IPv6 is the IPv4 address.
175 blocked, _ = l.Blocked(bg, "::ffff:203.0.113.5")
176 assert.True(t, blocked)
177}
178
179func TestClientAddr(t *testing.T) {
180 t.Parallel()
181 trusted := []netip.Prefix{netip.MustParsePrefix("10.0.0.0/8"), netip.MustParsePrefix("fd00::/8")}
182 tests := []struct {
183 name string
184 remote string
185 xff []string
186 trusted []netip.Prefix
187 want string
188 }{
189 {"no proxies configured", "203.0.113.5:4000", []string{"1.1.1.1"}, nil, "203.0.113.5"},
190 {"untrusted peer: header ignored", "203.0.113.5:4000", []string{"1.1.1.1"}, trusted, "203.0.113.5"},
191 {"trusted peer, one entry", "10.1.2.3:4000", []string{"198.51.100.7"}, trusted, "198.51.100.7"},
192 {"trusted peer, no header", "10.1.2.3:4000", nil, trusted, "10.1.2.3"},
193 {"right-most untrusted entry wins over a spoofed left one", "10.1.2.3:4000", []string{"6.6.6.6, 198.51.100.7"}, trusted, "198.51.100.7"},
194 {"trusted hops are skipped", "10.1.2.3:4000", []string{"198.51.100.7, 10.9.9.9, 10.8.8.8"}, trusted, "198.51.100.7"},
195 {"several header lines", "10.1.2.3:4000", []string{"6.6.6.6", "198.51.100.7"}, trusted, "198.51.100.7"},
196 {"all entries trusted", "10.1.2.3:4000", []string{"10.9.9.9"}, trusted, "10.1.2.3"},
197 {"garbage entry falls back to the peer", "10.1.2.3:4000", []string{"not-an-ip"}, trusted, "10.1.2.3"},
198 {"garbage left of a good entry is never reached", "10.1.2.3:4000", []string{"junk, 198.51.100.7"}, trusted, "198.51.100.7"},
199 {"ipv6 peer", "[fd00::1]:4000", []string{"2001:db8::7"}, trusted, "2001:db8::7"},
200 {"v4-mapped peer is unmapped", "[::ffff:10.1.2.3]:4000", []string{"198.51.100.7"}, trusted, "198.51.100.7"},
201 {"remote without port", "203.0.113.5", nil, trusted, "203.0.113.5"},
202 }
203 for _, tt := range tests {
204 t.Run(tt.name, func(t *testing.T) {
205 t.Parallel()
206 r := &http.Request{RemoteAddr: tt.remote, Header: http.Header{}}
207 for _, v := range tt.xff {
208 r.Header.Add("X-Forwarded-For", v)
209 }
210 assert.Equal(t, tt.want, clientlimit.ClientAddr(r, tt.trusted))
211 })
212 }
213}