Initial commit

core/connect/stun.go

@@ -0,0 +1,119 @@
+package connect
+
+import (
+	"fmt"
+	"net"
+	"time"
+
+	"github.com/pion/stun"
+	"go.uber.org/zap"
+)
+
+// STUNClient STUN 客户端 - 用于 NAT 探测和候选地址采集
+type STUNClient struct {
+	servers []string
+	logger  *zap.Logger
+	timeout time.Duration
+}
+
+// NewSTUNClient 创建 STUN 客户端
+func NewSTUNClient(servers []string, logger *zap.Logger) *STUNClient {
+	return &STUNClient{
+		servers: servers,
+		logger:  logger,
+		timeout: 5 * time.Second,
+	}
+}
+
+// DiscoverAddress 发现外部地址（通过单个 STUN 服务器）
+func (c *STUNClient) DiscoverAddress(server string) (*net.UDPAddr, error) {
+	host, port, err := net.SplitHostPort(server)
+	if err != nil {
+		return nil, fmt.Errorf("STUN 服务器地址格式错误：%w", err)
+	}
+
+	udpAddr, err := net.ResolveUDPAddr("udp4", net.JoinHostPort(host, port))
+	if err != nil {
+		return nil, fmt.Errorf("解析 UDP 地址失败：%w", err)
+	}
+
+	conn, err := net.DialUDP("udp4", nil, udpAddr)
+	if err != nil {
+		return nil, fmt.Errorf("连接 STUN 服务器失败：%w", err)
+	}
+	defer conn.Close()
+
+	conn.SetDeadline(time.Now().Add(c.timeout))
+
+	// 构建 STUN Binding Request
+	msg, err := stun.Build(stun.BindingRequest, stun.TransactionID)
+	if err != nil {
+		return nil, fmt.Errorf("构建 STUN 请求失败：%w", err)
+	}
+
+	// 发送请求
+	if _, err := conn.Write(msg.Raw); err != nil {
+		return nil, fmt.Errorf("发送 STUN 请求失败：%w", err)
+	}
+
+	// 读取响应
+	buf := make([]byte, 1024)
+	n, err := conn.Read(buf)
+	if err != nil {
+		return nil, fmt.Errorf("读取 STUN 响应失败：%w", err)
+	}
+
+	// 解析响应
+	res := &stun.Message{Raw: buf[:n]}
+	if err := res.Decode(); err != nil {
+		return nil, fmt.Errorf("解码 STUN 响应失败：%w", err)
+	}
+
+	// 提取 XOR-MAPPED-ADDRESS
+	var xorAddr stun.XORMappedAddress
+	if err := xorAddr.GetFrom(res); err != nil {
+		return nil, fmt.Errorf("提取外部地址失败：%w", err)
+	}
+
+	c.logger.Debug("STUN 查询成功",
+		zap.String("server", server),
+		zap.String("external_addr", xorAddr.String()))
+
+	return &net.UDPAddr{
+		IP:   xorAddr.IP,
+		Port: xorAddr.Port,
+	}, nil
+}
+
+// CollectCandidates 收集候选地址（通过多个 STUN 服务器）
+func (c *STUNClient) CollectCandidates() []string {
+	var candidates []string
+
+	for _, server := range c.servers {
+		addr, err := c.DiscoverAddress(server)
+		if err != nil {
+			c.logger.Debug("STUN 服务器查询失败",
+				zap.String("server", server),
+				zap.Error(err))
+			continue
+		}
+
+		candidates = append(candidates, addr.String())
+		c.logger.Debug("收集到候选地址",
+			zap.String("server", server),
+			zap.String("candidate", addr.String()))
+	}
+
+	return candidates
+}
+
+// GetExternalIP 获取外部 IP（兼容旧 API）
+func (c *STUNClient) GetExternalIP() (string, error) {
+	for _, server := range c.servers {
+		addr, err := c.DiscoverAddress(server)
+		if err == nil {
+			return addr.IP.String(), nil
+		}
+	}
+	return "", fmt.Errorf("所有 STUN 服务器均查询失败")
+}