fix(dot/network): split stored streams and handshakeData into inbound…

… and outbound (#1553)
ChainSafe · May 4, 2021 · 637050b · 637050b
1 parent 88b88f2
commit 637050b
Show file tree

Hide file tree

Showing 11 changed files with 191 additions and 174 deletions.
diff --git a/dot/network/block_announce.go b/dot/network/block_announce.go
@@ -220,18 +220,14 @@ func (s *Service) validateBlockAnnounceHandshake(peer peer.ID, hs Handshake) err
 
 	// don't need to lock here, since function is always called inside the func returned by
 	// `createNotificationsMessageHandler` which locks the map beforehand.
-	data, ok := np.getHandshakeData(peer)
-	if !ok {
-		np.handshakeData.Store(peer, handshakeData{
-			received:  true,
-			validated: true,
-		})
-		data, _ = np.getHandshakeData(peer)
+	data, ok := np.getHandshakeData(peer, true)
+	if ok {
+		data.handshake = hs
+		// TODO: since this is used only for rpc system_peers only,
+		// we can just set the inbound handshake and use that in Peers()
+		np.inboundHandshakeData.Store(peer, data)
 	}
 
-	data.handshake = hs
-	np.handshakeData.Store(peer, data)
-
 	// if peer has higher best block than us, begin syncing
 	latestHeader, err := s.blockState.BestBlockHeader()
 	if err != nil {

diff --git a/dot/network/block_announce_test.go b/dot/network/block_announce_test.go
@@ -117,10 +117,10 @@ func TestValidateBlockAnnounceHandshake(t *testing.T) {
 	nodeA := createTestService(t, configA)
 	nodeA.noGossip = true
 	nodeA.notificationsProtocols[BlockAnnounceMsgType] = &notificationsProtocol{
-		handshakeData: new(sync.Map),
+		inboundHandshakeData: new(sync.Map),
 	}
 	testPeerID := peer.ID("noot")
-	nodeA.notificationsProtocols[BlockAnnounceMsgType].handshakeData.Store(testPeerID, handshakeData{})
+	nodeA.notificationsProtocols[BlockAnnounceMsgType].inboundHandshakeData.Store(testPeerID, handshakeData{})
 
 	err := nodeA.validateBlockAnnounceHandshake(testPeerID, &BlockAnnounceHandshake{
 		BestBlockNumber: 100,

diff --git a/dot/network/host_test.go b/dot/network/host_test.go
@@ -352,13 +352,13 @@ func TestStreamCloseMetadataCleanup(t *testing.T) {
 	info := nodeA.notificationsProtocols[BlockAnnounceMsgType]
 
 	// Set handshake data to received
-	info.handshakeData.Store(nodeB.host.id(), handshakeData{
+	info.inboundHandshakeData.Store(nodeB.host.id(), handshakeData{
 		received:  true,
 		validated: true,
 	})
 
 	// Verify that handshake data exists.
-	_, ok := info.getHandshakeData(nodeB.host.id())
+	_, ok := info.getHandshakeData(nodeB.host.id(), true)
 	require.True(t, ok)
 
 	time.Sleep(time.Second)
@@ -368,7 +368,7 @@ func TestStreamCloseMetadataCleanup(t *testing.T) {
 	time.Sleep(time.Second)
 
 	// Verify that handshake data is cleared.
-	_, ok = info.getHandshakeData(nodeB.host.id())
+	_, ok = info.getHandshakeData(nodeB.host.id(), true)
 	require.False(t, ok)
 }
 

diff --git a/dot/network/light_test.go b/dot/network/light_test.go
@@ -22,7 +22,7 @@ func TestDecodeLightMessage(t *testing.T) {
 	reqEnc, err := testLightRequest.Encode()
 	require.NoError(t, err)
 
-	msg, err := s.decodeLightMessage(reqEnc, testPeer)
+	msg, err := s.decodeLightMessage(reqEnc, testPeer, true)
 	require.NoError(t, err)
 
 	req, ok := msg.(*LightRequest)
@@ -36,7 +36,7 @@ func TestDecodeLightMessage(t *testing.T) {
 	respEnc, err := testLightResponse.Encode()
 	require.NoError(t, err)
 
-	msg, err = s.decodeLightMessage(respEnc, testPeer)
+	msg, err = s.decodeLightMessage(respEnc, testPeer, true)
 	require.NoError(t, err)
 	resp, ok := msg.(*LightResponse)
 	require.True(t, ok)