usb-server/internal/client/use.go

package client

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

	"github.com/duffy/usb-server/internal/config"
	"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"`
}

// 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]chan *protocol.DeviceGranted // requestID -> response channel
	forceDetachable map[string]bool                // clientID -> allow_force_detach
}

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, 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]chan *protocol.DeviceGranted),
		forceDetachable: make(map[string]bool),
	}

	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 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] = 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, ok := <-respChan:
		if !ok || granted == nil {
			return fmt.Errorf("device request denied")
		}
		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 {
	// 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()

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

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

	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 -> relay)
	go um.tunnelReadLoop(tunnel)

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

	// Fix permissions on newly created device nodes (e.g. /dev/video*)
	// VHCI-created devices don't get normal udev permissions
	go 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 err := um.client.SendTunnelData(tunnel.id, 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, msg.ClientID[:8], 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,
	})
}

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)

	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
		}

		// 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 for force-released device: %v", err)
			}
		}

		delete(um.attached, key)
		log.Printf("[use] device %s cleaned up (force-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 {
		log.Printf("[use] tunnel data for unknown tunnel %s (%d bytes)", tunnelID[:8], len(data))
		return
	}

	// Write to the tunnel socket (relay -> VHCI)
	n, err := tunnel.conn.Write(data)
	if err != nil {
		log.Printf("[use] tunnel write error: %v", err)
	} else if n != len(data) {
		log.Printf("[use] tunnel short write: %d/%d", n, len(data))
	}
}

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 {
			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)
}