prysm/beacon-chain/sync/rate_limiter_test.go

package sync

import (
	"sync"
	"testing"
	"time"

	"github.com/OffchainLabs/prysm/v7/beacon-chain/p2p"
	mockp2p "github.com/OffchainLabs/prysm/v7/beacon-chain/p2p/testing"
	p2ptypes "github.com/OffchainLabs/prysm/v7/beacon-chain/p2p/types"
	"github.com/OffchainLabs/prysm/v7/testing/assert"
	"github.com/OffchainLabs/prysm/v7/testing/require"
	"github.com/OffchainLabs/prysm/v7/testing/util"
	"github.com/libp2p/go-libp2p/core/network"
	"github.com/libp2p/go-libp2p/core/protocol"
)

func TestNewRateLimiter(t *testing.T) {
	rlimiter := newRateLimiter(mockp2p.NewTestP2P(t))
	assert.Equal(t, len(rlimiter.limiterMap), 20, "correct number of topics not registered")
}

func TestNewRateLimiter_FreeCorrectly(t *testing.T) {
	rlimiter := newRateLimiter(mockp2p.NewTestP2P(t))
	rlimiter.free()
	assert.Equal(t, len(rlimiter.limiterMap), 0, "rate limiter not freed correctly")
}

func TestRateLimiter_ExceedCapacity(t *testing.T) {
	p1 := mockp2p.NewTestP2P(t)
	p2 := mockp2p.NewTestP2P(t)
	p1.Connect(p2)
	rlimiter := newRateLimiter(p1)

	// BlockByRange
	topic := p2p.RPCBlocksByRangeTopicV1 + p1.Encoding().ProtocolSuffix()

	wg := sync.WaitGroup{}
	p2.BHost.SetStreamHandler(protocol.ID(topic), func(stream network.Stream) {
		defer wg.Done()
		code, errMsg, err := readStatusCodeNoDeadline(stream, p2.Encoding())
		require.NoError(t, err, "could not read incoming stream")
		assert.Equal(t, responseCodeInvalidRequest, code, "not equal response codes")
		assert.Equal(t, p2ptypes.ErrRateLimited.Error(), errMsg, "not equal errors")
	})
	wg.Add(1)
	stream, err := p1.BHost.NewStream(t.Context(), p2.PeerID(), protocol.ID(topic))
	require.NoError(t, err, "could not create stream")

	err = rlimiter.validateRequest(stream, 64)
	require.NoError(t, err, "could not validate incoming request")

	// Attempt to create an error, rate limit and lead to disconnect
	err = rlimiter.validateRequest(stream, 1000)
	require.NotNil(t, err, "could not get error from leaky bucket")

	require.NoError(t, stream.Close(), "could not close stream")

	if util.WaitTimeout(&wg, 1*time.Second) {
		t.Fatal("Did not receive stream within 1 sec")
	}
}

func TestRateLimiter_ExceedRawCapacity(t *testing.T) {
	p1 := mockp2p.NewTestP2P(t)
	p2 := mockp2p.NewTestP2P(t)
	p1.Connect(p2)
	p1.Peers().Add(nil, p2.PeerID(), p2.BHost.Addrs()[0], network.DirOutbound)

	rlimiter := newRateLimiter(p1)

	// BlockByRange
	topic := p2p.RPCBlocksByRangeTopicV1 + p1.Encoding().ProtocolSuffix()

	wg := sync.WaitGroup{}
	p2.BHost.SetStreamHandler(protocol.ID(topic), func(stream network.Stream) {
		defer wg.Done()
		code, errMsg, err := readStatusCodeNoDeadline(stream, p2.Encoding())
		require.NoError(t, err, "could not read incoming stream")
		assert.Equal(t, responseCodeInvalidRequest, code, "not equal response codes")
		assert.Equal(t, p2ptypes.ErrRateLimited.Error(), errMsg, "not equal errors")
	})
	wg.Add(1)
	stream, err := p1.BHost.NewStream(t.Context(), p2.PeerID(), protocol.ID(topic))
	require.NoError(t, err, "could not create stream")

	for range 2 * defaultBurstLimit {
		err = rlimiter.validateRawRpcRequest(stream, 1)
		rlimiter.addRawStream(stream)
		require.NoError(t, err, "could not validate incoming request")
	}
	// Triggers rate limit error on burst.
	assert.ErrorContains(t, p2ptypes.ErrRateLimited.Error(), rlimiter.validateRawRpcRequest(stream, 1))

	// Make Peer bad.
	for range defaultBurstLimit {
		assert.ErrorContains(t, p2ptypes.ErrRateLimited.Error(), rlimiter.validateRawRpcRequest(stream, 1))
	}
	assert.NotNil(t, p1.Peers().IsBad(p2.PeerID()), "peer is not marked as a bad peer")
	require.NoError(t, stream.Close(), "could not close stream")

	if util.WaitTimeout(&wg, 1*time.Second) {
		t.Fatal("Did not receive stream within 1 sec")
	}
}

func Test_limiter_retrieveCollector_requiresLock(t *testing.T) {
	l := limiter{}
	_, err := l.retrieveCollector("")
	require.ErrorContains(t, "caller must hold read/write lock", err)
}