Set external IP and mapping port to es-node ENR when it behind a NAT and the advertise IP is not set

ethstorage/p2p/discovery.go

@@ -23,6 +23,7 @@ import (
 	"github.com/libp2p/go-libp2p/core/crypto"
 	"github.com/libp2p/go-libp2p/core/network"
 	"github.com/libp2p/go-libp2p/core/peer"
+	"github.com/libp2p/go-libp2p/p2p/net/nat"
 	"github.com/multiformats/go-multiaddr"
 	ma "github.com/multiformats/go-multiaddr"
 )
@@ -49,9 +50,24 @@ func (conf *Config) Discovery(log log.Logger, l1ChainID uint64, tcpPort uint16,
 	// use the geth curve definition. Same crypto, but geth needs to detect it as *their* definition of the curve.
 	priv.Curve = gcrypto.S256()
 	localNode := enode.NewLocalNode(conf.DiscoveryDB, priv)
+	if conf.AdvertiseUDPPort != 0 { // explicitly advertised port gets priority
+		localNode.SetFallbackUDP(int(conf.AdvertiseUDPPort))
+	} else if conf.ListenUDPPort != 0 { // otherwise default to the port we configured it to listen on
+		localNode.SetFallbackUDP(int(conf.ListenUDPPort))
+	}
+	if conf.AdvertiseTCPPort != 0 { // explicitly advertised port gets priority
+		localNode.Set(enr.TCP(conf.AdvertiseTCPPort))
+	} else if tcpPort != 0 { // otherwise try to pick up whatever port LibP2P binded to (listen port, or dynamically picked)
+		localNode.Set(enr.TCP(tcpPort))
+	} else if conf.ListenTCPPort != 0 { // otherwise default to the port we configured it to listen on
+		localNode.Set(enr.TCP(conf.ListenTCPPort))
+	} else {
+		return nil, nil, fmt.Errorf("no TCP port to put in discovery record")
+	}
 	if conf.AdvertiseIP != nil {
 		localNode.SetStaticIP(conf.AdvertiseIP)
 	} else {
+		end := false
 		for _, addr := range addrs {
 			ipStr, err := addr.ValueForProtocol(4)
 			if err != nil {
@@ -62,23 +78,24 @@ func (conf *Config) Discovery(log log.Logger, l1ChainID uint64, tcpPort uint16,
 				continue
 			}
 			localNode.SetStaticIP(ip)
+			end = true
 			break
 		}
-		// TODO: if no external IP found, use NAT protocol to find external IP
-	}
-	if conf.AdvertiseUDPPort != 0 { // explicitly advertised port gets priority
-		localNode.SetFallbackUDP(int(conf.AdvertiseUDPPort))
-	} else if conf.ListenUDPPort != 0 { // otherwise default to the port we configured it to listen on
-		localNode.SetFallbackUDP(int(conf.ListenUDPPort))
-	}
-	if conf.AdvertiseTCPPort != 0 { // explicitly advertised port gets priority
-		localNode.Set(enr.TCP(conf.AdvertiseTCPPort))
-	} else if tcpPort != 0 { // otherwise try to pick up whatever port LibP2P binded to (listen port, or dynamically picked)
-		localNode.Set(enr.TCP(tcpPort))
-	} else if conf.ListenTCPPort != 0 { // otherwise default to the port we configured it to listen on
-		localNode.Set(enr.TCP(conf.ListenTCPPort))
-	} else {
-		return nil, nil, fmt.Errorf("no TCP port to put in discovery record")
+
+		if !end {
+			intTcpPort := localNode.Node().TCP()
+			intUdpPort := localNode.Node().UDP()
+			ip, extTcpPort, extUdpPort, err := addNATMappings(intTcpPort, intUdpPort)
+			if err == nil {
+				localNode.SetStaticIP(ip)
+				localNode.Set(enr.TCP(extTcpPort))
+				localNode.Set(enr.UDP(extUdpPort))
+				log.Info("Add mappings to NAT", "external IP", ip.String(), "internal tcp port", intTcpPort,
+					"external tcp port", extTcpPort, "internal udp port", intUdpPort, "external udp port", extUdpPort)
+			} else {
+				log.Warn("", "error", err.Error())
+			}
+		}
 	}
 	dat := protocol.EthStorageENRData{
 		ChainID: l1ChainID,
@@ -149,6 +166,30 @@ func (v *Secp256k1) DecodeRLP(s *rlp.Stream) error {
 	return nil
 }
 
+func addNATMappings(tcpPort, udpPort int) (net.IP, int, int, error) {
+	n, err := nat.DiscoverNAT(context.Background())
+	if err != nil {
+		return nil, 0, 0, fmt.Errorf("no external address set to ENR as no NAT service found: %s", err.Error())
+	}
+	tcpMapping, err := n.NewMapping("tcp", tcpPort)
+	if err != nil {
+		return nil, 0, 0, fmt.Errorf("fail to add tcpMapping for tcp port to NAT: %s", err.Error())
+	}
+	addr, err := tcpMapping.ExternalAddr()
+	if err != nil {
+		return nil, 0, 0, fmt.Errorf("get external IP from NAT tcpMapping fail: %s", err.Error())
+	}
+	tcpAddr, err := net.ResolveTCPAddr("tcp", addr.String())
+	if err != nil {
+		return nil, 0, 0, fmt.Errorf("fail to resolve external address: %s", err.Error())
+	}
+	udpMapping, err := n.NewMapping("udp", udpPort)
+	if err != nil {
+		return nil, 0, 0, fmt.Errorf("fail to add tcpMapping for udp port to NAT: %s", err.Error())
+	}
+	return tcpAddr.IP, tcpMapping.ExternalPort(), udpMapping.ExternalPort(), nil
+}
+
 func enrToAddrInfo(r *enode.Node) (*peer.AddrInfo, *crypto.Secp256k1PublicKey, error) {
 	ip := r.IP()
 	ipScheme := "ip4"