package peer import ( "os" "path/filepath" "reflect" "strings" "sync" "testing" kcp "github.com/xtaci/kcp-go/v5" "omnisocketgo/cmd/internal/latencylog" "omnisocketgo/cmd/internal/protocol" "omnisocketgo/cmd/internal/server" "omnisocketgo/cmd/internal/transport" ) func TestKCPDialRegistersPeer(t *testing.T) { hub := server.NewKCPHub() serverAddr, cleanup := startRealKCPHubServer(t, hub) defer cleanup() client, err := DialKCP(serverAddr, "peer-a") if err != nil { t.Fatalf("DialKCP() error = %v", err) } defer func() { _ = client.Close() }() waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered") } func TestKCPDialRejectsInvalidBindIP(t *testing.T) { _, err := DialKCP("127.0.0.1:9002", "peer-a", WithBindIP("not-an-ip")) if err == nil { t.Fatal("DialKCP() error = nil, want invalid bind ip error") } if !strings.Contains(err.Error(), `invalid bind ip "not-an-ip"`) { t.Fatalf("DialKCP() error = %v, want invalid bind ip error", err) } } func TestKCPClientsExchangeTextAndFileMessages(t *testing.T) { hub := server.NewKCPHub() serverAddr, cleanup := startRealKCPHubServer(t, hub) defer cleanup() peerA, err := DialKCP(serverAddr, "peer-a") if err != nil { t.Fatalf("DialKCP(peer-a) error = %v", err) } defer func() { _ = peerA.Close() }() peerB, err := DialKCP(serverAddr, "peer-b") if err != nil { t.Fatalf("DialKCP(peer-b) error = %v", err) } defer func() { _ = peerB.Close() }() waitFor(t, func() bool { return hub.HasPeer("peer-a") && hub.HasPeer("peer-b") }, "both peers to be registered") received := make(chan protocol.Message, 2) receiveErr := make(chan error, 1) go func() { for i := 0; i < 2; i++ { msg, err := peerB.Receive() if err != nil { receiveErr <- err return } received <- msg } receiveErr <- nil }() if err := peerA.SendText("peer-b", "hello over kcp"); err != nil { t.Fatalf("SendText() error = %v", err) } fileBody := []byte{0x01, 0x02, 0x03} if err := peerA.SendFile("peer-b", "payload.bin", fileBody); err != nil { t.Fatalf("SendFile() error = %v", err) } if err := <-receiveErr; err != nil { t.Fatalf("peerB.Receive() error = %v", err) } gotFirst := <-received wantFirst := protocol.Message{ Type: protocol.MessageTypeText, ID: 1, From: "peer-a", To: "peer-b", Body: []byte("hello over kcp"), } if !reflect.DeepEqual(gotFirst, wantFirst) { t.Fatalf("first message mismatch: got %+v want %+v", gotFirst, wantFirst) } gotSecond := <-received wantSecond := protocol.Message{ Type: protocol.MessageTypeFile, ID: 2, From: "peer-a", To: "peer-b", FileName: "payload.bin", Body: fileBody, } if !reflect.DeepEqual(gotSecond, wantSecond) { t.Fatalf("second message mismatch: got %+v want %+v", gotSecond, wantSecond) } } func TestKCPClientReceivesServerErrorForUnknownTarget(t *testing.T) { hub := server.NewKCPHub() serverAddr, cleanup := startRealKCPHubServer(t, hub) defer cleanup() client, err := DialKCP(serverAddr, "peer-a") if err != nil { t.Fatalf("DialKCP() error = %v", err) } defer func() { _ = client.Close() }() waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered") if err := client.SendText("missing-peer", "hello"); err != nil { t.Fatalf("SendText() error = %v", err) } got, err := client.Receive() if err != nil { t.Fatalf("Receive() error = %v", err) } if got.Type != protocol.MessageTypeError { t.Fatalf("got type %s, want %s", got.Type, protocol.MessageTypeError) } if string(got.Body) != "unknown target: missing-peer" { t.Fatalf("error body = %q, want unknown target message", got.Body) } } func TestKCPClientsExchangeMessagesWithLatencyLogs(t *testing.T) { hub := server.NewKCPHub() serverAddr, cleanup := startRealKCPHubServer(t, hub) defer cleanup() peerALogger := &recordingLogger{} peerA, err := DialKCP(serverAddr, "peer-a", WithLogger(peerALogger)) if err != nil { t.Fatalf("DialKCP(peer-a) error = %v", err) } defer func() { _ = peerA.Close() }() peerBLogger := &recordingLogger{} peerB, err := DialKCP(serverAddr, "peer-b", WithLogger(peerBLogger)) if err != nil { t.Fatalf("DialKCP(peer-b) error = %v", err) } defer func() { _ = peerB.Close() }() inboxDir := t.TempDir() waitFor(t, func() bool { return hub.HasPeer("peer-a") && hub.HasPeer("peer-b") }, "both peers to be registered") if err := peerA.SendText("peer-b", "hello"); err != nil { t.Fatalf("SendText() error = %v", err) } textMsg, err := peerB.Receive() if err != nil { t.Fatalf("peerB.Receive(text) error = %v", err) } if _, err := peerB.PersistMessage(textMsg, inboxDir); err != nil { t.Fatalf("peerB.PersistMessage(text) error = %v", err) } filePath := filepath.Join(t.TempDir(), "payload.bin") if err := os.WriteFile(filePath, []byte{0x01, 0x02, 0x03}, 0o644); err != nil { t.Fatalf("os.WriteFile() error = %v", err) } if err := peerA.SendFilePath("peer-b", filePath); err != nil { t.Fatalf("SendFilePath() error = %v", err) } fileMsg, err := peerB.Receive() if err != nil { t.Fatalf("peerB.Receive(file) error = %v", err) } if _, err := peerB.PersistMessage(fileMsg, inboxDir); err != nil { t.Fatalf("peerB.PersistMessage(file) error = %v", err) } waitFor(t, func() bool { return len(peerALogger.Events()) == 6 }, "peer-a latency events") waitFor(t, func() bool { return len(peerBLogger.Events()) == 6 }, "peer-b latency events") assertEventSequencesByMessage(t, peerALogger.Events(), map[uint64][]string{ 1: {latencylog.EventAAppPrepBegin, latencylog.EventSendHandoffBegin, latencylog.EventSendHandoffEnd}, 2: {latencylog.EventAAppPrepBegin, latencylog.EventSendHandoffBegin, latencylog.EventSendHandoffEnd}, }) assertEventSequencesByMessage(t, peerBLogger.Events(), map[uint64][]string{ 1: {latencylog.EventBAppRecv, latencylog.EventBPersistBegin, latencylog.EventBPersistEnd}, 2: {latencylog.EventBAppRecv, latencylog.EventBPersistBegin, latencylog.EventBPersistEnd}, }) } func startRealKCPHubServer(t *testing.T, hub *server.KCPHub) (string, func()) { t.Helper() 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 && !isExpectedKCPHubServeExit(serveErr) { t.Logf("hub.ServeSession() ended with %v", serveErr) } }(session) } }() cleanup := func() { close(stop) _ = listener.Close() _ = packetConn.Close() wg.Wait() } return listener.Addr().String(), cleanup } func isExpectedKCPHubServeExit(err error) bool { if err == nil { return true } message := err.Error() return strings.Contains(message, "closed") || strings.Contains(message, "broken pipe") || strings.Contains(message, "io: read/write on closed pipe") }