cmd/devp2p: add snap protocol tests for trie node retrieval

cmd/devp2p/internal/ethtest/snap.go

@@ -186,6 +186,70 @@ func (s *Suite) TestSnapGetByteCodes(t *utesting.T) {
 	}
 }
 
+type trieNodesTest struct {
+	root   common.Hash
+	paths  []snap.TrieNodePathSet
+	nBytes uint64
+
+	expHashes []common.Hash
+	expReject bool
+}
+
+// TestSnapTrieNodes various forms of GetTrieNodes requests.
+func (s *Suite) TestSnapTrieNodes(t *utesting.T) {
+
+	for i, tc := range []trieNodesTest{
+		{
+			root:      s.chain.RootAt(999),
+			paths:     nil,
+			nBytes:    500,
+			expHashes: nil,
+		},
+		{
+			root: s.chain.RootAt(999),
+			paths: []snap.TrieNodePathSet{
+				snap.TrieNodePathSet{}, // zero-length pathset should 'abort' and kick us off
+				snap.TrieNodePathSet{[]byte{0}},
+			},
+			nBytes:    5000,
+			expHashes: []common.Hash{},
+			expReject: true,
+		},
+		{
+			root: s.chain.RootAt(999),
+			paths: []snap.TrieNodePathSet{
+				snap.TrieNodePathSet{[]byte{0}},
+				snap.TrieNodePathSet{[]byte{1}, []byte{0}},
+			},
+			nBytes: 5000,
+			//0x6b3724a41b8c38b46d4d02fba2bb2074c47a507eb16a9a4b978f91d32e406faf
+			expHashes: []common.Hash{s.chain.RootAt(999)},
+		},
+		{ // nonsensically long path
+			root: s.chain.RootAt(999),
+			paths: []snap.TrieNodePathSet{
+				snap.TrieNodePathSet{[]byte{0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 1, 2, 3, 4, 5, 6, 7, 8,
+					0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 1, 2, 3, 4, 5, 6, 7, 8, 0, 1, 2, 3, 4, 5, 6, 7, 8}},
+			},
+			nBytes:    5000,
+			expHashes: []common.Hash{common.HexToHash("0xc5d2460186f7233c927e7db2dcc703c0e500b653ca82273b7bfad8045d85a470")},
+		},
+		{
+			root: s.chain.RootAt(0),
+			paths: []snap.TrieNodePathSet{
+				snap.TrieNodePathSet{[]byte{0}},
+				snap.TrieNodePathSet{[]byte{1}, []byte{0}},
+			},
+			nBytes:    5000,
+			expHashes: []common.Hash{},
+		},
+	} {
+		if err := s.snapGetTrieNodes(t, &tc); err != nil {
+			t.Errorf("test %d \n #hashes %x\n root: %#x\n bytes: %d\nfailed: %v", i, len(tc.expHashes), tc.root, tc.nBytes, err)
+		}
+	}
+}
+
 func (s *Suite) snapGetAccountRange(t *utesting.T, tc *accRangeTest) error {
 	conn, err := s.dialSnap()
 	if err != nil {
@@ -321,3 +385,55 @@ func (s *Suite) snapGetByteCodes(t *utesting.T, tc *byteCodesTest) error {
 
 	return nil
 }
+
+func (s *Suite) snapGetTrieNodes(t *utesting.T, tc *trieNodesTest) error {
+	conn, err := s.dialSnap()
+	if err != nil {
+		t.Fatalf("dial failed: %v", err)
+	}
+	defer conn.Close()
+	if err = conn.peer(s.chain, nil); err != nil {
+		t.Fatalf("peering failed: %v", err)
+	}
+	// write request
+	req := &GetTrieNodes{
+		ID:    uint64(rand.Int63()),
+		Root:  tc.root,
+		Paths: tc.paths,
+		Bytes: tc.nBytes,
+	}
+	resp, err := conn.snapRequest(req, req.ID, s.chain)
+	if err != nil {
+		if tc.expReject {
+			return nil
+		}
+		return fmt.Errorf("trienodes  request failed: %v", err)
+	}
+	var res *snap.TrieNodesPacket
+	if r, ok := resp.(*TrieNodes); !ok {
+		return fmt.Errorf("trienodes response wrong: %T %v", resp, resp)
+	} else {
+		res = (*snap.TrieNodesPacket)(r)
+	}
+
+	// Check the correctness
+
+	// Cross reference the requested trienodes with the response to find gaps
+	// that the serving node is missing
+	hasher := sha3.NewLegacyKeccak256().(crypto.KeccakState)
+	hash := make([]byte, 32)
+	trienodes := res.Nodes
+	if got, want := len(trienodes), len(tc.expHashes); got != want {
+		return fmt.Errorf("wrong trienode count, got %d, want %d\n", got, want)
+	}
+
+	for i, trienode := range trienodes {
+		hasher.Reset()
+		hasher.Write(trienode)
+		hasher.Read(hash)
+		if got, want := hash, tc.expHashes[i]; !bytes.Equal(got, want[:]) {
+			return fmt.Errorf("hash %d wrong, got %#x, want %#x", i, got, want)
+		}
+	}
+	return nil
+}

cmd/devp2p/internal/ethtest/suite.go

@@ -130,6 +130,7 @@ func (s *Suite) SnapTests() []utesting.Test {
 		{Name: "TestSnapStatus", Fn: s.TestSnapStatus},
 		{Name: "TestSnapAccountRange", Fn: s.TestSnapGetAccountRange},
 		{Name: "TestSnapGetByteCodes", Fn: s.TestSnapGetByteCodes},
+		{Name: "TestSnapGetByteCodes", Fn: s.TestSnapTrieNodes},
 	}
 }
 

cmd/devp2p/internal/ethtest/types.go

@@ -308,7 +308,8 @@ func (c *Conn) ReadSnap(id uint64) (Message, error) {
 		case (TrieNodes{}).Code():
 			snpMsg = new(TrieNodes)
 		default:
-			return nil, fmt.Errorf("invalid message code: %d", code)
+			//return nil, fmt.Errorf("invalid message code: %d", code)
+			continue
 		}
 		if err := rlp.DecodeBytes(rawData, snpMsg); err != nil {
 			return nil, fmt.Errorf("could not rlp decode message: %v", err)