github.com/ethereum/go-ethereum@v1.16.1/p2p/discover/v5_udp_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 discover
    18  
    19  import (
    20  	"bytes"
    21  	"crypto/ecdsa"
    22  	"encoding/binary"
    23  	"fmt"
    24  	"math/rand"
    25  	"net"
    26  	"net/netip"
    27  	"reflect"
    28  	"slices"
    29  	"testing"
    30  	"time"
    31  
    32  	"github.com/ethereum/go-ethereum/internal/testlog"
    33  	"github.com/ethereum/go-ethereum/log"
    34  	"github.com/ethereum/go-ethereum/p2p/discover/v4wire"
    35  	"github.com/ethereum/go-ethereum/p2p/discover/v5wire"
    36  	"github.com/ethereum/go-ethereum/p2p/enode"
    37  	"github.com/ethereum/go-ethereum/p2p/enr"
    38  	"github.com/ethereum/go-ethereum/rlp"
    39  	"github.com/stretchr/testify/require"
    40  )
    41  
    42  // Real sockets, real crypto: this test checks end-to-end connectivity for UDPv5.
    43  func TestUDPv5_lookupE2E(t *testing.T) {
    44  	t.Parallel()
    45  
    46  	const N = 5
    47  	var nodes []*UDPv5
    48  	for i := 0; i < N; i++ {
    49  		var cfg Config
    50  		if len(nodes) > 0 {
    51  			bn := nodes[0].Self()
    52  			cfg.Bootnodes = []*enode.Node{bn}
    53  		}
    54  		node := startLocalhostV5(t, cfg)
    55  		nodes = append(nodes, node)
    56  		defer node.Close()
    57  	}
    58  	last := nodes[N-1]
    59  	target := nodes[rand.Intn(N-2)].Self()
    60  
    61  	// It is expected that all nodes can be found.
    62  	expectedResult := make([]*enode.Node, len(nodes))
    63  	for i := range nodes {
    64  		expectedResult[i] = nodes[i].Self()
    65  	}
    66  	slices.SortFunc(expectedResult, func(a, b *enode.Node) int {
    67  		return enode.DistCmp(target.ID(), a.ID(), b.ID())
    68  	})
    69  
    70  	// Do the lookup.
    71  	results := last.Lookup(target.ID())
    72  	if err := checkNodesEqual(results, expectedResult); err != nil {
    73  		t.Fatalf("lookup returned wrong results: %v", err)
    74  	}
    75  }
    76  
    77  func startLocalhostV5(t *testing.T, cfg Config) *UDPv5 {
    78  	cfg.PrivateKey = newkey()
    79  	db, _ := enode.OpenDB("")
    80  	ln := enode.NewLocalNode(db, cfg.PrivateKey)
    81  
    82  	// Prefix logs with node ID.
    83  	lprefix := fmt.Sprintf("(%s)", ln.ID().TerminalString())
    84  	cfg.Log = testlog.Logger(t, log.LevelTrace).With("node-id", lprefix)
    85  
    86  	// Listen.
    87  	socket, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IP{127, 0, 0, 1}})
    88  	if err != nil {
    89  		t.Fatal(err)
    90  	}
    91  	realaddr := socket.LocalAddr().(*net.UDPAddr)
    92  	ln.SetStaticIP(realaddr.IP)
    93  	ln.Set(enr.UDP(realaddr.Port))
    94  	udp, err := ListenV5(socket, ln, cfg)
    95  	if err != nil {
    96  		t.Fatal(err)
    97  	}
    98  	return udp
    99  }
   100  
   101  // This test checks that incoming PING calls are handled correctly.
   102  func TestUDPv5_pingHandling(t *testing.T) {
   103  	t.Parallel()
   104  	test := newUDPV5Test(t)
   105  	defer test.close()
   106  
   107  	test.packetIn(&v5wire.Ping{ReqID: []byte("foo")})
   108  	test.waitPacketOut(func(p *v5wire.Pong, addr netip.AddrPort, _ v5wire.Nonce) {
   109  		if !bytes.Equal(p.ReqID, []byte("foo")) {
   110  			t.Error("wrong request ID in response:", p.ReqID)
   111  		}
   112  		if p.ENRSeq != test.table.self().Seq() {
   113  			t.Error("wrong ENR sequence number in response:", p.ENRSeq)
   114  		}
   115  	})
   116  }
   117  
   118  // This test checks that incoming 'unknown' packets trigger the handshake.
   119  func TestUDPv5_unknownPacket(t *testing.T) {
   120  	t.Parallel()
   121  	test := newUDPV5Test(t)
   122  	defer test.close()
   123  
   124  	nonce := v5wire.Nonce{1, 2, 3}
   125  	check := func(p *v5wire.Whoareyou, wantSeq uint64) {
   126  		t.Helper()
   127  		if p.Nonce != nonce {
   128  			t.Error("wrong nonce in WHOAREYOU:", p.Nonce, nonce)
   129  		}
   130  		if p.IDNonce == ([16]byte{}) {
   131  			t.Error("all zero ID nonce")
   132  		}
   133  		if p.RecordSeq != wantSeq {
   134  			t.Errorf("wrong record seq %d in WHOAREYOU, want %d", p.RecordSeq, wantSeq)
   135  		}
   136  	}
   137  
   138  	// Unknown packet from unknown node.
   139  	test.packetIn(&v5wire.Unknown{Nonce: nonce})
   140  	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, _ v5wire.Nonce) {
   141  		check(p, 0)
   142  	})
   143  }
   144  
   145  func TestUDPv5_unknownPacketKnownNode(t *testing.T) {
   146  	t.Parallel()
   147  	test := newUDPV5Test(t)
   148  	defer test.close()
   149  
   150  	nonce := v5wire.Nonce{1, 2, 3}
   151  	check := func(p *v5wire.Whoareyou, wantSeq uint64) {
   152  		t.Helper()
   153  		if p.Nonce != nonce {
   154  			t.Error("wrong nonce in WHOAREYOU:", p.Nonce, nonce)
   155  		}
   156  		if p.IDNonce == ([16]byte{}) {
   157  			t.Error("all zero ID nonce")
   158  		}
   159  		if p.RecordSeq != wantSeq {
   160  			t.Errorf("wrong record seq %d in WHOAREYOU, want %d", p.RecordSeq, wantSeq)
   161  		}
   162  	}
   163  
   164  	// Make node known.
   165  	n := test.getNode(test.remotekey, test.remoteaddr).Node()
   166  	test.table.addFoundNode(n, false)
   167  
   168  	test.packetIn(&v5wire.Unknown{Nonce: nonce})
   169  	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, _ v5wire.Nonce) {
   170  		check(p, n.Seq())
   171  	})
   172  }
   173  
   174  // This test checks that, when multiple 'unknown' packets are received during a handshake,
   175  // the node sticks to the first handshake attempt.
   176  func TestUDPv5_handshakeRepeatChallenge(t *testing.T) {
   177  	t.Parallel()
   178  	test := newUDPV5Test(t)
   179  	defer test.close()
   180  
   181  	nonce1 := v5wire.Nonce{1}
   182  	nonce2 := v5wire.Nonce{2}
   183  	nonce3 := v5wire.Nonce{3}
   184  	var firstAuthTag *v5wire.Nonce
   185  	check := func(p *v5wire.Whoareyou, authTag, wantNonce v5wire.Nonce) {
   186  		t.Helper()
   187  		if p.Nonce != wantNonce {
   188  			t.Error("wrong nonce in WHOAREYOU:", p.Nonce, "want:", wantNonce)
   189  		}
   190  		if firstAuthTag == nil {
   191  			firstAuthTag = &authTag
   192  		} else if authTag != *firstAuthTag {
   193  			t.Error("wrong auth tag in WHOAREYOU header:", authTag, "want:", *firstAuthTag)
   194  		}
   195  	}
   196  
   197  	// Unknown packet from unknown node.
   198  	test.packetIn(&v5wire.Unknown{Nonce: nonce1})
   199  	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, authTag v5wire.Nonce) {
   200  		check(p, authTag, nonce1)
   201  	})
   202  
   203  	// Second unknown packet. Here we expect the response to reference the
   204  	// first unknown packet.
   205  	test.packetIn(&v5wire.Unknown{Nonce: nonce2})
   206  	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, authTag v5wire.Nonce) {
   207  		check(p, authTag, nonce1)
   208  	})
   209  	// Third unknown packet. This should still return the first nonce.
   210  	test.packetIn(&v5wire.Unknown{Nonce: nonce3})
   211  	test.waitPacketOut(func(p *v5wire.Whoareyou, addr netip.AddrPort, authTag v5wire.Nonce) {
   212  		check(p, authTag, nonce1)
   213  	})
   214  }
   215  
   216  // This test checks that incoming FINDNODE calls are handled correctly.
   217  func TestUDPv5_findnodeHandling(t *testing.T) {
   218  	t.Parallel()
   219  	test := newUDPV5Test(t)
   220  	defer test.close()
   221  
   222  	// Create test nodes and insert them into the table.
   223  	nodes253 := nodesAtDistance(test.table.self().ID(), 253, 16)
   224  	nodes249 := nodesAtDistance(test.table.self().ID(), 249, 4)
   225  	nodes248 := nodesAtDistance(test.table.self().ID(), 248, 10)
   226  	fillTable(test.table, nodes253, true)
   227  	fillTable(test.table, nodes249, true)
   228  	fillTable(test.table, nodes248, true)
   229  
   230  	// Requesting with distance zero should return the node's own record.
   231  	test.packetIn(&v5wire.Findnode{ReqID: []byte{0}, Distances: []uint{0}})
   232  	test.expectNodes([]byte{0}, 1, []*enode.Node{test.udp.Self()})
   233  
   234  	// Requesting with distance > 256 shouldn't crash.
   235  	test.packetIn(&v5wire.Findnode{ReqID: []byte{1}, Distances: []uint{4234098}})
   236  	test.expectNodes([]byte{1}, 1, nil)
   237  
   238  	// Requesting with empty distance list shouldn't crash either.
   239  	test.packetIn(&v5wire.Findnode{ReqID: []byte{2}, Distances: []uint{}})
   240  	test.expectNodes([]byte{2}, 1, nil)
   241  
   242  	// This request gets no nodes because the corresponding bucket is empty.
   243  	test.packetIn(&v5wire.Findnode{ReqID: []byte{3}, Distances: []uint{254}})
   244  	test.expectNodes([]byte{3}, 1, nil)
   245  
   246  	// This request gets all the distance-253 nodes.
   247  	test.packetIn(&v5wire.Findnode{ReqID: []byte{4}, Distances: []uint{253}})
   248  	test.expectNodes([]byte{4}, 2, nodes253)
   249  
   250  	// This request gets all the distance-249 nodes and some more at 248 because
   251  	// the bucket at 249 is not full.
   252  	test.packetIn(&v5wire.Findnode{ReqID: []byte{5}, Distances: []uint{249, 248}})
   253  	var nodes []*enode.Node
   254  	nodes = append(nodes, nodes249...)
   255  	nodes = append(nodes, nodes248[:10]...)
   256  	test.expectNodes([]byte{5}, 1, nodes)
   257  }
   258  
   259  func (test *udpV5Test) expectNodes(wantReqID []byte, wantTotal uint8, wantNodes []*enode.Node) {
   260  	nodeSet := make(map[enode.ID]*enr.Record, len(wantNodes))
   261  	for _, n := range wantNodes {
   262  		nodeSet[n.ID()] = n.Record()
   263  	}
   264  
   265  	for {
   266  		test.waitPacketOut(func(p *v5wire.Nodes, addr netip.AddrPort, _ v5wire.Nonce) {
   267  			if !bytes.Equal(p.ReqID, wantReqID) {
   268  				test.t.Fatalf("wrong request ID %v in response, want %v", p.ReqID, wantReqID)
   269  			}
   270  			if p.RespCount != wantTotal {
   271  				test.t.Fatalf("wrong total response count %d, want %d", p.RespCount, wantTotal)
   272  			}
   273  			for _, record := range p.Nodes {
   274  				n, _ := enode.New(enode.ValidSchemesForTesting, record)
   275  				want := nodeSet[n.ID()]
   276  				if want == nil {
   277  					test.t.Fatalf("unexpected node in response: %v", n)
   278  				}
   279  				if !reflect.DeepEqual(record, want) {
   280  					test.t.Fatalf("wrong record in response: %v", n)
   281  				}
   282  				delete(nodeSet, n.ID())
   283  			}
   284  		})
   285  		if len(nodeSet) == 0 {
   286  			return
   287  		}
   288  	}
   289  }
   290  
   291  // This test checks that outgoing PING calls work.
   292  func TestUDPv5_pingCall(t *testing.T) {
   293  	t.Parallel()
   294  	test := newUDPV5Test(t)
   295  	defer test.close()
   296  
   297  	remote := test.getNode(test.remotekey, test.remoteaddr).Node()
   298  	done := make(chan error, 1)
   299  
   300  	// This ping times out.
   301  	go func() {
   302  		_, err := test.udp.ping(remote)
   303  		done <- err
   304  	}()
   305  	test.waitPacketOut(func(p *v5wire.Ping, addr netip.AddrPort, _ v5wire.Nonce) {})
   306  	if err := <-done; err != errTimeout {
   307  		t.Fatalf("want errTimeout, got %q", err)
   308  	}
   309  
   310  	// This ping works.
   311  	go func() {
   312  		_, err := test.udp.ping(remote)
   313  		done <- err
   314  	}()
   315  	test.waitPacketOut(func(p *v5wire.Ping, addr netip.AddrPort, _ v5wire.Nonce) {
   316  		test.packetInFrom(test.remotekey, test.remoteaddr, &v5wire.Pong{ReqID: p.ReqID})
   317  	})
   318  	if err := <-done; err != nil {
   319  		t.Fatal(err)
   320  	}
   321  
   322  	// This ping gets a reply from the wrong endpoint.
   323  	go func() {
   324  		_, err := test.udp.ping(remote)
   325  		done <- err
   326  	}()
   327  	test.waitPacketOut(func(p *v5wire.Ping, addr netip.AddrPort, _ v5wire.Nonce) {
   328  		wrongAddr := netip.MustParseAddrPort("33.44.55.22:10101")
   329  		test.packetInFrom(test.remotekey, wrongAddr, &v5wire.Pong{ReqID: p.ReqID})
   330  	})
   331  	if err := <-done; err != errTimeout {
   332  		t.Fatalf("want errTimeout for reply from wrong IP, got %q", err)
   333  	}
   334  }
   335  
   336  // This test checks that outgoing FINDNODE calls work and multiple NODES
   337  // replies are aggregated.
   338  func TestUDPv5_findnodeCall(t *testing.T) {
   339  	t.Parallel()
   340  	test := newUDPV5Test(t)
   341  	defer test.close()
   342  
   343  	// Launch the request:
   344  	var (
   345  		distances = []uint{230}
   346  		remote    = test.getNode(test.remotekey, test.remoteaddr).Node()
   347  		nodes     = nodesAtDistance(remote.ID(), int(distances[0]), 8)
   348  		done      = make(chan error, 1)
   349  		response  []*enode.Node
   350  	)
   351  	go func() {
   352  		var err error
   353  		response, err = test.udp.Findnode(remote, distances)
   354  		done <- err
   355  	}()
   356  
   357  	// Serve the responses:
   358  	test.waitPacketOut(func(p *v5wire.Findnode, addr netip.AddrPort, _ v5wire.Nonce) {
   359  		if !reflect.DeepEqual(p.Distances, distances) {
   360  			t.Fatalf("wrong distances in request: %v", p.Distances)
   361  		}
   362  		test.packetIn(&v5wire.Nodes{
   363  			ReqID:     p.ReqID,
   364  			RespCount: 2,
   365  			Nodes:     nodesToRecords(nodes[:4]),
   366  		})
   367  		test.packetIn(&v5wire.Nodes{
   368  			ReqID:     p.ReqID,
   369  			RespCount: 2,
   370  			Nodes:     nodesToRecords(nodes[4:]),
   371  		})
   372  	})
   373  
   374  	// Check results:
   375  	if err := <-done; err != nil {
   376  		t.Fatalf("unexpected error: %v", err)
   377  	}
   378  	if !reflect.DeepEqual(response, nodes) {
   379  		t.Fatalf("wrong nodes in response")
   380  	}
   381  
   382  	// TODO: check invalid IPs
   383  	// TODO: check invalid/unsigned record
   384  }
   385  
   386  // This test checks that pending calls are re-sent when a handshake happens.
   387  func TestUDPv5_callResend(t *testing.T) {
   388  	t.Parallel()
   389  	test := newUDPV5Test(t)
   390  	defer test.close()
   391  
   392  	remote := test.getNode(test.remotekey, test.remoteaddr).Node()
   393  	done := make(chan error, 2)
   394  	go func() {
   395  		_, err := test.udp.ping(remote)
   396  		done <- err
   397  	}()
   398  	go func() {
   399  		_, err := test.udp.ping(remote)
   400  		done <- err
   401  	}()
   402  
   403  	// Ping answered by WHOAREYOU.
   404  	test.waitPacketOut(func(p *v5wire.Ping, addr netip.AddrPort, nonce v5wire.Nonce) {
   405  		test.packetIn(&v5wire.Whoareyou{Nonce: nonce})
   406  	})
   407  	// Ping should be re-sent.
   408  	test.waitPacketOut(func(p *v5wire.Ping, addr netip.AddrPort, _ v5wire.Nonce) {
   409  		test.packetIn(&v5wire.Pong{ReqID: p.ReqID})
   410  	})
   411  	// Answer the other ping.
   412  	test.waitPacketOut(func(p *v5wire.Ping, addr netip.AddrPort, _ v5wire.Nonce) {
   413  		test.packetIn(&v5wire.Pong{ReqID: p.ReqID})
   414  	})
   415  	if err := <-done; err != nil {
   416  		t.Fatalf("unexpected ping error: %v", err)
   417  	}
   418  	if err := <-done; err != nil {
   419  		t.Fatalf("unexpected ping error: %v", err)
   420  	}
   421  }
   422  
   423  // This test ensures we don't allow multiple rounds of WHOAREYOU for a single call.
   424  func TestUDPv5_multipleHandshakeRounds(t *testing.T) {
   425  	t.Parallel()
   426  	test := newUDPV5Test(t)
   427  	defer test.close()
   428  
   429  	remote := test.getNode(test.remotekey, test.remoteaddr).Node()
   430  	done := make(chan error, 1)
   431  	go func() {
   432  		_, err := test.udp.ping(remote)
   433  		done <- err
   434  	}()
   435  
   436  	// Ping answered by WHOAREYOU.
   437  	test.waitPacketOut(func(p *v5wire.Ping, addr netip.AddrPort, nonce v5wire.Nonce) {
   438  		test.packetIn(&v5wire.Whoareyou{Nonce: nonce})
   439  	})
   440  	// Ping answered by WHOAREYOU again.
   441  	test.waitPacketOut(func(p *v5wire.Ping, addr netip.AddrPort, nonce v5wire.Nonce) {
   442  		test.packetIn(&v5wire.Whoareyou{Nonce: nonce})
   443  	})
   444  	if err := <-done; err != errTimeout {
   445  		t.Fatalf("unexpected ping error: %q", err)
   446  	}
   447  }
   448  
   449  // This test checks that calls with n replies may take up to n * respTimeout.
   450  func TestUDPv5_callTimeoutReset(t *testing.T) {
   451  	t.Parallel()
   452  	test := newUDPV5Test(t)
   453  	defer test.close()
   454  
   455  	// Launch the request:
   456  	var (
   457  		distance = uint(230)
   458  		remote   = test.getNode(test.remotekey, test.remoteaddr).Node()
   459  		nodes    = nodesAtDistance(remote.ID(), int(distance), 8)
   460  		done     = make(chan error, 1)
   461  	)
   462  	go func() {
   463  		_, err := test.udp.Findnode(remote, []uint{distance})
   464  		done <- err
   465  	}()
   466  
   467  	// Serve two responses, slowly.
   468  	test.waitPacketOut(func(p *v5wire.Findnode, addr netip.AddrPort, _ v5wire.Nonce) {
   469  		time.Sleep(respTimeout - 50*time.Millisecond)
   470  		test.packetIn(&v5wire.Nodes{
   471  			ReqID:     p.ReqID,
   472  			RespCount: 2,
   473  			Nodes:     nodesToRecords(nodes[:4]),
   474  		})
   475  
   476  		time.Sleep(respTimeout - 50*time.Millisecond)
   477  		test.packetIn(&v5wire.Nodes{
   478  			ReqID:     p.ReqID,
   479  			RespCount: 2,
   480  			Nodes:     nodesToRecords(nodes[4:]),
   481  		})
   482  	})
   483  	if err := <-done; err != nil {
   484  		t.Fatalf("unexpected error: %q", err)
   485  	}
   486  }
   487  
   488  // This test checks that TALKREQ calls the registered handler function.
   489  func TestUDPv5_talkHandling(t *testing.T) {
   490  	t.Parallel()
   491  	test := newUDPV5Test(t)
   492  	defer test.close()
   493  
   494  	var recvMessage []byte
   495  	test.udp.RegisterTalkHandler("test", func(n *enode.Node, addr *net.UDPAddr, message []byte) []byte {
   496  		recvMessage = message
   497  		return []byte("test response")
   498  	})
   499  
   500  	// Successful case:
   501  	test.packetIn(&v5wire.TalkRequest{
   502  		ReqID:    []byte("foo"),
   503  		Protocol: "test",
   504  		Message:  []byte("test request"),
   505  	})
   506  	test.waitPacketOut(func(p *v5wire.TalkResponse, addr netip.AddrPort, _ v5wire.Nonce) {
   507  		if !bytes.Equal(p.ReqID, []byte("foo")) {
   508  			t.Error("wrong request ID in response:", p.ReqID)
   509  		}
   510  		if string(p.Message) != "test response" {
   511  			t.Errorf("wrong talk response message: %q", p.Message)
   512  		}
   513  		if string(recvMessage) != "test request" {
   514  			t.Errorf("wrong message received in handler: %q", recvMessage)
   515  		}
   516  	})
   517  
   518  	// Check that empty response is returned for unregistered protocols.
   519  	recvMessage = nil
   520  	test.packetIn(&v5wire.TalkRequest{
   521  		ReqID:    []byte("2"),
   522  		Protocol: "wrong",
   523  		Message:  []byte("test request"),
   524  	})
   525  	test.waitPacketOut(func(p *v5wire.TalkResponse, addr netip.AddrPort, _ v5wire.Nonce) {
   526  		if !bytes.Equal(p.ReqID, []byte("2")) {
   527  			t.Error("wrong request ID in response:", p.ReqID)
   528  		}
   529  		if string(p.Message) != "" {
   530  			t.Errorf("wrong talk response message: %q", p.Message)
   531  		}
   532  		if recvMessage != nil {
   533  			t.Errorf("handler was called for wrong protocol: %q", recvMessage)
   534  		}
   535  	})
   536  }
   537  
   538  // This test checks that outgoing TALKREQ calls work.
   539  func TestUDPv5_talkRequest(t *testing.T) {
   540  	t.Parallel()
   541  	test := newUDPV5Test(t)
   542  	defer test.close()
   543  
   544  	remote := test.getNode(test.remotekey, test.remoteaddr).Node()
   545  	done := make(chan error, 1)
   546  
   547  	// This request times out.
   548  	go func() {
   549  		_, err := test.udp.TalkRequest(remote, "test", []byte("test request"))
   550  		done <- err
   551  	}()
   552  	test.waitPacketOut(func(p *v5wire.TalkRequest, addr netip.AddrPort, _ v5wire.Nonce) {})
   553  	if err := <-done; err != errTimeout {
   554  		t.Fatalf("want errTimeout, got %q", err)
   555  	}
   556  
   557  	// This request works.
   558  	go func() {
   559  		_, err := test.udp.TalkRequest(remote, "test", []byte("test request"))
   560  		done <- err
   561  	}()
   562  	test.waitPacketOut(func(p *v5wire.TalkRequest, addr netip.AddrPort, _ v5wire.Nonce) {
   563  		if p.Protocol != "test" {
   564  			t.Errorf("wrong protocol ID in talk request: %q", p.Protocol)
   565  		}
   566  		if string(p.Message) != "test request" {
   567  			t.Errorf("wrong message talk request: %q", p.Message)
   568  		}
   569  		test.packetInFrom(test.remotekey, test.remoteaddr, &v5wire.TalkResponse{
   570  			ReqID:   p.ReqID,
   571  			Message: []byte("test response"),
   572  		})
   573  	})
   574  	if err := <-done; err != nil {
   575  		t.Fatal(err)
   576  	}
   577  
   578  	// Also check requesting without ENR.
   579  	go func() {
   580  		_, err := test.udp.TalkRequestToID(remote.ID(), test.remoteaddr, "test", []byte("test request 2"))
   581  		done <- err
   582  	}()
   583  	test.waitPacketOut(func(p *v5wire.TalkRequest, addr netip.AddrPort, _ v5wire.Nonce) {
   584  		if p.Protocol != "test" {
   585  			t.Errorf("wrong protocol ID in talk request: %q", p.Protocol)
   586  		}
   587  		if string(p.Message) != "test request 2" {
   588  			t.Errorf("wrong message talk request: %q", p.Message)
   589  		}
   590  		test.packetInFrom(test.remotekey, test.remoteaddr, &v5wire.TalkResponse{
   591  			ReqID:   p.ReqID,
   592  			Message: []byte("test response 2"),
   593  		})
   594  	})
   595  	if err := <-done; err != nil {
   596  		t.Fatal(err)
   597  	}
   598  }
   599  
   600  // This test checks that lookupDistances works.
   601  func TestUDPv5_lookupDistances(t *testing.T) {
   602  	test := newUDPV5Test(t)
   603  	lnID := test.table.self().ID()
   604  
   605  	t.Run("target distance of 1", func(t *testing.T) {
   606  		node := nodeAtDistance(lnID, 1, intIP(0))
   607  		dists := lookupDistances(lnID, node.ID())
   608  		require.Equal(t, []uint{1, 2, 3}, dists)
   609  	})
   610  
   611  	t.Run("target distance of 2", func(t *testing.T) {
   612  		node := nodeAtDistance(lnID, 2, intIP(0))
   613  		dists := lookupDistances(lnID, node.ID())
   614  		require.Equal(t, []uint{2, 3, 1}, dists)
   615  	})
   616  
   617  	t.Run("target distance of 128", func(t *testing.T) {
   618  		node := nodeAtDistance(lnID, 128, intIP(0))
   619  		dists := lookupDistances(lnID, node.ID())
   620  		require.Equal(t, []uint{128, 129, 127}, dists)
   621  	})
   622  
   623  	t.Run("target distance of 255", func(t *testing.T) {
   624  		node := nodeAtDistance(lnID, 255, intIP(0))
   625  		dists := lookupDistances(lnID, node.ID())
   626  		require.Equal(t, []uint{255, 256, 254}, dists)
   627  	})
   628  
   629  	t.Run("target distance of 256", func(t *testing.T) {
   630  		node := nodeAtDistance(lnID, 256, intIP(0))
   631  		dists := lookupDistances(lnID, node.ID())
   632  		require.Equal(t, []uint{256, 255, 254}, dists)
   633  	})
   634  }
   635  
   636  // This test checks that lookup works.
   637  func TestUDPv5_lookup(t *testing.T) {
   638  	t.Parallel()
   639  	test := newUDPV5Test(t)
   640  
   641  	// Lookup on empty table returns no nodes.
   642  	if results := test.udp.Lookup(lookupTestnet.target.ID()); len(results) > 0 {
   643  		t.Fatalf("lookup on empty table returned %d results: %#v", len(results), results)
   644  	}
   645  
   646  	// Ensure the tester knows all nodes in lookupTestnet by IP.
   647  	for d, nn := range lookupTestnet.dists {
   648  		for i, key := range nn {
   649  			n := lookupTestnet.node(d, i)
   650  			addr, _ := n.UDPEndpoint()
   651  			test.getNode(key, addr)
   652  		}
   653  	}
   654  
   655  	// Seed table with initial node.
   656  	initialNode := lookupTestnet.node(256, 0)
   657  	fillTable(test.table, []*enode.Node{initialNode}, true)
   658  
   659  	// Start the lookup.
   660  	resultC := make(chan []*enode.Node, 1)
   661  	go func() {
   662  		resultC <- test.udp.Lookup(lookupTestnet.target.ID())
   663  		test.close()
   664  	}()
   665  
   666  	// Answer lookup packets.
   667  	asked := make(map[enode.ID]bool)
   668  	for done := false; !done; {
   669  		done = test.waitPacketOut(func(p v5wire.Packet, to netip.AddrPort, _ v5wire.Nonce) {
   670  			recipient, key := lookupTestnet.nodeByAddr(to)
   671  			switch p := p.(type) {
   672  			case *v5wire.Ping:
   673  				test.packetInFrom(key, to, &v5wire.Pong{ReqID: p.ReqID})
   674  			case *v5wire.Findnode:
   675  				if asked[recipient.ID()] {
   676  					t.Error("Asked node", recipient.ID(), "twice")
   677  				}
   678  				asked[recipient.ID()] = true
   679  				nodes := lookupTestnet.neighborsAtDistances(recipient, p.Distances, 16)
   680  				t.Logf("Got FINDNODE for %v, returning %d nodes", p.Distances, len(nodes))
   681  				for _, resp := range packNodes(p.ReqID, nodes) {
   682  					test.packetInFrom(key, to, resp)
   683  				}
   684  			}
   685  		})
   686  	}
   687  
   688  	// Verify result nodes.
   689  	results := <-resultC
   690  	checkLookupResults(t, lookupTestnet, results)
   691  }
   692  
   693  // This test checks the local node can be utilised to set key-values.
   694  func TestUDPv5_LocalNode(t *testing.T) {
   695  	t.Parallel()
   696  	var cfg Config
   697  	node := startLocalhostV5(t, cfg)
   698  	defer node.Close()
   699  	localNd := node.LocalNode()
   700  
   701  	// set value in node's local record
   702  	testVal := [4]byte{'A', 'B', 'C', 'D'}
   703  	localNd.Set(enr.WithEntry("testing", &testVal))
   704  
   705  	// retrieve the value from self to make sure it matches.
   706  	outputVal := [4]byte{}
   707  	if err := node.Self().Load(enr.WithEntry("testing", &outputVal)); err != nil {
   708  		t.Errorf("Could not load value from record: %v", err)
   709  	}
   710  	if testVal != outputVal {
   711  		t.Errorf("Wanted %#x to be retrieved from the record but instead got %#x", testVal, outputVal)
   712  	}
   713  }
   714  
   715  func TestUDPv5_PingWithIPV4MappedAddress(t *testing.T) {
   716  	t.Parallel()
   717  	test := newUDPV5Test(t)
   718  	defer test.close()
   719  
   720  	rawIP := netip.AddrFrom4([4]byte{0xFF, 0x12, 0x33, 0xE5})
   721  	test.remoteaddr = netip.AddrPortFrom(netip.AddrFrom16(rawIP.As16()), 0)
   722  	remote := test.getNode(test.remotekey, test.remoteaddr).Node()
   723  	done := make(chan struct{}, 1)
   724  
   725  	// This handler will truncate the ipv4-mapped in ipv6 address.
   726  	go func() {
   727  		test.udp.handlePing(&v5wire.Ping{ENRSeq: 1}, remote.ID(), test.remoteaddr)
   728  		done <- struct{}{}
   729  	}()
   730  	test.waitPacketOut(func(p *v5wire.Pong, addr netip.AddrPort, _ v5wire.Nonce) {
   731  		if len(p.ToIP) == net.IPv6len {
   732  			t.Error("Received untruncated ip address")
   733  		}
   734  		if len(p.ToIP) != net.IPv4len {
   735  			t.Errorf("Received ip address with incorrect length: %d", len(p.ToIP))
   736  		}
   737  		if !p.ToIP.Equal(rawIP.AsSlice()) {
   738  			t.Errorf("Received incorrect ip address: wanted %s but received %s", rawIP.String(), p.ToIP.String())
   739  		}
   740  	})
   741  	<-done
   742  }
   743  
   744  // udpV5Test is the framework for all tests above.
   745  // It runs the UDPv5 transport on a virtual socket and allows testing outgoing packets.
   746  type udpV5Test struct {
   747  	t                   *testing.T
   748  	pipe                *dgramPipe
   749  	table               *Table
   750  	db                  *enode.DB
   751  	udp                 *UDPv5
   752  	localkey, remotekey *ecdsa.PrivateKey
   753  	remoteaddr          netip.AddrPort
   754  	nodesByID           map[enode.ID]*enode.LocalNode
   755  	nodesByIP           map[netip.Addr]*enode.LocalNode
   756  }
   757  
   758  // testCodec is the packet encoding used by protocol tests. This codec does not perform encryption.
   759  type testCodec struct {
   760  	test *udpV5Test
   761  	id   enode.ID
   762  	ctr  uint64
   763  
   764  	sentChallenges map[enode.ID]*v5wire.Whoareyou
   765  }
   766  
   767  type testCodecFrame struct {
   768  	NodeID  enode.ID
   769  	AuthTag v5wire.Nonce
   770  	Ptype   byte
   771  	Packet  rlp.RawValue
   772  }
   773  
   774  func (c *testCodec) Encode(toID enode.ID, addr string, p v5wire.Packet, _ *v5wire.Whoareyou) ([]byte, v5wire.Nonce, error) {
   775  	// To match the behavior of v5wire.Codec, we return the cached encoding of
   776  	// WHOAREYOU challenges.
   777  	if wp, ok := p.(*v5wire.Whoareyou); ok && len(wp.Encoded) > 0 {
   778  		return wp.Encoded, wp.Nonce, nil
   779  	}
   780  
   781  	c.ctr++
   782  	var authTag v5wire.Nonce
   783  	binary.BigEndian.PutUint64(authTag[:], c.ctr)
   784  	penc, _ := rlp.EncodeToBytes(p)
   785  	frame, err := rlp.EncodeToBytes(testCodecFrame{c.id, authTag, p.Kind(), penc})
   786  	if err != nil {
   787  		return frame, authTag, err
   788  	}
   789  
   790  	// Store recently sent challenges.
   791  	if w, ok := p.(*v5wire.Whoareyou); ok {
   792  		w.Nonce = authTag
   793  		w.Encoded = frame
   794  		if c.sentChallenges == nil {
   795  			c.sentChallenges = make(map[enode.ID]*v5wire.Whoareyou)
   796  		}
   797  		c.sentChallenges[toID] = w
   798  	}
   799  	return frame, authTag, err
   800  }
   801  
   802  func (c *testCodec) CurrentChallenge(id enode.ID, addr string) *v5wire.Whoareyou {
   803  	return c.sentChallenges[id]
   804  }
   805  
   806  func (c *testCodec) Decode(input []byte, addr string) (enode.ID, *enode.Node, v5wire.Packet, error) {
   807  	frame, p, err := c.decodeFrame(input)
   808  	if err != nil {
   809  		return enode.ID{}, nil, nil, err
   810  	}
   811  	return frame.NodeID, nil, p, nil
   812  }
   813  
   814  func (c *testCodec) SessionNode(id enode.ID, addr string) *enode.Node {
   815  	return c.test.nodesByID[id].Node()
   816  }
   817  
   818  func (c *testCodec) decodeFrame(input []byte) (frame testCodecFrame, p v5wire.Packet, err error) {
   819  	if err = rlp.DecodeBytes(input, &frame); err != nil {
   820  		return frame, nil, fmt.Errorf("invalid frame: %v", err)
   821  	}
   822  	switch frame.Ptype {
   823  	case v5wire.UnknownPacket:
   824  		dec := new(v5wire.Unknown)
   825  		err = rlp.DecodeBytes(frame.Packet, &dec)
   826  		p = dec
   827  	case v5wire.WhoareyouPacket:
   828  		dec := new(v5wire.Whoareyou)
   829  		err = rlp.DecodeBytes(frame.Packet, &dec)
   830  		p = dec
   831  	default:
   832  		p, err = v5wire.DecodeMessage(frame.Ptype, frame.Packet)
   833  	}
   834  	return frame, p, err
   835  }
   836  
   837  func newUDPV5Test(t *testing.T) *udpV5Test {
   838  	test := &udpV5Test{
   839  		t:          t,
   840  		pipe:       newpipe(),
   841  		localkey:   newkey(),
   842  		remotekey:  newkey(),
   843  		remoteaddr: netip.MustParseAddrPort("10.0.1.99:30303"),
   844  		nodesByID:  make(map[enode.ID]*enode.LocalNode),
   845  		nodesByIP:  make(map[netip.Addr]*enode.LocalNode),
   846  	}
   847  	test.db, _ = enode.OpenDB("")
   848  	ln := enode.NewLocalNode(test.db, test.localkey)
   849  	ln.SetStaticIP(net.IP{10, 0, 0, 1})
   850  	ln.Set(enr.UDP(30303))
   851  	test.udp, _ = ListenV5(test.pipe, ln, Config{
   852  		PrivateKey:   test.localkey,
   853  		Log:          testlog.Logger(t, log.LvlTrace),
   854  		ValidSchemes: enode.ValidSchemesForTesting,
   855  	})
   856  	test.udp.codec = &testCodec{test: test, id: ln.ID()}
   857  	test.table = test.udp.tab
   858  	test.nodesByID[ln.ID()] = ln
   859  	// Wait for initial refresh so the table doesn't send unexpected findnode.
   860  	<-test.table.initDone
   861  	return test
   862  }
   863  
   864  // handles a packet as if it had been sent to the transport.
   865  func (test *udpV5Test) packetIn(packet v5wire.Packet) {
   866  	test.t.Helper()
   867  	test.packetInFrom(test.remotekey, test.remoteaddr, packet)
   868  }
   869  
   870  // packetInFrom handles a packet as if it had been sent to the transport by the key/endpoint.
   871  func (test *udpV5Test) packetInFrom(key *ecdsa.PrivateKey, addr netip.AddrPort, packet v5wire.Packet) {
   872  	test.t.Helper()
   873  
   874  	ln := test.getNode(key, addr)
   875  	codec := &testCodec{test: test, id: ln.ID()}
   876  	enc, _, err := codec.Encode(test.udp.Self().ID(), addr.String(), packet, nil)
   877  	if err != nil {
   878  		test.t.Errorf("%s encode error: %v", packet.Name(), err)
   879  	}
   880  	if test.udp.dispatchReadPacket(addr, enc) {
   881  		<-test.udp.readNextCh // unblock UDPv5.dispatch
   882  	}
   883  }
   884  
   885  // getNode ensures the test knows about a node at the given endpoint.
   886  func (test *udpV5Test) getNode(key *ecdsa.PrivateKey, addr netip.AddrPort) *enode.LocalNode {
   887  	id := v4wire.EncodePubkey(&key.PublicKey).ID()
   888  	ln := test.nodesByID[id]
   889  	if ln == nil {
   890  		db, _ := enode.OpenDB("")
   891  		ln = enode.NewLocalNode(db, key)
   892  		ln.SetStaticIP(addr.Addr().AsSlice())
   893  		ln.Set(enr.UDP(addr.Port()))
   894  		test.nodesByID[id] = ln
   895  	}
   896  	test.nodesByIP[addr.Addr()] = ln
   897  	return ln
   898  }
   899  
   900  // waitPacketOut waits for the next output packet and handles it using the given 'validate'
   901  // function. The function must be of type func (X, netip.AddrPort, v5wire.Nonce) where X is
   902  // assignable to packetV5.
   903  func (test *udpV5Test) waitPacketOut(validate interface{}) (closed bool) {
   904  	test.t.Helper()
   905  
   906  	fn := reflect.ValueOf(validate)
   907  	exptype := fn.Type().In(0)
   908  
   909  	dgram, err := test.pipe.receive()
   910  	if err == errClosed {
   911  		return true
   912  	}
   913  	if err == errTimeout {
   914  		test.t.Fatalf("timed out waiting for %v", exptype)
   915  		return false
   916  	}
   917  	ln := test.nodesByIP[dgram.to.Addr()]
   918  	if ln == nil {
   919  		test.t.Fatalf("attempt to send to non-existing node %v", &dgram.to)
   920  		return false
   921  	}
   922  	codec := &testCodec{test: test, id: ln.ID()}
   923  	frame, p, err := codec.decodeFrame(dgram.data)
   924  	if err != nil {
   925  		test.t.Errorf("sent packet decode error: %v", err)
   926  		return false
   927  	}
   928  	if !reflect.TypeOf(p).AssignableTo(exptype) {
   929  		test.t.Errorf("sent packet type mismatch, got: %v, want: %v", reflect.TypeOf(p), exptype)
   930  		return false
   931  	}
   932  	fn.Call([]reflect.Value{reflect.ValueOf(p), reflect.ValueOf(dgram.to), reflect.ValueOf(frame.AuthTag)})
   933  	return false
   934  }
   935  
   936  func (test *udpV5Test) close() {
   937  	test.t.Helper()
   938  
   939  	test.udp.Close()
   940  	test.db.Close()
   941  	for id, n := range test.nodesByID {
   942  		if id != test.udp.Self().ID() {
   943  			n.Database().Close()
   944  		}
   945  	}
   946  	if len(test.pipe.queue) != 0 {
   947  		test.t.Fatalf("%d unmatched UDP packets in queue", len(test.pipe.queue))
   948  	}
   949  }