github.com/noisysockets/noisysockets@v0.21.2-0.20240515114641-7f467e651c90/internal/ratelimiter/ratelimiter_test.go (about)

     1  // SPDX-License-Identifier: MPL-2.0
     2  /*
     3   * Copyright (C) 2024 The Noisy Sockets Authors.
     4   *
     5   * This Source Code Form is subject to the terms of the Mozilla Public
     6   * License, v. 2.0. If a copy of the MPL was not distributed with this
     7   * file, You can obtain one at http://mozilla.org/MPL/2.0/.
     8   *
     9   * Portions of this file are based on code originally from wireguard-go,
    10   *
    11   * Copyright (C) 2017-2023 WireGuard LLC. All Rights Reserved.
    12   *
    13   * Permission is hereby granted, free of charge, to any person obtaining a copy of
    14   * this software and associated documentation files (the "Software"), to deal in
    15   * the Software without restriction, including without limitation the rights to
    16   * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies
    17   * of the Software, and to permit persons to whom the Software is furnished to do
    18   * so, subject to the following conditions:
    19   *
    20   * The above copyright notice and this permission notice shall be included in all
    21   * copies or substantial portions of the Software.
    22   *
    23   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    24   * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    25   * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    26   * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    27   * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    28   * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    29   * SOFTWARE.
    30   */
    31  
    32  package ratelimiter
    33  
    34  import (
    35  	"net/netip"
    36  	"testing"
    37  	"time"
    38  )
    39  
    40  type result struct {
    41  	allowed bool
    42  	text    string
    43  	wait    time.Duration
    44  }
    45  
    46  func TestRatelimiter(t *testing.T) {
    47  	var rate Ratelimiter
    48  	var expectedResults []result
    49  
    50  	nano := func(nano int64) time.Duration {
    51  		return time.Nanosecond * time.Duration(nano)
    52  	}
    53  
    54  	add := func(res result) {
    55  		expectedResults = append(
    56  			expectedResults,
    57  			res,
    58  		)
    59  	}
    60  
    61  	for i := 0; i < packetsBurstable; i++ {
    62  		add(result{
    63  			allowed: true,
    64  			text:    "initial burst",
    65  		})
    66  	}
    67  
    68  	add(result{
    69  		allowed: false,
    70  		text:    "after burst",
    71  	})
    72  
    73  	add(result{
    74  		allowed: true,
    75  		wait:    nano(time.Second.Nanoseconds() / packetsPerSecond),
    76  		text:    "filling tokens for single packet",
    77  	})
    78  
    79  	add(result{
    80  		allowed: false,
    81  		text:    "not having refilled enough",
    82  	})
    83  
    84  	add(result{
    85  		allowed: true,
    86  		wait:    2 * (nano(time.Second.Nanoseconds() / packetsPerSecond)),
    87  		text:    "filling tokens for two packet burst",
    88  	})
    89  
    90  	add(result{
    91  		allowed: true,
    92  		text:    "second packet in 2 packet burst",
    93  	})
    94  
    95  	add(result{
    96  		allowed: false,
    97  		text:    "packet following 2 packet burst",
    98  	})
    99  
   100  	ips := []netip.Addr{
   101  		netip.MustParseAddr("127.0.0.1"),
   102  		netip.MustParseAddr("192.168.1.1"),
   103  		netip.MustParseAddr("172.167.2.3"),
   104  		netip.MustParseAddr("97.231.252.215"),
   105  		netip.MustParseAddr("248.97.91.167"),
   106  		netip.MustParseAddr("188.208.233.47"),
   107  		netip.MustParseAddr("104.2.183.179"),
   108  		netip.MustParseAddr("72.129.46.120"),
   109  		netip.MustParseAddr("2001:0db8:0a0b:12f0:0000:0000:0000:0001"),
   110  		netip.MustParseAddr("f5c2:818f:c052:655a:9860:b136:6894:25f0"),
   111  		netip.MustParseAddr("b2d7:15ab:48a7:b07c:a541:f144:a9fe:54fc"),
   112  		netip.MustParseAddr("a47b:786e:1671:a22b:d6f9:4ab0:abc7:c918"),
   113  		netip.MustParseAddr("ea1e:d155:7f7a:98fb:2bf5:9483:80f6:5445"),
   114  		netip.MustParseAddr("3f0e:54a2:f5b4:cd19:a21d:58e1:3746:84c4"),
   115  	}
   116  
   117  	now := time.Now()
   118  	rate.timeNow = func() time.Time {
   119  		return now
   120  	}
   121  	defer func() {
   122  		// Lock to avoid data race with cleanup goroutine from Init.
   123  		rate.mu.Lock()
   124  		defer rate.mu.Unlock()
   125  
   126  		rate.timeNow = time.Now
   127  	}()
   128  	timeSleep := func(d time.Duration) {
   129  		now = now.Add(d + 1)
   130  		rate.cleanup()
   131  	}
   132  
   133  	rate.Init()
   134  	t.Cleanup(func() {
   135  		if err := rate.Close(); err != nil {
   136  			t.Fatalf("rate.Close()=%v, want nil", err)
   137  		}
   138  	})
   139  
   140  	for i, res := range expectedResults {
   141  		timeSleep(res.wait)
   142  		for _, ip := range ips {
   143  			allowed := rate.Allow(ip)
   144  			if allowed != res.allowed {
   145  				t.Fatalf("%d: %s: rate.Allow(%q)=%v, want %v", i, res.text, ip, allowed, res.allowed)
   146  			}
   147  		}
   148  	}
   149  }