github.com/theQRL/go-zond@v0.1.1/p2p/discover/v5wire/encoding_test.go (about)

     1  // Copyright 2020 The go-ethereum Authors
     2  // This file is part of the go-ethereum library.
     3  //
     4  // The go-ethereum library is free software: you can redistribute it and/or modify
     5  // it under the terms of the GNU Lesser General Public License as published by
     6  // the Free Software Foundation, either version 3 of the License, or
     7  // (at your option) any later version.
     8  //
     9  // The go-ethereum library is distributed in the hope that it will be useful,
    10  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    11  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
    12  // GNU Lesser General Public License for more details.
    13  //
    14  // You should have received a copy of the GNU Lesser General Public License
    15  // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>.
    16  
    17  package v5wire
    18  
    19  import (
    20  	"bytes"
    21  	"crypto/ecdsa"
    22  	"encoding/hex"
    23  	"errors"
    24  	"flag"
    25  	"fmt"
    26  	"net"
    27  	"os"
    28  	"path/filepath"
    29  	"strings"
    30  	"testing"
    31  
    32  	"github.com/davecgh/go-spew/spew"
    33  	"github.com/theQRL/go-zond/common/hexutil"
    34  	"github.com/theQRL/go-zond/common/mclock"
    35  	"github.com/theQRL/go-zond/crypto"
    36  	"github.com/theQRL/go-zond/p2p/enode"
    37  )
    38  
    39  // To regenerate discv5 test vectors, run
    40  //
    41  //	go test -run TestVectors -write-test-vectors
    42  var writeTestVectorsFlag = flag.Bool("write-test-vectors", false, "Overwrite discv5 test vectors in testdata/")
    43  
    44  var (
    45  	testKeyA, _   = crypto.HexToECDSA("eef77acb6c6a6eebc5b363a475ac583ec7eccdb42b6481424c60f59aa326547f")
    46  	testKeyB, _   = crypto.HexToECDSA("66fb62bfbd66b9177a138c1e5cddbe4f7c30c343e94e68df8769459cb1cde628")
    47  	testEphKey, _ = crypto.HexToECDSA("0288ef00023598499cb6c940146d050d2b1fb914198c327f76aad590bead68b6")
    48  	testIDnonce   = [16]byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}
    49  )
    50  
    51  // This test checks that the minPacketSize and randomPacketMsgSize constants are well-defined.
    52  func TestMinSizes(t *testing.T) {
    53  	var (
    54  		gcmTagSize = 16
    55  		emptyMsg   = sizeofMessageAuthData + gcmTagSize
    56  	)
    57  	t.Log("static header size", sizeofStaticPacketData)
    58  	t.Log("whoareyou size", sizeofStaticPacketData+sizeofWhoareyouAuthData)
    59  	t.Log("empty msg size", sizeofStaticPacketData+emptyMsg)
    60  	if want := emptyMsg; minMessageSize != want {
    61  		t.Fatalf("wrong minMessageSize %d, want %d", minMessageSize, want)
    62  	}
    63  	if sizeofMessageAuthData+randomPacketMsgSize < minMessageSize {
    64  		t.Fatalf("randomPacketMsgSize %d too small", randomPacketMsgSize)
    65  	}
    66  }
    67  
    68  // This test checks the basic handshake flow where A talks to B and A has no secrets.
    69  func TestHandshake(t *testing.T) {
    70  	t.Parallel()
    71  	net := newHandshakeTest()
    72  	defer net.close()
    73  
    74  	// A -> B   RANDOM PACKET
    75  	packet, _ := net.nodeA.encode(t, net.nodeB, &Findnode{})
    76  	resp := net.nodeB.expectDecode(t, UnknownPacket, packet)
    77  
    78  	// A <- B   WHOAREYOU
    79  	challenge := &Whoareyou{
    80  		Nonce:     resp.(*Unknown).Nonce,
    81  		IDNonce:   testIDnonce,
    82  		RecordSeq: 0,
    83  	}
    84  	whoareyou, _ := net.nodeB.encode(t, net.nodeA, challenge)
    85  	net.nodeA.expectDecode(t, WhoareyouPacket, whoareyou)
    86  
    87  	// A -> B   FINDNODE (handshake packet)
    88  	findnode, _ := net.nodeA.encodeWithChallenge(t, net.nodeB, challenge, &Findnode{})
    89  	net.nodeB.expectDecode(t, FindnodeMsg, findnode)
    90  	if len(net.nodeB.c.sc.handshakes) > 0 {
    91  		t.Fatalf("node B didn't remove handshake from challenge map")
    92  	}
    93  
    94  	// A <- B   NODES
    95  	nodes, _ := net.nodeB.encode(t, net.nodeA, &Nodes{RespCount: 1})
    96  	net.nodeA.expectDecode(t, NodesMsg, nodes)
    97  }
    98  
    99  // This test checks that handshake attempts are removed within the timeout.
   100  func TestHandshake_timeout(t *testing.T) {
   101  	t.Parallel()
   102  	net := newHandshakeTest()
   103  	defer net.close()
   104  
   105  	// A -> B   RANDOM PACKET
   106  	packet, _ := net.nodeA.encode(t, net.nodeB, &Findnode{})
   107  	resp := net.nodeB.expectDecode(t, UnknownPacket, packet)
   108  
   109  	// A <- B   WHOAREYOU
   110  	challenge := &Whoareyou{
   111  		Nonce:     resp.(*Unknown).Nonce,
   112  		IDNonce:   testIDnonce,
   113  		RecordSeq: 0,
   114  	}
   115  	whoareyou, _ := net.nodeB.encode(t, net.nodeA, challenge)
   116  	net.nodeA.expectDecode(t, WhoareyouPacket, whoareyou)
   117  
   118  	// A -> B   FINDNODE (handshake packet) after timeout
   119  	net.clock.Run(handshakeTimeout + 1)
   120  	findnode, _ := net.nodeA.encodeWithChallenge(t, net.nodeB, challenge, &Findnode{})
   121  	net.nodeB.expectDecodeErr(t, errUnexpectedHandshake, findnode)
   122  }
   123  
   124  // This test checks handshake behavior when no record is sent in the auth response.
   125  func TestHandshake_norecord(t *testing.T) {
   126  	t.Parallel()
   127  	net := newHandshakeTest()
   128  	defer net.close()
   129  
   130  	// A -> B   RANDOM PACKET
   131  	packet, _ := net.nodeA.encode(t, net.nodeB, &Findnode{})
   132  	resp := net.nodeB.expectDecode(t, UnknownPacket, packet)
   133  
   134  	// A <- B   WHOAREYOU
   135  	nodeA := net.nodeA.n()
   136  	if nodeA.Seq() == 0 {
   137  		t.Fatal("need non-zero sequence number")
   138  	}
   139  	challenge := &Whoareyou{
   140  		Nonce:     resp.(*Unknown).Nonce,
   141  		IDNonce:   testIDnonce,
   142  		RecordSeq: nodeA.Seq(),
   143  		Node:      nodeA,
   144  	}
   145  	whoareyou, _ := net.nodeB.encode(t, net.nodeA, challenge)
   146  	net.nodeA.expectDecode(t, WhoareyouPacket, whoareyou)
   147  
   148  	// A -> B   FINDNODE
   149  	findnode, _ := net.nodeA.encodeWithChallenge(t, net.nodeB, challenge, &Findnode{})
   150  	net.nodeB.expectDecode(t, FindnodeMsg, findnode)
   151  
   152  	// A <- B   NODES
   153  	nodes, _ := net.nodeB.encode(t, net.nodeA, &Nodes{RespCount: 1})
   154  	net.nodeA.expectDecode(t, NodesMsg, nodes)
   155  }
   156  
   157  // In this test, A tries to send FINDNODE with existing secrets but B doesn't know
   158  // anything about A.
   159  func TestHandshake_rekey(t *testing.T) {
   160  	t.Parallel()
   161  	net := newHandshakeTest()
   162  	defer net.close()
   163  
   164  	session := &session{
   165  		readKey:  []byte("BBBBBBBBBBBBBBBB"),
   166  		writeKey: []byte("AAAAAAAAAAAAAAAA"),
   167  	}
   168  	net.nodeA.c.sc.storeNewSession(net.nodeB.id(), net.nodeB.addr(), session)
   169  
   170  	// A -> B   FINDNODE (encrypted with zero keys)
   171  	findnode, authTag := net.nodeA.encode(t, net.nodeB, &Findnode{})
   172  	net.nodeB.expectDecode(t, UnknownPacket, findnode)
   173  
   174  	// A <- B   WHOAREYOU
   175  	challenge := &Whoareyou{Nonce: authTag, IDNonce: testIDnonce}
   176  	whoareyou, _ := net.nodeB.encode(t, net.nodeA, challenge)
   177  	net.nodeA.expectDecode(t, WhoareyouPacket, whoareyou)
   178  
   179  	// Check that new keys haven't been stored yet.
   180  	sa := net.nodeA.c.sc.session(net.nodeB.id(), net.nodeB.addr())
   181  	if !bytes.Equal(sa.writeKey, session.writeKey) || !bytes.Equal(sa.readKey, session.readKey) {
   182  		t.Fatal("node A stored keys too early")
   183  	}
   184  	if s := net.nodeB.c.sc.session(net.nodeA.id(), net.nodeA.addr()); s != nil {
   185  		t.Fatal("node B stored keys too early")
   186  	}
   187  
   188  	// A -> B   FINDNODE encrypted with new keys
   189  	findnode, _ = net.nodeA.encodeWithChallenge(t, net.nodeB, challenge, &Findnode{})
   190  	net.nodeB.expectDecode(t, FindnodeMsg, findnode)
   191  
   192  	// A <- B   NODES
   193  	nodes, _ := net.nodeB.encode(t, net.nodeA, &Nodes{RespCount: 1})
   194  	net.nodeA.expectDecode(t, NodesMsg, nodes)
   195  }
   196  
   197  // In this test A and B have different keys before the handshake.
   198  func TestHandshake_rekey2(t *testing.T) {
   199  	t.Parallel()
   200  	net := newHandshakeTest()
   201  	defer net.close()
   202  
   203  	initKeysA := &session{
   204  		readKey:  []byte("BBBBBBBBBBBBBBBB"),
   205  		writeKey: []byte("AAAAAAAAAAAAAAAA"),
   206  	}
   207  	initKeysB := &session{
   208  		readKey:  []byte("CCCCCCCCCCCCCCCC"),
   209  		writeKey: []byte("DDDDDDDDDDDDDDDD"),
   210  	}
   211  	net.nodeA.c.sc.storeNewSession(net.nodeB.id(), net.nodeB.addr(), initKeysA)
   212  	net.nodeB.c.sc.storeNewSession(net.nodeA.id(), net.nodeA.addr(), initKeysB)
   213  
   214  	// A -> B   FINDNODE encrypted with initKeysA
   215  	findnode, authTag := net.nodeA.encode(t, net.nodeB, &Findnode{Distances: []uint{3}})
   216  	net.nodeB.expectDecode(t, UnknownPacket, findnode)
   217  
   218  	// A <- B   WHOAREYOU
   219  	challenge := &Whoareyou{Nonce: authTag, IDNonce: testIDnonce}
   220  	whoareyou, _ := net.nodeB.encode(t, net.nodeA, challenge)
   221  	net.nodeA.expectDecode(t, WhoareyouPacket, whoareyou)
   222  
   223  	// A -> B   FINDNODE (handshake packet)
   224  	findnode, _ = net.nodeA.encodeWithChallenge(t, net.nodeB, challenge, &Findnode{})
   225  	net.nodeB.expectDecode(t, FindnodeMsg, findnode)
   226  
   227  	// A <- B   NODES
   228  	nodes, _ := net.nodeB.encode(t, net.nodeA, &Nodes{RespCount: 1})
   229  	net.nodeA.expectDecode(t, NodesMsg, nodes)
   230  }
   231  
   232  func TestHandshake_BadHandshakeAttack(t *testing.T) {
   233  	t.Parallel()
   234  	net := newHandshakeTest()
   235  	defer net.close()
   236  
   237  	// A -> B   RANDOM PACKET
   238  	packet, _ := net.nodeA.encode(t, net.nodeB, &Findnode{})
   239  	resp := net.nodeB.expectDecode(t, UnknownPacket, packet)
   240  
   241  	// A <- B   WHOAREYOU
   242  	challenge := &Whoareyou{
   243  		Nonce:     resp.(*Unknown).Nonce,
   244  		IDNonce:   testIDnonce,
   245  		RecordSeq: 0,
   246  	}
   247  	whoareyou, _ := net.nodeB.encode(t, net.nodeA, challenge)
   248  	net.nodeA.expectDecode(t, WhoareyouPacket, whoareyou)
   249  
   250  	// A -> B   FINDNODE
   251  	incorrect_challenge := &Whoareyou{
   252  		IDNonce:   [16]byte{5, 6, 7, 8, 9, 6, 11, 12},
   253  		RecordSeq: challenge.RecordSeq,
   254  		Node:      challenge.Node,
   255  		sent:      challenge.sent,
   256  	}
   257  	incorrect_findnode, _ := net.nodeA.encodeWithChallenge(t, net.nodeB, incorrect_challenge, &Findnode{})
   258  	incorrect_findnode2 := make([]byte, len(incorrect_findnode))
   259  	copy(incorrect_findnode2, incorrect_findnode)
   260  
   261  	net.nodeB.expectDecodeErr(t, errInvalidNonceSig, incorrect_findnode)
   262  
   263  	// Reject new findnode as previous handshake is now deleted.
   264  	net.nodeB.expectDecodeErr(t, errUnexpectedHandshake, incorrect_findnode2)
   265  
   266  	// The findnode packet is again rejected even with a valid challenge this time.
   267  	findnode, _ := net.nodeA.encodeWithChallenge(t, net.nodeB, challenge, &Findnode{})
   268  	net.nodeB.expectDecodeErr(t, errUnexpectedHandshake, findnode)
   269  }
   270  
   271  // This test checks some malformed packets.
   272  func TestDecodeErrorsV5(t *testing.T) {
   273  	t.Parallel()
   274  	net := newHandshakeTest()
   275  	defer net.close()
   276  
   277  	b := make([]byte, 0)
   278  	net.nodeA.expectDecodeErr(t, errTooShort, b)
   279  
   280  	b = make([]byte, 62)
   281  	net.nodeA.expectDecodeErr(t, errTooShort, b)
   282  
   283  	b = make([]byte, 63)
   284  	net.nodeA.expectDecodeErr(t, errInvalidHeader, b)
   285  
   286  	// TODO some more tests would be nice :)
   287  	// - check invalid authdata sizes
   288  	// - check invalid handshake data sizes
   289  }
   290  
   291  // This test checks that all test vectors can be decoded.
   292  func TestTestVectorsV5(t *testing.T) {
   293  	var (
   294  		idA     = enode.PubkeyToIDV4(&testKeyA.PublicKey)
   295  		idB     = enode.PubkeyToIDV4(&testKeyB.PublicKey)
   296  		addr    = "127.0.0.1"
   297  		session = &session{
   298  			writeKey: hexutil.MustDecode("0x00000000000000000000000000000000"),
   299  			readKey:  hexutil.MustDecode("0x01010101010101010101010101010101"),
   300  		}
   301  		challenge0A, challenge1A, challenge0B Whoareyou
   302  	)
   303  
   304  	// Create challenge packets.
   305  	c := Whoareyou{
   306  		Nonce:   Nonce{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12},
   307  		IDNonce: testIDnonce,
   308  	}
   309  	challenge0A, challenge1A, challenge0B = c, c, c
   310  	challenge1A.RecordSeq = 1
   311  	net := newHandshakeTest()
   312  	challenge0A.Node = net.nodeA.n()
   313  	challenge0B.Node = net.nodeB.n()
   314  	challenge1A.Node = net.nodeA.n()
   315  	net.close()
   316  
   317  	type testVectorTest struct {
   318  		name      string               // test vector name
   319  		packet    Packet               // the packet to be encoded
   320  		challenge *Whoareyou           // handshake challenge passed to encoder
   321  		prep      func(*handshakeTest) // called before encode/decode
   322  	}
   323  	tests := []testVectorTest{
   324  		{
   325  			name:   "v5.1-whoareyou",
   326  			packet: &challenge0B,
   327  		},
   328  		{
   329  			name: "v5.1-ping-message",
   330  			packet: &Ping{
   331  				ReqID:  []byte{0, 0, 0, 1},
   332  				ENRSeq: 2,
   333  			},
   334  			prep: func(net *handshakeTest) {
   335  				net.nodeA.c.sc.storeNewSession(idB, addr, session)
   336  				net.nodeB.c.sc.storeNewSession(idA, addr, session.keysFlipped())
   337  			},
   338  		},
   339  		{
   340  			name: "v5.1-ping-handshake-enr",
   341  			packet: &Ping{
   342  				ReqID:  []byte{0, 0, 0, 1},
   343  				ENRSeq: 1,
   344  			},
   345  			challenge: &challenge0A,
   346  			prep: func(net *handshakeTest) {
   347  				// Update challenge.Header.AuthData.
   348  				net.nodeA.c.Encode(idB, "", &challenge0A, nil)
   349  				net.nodeB.c.sc.storeSentHandshake(idA, addr, &challenge0A)
   350  			},
   351  		},
   352  		{
   353  			name: "v5.1-ping-handshake",
   354  			packet: &Ping{
   355  				ReqID:  []byte{0, 0, 0, 1},
   356  				ENRSeq: 1,
   357  			},
   358  			challenge: &challenge1A,
   359  			prep: func(net *handshakeTest) {
   360  				// Update challenge data.
   361  				net.nodeA.c.Encode(idB, "", &challenge1A, nil)
   362  				net.nodeB.c.sc.storeSentHandshake(idA, addr, &challenge1A)
   363  			},
   364  		},
   365  	}
   366  
   367  	for _, test := range tests {
   368  		test := test
   369  		t.Run(test.name, func(t *testing.T) {
   370  			net := newHandshakeTest()
   371  			defer net.close()
   372  
   373  			// Override all random inputs.
   374  			net.nodeA.c.sc.nonceGen = func(counter uint32) (Nonce, error) {
   375  				return Nonce{0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF, 0xFF}, nil
   376  			}
   377  			net.nodeA.c.sc.maskingIVGen = func(buf []byte) error {
   378  				return nil // all zero
   379  			}
   380  			net.nodeA.c.sc.ephemeralKeyGen = func() (*ecdsa.PrivateKey, error) {
   381  				return testEphKey, nil
   382  			}
   383  
   384  			// Prime the codec for encoding/decoding.
   385  			if test.prep != nil {
   386  				test.prep(net)
   387  			}
   388  
   389  			file := filepath.Join("testdata", test.name+".txt")
   390  			if *writeTestVectorsFlag {
   391  				// Encode the packet.
   392  				d, nonce := net.nodeA.encodeWithChallenge(t, net.nodeB, test.challenge, test.packet)
   393  				comment := testVectorComment(net, test.packet, test.challenge, nonce)
   394  				writeTestVector(file, comment, d)
   395  			}
   396  			enc := hexFile(file)
   397  			net.nodeB.expectDecode(t, test.packet.Kind(), enc)
   398  		})
   399  	}
   400  }
   401  
   402  // testVectorComment creates the commentary for discv5 test vector files.
   403  func testVectorComment(net *handshakeTest, p Packet, challenge *Whoareyou, nonce Nonce) string {
   404  	o := new(strings.Builder)
   405  	printWhoareyou := func(p *Whoareyou) {
   406  		fmt.Fprintf(o, "whoareyou.challenge-data = %#x\n", p.ChallengeData)
   407  		fmt.Fprintf(o, "whoareyou.request-nonce = %#x\n", p.Nonce[:])
   408  		fmt.Fprintf(o, "whoareyou.id-nonce = %#x\n", p.IDNonce[:])
   409  		fmt.Fprintf(o, "whoareyou.enr-seq = %d\n", p.RecordSeq)
   410  	}
   411  
   412  	fmt.Fprintf(o, "src-node-id = %#x\n", net.nodeA.id().Bytes())
   413  	fmt.Fprintf(o, "dest-node-id = %#x\n", net.nodeB.id().Bytes())
   414  	switch p := p.(type) {
   415  	case *Whoareyou:
   416  		// WHOAREYOU packet.
   417  		printWhoareyou(p)
   418  	case *Ping:
   419  		fmt.Fprintf(o, "nonce = %#x\n", nonce[:])
   420  		fmt.Fprintf(o, "read-key = %#x\n", net.nodeA.c.sc.session(net.nodeB.id(), net.nodeB.addr()).writeKey)
   421  		fmt.Fprintf(o, "ping.req-id = %#x\n", p.ReqID)
   422  		fmt.Fprintf(o, "ping.enr-seq = %d\n", p.ENRSeq)
   423  		if challenge != nil {
   424  			// Handshake message packet.
   425  			fmt.Fprint(o, "\nhandshake inputs:\n\n")
   426  			printWhoareyou(challenge)
   427  			fmt.Fprintf(o, "ephemeral-key = %#x\n", testEphKey.D.Bytes())
   428  			fmt.Fprintf(o, "ephemeral-pubkey = %#x\n", crypto.CompressPubkey(&testEphKey.PublicKey))
   429  		}
   430  	default:
   431  		panic(fmt.Errorf("unhandled packet type %T", p))
   432  	}
   433  	return o.String()
   434  }
   435  
   436  // This benchmark checks performance of handshake packet decoding.
   437  func BenchmarkV5_DecodeHandshakePingSecp256k1(b *testing.B) {
   438  	net := newHandshakeTest()
   439  	defer net.close()
   440  
   441  	var (
   442  		idA       = net.nodeA.id()
   443  		challenge = &Whoareyou{Node: net.nodeB.n()}
   444  		message   = &Ping{ReqID: []byte("reqid")}
   445  	)
   446  	enc, _, err := net.nodeA.c.Encode(net.nodeB.id(), "", message, challenge)
   447  	if err != nil {
   448  		b.Fatal("can't encode handshake packet")
   449  	}
   450  	challenge.Node = nil // force ENR signature verification in decoder
   451  	b.ResetTimer()
   452  
   453  	input := make([]byte, len(enc))
   454  	for i := 0; i < b.N; i++ {
   455  		copy(input, enc)
   456  		net.nodeB.c.sc.storeSentHandshake(idA, "", challenge)
   457  		_, _, _, err := net.nodeB.c.Decode(input, "")
   458  		if err != nil {
   459  			b.Fatal(err)
   460  		}
   461  	}
   462  }
   463  
   464  // This benchmark checks how long it takes to decode an encrypted ping packet.
   465  func BenchmarkV5_DecodePing(b *testing.B) {
   466  	net := newHandshakeTest()
   467  	defer net.close()
   468  
   469  	session := &session{
   470  		readKey:  []byte{233, 203, 93, 195, 86, 47, 177, 186, 227, 43, 2, 141, 244, 230, 120, 17},
   471  		writeKey: []byte{79, 145, 252, 171, 167, 216, 252, 161, 208, 190, 176, 106, 214, 39, 178, 134},
   472  	}
   473  	net.nodeA.c.sc.storeNewSession(net.nodeB.id(), net.nodeB.addr(), session)
   474  	net.nodeB.c.sc.storeNewSession(net.nodeA.id(), net.nodeA.addr(), session.keysFlipped())
   475  	addrB := net.nodeA.addr()
   476  	ping := &Ping{ReqID: []byte("reqid"), ENRSeq: 5}
   477  	enc, _, err := net.nodeA.c.Encode(net.nodeB.id(), addrB, ping, nil)
   478  	if err != nil {
   479  		b.Fatalf("can't encode: %v", err)
   480  	}
   481  	b.ResetTimer()
   482  
   483  	input := make([]byte, len(enc))
   484  	for i := 0; i < b.N; i++ {
   485  		copy(input, enc)
   486  		_, _, packet, _ := net.nodeB.c.Decode(input, addrB)
   487  		if _, ok := packet.(*Ping); !ok {
   488  			b.Fatalf("wrong packet type %T", packet)
   489  		}
   490  	}
   491  }
   492  
   493  var pp = spew.NewDefaultConfig()
   494  
   495  type handshakeTest struct {
   496  	nodeA, nodeB handshakeTestNode
   497  	clock        mclock.Simulated
   498  }
   499  
   500  type handshakeTestNode struct {
   501  	ln *enode.LocalNode
   502  	c  *Codec
   503  }
   504  
   505  func newHandshakeTest() *handshakeTest {
   506  	t := new(handshakeTest)
   507  	t.nodeA.init(testKeyA, net.IP{127, 0, 0, 1}, &t.clock, DefaultProtocolID)
   508  	t.nodeB.init(testKeyB, net.IP{127, 0, 0, 1}, &t.clock, DefaultProtocolID)
   509  	return t
   510  }
   511  
   512  func (t *handshakeTest) close() {
   513  	t.nodeA.ln.Database().Close()
   514  	t.nodeB.ln.Database().Close()
   515  }
   516  
   517  func (n *handshakeTestNode) init(key *ecdsa.PrivateKey, ip net.IP, clock mclock.Clock, protocolID [6]byte) {
   518  	db, _ := enode.OpenDB("")
   519  	n.ln = enode.NewLocalNode(db, key)
   520  	n.ln.SetStaticIP(ip)
   521  	n.c = NewCodec(n.ln, key, clock, nil)
   522  }
   523  
   524  func (n *handshakeTestNode) encode(t testing.TB, to handshakeTestNode, p Packet) ([]byte, Nonce) {
   525  	t.Helper()
   526  	return n.encodeWithChallenge(t, to, nil, p)
   527  }
   528  
   529  func (n *handshakeTestNode) encodeWithChallenge(t testing.TB, to handshakeTestNode, c *Whoareyou, p Packet) ([]byte, Nonce) {
   530  	t.Helper()
   531  
   532  	// Copy challenge and add destination node. This avoids sharing 'c' among the two codecs.
   533  	var challenge *Whoareyou
   534  	if c != nil {
   535  		challengeCopy := *c
   536  		challenge = &challengeCopy
   537  		challenge.Node = to.n()
   538  	}
   539  	// Encode to destination.
   540  	enc, nonce, err := n.c.Encode(to.id(), to.addr(), p, challenge)
   541  	if err != nil {
   542  		t.Fatal(fmt.Errorf("(%s) %v", n.ln.ID().TerminalString(), err))
   543  	}
   544  	t.Logf("(%s) -> (%s)   %s\n%s", n.ln.ID().TerminalString(), to.id().TerminalString(), p.Name(), hex.Dump(enc))
   545  	return enc, nonce
   546  }
   547  
   548  func (n *handshakeTestNode) expectDecode(t *testing.T, ptype byte, p []byte) Packet {
   549  	t.Helper()
   550  
   551  	dec, err := n.decode(p)
   552  	if err != nil {
   553  		t.Fatal(fmt.Errorf("(%s) %v", n.ln.ID().TerminalString(), err))
   554  	}
   555  	t.Logf("(%s) %#v", n.ln.ID().TerminalString(), pp.NewFormatter(dec))
   556  	if dec.Kind() != ptype {
   557  		t.Fatalf("expected packet type %d, got %d", ptype, dec.Kind())
   558  	}
   559  	return dec
   560  }
   561  
   562  func (n *handshakeTestNode) expectDecodeErr(t *testing.T, wantErr error, p []byte) {
   563  	t.Helper()
   564  	if _, err := n.decode(p); !errors.Is(err, wantErr) {
   565  		t.Fatal(fmt.Errorf("(%s) got err %q, want %q", n.ln.ID().TerminalString(), err, wantErr))
   566  	}
   567  }
   568  
   569  func (n *handshakeTestNode) decode(input []byte) (Packet, error) {
   570  	_, _, p, err := n.c.Decode(input, "127.0.0.1")
   571  	return p, err
   572  }
   573  
   574  func (n *handshakeTestNode) n() *enode.Node {
   575  	return n.ln.Node()
   576  }
   577  
   578  func (n *handshakeTestNode) addr() string {
   579  	return n.ln.Node().IP().String()
   580  }
   581  
   582  func (n *handshakeTestNode) id() enode.ID {
   583  	return n.ln.ID()
   584  }
   585  
   586  // hexFile reads the given file and decodes the hex data contained in it.
   587  // Whitespace and any lines beginning with the # character are ignored.
   588  func hexFile(file string) []byte {
   589  	fileContent, err := os.ReadFile(file)
   590  	if err != nil {
   591  		panic(err)
   592  	}
   593  
   594  	// Gather hex data, ignore comments.
   595  	var text []byte
   596  	for _, line := range bytes.Split(fileContent, []byte("\n")) {
   597  		line = bytes.TrimSpace(line)
   598  		if len(line) > 0 && line[0] == '#' {
   599  			continue
   600  		}
   601  		text = append(text, line...)
   602  	}
   603  
   604  	// Parse the hex.
   605  	if bytes.HasPrefix(text, []byte("0x")) {
   606  		text = text[2:]
   607  	}
   608  	data := make([]byte, hex.DecodedLen(len(text)))
   609  	if _, err := hex.Decode(data, text); err != nil {
   610  		panic("invalid hex in " + file)
   611  	}
   612  	return data
   613  }
   614  
   615  // writeTestVector writes a test vector file with the given commentary and binary data.
   616  func writeTestVector(file, comment string, data []byte) {
   617  	fd, err := os.OpenFile(file, os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
   618  	if err != nil {
   619  		panic(err)
   620  	}
   621  	defer fd.Close()
   622  
   623  	if len(comment) > 0 {
   624  		for _, line := range strings.Split(strings.TrimSpace(comment), "\n") {
   625  			fmt.Fprintf(fd, "# %s\n", line)
   626  		}
   627  		fmt.Fprintln(fd)
   628  	}
   629  	for len(data) > 0 {
   630  		var chunk []byte
   631  		if len(data) < 32 {
   632  			chunk = data
   633  		} else {
   634  			chunk = data[:32]
   635  		}
   636  		data = data[len(chunk):]
   637  		fmt.Fprintf(fd, "%x\n", chunk)
   638  	}
   639  }