From eba7a41a34f0325184c145dff6758c11ebacb888 Mon Sep 17 00:00:00 2001
From: Felix Lange <fjl@twurst.com>
Date: Mon, 1 Oct 2018 16:25:43 +0200
Subject: [PATCH] p2p/enode: clean up LocalNode endpoint setters

---
 p2p/enode/localnode.go | 43 ++++++++++++++++++------------------------
 p2p/server.go          |  3 ++-
 2 files changed, 20 insertions(+), 26 deletions(-)

diff --git a/p2p/enode/localnode.go b/p2p/enode/localnode.go
index d3944f4ef8..d4da0aa731 100644
--- a/p2p/enode/localnode.go
+++ b/p2p/enode/localnode.go
@@ -38,8 +38,9 @@ const (
 	iptrackContactWindow = 10 * time.Minute
 )
 
-// LocalNode produces the signed node record of a local node, i.e. a node
-// run in the current process.
+// LocalNode produces the signed node record of a local node, i.e. a node run in the
+// current process. Setting ENR entries via the Set method updates the record. A new version
+// of the record is signed on demand when the Node method is called.
 type LocalNode struct {
 	cur atomic.Value // holds a non-nil node pointer while the record is up-to-date.
 	id  ID
@@ -52,11 +53,11 @@ type LocalNode struct {
 	entries     map[string]enr.Entry
 	udpTrack    *netutil.IPTracker // predicts external UDP endpoint
 	staticIP    net.IP
-	staticTCP   int
 	fallbackIP  net.IP
 	fallbackUDP int
 }
 
+// NewLocalNode creates a local node.
 func NewLocalNode(db *DB, key *ecdsa.PrivateKey) *LocalNode {
 	ln := &LocalNode{
 		id:       PubkeyToIDV4(&key.PublicKey),
@@ -126,6 +127,8 @@ func (ln *LocalNode) delete(e enr.Entry) {
 	}
 }
 
+// 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()
@@ -134,6 +137,8 @@ func (ln *LocalNode) SetStaticIP(ip net.IP) {
 	ln.updateEndpoints()
 }
 
+// SetFallbackIP sets the last-resort IP address. This address is used
+// if no endpoint prediction can be made and no static IP is set.
 func (ln *LocalNode) SetFallbackIP(ip net.IP) {
 	ln.mu.Lock()
 	defer ln.mu.Unlock()
@@ -142,6 +147,8 @@ func (ln *LocalNode) SetFallbackIP(ip net.IP) {
 	ln.updateEndpoints()
 }
 
+// SetFallbackUDP sets the last-resort UDP 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()
@@ -150,14 +157,6 @@ func (ln *LocalNode) SetFallbackUDP(port int) {
 	ln.updateEndpoints()
 }
 
-func (ln *LocalNode) SetStaticTCP(port int) {
-	ln.mu.Lock()
-	defer ln.mu.Unlock()
-
-	ln.staticTCP = port
-	ln.updateEndpoints()
-}
-
 // UDPEndpointStatement should be called whenever a statement about the local node's
 // UDP endpoint is received. It feeds the local endpoint predictor.
 func (ln *LocalNode) UDPEndpointStatement(fromaddr, endpoint *net.UDPAddr) {
@@ -182,30 +181,24 @@ func (ln *LocalNode) updateEndpoints() {
 	// Determine the endpoints.
 	newIP := ln.fallbackIP
 	newUDP := ln.fallbackUDP
-	if ip, port := predictAddr(ln.udpTrack); ip != nil {
-		newIP = ip
-		newUDP = port
-	}
 	if ln.staticIP != nil {
 		newIP = ln.staticIP
+	} else if ip, port := predictAddr(ln.udpTrack); ip != nil {
+		newIP = ip
+		newUDP = port
 	}
 
 	// Update the record.
 	if newIP != nil && !newIP.IsUnspecified() {
 		ln.set(enr.IP(newIP))
+		if newUDP != 0 {
+			ln.set(enr.UDP(newUDP))
+		} else {
+			ln.delete(enr.UDP(0))
+		}
 	} else {
 		ln.delete(enr.IP{})
 	}
-	if newUDP != 0 {
-		ln.set(enr.UDP(newUDP))
-	} else {
-		ln.delete(enr.UDP(0))
-	}
-	if ln.staticTCP != 0 {
-		ln.set(enr.TCP(ln.staticTCP))
-	} else {
-		ln.delete(enr.TCP(0))
-	}
 }
 
 // predictAddr wraps IPTracker.PredictEndpoint, converting from its string-based
diff --git a/p2p/server.go b/p2p/server.go
index 8e908f2e3a..c62a0e77b8 100644
--- a/p2p/server.go
+++ b/p2p/server.go
@@ -37,6 +37,7 @@ import (
 	"github.com/ethereum/go-ethereum/p2p/discover"
 	"github.com/ethereum/go-ethereum/p2p/discv5"
 	"github.com/ethereum/go-ethereum/p2p/enode"
+	"github.com/ethereum/go-ethereum/p2p/enr"
 	"github.com/ethereum/go-ethereum/p2p/nat"
 	"github.com/ethereum/go-ethereum/p2p/netutil"
 	"github.com/ethereum/go-ethereum/rlp"
@@ -572,7 +573,7 @@ func (srv *Server) setupListening() error {
 	laddr := listener.Addr().(*net.TCPAddr)
 	srv.ListenAddr = laddr.String()
 	srv.listener = listener
-	srv.localnode.SetStaticTCP(laddr.Port)
+	srv.localnode.Set(enr.TCP(laddr.Port))
 
 	srv.loopWG.Add(1)
 	go srv.listenLoop()