feat:新增server upd转发功能
This commit is contained in:
nnbcccscdscdsc
@@ -1,103 +1,291 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"net"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
kcp "github.com/xtaci/kcp-go/v5"
|
||||
|
||||
"omnisocketgo/cmd/internal/latencylog"
|
||||
"omnisocketgo/cmd/internal/protocol"
|
||||
"omnisocketgo/cmd/internal/transport"
|
||||
)
|
||||
|
||||
func TestUDPRelayRoutesPacketsByKCPConversationID(t *testing.T) {
|
||||
remote, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("ListenPacket(remote) error = %v", err)
|
||||
}
|
||||
defer remote.Close()
|
||||
// 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()
|
||||
|
||||
relayConn, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("ListenPacket(relay) error = %v", err)
|
||||
// 启动 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)
|
||||
}
|
||||
|
||||
relay, err := NewUDPRelay(relayConn, remote.LocalAddr().(*net.UDPAddr))
|
||||
if err != nil {
|
||||
_ = relayConn.Close()
|
||||
t.Fatalf("NewUDPRelay() error = %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)
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
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()
|
||||
if serveErr := relay.Serve(); serveErr != nil {
|
||||
t.Errorf("relay.Serve() error = %v", serveErr)
|
||||
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)
|
||||
}
|
||||
}()
|
||||
defer func() {
|
||||
_ = relayConn.Close()
|
||||
|
||||
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()
|
||||
}()
|
||||
|
||||
client1, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("ListenPacket(client1) error = %v", err)
|
||||
}
|
||||
defer client1.Close()
|
||||
t.Cleanup(func() {
|
||||
_ = relay.Close()
|
||||
})
|
||||
|
||||
client2, err := net.ListenPacket("udp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatalf("ListenPacket(client2) error = %v", err)
|
||||
}
|
||||
defer client2.Close()
|
||||
|
||||
relayAddr := relayConn.LocalAddr()
|
||||
|
||||
sendPacket(t, client1, relayAddr, buildRelayTestPacket(1, []byte("client-one")))
|
||||
assertPacketReceived(t, remote, buildRelayTestPacket(1, []byte("client-one")))
|
||||
|
||||
sendPacket(t, client2, relayAddr, buildRelayTestPacket(2, []byte("client-two")))
|
||||
assertPacketReceived(t, remote, buildRelayTestPacket(2, []byte("client-two")))
|
||||
|
||||
sendPacket(t, remote, relayAddr, buildRelayTestPacket(2, []byte("reply-two")))
|
||||
assertPacketReceived(t, client2, buildRelayTestPacket(2, []byte("reply-two")))
|
||||
|
||||
sendPacket(t, remote, relayAddr, buildRelayTestPacket(1, []byte("reply-one")))
|
||||
assertPacketReceived(t, client1, buildRelayTestPacket(1, []byte("reply-one")))
|
||||
return conn.LocalAddr().String()
|
||||
}
|
||||
|
||||
func buildRelayTestPacket(convID uint32, body []byte) []byte {
|
||||
packet := make([]byte, 4+len(body))
|
||||
binary.LittleEndian.PutUint32(packet[:4], convID)
|
||||
copy(packet[4:], body)
|
||||
return packet
|
||||
}
|
||||
|
||||
func sendPacket(t *testing.T, conn net.PacketConn, addr net.Addr, payload []byte) {
|
||||
// waitForRelay 轮询等待条件满足，超时则 fail。
|
||||
func waitForRelay(t *testing.T, condition func() bool, description string) {
|
||||
t.Helper()
|
||||
|
||||
if err := conn.SetWriteDeadline(time.Now().Add(2 * time.Second)); err != nil {
|
||||
t.Fatalf("SetWriteDeadline() error = %v", err)
|
||||
}
|
||||
if _, err := conn.WriteTo(payload, addr); err != nil {
|
||||
t.Fatalf("WriteTo(%s) error = %v", addr, err)
|
||||
}
|
||||
}
|
||||
|
||||
func assertPacketReceived(t *testing.T, conn net.PacketConn, want []byte) {
|
||||
t.Helper()
|
||||
|
||||
if err := conn.SetReadDeadline(time.Now().Add(2 * time.Second)); err != nil {
|
||||
t.Fatalf("SetReadDeadline() error = %v", err)
|
||||
}
|
||||
|
||||
buffer := make([]byte, 1024)
|
||||
n, _, err := conn.ReadFrom(buffer)
|
||||
if err != nil {
|
||||
t.Fatalf("ReadFrom() error = %v", err)
|
||||
}
|
||||
got := buffer[:n]
|
||||
if string(got) != string(want) {
|
||||
t.Fatalf("packet = %v, want %v", got, want)
|
||||
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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user