diff options
Diffstat (limited to 'ratelimiter/ratelimiter_test.go')
| -rw-r--r-- | ratelimiter/ratelimiter_test.go | 90 |
1 files changed, 53 insertions, 37 deletions
diff --git a/ratelimiter/ratelimiter_test.go b/ratelimiter/ratelimiter_test.go index a18a097..0bfa3af 100644 --- a/ratelimiter/ratelimiter_test.go +++ b/ratelimiter/ratelimiter_test.go @@ -1,32 +1,31 @@ /* SPDX-License-Identifier: MIT * - * Copyright (C) 2017-2019 WireGuard LLC. All Rights Reserved. + * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved. */ package ratelimiter import ( - "net" + "net/netip" "testing" "time" ) -type RatelimiterResult struct { +type result struct { allowed bool text string wait time.Duration } func TestRatelimiter(t *testing.T) { + var rate Ratelimiter + var expectedResults []result - var ratelimiter Ratelimiter - var expectedResults []RatelimiterResult - - Nano := func(nano int64) time.Duration { + nano := func(nano int64) time.Duration { return time.Nanosecond * time.Duration(nano) } - Add := func(res RatelimiterResult) { + add := func(res result) { expectedResults = append( expectedResults, res, @@ -34,69 +33,86 @@ func TestRatelimiter(t *testing.T) { } for i := 0; i < packetsBurstable; i++ { - Add(RatelimiterResult{ + add(result{ allowed: true, - text: "inital burst", + text: "initial burst", }) } - Add(RatelimiterResult{ + add(result{ allowed: false, text: "after burst", }) - Add(RatelimiterResult{ + add(result{ allowed: true, - wait: Nano(time.Second.Nanoseconds() / packetsPerSecond), + wait: nano(time.Second.Nanoseconds() / packetsPerSecond), text: "filling tokens for single packet", }) - Add(RatelimiterResult{ + add(result{ allowed: false, text: "not having refilled enough", }) - Add(RatelimiterResult{ + add(result{ allowed: true, - wait: 2 * (Nano(time.Second.Nanoseconds() / packetsPerSecond)), + wait: 2 * (nano(time.Second.Nanoseconds() / packetsPerSecond)), text: "filling tokens for two packet burst", }) - Add(RatelimiterResult{ + add(result{ allowed: true, text: "second packet in 2 packet burst", }) - Add(RatelimiterResult{ + add(result{ allowed: false, text: "packet following 2 packet burst", }) - ips := []net.IP{ - net.ParseIP("127.0.0.1"), - net.ParseIP("192.168.1.1"), - net.ParseIP("172.167.2.3"), - net.ParseIP("97.231.252.215"), - net.ParseIP("248.97.91.167"), - net.ParseIP("188.208.233.47"), - net.ParseIP("104.2.183.179"), - net.ParseIP("72.129.46.120"), - net.ParseIP("2001:0db8:0a0b:12f0:0000:0000:0000:0001"), - net.ParseIP("f5c2:818f:c052:655a:9860:b136:6894:25f0"), - net.ParseIP("b2d7:15ab:48a7:b07c:a541:f144:a9fe:54fc"), - net.ParseIP("a47b:786e:1671:a22b:d6f9:4ab0:abc7:c918"), - net.ParseIP("ea1e:d155:7f7a:98fb:2bf5:9483:80f6:5445"), - net.ParseIP("3f0e:54a2:f5b4:cd19:a21d:58e1:3746:84c4"), + ips := []netip.Addr{ + netip.MustParseAddr("127.0.0.1"), + netip.MustParseAddr("192.168.1.1"), + netip.MustParseAddr("172.167.2.3"), + netip.MustParseAddr("97.231.252.215"), + netip.MustParseAddr("248.97.91.167"), + netip.MustParseAddr("188.208.233.47"), + netip.MustParseAddr("104.2.183.179"), + netip.MustParseAddr("72.129.46.120"), + netip.MustParseAddr("2001:0db8:0a0b:12f0:0000:0000:0000:0001"), + netip.MustParseAddr("f5c2:818f:c052:655a:9860:b136:6894:25f0"), + netip.MustParseAddr("b2d7:15ab:48a7:b07c:a541:f144:a9fe:54fc"), + netip.MustParseAddr("a47b:786e:1671:a22b:d6f9:4ab0:abc7:c918"), + netip.MustParseAddr("ea1e:d155:7f7a:98fb:2bf5:9483:80f6:5445"), + netip.MustParseAddr("3f0e:54a2:f5b4:cd19:a21d:58e1:3746:84c4"), + } + + now := time.Now() + rate.timeNow = func() time.Time { + return now + } + defer func() { + // Lock to avoid data race with cleanup goroutine from Init. + rate.mu.Lock() + defer rate.mu.Unlock() + + rate.timeNow = time.Now + }() + timeSleep := func(d time.Duration) { + now = now.Add(d + 1) + rate.cleanup() } - ratelimiter.Init() + rate.Init() + defer rate.Close() for i, res := range expectedResults { - time.Sleep(res.wait) + timeSleep(res.wait) for _, ip := range ips { - allowed := ratelimiter.Allow(ip) + allowed := rate.Allow(ip) if allowed != res.allowed { - t.Fatal("Test failed for", ip.String(), ", on:", i, "(", res.text, ")", "expected:", res.allowed, "got:", allowed) + t.Fatalf("%d: %s: rate.Allow(%q)=%v, want %v", i, res.text, ip, allowed, res.allowed) } } } |