usb-server/internal/client/use.go

package client

import (
	"fmt"
	"log"
	"net"
	"strings"
	"sync"
	"time"

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

// attachRequestTimeout bounds how long we wait for a share client to answer a
// device request before giving up.
const attachRequestTimeout = 30 * time.Second

// 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"`
}

// UseManager handles receiving/using remote USB devices
type UseManager struct {
	client          *Client
	cfg             *config.Config
	cfgPath         string
	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]*pendingRequest // requestID -> in-flight attach
	forceDetachable map[string]bool            // clientID -> allow_force_detach
}

// pendingRequest tracks an attach request waiting for the share client's reply.
// It carries the target so that a reply arriving after the caller gave up can
// still be undone — otherwise the share side would hold the device open for a
// user who is no longer waiting for it.
type pendingRequest struct {
	clientID string
	busID    string
	resp     chan *protocol.DeviceGranted
}

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

	// codec seals outgoing and opens incoming payloads, whichever transport
	// carries them.
	codec *tunnelCodec

	// send delivers an encoded frame; direct is non-nil when this tunnel
	// bypasses the relay.
	send   tunnelSender
	direct *directConn

	// done is closed exactly once. Several paths can tear down the same
	// tunnel — an explicit detach, a release from the share side, the peer
	// leaving, a relay disconnect — and closing it twice would panic.
	done      chan struct{}
	closeOnce sync.Once
}

// tryDirect attempts a direct connection to the granting peer, returning nil
// if none can be established.
//
// Every failure here is benign: the tunnel simply runs through the relay, the
// way it always did. Only the latency improves when this succeeds.
func (um *UseManager) tryDirect(granted *protocol.DeviceGranted, secret *crypto.TunnelSecret) *directConn {
	if um.cfg.DisableDirect || secret == nil || len(granted.Endpoints) == 0 {
		return nil
	}

	conn, addr, err := dialDirect(granted.Endpoints, granted.TunnelID, secret.PeerToken(granted.TunnelID))
	if err != nil {
		log.Printf("[use] no direct route to %s (%v), using the relay", granted.BusID, err)
		return nil
	}

	log.Printf("[use] direct connection to %s established for %s", addr, granted.BusID)
	return conn
}

// NewUseManager creates a use manager
func NewUseManager(client *Client, cfg *config.Config, cfgPath string) *UseManager {
	um := &UseManager{
		client:          client,
		cfg:             cfg,
		cfgPath:         cfgPath,
		available:       make(map[string][]RemoteDevice),
		attached:        make(map[string]*AttachedDevice),
		tunnels:         make(map[string]*useTunnel),
		pending:         make(map[string]*pendingRequest),
		forceDetachable: make(map[string]bool),
	}

	// Use-side messages only this manager handles.
	client.OnDeviceList = um.handleDeviceList
	client.OnDeviceGranted = um.handleDeviceGranted
	client.OnDeviceDenied = um.handleDeviceDenied
	client.OnDeviceReleased = um.handleDeviceReleased

	// Shared with the share manager in "both" mode, hence multicast.
	client.AddTunnelHandler(um.handleTunnelData)
	client.AddClientLeftHandler(um.handleClientLeft)
	client.AddDisconnectHandler(um.handleRelayDisconnect)

	return um
}

// handleRelayDisconnect detaches everything after the relay link drops.
// The relay discarded those tunnels, so the devices are dead: without this
// they would stay listed as attached while no traffic could reach them.
func (um *UseManager) handleRelayDisconnect() {
	um.mu.Lock()
	defer um.mu.Unlock()

	if n := len(um.attached); n > 0 {
		log.Printf("[use] detaching %d device(s) (relay connection lost)", n)
		for key, dev := range um.attached {
			um.closeAttachedLocked(key, dev)
		}
	}

	// The peers that advertised these are unreachable, and their device lists
	// are re-sent on reconnect. Keeping stale entries would show devices the
	// UI cannot actually attach.
	um.available = make(map[string][]RemoteDevice)
	um.forceDetachable = make(map[string]bool)
}

// closeAttachedLocked tears down one attached device: its tunnel, its VHCI
// port and its bookkeeping. Callers must hold um.mu.
func (um *UseManager) closeAttachedLocked(key string, dev *AttachedDevice) {
	if tunnel, ok := um.tunnels[dev.TunnelID]; ok {
		tunnel.closeOnce.Do(func() { close(tunnel.done) })
		if tunnel.conn != nil {
			tunnel.conn.Close()
		}
		if tunnel.direct != nil {
			tunnel.direct.Close()
		}
		delete(um.tunnels, dev.TunnelID)
	}

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

	delete(um.attached, key)
}

// 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 err := usbip.VHCIUnavailableError(); err != nil {
		return err
	}

	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] = &pendingRequest{clientID: clientID, busID: busID, resp: 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 a grant or denial. Without the timeout a share client that
	// never answers — because it crashed, or the relay dropped the message —
	// would leave this call blocked forever, and with it the HTTP request or
	// auto-connect goroutine that made it.
	timer := time.NewTimer(attachRequestTimeout)
	defer timer.Stop()

	select {
	case granted, ok := <-respChan:
		if !ok || granted == nil {
			return fmt.Errorf("device request denied")
		}
		return um.setupVHCI(clientID, busID, granted)
	case <-timer.C:
		return fmt.Errorf("no response from %s for device %s after %s",
			protocol.ShortID(clientID), busID, attachRequestTimeout)
	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)
	}
	um.closeAttachedLocked(key, dev)
	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 {
	// Look up device info from available list (needed for Windows management phase)
	var devInfo *RemoteDevice
	um.mu.RLock()
	for _, d := range um.available[clientID] {
		if d.BusID == busID {
			cp := d
			devInfo = &cp
			break
		}
	}
	um.mu.RUnlock()

	// The granting side tells us whether it encrypts. Both ends must agree:
	// a mismatch would turn ciphertext into garbage USB traffic.
	secret := um.client.TunnelSecret()
	if granted.Encrypted && secret == nil {
		return fmt.Errorf("%s encrypts its tunnels but this client has no tokens configured, "+
			"only a group hash — copy the three tokens over to connect", protocol.ShortID(clientID))
	}
	if !granted.Encrypted {
		if secret != nil {
			log.Printf("[use] warning: %s does not encrypt tunnel traffic for %s",
				protocol.ShortID(clientID), busID)
		}
		secret = nil
	}

	codec, err := newTunnelCodec(secret, granted.TunnelID, crypto.DirUseToShare)
	if err != nil {
		return fmt.Errorf("setting up tunnel encryption: %w", err)
	}

	// Try to reach the peer directly before falling back to the relay. This
	// is where the latency win comes from: two machines on the same network
	// otherwise send every USB transfer out to the relay and back.
	direct := um.tryDirect(granted, secret)

	// Platform-specific VHCI attachment (Linux: socketpair+sysfs, Windows: TCP proxy+usbip.exe)
	tunnelConn, vhciPort, err := createVHCIAttachment(um.client.ctx, granted, devInfo)
	if err != nil {
		if direct != nil {
			direct.Close()
		}
		return fmt.Errorf("VHCI attachment: %w", err)
	}

	tunnel := &useTunnel{
		id:       granted.TunnelID,
		busID:    busID,
		clientID: clientID,
		conn:     tunnelConn,
		done:     make(chan struct{}),
		codec:    codec,
		direct:   direct,
	}
	if direct != nil {
		tunnel.send = directSender(direct)
	} else {
		tunnel.send = relaySender(um.client, granted.TunnelID)
	}

	key := busID + "@" + clientID
	remDev := RemoteDevice{
		USBDevice: protocol.USBDevice{BusID: busID},
		ClientID:  clientID,
	}
	if devInfo != nil {
		remDev = *devInfo
	}

	um.mu.Lock()
	um.tunnels[granted.TunnelID] = tunnel
	um.attached[key] = &AttachedDevice{
		RemoteDevice: remDev,
		TunnelID:     granted.TunnelID,
		VHCIPort:     vhciPort,
	}
	um.mu.Unlock()

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

	// On a direct connection, incoming frames arrive here instead of through
	// the relay's tunnel-data callback.
	if direct != nil {
		go func() {
			receiveLoop(direct, tunnel.codec, func(payload []byte) error {
				_, err := tunnel.conn.Write(payload)
				return err
			}, tunnel.done, "use/"+busID)

			// Losing the direct connection ends the tunnel: the USB/IP stream
			// cannot be resumed on the relay mid-conversation.
			select {
			case <-tunnel.done:
			default:
				log.Printf("[use] direct connection for %s ended, detaching", key)
				um.DetachDevice(clientID, busID)
			}
		}()
	}

	transport := "relay"
	if direct != nil {
		transport = "direct " + direct.RemoteAddr()
	}
	log.Printf("[use] device %s attached on VHCI port %d (devID=0x%08x speed=%d, %s, encrypted=%v)",
		key, vhciPort, granted.DevID, granted.Speed, transport, codec.encrypted())

	// Check device status and fix permissions on newly created device nodes
	// VHCI-created devices don't get normal udev permissions
	go func() {
		logVHCIDeviceStatus(vhciPort)
		fixVHCIDevicePermissions(vhciPort)
	}()

	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 protocol.Debug {
			usbip.TraceRequest("use-tunnel", buf[:n])
		}

		if err := send(tunnel.codec, tunnel.send, buf[:n]); err != nil {
			log.Printf("[use] tunnel send error: %v", err)
			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.forceDetachable[msg.ClientID] = msg.AllowForceDetach

	// Collect devices to auto-connect (while holding the lock to check attached map)
	var toAutoConnect []RemoteDevice
	for _, dev := range remoteDevs {
		if dev.Status != protocol.StatusAvailable {
			continue
		}
		key := dev.BusID + "@" + msg.ClientID
		if _, attached := um.attached[key]; attached {
			continue
		}
		if um.matchesAutoConnect(dev) {
			toAutoConnect = append(toAutoConnect, dev)
		}
	}
	um.mu.Unlock()

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

	// Auto-connect matching devices (outside lock, each in its own goroutine)
	for _, dev := range toAutoConnect {
		log.Printf("[use] auto-connecting %s (%s:%s) from %s",
			dev.Name, dev.VendorID, dev.ProductID, msg.ClientName)
		go um.AttachDevice(msg.ClientID, dev.BusID)
	}
}

// matchesAutoConnect checks if a device matches any auto-connect rule.
// Must be called with um.mu held (at least RLock).
func (um *UseManager) matchesAutoConnect(dev RemoteDevice) bool {
	for _, rule := range um.cfg.AutoConnect {
		if rule.VendorID != "" && !strings.EqualFold(rule.VendorID, dev.VendorID) {
			continue
		}
		if rule.ProductID != "" && !strings.EqualFold(rule.ProductID, dev.ProductID) {
			continue
		}
		if rule.BusID != "" && rule.BusID != dev.BusID {
			continue
		}
		if rule.ClientName != "" && rule.ClientName != dev.ClientName {
			continue
		}
		return true // all specified fields match
	}
	return false
}

// SetAutoConnect adds or removes an auto-connect rule for a VendorID:ProductID pair.
func (um *UseManager) SetAutoConnect(vendorID, productID string, enabled bool) error {
	um.mu.Lock()
	defer um.mu.Unlock()

	if enabled {
		// Check if rule already exists
		for _, rule := range um.cfg.AutoConnect {
			if strings.EqualFold(rule.VendorID, vendorID) && strings.EqualFold(rule.ProductID, productID) {
				return nil // already exists
			}
		}
		um.cfg.AutoConnect = append(um.cfg.AutoConnect, config.AutoConnectRule{
			VendorID:  vendorID,
			ProductID: productID,
		})
	} else {
		// Remove matching rule
		filtered := um.cfg.AutoConnect[:0]
		for _, rule := range um.cfg.AutoConnect {
			if strings.EqualFold(rule.VendorID, vendorID) && strings.EqualFold(rule.ProductID, productID) {
				continue
			}
			filtered = append(filtered, rule)
		}
		um.cfg.AutoConnect = filtered
	}

	if err := um.cfg.Save(um.cfgPath); err != nil {
		return fmt.Errorf("saving config: %w", err)
	}

	log.Printf("[use] auto-connect %s:%s = %v", vendorID, productID, enabled)
	return nil
}

// IsAutoConnect checks if there is an auto-connect rule for this VendorID:ProductID.
func (um *UseManager) IsAutoConnect(vendorID, productID string) bool {
	um.mu.RLock()
	defer um.mu.RUnlock()

	for _, rule := range um.cfg.AutoConnect {
		if strings.EqualFold(rule.VendorID, vendorID) && strings.EqualFold(rule.ProductID, productID) {
			return true
		}
	}
	return false
}

// IsForceDetachable checks if a share client allows force-detach
func (um *UseManager) IsForceDetachable(clientID string) bool {
	um.mu.RLock()
	defer um.mu.RUnlock()
	return um.forceDetachable[clientID]
}

// ForceDetachDevice sends a force-release request to the share client
func (um *UseManager) ForceDetachDevice(clientID, busID string) error {
	return um.client.SendJSON(&protocol.ForceRelease{
		Type:         protocol.MsgForceRelease,
		TargetClient: clientID,
		BusID:        busID,
	})
}

// resolvePending hands a response to the waiting AttachDevice call and removes
// the request, so that a duplicate or late reply cannot reach the channel
// twice — a grant arriving after a denial closed it would panic.
func (um *UseManager) resolvePending(requestID string) (*pendingRequest, bool) {
	um.mu.Lock()
	defer um.mu.Unlock()

	req, exists := um.pending[requestID]
	if exists {
		delete(um.pending, requestID)
	}
	return req, exists
}

func (um *UseManager) handleDeviceGranted(msg *protocol.DeviceGranted) {
	req, exists := um.resolvePending(msg.RequestID)
	if !exists {
		// Nobody is waiting any more — the request timed out, or the caller
		// gave up. The share client has already claimed the device for us, so
		// hand it back instead of leaving it stuck in "in use".
		log.Printf("[use] late grant for %s, releasing it again", msg.BusID)
		um.releaseOrphanedGrant(msg)
		return
	}

	// The channel is buffered with capacity 1 and we are the only sender for
	// this request ID, so this never blocks.
	req.resp <- msg
}

// releaseOrphanedGrant tells the share client to take back a device that was
// granted to a request nobody is waiting for.
func (um *UseManager) releaseOrphanedGrant(msg *protocol.DeviceGranted) {
	// Find who owns this bus ID; the grant message does not name the sender.
	um.mu.RLock()
	var owner string
	for clientID, devs := range um.available {
		for _, d := range devs {
			if d.BusID == msg.BusID {
				owner = clientID
				break
			}
		}
		if owner != "" {
			break
		}
	}
	um.mu.RUnlock()

	if owner == "" {
		return
	}

	um.client.SendJSON(&protocol.ReleaseDevice{
		Type:         protocol.MsgReleaseDevice,
		TargetClient: owner,
		BusID:        msg.BusID,
	})
}

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

	if req, exists := um.resolvePending(msg.RequestID); exists {
		close(req.resp) // a closed channel reads as a denial
	}
}

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

	um.mu.Lock()
	// Find and clean up any attached device matching this BusID (and ClientID if provided)
	for key, dev := range um.attached {
		if dev.BusID != msg.BusID {
			continue
		}
		if msg.ClientID != "" && dev.ClientID != msg.ClientID {
			continue
		}

		um.closeAttachedLocked(key, dev)
		log.Printf("[use] device %s cleaned up (released by share client)", key)
		break
	}
	um.mu.Unlock()
}

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

	if !exists {
		// In "both" mode the share manager sees the same frames and owns the
		// other tunnels, so an unknown ID here is normal, not an error.
		return
	}

	// A tunnel running directly gets its frames from that connection; anything
	// arriving via the relay for it is stale or spoofed.
	if tunnel.direct != nil {
		return
	}

	payload, err := tunnel.codec.decode(data)
	if err != nil {
		log.Printf("[use] tunnel %s: rejecting relayed frame: %v", protocol.ShortID(tunnelID), err)
		tunnel.closeOnce.Do(func() { close(tunnel.done) })
		tunnel.conn.Close()
		return
	}

	if protocol.Debug {
		usbip.TraceResponse("use-tunnel", payload)
	}

	// Write to the tunnel socket (peer -> VHCI). A failed write would desync
	// the USB/IP stream permanently, so treat it as fatal for this tunnel.
	if _, err := tunnel.conn.Write(payload); err != nil {
		log.Printf("[use] tunnel %s write error: %v", protocol.ShortID(tunnelID), err)
		tunnel.closeOnce.Do(func() { close(tunnel.done) })
		tunnel.conn.Close()
	}
}

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

	// Detach any devices from this client
	for key, dev := range um.attached {
		if dev.ClientID == msg.ClientID {
			um.closeAttachedLocked(key, dev)
			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 {
		um.closeAttachedLocked(key, dev)
		log.Printf("[use] cleaned up device %s", key)
	}

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