p2p/enode: track IPv4 and IPv6 address separately

LocalNode predicts the local node's UDP endpoint and updates the record.
This change makes it predict IPv4 and IPv6 endpoints separately since
they can now be in the record at the same time.

p2p/enode/localnode.go

@@ -51,9 +51,13 @@ type LocalNode struct {
 	mu        sync.Mutex
 	seq       uint64
 	entries   map[string]enr.Entry
-	udpTrack    *netutil.IPTracker // predicts external UDP endpoint
+	endpoint4 lnEndpoint
-	staticIP    net.IP
+	endpoint6 lnEndpoint
-	fallbackIP  net.IP
+}
+
+type lnEndpoint struct {
+	track                *netutil.IPTracker
+	staticIP, fallbackIP net.IP
 	fallbackUDP          int
 }
 
@@ -63,8 +67,13 @@ func NewLocalNode(db *DB, key *ecdsa.PrivateKey) *LocalNode {
 		id:      PubkeyToIDV4(&key.PublicKey),
 		db:      db,
 		key:     key,
-		udpTrack: netutil.NewIPTracker(iptrackWindow, iptrackContactWindow, iptrackMinStatements),
 		entries: make(map[string]enr.Entry),
+		endpoint4: lnEndpoint{
+			track: netutil.NewIPTracker(iptrackWindow, iptrackContactWindow, iptrackMinStatements),
+		},
+		endpoint6: lnEndpoint{
+			track: netutil.NewIPTracker(iptrackWindow, iptrackContactWindow, iptrackMinStatements),
+		},
 	}
 	ln.seq = db.localSeq(ln.id)
 	ln.invalidate()
@@ -89,13 +98,22 @@ func (ln *LocalNode) Node() *Node {
 	return ln.cur.Load().(*Node)
 }
 
+// Seq returns the current sequence number of the local node record.
+func (ln *LocalNode) Seq() uint64 {
+	ln.mu.Lock()
+	defer ln.mu.Unlock()
+
+	return ln.seq
+}
+
 // ID returns the local node ID.
 func (ln *LocalNode) ID() ID {
 	return ln.id
 }
 
-// Set puts the given entry into the local record, overwriting
+// Set puts the given entry into the local record, overwriting any existing value.
-// any existing value.
+// Use Set*IP and SetFallbackUDP to set IP addresses and UDP port, otherwise they'll
+// be overwritten by the endpoint predictor.
 func (ln *LocalNode) Set(e enr.Entry) {
 	ln.mu.Lock()
 	defer ln.mu.Unlock()
@@ -127,13 +145,20 @@ func (ln *LocalNode) delete(e enr.Entry) {
 	}
 }
 
+func (ln *LocalNode) endpointForIP(ip net.IP) *lnEndpoint {
+	if ip.To4() != nil {
+		return &ln.endpoint4
+	}
+	return &ln.endpoint6
+}
+
 // SetStaticIP sets the local IP to the given one unconditionally.
 // This disables endpoint prediction.
 func (ln *LocalNode) SetStaticIP(ip net.IP) {
 	ln.mu.Lock()
 	defer ln.mu.Unlock()
 
-	ln.staticIP = ip
+	ln.endpointForIP(ip).staticIP = ip
 	ln.updateEndpoints()
 }
 
@@ -143,17 +168,18 @@ func (ln *LocalNode) SetFallbackIP(ip net.IP) {
 	ln.mu.Lock()
 	defer ln.mu.Unlock()
 
-	ln.fallbackIP = ip
+	ln.endpointForIP(ip).fallbackIP = ip
 	ln.updateEndpoints()
 }
 
-// SetFallbackUDP sets the last-resort UDP port. This port is used
+// SetFallbackUDP sets the last-resort UDP-on-IPv4 port. This port is used
 // if no endpoint prediction can be made.
 func (ln *LocalNode) SetFallbackUDP(port int) {
 	ln.mu.Lock()
 	defer ln.mu.Unlock()
 
-	ln.fallbackUDP = port
+	ln.endpoint4.fallbackUDP = port
+	ln.endpoint6.fallbackUDP = port
 	ln.updateEndpoints()
 }
 
@@ -163,7 +189,7 @@ func (ln *LocalNode) UDPEndpointStatement(fromaddr, endpoint *net.UDPAddr) {
 	ln.mu.Lock()
 	defer ln.mu.Unlock()
 
-	ln.udpTrack.AddStatement(fromaddr.String(), endpoint.String())
+	ln.endpointForIP(endpoint.IP).track.AddStatement(fromaddr.String(), endpoint.String())
 	ln.updateEndpoints()
 }
 
@@ -173,34 +199,52 @@ func (ln *LocalNode) UDPContact(toaddr *net.UDPAddr) {
 	ln.mu.Lock()
 	defer ln.mu.Unlock()
 
-	ln.udpTrack.AddContact(toaddr.String())
+	ln.endpointForIP(toaddr.IP).track.AddContact(toaddr.String())
 	ln.updateEndpoints()
 }
 
+// updateEndpoints updates the record with predicted endpoints.
 func (ln *LocalNode) updateEndpoints() {
-	// Determine the endpoints.
+	ip4, udp4 := ln.endpoint4.get()
-	newIP := ln.fallbackIP
+	ip6, udp6 := ln.endpoint6.get()
-	newUDP := ln.fallbackUDP
-	if ln.staticIP != nil {
-		newIP = ln.staticIP
-	} else if ip, port := predictAddr(ln.udpTrack); ip != nil {
-		newIP = ip
-		newUDP = port
-	}
 
-	// Update the record.
+	if ip4 != nil && !ip4.IsUnspecified() {
-	if newIP != nil && !newIP.IsUnspecified() {
+		ln.set(enr.IPv4(ip4))
-		ln.set(enr.IP(newIP))
+	} else {
-		if newUDP != 0 {
+		ln.delete(enr.IPv4{})
-			ln.set(enr.UDP(newUDP))
+	}
+	if ip6 != nil && !ip6.IsUnspecified() {
+		ln.set(enr.IPv6(ip6))
+	} else {
+		ln.delete(enr.IPv6{})
+	}
+	if udp4 != 0 {
+		ln.set(enr.UDP(udp4))
 	} else {
 		ln.delete(enr.UDP(0))
 	}
+	if udp6 != 0 && udp6 != udp4 {
+		ln.set(enr.UDP6(udp6))
 	} else {
-		ln.delete(enr.IP{})
+		ln.delete(enr.UDP6(0))
 	}
 }
 
+// get returns the endpoint with highest precedence.
+func (e *lnEndpoint) get() (newIP net.IP, newPort int) {
+	newPort = e.fallbackUDP
+	if e.fallbackIP != nil {
+		newIP = e.fallbackIP
+	}
+	if e.staticIP != nil {
+		newIP = e.staticIP
+	} else if ip, port := predictAddr(e.track); ip != nil {
+		newIP = ip
+		newPort = port
+	}
+	return newIP, newPort
+}
+
 // predictAddr wraps IPTracker.PredictEndpoint, converting from its string-based
 // endpoint representation to IP and port types.
 func predictAddr(t *netutil.IPTracker) (net.IP, int) {

p2p/enode/localnode_test.go

@@ -17,10 +17,13 @@
 package enode
 
 import (
+	"math/rand"
+	"net"
 	"testing"
 
 	"github.com/ethereum/go-ethereum/crypto"
 	"github.com/ethereum/go-ethereum/p2p/enr"
+	"github.com/stretchr/testify/assert"
 )
 
 func newLocalNodeForTesting() (*LocalNode, *DB) {
@@ -74,3 +77,46 @@ func TestLocalNodeSeqPersist(t *testing.T) {
 		t.Fatalf("wrong seq %d on instance with changed key, want 1", s)
 	}
 }
+
+// This test checks behavior of the endpoint predictor.
+func TestLocalNodeEndpoint(t *testing.T) {
+	var (
+		fallback  = &net.UDPAddr{IP: net.IP{127, 0, 0, 1}, Port: 80}
+		predicted = &net.UDPAddr{IP: net.IP{127, 0, 1, 2}, Port: 81}
+		staticIP  = net.IP{127, 0, 1, 2}
+	)
+	ln, db := newLocalNodeForTesting()
+	defer db.Close()
+
+	// Nothing is set initially.
+	assert.Equal(t, net.IP(nil), ln.Node().IP())
+	assert.Equal(t, 0, ln.Node().UDP())
+	assert.Equal(t, uint64(1), ln.Node().Seq())
+
+	// Set up fallback address.
+	ln.SetFallbackIP(fallback.IP)
+	ln.SetFallbackUDP(fallback.Port)
+	assert.Equal(t, fallback.IP, ln.Node().IP())
+	assert.Equal(t, fallback.Port, ln.Node().UDP())
+	assert.Equal(t, uint64(2), ln.Node().Seq())
+
+	// Add endpoint statements from random hosts.
+	for i := 0; i < iptrackMinStatements; i++ {
+		assert.Equal(t, fallback.IP, ln.Node().IP())
+		assert.Equal(t, fallback.Port, ln.Node().UDP())
+		assert.Equal(t, uint64(2), ln.Node().Seq())
+
+		from := &net.UDPAddr{IP: make(net.IP, 4), Port: 90}
+		rand.Read(from.IP)
+		ln.UDPEndpointStatement(from, predicted)
+	}
+	assert.Equal(t, predicted.IP, ln.Node().IP())
+	assert.Equal(t, predicted.Port, ln.Node().UDP())
+	assert.Equal(t, uint64(3), ln.Node().Seq())
+
+	// Static IP overrides prediction.
+	ln.SetStaticIP(staticIP)
+	assert.Equal(t, staticIP, ln.Node().IP())
+	assert.Equal(t, fallback.Port, ln.Node().UDP())
+	assert.Equal(t, uint64(4), ln.Node().Seq())
+}

p2p/enode/node.go

@@ -68,11 +68,19 @@ func (n *Node) Load(k enr.Entry) error {
 	return n.r.Load(k)
 }
 
-// IP returns the IP address of the node.
+// IP returns the IP address of the node. This prefers IPv4 addresses.
 func (n *Node) IP() net.IP {
-	var ip net.IP
+	var (
-	n.Load((*enr.IP)(&ip))
+		ip4 enr.IPv4
-	return ip
+		ip6 enr.IPv6
+	)
+	if n.Load(&ip4) == nil {
+		return net.IP(ip4)
+	}
+	if n.Load(&ip6) == nil {
+		return net.IP(ip6)
+	}
+	return nil
 }
 
 // UDP returns the UDP port of the node.

p2p/enode/node_test.go

@@ -46,7 +46,7 @@ func TestPythonInterop(t *testing.T) {
 	var (
 		wantID  = HexID("a448f24c6d18e575453db13171562b71999873db5b286df957af199ec94617f7")
 		wantSeq = uint64(1)
-		wantIP  = enr.IP{127, 0, 0, 1}
+		wantIP  = enr.IPv4{127, 0, 0, 1}
 		wantUDP = enr.UDP(30303)
 	)
 	if n.Seq() != wantSeq {
@@ -55,7 +55,7 @@ func TestPythonInterop(t *testing.T) {
 	if n.ID() != wantID {
 		t.Errorf("wrong id: got %x, want %x", n.ID(), wantID)
 	}
-	want := map[enr.Entry]interface{}{new(enr.IP): &wantIP, new(enr.UDP): &wantUDP}
+	want := map[enr.Entry]interface{}{new(enr.IPv4): &wantIP, new(enr.UDP): &wantUDP}
 	for k, v := range want {
 		desc := fmt.Sprintf("loading key %q", k.ENRKey())
 		if assert.NoError(t, n.Load(k), desc) {

p2p/enode/urlv4.go

@@ -81,7 +81,7 @@ func ParseV4(rawurl string) (*Node, error) {
 // contained in the node has a zero-length signature.
 func NewV4(pubkey *ecdsa.PublicKey, ip net.IP, tcp, udp int) *Node {
 	var r enr.Record
-	if ip != nil {
+	if len(ip) > 0 {
 		r.Set(enr.IP(ip))
 	}
 	if udp != 0 {
@@ -126,10 +126,6 @@ func parseComplete(rawurl string) (*Node, error) {
 	if ip = net.ParseIP(host); ip == nil {
 		return nil, errors.New("invalid IP address")
 	}
-	// Ensure the IP is 4 bytes long for IPv4 addresses.
-	if ipv4 := ip.To4(); ipv4 != nil {
-		ip = ipv4
-	}
 	// Parse the port numbers.
 	if tcpPort, err = strconv.ParseUint(port, 10, 16); err != nil {
 		return nil, errors.New("invalid port")