cmd/devp2p: implement snap protocol testing

cmd/devp2p/internal/ethtest/chain.go

@@ -26,6 +26,7 @@ import (
 	"os"
 	"strings"
 
+	"github.com/ethereum/go-ethereum/common"
 	"github.com/ethereum/go-ethereum/core"
 	"github.com/ethereum/go-ethereum/core/forkid"
 	"github.com/ethereum/go-ethereum/core/types"
@@ -67,6 +68,13 @@ func (c *Chain) TotalDifficultyAt(height int) *big.Int {
 	return sum
 }
 
+func (c *Chain) RootAt(height int) common.Hash {
+	if height < c.Len() {
+		return c.blocks[height].Root()
+	}
+	return common.Hash{}
+}
+
 // ForkID gets the fork id of the chain.
 func (c *Chain) ForkID() forkid.ID {
 	return forkid.NewID(c.chainConfig, c.blocks[0].Hash(), uint64(c.Len()))

cmd/devp2p/internal/ethtest/helpers.go

@@ -96,6 +96,19 @@ func (s *Suite) dial66() (*Conn, error) {
 	return conn, nil
 }
 
+// dial66 attempts to dial the given node and perform a handshake,
+// returning the created Conn with additional snap/1 capabilities if
+// successful.
+func (s *Suite) dialSnap() (*Conn, error) {
+	conn, err := s.dial66()
+	if err != nil {
+		return nil, fmt.Errorf("dial failed: %v", err)
+	}
+	conn.caps = append(conn.caps, p2p.Cap{Name: "snap", Version: 1})
+	conn.ourHighestSnapProtoVersion = 1
+	return conn, nil
+}
+
 // peer performs both the protocol handshake and the status message
 // exchange with the node in order to peer with it.
 func (c *Conn) peer(chain *Chain, status *Status) error {
@@ -131,7 +144,11 @@ func (c *Conn) handshake() error {
 		}
 		c.negotiateEthProtocol(msg.Caps)
 		if c.negotiatedProtoVersion == 0 {
-			return fmt.Errorf("could not negotiate protocol (remote caps: %v, local eth version: %v)", msg.Caps, c.ourHighestProtoVersion)
+			return fmt.Errorf("could not negotiate eth protocol (remote caps: %v, local eth version: %v)", msg.Caps, c.ourHighestProtoVersion)
+		}
+		// If we require snap, verify that it was negotiated
+		if c.ourHighestSnapProtoVersion != c.negotiatedSnapProtoVersion {
+			return fmt.Errorf("could not negotiate snap protocol (remote caps: %v, local snap version: %v)", msg.Caps, c.ourHighestSnapProtoVersion)
 		}
 		return nil
 	default:
@@ -143,15 +160,21 @@ func (c *Conn) handshake() error {
 // advertised capability from peer.
 func (c *Conn) negotiateEthProtocol(caps []p2p.Cap) {
 	var highestEthVersion uint
+	var highestSnapVersion uint
 	for _, capability := range caps {
-		if capability.Name != "eth" {
-			continue
-		}
-		if capability.Version > highestEthVersion && capability.Version <= c.ourHighestProtoVersion {
-			highestEthVersion = capability.Version
+		switch capability.Name {
+		case "eth":
+			if capability.Version > highestEthVersion && capability.Version <= c.ourHighestProtoVersion {
+				highestEthVersion = capability.Version
+			}
+		case "snap":
+			if capability.Version > highestSnapVersion && capability.Version <= c.ourHighestSnapProtoVersion {
+				highestSnapVersion = capability.Version
+			}
 		}
 	}
 	c.negotiatedProtoVersion = highestEthVersion
+	c.negotiatedSnapProtoVersion = highestSnapVersion
 }
 
 // statusExchange performs a `Status` message exchange with the given node.
@@ -325,6 +348,15 @@ func (c *Conn) headersRequest(request *GetBlockHeaders, chain *Chain, isEth66 bo
 	}
 }
 
+func (c *Conn) snapRequest(msg Message, id uint64, chain *Chain) (Message, error) {
+	defer c.SetReadDeadline(time.Time{})
+	c.SetReadDeadline(time.Now().Add(5 * time.Second))
+	if err := c.Write(msg); err != nil {
+		return nil, fmt.Errorf("could not write to connection: %v", err)
+	}
+	return c.ReadSnap(id)
+}
+
 // getBlockHeaders66 executes the given `GetBlockHeaders` request over the eth66 protocol.
 func getBlockHeaders66(chain *Chain, conn *Conn, request *GetBlockHeaders, id uint64) (BlockHeaders, error) {
 	// write request