github.com/ava-labs/avalanchego@v1.11.11/network/throttling/outbound_msg_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  	"testing"
     8  
     9  	"github.com/prometheus/client_golang/prometheus"
    10  	"github.com/stretchr/testify/require"
    11  	"go.uber.org/mock/gomock"
    12  
    13  	"github.com/ava-labs/avalanchego/ids"
    14  	"github.com/ava-labs/avalanchego/message"
    15  	"github.com/ava-labs/avalanchego/message/messagemock"
    16  	"github.com/ava-labs/avalanchego/snow/validators"
    17  	"github.com/ava-labs/avalanchego/utils/constants"
    18  	"github.com/ava-labs/avalanchego/utils/logging"
    19  )
    20  
    21  func TestSybilOutboundMsgThrottler(t *testing.T) {
    22  	ctrl := gomock.NewController(t)
    23  	require := require.New(t)
    24  	config := MsgByteThrottlerConfig{
    25  		VdrAllocSize:        1024,
    26  		AtLargeAllocSize:    1024,
    27  		NodeMaxAtLargeBytes: 1024,
    28  	}
    29  	vdrs := validators.NewManager()
    30  	vdr1ID := ids.GenerateTestNodeID()
    31  	vdr2ID := ids.GenerateTestNodeID()
    32  	require.NoError(vdrs.AddStaker(constants.PrimaryNetworkID, vdr1ID, nil, ids.Empty, 1))
    33  	require.NoError(vdrs.AddStaker(constants.PrimaryNetworkID, vdr2ID, nil, ids.Empty, 1))
    34  	throttlerIntf, err := NewSybilOutboundMsgThrottler(
    35  		logging.NoLog{},
    36  		prometheus.NewRegistry(),
    37  		vdrs,
    38  		config,
    39  	)
    40  	require.NoError(err)
    41  
    42  	// Make sure NewSybilOutboundMsgThrottler works
    43  	throttler := throttlerIntf.(*outboundMsgThrottler)
    44  	require.Equal(config.VdrAllocSize, throttler.maxVdrBytes)
    45  	require.Equal(config.VdrAllocSize, throttler.remainingVdrBytes)
    46  	require.Equal(config.AtLargeAllocSize, throttler.remainingAtLargeBytes)
    47  	require.NotNil(throttler.nodeToVdrBytesUsed)
    48  	require.NotNil(throttler.log)
    49  	require.NotNil(throttler.vdrs)
    50  
    51  	// Take from at-large allocation.
    52  	msg := testMsgWithSize(ctrl, 1)
    53  	acquired := throttlerIntf.Acquire(msg, vdr1ID)
    54  	require.True(acquired)
    55  	require.Equal(config.AtLargeAllocSize-1, throttler.remainingAtLargeBytes)
    56  	require.Equal(config.VdrAllocSize, throttler.remainingVdrBytes)
    57  	require.Empty(throttler.nodeToVdrBytesUsed)
    58  	require.Len(throttler.nodeToAtLargeBytesUsed, 1)
    59  	require.Equal(uint64(1), throttler.nodeToAtLargeBytesUsed[vdr1ID])
    60  
    61  	// Release the bytes
    62  	throttlerIntf.Release(msg, vdr1ID)
    63  	require.Equal(config.AtLargeAllocSize, throttler.remainingAtLargeBytes)
    64  	require.Equal(config.VdrAllocSize, throttler.remainingVdrBytes)
    65  	require.Empty(throttler.nodeToVdrBytesUsed)
    66  	require.Empty(throttler.nodeToAtLargeBytesUsed)
    67  
    68  	// Use all the at-large allocation bytes and 1 of the validator allocation bytes
    69  	msg = testMsgWithSize(ctrl, config.AtLargeAllocSize+1)
    70  	acquired = throttlerIntf.Acquire(msg, vdr1ID)
    71  	require.True(acquired)
    72  	// vdr1 at-large bytes used: 1024. Validator bytes used: 1
    73  	require.Zero(throttler.remainingAtLargeBytes)
    74  	require.Equal(throttler.remainingVdrBytes, config.VdrAllocSize-1)
    75  	require.Equal(uint64(1), throttler.nodeToVdrBytesUsed[vdr1ID])
    76  	require.Len(throttler.nodeToVdrBytesUsed, 1)
    77  	require.Len(throttler.nodeToAtLargeBytesUsed, 1)
    78  	require.Equal(config.AtLargeAllocSize, throttler.nodeToAtLargeBytesUsed[vdr1ID])
    79  
    80  	// The other validator should be able to acquire half the validator allocation.
    81  	msg = testMsgWithSize(ctrl, config.AtLargeAllocSize/2)
    82  	acquired = throttlerIntf.Acquire(msg, vdr2ID)
    83  	require.True(acquired)
    84  	// vdr2 at-large bytes used: 0. Validator bytes used: 512
    85  	require.Equal(throttler.remainingVdrBytes, config.VdrAllocSize/2-1)
    86  	require.Equal(uint64(1), throttler.nodeToVdrBytesUsed[vdr1ID], 1)
    87  	require.Equal(config.VdrAllocSize/2, throttler.nodeToVdrBytesUsed[vdr2ID])
    88  	require.Len(throttler.nodeToVdrBytesUsed, 2)
    89  	require.Len(throttler.nodeToAtLargeBytesUsed, 1)
    90  
    91  	// vdr1 should be able to acquire the rest of the validator allocation
    92  	msg = testMsgWithSize(ctrl, config.VdrAllocSize/2-1)
    93  	acquired = throttlerIntf.Acquire(msg, vdr1ID)
    94  	require.True(acquired)
    95  	// vdr1 at-large bytes used: 1024. Validator bytes used: 512
    96  	require.Equal(throttler.nodeToVdrBytesUsed[vdr1ID], config.VdrAllocSize/2)
    97  	require.Len(throttler.nodeToAtLargeBytesUsed, 1)
    98  	require.Equal(config.AtLargeAllocSize, throttler.nodeToAtLargeBytesUsed[vdr1ID])
    99  
   100  	// Trying to take more bytes for either node should fail
   101  	msg = testMsgWithSize(ctrl, 1)
   102  	acquired = throttlerIntf.Acquire(msg, vdr1ID)
   103  	require.False(acquired)
   104  	acquired = throttlerIntf.Acquire(msg, vdr2ID)
   105  	require.False(acquired)
   106  	// Should also fail for non-validators
   107  	acquired = throttlerIntf.Acquire(msg, ids.GenerateTestNodeID())
   108  	require.False(acquired)
   109  
   110  	// Release config.MaxAtLargeBytes+1 (1025) bytes
   111  	// When the choice exists, bytes should be given back to the validator allocation
   112  	// rather than the at-large allocation.
   113  	// vdr1 at-large bytes used: 511. Validator bytes used: 0
   114  	msg = testMsgWithSize(ctrl, config.AtLargeAllocSize+1)
   115  	throttlerIntf.Release(msg, vdr1ID)
   116  
   117  	require.Equal(config.NodeMaxAtLargeBytes/2, throttler.remainingVdrBytes)
   118  	require.Len(throttler.nodeToAtLargeBytesUsed, 1) // vdr1
   119  	require.Equal(config.AtLargeAllocSize/2-1, throttler.nodeToAtLargeBytesUsed[vdr1ID])
   120  	require.Len(throttler.nodeToVdrBytesUsed, 1)
   121  	require.Equal(config.AtLargeAllocSize/2+1, throttler.remainingAtLargeBytes)
   122  
   123  	// Non-validator should be able to take the rest of the at-large bytes
   124  	// nonVdrID at-large bytes used: 513
   125  	nonVdrID := ids.GenerateTestNodeID()
   126  	msg = testMsgWithSize(ctrl, config.AtLargeAllocSize/2+1)
   127  	acquired = throttlerIntf.Acquire(msg, nonVdrID)
   128  	require.True(acquired)
   129  	require.Zero(throttler.remainingAtLargeBytes)
   130  	require.Equal(config.AtLargeAllocSize/2+1, throttler.nodeToAtLargeBytesUsed[nonVdrID])
   131  
   132  	// Non-validator shouldn't be able to acquire more since at-large allocation empty
   133  	msg = testMsgWithSize(ctrl, 1)
   134  	acquired = throttlerIntf.Acquire(msg, nonVdrID)
   135  	require.False(acquired)
   136  
   137  	// Release all of vdr2's messages
   138  	msg = testMsgWithSize(ctrl, config.AtLargeAllocSize/2)
   139  	throttlerIntf.Release(msg, vdr2ID)
   140  	require.Zero(throttler.nodeToAtLargeBytesUsed[vdr2ID])
   141  	require.Equal(config.VdrAllocSize, throttler.remainingVdrBytes)
   142  	require.Empty(throttler.nodeToVdrBytesUsed)
   143  	require.Zero(throttler.remainingAtLargeBytes)
   144  
   145  	// Release all of vdr1's messages
   146  	msg = testMsgWithSize(ctrl, config.VdrAllocSize/2-1)
   147  	throttlerIntf.Release(msg, vdr1ID)
   148  	require.Empty(throttler.nodeToVdrBytesUsed)
   149  	require.Equal(config.VdrAllocSize, throttler.remainingVdrBytes)
   150  	require.Equal(config.AtLargeAllocSize/2-1, throttler.remainingAtLargeBytes)
   151  	require.Zero(throttler.nodeToAtLargeBytesUsed[vdr1ID])
   152  
   153  	// Release nonVdr's messages
   154  	msg = testMsgWithSize(ctrl, config.AtLargeAllocSize/2+1)
   155  	throttlerIntf.Release(msg, nonVdrID)
   156  	require.Empty(throttler.nodeToVdrBytesUsed)
   157  	require.Equal(config.VdrAllocSize, throttler.remainingVdrBytes)
   158  	require.Equal(config.AtLargeAllocSize, throttler.remainingAtLargeBytes)
   159  	require.Empty(throttler.nodeToAtLargeBytesUsed)
   160  	require.Zero(throttler.nodeToAtLargeBytesUsed[nonVdrID])
   161  }
   162  
   163  // Ensure that the limit on taking from the at-large allocation is enforced
   164  func TestSybilOutboundMsgThrottlerMaxNonVdr(t *testing.T) {
   165  	ctrl := gomock.NewController(t)
   166  	require := require.New(t)
   167  	config := MsgByteThrottlerConfig{
   168  		VdrAllocSize:        100,
   169  		AtLargeAllocSize:    100,
   170  		NodeMaxAtLargeBytes: 10,
   171  	}
   172  	vdrs := validators.NewManager()
   173  	vdr1ID := ids.GenerateTestNodeID()
   174  	require.NoError(vdrs.AddStaker(constants.PrimaryNetworkID, vdr1ID, nil, ids.Empty, 1))
   175  	throttlerIntf, err := NewSybilOutboundMsgThrottler(
   176  		logging.NoLog{},
   177  		prometheus.NewRegistry(),
   178  		vdrs,
   179  		config,
   180  	)
   181  	require.NoError(err)
   182  	throttler := throttlerIntf.(*outboundMsgThrottler)
   183  	nonVdrNodeID1 := ids.GenerateTestNodeID()
   184  	msg := testMsgWithSize(ctrl, config.NodeMaxAtLargeBytes)
   185  	acquired := throttlerIntf.Acquire(msg, nonVdrNodeID1)
   186  	require.True(acquired)
   187  
   188  	// Acquiring more should fail
   189  	msg = testMsgWithSize(ctrl, 1)
   190  	acquired = throttlerIntf.Acquire(msg, nonVdrNodeID1)
   191  	require.False(acquired)
   192  
   193  	// A different non-validator should be able to acquire
   194  	nonVdrNodeID2 := ids.GenerateTestNodeID()
   195  	msg = testMsgWithSize(ctrl, config.NodeMaxAtLargeBytes)
   196  	acquired = throttlerIntf.Acquire(msg, nonVdrNodeID2)
   197  	require.True(acquired)
   198  
   199  	// Validator should only be able to take [MaxAtLargeBytes]
   200  	msg = testMsgWithSize(ctrl, config.NodeMaxAtLargeBytes+1)
   201  	throttlerIntf.Acquire(msg, vdr1ID)
   202  	require.Equal(config.NodeMaxAtLargeBytes, throttler.nodeToAtLargeBytesUsed[vdr1ID])
   203  	require.Equal(uint64(1), throttler.nodeToVdrBytesUsed[vdr1ID])
   204  	require.Equal(config.NodeMaxAtLargeBytes, throttler.nodeToAtLargeBytesUsed[nonVdrNodeID1])
   205  	require.Equal(config.NodeMaxAtLargeBytes, throttler.nodeToAtLargeBytesUsed[nonVdrNodeID2])
   206  	require.Equal(config.AtLargeAllocSize-config.NodeMaxAtLargeBytes*3, throttler.remainingAtLargeBytes)
   207  }
   208  
   209  // Ensure that the throttler honors requested bypasses
   210  func TestBypassThrottling(t *testing.T) {
   211  	ctrl := gomock.NewController(t)
   212  	require := require.New(t)
   213  	config := MsgByteThrottlerConfig{
   214  		VdrAllocSize:        100,
   215  		AtLargeAllocSize:    100,
   216  		NodeMaxAtLargeBytes: 10,
   217  	}
   218  	vdrs := validators.NewManager()
   219  	vdr1ID := ids.GenerateTestNodeID()
   220  	require.NoError(vdrs.AddStaker(constants.PrimaryNetworkID, vdr1ID, nil, ids.Empty, 1))
   221  	throttlerIntf, err := NewSybilOutboundMsgThrottler(
   222  		logging.NoLog{},
   223  		prometheus.NewRegistry(),
   224  		vdrs,
   225  		config,
   226  	)
   227  	require.NoError(err)
   228  	throttler := throttlerIntf.(*outboundMsgThrottler)
   229  	nonVdrNodeID1 := ids.GenerateTestNodeID()
   230  	msg := messagemock.NewOutboundMessage(ctrl)
   231  	msg.EXPECT().BypassThrottling().Return(true).AnyTimes()
   232  	msg.EXPECT().Op().Return(message.AppGossipOp).AnyTimes()
   233  	msg.EXPECT().Bytes().Return(make([]byte, config.NodeMaxAtLargeBytes)).AnyTimes()
   234  	acquired := throttlerIntf.Acquire(msg, nonVdrNodeID1)
   235  	require.True(acquired)
   236  
   237  	// Acquiring more should not fail
   238  	msg = messagemock.NewOutboundMessage(ctrl)
   239  	msg.EXPECT().BypassThrottling().Return(true).AnyTimes()
   240  	msg.EXPECT().Op().Return(message.AppGossipOp).AnyTimes()
   241  	msg.EXPECT().Bytes().Return(make([]byte, 1)).AnyTimes()
   242  	acquired = throttlerIntf.Acquire(msg, nonVdrNodeID1)
   243  	require.True(acquired)
   244  
   245  	// Acquiring more should not fail
   246  	msg2 := testMsgWithSize(ctrl, 1)
   247  	acquired = throttlerIntf.Acquire(msg2, nonVdrNodeID1)
   248  	require.True(acquired)
   249  
   250  	// Validator should only be able to take [MaxAtLargeBytes]
   251  	msg = messagemock.NewOutboundMessage(ctrl)
   252  	msg.EXPECT().BypassThrottling().Return(true).AnyTimes()
   253  	msg.EXPECT().Op().Return(message.AppGossipOp).AnyTimes()
   254  	msg.EXPECT().Bytes().Return(make([]byte, config.NodeMaxAtLargeBytes+1)).AnyTimes()
   255  	throttlerIntf.Acquire(msg, vdr1ID)
   256  	require.Zero(throttler.nodeToAtLargeBytesUsed[vdr1ID])
   257  	require.Zero(throttler.nodeToVdrBytesUsed[vdr1ID])
   258  	require.Equal(uint64(1), throttler.nodeToAtLargeBytesUsed[nonVdrNodeID1])
   259  	require.Equal(config.AtLargeAllocSize-1, throttler.remainingAtLargeBytes)
   260  }
   261  
   262  func testMsgWithSize(ctrl *gomock.Controller, size uint64) message.OutboundMessage {
   263  	msg := messagemock.NewOutboundMessage(ctrl)
   264  	msg.EXPECT().BypassThrottling().Return(false).AnyTimes()
   265  	msg.EXPECT().Op().Return(message.AppGossipOp).AnyTimes()
   266  	msg.EXPECT().Bytes().Return(make([]byte, size)).AnyTimes()
   267  	return msg
   268  }