usb-server/internal/relay/hub.go

package relay

import (
	"encoding/json"
	"log"
	"net"
	"strconv"
	"sync"

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

// sendQueueDepth bounds per-client outgoing backlog. A client that falls this
// far behind is not going to catch up, and buffering more would let one stuck
// peer consume the relay's memory.
const sendQueueDepth = 256

// outMsg is one queued WebSocket frame.
type outMsg struct {
	typ  int // websocket.TextMessage or websocket.BinaryMessage
	data []byte
}

// Client represents a connected WebSocket client
type Client struct {
	ID   string
	Hash string
	Mode string // "share", "use" or "both"
	Name string
	Conn *websocket.Conn

	// DirectPort is the port this client accepts direct tunnel connections on,
	// 0 if it accepts none.
	DirectPort int

	// PublicIP is the source address the relay sees this client connect from.
	// Peers cannot determine their own public address, so the relay supplies
	// it when passing on a grant — that is the whole reason it is involved in
	// setting up connections that then bypass it.
	PublicIP string

	// Send carries outgoing frames to this client's write pump. All writes go
	// through it: writing to the socket directly from another client's read
	// loop would block that peer — and, because the hub held its lock across
	// the write, every other client with it.
	Send chan outMsg

	closeOnce sync.Once
	dead      chan struct{}
}

// newClient creates a client with its outgoing queue ready.
func newClient(id, hash, mode, name string, conn *websocket.Conn) *Client {
	return &Client{
		ID:   id,
		Hash: hash,
		Mode: mode,
		Name: name,
		Conn: conn,
		Send: make(chan outMsg, sendQueueDepth),
		dead: make(chan struct{}),
	}
}

// enqueue queues a frame without blocking.
// It reports false when the client's queue is full or it is already gone; the
// caller should treat that as a disconnect rather than retrying.
func (c *Client) enqueue(typ int, data []byte) bool {
	select {
	case <-c.dead:
		return false
	default:
	}

	select {
	case c.Send <- outMsg{typ: typ, data: data}:
		return true
	case <-c.dead:
		return false
	default:
		log.Printf("[hub] send queue full for %s (%s), dropping client",
			protocol.ShortID(c.ID), c.Name)
		c.kill()
		return false
	}
}

// kill marks the client dead and wakes its write pump. Idempotent.
func (c *Client) kill() {
	c.closeOnce.Do(func() { close(c.dead) })
}

// enqueueJSON marshals and queues a JSON control message.
func (c *Client) enqueueJSON(v interface{}) bool {
	data, err := json.Marshal(v)
	if err != nil {
		return false
	}
	return c.enqueue(websocket.TextMessage, data)
}

// Hub manages all connected clients and routes messages between them
type Hub struct {
	mu      sync.RWMutex
	groups  map[string]map[string]*Client // hash -> client_id -> client
	tunnels map[string]*Tunnel            // tunnel_id -> tunnel info
}

// Tunnel tracks an active USB/IP tunnel between two clients
type Tunnel struct {
	ID          string
	ShareClient string
	UseClient   string
	BusID       string
}

// NewHub creates a new Hub
func NewHub() *Hub {
	return &Hub{
		groups:  make(map[string]map[string]*Client),
		tunnels: make(map[string]*Tunnel),
	}
}

// peers returns a snapshot of the clients in a hash group, excluding one ID.
//
// Taking a snapshot and releasing the lock before doing anything with the
// clients is deliberate: holding the hub lock across a send is what let a
// single slow peer stall registration and routing for everyone.
func (h *Hub) peers(hash, excludeID string) []*Client {
	h.mu.RLock()
	defer h.mu.RUnlock()

	group := h.groups[hash]
	result := make([]*Client, 0, len(group))
	for _, c := range group {
		if c.ID != excludeID {
			result = append(result, c)
		}
	}
	return result
}

// peer looks up a single client in a hash group.
func (h *Hub) peer(hash, clientID string) *Client {
	h.mu.RLock()
	defer h.mu.RUnlock()

	group := h.groups[hash]
	if group == nil {
		return nil
	}
	return group[clientID]
}

// Register adds a client to its hash group
func (h *Hub) Register(client *Client) {
	h.mu.Lock()
	if h.groups[client.Hash] == nil {
		h.groups[client.Hash] = make(map[string]*Client)
	}
	// A reconnecting client reuses its ID; drop the stale entry so its
	// write pump exits instead of lingering with a dead socket.
	if old, exists := h.groups[client.Hash][client.ID]; exists && old != client {
		old.kill()
	}
	h.groups[client.Hash][client.ID] = client
	h.mu.Unlock()

	log.Printf("[hub] client registered: id=%s hash=%s mode=%s name=%s",
		protocol.ShortID(client.ID), protocol.ShortID(client.Hash), client.Mode, client.Name)

	// Notify other clients in the group
	joined := &protocol.ClientJoined{
		Type:     protocol.MsgClientJoined,
		ClientID: client.ID,
		Mode:     client.Mode,
		Name:     client.Name,
	}
	for _, peer := range h.peers(client.Hash, client.ID) {
		peer.enqueueJSON(joined)
	}
}

// Unregister removes a client and cleans up its tunnels
func (h *Hub) Unregister(client *Client) {
	h.mu.Lock()
	group := h.groups[client.Hash]
	if group == nil {
		h.mu.Unlock()
		return
	}

	// Only remove this exact client: a reconnect may already have installed a
	// newer connection under the same ID.
	if group[client.ID] == client {
		delete(group, client.ID)
	}
	if len(group) == 0 {
		delete(h.groups, client.Hash)
	}

	// Clean up tunnels involving this client
	for tid, tunnel := range h.tunnels {
		if tunnel.ShareClient == client.ID || tunnel.UseClient == client.ID {
			delete(h.tunnels, tid)
		}
	}
	h.mu.Unlock()

	client.kill()

	log.Printf("[hub] client unregistered: id=%s name=%s", protocol.ShortID(client.ID), client.Name)

	left := &protocol.ClientLeft{
		Type:     protocol.MsgClientLeft,
		ClientID: client.ID,
	}
	for _, peer := range h.peers(client.Hash, client.ID) {
		peer.enqueueJSON(left)
	}
}

// HandleTextMessage processes a JSON control message
func (h *Hub) HandleTextMessage(sender *Client, data []byte) {
	var env protocol.Envelope
	if err := json.Unmarshal(data, &env); err != nil {
		log.Printf("[hub] invalid message from %s: %v", protocol.ShortID(sender.ID), err)
		return
	}

	switch env.Type {
	case protocol.MsgDeviceList:
		h.handleDeviceList(sender, data)
	case protocol.MsgRequestDevice:
		h.handleRequestDevice(sender, data)
	case protocol.MsgDeviceGranted:
		h.handleDeviceGranted(sender, data)
	case protocol.MsgDeviceDenied:
		h.handleDeviceDenied(sender, data)
	case protocol.MsgForceRelease:
		h.handleForceRelease(sender, data)
	case protocol.MsgReleaseDevice:
		h.handleReleaseDevice(sender, data)
	case protocol.MsgDeviceReleased:
		h.handleDeviceReleased(sender, data)
	case protocol.MsgPing:
		sender.enqueueJSON(&protocol.Pong{Type: protocol.MsgPong})
	default:
		log.Printf("[hub] unknown message type from %s: %s", protocol.ShortID(sender.ID), env.Type)
	}
}

// HandleBinaryMessage forwards tunnel data to the other end
func (h *Hub) HandleBinaryMessage(sender *Client, data []byte) {
	if len(data) < protocol.TunnelHeaderSize {
		return
	}

	tunnelID := string(data[:protocol.TunnelHeaderSize])

	h.mu.RLock()
	tunnel := h.tunnels[tunnelID]
	h.mu.RUnlock()

	if tunnel == nil {
		return
	}

	// Forward to the other end of the tunnel
	var targetID string
	switch sender.ID {
	case tunnel.ShareClient:
		targetID = tunnel.UseClient
	case tunnel.UseClient:
		targetID = tunnel.ShareClient
	default:
		return
	}

	if target := h.peer(sender.Hash, targetID); target != nil {
		target.enqueue(websocket.BinaryMessage, data)
	}
}

// handleDeviceList broadcasts a device list to every client in the group that
// can consume devices.
func (h *Hub) handleDeviceList(sender *Client, data []byte) {
	if !protocol.CanShare(sender.Mode) {
		return
	}

	for _, client := range h.peers(sender.Hash, sender.ID) {
		if protocol.CanUse(client.Mode) {
			client.enqueue(websocket.TextMessage, data)
		}
	}
}

// handleRequestDevice forwards a device request to the target share client
func (h *Hub) handleRequestDevice(sender *Client, data []byte) {
	var msg protocol.RequestDevice
	if err := json.Unmarshal(data, &msg); err != nil {
		return
	}

	target := h.peer(sender.Hash, msg.TargetClient)
	if target == nil || !protocol.CanShare(target.Mode) {
		return
	}

	// Add the sender's ID so the share client knows who's requesting
	target.enqueueJSON(map[string]interface{}{
		"type":          protocol.MsgRequestDevice,
		"target_client": msg.TargetClient,
		"bus_id":        msg.BusID,
		"request_id":    msg.RequestID,
		"from_client":   sender.ID,
	})
}

// handleDeviceGranted registers the tunnel and forwards to the requesting client
func (h *Hub) handleDeviceGranted(sender *Client, data []byte) {
	var granted struct {
		protocol.DeviceGranted
		TargetClient string `json:"target_client"`
	}
	if err := json.Unmarshal(data, &granted); err != nil {
		return
	}
	if granted.TunnelID == "" {
		return
	}

	h.mu.Lock()
	h.tunnels[granted.TunnelID] = &Tunnel{
		ID:          granted.TunnelID,
		ShareClient: sender.ID,
		UseClient:   granted.TargetClient,
		BusID:       granted.BusID,
	}
	h.mu.Unlock()

	log.Printf("[hub] tunnel created: %s (share=%s, use=%s, device=%s)",
		granted.TunnelID, protocol.ShortID(sender.ID), protocol.ShortID(granted.TargetClient), granted.BusID)

	target := h.peer(sender.Hash, granted.TargetClient)
	if target == nil {
		return
	}

	// Add the address we see the granting client at. It cannot know its own
	// public address, and this is what lets the two peers connect directly
	// across NAT and take their USB traffic off this relay entirely.
	out := data
	if extra := publicEndpoint(sender); extra != "" {
		granted.Endpoints = appendUnique(granted.Endpoints, extra)
		if reencoded, err := json.Marshal(granted); err == nil {
			out = reencoded
		}
	}

	target.enqueue(websocket.TextMessage, out)
}

// publicEndpoint builds the host:port at which a client's direct listener
// should be reachable from outside, or "" if it accepts no direct connections.
func publicEndpoint(c *Client) string {
	if c.DirectPort == 0 || c.PublicIP == "" {
		return ""
	}
	return net.JoinHostPort(c.PublicIP, strconv.Itoa(c.DirectPort))
}

// appendUnique adds an entry unless it is already present.
func appendUnique(list []string, item string) []string {
	for _, existing := range list {
		if existing == item {
			return list
		}
	}
	return append(list, item)
}

// handleDeviceDenied forwards denial to the requesting client
func (h *Hub) handleDeviceDenied(sender *Client, data []byte) {
	var denied struct {
		protocol.DeviceDenied
		TargetClient string `json:"target_client"`
	}
	if err := json.Unmarshal(data, &denied); err != nil {
		return
	}

	if target := h.peer(sender.Hash, denied.TargetClient); target != nil {
		target.enqueue(websocket.TextMessage, data)
	}
}

// handleReleaseDevice forwards a release to the share client
func (h *Hub) handleReleaseDevice(sender *Client, data []byte) {
	var msg protocol.ReleaseDevice
	if err := json.Unmarshal(data, &msg); err != nil {
		return
	}

	// Clean up tunnel
	h.mu.Lock()
	for tid, tunnel := range h.tunnels {
		if tunnel.UseClient == sender.ID && tunnel.BusID == msg.BusID {
			delete(h.tunnels, tid)
			log.Printf("[hub] tunnel closed: %s", tid)
			break
		}
	}
	h.mu.Unlock()

	if target := h.peer(sender.Hash, msg.TargetClient); target != nil {
		target.enqueueJSON(map[string]interface{}{
			"type":          protocol.MsgReleaseDevice,
			"target_client": msg.TargetClient,
			"bus_id":        msg.BusID,
			"from_client":   sender.ID,
		})
	}
}

// handleForceRelease forwards a force-release request to the target share client
func (h *Hub) handleForceRelease(sender *Client, data []byte) {
	var msg protocol.ForceRelease
	if err := json.Unmarshal(data, &msg); err != nil {
		return
	}

	// Clean up tunnel for this device (by BusID, regardless of who owns it)
	h.mu.Lock()
	for tid, tunnel := range h.tunnels {
		if tunnel.BusID == msg.BusID && tunnel.ShareClient == msg.TargetClient {
			delete(h.tunnels, tid)
			log.Printf("[hub] tunnel force-closed: %s", tid)
			break
		}
	}
	h.mu.Unlock()

	target := h.peer(sender.Hash, msg.TargetClient)
	if target == nil || !protocol.CanShare(target.Mode) {
		return
	}

	target.enqueueJSON(map[string]interface{}{
		"type":          protocol.MsgForceRelease,
		"target_client": msg.TargetClient,
		"bus_id":        msg.BusID,
		"from_client":   sender.ID,
	})
}

// handleDeviceReleased broadcasts device released notification
func (h *Hub) handleDeviceReleased(sender *Client, data []byte) {
	for _, client := range h.peers(sender.Hash, sender.ID) {
		if protocol.CanUse(client.Mode) {
			client.enqueue(websocket.TextMessage, data)
		}
	}
}

// GroupStats reports the number of clients per hash group, for diagnostics.
func (h *Hub) GroupStats() map[string]int {
	h.mu.RLock()
	defer h.mu.RUnlock()

	stats := make(map[string]int, len(h.groups))
	for hash, group := range h.groups {
		stats[protocol.ShortID(hash)] = len(group)
	}
	return stats
}