Separate Handshake and Notification message interfaces and other tiny changes (#2696)

* Added checks to avoid decoding handshakes as messages
* Use roles to check if a message is a handshake or not
* break createNotificationsMessageHandle into smaller functions
* fixing IsHandshake for block announce
* Separated Handshake and NotificationMessage interfaces
* remove IsValidHandshake from NotificationsMessageInterface
* renamed IsHandshake to IsValid
* removed Type() from Handshake

Author: kishansagathiya

dot/network/block_announce.go

@@ -15,10 +15,9 @@ import (
 	"github.com/libp2p/go-libp2p-core/peer"
 )
 
-var errInvalidRole = errors.New("invalid role")
 var (
 	_ NotificationsMessage = &BlockAnnounceMessage{}
-	_ NotificationsMessage = &BlockAnnounceHandshake{}
+	_ Handshake            = (*BlockAnnounceHandshake)(nil)
 )
 
 // BlockAnnounceMessage is a state block header
@@ -31,9 +30,9 @@ type BlockAnnounceMessage struct {
 	BestBlock      bool
 }
 
-// Type returns BlockAnnounceMsgType
+// Type returns blockAnnounceMsgType
 func (*BlockAnnounceMessage) Type() byte {
-	return BlockAnnounceMsgType
+	return blockAnnounceMsgType
 }
 
 // string formats a BlockAnnounceMessage as a string
@@ -75,11 +74,6 @@ func (bm *BlockAnnounceMessage) Hash() (common.Hash, error) {
 	return common.Blake2bHash(encMsg)
 }
 
-// IsHandshake returns false
-func (*BlockAnnounceMessage) IsHandshake() bool {
-	return false
-}
-
 func decodeBlockAnnounceHandshake(in []byte) (Handshake, error) {
 	hs := BlockAnnounceHandshake{}
 	err := scale.Unmarshal(in, &hs)
@@ -133,25 +127,14 @@ func (hs *BlockAnnounceHandshake) Decode(in []byte) error {
 	return nil
 }
 
-// Type ...
-func (*BlockAnnounceHandshake) Type() byte {
-	return 0
-}
-
-// Hash returns blake2b hash of block announce handshake.
-func (hs *BlockAnnounceHandshake) Hash() (common.Hash, error) {
-	// scale encode each extrinsic
-	encMsg, err := hs.Encode()
-	if err != nil {
-		return common.Hash{}, fmt.Errorf("cannot encode handshake: %w", err)
+// IsValid returns true if handshakes's role is valid.
+func (hs *BlockAnnounceHandshake) IsValid() bool {
+	switch hs.Roles {
+	case common.AuthorityRole, common.FullNodeRole, common.LightClientRole:
+		return true
+	default:
+		return false
 	}
-
-	return common.Blake2bHash(encMsg)
-}
-
-// IsHandshake returns true
-func (*BlockAnnounceHandshake) IsHandshake() bool {
-	return true
 }
 
 func (s *Service) getBlockAnnounceHandshake() (Handshake, error) {
@@ -188,7 +171,7 @@ func (s *Service) validateBlockAnnounceHandshake(from peer.ID, hs Handshake) err
 		return errors.New("genesis hash mismatch")
 	}
 
-	np, ok := s.notificationsProtocols[BlockAnnounceMsgType]
+	np, ok := s.notificationsProtocols[blockAnnounceMsgType]
 	if !ok {
 		// this should never happen.
 		return nil

dot/network/block_announce_test.go

@@ -158,11 +158,11 @@ func TestValidateBlockAnnounceHandshake(t *testing.T) {
 
 	nodeA := createTestService(t, configA)
 	nodeA.noGossip = true
-	nodeA.notificationsProtocols[BlockAnnounceMsgType] = &notificationsProtocol{
+	nodeA.notificationsProtocols[blockAnnounceMsgType] = &notificationsProtocol{
 		peersData: newPeersData(),
 	}
 	testPeerID := peer.ID("noot")
-	nodeA.notificationsProtocols[BlockAnnounceMsgType].peersData.setInboundHandshakeData(testPeerID, &handshakeData{})
+	nodeA.notificationsProtocols[blockAnnounceMsgType].peersData.setInboundHandshakeData(testPeerID, &handshakeData{})
 
 	err := nodeA.validateBlockAnnounceHandshake(testPeerID, &BlockAnnounceHandshake{
 		Roles:           common.FullNodeRole,

dot/network/errors.go

@@ -10,9 +10,10 @@ import (
 var (
 	errCannotValidateHandshake       = errors.New("failed to validate handshake")
 	errMessageTypeNotValid           = errors.New("message type is not valid")
-	errMessageIsNotHandshake         = errors.New("failed to convert message to Handshake")
 	errInvalidHandshakeForPeer       = errors.New("peer previously sent invalid handshake")
 	errHandshakeTimeout              = errors.New("handshake timeout reached")
 	errBlockRequestFromNumberInvalid = errors.New("block request message From number is not valid")
 	errInvalidStartingBlockType      = errors.New("invalid StartingBlock in messsage")
+	errInboundHanshakeExists         = errors.New("an inbound handshake already exists for given peer")
+	errInvalidRole                   = errors.New("invalid role")
 )

dot/network/gossip.go

@@ -40,9 +40,7 @@ func (g *gossip) hasSeen(msg NotificationsMessage) (bool, error) {
 	_, ok := g.seenMap[msgHash]
 	if !ok {
 		// set message to has been seen
-		if !msg.IsHandshake() {
-			g.seenMap[msgHash] = struct{}{}
-		}
+		g.seenMap[msgHash] = struct{}{}
 		return false, nil
 	}
 

dot/network/host_test.go

@@ -344,7 +344,7 @@ func TestStreamCloseMetadataCleanup(t *testing.T) {
 	_, err = nodeA.host.send(nodeB.host.id(), nodeB.host.protocolID+blockAnnounceID, testHandshake)
 	require.NoError(t, err)
 
-	info := nodeA.notificationsProtocols[BlockAnnounceMsgType]
+	info := nodeA.notificationsProtocols[blockAnnounceMsgType]
 
 	// Set handshake data to received
 	info.peersData.setInboundHandshakeData(nodeB.host.id(), &handshakeData{

dot/network/message.go

@@ -19,8 +19,8 @@ import (
 
 // Message types for notifications protocol messages. Used internally to map message to protocol.
 const (
-	BlockAnnounceMsgType byte = 3
-	TransactionMsgType   byte = 4
+	blockAnnounceMsgType byte = 3
+	transactionMsgType   byte = 4
 	ConsensusMsgType     byte = 5
 )
 
@@ -36,7 +36,6 @@ type NotificationsMessage interface {
 	Message
 	Type() byte
 	Hash() (common.Hash, error)
-	IsHandshake() bool
 }
 
 //nolint:revive
@@ -381,8 +380,3 @@ func (cm *ConsensusMessage) Hash() (common.Hash, error) {
 	}
 	return common.Blake2bHash(encMsg)
 }
-
-// IsHandshake returns false
-func (cm *ConsensusMessage) IsHandshake() bool {
-	return false
-}

dot/network/message_cache.go

@@ -4,7 +4,6 @@
 package network
 
 import (
-	"errors"
 	"fmt"
 	"time"
 
@@ -65,10 +64,6 @@ func (m *messageCache) exists(peer peer.ID, msg NotificationsMessage) bool {
 }
 
 func generateCacheKey(peer peer.ID, msg NotificationsMessage) ([]byte, error) {
-	if msg.IsHandshake() {
-		return nil, errors.New("cache does not support handshake messages")
-	}
-
 	msgHash, err := msg.Hash()
 	if err != nil {
 		return nil, fmt.Errorf("cannot hash notification message: %w", err)

dot/network/message_cache_test.go

@@ -51,30 +51,3 @@ func TestMessageCache(t *testing.T) {
 	ok = msgCache.exists(peerID, msg)
 	require.False(t, ok)
 }
-
-func TestMessageCacheError(t *testing.T) {
-	t.Parallel()
-
-	cacheSize := 64 << 20 // 64 MB
-	msgCache, err := newMessageCache(ristretto.Config{
-		NumCounters: int64(float64(cacheSize) * 0.05 * 2),
-		MaxCost:     int64(float64(cacheSize) * 0.95),
-		BufferItems: 64,
-		Cost: func(value interface{}) int64 {
-			return int64(1)
-		},
-	}, 800*time.Millisecond)
-	require.NoError(t, err)
-
-	peerID := peer.ID("gossamer")
-	msg := &BlockAnnounceHandshake{
-		Roles:           4,
-		BestBlockNumber: 77,
-		BestBlockHash:   common.Hash{1},
-		GenesisHash:     common.Hash{2},
-	}
-
-	ok, err := msgCache.put(peerID, msg)
-	require.Error(t, err, "cache does not support handshake messages")
-	require.False(t, ok)
-}

dot/network/notifications.go

@@ -20,7 +20,8 @@ const handshakeTimeout = time.Second * 10
 
 // Handshake is the interface all handshakes for notifications protocols must implement
 type Handshake interface {
-	NotificationsMessage
+	Message
+	IsValid() bool
 }
 
 // the following are used for RegisterNotificationsProtocol
@@ -91,6 +92,9 @@ func newHandshakeData(received, validated bool, stream network.Stream) *handshak
 	}
 }
 
+// createDecoder combines the notification message decoder and the handshake decoder. The combined
+// decoder decodes using the handshake decoder if we already have handshake data stored for a given
+// peer, otherwise it decodes using the notification message decoder.
 func createDecoder(info *notificationsProtocol, handshakeDecoder HandshakeDecoder,
 	messageDecoder MessageDecoder) messageDecoder {
 	return func(in []byte, peer peer.ID, inbound bool) (Message, error) {
@@ -130,6 +134,18 @@ func (s *Service) createNotificationsMessageHandler(
 			peer = stream.Conn().RemotePeer()
 		)
 
+		hs, ok := m.(Handshake)
+		if ok {
+			if !hs.IsValid() {
+				return errInvalidRole
+			}
+			err := s.handleHandshake(info, stream, hs, peer)
+			if err != nil {
+				return fmt.Errorf("handling handshake: %w", err)
+			}
+			return nil
+		}
+
 		if msg, ok = m.(NotificationsMessage); !ok {
 			return fmt.Errorf("%w: expected %T but got %T", errMessageTypeNotValid, (NotificationsMessage)(nil), msg)
 		}
@@ -148,62 +164,6 @@ func (s *Service) createNotificationsMessageHandler(
 			return nil
 		}
 
-		if msg.IsHandshake() {
-			logger.Tracef("received handshake on notifications sub-protocol %s from peer %s, message is: %s",
-				info.protocolID, stream.Conn().RemotePeer(), msg)
-
-			hs, ok := msg.(Handshake)
-			if !ok {
-				// NOTE: As long as, Handshake interface and NotificationMessage interfaces are same,
-				// this error would never happen.
-				return errMessageIsNotHandshake
-			}
-
-			// if we are the receiver and haven't received the handshake already, validate it
-			// note: if this function is being called, it's being called via SetStreamHandler,
-			// ie it is an inbound stream and we only send the handshake over it.
-			// we do not send any other data over this stream, we would need to open a new outbound stream.
-			hsData := info.peersData.getInboundHandshakeData(peer)
-			if hsData == nil {
-				logger.Tracef("receiver: validating handshake using protocol %s", info.protocolID)
-
-				hsData = newHandshakeData(true, false, stream)
-				info.peersData.setInboundHandshakeData(peer, hsData)
-
-				err := info.handshakeValidator(peer, hs)
-				if err != nil {
-					logger.Tracef(
-						"failed to validate handshake from peer %s using protocol %s: %s",
-						peer, info.protocolID, err)
-					return errCannotValidateHandshake
-				}
-
-				hsData.validated = true
-				info.peersData.setInboundHandshakeData(peer, hsData)
-
-				// once validated, send back a handshake
-				resp, err := info.getHandshake()
-				if err != nil {
-					logger.Warnf("failed to get handshake using protocol %s: %s", info.protocolID, err)
-					return err
-				}
-
-				err = s.host.writeToStream(stream, resp)
-				if err != nil {
-					logger.Tracef("failed to send handshake to peer %s using protocol %s: %s", peer, info.protocolID, err)
-					return err
-				}
-
-				logger.Tracef("receiver: sent handshake to peer %s using protocol %s", peer, info.protocolID)
-
-				if err := stream.CloseWrite(); err != nil {
-					logger.Tracef("failed to close stream for writing: %s", err)
-				}
-			}
-
-			return nil
-		}
-
 		logger.Tracef("received message on notifications sub-protocol %s from peer %s, message is: %s",
 			info.protocolID, stream.Conn().RemotePeer(), msg)
 
@@ -226,6 +186,54 @@ func (s *Service) createNotificationsMessageHandler(
 	}
 }
 
+func (s *Service) handleHandshake(info *notificationsProtocol, stream network.Stream,
+	hs Handshake, peer peer.ID) error {
+	logger.Tracef("received handshake on notifications sub-protocol %s from peer %s, message is: %s",
+		info.protocolID, stream.Conn().RemotePeer(), hs)
+
+	// if we are the receiver and haven't received the handshake already, validate it
+	// note: if this function is being called, it's being called via SetStreamHandler,
+	// ie it is an inbound stream and we only send the handshake over it.
+	// we do not send any other data over this stream, we would need to open a new outbound stream.
+	hsData := info.peersData.getInboundHandshakeData(peer)
+	if hsData != nil {
+		return fmt.Errorf("%w: for peer id %s", errInboundHanshakeExists, peer)
+	}
+
+	logger.Tracef("receiver: validating handshake using protocol %s", info.protocolID)
+
+	hsData = newHandshakeData(true, false, stream)
+	info.peersData.setInboundHandshakeData(peer, hsData)
+
+	err := info.handshakeValidator(peer, hs)
+	if err != nil {
+		return fmt.Errorf("%w from peer %s using protocol %s: %s",
+			errCannotValidateHandshake, peer, info.protocolID, err)
+	}
+
+	hsData.validated = true
+	info.peersData.setInboundHandshakeData(peer, hsData)
+
+	// once validated, send back a handshake
+	resp, err := info.getHandshake()
+	if err != nil {
+		return fmt.Errorf("failed to get handshake using protocol %s: %s", info.protocolID, err)
+	}
+
+	err = s.host.writeToStream(stream, resp)
+	if err != nil {
+		return fmt.Errorf("failed to send handshake to peer %s using protocol %s: %w", peer, info.protocolID, err)
+	}
+
+	logger.Tracef("receiver: sent handshake to peer %s using protocol %s", peer, info.protocolID)
+
+	if err := stream.CloseWrite(); err != nil {
+		return fmt.Errorf("failed to close stream for writing: %s", err)
+	}
+
+	return nil
+}
+
 func closeOutboundStream(info *notificationsProtocol, peerID peer.ID, stream network.Stream) {
 	logger.Debugf(
 		"cleaning up outbound handshake data for protocol=%s, peer=%s",

dot/network/notifications_test.go

@@ -186,7 +186,7 @@ func TestCreateNotificationsMessageHandler_BlockAnnounceHandshake(t *testing.T)
 
 	// try invalid handshake
 	testHandshake := &BlockAnnounceHandshake{
-		Roles:           4,
+		Roles:           common.AuthorityRole,
 		BestBlockNumber: 77,
 		BestBlockHash:   common.Hash{1},
 		// we are using a different genesis here, thus this
@@ -195,7 +195,7 @@ func TestCreateNotificationsMessageHandler_BlockAnnounceHandshake(t *testing.T)
 	}
 
 	err = handler(stream, testHandshake)
-	require.Equal(t, errCannotValidateHandshake, err)
+	require.ErrorIs(t, err, errCannotValidateHandshake)
 	data := info.peersData.getInboundHandshakeData(testPeerID)
 	require.NotNil(t, data)
 	require.True(t, data.received)