limingjie/OmniSocketGo
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: 将架构改成支持中间两个server

Browse Source
This commit is contained in:
nnbcccscdscdsc
2026-03-27 19:08:22 +08:00
parent 8cec6a0766
commit 5be3ff670f
4 changed files with 609 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

283
cmd/internal/peer/kcp_client_test.go
View File

@@ -1,11 +1,14 @@
package peer
import (
"bytes"
"net"
"os"
"path/filepath"
"reflect"
"strings"
"sync"
"sync/atomic"
"testing"
kcp "github.com/xtaci/kcp-go/v5"
@@ -202,6 +205,185 @@ func TestKCPClientsExchangeMessagesWithLatencyLogs(t *testing.T) {
})
}
func TestKCPClientsExchangeMessagesAcrossRelayedServers(t *testing.T) {
fixture := startRelayedKCPHubs(t)
defer fixture.cleanup()
peerA, err := DialKCP(fixture.serverCAddr, "peer-a")
if err != nil {
t.Fatalf("DialKCP(peer-a) error = %v", err)
}
defer func() { _ = peerA.Close() }()
peerB, err := DialKCP(fixture.serverDAddr, "peer-b")
if err != nil {
t.Fatalf("DialKCP(peer-b) error = %v", err)
}
defer func() { _ = peerB.Close() }()
waitFor(t, func() bool { return fixture.hubC.HasPeer("peer-a") && fixture.hubD.HasPeer("peer-b") }, "both relayed peers to be registered")
if err := peerA.SendText("peer-b", "hello via relay"); err != nil {
t.Fatalf("peerA.SendText() error = %v", err)
}
gotAtB, err := peerB.Receive()
if err != nil {
t.Fatalf("peerB.Receive() error = %v", err)
}
wantAtB := protocol.Message{
Type: protocol.MessageTypeText,
ID: 1,
From: "peer-a",
To: "peer-b",
Body: []byte("hello via relay"),
}
if !reflect.DeepEqual(gotAtB, wantAtB) {
t.Fatalf("peerB received %+v, want %+v", gotAtB, wantAtB)
}
if err := peerB.SendText("peer-a", "hello back"); err != nil {
t.Fatalf("peerB.SendText() error = %v", err)
}
gotAtA, err := peerA.Receive()
if err != nil {
t.Fatalf("peerA.Receive() error = %v", err)
}
wantAtA := protocol.Message{
Type: protocol.MessageTypeText,
ID: 1,
From: "peer-b",
To: "peer-a",
Body: []byte("hello back"),
}
if !reflect.DeepEqual(gotAtA, wantAtA) {
t.Fatalf("peerA received %+v, want %+v", gotAtA, wantAtA)
}
if got := fixture.relayC.WriteCount(); got != 1 {
t.Fatalf("relayC write count = %d, want 1", got)
}
if got := fixture.relayD.WriteCount(); got != 1 {
t.Fatalf("relayD write count = %d, want 1", got)
}
}
func TestKCPHubPrefersLocalPeerBeforeRelay(t *testing.T) {
fixture := startRelayedKCPHubs(t)
defer fixture.cleanup()
peerA, err := DialKCP(fixture.serverCAddr, "peer-a")
if err != nil {
t.Fatalf("DialKCP(peer-a) error = %v", err)
}
defer func() { _ = peerA.Close() }()
peerB, err := DialKCP(fixture.serverCAddr, "peer-b")
if err != nil {
t.Fatalf("DialKCP(peer-b) error = %v", err)
}
defer func() { _ = peerB.Close() }()
waitFor(t, func() bool { return fixture.hubC.HasPeer("peer-a") && fixture.hubC.HasPeer("peer-b") }, "local peers on hubC to be registered")
if err := peerA.SendText("peer-b", "local delivery"); err != nil {
t.Fatalf("peerA.SendText() error = %v", err)
}
got, err := peerB.Receive()
if err != nil {
t.Fatalf("peerB.Receive() error = %v", err)
}
want := protocol.Message{
Type: protocol.MessageTypeText,
ID: 1,
From: "peer-a",
To: "peer-b",
Body: []byte("local delivery"),
}
if !reflect.DeepEqual(got, want) {
t.Fatalf("peerB received %+v, want %+v", got, want)
}
if got := fixture.relayC.WriteCount(); got != 0 {
t.Fatalf("relayC write count = %d, want 0 for local delivery", got)
}
if got := fixture.relayD.WriteCount(); got != 0 {
t.Fatalf("relayD write count = %d, want 0 for local delivery", got)
}
}
func TestKCPRelayedUnknownTargetReturnsErrorToOriginalSender(t *testing.T) {
fixture := startRelayedKCPHubs(t)
defer fixture.cleanup()
peerA, err := DialKCP(fixture.serverCAddr, "peer-a")
if err != nil {
t.Fatalf("DialKCP(peer-a) error = %v", err)
}
defer func() { _ = peerA.Close() }()
waitFor(t, func() bool { return fixture.hubC.HasPeer("peer-a") }, "peer-a to be registered on hubC")
if err := peerA.SendText("remote-missing", "hello"); err != nil {
t.Fatalf("peerA.SendText() error = %v", err)
}
got, err := peerA.Receive()
if err != nil {
t.Fatalf("peerA.Receive() error = %v", err)
}
if got.Type != protocol.MessageTypeError {
t.Fatalf("got type %s, want %s", got.Type, protocol.MessageTypeError)
}
if got.From != protocol.ServerPeerID {
t.Fatalf("error from = %s, want %s", got.From, protocol.ServerPeerID)
}
if got.To != "peer-a" {
t.Fatalf("error to = %s, want peer-a", got.To)
}
if string(got.Body) != "unknown target: remote-missing" {
t.Fatalf("error body = %q, want unknown target from relayed hub", got.Body)
}
if got := fixture.relayC.WriteCount(); got != 1 {
t.Fatalf("relayC write count = %d, want 1 for outbound relay", got)
}
if got := fixture.relayD.WriteCount(); got != 1 {
t.Fatalf("relayD write count = %d, want 1 for return error relay", got)
}
}
func TestKCPHubRejectsOversizeRelayedMessage(t *testing.T) {
fixture := startRelayedKCPHubs(t)
defer fixture.cleanup()
peerA, err := DialKCP(fixture.serverCAddr, "peer-a")
if err != nil {
t.Fatalf("DialKCP(peer-a) error = %v", err)
}
defer func() { _ = peerA.Close() }()
waitFor(t, func() bool { return fixture.hubC.HasPeer("peer-a") }, "peer-a to be registered on hubC")
body := bytes.Repeat([]byte("a"), 70*1024)
if err := peerA.SendFile("remote-peer", "payload.bin", body); err != nil {
t.Fatalf("peerA.SendFile() error = %v", err)
}
got, err := peerA.Receive()
if err != nil {
t.Fatalf("peerA.Receive() error = %v", err)
}
if got.Type != protocol.MessageTypeError {
t.Fatalf("got type %s, want %s", got.Type, protocol.MessageTypeError)
}
if string(got.Body) != "message too large for relay udp" {
t.Fatalf("error body = %q, want oversize relay error", got.Body)
}
if got := fixture.relayC.WriteCount(); got != 0 {
t.Fatalf("relayC write count = %d, want 0 when relay rejects oversize payload", got)
}
}
func startRealKCPHubServer(t *testing.T, hub *server.KCPHub) (string, func()) {
t.Helper()
@@ -253,6 +435,98 @@ func startRealKCPHubServer(t *testing.T, hub *server.KCPHub) (string, func()) {
return listener.Addr().String(), cleanup
}
type relayedKCPHubFixture struct {
hubC *server.KCPHub
hubD *server.KCPHub
serverCAddr string
serverDAddr string
relayC *countingPacketConn
relayD *countingPacketConn
cleanup func()
}
func startRelayedKCPHubs(t *testing.T) relayedKCPHubFixture {
t.Helper()
hubC := server.NewKCPHub()
serverCAddr, cleanupC := startRealKCPHubServer(t, hubC)
hubD := server.NewKCPHub()
serverDAddr, cleanupD := startRealKCPHubServer(t, hubD)
baseRelayC, err := net.ListenPacket("udp", "127.0.0.1:0")
if err != nil {
cleanupD()
cleanupC()
t.Fatalf("ListenPacket(relayC) error = %v", err)
}
relayC := &countingPacketConn{PacketConn: baseRelayC}
baseRelayD, err := net.ListenPacket("udp", "127.0.0.1:0")
if err != nil {
_ = relayC.Close()
cleanupD()
cleanupC()
t.Fatalf("ListenPacket(relayD) error = %v", err)
}
relayD := &countingPacketConn{PacketConn: baseRelayD}
hubC.SetRelaySocket(relayC, relayD.LocalAddr(), false)
hubD.SetRelaySocket(relayD, relayC.LocalAddr(), false)
stopRelayC := startRelayLoop(t, hubC, relayC)
stopRelayD := startRelayLoop(t, hubD, relayD)
cleanup := func() {
stopRelayC()
stopRelayD()
cleanupD()
cleanupC()
}
return relayedKCPHubFixture{
hubC: hubC,
hubD: hubD,
serverCAddr: serverCAddr,
serverDAddr: serverDAddr,
relayC: relayC,
relayD: relayD,
cleanup: cleanup,
}
}
func startRelayLoop(t *testing.T, hub *server.KCPHub, conn net.PacketConn) func() {
t.Helper()
var wg sync.WaitGroup
wg.Add(1)
go func() {
defer wg.Done()
if err := hub.ServeRelay(); err != nil && !isExpectedKCPRelayServeExit(err) {
t.Errorf("hub.ServeRelay() error = %v", err)
}
}()
return func() {
_ = conn.Close()
wg.Wait()
}
}
type countingPacketConn struct {
net.PacketConn
writeCount int32
}
func (c *countingPacketConn) WriteTo(p []byte, addr net.Addr) (int, error) {
atomic.AddInt32(&c.writeCount, 1)
return c.PacketConn.WriteTo(p, addr)
}
func (c *countingPacketConn) WriteCount() int {
return int(atomic.LoadInt32(&c.writeCount))
}
func isExpectedKCPHubServeExit(err error) bool {
if err == nil {
return true
@@ -261,3 +535,12 @@ func isExpectedKCPHubServeExit(err error) bool {
message := err.Error()
return strings.Contains(message, "closed") || strings.Contains(message, "broken pipe") || strings.Contains(message, "io: read/write on closed pipe")
}
func isExpectedKCPRelayServeExit(err error) bool {
if err == nil {
return true
}
message := err.Error()
return strings.Contains(message, "closed") || strings.Contains(message, "use of closed network connection")
}