1package ratelimit
2
3import (
4 "net"
5 "testing"
6 "time"
7)
8
9func TestLimiter(t *testing.T) {
10 l := &Limiter{
11 WindowLimits: []WindowLimit{
12 {
13 Window: time.Minute,
14 Limits: [...]int64{2, 4, 6},
15 },
16 },
17 }
18
19 now := time.Now()
20 check := func(exp bool, ip net.IP, tm time.Time, n int64) {
21 t.Helper()
22 ok := l.CanAdd(ip, tm, n)
23 if ok != exp {
24 t.Fatalf("canadd, got %v, expected %v", ok, exp)
25 }
26 ok = l.Add(ip, tm, n)
27 if ok != exp {
28 t.Fatalf("add, got %v, expected %v", ok, exp)
29 }
30 }
31 check(false, net.ParseIP("10.0.0.1"), now, 3) // past limit
32 check(true, net.ParseIP("10.0.0.1"), now, 1)
33 check(false, net.ParseIP("10.0.0.1"), now, 2) // now past limit
34 check(true, net.ParseIP("10.0.0.1"), now, 1)
35 check(false, net.ParseIP("10.0.0.1"), now, 1) // now past limit
36
37 next := now.Add(time.Minute)
38 check(true, net.ParseIP("10.0.0.1"), next, 2) // next minute, should have reset
39 check(true, net.ParseIP("10.0.0.2"), next, 2) // other ip
40 check(false, net.ParseIP("10.0.0.3"), next, 2) // yet another ip, ipmasked2 was consumed
41 check(true, net.ParseIP("10.0.1.4"), next, 2) // using ipmasked3
42 check(false, net.ParseIP("10.0.2.4"), next, 2) // ipmasked3 consumed
43 l.Reset(net.ParseIP("10.0.1.4"), next)
44 if !l.CanAdd(net.ParseIP("10.0.1.4"), next, 2) {
45 t.Fatalf("reset did not free up count for ip")
46 }
47 check(true, net.ParseIP("10.0.2.4"), next, 2) // ipmasked3 available again
48
49 l = &Limiter{
50 WindowLimits: []WindowLimit{
51 {
52 Window: time.Minute,
53 Limits: [...]int64{1, 2, 3},
54 },
55 {
56 Window: time.Hour,
57 Limits: [...]int64{2, 3, 4},
58 },
59 },
60 }
61
62 min1 := time.UnixMilli((time.Now().UnixNano() / int64(time.Hour)) * int64(time.Hour) / int64(time.Millisecond))
63 min2 := min1.Add(time.Minute)
64 min3 := min1.Add(2 * time.Minute)
65 check(true, net.ParseIP("10.0.0.1"), min1, 1)
66 check(true, net.ParseIP("10.0.0.1"), min2, 1)
67 check(false, net.ParseIP("10.0.0.1"), min3, 1)
68 check(true, net.ParseIP("10.0.0.255"), min3, 1) // ipmasked2 still ok
69 check(false, net.ParseIP("10.0.0.255"), min3, 1) // ipmasked2 also full
70 check(true, net.ParseIP("10.0.1.1"), min3, 1) // ipmasked3 still ok
71 check(false, net.ParseIP("10.0.1.255"), min3, 1) // ipmasked3 also full
72}
73