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