SimpleRemoter/server/go/hub/hub_test.go

package hub

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"github.com/yuanyuanxiang/SimpleRemoter/server/go/connection"
)

// stubCtx returns a non-nil *connection.Context useful only as a sentinel.
// Tests never invoke Send / Close on it.
func stubCtx() *connection.Context { return &connection.Context{} }

func TestHubRegisterListUnregister(t *testing.T) {
	h := New()
	if got := h.Count(); got != 0 {
		t.Fatalf("empty hub: want Count=0, got %d", got)
	}

	h.Register(&Device{ID: "a", Name: "Alice", ConnectedAt: time.Now()}, stubCtx())
	h.Register(&Device{ID: "b", Name: "Bob", ConnectedAt: time.Now()}, stubCtx())
	if got := h.Count(); got != 2 {
		t.Fatalf("after 2 registers: want Count=2, got %d", got)
	}

	list := h.ListDevices()
	if len(list) != 2 {
		t.Fatalf("want 2 devices in list, got %d", len(list))
	}

	h.Unregister("a")
	if got := h.Count(); got != 1 {
		t.Fatalf("after unregister: want Count=1, got %d", got)
	}

	// Unregister non-existent ID is a no-op
	h.Unregister("ghost")
	if got := h.Count(); got != 1 {
		t.Fatalf("after no-op unregister: want Count=1, got %d", got)
	}
}

func TestHubNilAndEmptyIgnored(t *testing.T) {
	h := New()
	h.Register(nil, stubCtx())
	h.Register(&Device{ID: ""}, stubCtx())
	h.Register(&Device{ID: "valid"}, nil) // nil conn should also be rejected
	h.Unregister("")
	if got := h.Count(); got != 0 {
		t.Fatalf("nil/empty register should be no-op, got Count=%d", got)
	}
}

type captureHandler struct {
	mu      sync.Mutex
	online  []string
	offline []string
	updates []string // formatted "id:rtt"
}

func (c *captureHandler) OnDeviceOnline(d DeviceInfo) {
	c.mu.Lock()
	c.online = append(c.online, d.ID)
	c.mu.Unlock()
}

func (c *captureHandler) OnDeviceOffline(id string) {
	c.mu.Lock()
	c.offline = append(c.offline, id)
	c.mu.Unlock()
}

func (c *captureHandler) OnDeviceUpdate(id string, rtt int, _ string) {
	c.mu.Lock()
	c.updates = append(c.updates, fmt.Sprintf("%s:%d", id, rtt))
	c.mu.Unlock()
}

func (c *captureHandler) OnScreenFrame(_ string, _ []byte, _ bool) {}

func (c *captureHandler) OnResolutionChange(_ string, _, _ int) {}

func (c *captureHandler) OnCursorChange(_ string, _ byte) {}

func (c *captureHandler) OnClipboardText(_ string, _ string) {}

func (c *captureHandler) OnTerminalReady(_ string, _ bool) {}

func (c *captureHandler) OnTerminalData(_ string, _ []byte) {}

func (c *captureHandler) OnTerminalClosed(_ string, _ string) {}

func TestHubSubscribeEvents(t *testing.T) {
	h := New()
	c := &captureHandler{}
	unsub := h.Subscribe(c)

	h.Register(&Device{ID: "x", Name: "x"}, stubCtx())
	h.Register(&Device{ID: "y", Name: "y"}, stubCtx())
	h.Unregister("x")
	h.Unregister("nonexistent") // no event

	if len(c.online) != 2 || c.online[0] != "x" || c.online[1] != "y" {
		t.Fatalf("online events: %+v", c.online)
	}
	if len(c.offline) != 1 || c.offline[0] != "x" {
		t.Fatalf("offline events: %+v", c.offline)
	}

	unsub()
	h.Register(&Device{ID: "z"}, stubCtx())
	if len(c.online) != 2 {
		t.Fatalf("after unsubscribe should not receive events: %+v", c.online)
	}
}

func TestHubUpdateLive(t *testing.T) {
	h := New()
	c := &captureHandler{}
	h.Subscribe(c)

	h.Register(&Device{ID: "x", Name: "x"}, stubCtx())
	h.UpdateLive("x", 42, "Notepad")
	h.UpdateLive("ghost", 999, "should be ignored") // unknown id, no event

	if len(c.updates) != 1 || c.updates[0] != "x:42" {
		t.Fatalf("updates: %+v", c.updates)
	}

	list := h.ListDevices()
	if list[0].RTT != 42 || list[0].ActiveWindow != "Notepad" {
		t.Fatalf("live fields not applied: %+v", list[0])
	}
}

func TestHubRegisterOverwrites(t *testing.T) {
	h := New()
	h.Register(&Device{ID: "x", Name: "first"}, stubCtx())
	h.Register(&Device{ID: "x", Name: "second"}, stubCtx())
	list := h.ListDevices()
	if len(list) != 1 || list[0].Name != "second" {
		t.Fatalf("re-register should overwrite, got %+v", list)
	}
}

// ctxFromIP builds a *connection.Context whose source IP resolves to ip.
// BindScreenConn/BindTerminalConn compare this against Device.PeerIP.
func ctxFromIP(ip string) *connection.Context {
	return &connection.Context{RemoteAddr: ip + ":54321"}
}

// Sub-conns (screen/terminal) are opened device→server from the same host as
// the login conn. A caller claiming a MasterID from a different source IP is
// not that device and must not be able to bind, replace, or hijack the stream.
func TestHubBindScreenConnIPBinding(t *testing.T) {
	h := New()
	h.Register(&Device{ID: "dev", PeerIP: "10.0.0.1"}, stubCtx())

	// First bind from the device's own IP succeeds.
	if !h.BindScreenConn("dev", ctxFromIP("10.0.0.1")) {
		t.Fatal("same-IP screen bind should succeed")
	}
	// A second conn claiming the same MasterID from a foreign IP is refused,
	// so it cannot displace the legitimate stream.
	if h.BindScreenConn("dev", ctxFromIP("203.0.113.9")) {
		t.Fatal("foreign-IP screen bind should be rejected")
	}

	// A device that has never bound a screen conn is equally protected.
	h.Register(&Device{ID: "fresh", PeerIP: "10.0.0.2"}, stubCtx())
	if h.BindScreenConn("fresh", ctxFromIP("203.0.113.9")) {
		t.Fatal("foreign-IP first screen bind should be rejected")
	}
}

func TestHubBindTerminalConnIPBinding(t *testing.T) {
	h := New()
	h.Register(&Device{ID: "dev", PeerIP: "10.0.0.1"}, stubCtx())

	openPending := func(id string) {
		h.mu.Lock()
		h.devices[id].terminalPending = true
		h.mu.Unlock()
	}
	openPending("dev")

	// Same-IP terminal bind succeeds and consumes the pending slot.
	if !h.BindTerminalConn("dev", ctxFromIP("10.0.0.1"), true) {
		t.Fatal("same-IP terminal bind should succeed")
	}

	// A second device with a pending session cannot be claimed from elsewhere.
	h.Register(&Device{ID: "dev2", PeerIP: "10.0.0.2"}, stubCtx())
	openPending("dev2")
	if h.BindTerminalConn("dev2", ctxFromIP("203.0.113.9"), true) {
		t.Fatal("foreign-IP terminal bind should be rejected")
	}
}

// Race detector should not fire under `go test -race ./hub/...`.
func TestHubConcurrent(t *testing.T) {
	h := New()
	const goroutines = 50
	const opsPer = 100

	var wg sync.WaitGroup
	for g := range goroutines {
		wg.Add(1)
		go func(g int) {
			defer wg.Done()
			for i := range opsPer {
				id := fmt.Sprintf("g%d-%d", g, i)
				h.Register(&Device{ID: id, Name: id, ConnectedAt: time.Now()}, stubCtx())
				_ = h.ListDevices()
				_ = h.Count()
				h.Unregister(id)
			}
		}(g)
	}
	wg.Wait()

	if got := h.Count(); got != 0 {
		t.Fatalf("after all unregisters: want 0, got %d", got)
	}
}