From 4cc6cde8c6955454a38cf85ae24f5767a7eab1b9 Mon Sep 17 00:00:00 2001 From: fearlseefe <505380967@qq.com> Date: Mon, 5 Feb 2024 16:45:35 +0800 Subject: [PATCH] feat: add test for message --- p2p/discover/api.go | 4 +- p2p/discover/portal_protocol.go | 14 +- p2p/discover/portalwire/messages.go | 18 +-- p2p/discover/portalwire/messages_test.go | 188 +++++++++++++++++++++++ 4 files changed, 203 insertions(+), 21 deletions(-) create mode 100644 p2p/discover/portalwire/messages_test.go diff --git a/p2p/discover/api.go b/p2p/discover/api.go index 0fcc316918..1b14502b65 100644 --- a/p2p/discover/api.go +++ b/p2p/discover/api.go @@ -38,7 +38,7 @@ type DiscV5PongResp struct { } type PortalPongResp struct { - EnrSeq uint64 `json:"enrSeq"` + EnrSeq uint32 `json:"enrSeq"` DataRadius string `json:"dataRadius"` } @@ -327,7 +327,7 @@ func (p *PortalAPI) HistoryPing(enr string) (*PortalPongResp, error) { } return &PortalPongResp{ - EnrSeq: pong.EnrSeq, + EnrSeq: uint32(pong.EnrSeq), DataRadius: nodeRadius.Hex(), }, nil } diff --git a/p2p/discover/portal_protocol.go b/p2p/discover/portal_protocol.go index e6717927c6..2a892a5e13 100644 --- a/p2p/discover/portal_protocol.go +++ b/p2p/discover/portal_protocol.go @@ -84,6 +84,12 @@ var ErrNilContentKey = errors.New("content key cannot be nil") var ContentNotFound = storage.ErrContentNotFound +type clientTag string + +func (c clientTag) ENRKey() string { return "c" } + +const tag clientTag = "shisui" + type ContentElement struct { Node enode.ID ContentKeys [][]byte @@ -184,6 +190,7 @@ func NewPortalProtocol(config *PortalProtocolConfig, protocolId string, privateK localNode := enode.NewLocalNode(nodeDB, privateKey) localNode.SetFallbackIP(net.IP{127, 0, 0, 1}) + localNode.Set(tag) addrs, err := net.InterfaceAddrs() if err != nil { @@ -191,8 +198,7 @@ func NewPortalProtocol(config *PortalProtocolConfig, protocolId string, privateK } for _, address := range addrs { - - // check ip addr is loopback addr + // check ip addr is loopback addr if ipnet, ok := address.(*net.IPNet); ok && !ipnet.IP.IsLoopback() { if ipnet.IP.To4() != nil { localNode.SetStaticIP(ipnet.IP) @@ -396,6 +402,8 @@ func (p *PortalProtocol) pingInner(node *enode.Node) (*portalwire.Pong, error) { return nil, err } + p.log.Trace("Reveice ping response", "source", p.Self().ID(), "target", node.ID(), "res", talkResp) + return p.processPong(node, talkResp) } @@ -755,7 +763,7 @@ func (p *PortalProtocol) handleUtpTalkRequest(id enode.ID, addr *net.UDPAddr, ms } func (p *PortalProtocol) handleTalkRequest(id enode.ID, addr *net.UDPAddr, msg []byte) []byte { - p.log.Error("handleTalkRequest", "id", id, "addr", addr) + p.log.Trace("handleTalkRequest", "id", id, "addr", addr) if n := p.DiscV5.getNode(id); n != nil { p.table.addSeenNode(wrapNode(n)) } diff --git a/p2p/discover/portalwire/messages.go b/p2p/discover/portalwire/messages.go index 3ae38a9334..f996be6f4d 100644 --- a/p2p/discover/portalwire/messages.go +++ b/p2p/discover/portalwire/messages.go @@ -1,5 +1,7 @@ package portalwire +//go:generate sszgen --path p2p/discover/portalwire/messages.go --exclude-objs BlockHeaderProof,PortalReceipts + // Protocol IDs for the portal protocol. const ( StateNetwork = "0x500a" @@ -106,19 +108,3 @@ type ( ContentKeys []byte `ssz:"bitlist" ssz-max:"64"` } ) - -//func getTalkReqOverheadByLen(protocolIdLen int) int { -// return 16 + // IV size -// 55 + // header size -// 1 + // talkReq msg id -// 3 + // rlp encoding outer list, max length will be encoded in 2 bytes -// 9 + // request id (max = 8) + 1 byte from rlp encoding byte string -// protocolIdLen + 1 + // + 1 is necessary due to rlp encoding of byte string -// 3 + // rlp encoding response byte string, max length in 2 bytes -// 16 // HMAC -//} -// -//func getTalkReqOverhead(protocolId string) int { -// protocolIdBytes, _ := hexutil.Decode(protocolId) -// return getTalkReqOverheadByLen(len(protocolIdBytes)) -//} diff --git a/p2p/discover/portalwire/messages_test.go b/p2p/discover/portalwire/messages_test.go new file mode 100644 index 0000000000..c0e5b23efd --- /dev/null +++ b/p2p/discover/portalwire/messages_test.go @@ -0,0 +1,188 @@ +package portalwire + +import ( + "fmt" + "testing" + + "github.com/ethereum/go-ethereum/common/hexutil" + "github.com/ethereum/go-ethereum/p2p/enode" + "github.com/ethereum/go-ethereum/rlp" + ssz "github.com/ferranbt/fastssz" + "github.com/holiman/uint256" + "github.com/stretchr/testify/assert" +) + +var maxUint256 = uint256.MustFromHex("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff") + +// https://github.com/ethereum/portal-network-specs/blob/master/portal-wire-test-vectors.md +// we remove the message type here +func TestPingMessage(t *testing.T) { + dataRadius := maxUint256.Sub(maxUint256, uint256.NewInt(1)) + + reverseBytes := ReverseBytes(dataRadius.Bytes()) + customData := &PingPongCustomData{ + Radius: reverseBytes, + } + dataBytes, err := customData.MarshalSSZ() + assert.NoError(t, err) + ping := &Ping{ + EnrSeq: 1, + CustomPayload: dataBytes, + } + + expected := "0x01000000000000000c000000feffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff" + + data, err := ping.MarshalSSZ() + assert.NoError(t, err) + assert.Equal(t, expected, fmt.Sprintf("0x%x", data)) +} + +func TestPongMessage(t *testing.T) { + dataRadius := maxUint256.Div(maxUint256, uint256.NewInt(2)) + reverseBytes := ReverseBytes(dataRadius.Bytes()) + customData := &PingPongCustomData{ + Radius: reverseBytes, + } + + dataBytes, err := customData.MarshalSSZ() + assert.NoError(t, err) + pong := &Pong{ + EnrSeq: 1, + CustomPayload: dataBytes, + } + + expected := "0x01000000000000000c000000ffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff7f" + + data, err := pong.MarshalSSZ() + assert.NoError(t, err) + assert.Equal(t, expected, fmt.Sprintf("0x%x", data)) +} + +func TestFindNodesMessage(t *testing.T) { + distances := []uint16{256, 255} + + distancesBytes := make([][2]byte, len(distances)) + for i, distance := range distances { + copy(distancesBytes[i][:], ssz.MarshalUint16(make([]byte, 0), distance)) + } + + findNode := &FindNodes{ + Distances: distancesBytes, + } + + data, err := findNode.MarshalSSZ() + expected := "0x040000000001ff00" + assert.NoError(t, err) + assert.Equal(t, expected, fmt.Sprintf("0x%x", data)) +} + +func TestNodes(t *testing.T) { + enrs := []string{ + "enr:-HW4QBzimRxkmT18hMKaAL3IcZF1UcfTMPyi3Q1pxwZZbcZVRI8DC5infUAB_UauARLOJtYTxaagKoGmIjzQxO2qUygBgmlkgnY0iXNlY3AyNTZrMaEDymNMrg1JrLQB2KTGtv6MVbcNEVv0AHacwUAPMljNMTg", + "enr:-HW4QNfxw543Ypf4HXKXdYxkyzfcxcO-6p9X986WldfVpnVTQX1xlTnWrktEWUbeTZnmgOuAY_KUhbVV1Ft98WoYUBMBgmlkgnY0iXNlY3AyNTZrMaEDDiy3QkHAxPyOgWbxp5oF1bDdlYE6dLCUUp8xfVw50jU", + } + + enrsBytes := make([][]byte, 0) + for _, enr := range enrs { + n, err := enode.Parse(enode.ValidSchemes, enr) + assert.NoError(t, err) + + enrBytes, err := rlp.EncodeToBytes(n.Record()) + assert.NoError(t, err) + enrsBytes = append(enrsBytes, enrBytes) + } + + testCases := []struct { + name string + input [][]byte + expected string + }{ + { + name: "empty nodes", + input: make([][]byte, 0), + expected: "0x0105000000", + }, + { + name: "two nodes", + input: enrsBytes, + expected: "0x0105000000080000007f000000f875b8401ce2991c64993d7c84c29a00bdc871917551c7d330fca2dd0d69c706596dc655448f030b98a77d4001fd46ae0112ce26d613c5a6a02a81a6223cd0c4edaa53280182696482763489736563703235366b31a103ca634cae0d49acb401d8a4c6b6fe8c55b70d115bf400769cc1400f3258cd3138f875b840d7f1c39e376297f81d7297758c64cb37dcc5c3beea9f57f7ce9695d7d5a67553417d719539d6ae4b445946de4d99e680eb8063f29485b555d45b7df16a1850130182696482763489736563703235366b31a1030e2cb74241c0c4fc8e8166f1a79a05d5b0dd95813a74b094529f317d5c39d235", + }, + } + + for _, test := range testCases { + t.Run(test.name, func(t *testing.T) { + nodes := &Nodes{ + Total: 1, + Enrs: test.input, + } + + data, err := nodes.MarshalSSZ() + assert.NoError(t, err) + assert.Equal(t, test.expected, fmt.Sprintf("0x%x", data)) + }) + } +} + +func TestContent(t *testing.T) { + contentKey := "0x706f7274616c" + + content := &FindContent{ + ContentKey: hexutil.MustDecode(contentKey), + } + expected := "0x04000000706f7274616c" + data, err := content.MarshalSSZ() + assert.NoError(t, err) + assert.Equal(t, expected, fmt.Sprintf("0x%x", data)) + + // TODO content response is Union type + + // idBuffer := make([]byte, 0, 2) + // idBuffer = append(idBuffer, 0x01) + // idBuffer = append(idBuffer, 0x02) + // // binary.BigEndian.PutUint16() + // // binary.BigEndian.PutUint16(idBuffer, uint16(0x0102)) + // cIds := &ConnectionId{ + // Id: idBuffer, + // } + // expected = "0x000102" + // data, err = cIds.MarshalSSZ() + // assert.NoError(t, err) + // assert.Equal(t, expected, fmt.Sprintf("0x%x", data)) +} + +func TestOfferAndAcceptMessage(t *testing.T) { + contentKey := "0x010203" + contentBytes := hexutil.MustDecode(contentKey) + contentKeys := [][]byte{contentBytes} + offer := &Offer{ + ContentKeys: contentKeys, + } + + expected := "0x0400000004000000010203" + + data, err := offer.MarshalSSZ() + assert.NoError(t, err) + assert.Equal(t, expected, fmt.Sprintf("0x%x", data)) + + // TODO test for accept + + accept := &Accept{ + ConnectionId: []byte{0x01, 0x02}, + ContentKeys: []byte{1, 0, 0, 0, 0, 0, 0, 0}, + } + + expected = "0x0102060000000101" + + data, err = accept.MarshalSSZ() + assert.NoError(t, err) + assert.Equal(t, expected, fmt.Sprintf("0x%x", data)) +} + +func ReverseBytes(src []byte) []byte { + lenth := len(src) + dst := make([]byte, lenth) + for i := 0; i < len(src); i++ { + dst[lenth-1-i] = src[i] + } + return dst +}