From 5765be19c39da350a4f9776d7827d78a010a0ff4 Mon Sep 17 00:00:00 2001
From: fearlessfe <505380967@qq.com>
Date: Sat, 14 Dec 2024 19:57:58 +0800
Subject: [PATCH] feat: add stun server list

---
 p2p/nat/nat_stun.go      | 101 +++++++++++++++++++++++++++++++++------
 p2p/nat/nat_stun_test.go |  23 +++++++++
 p2p/nat/nat_test.go      |  28 ++++++-----
 3 files changed, 127 insertions(+), 25 deletions(-)
 create mode 100644 p2p/nat/nat_stun_test.go

diff --git a/p2p/nat/nat_stun.go b/p2p/nat/nat_stun.go
index 6716d5248716..7277aa36e083 100644
--- a/p2p/nat/nat_stun.go
+++ b/p2p/nat/nat_stun.go
@@ -24,28 +24,58 @@ import (
 	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"
+var stunDefaultServerList = []string{
+	"159.223.0.83:3478",
+	"stun.l.google.com:19302",
+	"stun1.l.google.com:19302",
+	"stun2.l.google.com:19302",
+	"stun3.l.google.com:19302",
+	"stun4.l.google.com:19302",
+	"stun01.sipphone.com",
+	"stun.ekiga.net",
+	"stun.fwdnet.net",
+	"stun.ideasip.com",
+	"stun.iptel.org",
+	"stun.rixtelecom.se",
+	"stun.schlund.de",
+	"stunserver.org",
+	"stun.softjoys.com",
+	"stun.voiparound.com",
+	"stun.voipbuster.com",
+	"stun.voipstunt.com",
+	"stun.voxgratia.org",
+	"stun.xten.com",
+}
+
+const requestLimit = 3
 
 type stun struct {
-	server *net.UDPAddr
+	serverList      []string
+	activeIndex     int // the server index which return the IP
+	pendingRequests int // request in flight
+	askedIndex      map[int]struct{}
+	replyCh         chan stunResponse
 }
 
 func newSTUN(serverAddr string) (Interface, error) {
+	serverList := make([]string, 0)
 	if serverAddr == "default" {
-		serverAddr = stunDefaultServerAddr
-	}
-	addr, err := net.ResolveUDPAddr("udp4", serverAddr)
-	if err != nil {
-		return nil, err
+		serverList = stunDefaultServerList
+	} else {
+		_, err := net.ResolveUDPAddr("udp4", serverAddr)
+		if err != nil {
+			return nil, err
+		}
+		serverList = append(serverList, serverAddr)
 	}
-	return stun{server: addr}, nil
+
+	return &stun{
+		serverList: serverList,
+	}, nil
 }
 
 func (s stun) String() string {
-	return fmt.Sprintf("STUN(%s)", s.server)
+	return fmt.Sprintf("STUN(%s)", s.serverList[s.activeIndex])
 }
 
 func (stun) SupportsMapping() bool {
@@ -60,8 +90,51 @@ func (stun) DeleteMapping(string, int, int) error {
 	return nil
 }
 
-func (s stun) ExternalIP() (net.IP, error) {
-	conn, err := stunV2.Dial("udp4", s.server.String())
+type stunResponse struct {
+	ip    net.IP
+	err   error
+	index int
+}
+
+func (s *stun) ExternalIP() (net.IP, error) {
+	var err error
+	s.replyCh = make(chan stunResponse, requestLimit)
+	s.askedIndex = make(map[int]struct{})
+	for s.startQueries() {
+		response := <-s.replyCh
+		s.pendingRequests--
+		if response.err != nil {
+			err = response.err
+			continue
+		}
+		s.activeIndex = response.index
+		return response.ip, nil
+	}
+	return nil, err
+}
+
+func (s *stun) startQueries() bool {
+	for i := 0; s.pendingRequests < requestLimit && i < len(s.serverList); i++ {
+		_, exist := s.askedIndex[i]
+		if exist {
+			continue
+		}
+		s.pendingRequests++
+		s.askedIndex[i] = struct{}{}
+		go func(index int, server string) {
+			ip, err := externalIP(server)
+			s.replyCh <- stunResponse{
+				ip:    ip,
+				index: index,
+				err:   err,
+			}
+		}(i, s.serverList[i])
+	}
+	return s.pendingRequests > 0
+}
+
+func externalIP(server string) (net.IP, error) {
+	conn, err := stunV2.Dial("udp4", server)
 	if err != nil {
 		return nil, err
 	}
diff --git a/p2p/nat/nat_stun_test.go b/p2p/nat/nat_stun_test.go
new file mode 100644
index 000000000000..b681ba56b1f3
--- /dev/null
+++ b/p2p/nat/nat_stun_test.go
@@ -0,0 +1,23 @@
+package nat
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNatStun(t *testing.T) {
+	nat, err := newSTUN("default")
+	assert.NoError(t, err)
+	_, err = nat.ExternalIP()
+	assert.NoError(t, err)
+}
+
+func TestUnreachedNatServer(t *testing.T) {
+	stun := &stun{
+		serverList: []string{"1.2.3.4:1234", "1.2.3.4:1234", "1.2.3.4:1234"},
+	}
+	stun.serverList = append(stun.serverList, stunDefaultServerList...)
+	_, err := stun.ExternalIP()
+	assert.NoError(t, err)
+}
diff --git a/p2p/nat/nat_test.go b/p2p/nat/nat_test.go
index 3659af452454..6279566ad985 100644
--- a/p2p/nat/nat_test.go
+++ b/p2p/nat/nat_test.go
@@ -20,6 +20,8 @@ import (
 	"net"
 	"testing"
 	"time"
+
+	"github.com/stretchr/testify/assert"
 )
 
 // This test checks that autodisc doesn't hang and returns
@@ -63,17 +65,21 @@ 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.server.String() != stunDefaultServerAddr {
-		t.Errorf("want addr %s, got addr %s", stunDefaultServerAddr, stun.server.String())
+func TestParseStun(t *testing.T) {
+	testcases := []struct {
+		natStr string
+		want   *stun
+	}{
+		{"stun:default", &stun{serverList: stunDefaultServerList}},
+		{"stun:1.2.3.4:1234", &stun{serverList: []string{"1.2.3.4:1234"}}},
 	}
-	_, err = stun.ExternalIP()
-	if err != nil {
-		t.Errorf("get err: %v", err)
+
+	for _, tc := range testcases {
+		nat, err := Parse(tc.natStr)
+		if err != nil {
+			t.Errorf("should no err, but get %v", err)
+		}
+		stun := nat.(*stun)
+		assert.Equal(t, stun.serverList, tc.want.serverList)
 	}
 }