From 9bccaac318e615ed17f09656ae04642bfcfa5e7a Mon Sep 17 00:00:00 2001
From: pengzhen <pengzhen@cicvdata.com>
Date: Tue, 10 Dec 2024 10:26:56 +0800
Subject: [PATCH] feat: make stun server private

---
 cmd/utils/flags.go  |  2 +-
 p2p/nat/nat.go      |  2 +-
 p2p/nat/nat_stun.go | 36 ++++++++++++++++++------------------
 p2p/nat/nat_test.go |  7 ++++---
 4 files changed, 24 insertions(+), 23 deletions(-)

diff --git a/cmd/utils/flags.go b/cmd/utils/flags.go
index 48143c5acb..b539ae6dfb 100644
--- a/cmd/utils/flags.go
+++ b/cmd/utils/flags.go
@@ -760,7 +760,7 @@ var (
 	}
 	NATFlag = &cli.StringFlag{
 		Name:     "nat",
-		Usage:    "NAT port mapping mechanism (any|none|upnp|pmp|pmp:<IP>|extip:<IP>|stun:default|stun:<IP>)",
+		Usage:    "NAT port mapping mechanism (any|none|upnp|pmp|pmp:<IP>|extip:<IP>|stun:default|stun:<IP:PORT>)",
 		Value:    "any",
 		Category: flags.NetworkingCategory,
 	}
diff --git a/p2p/nat/nat.go b/p2p/nat/nat.go
index 387ddebca7..0a84eb3090 100644
--- a/p2p/nat/nat.go
+++ b/p2p/nat/nat.go
@@ -89,7 +89,7 @@ func Parse(spec string) (Interface, error) {
 	case "pmp", "natpmp", "nat-pmp":
 		return PMP(ip), nil
 	case "stun":
-		return NewSTUN(after)
+		return newSTUN(after)
 	default:
 		return nil, fmt.Errorf("unknown mechanism %q", before)
 	}
diff --git a/p2p/nat/nat_stun.go b/p2p/nat/nat_stun.go
index 8b7bdaf268..690610a65e 100644
--- a/p2p/nat/nat_stun.go
+++ b/p2p/nat/nat_stun.go
@@ -21,58 +21,58 @@ import (
 	"net"
 	"time"
 
-	"github.com/pion/stun/v2"
+	stunV2 "github.com/pion/stun/v2"
 )
 
 // The code are from erigon p2p/nat/nat_stun.go
 // This stun server is part of the mainnet infrastructure.
 // The addr are from https://github.com/ethereum/trin/blob/master/portalnet/src/socket.rs
-const STUNDefaultServerAddr = "159.223.0.83:3478"
+const stunDefaultServerAddr = "159.223.0.83:3478"
 
-type STUN struct {
-	serverAddr *net.UDPAddr
+type stun struct {
+	server *net.UDPAddr
 }
 
-func NewSTUN(serverAddr string) (Interface, error) {
+func newSTUN(serverAddr string) (Interface, error) {
 	if serverAddr == "default" {
-		serverAddr = STUNDefaultServerAddr
+		serverAddr = stunDefaultServerAddr
 	}
 	addr, err := net.ResolveUDPAddr("udp4", serverAddr)
 	if err != nil {
 		return nil, err
 	}
-	return STUN{serverAddr: addr}, nil
+	return stun{server: addr}, nil
 }
 
-func (s STUN) String() string {
-	return fmt.Sprintf("STUN(%s)", s.serverAddr)
+func (s stun) String() string {
+	return fmt.Sprintf("STUN(%s)", s.server)
 }
 
-func (STUN) SupportsMapping() bool {
+func (stun) SupportsMapping() bool {
 	return false
 }
 
-func (STUN) AddMapping(protocol string, extport, intport int, name string, lifetime time.Duration) (uint16, error) {
+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 {
+func (stun) DeleteMapping(string, int, int) error {
 	return nil
 }
 
-func (s STUN) ExternalIP() (net.IP, error) {
-	conn, err := stun.Dial("udp4", s.serverAddr.String())
+func (s stun) ExternalIP() (net.IP, error) {
+	conn, err := stunV2.Dial("udp4", s.server.String())
 	if err != nil {
 		return nil, err
 	}
 	defer conn.Close()
 
-	message, err := stun.Build(stun.TransactionID, stun.BindingRequest)
+	message, err := stunV2.Build(stunV2.TransactionID, stunV2.BindingRequest)
 	if err != nil {
 		return nil, err
 	}
-	var response *stun.Event
-	err = conn.Do(message, func(event stun.Event) {
+	var response *stunV2.Event
+	err = conn.Do(message, func(event stunV2.Event) {
 		response = &event
 	})
 	if err != nil {
@@ -82,7 +82,7 @@ func (s STUN) ExternalIP() (net.IP, error) {
 		return nil, response.Error
 	}
 
-	var mappedAddr stun.XORMappedAddress
+	var mappedAddr stunV2.XORMappedAddress
 	if err := mappedAddr.GetFrom(response.Message); err != nil {
 		return nil, err
 	}
diff --git a/p2p/nat/nat_test.go b/p2p/nat/nat_test.go
index edf501624a..3659af4524 100644
--- a/p2p/nat/nat_test.go
+++ b/p2p/nat/nat_test.go
@@ -62,14 +62,15 @@ func TestAutoDiscRace(t *testing.T) {
 	}
 }
 
+// stun:default should work well
 func TestStunDefault(t *testing.T) {
 	nat, err := Parse("stun:default")
 	if err != nil {
 		t.Errorf("should no err, but get %v", err)
 	}
-	stun := nat.(STUN)
-	if stun.serverAddr.String() != STUNDefaultServerAddr {
-		t.Errorf("want addr %s, got addr %s", STUNDefaultServerAddr, stun.serverAddr.String())
+	stun := nat.(stun)
+	if stun.server.String() != stunDefaultServerAddr {
+		t.Errorf("want addr %s, got addr %s", stunDefaultServerAddr, stun.server.String())
 	}
 	_, err = stun.ExternalIP()
 	if err != nil {