OmniSocketGo/cmd/internal/server/hub_test.go

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)
}