fix(nathole): migrate STUN client to golib (#5419)

2026-07-21 06:59:17 +08:00 · 2026-07-16 13:11:45 +08:00
parent 2886393f5b
commit 269a26c5d5
5 changed files with 410 additions and 99 deletions
--- a/pkg/nathole/discovery_test.go
+++ b/pkg/nathole/discovery_test.go
@@ -0,0 +1,382 @@
+package nathole
+
+import (
+	"encoding/binary"
+	"errors"
+	"fmt"
+	"net"
+	"testing"
+	"time"
+
+	"github.com/fatedier/golib/net/stun"
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	testBindingRequest  = 0x0001
+	testBindingSuccess  = 0x0101
+	testBindingError    = 0x0111
+	testMagicCookie     = 0x2112a442
+	testAttrMapped      = 0x0001
+	testAttrChanged     = 0x0005
+	testAttrErrorCode   = 0x0009
+	testAttrXORMapped   = 0x0020
+	testAttrOther       = 0x802c
+	testSTUNHeaderSize  = 20
+	testSTUNServerLimit = time.Second
+)
+
+type testSTUNAttribute struct {
+	typ   uint16
+	value []byte
+}
+
+type testSTUNExchange struct {
+	source *net.UDPAddr
+	err    error
+}
+
+func listenTestUDP4(t *testing.T) *net.UDPConn {
+	t.Helper()
+
+	conn, err := net.ListenUDP("udp4", &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)})
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+	return conn
+}
+
+func serveOneSTUNRequest(
+	server *net.UDPConn,
+	buildResponse func([]byte, *net.UDPAddr) ([]byte, error),
+) <-chan testSTUNExchange {
+	done := make(chan testSTUNExchange, 1)
+	go func() {
+		if err := server.SetDeadline(time.Now().Add(testSTUNServerLimit)); err != nil {
+			done <- testSTUNExchange{err: err}
+			return
+		}
+		buffer := make([]byte, 1024)
+		n, source, err := server.ReadFromUDP(buffer)
+		if err == nil && buildResponse != nil {
+			var response []byte
+			response, err = buildResponse(buffer[:n], source)
+			if err == nil && response != nil {
+				_, err = server.WriteToUDP(response, source)
+			}
+		}
+		done <- testSTUNExchange{source: source, err: err}
+	}()
+	return done
+}
+
+func waitSTUNExchange(t *testing.T, done <-chan testSTUNExchange) *net.UDPAddr {
+	t.Helper()
+
+	select {
+	case exchange := <-done:
+		require.NoError(t, exchange.err)
+		return exchange.source
+	case <-time.After(testSTUNServerLimit):
+		t.Fatal("timed out waiting for local STUN server")
+		return nil
+	}
+}
+
+func makeTestSTUNResponse(request []byte, typ uint16, attributes ...testSTUNAttribute) ([]byte, error) {
+	if len(request) != testSTUNHeaderSize || binary.BigEndian.Uint16(request[0:2]) != testBindingRequest ||
+		binary.BigEndian.Uint32(request[4:8]) != testMagicCookie {
+		return nil, fmt.Errorf("invalid Binding request")
+	}
+
+	length := 0
+	for _, attribute := range attributes {
+		length += 4 + (len(attribute.value)+3)&^3
+	}
+	response := make([]byte, testSTUNHeaderSize, testSTUNHeaderSize+length)
+	binary.BigEndian.PutUint16(response[0:2], typ)
+	binary.BigEndian.PutUint16(response[2:4], uint16(length))
+	binary.BigEndian.PutUint32(response[4:8], testMagicCookie)
+	copy(response[8:20], request[8:20])
+
+	for _, attribute := range attributes {
+		start := len(response)
+		paddedLength := (len(attribute.value) + 3) &^ 3
+		response = append(response, make([]byte, 4+paddedLength)...)
+		binary.BigEndian.PutUint16(response[start:start+2], attribute.typ)
+		binary.BigEndian.PutUint16(response[start+2:start+4], uint16(len(attribute.value)))
+		copy(response[start+4:], attribute.value)
+	}
+	return response, nil
+}
+
+func testIPv4AddressValue(ip net.IP, port int, xor bool) []byte {
+	value := make([]byte, 8)
+	value[1] = 0x01
+	binary.BigEndian.PutUint16(value[2:4], uint16(port))
+	copy(value[4:], ip.To4())
+	if xor {
+		binary.BigEndian.PutUint16(value[2:4], binary.BigEndian.Uint16(value[2:4])^uint16(testMagicCookie>>16))
+		for i := range 4 {
+			value[4+i] ^= byte(uint32(testMagicCookie) >> uint(24-8*i))
+		}
+	}
+	return value
+}
+
+func TestDiscoverReusesLocalPortAndPreservesNATClassification(t *testing.T) {
+	tests := []struct {
+		name             string
+		secondMapped     string
+		secondMappedPort int
+		wantNATType      string
+		wantBehavior     string
+	}{
+		{
+			name:             "same mapped address",
+			secondMapped:     "198.51.100.10:40000",
+			secondMappedPort: 40000,
+			wantNATType:      EasyNAT,
+			wantBehavior:     BehaviorNoChange,
+		},
+		{
+			name:             "different mapped port",
+			secondMapped:     "198.51.100.10:40001",
+			secondMappedPort: 40001,
+			wantNATType:      HardNAT,
+			wantBehavior:     BehaviorPortChanged,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			primary := listenTestUDP4(t)
+			alternate := listenTestUDP4(t)
+			alternateAddr := alternate.LocalAddr().(*net.UDPAddr)
+
+			primaryDone := serveOneSTUNRequest(primary, func(request []byte, _ *net.UDPAddr) ([]byte, error) {
+				return makeTestSTUNResponse(request, testBindingSuccess,
+					testSTUNAttribute{typ: testAttrXORMapped, value: testIPv4AddressValue(net.ParseIP("198.51.100.10"), 40000, true)},
+					testSTUNAttribute{typ: testAttrOther, value: testIPv4AddressValue(alternateAddr.IP, alternateAddr.Port, false)},
+				)
+			})
+			alternateDone := serveOneSTUNRequest(alternate, func(request []byte, _ *net.UDPAddr) ([]byte, error) {
+				return makeTestSTUNResponse(request, testBindingSuccess,
+					testSTUNAttribute{typ: testAttrXORMapped, value: testIPv4AddressValue(net.ParseIP("198.51.100.10"), tt.secondMappedPort, true)},
+				)
+			})
+
+			addresses, localAddr, err := Discover([]string{primary.LocalAddr().String()}, "")
+			require.NoError(t, err)
+			require.Equal(t, []string{"198.51.100.10:40000", tt.secondMapped}, addresses)
+
+			primarySource := waitSTUNExchange(t, primaryDone)
+			alternateSource := waitSTUNExchange(t, alternateDone)
+			require.Equal(t, primarySource.Port, alternateSource.Port)
+			require.Equal(t, localAddr.(*net.UDPAddr).Port, primarySource.Port)
+
+			feature, err := ClassifyNATFeature(addresses, nil)
+			require.NoError(t, err)
+			require.Equal(t, tt.wantNATType, feature.NatType)
+			require.Equal(t, tt.wantBehavior, feature.Behavior)
+		})
+	}
+}
+
+func TestDoSTUNRequestMapsLegacyAndModernAddresses(t *testing.T) {
+	tests := []struct {
+		name         string
+		attributes   []testSTUNAttribute
+		wantExternal string
+		wantOther    string
+	}{
+		{
+			name: "legacy",
+			attributes: []testSTUNAttribute{
+				{typ: testAttrMapped, value: testIPv4AddressValue(net.ParseIP("192.0.2.1"), 1000, false)},
+				{typ: testAttrChanged, value: testIPv4AddressValue(net.ParseIP("192.0.2.2"), 2000, false)},
+			},
+			wantExternal: "192.0.2.1:1000",
+			wantOther:    "192.0.2.2:2000",
+		},
+		{
+			name: "modern takes precedence",
+			attributes: []testSTUNAttribute{
+				{typ: testAttrMapped, value: testIPv4AddressValue(net.ParseIP("192.0.2.1"), 1000, false)},
+				{typ: testAttrXORMapped, value: testIPv4AddressValue(net.ParseIP("198.51.100.1"), 3000, true)},
+				{typ: testAttrChanged, value: testIPv4AddressValue(net.ParseIP("192.0.2.2"), 2000, false)},
+				{typ: testAttrOther, value: testIPv4AddressValue(net.ParseIP("198.51.100.2"), 4000, false)},
+			},
+			wantExternal: "198.51.100.1:3000",
+			wantOther:    "198.51.100.2:4000",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			server := listenTestUDP4(t)
+			done := serveOneSTUNRequest(server, func(request []byte, _ *net.UDPAddr) ([]byte, error) {
+				return makeTestSTUNResponse(request, testBindingSuccess, tt.attributes...)
+			})
+			conn, err := listen("")
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = conn.Close() })
+
+			response, err := conn.doSTUNRequest(server.LocalAddr().String())
+			require.NoError(t, err)
+			require.Equal(t, tt.wantExternal, response.externalAddr)
+			require.Equal(t, tt.wantOther, response.otherAddr)
+			waitSTUNExchange(t, done)
+		})
+	}
+}
+
+func TestSTUNResponseErrorsAndMissingAddresses(t *testing.T) {
+	tests := []struct {
+		name          string
+		buildResponse func([]byte, *net.UDPAddr) ([]byte, error)
+		request       func(*discoverConn, string) error
+		checkError    func(*testing.T, error)
+	}{
+		{
+			name: "correlated malformed response",
+			buildResponse: func(request []byte, _ *net.UDPAddr) ([]byte, error) {
+				response, err := makeTestSTUNResponse(request, testBindingSuccess)
+				if err == nil {
+					binary.BigEndian.PutUint16(response[2:4], 4)
+				}
+				return response, err
+			},
+			request: func(conn *discoverConn, server string) error {
+				_, err := conn.doSTUNRequest(server)
+				return err
+			},
+			checkError: func(t *testing.T, err error) {
+				require.ErrorIs(t, err, stun.ErrMalformedResponse)
+			},
+		},
+		{
+			name: "Binding error response",
+			buildResponse: func(request []byte, _ *net.UDPAddr) ([]byte, error) {
+				return makeTestSTUNResponse(request, testBindingError, testSTUNAttribute{
+					typ:   testAttrErrorCode,
+					value: []byte{0, 0, 4, 20, 'U', 'n', 'k', 'n', 'o', 'w', 'n'},
+				})
+			},
+			request: func(conn *discoverConn, server string) error {
+				_, err := conn.doSTUNRequest(server)
+				return err
+			},
+			checkError: func(t *testing.T, err error) {
+				var responseErr *stun.ResponseError
+				require.ErrorAs(t, err, &responseErr)
+				require.Equal(t, 420, responseErr.Code)
+			},
+		},
+		{
+			name: "missing mapped address",
+			buildResponse: func(request []byte, _ *net.UDPAddr) ([]byte, error) {
+				return makeTestSTUNResponse(request, testBindingSuccess,
+					testSTUNAttribute{typ: testAttrOther, value: testIPv4AddressValue(net.ParseIP("192.0.2.2"), 2000, false)},
+				)
+			},
+			request: func(conn *discoverConn, server string) error {
+				_, err := conn.discoverFromStunServer(server)
+				return err
+			},
+			checkError: func(t *testing.T, err error) {
+				require.EqualError(t, err, "no external address found")
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			server := listenTestUDP4(t)
+			done := serveOneSTUNRequest(server, tt.buildResponse)
+			conn, err := listen("")
+			require.NoError(t, err)
+			t.Cleanup(func() { _ = conn.Close() })
+
+			err = tt.request(conn, server.LocalAddr().String())
+			tt.checkError(t, err)
+			waitSTUNExchange(t, done)
+		})
+	}
+
+	t.Run("missing other address", func(t *testing.T) {
+		server := listenTestUDP4(t)
+		done := serveOneSTUNRequest(server, func(request []byte, _ *net.UDPAddr) ([]byte, error) {
+			return makeTestSTUNResponse(request, testBindingSuccess,
+				testSTUNAttribute{typ: testAttrXORMapped, value: testIPv4AddressValue(net.ParseIP("198.51.100.1"), 3000, true)},
+			)
+		})
+
+		_, err := Prepare([]string{server.LocalAddr().String()}, PrepareOptions{})
+		require.EqualError(t, err, "discover error: not enough addresses")
+		waitSTUNExchange(t, done)
+	})
+}
+
+func TestSTUNTimeoutUsesCallerDeadlineWithoutRetry(t *testing.T) {
+	originalTimeout := responseTimeout
+	responseTimeout = 50 * time.Millisecond
+	t.Cleanup(func() { responseTimeout = originalTimeout })
+
+	server := listenTestUDP4(t)
+	done := serveOneSTUNRequest(server, nil)
+	conn, err := listen("")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+
+	_, err = conn.doSTUNRequest(server.LocalAddr().String())
+	require.EqualError(t, err, "wait response from stun server timeout")
+	waitSTUNExchange(t, done)
+
+	require.NoError(t, server.SetReadDeadline(time.Now().Add(50*time.Millisecond)))
+	_, _, err = server.ReadFromUDP(make([]byte, 1))
+	var netErr net.Error
+	require.ErrorAs(t, err, &netErr)
+	require.True(t, netErr.Timeout())
+}
+
+func TestSTUNClientLeavesSocketAndDeadlineWithCaller(t *testing.T) {
+	originalTimeout := responseTimeout
+	responseTimeout = 100 * time.Millisecond
+	t.Cleanup(func() { responseTimeout = originalTimeout })
+
+	server := listenTestUDP4(t)
+	unrelated := listenTestUDP4(t)
+	done := serveOneSTUNRequest(server, func(request []byte, source *net.UDPAddr) ([]byte, error) {
+		response, err := makeTestSTUNResponse(request, testBindingSuccess,
+			testSTUNAttribute{typ: testAttrXORMapped, value: testIPv4AddressValue(net.ParseIP("198.51.100.5"), 5000, true)},
+		)
+		if err != nil {
+			return nil, err
+		}
+		if _, err := unrelated.WriteToUDP(response, source); err != nil {
+			return nil, err
+		}
+		return response, nil
+	})
+	conn, err := listen("")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = conn.Close() })
+
+	response, err := conn.doSTUNRequest(server.LocalAddr().String())
+	require.NoError(t, err)
+	require.Equal(t, "198.51.100.5:5000", response.externalAddr)
+	waitSTUNExchange(t, done)
+
+	_, _, err = conn.conn.ReadFromUDP(make([]byte, 1))
+	var netErr net.Error
+	require.True(t, errors.As(err, &netErr))
+	require.True(t, netErr.Timeout())
+
+	require.NoError(t, conn.conn.SetDeadline(time.Time{}))
+	require.NoError(t, server.SetReadDeadline(time.Now().Add(testSTUNServerLimit)))
+	_, err = conn.conn.WriteToUDP([]byte{1}, server.LocalAddr().(*net.UDPAddr))
+	require.NoError(t, err)
+	_, source, err := server.ReadFromUDP(make([]byte, 1))
+	require.NoError(t, err)
+	require.Equal(t, conn.localAddr.(*net.UDPAddr).Port, source.Port)
+}