usb-server/internal/relay/hub_test.go

package relay

import (
	"encoding/json"
	"testing"
	"time"

	"github.com/duffy/usb-server/internal/protocol"
	"github.com/gorilla/websocket"
)

// newTestClient builds a client without a socket. Nothing in the routing path
// touches Conn — only the write pump does, and these tests read Send directly.
func newTestClient(id, hash, mode string) *Client {
	return newClient(id, hash, mode, "test-"+id, nil)
}

// drain collects everything queued for a client without blocking.
func drain(c *Client) []outMsg {
	var msgs []outMsg
	for {
		select {
		case m := <-c.Send:
			msgs = append(msgs, m)
		default:
			return msgs
		}
	}
}

// typeOf extracts the "type" field of a queued JSON control message.
func typeOf(t *testing.T, m outMsg) string {
	t.Helper()
	var env protocol.Envelope
	if err := json.Unmarshal(m.data, &env); err != nil {
		t.Fatalf("queued message is not JSON: %v", err)
	}
	return env.Type
}

// countType drains a client and reports how many messages of one type it got.
// Counting by type rather than total keeps these assertions independent of the
// client_joined notifications registration produces.
func countType(t *testing.T, c *Client, msgType string) int {
	t.Helper()
	n := 0
	for _, m := range drain(c) {
		if typeOf(t, m) == msgType {
			n++
		}
	}
	return n
}

// registerAll registers every client, then drains them, so that no client is
// left holding join notifications from a peer that registered after it.
func registerAll(h *Hub, clients ...*Client) {
	for _, c := range clients {
		h.Register(c)
	}
	for _, c := range clients {
		drain(c)
	}
}

func TestDeviceListReachesUseAndBothButNotShare(t *testing.T) {
	h := NewHub()

	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
	user := newTestClient("user", "grp", protocol.ModeUse)
	both := newTestClient("both", "grp", protocol.ModeBoth)
	otherSharer := newTestClient("sharer2", "grp", protocol.ModeShare)

	registerAll(h, sharer, user, both, otherSharer)

	list, _ := json.Marshal(&protocol.DeviceList{
		Type:     protocol.MsgDeviceList,
		ClientID: sharer.ID,
		Devices:  []protocol.USBDevice{{BusID: "1-1"}},
	})
	h.HandleTextMessage(sharer, list)

	if got := countType(t, user, protocol.MsgDeviceList); got != 1 {
		t.Errorf("use client received %d device lists, want 1", got)
	}
	if got := countType(t, both, protocol.MsgDeviceList); got != 1 {
		t.Errorf("both client received %d device lists, want 1", got)
	}
	if got := countType(t, otherSharer, protocol.MsgDeviceList); got != 0 {
		t.Errorf("share-only client received %d device lists, want 0", got)
	}
	if got := countType(t, sharer, protocol.MsgDeviceList); got != 0 {
		t.Errorf("sender received %d copies of its own list, want 0", got)
	}
}

// A "both" client must be able to offer devices, which means its device list
// has to be routed like any share client's.
func TestBothClientCanShare(t *testing.T) {
	h := NewHub()

	both := newTestClient("both", "grp", protocol.ModeBoth)
	user := newTestClient("user", "grp", protocol.ModeUse)
	registerAll(h, both, user)

	list, _ := json.Marshal(&protocol.DeviceList{
		Type:     protocol.MsgDeviceList,
		ClientID: both.ID,
		Devices:  []protocol.USBDevice{{BusID: "2-1"}},
	})
	h.HandleTextMessage(both, list)

	if got := countType(t, user, protocol.MsgDeviceList); got != 1 {
		t.Fatalf("use client received %d lists from a both-mode sharer, want 1", got)
	}
}

func TestRequestDeviceReachesShareCapableTargetsOnly(t *testing.T) {
	h := NewHub()

	requester := newTestClient("req", "grp", protocol.ModeUse)
	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
	useOnly := newTestClient("useonly", "grp", protocol.ModeUse)

	registerAll(h, requester, sharer, useOnly)

	req, _ := json.Marshal(&protocol.RequestDevice{
		Type:         protocol.MsgRequestDevice,
		TargetClient: sharer.ID,
		BusID:        "1-1",
		RequestID:    "r1",
	})
	h.HandleTextMessage(requester, req)

	var msgs []outMsg
	for _, m := range drain(sharer) {
		if typeOf(t, m) == protocol.MsgRequestDevice {
			msgs = append(msgs, m)
		}
	}
	if len(msgs) != 1 {
		t.Fatalf("share client received %d requests, want 1", len(msgs))
	}

	// The relay must stamp in who is asking; the share side needs it to reply.
	var got map[string]interface{}
	json.Unmarshal(msgs[0].data, &got)
	if got["from_client"] != requester.ID {
		t.Errorf("from_client = %v, want %q", got["from_client"], requester.ID)
	}

	// A use-only client is not a valid target.
	req2, _ := json.Marshal(&protocol.RequestDevice{
		Type:         protocol.MsgRequestDevice,
		TargetClient: useOnly.ID,
		BusID:        "1-1",
		RequestID:    "r2",
	})
	h.HandleTextMessage(requester, req2)

	if got := countType(t, useOnly, protocol.MsgRequestDevice); got != 0 {
		t.Errorf("use-only client received %d device requests, want 0", got)
	}
}

func TestGroupsAreIsolatedByHash(t *testing.T) {
	h := NewHub()

	a := newTestClient("a", "hash-a", protocol.ModeShare)
	b := newTestClient("b", "hash-b", protocol.ModeUse)
	registerAll(h, a, b)

	list, _ := json.Marshal(&protocol.DeviceList{
		Type: protocol.MsgDeviceList, ClientID: a.ID,
	})
	h.HandleTextMessage(a, list)

	if got := len(drain(b)); got != 0 {
		t.Errorf("client in another hash group received %d messages, want 0", got)
	}
}

func TestTunnelForwardsBothWays(t *testing.T) {
	h := NewHub()

	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
	user := newTestClient("user", "grp", protocol.ModeUse)
	registerAll(h, sharer, user)

	tunnelID := "0123456789abcdef" // exactly TunnelHeaderSize
	granted, _ := json.Marshal(map[string]interface{}{
		"type":          protocol.MsgDeviceGranted,
		"bus_id":        "1-1",
		"tunnel_id":     tunnelID,
		"request_id":    "r1",
		"target_client": user.ID,
	})
	h.HandleTextMessage(sharer, granted)

	if msgs := drain(user); len(msgs) != 1 || typeOf(t, msgs[0]) != protocol.MsgDeviceGranted {
		t.Fatalf("grant was not forwarded to the use client: %v", msgs)
	}

	// use -> share
	frame := append([]byte(tunnelID), 0xAA, 0xBB)
	h.HandleBinaryMessage(user, frame)
	msgs := drain(sharer)
	if len(msgs) != 1 {
		t.Fatalf("share client received %d tunnel frames, want 1", len(msgs))
	}
	if msgs[0].typ != websocket.BinaryMessage {
		t.Errorf("tunnel frame sent as type %d, want binary", msgs[0].typ)
	}

	// share -> use
	h.HandleBinaryMessage(sharer, frame)
	if got := len(drain(user)); got != 1 {
		t.Errorf("use client received %d tunnel frames, want 1", got)
	}
}

func TestTunnelFramesForUnknownTunnelAreDropped(t *testing.T) {
	h := NewHub()

	a := newTestClient("a", "grp", protocol.ModeShare)
	b := newTestClient("b", "grp", protocol.ModeUse)
	registerAll(h, a, b)

	h.HandleBinaryMessage(a, append([]byte("nonexistenttunnl"), 0x01))

	if got := len(drain(b)); got != 0 {
		t.Errorf("frame for an unknown tunnel was forwarded (%d messages)", got)
	}
}

func TestUnregisterNotifiesPeersAndDropsTunnels(t *testing.T) {
	h := NewHub()

	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
	user := newTestClient("user", "grp", protocol.ModeUse)
	registerAll(h, sharer, user)

	tunnelID := "0123456789abcdef"
	granted, _ := json.Marshal(map[string]interface{}{
		"type": protocol.MsgDeviceGranted, "bus_id": "1-1",
		"tunnel_id": tunnelID, "target_client": user.ID,
	})
	h.HandleTextMessage(sharer, granted)
	drain(user)

	h.Unregister(sharer)

	msgs := drain(user)
	if len(msgs) != 1 || typeOf(t, msgs[0]) != protocol.MsgClientLeft {
		t.Fatalf("peer was not told about the disconnect: %v", msgs)
	}

	h.mu.RLock()
	_, stillThere := h.tunnels[tunnelID]
	h.mu.RUnlock()
	if stillThere {
		t.Error("tunnel survived the share client leaving")
	}
}

// Registration must not panic on short or empty identifiers: the relay
// truncated hashes for logging, so a client with a 3-character hash used to
// take the whole server down.
func TestRegisterSurvivesShortIdentifiers(t *testing.T) {
	h := NewHub()

	for _, c := range []*Client{
		newTestClient("", "", protocol.ModeUse),
		newTestClient("x", "ab", protocol.ModeShare),
		newTestClient("y", "abc", protocol.ModeBoth),
	} {
		h.Register(c)
		h.Unregister(c)
	}
}

// A client that stops draining must be dropped rather than allowed to consume
// unbounded memory or block the peer producing the traffic.
func TestFullSendQueueDropsClient(t *testing.T) {
	h := NewHub()

	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
	slow := newTestClient("slow", "grp", protocol.ModeUse)
	registerAll(h, sharer, slow)

	list, _ := json.Marshal(&protocol.DeviceList{
		Type: protocol.MsgDeviceList, ClientID: sharer.ID,
	})

	done := make(chan struct{})
	go func() {
		defer close(done)
		for i := 0; i < sendQueueDepth+50; i++ {
			h.HandleTextMessage(sharer, list)
		}
	}()

	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("routing blocked on a client that never reads")
	}

	select {
	case <-slow.dead:
	default:
		t.Error("client with a full queue was not dropped")
	}
}

// Reconnecting with the same ID must retire the stale entry, not leave two.
func TestReRegisterReplacesStaleClient(t *testing.T) {
	h := NewHub()

	first := newTestClient("dup", "grp", protocol.ModeUse)
	h.Register(first)

	second := newTestClient("dup", "grp", protocol.ModeUse)
	h.Register(second)

	select {
	case <-first.dead:
	default:
		t.Error("stale connection was not killed on re-registration")
	}

	if got := h.GroupStats()[protocol.ShortID("grp")]; got != 1 {
		t.Errorf("group holds %d clients, want 1", got)
	}
}

func TestValidModeAndCapabilities(t *testing.T) {
	tests := []struct {
		mode                    string
		valid, canShare, canUse bool
	}{
		{protocol.ModeShare, true, true, false},
		{protocol.ModeUse, true, false, true},
		{protocol.ModeBoth, true, true, true},
		{"", false, false, false},
		{"admin", false, false, false},
	}

	for _, tt := range tests {
		if got := protocol.ValidMode(tt.mode); got != tt.valid {
			t.Errorf("ValidMode(%q) = %v, want %v", tt.mode, got, tt.valid)
		}
		if got := protocol.CanShare(tt.mode); got != tt.canShare {
			t.Errorf("CanShare(%q) = %v, want %v", tt.mode, got, tt.canShare)
		}
		if got := protocol.CanUse(tt.mode); got != tt.canUse {
			t.Errorf("CanUse(%q) = %v, want %v", tt.mode, got, tt.canUse)
		}
	}
}

// The relay is the only party that knows a client's public address, so it
// must add it to a grant. Without this, two peers behind NAT could never find
// each other and every tunnel would stay relayed.
func TestGrantGetsPublicEndpointAppended(t *testing.T) {
	h := NewHub()

	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
	sharer.DirectPort = 41000
	sharer.PublicIP = "203.0.113.7"
	user := newTestClient("user", "grp", protocol.ModeUse)
	registerAll(h, sharer, user)

	granted, _ := json.Marshal(map[string]interface{}{
		"type":          protocol.MsgDeviceGranted,
		"bus_id":        "1-1",
		"tunnel_id":     "0123456789abcdef",
		"target_client": user.ID,
		"endpoints":     []string{"192.168.1.5:41000"},
		"encrypted":     true,
	})
	h.HandleTextMessage(sharer, granted)

	msgs := drain(user)
	if len(msgs) != 1 {
		t.Fatalf("use client received %d messages, want 1", len(msgs))
	}

	var got protocol.DeviceGranted
	if err := json.Unmarshal(msgs[0].data, &got); err != nil {
		t.Fatalf("decoding forwarded grant: %v", err)
	}

	want := "203.0.113.7:41000"
	var found, keptLocal bool
	for _, ep := range got.Endpoints {
		if ep == want {
			found = true
		}
		if ep == "192.168.1.5:41000" {
			keptLocal = true
		}
	}
	if !found {
		t.Errorf("endpoints %v do not include the public address %q", got.Endpoints, want)
	}
	if !keptLocal {
		t.Errorf("endpoints %v lost the sharer's own local address", got.Endpoints)
	}
	if !got.Encrypted {
		t.Error("the encrypted flag did not survive re-encoding")
	}
}

// A client that accepts no direct connections must not have a bogus endpoint
// invented for it.
func TestGrantWithoutDirectPortIsUnchanged(t *testing.T) {
	h := NewHub()

	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
	sharer.PublicIP = "203.0.113.7" // reachable, but no listener
	user := newTestClient("user", "grp", protocol.ModeUse)
	registerAll(h, sharer, user)

	granted, _ := json.Marshal(map[string]interface{}{
		"type": protocol.MsgDeviceGranted, "bus_id": "1-1",
		"tunnel_id": "0123456789abcdef", "target_client": user.ID,
	})
	h.HandleTextMessage(sharer, granted)

	msgs := drain(user)
	if len(msgs) != 1 {
		t.Fatalf("use client received %d messages, want 1", len(msgs))
	}

	var got protocol.DeviceGranted
	json.Unmarshal(msgs[0].data, &got)
	if len(got.Endpoints) != 0 {
		t.Errorf("endpoints = %v, want none for a client with no direct port", got.Endpoints)
	}
}

func TestPublicEndpointRequiresBothParts(t *testing.T) {
	tests := []struct {
		name string
		port int
		ip   string
		want string
	}{
		{"both present", 41000, "203.0.113.7", "203.0.113.7:41000"},
		{"no port", 0, "203.0.113.7", ""},
		{"no ip", 41000, "", ""},
		{"neither", 0, "", ""},
		{"ipv6", 41000, "2001:db8::1", "[2001:db8::1]:41000"},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			c := &Client{DirectPort: tt.port, PublicIP: tt.ip}
			if got := publicEndpoint(c); got != tt.want {
				t.Errorf("publicEndpoint() = %q, want %q", got, tt.want)
			}
		})
	}
}