feat: add test

Author: fearlessfe

p2p/discover/portal_protocol.go

@@ -1518,15 +1518,35 @@ func (p *PortalProtocol) TraceContentLookup(contentKey []byte) (*TraceContentRes
 
 	traceContentRes := &TraceContentResult{}
 
+	selfHexId := "0x" + p.Self().ID().String()
+
 	trace := &Trace{
-		Origin:      p.Self().ID().String(),
+		Origin:      selfHexId,
 		TargetId:    hexutil.Encode(p.ToContentId(contentKey)),
 		StartedAtMs: int(time.Now().UnixMilli()),
 		Responses:   make(map[string][]string),
 		Metadata:    make(map[string]*NodeMetadata),
 		Cancelled:   make([]string, 0),
 	}
 
+	nodes := p.table.findnodeByID(enode.ID(p.ToContentId(contentKey)), bucketSize, false)
+
+	localResponse := make([]string, 0, len(nodes.entries))
+	for _, node := range nodes.entries {
+		id := "0x" + node.ID().String()
+		localResponse = append(localResponse, id)
+	}
+	trace.Responses[selfHexId] = localResponse
+
+	contentId := p.ToContentId(contentKey)
+
+	dis := p.Distance(p.Self().ID(), enode.ID(contentId))
+
+	trace.Metadata[selfHexId] = &NodeMetadata{
+		Enr:      p.Self().String(),
+		Distance: hexutil.Encode(dis[:]),
+	}
+
 	var wg sync.WaitGroup
 	wg.Add(2)
 	go func() {
@@ -1561,23 +1581,27 @@ func (p *PortalProtocol) TraceContentLookup(contentKey []byte) (*TraceContentRes
 
 	for _, node := range requestNode {
 		id := node.ID().String()
-		trace.Metadata[id] = &NodeMetadata{
-			Enr: node.String(),
-			//Distance: node.Seq(),
+		hexId := "0x" + id
+		dis := p.Distance(node.ID(), enode.ID(contentId))
+		trace.Metadata[hexId] = &NodeMetadata{
+			Enr:      node.String(),
+			Distance: hexutil.Encode(dis[:]),
 		}
 		if res, ok := requestRes[id]; ok {
 			if res.Flag == portalwire.ContentRawSelector || res.Flag == portalwire.ContentConnIdSelector {
-				trace.ReceivedFrom = res.Node.ID().String()
+				trace.ReceivedFrom = hexId
 				content := res.Content.([]byte)
 				traceContentRes.Content = hexutil.Encode(content)
 				traceContentRes.UtpTransfer = res.UtpTransfer
+				trace.Responses[hexId] = make([]string, 0)
 			} else {
 				content := res.Content.([]*enode.Node)
 				ids := make([]string, 0)
 				for _, n := range content {
-					ids = append(ids, n.ID().String())
+					hexId := "0x" + n.ID().String()
+					ids = append(ids, hexId)
 				}
-				trace.Responses[id] = ids
+				trace.Responses[hexId] = ids
 			}
 		} else {
 			trace.Cancelled = append(trace.Cancelled, id)
@@ -1793,6 +1817,14 @@ func (p *PortalProtocol) RandomGossip(srcNodeId *enode.ID, contentKeys [][]byte,
 	return len(nodes), nil
 }
 
+func (p *PortalProtocol) Distance(a, b enode.ID) enode.ID {
+	res := [32]byte{}
+	for i := range a {
+		res[i] = a[i] ^ b[i]
+	}
+	return enode.ID(res)
+}
+
 func inRange(nodeId enode.ID, nodeRadius *uint256.Int, contentId []byte) bool {
 	distance := enode.LogDist(nodeId, enode.ID(contentId))
 	disBig := new(big.Int).SetInt64(int64(distance))

p2p/discover/portal_protocol_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/prysmaticlabs/go-bitfield"
 	"golang.org/x/exp/slices"
 
+	"github.com/ethereum/go-ethereum/common/hexutil"
 	"github.com/ethereum/go-ethereum/internal/testlog"
 	"github.com/ethereum/go-ethereum/log"
 	"github.com/ethereum/go-ethereum/p2p/discover/portalwire"
@@ -410,3 +411,72 @@ func TestContentLookup(t *testing.T) {
 	assert.Equal(t, ContentNotFound, err)
 	assert.Nil(t, res)
 }
+
+func TestTraceContentLookup(t *testing.T) {
+	node1, err := setupLocalPortalNode(":17787", nil)
+	assert.NoError(t, err)
+	node1.log = testlog.Logger(t, log.LvlTrace)
+	err = node1.Start()
+	assert.NoError(t, err)
+
+	node2, err := setupLocalPortalNode(":17788", []*enode.Node{node1.localNode.Node()})
+	assert.NoError(t, err)
+	node2.log = testlog.Logger(t, log.LvlTrace)
+	err = node2.Start()
+	assert.NoError(t, err)
+
+	node3, err := setupLocalPortalNode(":17789", []*enode.Node{node2.localNode.Node()})
+	assert.NoError(t, err)
+	node3.log = testlog.Logger(t, log.LvlTrace)
+	err = node3.Start()
+	assert.NoError(t, err)
+
+	contentKey := []byte{0x3, 0x4}
+	content := []byte{0x1, 0x2}
+	contentId := node1.toContentId(contentKey)
+
+	err = node1.storage.Put(contentId, content)
+	assert.NoError(t, err)
+
+	node1Id := hexutil.Encode(node1.Self().ID().Bytes())
+	node2Id := hexutil.Encode(node2.Self().ID().Bytes())
+	node3Id := hexutil.Encode(node3.Self().ID().Bytes())
+
+	res, err := node3.TraceContentLookup(contentKey)
+	assert.NoError(t, err)
+	assert.Equal(t, res.Content, hexutil.Encode(content))
+	assert.Equal(t, res.UtpTransfer, false)
+	assert.Equal(t, res.Trace.Origin, node3Id)
+	assert.Equal(t, res.Trace.TargetId, hexutil.Encode(contentId))
+	assert.Equal(t, res.Trace.ReceivedFrom, node1Id)
+
+	// check nodeMeta
+	node1Meta := res.Trace.Metadata[node1Id]
+	assert.Equal(t, node1Meta.Enr, node1.Self().String())
+	dis := node1.Distance(node1.Self().ID(), enode.ID(contentId))
+	assert.Equal(t, node1Meta.Distance, hexutil.Encode(dis[:]))
+
+	node2Meta := res.Trace.Metadata[node2Id]
+	assert.Equal(t, node2Meta.Enr, node2.Self().String())
+	dis = node2.Distance(node2.Self().ID(), enode.ID(contentId))
+	assert.Equal(t, node2Meta.Distance, hexutil.Encode(dis[:]))
+
+	node3Meta := res.Trace.Metadata[node3Id]
+	assert.Equal(t, node3Meta.Enr, node3.Self().String())
+	dis = node3.Distance(node3.Self().ID(), enode.ID(contentId))
+	assert.Equal(t, node3Meta.Distance, hexutil.Encode(dis[:]))
+
+	// check response
+	node3Response := res.Trace.Responses[node3Id]
+	assert.Equal(t, node3Response, []string{node2Id})
+
+	node2Response := res.Trace.Responses[node2Id]
+	assert.Equal(t, node2Response, []string{node1Id})
+
+	node1Response := res.Trace.Responses[node1Id]
+	assert.Equal(t, node1Response, []string{})
+
+	// res, _, err = node1.ContentLookup([]byte{0x2, 0x4})
+	// assert.Equal(t, ContentNotFound, err)
+	// assert.Nil(t, res)
+}