package server import ( "net" "reflect" "strings" "sync" "testing" "time" "omnisocketgo/cmd/internal/latencylog" "omnisocketgo/cmd/internal/protocol" "omnisocketgo/cmd/internal/transport" ) type recordingLogger struct { mu sync.Mutex events []latencylog.Event } func (l *recordingLogger) LogEvent(event latencylog.Event) error { l.mu.Lock() defer l.mu.Unlock() l.events = append(l.events, event) return nil } func (l *recordingLogger) Events() []latencylog.Event { l.mu.Lock() defer l.mu.Unlock() return append([]latencylog.Event(nil), l.events...) } func TestServeConnRegistersPeer(t *testing.T) { hub := NewHub() client, done := startHubConn(t, hub) if err := client.Send(protocol.Message{ Type: protocol.MessageTypeRegister, From: "peer-a", To: protocol.ServerPeerID, }); err != nil { t.Fatalf("Send(register) error = %v", err) } waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered") if err := client.Close(); err != nil { t.Fatalf("client.Close() error = %v", err) } err := <-done if err == nil || !strings.Contains(err.Error(), "receive loop read") { t.Fatalf("ServeConn() error = %v, want read-loop shutdown error", err) } } func TestServeConnRejectsDuplicatePeerID(t *testing.T) { hub := NewHub() first, firstDone := startHubConn(t, hub) registerPeer(t, first, "peer-a") waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered") second, secondDone := startHubConn(t, hub) registerPeer(t, second, "peer-a") got, err := second.Receive() if err != nil { t.Fatalf("second.Receive() error = %v", err) } if got.Type != protocol.MessageTypeError { t.Fatalf("got message type %s, want %s", got.Type, protocol.MessageTypeError) } if string(got.Body) != "duplicate peer id: peer-a" { t.Fatalf("error body = %q, want duplicate peer message", got.Body) } waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "original peer-a to remain registered") if err := first.Close(); err != nil { t.Fatalf("first.Close() error = %v", err) } if err := <-secondDone; err == nil || !strings.Contains(err.Error(), "duplicate peer id") { t.Fatalf("second ServeConn() error = %v, want duplicate peer id error", err) } if err := <-firstDone; err == nil || !strings.Contains(err.Error(), "receive loop read") { t.Fatalf("first ServeConn() error = %v, want read-loop shutdown error", err) } } func TestServeConnForwardsMessages(t *testing.T) { tests := []struct { name string msg protocol.Message }{ { name: "text", msg: protocol.Message{ Type: protocol.MessageTypeText, ID: 1, From: "spoofed", To: "peer-b", Body: []byte("hello"), }, }, { name: "file", msg: protocol.Message{ Type: protocol.MessageTypeFile, ID: 2, From: "spoofed", To: "peer-b", FileName: "payload.bin", Body: []byte{0x01, 0x02, 0x03}, }, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { hub := NewHub() sender, senderDone := startHubConn(t, hub) receiver, receiverDone := startHubConn(t, hub) registerPeer(t, sender, "peer-a") registerPeer(t, receiver, "peer-b") waitFor(t, func() bool { return hub.HasPeer("peer-a") && hub.HasPeer("peer-b") }, "both peers to be registered") if err := sender.Send(tt.msg); err != nil { t.Fatalf("sender.Send() error = %v", err) } got, err := receiver.Receive() if err != nil { t.Fatalf("receiver.Receive() error = %v", err) } want := tt.msg want.From = "peer-a" if !reflect.DeepEqual(got, want) { t.Fatalf("forwarded message mismatch: got %+v want %+v", got, want) } _ = sender.Close() _ = receiver.Close() <-senderDone <-receiverDone }) } } func TestServeConnReturnsErrorForUnknownTarget(t *testing.T) { hub := NewHub() client, done := startHubConn(t, hub) registerPeer(t, client, "peer-a") waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered") if err := client.Send(protocol.Message{ Type: protocol.MessageTypeText, ID: 1, From: "peer-a", To: "missing-peer", Body: []byte("hello"), }); err != nil { t.Fatalf("Send(text) error = %v", err) } got, err := client.Receive() if err != nil { t.Fatalf("Receive() error = %v", err) } if got.Type != protocol.MessageTypeError { t.Fatalf("got message 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) } if !hub.HasPeer("peer-a") { t.Fatal("peer-a should remain registered after unknown target error") } _ = client.Close() <-done } func TestServeConnRejectsRegisterAfterRegistration(t *testing.T) { hub := NewHub() client, done := startHubConn(t, hub) registerPeer(t, client, "peer-a") waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered") if err := client.Send(protocol.Message{ Type: protocol.MessageTypeRegister, From: "peer-a", To: protocol.ServerPeerID, }); err != nil { t.Fatalf("Send(register again) error = %v", err) } got, err := client.Receive() if err != nil { t.Fatalf("Receive() error = %v", err) } if got.Type != protocol.MessageTypeError { t.Fatalf("got message type %s, want %s", got.Type, protocol.MessageTypeError) } if string(got.Body) != "registered peers can only send text or file messages" { t.Fatalf("error body = %q, want registered-peer protocol error", got.Body) } if err := <-done; err == nil || !strings.Contains(err.Error(), "unexpected message type from peer peer-a: register") { t.Fatalf("ServeConn() error = %v, want unexpected register message error", err) } } func TestServeConnUnregistersPeerOnClose(t *testing.T) { hub := NewHub() client, done := startHubConn(t, hub) registerPeer(t, client, "peer-a") waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered") if err := client.Close(); err != nil { t.Fatalf("client.Close() error = %v", err) } <-done waitFor(t, func() bool { return !hub.HasPeer("peer-a") }, "peer-a to be unregistered") } func TestServeConnDoesNotEmitEndpointLatencyEventsOnForward(t *testing.T) { logger := &recordingLogger{} hub := NewHub(WithLogger(logger)) sender, senderDone := startHubConn(t, hub) receiver, receiverDone := startHubConn(t, hub) registerPeer(t, sender, "peer-a") registerPeer(t, receiver, "peer-b") waitFor(t, func() bool { return hub.HasPeer("peer-a") && hub.HasPeer("peer-b") }, "both peers to be registered") msg := protocol.Message{ Type: protocol.MessageTypeText, ID: 11, From: "spoofed", To: "peer-b", Body: []byte("hello"), } if err := sender.Send(msg); err != nil { t.Fatalf("sender.Send() error = %v", err) } got, err := receiver.Receive() if err != nil { t.Fatalf("receiver.Receive() error = %v", err) } msg.From = "peer-a" if !reflect.DeepEqual(got, msg) { t.Fatalf("forwarded message mismatch: got %+v want %+v", got, msg) } events := logger.Events() if len(events) != 0 { t.Fatalf("event count = %d, want 0 because server is a black-box relay", len(events)) } _ = sender.Close() _ = receiver.Close() <-senderDone <-receiverDone } func TestServeConnDoesNotLogLatencyEventsForUnknownTarget(t *testing.T) { logger := &recordingLogger{} hub := NewHub(WithLogger(logger)) client, done := startHubConn(t, hub) registerPeer(t, client, "peer-a") waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered") if err := client.Send(protocol.Message{ Type: protocol.MessageTypeText, ID: 15, From: "peer-a", To: "missing-peer", Body: []byte("hello"), }); err != nil { t.Fatalf("Send(text) error = %v", err) } got, err := client.Receive() if err != nil { t.Fatalf("Receive() error = %v", err) } if got.Type != protocol.MessageTypeError { t.Fatalf("got message type %s, want %s", got.Type, protocol.MessageTypeError) } if events := logger.Events(); len(events) != 0 { t.Fatalf("event count = %d, want 0 for unknown target path", len(events)) } _ = client.Close() <-done } func TestServeConnDoesNotLogLatencyEventsForDuplicateRegister(t *testing.T) { logger := &recordingLogger{} hub := NewHub(WithLogger(logger)) first, firstDone := startHubConn(t, hub) registerPeer(t, first, "peer-a") waitFor(t, func() bool { return hub.HasPeer("peer-a") }, "peer-a to be registered") second, secondDone := startHubConn(t, hub) registerPeer(t, second, "peer-a") got, err := second.Receive() if err != nil { t.Fatalf("second.Receive() error = %v", err) } if got.Type != protocol.MessageTypeError { t.Fatalf("got type %s, want %s", got.Type, protocol.MessageTypeError) } if events := logger.Events(); len(events) != 0 { t.Fatalf("event count = %d, want 0 for duplicate register path", len(events)) } _ = first.Close() <-secondDone <-firstDone } func startHubConn(t *testing.T, hub *Hub) (*transport.TCPConn, <-chan error) { t.Helper() listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { t.Fatalf("net.Listen() error = %v", err) } done := make(chan error, 1) go func() { serverSide, acceptErr := listener.Accept() if acceptErr != nil { done <- acceptErr return } done <- hub.ServeConn(serverSide) }() clientSide, err := net.Dial("tcp", listener.Addr().String()) if err != nil { _ = listener.Close() t.Fatalf("net.Dial() error = %v", err) } if err := listener.Close(); err != nil { t.Fatalf("listener.Close() error = %v", err) } conn, err := transport.NewTCPConn(clientSide) if err != nil { _ = clientSide.Close() t.Fatalf("transport.NewTCPConn() error = %v", err) } return conn, done } func registerPeer(t *testing.T, conn *transport.TCPConn, peerID string) { t.Helper() if err := conn.Send(protocol.Message{ Type: protocol.MessageTypeRegister, From: peerID, To: protocol.ServerPeerID, }); err != nil { t.Fatalf("Send(register %s) error = %v", peerID, err) } } func waitFor(t *testing.T, condition func() bool, description string) { t.Helper() deadline := time.Now().Add(500 * time.Millisecond) for time.Now().Before(deadline) { if condition() { return } time.Sleep(10 * time.Millisecond) } t.Fatalf("timed out waiting for %s", description) }