p2p/discover: refactor node and endpoint representation

cmd/devp2p/internal/v4test/framework.go

@@ -110,7 +110,7 @@ func (te *testenv) localEndpoint(c net.PacketConn) v4wire.Endpoint {
 }
 
 func (te *testenv) remoteEndpoint() v4wire.Endpoint {
-	return v4wire.NewEndpoint(te.remoteAddr, 0)
+	return v4wire.NewEndpoint(te.remoteAddr.AddrPort(), 0)
 }
 
 func contains(ns []v4wire.Node, key v4wire.Pubkey) bool {

p2p/discover/common.go

@@ -22,6 +22,7 @@ import (
 	"encoding/binary"
 	"math/rand"
 	"net"
+	"net/netip"
 	"sync"
 	"time"
 
@@ -34,8 +35,8 @@ import (
 
 // UDPConn is a network connection on which discovery can operate.
 type UDPConn interface {
-	ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error)
-	WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error)
+	ReadFromUDPAddrPort(b []byte) (n int, addr netip.AddrPort, err error)
+	WriteToUDPAddrPort(b []byte, addr netip.AddrPort) (n int, err error)
 	Close() error
 	LocalAddr() net.Addr
 }
@@ -94,7 +95,7 @@ func ListenUDP(c UDPConn, ln *enode.LocalNode, cfg Config) (*UDPv4, error) {
 // channel if configured.
 type ReadPacket struct {
 	Data []byte
-	Addr *net.UDPAddr
+	Addr netip.AddrPort
 }
 
 type randomSource interface {

p2p/discover/lookup.go

@@ -29,16 +29,16 @@ import (
 // not need to be an actual node identifier.
 type lookup struct {
 	tab         *Table
-	queryfunc   func(*node) ([]*node, error)
-	replyCh     chan []*node
+	queryfunc   queryFunc
+	replyCh     chan []*enode.Node
 	cancelCh    <-chan struct{}
 	asked, seen map[enode.ID]bool
 	result      nodesByDistance
-	replyBuffer []*node
+	replyBuffer []*enode.Node
 	queries     int
 }
 
-type queryFunc func(*node) ([]*node, error)
+type queryFunc func(*enode.Node) ([]*enode.Node, error)
 
 func newLookup(ctx context.Context, tab *Table, target enode.ID, q queryFunc) *lookup {
 	it := &lookup{
@@ -47,7 +47,7 @@ func newLookup(ctx context.Context, tab *Table, target enode.ID, q queryFunc) *l
 		asked:     make(map[enode.ID]bool),
 		seen:      make(map[enode.ID]bool),
 		result:    nodesByDistance{target: target},
-		replyCh:   make(chan []*node, alpha),
+		replyCh:   make(chan []*enode.Node, alpha),
 		cancelCh:  ctx.Done(),
 		queries:   -1,
 	}
@@ -61,7 +61,7 @@ func newLookup(ctx context.Context, tab *Table, target enode.ID, q queryFunc) *l
 func (it *lookup) run() []*enode.Node {
 	for it.advance() {
 	}
-	return unwrapNodes(it.result.entries)
+	return it.result.entries
 }
 
 // advance advances the lookup until any new nodes have been found.
@@ -139,7 +139,7 @@ func (it *lookup) slowdown() {
 	}
 }
 
-func (it *lookup) query(n *node, reply chan<- []*node) {
+func (it *lookup) query(n *enode.Node, reply chan<- []*enode.Node) {
 	r, err := it.queryfunc(n)
 	if !errors.Is(err, errClosed) { // avoid recording failures on shutdown.
 		success := len(r) > 0
@@ -154,7 +154,7 @@ func (it *lookup) query(n *node, reply chan<- []*node) {
 // lookupIterator performs lookup operations and iterates over all seen nodes.
 // When a lookup finishes, a new one is created through nextLookup.
 type lookupIterator struct {
-	buffer     []*node
+	buffer     []*enode.Node
 	nextLookup lookupFunc
 	ctx        context.Context
 	cancel     func()
@@ -173,7 +173,7 @@ func (it *lookupIterator) Node() *enode.Node {
 	if len(it.buffer) == 0 {
 		return nil
 	}
-	return unwrapNode(it.buffer[0])
+	return it.buffer[0]
 }
 
 // Next moves to the next node.

p2p/discover/metrics.go

@@ -18,7 +18,7 @@ package discover
 
 import (
 	"fmt"
-	"net"
+	"net/netip"
 
 	"github.com/ethereum/go-ethereum/metrics"
 )
@@ -58,16 +58,16 @@ func newMeteredConn(conn UDPConn) UDPConn {
 	return &meteredUdpConn{UDPConn: conn}
 }
 
-// ReadFromUDP delegates a network read to the underlying connection, bumping the udp ingress traffic meter along the way.
-func (c *meteredUdpConn) ReadFromUDP(b []byte) (n int, addr *net.UDPAddr, err error) {
-	n, addr, err = c.UDPConn.ReadFromUDP(b)
+// ReadFromUDPAddrPort delegates a network read to the underlying connection, bumping the udp ingress traffic meter along the way.
+func (c *meteredUdpConn) ReadFromUDPAddrPort(b []byte) (n int, addr netip.AddrPort, err error) {
+	n, addr, err = c.UDPConn.ReadFromUDPAddrPort(b)
 	ingressTrafficMeter.Mark(int64(n))
 	return n, addr, err
 }
 
-// Write delegates a network write to the underlying connection, bumping the udp egress traffic meter along the way.
-func (c *meteredUdpConn) WriteToUDP(b []byte, addr *net.UDPAddr) (n int, err error) {
-	n, err = c.UDPConn.WriteToUDP(b, addr)
+// WriteToUDP delegates a network write to the underlying connection, bumping the udp egress traffic meter along the way.
+func (c *meteredUdpConn) WriteToUDP(b []byte, addr netip.AddrPort) (n int, err error) {
+	n, err = c.UDPConn.WriteToUDPAddrPort(b, addr)
 	egressTrafficMeter.Mark(int64(n))
 	return n, err
 }

p2p/discover/node.go

@@ -21,7 +21,8 @@ import (
 	"crypto/elliptic"
 	"errors"
 	"math/big"
-	"net"
+	"slices"
+	"sort"
 	"time"
 
 	"github.com/ethereum/go-ethereum/common/math"
@@ -37,9 +38,8 @@ type BucketNode struct {
 	Live          bool        `json:"live"`
 }
 
-// node represents a host on the network.
-// The fields of Node may not be modified.
-type node struct {
+// tableNode is an entry in Table.
+type tableNode struct {
 	*enode.Node
 	revalList       *revalidationList
 	addedToTable    time.Time // first time node was added to bucket or replacement list
@@ -75,34 +75,59 @@ func (e encPubkey) id() enode.ID {
 	return enode.ID(crypto.Keccak256Hash(e[:]))
 }
 
-func wrapNode(n *enode.Node) *node {
-	return &node{Node: n}
-}
-
-func wrapNodes(ns []*enode.Node) []*node {
-	result := make([]*node, len(ns))
+func unwrapNodes(ns []*tableNode) []*enode.Node {
+	result := make([]*enode.Node, len(ns))
 	for i, n := range ns {
-		result[i] = wrapNode(n)
+		result[i] = n.Node
 	}
 	return result
 }
 
-func unwrapNode(n *node) *enode.Node {
-	return n.Node
+func (n *tableNode) String() string {
+	return n.Node.String()
+}
+
+// nodesByDistance is a list of nodes, ordered by distance to target.
+type nodesByDistance struct {
+	entries []*enode.Node
+	target  enode.ID
 }
 
-func unwrapNodes(ns []*node) []*enode.Node {
-	result := make([]*enode.Node, len(ns))
-	for i, n := range ns {
-		result[i] = unwrapNode(n)
+// push adds the given node to the list, keeping the total size below maxElems.
+func (h *nodesByDistance) push(n *enode.Node, maxElems int) {
+	ix := sort.Search(len(h.entries), func(i int) bool {
+		return enode.DistCmp(h.target, h.entries[i].ID(), n.ID()) > 0
+	})
+
+	end := len(h.entries)
+	if len(h.entries) < maxElems {
+		h.entries = append(h.entries, n)
+	}
+	if ix < end {
+		// Slide existing entries down to make room.
+		// This will overwrite the entry we just appended.
+		copy(h.entries[ix+1:], h.entries[ix:])
+		h.entries[ix] = n
 	}
-	return result
 }
 
-func (n *node) addr() *net.UDPAddr {
-	return &net.UDPAddr{IP: n.IP(), Port: n.UDP()}
+type nodeType interface {
+	ID() enode.ID
 }
 
-func (n *node) String() string {
-	return n.Node.String()
+// containsID reports whether ns contains a node with the given ID.
+func containsID[N nodeType](ns []N, id enode.ID) bool {
+	for _, n := range ns {
+		if n.ID() == id {
+			return true
+		}
+	}
+	return false
+}
+
+// deleteNode removes a node from the list.
+func deleteNode[N nodeType](list []N, id enode.ID) []N {
+	return slices.DeleteFunc(list, func(n N) bool {
+		return n.ID() == id
+	})
 }