usb-server/internal/client/use.go

package client

import (
	"fmt"
	"log"
	"net"
	"sync"

	"github.com/duffy/usb-server/internal/protocol"
	"github.com/duffy/usb-server/internal/usbip"
	"github.com/google/uuid"
)

// RemoteDevice represents a USB device available from a share client
type RemoteDevice struct {
	protocol.USBDevice
	ClientID   string `json:"client_id"`
	ClientName string `json:"client_name"`
}

// AttachedDevice represents a device currently attached via VHCI
type AttachedDevice struct {
	RemoteDevice
	TunnelID string `json:"tunnel_id"`
	VHCIPort int    `json:"vhci_port"`
	SocketFD int    `json:"socket_fd"`
}

// UseManager handles receiving/using remote USB devices
type UseManager struct {
	client    *Client
	mu        sync.RWMutex
	available map[string][]RemoteDevice      // clientID -> devices
	attached  map[string]*AttachedDevice     // busID@clientID -> attached info
	tunnels   map[string]*useTunnel          // tunnelID -> tunnel
	pending   map[string]chan *protocol.DeviceGranted // requestID -> response channel
}

type useTunnel struct {
	id       string
	busID    string
	clientID string
	conn     net.Conn // our end of the socketpair
	done     chan struct{}
}

// NewUseManager creates a use manager
func NewUseManager(client *Client) *UseManager {
	um := &UseManager{
		client:    client,
		available: make(map[string][]RemoteDevice),
		attached:  make(map[string]*AttachedDevice),
		tunnels:   make(map[string]*useTunnel),
		pending:   make(map[string]chan *protocol.DeviceGranted),
	}

	client.OnDeviceList = um.handleDeviceList
	client.OnDeviceGranted = um.handleDeviceGranted
	client.OnDeviceDenied = um.handleDeviceDenied
	client.OnDeviceReleased = um.handleDeviceReleased
	client.OnTunnelData = um.handleTunnelData
	client.OnClientLeft = um.handleClientLeft

	return um
}

// GetAvailableDevices returns all available remote devices
func (um *UseManager) GetAvailableDevices() []RemoteDevice {
	um.mu.RLock()
	defer um.mu.RUnlock()

	var all []RemoteDevice
	for _, devs := range um.available {
		all = append(all, devs...)
	}
	return all
}

// GetAttachedDevices returns currently attached devices
func (um *UseManager) GetAttachedDevices() []*AttachedDevice {
	um.mu.RLock()
	defer um.mu.RUnlock()

	var result []*AttachedDevice
	for _, dev := range um.attached {
		result = append(result, dev)
	}
	return result
}

// AttachDevice requests and attaches a remote USB device
func (um *UseManager) AttachDevice(clientID, busID string) error {
	// Check if VHCI is available
	if !usbip.IsVHCIAvailable() {
		return fmt.Errorf("vhci-hcd kernel module not loaded (run: sudo modprobe vhci-hcd)")
	}

	key := busID + "@" + clientID
	um.mu.RLock()
	if _, already := um.attached[key]; already {
		um.mu.RUnlock()
		return fmt.Errorf("device %s already attached", key)
	}
	um.mu.RUnlock()

	// Create request
	requestID := uuid.New().String()
	respChan := make(chan *protocol.DeviceGranted, 1)

	um.mu.Lock()
	um.pending[requestID] = respChan
	um.mu.Unlock()

	defer func() {
		um.mu.Lock()
		delete(um.pending, requestID)
		um.mu.Unlock()
	}()

	// Send request to relay
	err := um.client.SendJSON(&protocol.RequestDevice{
		Type:         protocol.MsgRequestDevice,
		TargetClient: clientID,
		BusID:        busID,
		RequestID:    requestID,
	})
	if err != nil {
		return fmt.Errorf("sending request: %w", err)
	}

	log.Printf("[use] requesting device %s from %s", busID, clientID)

	// Wait for response (with timeout via context)
	select {
	case granted := <-respChan:
		return um.setupVHCI(clientID, busID, granted)
	case <-um.client.ctx.Done():
		return fmt.Errorf("client shutting down")
	}
}

// DetachDevice releases an attached device
func (um *UseManager) DetachDevice(clientID, busID string) error {
	key := busID + "@" + clientID

	um.mu.Lock()
	dev, exists := um.attached[key]
	if !exists {
		um.mu.Unlock()
		return fmt.Errorf("device %s not attached", key)
	}

	// Clean up tunnel
	if tunnel, ok := um.tunnels[dev.TunnelID]; ok {
		close(tunnel.done)
		if tunnel.conn != nil {
			tunnel.conn.Close()
		}
		delete(um.tunnels, dev.TunnelID)
	}

	// Detach from VHCI
	if dev.VHCIPort >= 0 {
		if err := usbip.DetachDevice(dev.VHCIPort); err != nil {
			log.Printf("[use] warning: VHCI detach error: %v", err)
		}
	}

	delete(um.attached, key)
	um.mu.Unlock()

	// Notify share client
	um.client.SendJSON(&protocol.ReleaseDevice{
		Type:         protocol.MsgReleaseDevice,
		TargetClient: clientID,
		BusID:        busID,
	})

	log.Printf("[use] device %s detached", key)
	return nil
}

func (um *UseManager) setupVHCI(clientID, busID string, granted *protocol.DeviceGranted) error {
	// Create a socketpair - one end for VHCI, one for our tunnel
	fds, err := createSocketPair()
	if err != nil {
		return fmt.Errorf("creating socketpair: %w", err)
	}

	vhciFD := fds[0]
	tunnelFD := fds[1]

	// Find a free VHCI port
	port, err := usbip.FindFreePort(granted.Speed)
	if err != nil {
		closeFDs(fds)
		return fmt.Errorf("finding free VHCI port: %w", err)
	}

	// Attach to VHCI
	if err := usbip.AttachDevice(port, vhciFD, granted.DevID, granted.Speed); err != nil {
		closeFDs(fds)
		return fmt.Errorf("VHCI attach: %w", err)
	}

	// The VHCI driver now owns vhciFD, so we don't close it
	// Create a net.Conn from the tunnel FD
	tunnelFile := fdToFile(tunnelFD, "usb-tunnel")
	tunnelConn, err := net.FileConn(tunnelFile)
	tunnelFile.Close() // FileConn dups the fd
	if err != nil {
		usbip.DetachDevice(port)
		return fmt.Errorf("creating tunnel conn: %w", err)
	}

	tunnel := &useTunnel{
		id:       granted.TunnelID,
		busID:    busID,
		clientID: clientID,
		conn:     tunnelConn,
		done:     make(chan struct{}),
	}

	key := busID + "@" + clientID
	um.mu.Lock()
	um.tunnels[granted.TunnelID] = tunnel
	um.attached[key] = &AttachedDevice{
		RemoteDevice: RemoteDevice{
			USBDevice:  protocol.USBDevice{BusID: busID},
			ClientID:   clientID,
		},
		TunnelID: granted.TunnelID,
		VHCIPort: port,
		SocketFD: vhciFD,
	}
	um.mu.Unlock()

	// Start reading from the tunnel socket (VHCI -> relay)
	go um.tunnelReadLoop(tunnel)

	log.Printf("[use] device %s attached on VHCI port %d", key, port)
	return nil
}

// tunnelReadLoop reads from the VHCI socket and sends to relay
func (um *UseManager) tunnelReadLoop(tunnel *useTunnel) {
	buf := make([]byte, 65536)
	for {
		select {
		case <-tunnel.done:
			return
		default:
		}

		n, err := tunnel.conn.Read(buf)
		if err != nil {
			select {
			case <-tunnel.done:
				return
			default:
				log.Printf("[use] tunnel read error: %v", err)
				return
			}
		}

		if err := um.client.SendTunnelData(tunnel.id, buf[:n]); err != nil {
			return
		}
	}
}

func (um *UseManager) handleDeviceList(msg *protocol.DeviceList) {
	um.mu.Lock()
	var remoteDevs []RemoteDevice
	for _, dev := range msg.Devices {
		remoteDevs = append(remoteDevs, RemoteDevice{
			USBDevice:  dev,
			ClientID:   msg.ClientID,
			ClientName: msg.ClientName,
		})
	}
	um.available[msg.ClientID] = remoteDevs
	um.mu.Unlock()

	log.Printf("[use] received device list from %s (%s): %d devices",
		msg.ClientName, msg.ClientID[:8], len(msg.Devices))
}

func (um *UseManager) handleDeviceGranted(msg *protocol.DeviceGranted) {
	um.mu.RLock()
	ch, exists := um.pending[msg.RequestID]
	um.mu.RUnlock()

	if exists {
		ch <- msg
	}
}

func (um *UseManager) handleDeviceDenied(msg *protocol.DeviceDenied) {
	log.Printf("[use] device request denied: %s - %s", msg.BusID, msg.Reason)

	um.mu.RLock()
	ch, exists := um.pending[msg.RequestID]
	um.mu.RUnlock()

	if exists {
		close(ch) // signal denial by closing channel
	}
}

func (um *UseManager) handleDeviceReleased(msg *protocol.DeviceReleased) {
	log.Printf("[use] device released by share client: %s", msg.BusID)
}

func (um *UseManager) handleTunnelData(tunnelID string, data []byte) {
	um.mu.RLock()
	tunnel, exists := um.tunnels[tunnelID]
	um.mu.RUnlock()

	if !exists {
		return
	}

	// Write to the tunnel socket (relay -> VHCI)
	tunnel.conn.Write(data)
}

func (um *UseManager) handleClientLeft(msg *protocol.ClientLeft) {
	um.mu.Lock()
	delete(um.available, msg.ClientID)

	// Detach any devices from this client
	for key, dev := range um.attached {
		if dev.ClientID == msg.ClientID {
			if tunnel, ok := um.tunnels[dev.TunnelID]; ok {
				close(tunnel.done)
				tunnel.conn.Close()
				delete(um.tunnels, dev.TunnelID)
			}
			if dev.VHCIPort >= 0 {
				usbip.DetachDevice(dev.VHCIPort)
			}
			delete(um.attached, key)
			log.Printf("[use] device %s auto-detached (client left)", key)
		}
	}
	um.mu.Unlock()
}

// Cleanup releases all attached devices
func (um *UseManager) Cleanup() {
	um.mu.Lock()
	defer um.mu.Unlock()

	for key, dev := range um.attached {
		if tunnel, ok := um.tunnels[dev.TunnelID]; ok {
			close(tunnel.done)
			tunnel.conn.Close()
		}
		if dev.VHCIPort >= 0 {
			usbip.DetachDevice(dev.VHCIPort)
		}
		log.Printf("[use] cleaned up device %s", key)
	}

	um.attached = make(map[string]*AttachedDevice)
	um.tunnels = make(map[string]*useTunnel)
}