github.com/prysmaticlabs/prysm@v1.4.4/beacon-chain/sync/rpc_goodbye_test.go (about)

     1  package sync
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"testing"
     7  	"time"
     8  
     9  	"github.com/kevinms/leakybucket-go"
    10  	"github.com/libp2p/go-libp2p-core/network"
    11  	"github.com/libp2p/go-libp2p-core/protocol"
    12  	types "github.com/prysmaticlabs/eth2-types"
    13  	db "github.com/prysmaticlabs/prysm/beacon-chain/db/testing"
    14  	p2ptest "github.com/prysmaticlabs/prysm/beacon-chain/p2p/testing"
    15  	p2ptypes "github.com/prysmaticlabs/prysm/beacon-chain/p2p/types"
    16  	"github.com/prysmaticlabs/prysm/shared/testutil"
    17  	"github.com/prysmaticlabs/prysm/shared/testutil/assert"
    18  	"github.com/prysmaticlabs/prysm/shared/testutil/require"
    19  )
    20  
    21  func TestGoodByeRPCHandler_Disconnects_With_Peer(t *testing.T) {
    22  	p1 := p2ptest.NewTestP2P(t)
    23  	p2 := p2ptest.NewTestP2P(t)
    24  	p1.Connect(p2)
    25  	assert.Equal(t, 1, len(p1.BHost.Network().Peers()), "Expected peers to be connected")
    26  
    27  	// Set up a head state in the database with data we expect.
    28  	d := db.SetupDB(t)
    29  	r := &Service{
    30  		cfg: &Config{
    31  			DB:  d,
    32  			P2P: p1,
    33  		},
    34  		rateLimiter: newRateLimiter(p1),
    35  	}
    36  
    37  	// Setup streams
    38  	pcl := protocol.ID("/testing")
    39  	topic := string(pcl)
    40  	r.rateLimiter.limiterMap[topic] = leakybucket.NewCollector(1, 1, false)
    41  	var wg sync.WaitGroup
    42  	wg.Add(1)
    43  	p2.BHost.SetStreamHandler(pcl, func(stream network.Stream) {
    44  		defer wg.Done()
    45  		expectResetStream(t, stream)
    46  	})
    47  	stream1, err := p1.BHost.NewStream(context.Background(), p2.BHost.ID(), pcl)
    48  	require.NoError(t, err)
    49  	failureCode := p2ptypes.GoodbyeCodeClientShutdown
    50  
    51  	assert.NoError(t, r.goodbyeRPCHandler(context.Background(), &failureCode, stream1))
    52  
    53  	if testutil.WaitTimeout(&wg, 1*time.Second) {
    54  		t.Fatal("Did not receive stream within 1 sec")
    55  	}
    56  
    57  	conns := p1.BHost.Network().ConnsToPeer(p2.BHost.ID())
    58  	if len(conns) > 0 {
    59  		t.Error("Peer is still not disconnected despite sending a goodbye message")
    60  	}
    61  }
    62  
    63  func TestGoodByeRPCHandler_BackOffPeer(t *testing.T) {
    64  	p1 := p2ptest.NewTestP2P(t)
    65  	p2 := p2ptest.NewTestP2P(t)
    66  	p3 := p2ptest.NewTestP2P(t)
    67  
    68  	p1.Connect(p2)
    69  	p1.Connect(p3)
    70  	assert.Equal(t, 2, len(p1.BHost.Network().Peers()), "Expected peers to be connected")
    71  
    72  	// Set up a head state in the database with data we expect.
    73  	d := db.SetupDB(t)
    74  	r := &Service{
    75  		cfg: &Config{
    76  			DB:  d,
    77  			P2P: p1,
    78  		},
    79  		rateLimiter: newRateLimiter(p1),
    80  	}
    81  
    82  	// Setup streams
    83  	pcl := protocol.ID("/testing")
    84  	topic := string(pcl)
    85  	r.rateLimiter.limiterMap[topic] = leakybucket.NewCollector(1, 1, false)
    86  	var wg sync.WaitGroup
    87  	wg.Add(1)
    88  	p2.BHost.SetStreamHandler(pcl, func(stream network.Stream) {
    89  		defer wg.Done()
    90  		expectResetStream(t, stream)
    91  	})
    92  	stream1, err := p1.BHost.NewStream(context.Background(), p2.BHost.ID(), pcl)
    93  	require.NoError(t, err)
    94  	failureCode := p2ptypes.GoodbyeCodeClientShutdown
    95  
    96  	assert.NoError(t, r.goodbyeRPCHandler(context.Background(), &failureCode, stream1))
    97  
    98  	if testutil.WaitTimeout(&wg, 1*time.Second) {
    99  		t.Fatal("Did not receive stream within 1 sec")
   100  	}
   101  
   102  	conns := p1.BHost.Network().ConnsToPeer(p2.BHost.ID())
   103  	if len(conns) > 0 {
   104  		t.Error("Peer is still not disconnected despite sending a goodbye message")
   105  	}
   106  	valTime, err := p1.Peers().NextValidTime(p2.BHost.ID())
   107  	require.NoError(t, err)
   108  	expectedTime := time.Now().Add(backOffTime[failureCode])
   109  	diff := expectedTime.Sub(valTime)
   110  	// Add a little bit of allowance
   111  	require.Equal(t, true, diff.Seconds() <= 1)
   112  
   113  	wg.Add(1)
   114  	p3.BHost.SetStreamHandler(pcl, func(stream network.Stream) {
   115  		defer wg.Done()
   116  		expectResetStream(t, stream)
   117  	})
   118  
   119  	stream2, err := p1.BHost.NewStream(context.Background(), p3.BHost.ID(), pcl)
   120  	require.NoError(t, err)
   121  	failureCode = p2ptypes.GoodbyeCodeBanned
   122  
   123  	assert.NoError(t, r.goodbyeRPCHandler(context.Background(), &failureCode, stream2))
   124  
   125  	if testutil.WaitTimeout(&wg, 1*time.Second) {
   126  		t.Fatal("Did not receive stream within 1 sec")
   127  	}
   128  
   129  	conns = p1.BHost.Network().ConnsToPeer(p3.BHost.ID())
   130  	if len(conns) > 0 {
   131  		t.Error("Peer is still not disconnected despite sending a goodbye message")
   132  	}
   133  	valTime, err = p1.Peers().NextValidTime(p3.BHost.ID())
   134  	require.NoError(t, err)
   135  	expectedTime = time.Now().Add(backOffTime[failureCode])
   136  	diff = expectedTime.Sub(valTime)
   137  	// Add a little bit of allowance
   138  	require.Equal(t, true, diff.Seconds() <= 1)
   139  }
   140  
   141  func TestSendGoodbye_SendsMessage(t *testing.T) {
   142  	p1 := p2ptest.NewTestP2P(t)
   143  	p2 := p2ptest.NewTestP2P(t)
   144  	p1.Connect(p2)
   145  	assert.Equal(t, 1, len(p1.BHost.Network().Peers()), "Expected peers to be connected")
   146  
   147  	// Set up a head state in the database with data we expect.
   148  	d := db.SetupDB(t)
   149  	r := &Service{
   150  		cfg: &Config{
   151  			DB:  d,
   152  			P2P: p1,
   153  		},
   154  		rateLimiter: newRateLimiter(p1),
   155  	}
   156  	failureCode := p2ptypes.GoodbyeCodeClientShutdown
   157  
   158  	// Setup streams
   159  	pcl := protocol.ID("/eth2/beacon_chain/req/goodbye/1/ssz_snappy")
   160  	topic := string(pcl)
   161  	r.rateLimiter.limiterMap[topic] = leakybucket.NewCollector(1, 1, false)
   162  	var wg sync.WaitGroup
   163  	wg.Add(1)
   164  	p2.BHost.SetStreamHandler(pcl, func(stream network.Stream) {
   165  		defer wg.Done()
   166  		out := new(types.SSZUint64)
   167  		assert.NoError(t, r.cfg.P2P.Encoding().DecodeWithMaxLength(stream, out))
   168  		assert.Equal(t, failureCode, *out)
   169  		assert.NoError(t, stream.Close())
   170  	})
   171  
   172  	err := r.sendGoodByeMessage(context.Background(), failureCode, p2.BHost.ID())
   173  	assert.NoError(t, err)
   174  
   175  	if testutil.WaitTimeout(&wg, 1*time.Second) {
   176  		t.Fatal("Did not receive stream within 1 sec")
   177  	}
   178  
   179  	conns := p1.BHost.Network().ConnsToPeer(p1.BHost.ID())
   180  	if len(conns) > 0 {
   181  		t.Error("Peer is still not disconnected despite sending a goodbye message")
   182  	}
   183  }
   184  
   185  func TestSendGoodbye_DisconnectWithPeer(t *testing.T) {
   186  	p1 := p2ptest.NewTestP2P(t)
   187  	p2 := p2ptest.NewTestP2P(t)
   188  	p1.Connect(p2)
   189  	assert.Equal(t, 1, len(p1.BHost.Network().Peers()), "Expected peers to be connected")
   190  
   191  	// Set up a head state in the database with data we expect.
   192  	d := db.SetupDB(t)
   193  	r := &Service{
   194  		cfg: &Config{
   195  			DB:  d,
   196  			P2P: p1,
   197  		},
   198  		rateLimiter: newRateLimiter(p1),
   199  	}
   200  	failureCode := p2ptypes.GoodbyeCodeClientShutdown
   201  
   202  	// Setup streams
   203  	pcl := protocol.ID("/eth2/beacon_chain/req/goodbye/1/ssz_snappy")
   204  	topic := string(pcl)
   205  	r.rateLimiter.limiterMap[topic] = leakybucket.NewCollector(1, 1, false)
   206  	var wg sync.WaitGroup
   207  	wg.Add(1)
   208  	p2.BHost.SetStreamHandler(pcl, func(stream network.Stream) {
   209  		defer wg.Done()
   210  		out := new(types.SSZUint64)
   211  		assert.NoError(t, r.cfg.P2P.Encoding().DecodeWithMaxLength(stream, out))
   212  		assert.Equal(t, failureCode, *out)
   213  		assert.NoError(t, stream.Close())
   214  	})
   215  
   216  	assert.NoError(t, r.sendGoodByeAndDisconnect(context.Background(), failureCode, p2.BHost.ID()))
   217  	conns := p1.BHost.Network().ConnsToPeer(p2.BHost.ID())
   218  	if len(conns) > 0 {
   219  		t.Error("Peer is still not disconnected despite sending a goodbye message")
   220  	}
   221  
   222  	if testutil.WaitTimeout(&wg, 1*time.Second) {
   223  		t.Fatal("Did not receive stream within 1 sec")
   224  	}
   225  
   226  }