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