120 lines
2.9 KiB
Go
120 lines
2.9 KiB
Go
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 服务器均查询失败")
|
|
}
|