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 }