github.com/liloew/wireguard-go@v0.0.0-20220224014633-9cd745e6f114/ratelimiter/ratelimiter_test.go (about)

     1  /* SPDX-License-Identifier: MIT
     2   *
     3   * Copyright (C) 2017-2021 WireGuard LLC. All Rights Reserved.
     4   */
     5  
     6  package ratelimiter
     7  
     8  import (
     9  	"testing"
    10  	"time"
    11  
    12  	"golang.zx2c4.com/go118/netip"
    13  )
    14  
    15  type result struct {
    16  	allowed bool
    17  	text    string
    18  	wait    time.Duration
    19  }
    20  
    21  func TestRatelimiter(t *testing.T) {
    22  	var rate Ratelimiter
    23  	var expectedResults []result
    24  
    25  	nano := func(nano int64) time.Duration {
    26  		return time.Nanosecond * time.Duration(nano)
    27  	}
    28  
    29  	add := func(res result) {
    30  		expectedResults = append(
    31  			expectedResults,
    32  			res,
    33  		)
    34  	}
    35  
    36  	for i := 0; i < packetsBurstable; i++ {
    37  		add(result{
    38  			allowed: true,
    39  			text:    "initial burst",
    40  		})
    41  	}
    42  
    43  	add(result{
    44  		allowed: false,
    45  		text:    "after burst",
    46  	})
    47  
    48  	add(result{
    49  		allowed: true,
    50  		wait:    nano(time.Second.Nanoseconds() / packetsPerSecond),
    51  		text:    "filling tokens for single packet",
    52  	})
    53  
    54  	add(result{
    55  		allowed: false,
    56  		text:    "not having refilled enough",
    57  	})
    58  
    59  	add(result{
    60  		allowed: true,
    61  		wait:    2 * (nano(time.Second.Nanoseconds() / packetsPerSecond)),
    62  		text:    "filling tokens for two packet burst",
    63  	})
    64  
    65  	add(result{
    66  		allowed: true,
    67  		text:    "second packet in 2 packet burst",
    68  	})
    69  
    70  	add(result{
    71  		allowed: false,
    72  		text:    "packet following 2 packet burst",
    73  	})
    74  
    75  	ips := []netip.Addr{
    76  		netip.MustParseAddr("127.0.0.1"),
    77  		netip.MustParseAddr("192.168.1.1"),
    78  		netip.MustParseAddr("172.167.2.3"),
    79  		netip.MustParseAddr("97.231.252.215"),
    80  		netip.MustParseAddr("248.97.91.167"),
    81  		netip.MustParseAddr("188.208.233.47"),
    82  		netip.MustParseAddr("104.2.183.179"),
    83  		netip.MustParseAddr("72.129.46.120"),
    84  		netip.MustParseAddr("2001:0db8:0a0b:12f0:0000:0000:0000:0001"),
    85  		netip.MustParseAddr("f5c2:818f:c052:655a:9860:b136:6894:25f0"),
    86  		netip.MustParseAddr("b2d7:15ab:48a7:b07c:a541:f144:a9fe:54fc"),
    87  		netip.MustParseAddr("a47b:786e:1671:a22b:d6f9:4ab0:abc7:c918"),
    88  		netip.MustParseAddr("ea1e:d155:7f7a:98fb:2bf5:9483:80f6:5445"),
    89  		netip.MustParseAddr("3f0e:54a2:f5b4:cd19:a21d:58e1:3746:84c4"),
    90  	}
    91  
    92  	now := time.Now()
    93  	rate.timeNow = func() time.Time {
    94  		return now
    95  	}
    96  	defer func() {
    97  		// Lock to avoid data race with cleanup goroutine from Init.
    98  		rate.mu.Lock()
    99  		defer rate.mu.Unlock()
   100  
   101  		rate.timeNow = time.Now
   102  	}()
   103  	timeSleep := func(d time.Duration) {
   104  		now = now.Add(d + 1)
   105  		rate.cleanup()
   106  	}
   107  
   108  	rate.Init()
   109  	defer rate.Close()
   110  
   111  	for i, res := range expectedResults {
   112  		timeSleep(res.wait)
   113  		for _, ip := range ips {
   114  			allowed := rate.Allow(ip)
   115  			if allowed != res.allowed {
   116  				t.Fatalf("%d: %s: rate.Allow(%q)=%v, want %v", i, res.text, ip, allowed, res.allowed)
   117  			}
   118  		}
   119  	}
   120  }