github.com/ava-labs/avalanchego@v1.11.11/network/throttling/inbound_conn_upgrade_throttler_test.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package throttling
     5  
     6  import (
     7  	"net/netip"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/stretchr/testify/require"
    12  
    13  	"github.com/ava-labs/avalanchego/utils/logging"
    14  )
    15  
    16  var (
    17  	host1      = netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 2, 3, 4}), 9651)
    18  	host2      = netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 2, 3, 5}), 9653)
    19  	host3      = netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 2, 3, 6}), 9655)
    20  	host4      = netip.AddrPortFrom(netip.AddrFrom4([4]byte{1, 2, 3, 7}), 9657)
    21  	loopbackIP = netip.AddrPortFrom(netip.AddrFrom4([4]byte{127, 0, 0, 1}), 9657)
    22  )
    23  
    24  func TestNoInboundConnUpgradeThrottler(t *testing.T) {
    25  	require := require.New(t)
    26  
    27  	{
    28  		throttler := NewInboundConnUpgradeThrottler(
    29  			logging.NoLog{},
    30  			InboundConnUpgradeThrottlerConfig{
    31  				UpgradeCooldown:        0,
    32  				MaxRecentConnsUpgraded: 5,
    33  			},
    34  		)
    35  		// throttler should allow all
    36  		for i := 0; i < 10; i++ {
    37  			require.True(throttler.ShouldUpgrade(host1))
    38  		}
    39  	}
    40  	{
    41  		throttler := NewInboundConnUpgradeThrottler(
    42  			logging.NoLog{},
    43  			InboundConnUpgradeThrottlerConfig{
    44  				UpgradeCooldown:        time.Second,
    45  				MaxRecentConnsUpgraded: 0,
    46  			},
    47  		)
    48  		// throttler should allow all
    49  		for i := 0; i < 10; i++ {
    50  			require.True(throttler.ShouldUpgrade(host1))
    51  		}
    52  	}
    53  }
    54  
    55  func TestInboundConnUpgradeThrottler(t *testing.T) {
    56  	require := require.New(t)
    57  
    58  	cooldown := 5 * time.Second
    59  	throttlerIntf := NewInboundConnUpgradeThrottler(
    60  		logging.NoLog{},
    61  		InboundConnUpgradeThrottlerConfig{
    62  			UpgradeCooldown:        cooldown,
    63  			MaxRecentConnsUpgraded: 3,
    64  		},
    65  	)
    66  
    67  	// Allow should always return true
    68  	// when called with a given IP for the first time
    69  	require.True(throttlerIntf.ShouldUpgrade(host1))
    70  	require.True(throttlerIntf.ShouldUpgrade(host2))
    71  	require.True(throttlerIntf.ShouldUpgrade(host3))
    72  
    73  	// Shouldn't allow this IP because the number of connections
    74  	// within the last [cooldown] is at [MaxRecentConns]
    75  	require.False(throttlerIntf.ShouldUpgrade(host4))
    76  
    77  	// Shouldn't allow these IPs again until [cooldown] has passed
    78  	require.False(throttlerIntf.ShouldUpgrade(host1))
    79  	require.False(throttlerIntf.ShouldUpgrade(host2))
    80  	require.False(throttlerIntf.ShouldUpgrade(host3))
    81  
    82  	// Local host should never be rate-limited
    83  	require.True(throttlerIntf.ShouldUpgrade(loopbackIP))
    84  	require.True(throttlerIntf.ShouldUpgrade(loopbackIP))
    85  	require.True(throttlerIntf.ShouldUpgrade(loopbackIP))
    86  	require.True(throttlerIntf.ShouldUpgrade(loopbackIP))
    87  	require.True(throttlerIntf.ShouldUpgrade(loopbackIP))
    88  
    89  	// Make sure [throttler.done] isn't closed
    90  	throttler := throttlerIntf.(*inboundConnUpgradeThrottler)
    91  	select {
    92  	case <-throttler.done:
    93  		require.FailNow("shouldn't be done")
    94  	default:
    95  	}
    96  
    97  	throttler.Stop()
    98  
    99  	// Make sure [throttler.done] is closed
   100  	select {
   101  	case _, chanOpen := <-throttler.done:
   102  		require.False(chanOpen)
   103  	default:
   104  		require.FailNow("should be done")
   105  	}
   106  }