github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/network/p2p/unicast/ratelimit/bandwidth_rate_limiter_test.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"testing"
     5  	"time"
     6  
     7  	"github.com/onflow/flow-go/model/flow"
     8  	libp2pmessage "github.com/onflow/flow-go/model/libp2p/message"
     9  	"github.com/onflow/flow-go/network/message"
    10  
    11  	"github.com/stretchr/testify/require"
    12  	"golang.org/x/time/rate"
    13  
    14  	"github.com/onflow/flow-go/network/channels"
    15  	"github.com/onflow/flow-go/utils/unittest"
    16  )
    17  
    18  // TestBandWidthRateLimiter_Allow ensures rate limiter allows messages as expected.
    19  func TestBandWidthRateLimiter_Allow(t *testing.T) {
    20  	//limiter limit will be set up to 1000 bytes/sec
    21  	limit := rate.Limit(1000)
    22  
    23  	//burst per interval
    24  	burst := 1000
    25  
    26  	// setup bandwidth rate limiter
    27  	bandwidthRateLimiter := NewBandWidthRateLimiter(limit, burst, 1)
    28  
    29  	id, _ := unittest.IdentityWithNetworkingKeyFixture()
    30  	peerID, err := unittest.PeerIDFromFlowID(id)
    31  	require.NoError(t, err)
    32  
    33  	// create message with about 400bytes (300 random bytes + 100bytes message info)
    34  	b := make([]byte, 300)
    35  	for i := range b {
    36  		b[i] = byte('X')
    37  	}
    38  
    39  	sporkId := unittest.IdentifierFixture()
    40  	msg, err := message.NewOutgoingScope(
    41  		flow.IdentifierList{unittest.IdentifierFixture()},
    42  		channels.TopicFromChannel(channels.TestNetworkChannel, sporkId),
    43  		&libp2pmessage.TestMessage{
    44  			Text: string(b),
    45  		},
    46  		unittest.NetworkCodec().Encode,
    47  		message.ProtocolTypeUnicast)
    48  	require.NoError(t, err)
    49  
    50  	allowed := bandwidthRateLimiter.Allow(peerID, msg.Size())
    51  	require.True(t, allowed)
    52  	allowed = bandwidthRateLimiter.Allow(peerID, msg.Size())
    53  	require.True(t, allowed)
    54  	allowed = bandwidthRateLimiter.Allow(peerID, msg.Size())
    55  	require.False(t, allowed)
    56  
    57  	// wait for 1 second, the rate limiter should allow 3 messages again
    58  	time.Sleep(1 * time.Second)
    59  	allowed = bandwidthRateLimiter.Allow(peerID, msg.Size())
    60  	require.True(t, allowed)
    61  	allowed = bandwidthRateLimiter.Allow(peerID, msg.Size())
    62  	require.True(t, allowed)
    63  	allowed = bandwidthRateLimiter.Allow(peerID, msg.Size())
    64  	require.False(t, allowed)
    65  }
    66  
    67  // TestBandWidthRateLimiter_IsRateLimited ensures IsRateLimited returns true for peers that are rate limited.
    68  func TestBandWidthRateLimiter_IsRateLimited(t *testing.T) {
    69  	//limiter limit will be set up to 1000 bytes/sec
    70  	limit := rate.Limit(1000)
    71  
    72  	//burst per interval
    73  	burst := 1000
    74  
    75  	// setup bandwidth rate limiter
    76  	bandwidthRateLimiter := NewBandWidthRateLimiter(limit, burst, time.Second)
    77  
    78  	// for the duration of a simulated second we will send 3 messages. Each message is about
    79  	// 400 bytes, the 3rd message will put our limiter over the 1000 byte limit at 1200 bytes. Thus
    80  	// the 3rd message should be rate limited.
    81  	id, _ := unittest.IdentityWithNetworkingKeyFixture()
    82  	peerID, err := unittest.PeerIDFromFlowID(id)
    83  	require.NoError(t, err)
    84  
    85  	// create message with about 400bytes (300 random bytes + 100bytes message info)
    86  	b := make([]byte, 500000)
    87  	for i := range b {
    88  		b[i] = byte('X')
    89  	}
    90  
    91  	require.False(t, bandwidthRateLimiter.IsRateLimited(peerID))
    92  
    93  	sporkId := unittest.IdentifierFixture()
    94  	msg, err := message.NewOutgoingScope(
    95  		flow.IdentifierList{unittest.IdentifierFixture()},
    96  		channels.TopicFromChannel(channels.TestNetworkChannel, sporkId),
    97  		&libp2pmessage.TestMessage{
    98  			Text: string(b),
    99  		},
   100  		unittest.NetworkCodec().Encode,
   101  		message.ProtocolTypeUnicast)
   102  	require.NoError(t, err)
   103  
   104  	allowed := bandwidthRateLimiter.Allow(peerID, msg.Size())
   105  	require.False(t, allowed)
   106  	require.True(t, bandwidthRateLimiter.IsRateLimited(peerID))
   107  
   108  	// wait for 1 second, the rate limiter should reset
   109  	time.Sleep(1 * time.Second)
   110  	require.False(t, bandwidthRateLimiter.IsRateLimited(peerID))
   111  }