hearthforge

hearthforge/internal/ratelimit/ratelimit_test.go
ratelimit_test.go
⎇
Raw
1package ratelimit
2
3import (
4 "net/http"
5 "testing"
6 "time"
7)
8
9func TestAllowsUpToTheLimitThenBlocks(t *testing.T) {
10 l := New(3, time.Minute)
11 for i := range 3 {
12 if !l.Allow("1.1.1.1") {
13 t.Fatalf("call %d should be allowed", i)
14 }
15 }
16 if l.Allow("1.1.1.1") {
17 t.Fatal("call past the limit should be blocked")
18 }
19}
20
21func TestBucketsArePerKey(t *testing.T) {
22 l := New(1, time.Minute)
23 if !l.Allow("3.3.3.1") || l.Allow("3.3.3.1") {
24 t.Fatal("first key not limited correctly")
25 }
26 if !l.Allow("3.3.3.2") {
27 t.Fatal("second key should have its own bucket")
28 }
29}
30
31func TestLimitersAreIndependent(t *testing.T) {
32 login := New(1, time.Minute)
33 comment := New(1, time.Minute)
34 if !login.Allow("2.2.2.2") || login.Allow("2.2.2.2") {
35 t.Fatal("login limiter not limited correctly")
36 }
37 if !comment.Allow("2.2.2.2") {
38 t.Fatal("comment limiter should be independent")
39 }
40}
41
42func TestWindowResets(t *testing.T) {
43 l := New(1, 20*time.Millisecond)
44 if !l.Allow("4.4.4.4") || l.Allow("4.4.4.4") {
45 t.Fatal("limit not enforced")
46 }
47 time.Sleep(40 * time.Millisecond)
48 if !l.Allow("4.4.4.4") {
49 t.Fatal("window should have reset")
50 }
51}
52
53func TestCountingStaysCorrectWithManyKeys(t *testing.T) {
54 l := New(5, time.Minute)
55 for i := range 1200 {
56 l.Allow("5.5.5.5-" + string(rune('a'+i%26)) + string(rune('a'+i/26)))
57 }
58 for i := range 5 {
59 if !l.Allow("5.5.5.5") {
60 t.Fatalf("call %d should be allowed", i)
61 }
62 }
63 if l.Allow("5.5.5.5") {
64 t.Fatal("call past the limit should be blocked")
65 }
66}
67
68func TestClientIP(t *testing.T) {
69 r, err := http.NewRequest("GET", "http://x/", nil)
70 if err != nil {
71 t.Fatal(err)
72 }
73 r.Header.Set("X-Forwarded-For", "9.9.9.9, 10.0.0.1")
74 r.RemoteAddr = "10.0.0.1:5555"
75
76 if got := ClientIP(r, true); got != "9.9.9.9" {
77 t.Errorf("trusted proxy: got %q", got)
78 }
79 if got := ClientIP(r, false); got != "10.0.0.1" {
80 t.Errorf("untrusted: got %q", got)
81 }
82}
83
84func TestClientIPFallbackAndIPv6Prefix(t *testing.T) {
85 r, err := http.NewRequest("GET", "http://x/", nil)
86 if err != nil {
87 t.Fatal(err)
88 }
89 r.RemoteAddr = "10.0.0.1:5555"
90 if got := ClientIP(r, true); got != "10.0.0.1" {
91 t.Errorf("trusted proxy without header: got %q", got)
92 }
93
94 r.RemoteAddr = "[2001:db8:1:2:aaaa::1]:5555"
95 a := ClientIP(r, false)
96 r.RemoteAddr = "[2001:db8:1:2:bbbb::9]:5555"
97 if b := ClientIP(r, false); a != b || a != "2001:db8:1:2::/64" {
98 t.Errorf("same /64 got %q and %q", a, b)
99 }
100
101 r.RemoteAddr = "[::ffff:10.0.0.2]:5555"
102 if got := ClientIP(r, false); got != "10.0.0.2" {
103 t.Errorf("mapped IPv4: got %q", got)
104 }
105
106 r.RemoteAddr = ""
107 if got := ClientIP(r, true); got == "" {
108 t.Error("key must never be empty")
109 }
110}
111