292 lines
7.7 KiB
Go
292 lines
7.7 KiB
Go
package server
|
||
|
||
import (
|
||
"net"
|
||
"strings"
|
||
"sync"
|
||
"testing"
|
||
"time"
|
||
|
||
kcp "github.com/xtaci/kcp-go/v5"
|
||
|
||
"omnisocketgo/cmd/internal/latencylog"
|
||
"omnisocketgo/cmd/internal/protocol"
|
||
"omnisocketgo/cmd/internal/transport"
|
||
)
|
||
|
||
// TestUDPRelayKCPForwardAndReturn 验证 KCP 通过 UDP relay 的完整双向转发路径：
|
||
// peer-b -> D(KCP hub) -> C(UDP relay) -> peer-a 以及反向。
|
||
func TestUDPRelayKCPForwardAndReturn(t *testing.T) {
|
||
// 启动 D（KCP Hub）
|
||
hub, hubAddr, hubCleanup := startKCPHubForRelay(t)
|
||
defer hubCleanup()
|
||
|
||
// 启动 C（UDP Relay），upstream 指向 D
|
||
relayAddr := startUDPRelay(t, hubAddr)
|
||
|
||
// peer-b 直连 D（KCP）
|
||
peerBConn := dialKCPPeer(t, hubAddr)
|
||
// peer-a 连 C（通过 relay 间接连到 D）
|
||
peerAConn := dialKCPPeer(t, relayAddr)
|
||
|
||
// 注册 peer-b
|
||
if err := peerBConn.Send(protocol.Message{
|
||
Type: protocol.MessageTypeRegister,
|
||
From: "peer-b",
|
||
To: protocol.ServerPeerID,
|
||
}); err != nil {
|
||
t.Fatalf("peerB register: %v", err)
|
||
}
|
||
|
||
// 注册 peer-a（通过 relay）
|
||
if err := peerAConn.Send(protocol.Message{
|
||
Type: protocol.MessageTypeRegister,
|
||
From: "peer-a",
|
||
To: protocol.ServerPeerID,
|
||
}); err != nil {
|
||
t.Fatalf("peerA register: %v", err)
|
||
}
|
||
|
||
waitForRelay(t, func() bool {
|
||
return hub.HasPeer("peer-a") && hub.HasPeer("peer-b")
|
||
}, "both peers to be registered")
|
||
|
||
// peer-b -> peer-a（路径: B -> D -> C -> A）
|
||
if err := peerBConn.Send(protocol.Message{
|
||
Type: protocol.MessageTypeText,
|
||
ID: 1,
|
||
From: "peer-b",
|
||
To: "peer-a",
|
||
Body: []byte("hello from peer-b"),
|
||
}); err != nil {
|
||
t.Fatalf("peerB send text: %v", err)
|
||
}
|
||
|
||
msg, err := peerAConn.Receive()
|
||
if err != nil {
|
||
t.Fatalf("peerA receive: %v", err)
|
||
}
|
||
if msg.Type != protocol.MessageTypeText {
|
||
t.Fatalf("message type = %s, want text", msg.Type)
|
||
}
|
||
if msg.From != "peer-b" {
|
||
t.Fatalf("message from = %s, want peer-b", msg.From)
|
||
}
|
||
if string(msg.Body) != "hello from peer-b" {
|
||
t.Fatalf("message body = %q, want %q", string(msg.Body), "hello from peer-b")
|
||
}
|
||
|
||
// peer-a -> peer-b（路径: A -> C -> D -> B）
|
||
if err := peerAConn.Send(protocol.Message{
|
||
Type: protocol.MessageTypeText,
|
||
ID: 2,
|
||
From: "peer-a",
|
||
To: "peer-b",
|
||
Body: []byte("reply from peer-单个 downstream peer 通过 relay 连到 KCP server”这条
|
||
链路是成立的，转发逻辑本身没有明显的地址错误。cmd/internal/server/udp_relay.go 里就是原
|
||
样双向转发，下游来的包会记录 clientAddr 并写给上游，上游回来的包再写回这个 clientAddr。
|
||
关键代码在 cmd/internal/server/udp_relay.go:68 和 cmd/internal/server/udp_relay.go:89。
|
||
|
||
还有一个关键事实：kcppeer 里那句 connected to ... as ... (KCP) 不能证明 peer-a 真的在
|
||
hub 注册成功a"),
|
||
}); err != nil {
|
||
t.Fatalf("peerA send text: %v", err)
|
||
}
|
||
|
||
msg2, err := peerBConn.Receive()
|
||
if err != nil {
|
||
t.Fatalf("peerB receive: %v", err)
|
||
}
|
||
if msg2.Type != protocol.MessageTypeText {
|
||
t.Fatalf("message type = %s, want text", msg2.Type)
|
||
}
|
||
if msg2.From != "peer-a" {
|
||
t.Fatalf("message from = %s, want peer-a", msg2.From)
|
||
}
|
||
if string(msg2.Body) != "reply from peer-a" {
|
||
t.Fatalf("message body = %q, want %q", string(msg2.Body), "reply from peer-a")
|
||
}
|
||
}
|
||
|
||
// TestUDPRelayKCPFileMessage 验证通过 relay 转发 KCP 文件消息。
|
||
func TestUDPRelayKCPFileMessage(t *testing.T) {
|
||
hub, hubAddr, hubCleanup := startKCPHubForRelay(t)
|
||
defer hubCleanup()
|
||
|
||
relayAddr := startUDPRelay(t, hubAddr)
|
||
|
||
peerBConn := dialKCPPeer(t, hubAddr)
|
||
peerAConn := dialKCPPeer(t, relayAddr)
|
||
|
||
_ = peerBConn.Send(protocol.Message{
|
||
Type: protocol.MessageTypeRegister,
|
||
From: "peer-b",
|
||
To: protocol.ServerPeerID,
|
||
})
|
||
_ = peerAConn.Send(protocol.Message{
|
||
Type: protocol.MessageTypeRegister,
|
||
From: "peer-a",
|
||
To: protocol.ServerPeerID,
|
||
})
|
||
|
||
waitForRelay(t, func() bool {
|
||
return hub.HasPeer("peer-a") && hub.HasPeer("peer-b")
|
||
}, "both peers to be registered")
|
||
|
||
if err := peerBConn.Send(protocol.Message{
|
||
Type: protocol.MessageTypeFile,
|
||
ID: 1,
|
||
From: "peer-b",
|
||
To: "peer-a",
|
||
FileName: "test.bin",
|
||
Body: []byte{0xDE, 0xAD, 0xBE, 0xEF},
|
||
}); err != nil {
|
||
t.Fatalf("peerB send file: %v", err)
|
||
}
|
||
|
||
msg, err := peerAConn.Receive()
|
||
if err != nil {
|
||
t.Fatalf("peerA receive: %v", err)
|
||
}
|
||
if msg.Type != protocol.MessageTypeFile {
|
||
t.Fatalf("message type = %s, want file", msg.Type)
|
||
}
|
||
if msg.FileName != "test.bin" {
|
||
t.Fatalf("file name = %q, want %q", msg.FileName, "test.bin")
|
||
}
|
||
if len(msg.Body) != 4 || msg.Body[0] != 0xDE {
|
||
t.Fatalf("file body mismatch: 单个 downstream peer 通过 relay 连到 KCP server”这条
|
||
链路是成立的，转发逻辑本身没有明显的地址错误。cmd/internal/server/udp_relay.go 里就是原
|
||
样双向转发，下游来的包会记录 clientAddr 并写给上游，上游回来的包再写回这个 clientAddr。
|
||
关键代码在 cmd/internal/server/udp_relay.go:68 和 cmd/internal/server/udp_relay.go:89。
|
||
|
||
还有一个关键事实：kcppeer 里那句 connected to ... as ... (KCP) 不能证明 peer-a 真的在
|
||
hub 注册成功got %v", msg.Body)
|
||
}
|
||
}
|
||
|
||
// startKCPHubForRelay 启动一个 KCP hub server，返回 hub、监听地址和 cleanup 函数。
|
||
func startKCPHubForRelay(t *testing.T) (*KCPHub, string, func()) {
|
||
t.Helper()
|
||
|
||
hub := NewKCPHub()
|
||
|
||
listener, packetConn, err := transport.ListenKCPSessions("127.0.0.1:0", "", nil, latencylog.NodeRoleServer, "hub")
|
||
if err != nil {
|
||
t.Fatalf("ListenKCPSessions() error = %v", err)
|
||
}
|
||
|
||
var (
|
||
wg sync.WaitGroup
|
||
stop = make(chan struct{})
|
||
)
|
||
|
||
wg.Add(1)
|
||
go func() {
|
||
defer wg.Done()
|
||
for {
|
||
session, acceptErr := listener.AcceptKCP()
|
||
if acceptErr != nil {
|
||
select {
|
||
case <-stop:
|
||
return
|
||
default:
|
||
}
|
||
if strings.Contains(acceptErr.Error(), "closed") {
|
||
return
|
||
}
|
||
t.Errorf("AcceptKCP() error = %v", acceptErr)
|
||
return
|
||
}
|
||
|
||
wg.Add(1)
|
||
go func(sess *kcp.UDPSession) {
|
||
defer wg.Done()
|
||
if serveErr := hub.ServeSession(sess); serveErr != nil {
|
||
msg := serveErr.Error()
|
||
if !strings.Contains(msg, "closed") && !strings.Contains(msg, "broken pipe") {
|
||
t.Logf("hub.ServeSession() ended with %v", serveErr)
|
||
}
|
||
}
|
||
}(session)
|
||
}
|
||
}()
|
||
|
||
cleanup := func() {
|
||
close(stop)
|
||
_ = listener.Close()
|
||
_ = packetConn.Close()
|
||
wg.Wait()
|
||
}
|
||
|
||
return hub, listener.Addr().String(), cleanup
|
||
}
|
||
|
||
// dialKCPPeer 创建一条到指定地址的 KCP 连接，用于测试。
|
||
func dialKCPPeer(t *testing.T, serverAddr string) *transport.KCPConn {
|
||
t.Helper()
|
||
|
||
session, err := transport.DialKCPSession(serverAddr, "", "", nil, latencylog.NodeRolePeer, "test")
|
||
if err != nil {
|
||
t.Fatalf("DialKCPSession(%s) error = %v", serverAddr, err)
|
||
}
|
||
|
||
conn, err := transport.NewKCPConn(session)
|
||
if err != nil {
|
||
_ = session.Close()
|
||
t.Fatalf("NewKCPConn() error = %v", err)
|
||
}
|
||
|
||
t.Cleanup(func() {
|
||
_ = conn.Close()
|
||
})
|
||
|
||
return conn
|
||
}
|
||
|
||
// startUDPRelay 创建并启动一个 UDPRelay，返回其监听地址字符串。
|
||
func startUDPRelay(t *testing.T, upstreamAddr string) string {
|
||
t.Helper()
|
||
|
||
addr, err := net.ResolveUDPAddr("udp", "127.0.0.1:0")
|
||
if err != nil {
|
||
t.Fatalf("ResolveUDPAddr() error = %v", err)
|
||
}
|
||
|
||
conn, err := net.ListenUDP("udp", addr)
|
||
if err != nil {
|
||
t.Fatalf("ListenUDP() error = %v", err)
|
||
}
|
||
|
||
relay, err := NewUDPRelay(conn, upstreamAddr)
|
||
if err != nil {
|
||
_ = conn.Close()
|
||
t.Fatalf("NewUDPRelay() error = %v", err)
|
||
}
|
||
|
||
go func() {
|
||
_ = relay.Serve()
|
||
}()
|
||
|
||
t.Cleanup(func() {
|
||
_ = relay.Close()
|
||
})
|
||
|
||
return conn.LocalAddr().String()
|
||
}
|
||
|
||
// waitForRelay 轮询等待条件满足，超时则 fail。
|
||
func waitForRelay(t *testing.T, condition func() bool, description string) {
|
||
t.Helper()
|
||
|
||
deadline := time.Now().Add(2 * time.Second)
|
||
for time.Now().Before(deadline) {
|
||
if condition() {
|
||
return
|
||
}
|
||
time.Sleep(10 * time.Millisecond)
|
||
}
|
||
|
||
t.Fatalf("timed out waiting for %s", description)
|
||
}
|