From 120be2f1eedc70084c4382d7c7c034e4c79d4ac0 Mon Sep 17 00:00:00 2001
From: Rafael Sampaio <5679073+r4f4ss@users.noreply.github.com>
Date: Fri, 4 Oct 2024 16:43:56 -0300
Subject: [PATCH] add portMapping to localNode creation

---
 cmd/shisui/main.go  | 59 ++++++++++++++++++++++++++++++++++++++++++++-
 p2p/discover/nat.go |  6 +++++
 p2p/nat/nat_stun.go |  4 +--
 3 files changed, 66 insertions(+), 3 deletions(-)

diff --git a/cmd/shisui/main.go b/cmd/shisui/main.go
index c320842ac9..c4eb346707 100644
--- a/cmd/shisui/main.go
+++ b/cmd/shisui/main.go
@@ -13,6 +13,7 @@ import (
 	"slices"
 	"strings"
 	"syscall"
+	"time"
 
 	"os"
 
@@ -249,8 +250,13 @@ func initDiscV5(config Config, conn discover.UDPConn) (*discover.UDPv5, *enode.L
 	}
 
 	localNode := enode.NewLocalNode(nodeDB, config.PrivateKey)
-	localNode.SetFallbackIP(net.IP{127, 0, 0, 1})
+
 	localNode.Set(discover.Tag)
+	listenerAddr := conn.LocalAddr().(*net.UDPAddr)
+	nat := config.Protocol.NAT
+	if nat != nil && !listenerAddr.IP.IsLoopback() {
+		doPortMapping(nat, localNode, listenerAddr)
+	}
 
 	discV5, err := discover.ListenV5(conn, localNode, discCfg)
 	if err != nil {
@@ -259,6 +265,57 @@ func initDiscV5(config Config, conn discover.UDPConn) (*discover.UDPv5, *enode.L
 	return discV5, localNode, nil
 }
 
+func doPortMapping(natm nat.Interface, ln *enode.LocalNode, addr *net.UDPAddr) {
+	const (
+		protocol = "udp"
+		name     = "ethereum discovery"
+	)
+
+	var (
+		intport    = addr.Port
+		extaddr    = &net.UDPAddr{IP: addr.IP, Port: addr.Port}
+		mapTimeout = nat.DefaultMapTimeout
+	)
+	addMapping := func() {
+		// Get the external address.
+		var err error
+		extaddr.IP, err = natm.ExternalIP()
+		if err != nil {
+			log.Debug("Couldn't get external IP", "err", err)
+			return
+		}
+		// Create the mapping.
+		p, err := natm.AddMapping(protocol, extaddr.Port, intport, name, mapTimeout)
+		if err != nil {
+			log.Debug("Couldn't add port mapping", "err", err)
+			return
+		}
+		if p != uint16(extaddr.Port) {
+			extaddr.Port = int(p)
+			log.Info("NAT mapped alternative port")
+		} else {
+			log.Info("NAT mapped port")
+		}
+		// Update IP/port information of the local node.
+		ln.SetStaticIP(extaddr.IP)
+		ln.SetFallbackUDP(extaddr.Port)
+	}
+
+	// Perform mapping once, synchronously.
+	log.Info("Attempting port mapping")
+	addMapping()
+
+	// Refresh the mapping periodically.
+	go func() {
+		refresh := time.NewTimer(mapTimeout)
+		defer refresh.Stop()
+		for range refresh.C {
+			addMapping()
+			refresh.Reset(mapTimeout)
+		}
+	}()
+}
+
 func initHistory(config Config, server *rpc.Server, conn discover.UDPConn, localNode *enode.LocalNode, discV5 *discover.UDPv5) (*history.HistoryNetwork, error) {
 	networkName := portalwire.History.Name()
 	db, err := history.NewDB(config.DataDir, networkName)
diff --git a/p2p/discover/nat.go b/p2p/discover/nat.go
index 8cb100e81c..0202ef3c66 100644
--- a/p2p/discover/nat.go
+++ b/p2p/discover/nat.go
@@ -47,6 +47,12 @@ func (p *PortalProtocol) setupPortMapping() {
 		p.localNode.SetStaticIP(ip)
 		go p.consumePortMappingRequests()
 
+	case nat.STUN:
+		// STUN doesn't block, set the IP right away.
+		ip, _ := p.NAT.ExternalIP()
+		p.localNode.SetStaticIP(ip)
+		go p.consumePortMappingRequests()
+
 	default:
 		go p.portMappingLoop()
 	}
diff --git a/p2p/nat/nat_stun.go b/p2p/nat/nat_stun.go
index c658a5078e..82c6f50680 100644
--- a/p2p/nat/nat_stun.go
+++ b/p2p/nat/nat_stun.go
@@ -29,8 +29,8 @@ func (STUN) SupportsMapping() bool {
 	return false
 }
 
-func (STUN) AddMapping(string, int, int, string, time.Duration) (uint16, error) {
-	return 0, nil
+func (STUN) AddMapping(protocol string, extport, intport int, name string, lifetime time.Duration) (uint16, error) {
+	return uint16(extport), nil
 }
 
 func (STUN) DeleteMapping(string, int, int) error {