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