usb-server/internal/client/share.go

package client

import (
	"fmt"
	"io"
	"log"
	"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/usb"
	"github.com/duffy/usb-server/internal/usbip"
	"github.com/google/uuid"
)

// ShareManager handles sharing USB devices
type ShareManager struct {
	client  *Client
	cfg     *config.Config
	mu      sync.RWMutex
	devices []usb.Device
	active  map[string]*activeShare // busID -> active share
	tunnels map[string]*shareTunnel // tunnelID -> tunnel

	// secret derives per-tunnel keys. Nil when the client is configured with
	// only a group hash, in which case tunnels stay unencrypted.
	secret *crypto.TunnelSecret

	// listener accepts direct connections from peers. Nil when direct
	// connections are disabled or could not be set up, which just means every
	// tunnel goes through the relay.
	listener *directListener
}

type activeShare struct {
	device   *usb.Device
	server   *usbip.Server
	usedBy   string // client ID using this device
	tunnelID string
}

type shareTunnel struct {
	id    string
	busID string

	// in carries peer -> USB/IP server bytes. It buffers instead of blocking
	// so that feeding it from the WebSocket read loop cannot stall the client.
	in *streamBuffer

	// out carries USB/IP server -> peer bytes. This direction stays a pipe:
	// blocking there is real backpressure onto the USB reap loop, which is
	// what we want when the network cannot keep up.
	outPipe *io.PipeReader
	done    chan struct{}

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

	// sendMu guards swapping the transport when a direct connection takes
	// over from the relay mid-tunnel.
	sendMu sync.Mutex
	send   tunnelSender
	direct *directConn
}

// setTransport switches this tunnel to a new sender, closing the old direct
// connection if there was one.
func (t *shareTunnel) setTransport(sender tunnelSender, conn *directConn) {
	t.sendMu.Lock()
	defer t.sendMu.Unlock()
	if t.direct != nil && t.direct != conn {
		t.direct.Close()
	}
	t.send = sender
	t.direct = conn
}

// deliver encodes and transmits one payload over the current transport.
func (t *shareTunnel) deliver(payload []byte) error {
	t.sendMu.Lock()
	sender := t.send
	t.sendMu.Unlock()

	if sender == nil {
		return fmt.Errorf("tunnel %s has no transport", t.id)
	}
	return send(t.codec, sender, payload)
}

// closeDirect tears down any direct connection this tunnel holds.
func (t *shareTunnel) closeDirect() {
	t.sendMu.Lock()
	defer t.sendMu.Unlock()
	if t.direct != nil {
		t.direct.Close()
		t.direct = nil
	}
}

// NewShareManager creates a share manager
func NewShareManager(client *Client, cfg *config.Config) *ShareManager {
	sm := &ShareManager{
		client:  client,
		cfg:     cfg,
		active:  make(map[string]*activeShare),
		tunnels: make(map[string]*shareTunnel),
		secret:  client.TunnelSecret(),
	}

	sm.startDirectListener()

	// Share-side messages only this manager handles.
	client.OnRequestDevice = sm.handleRequestDevice
	client.OnReleaseDevice = sm.handleReleaseDevice
	client.OnForceRelease = sm.handleForceRelease

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

	return sm
}

// startDirectListener opens the port peers connect to for direct tunnels.
//
// Failure is never fatal: without a listener every tunnel simply goes through
// the relay, which is exactly how the system worked before.
func (sm *ShareManager) startDirectListener() {
	if sm.cfg.DisableDirect {
		log.Printf("[share] direct connections disabled by configuration")
		return
	}
	if sm.secret == nil {
		log.Printf("[share] direct connections unavailable: no tokens configured, only a group hash")
		return
	}

	listener, err := newDirectListener(sm.cfg.DirectPort, sm.secret)
	if err != nil {
		log.Printf("[share] direct connections unavailable: %v (falling back to relay)", err)
		return
	}

	sm.listener = listener
	sm.client.SetDirectPort(listener.Port())
}

// DirectPort reports the port peers can reach for direct tunnels, 0 if none.
func (sm *ShareManager) DirectPort() int {
	if sm.listener == nil {
		return 0
	}
	return sm.listener.Port()
}

// awaitDirect waits for the peer to connect directly and, when it does, moves
// the tunnel off the relay.
//
// The switch is safe at any moment because USB/IP is a stream of complete
// messages and each tunnel frame carries one chunk of it: frames sent before
// the switch travel via the relay, frames after it travel directly, and both
// arrive in order at the same reader. Nothing is in flight in pieces.
func (sm *ShareManager) awaitDirect(tunnel *shareTunnel, accepted <-chan *directConn) {
	select {
	case conn := <-accepted:
		if conn == nil {
			return
		}

		select {
		case <-tunnel.done:
			conn.Close()
			return
		default:
		}

		log.Printf("[share] tunnel %s now direct with %s, bypassing the relay",
			tunnel.id, conn.RemoteAddr())
		tunnel.setTransport(directSender(conn), conn)

		// Incoming frames now arrive on this connection instead of the relay.
		receiveLoop(conn, tunnel.codec, func(payload []byte) error {
			_, err := tunnel.in.Write(payload)
			return err
		}, tunnel.done, "share/"+tunnel.busID)

		// The direct connection ended. The USB/IP stream cannot resume on the
		// relay mid-conversation — the peer's VHCI has torn down its side —
		// so release the device and let it be requested again.
		select {
		case <-tunnel.done:
		default:
			log.Printf("[share] direct connection for %s ended, releasing device", tunnel.busID)
			go sm.handleReleaseDevice(tunnel.busID, "")
		}

	case <-tunnel.done:
	}
}

// handleRelayDisconnect releases every active share after the relay link
// drops. The relay discarded those tunnels, so the remote side is gone and
// the local device would otherwise stay claimed and unusable.
func (sm *ShareManager) handleRelayDisconnect() {
	sm.mu.RLock()
	busIDs := make([]string, 0, len(sm.active))
	for busID := range sm.active {
		busIDs = append(busIDs, busID)
	}
	sm.mu.RUnlock()

	for _, busID := range busIDs {
		log.Printf("[share] releasing %s (relay connection lost)", busID)
		sm.handleReleaseDevice(busID, "")
	}
}

// Run starts the share manager: periodic device enumeration + event handling
func (sm *ShareManager) Run() error {
	// Initial enumeration
	sm.refreshDevices()
	sm.broadcastDeviceList()

	// Periodic refresh
	ticker := time.NewTicker(5 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-ticker.C:
			sm.refreshDevices()
			sm.broadcastDeviceList()
		case <-sm.client.ctx.Done():
			sm.cleanup()
			return nil
		}
	}
}

// GetDevices returns the current device list
func (sm *ShareManager) GetDevices() []usb.Device {
	sm.mu.RLock()
	defer sm.mu.RUnlock()
	result := make([]usb.Device, len(sm.devices))
	copy(result, sm.devices)
	return result
}

// RefreshNow re-enumerates and announces immediately, rather than waiting for
// the next poll. Used when devices appear through the bridge, where the
// change is known the instant it happens.
func (sm *ShareManager) RefreshNow() {
	sm.refreshDevices()
	sm.broadcastDeviceList()
}

func (sm *ShareManager) refreshDevices() {
	devices, err := usb.Enumerate()
	if err != nil {
		log.Printf("[share] USB enumeration error: %v", err)
		return
	}

	sm.mu.Lock()
	sm.devices = devices
	sm.mu.Unlock()
}

func (sm *ShareManager) broadcastDeviceList() {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	var protoDevices []protocol.USBDevice
	for _, dev := range sm.devices {
		status := protocol.StatusAvailable
		usedBy := ""
		if share, ok := sm.active[dev.BusID]; ok {
			status = protocol.StatusInUse
			usedBy = share.usedBy
		}

		protoDevices = append(protoDevices, protocol.USBDevice{
			BusID:         dev.BusID,
			BusNum:        dev.BusNum,
			DevNum:        dev.DevNum,
			Speed:         dev.Speed,
			VendorID:      fmt.Sprintf("%04x", dev.VendorID),
			ProductID:     fmt.Sprintf("%04x", dev.ProductID),
			Class:         dev.DeviceClass,
			SubClass:      dev.DeviceSubClass,
			Protocol:      dev.DeviceProtocol,
			Name:          dev.DisplayName(),
			Manufacturer:  dev.Manufacturer,
			NumInterfaces: uint8(len(dev.Interfaces)),
			Status:        status,
			UsedBy:        usedBy,
		})
	}

	msg := &protocol.DeviceList{
		Type:             protocol.MsgDeviceList,
		ClientID:         sm.client.ID(),
		ClientName:       sm.client.Config().Name,
		Devices:          protoDevices,
		AllowForceDetach: sm.cfg.AllowForceDetach,
	}

	sm.client.SendJSON(msg)
}

func (sm *ShareManager) handleRequestDevice(targetClient, fromClient, busID, requestID string) {
	log.Printf("[share] device request: busID=%s from=%s", busID, fromClient)

	sm.mu.Lock()

	// Check if device exists
	var dev *usb.Device
	for i := range sm.devices {
		if sm.devices[i].BusID == busID {
			dev = &sm.devices[i]
			break
		}
	}

	if dev == nil {
		sm.mu.Unlock()
		sm.client.SendJSON(map[string]interface{}{
			"type":          protocol.MsgDeviceDenied,
			"bus_id":        busID,
			"request_id":    requestID,
			"reason":        "device not found",
			"target_client": fromClient,
		})
		return
	}

	// Check if already in use
	if _, inUse := sm.active[busID]; inUse {
		sm.mu.Unlock()
		sm.client.SendJSON(map[string]interface{}{
			"type":          protocol.MsgDeviceDenied,
			"bus_id":        busID,
			"request_id":    requestID,
			"reason":        "device already in use",
			"target_client": fromClient,
		})
		return
	}

	// Create USB/IP server for this device
	server := usbip.NewServer(dev)
	if err := server.Attach(); err != nil {
		sm.mu.Unlock()
		log.Printf("[share] failed to attach device %s: %v", busID, err)
		sm.client.SendJSON(map[string]interface{}{
			"type":          protocol.MsgDeviceDenied,
			"bus_id":        busID,
			"request_id":    requestID,
			"reason":        fmt.Sprintf("attach failed: %v", err),
			"target_client": fromClient,
		})
		return
	}

	tunnelID := uuid.New().String()[:16] // 16 chars for tunnel header
	for len(tunnelID) < 16 {
		tunnelID += "0"
	}

	codec, err := newTunnelCodec(sm.secret, tunnelID, crypto.DirShareToUse)
	if err != nil {
		sm.mu.Unlock()
		server.Detach()
		log.Printf("[share] failed to set up tunnel encryption for %s: %v", busID, err)
		sm.client.SendJSON(map[string]interface{}{
			"type":          protocol.MsgDeviceDenied,
			"bus_id":        busID,
			"request_id":    requestID,
			"reason":        fmt.Sprintf("tunnel setup failed: %v", err),
			"target_client": fromClient,
		})
		return
	}

	in := newStreamBuffer()
	outReader, outWriter := io.Pipe()

	tunnel := &shareTunnel{
		id:      tunnelID,
		busID:   busID,
		in:      in,
		outPipe: outReader,
		done:    make(chan struct{}),
		codec:   codec,
	}
	// Start on the relay. If the peer reaches us directly, the transport is
	// swapped underneath without the USB/IP layer noticing.
	tunnel.setTransport(relaySender(sm.client, tunnelID), nil)

	// Register the tunnel before announcing it, so a peer that connects
	// immediately after receiving the grant is not rejected as unknown.
	var accepted <-chan *directConn
	if sm.listener != nil {
		var err error
		if accepted, err = sm.listener.Expect(tunnelID); err != nil {
			log.Printf("[share] cannot expect a direct connection for %s: %v", busID, err)
		}
	}

	share := &activeShare{
		device:   dev,
		server:   server,
		usedBy:   fromClient,
		tunnelID: tunnelID,
	}

	sm.active[busID] = share
	sm.tunnels[tunnelID] = tunnel
	sm.mu.Unlock()

	// Start USB/IP protocol handler in background.
	// The tunnel carries the USB/IP transfer phase directly: on Linux the
	// use side hands the socket straight to VHCI, and on Windows usbip.exe's
	// management phase is answered locally, so there is no import request here.
	go func() {
		defer func() {
			close(tunnel.done)
			in.Close()
			outWriter.Close()
			outReader.Close()
		}()

		err := server.HandleConnection(in, outWriter)
		if err != nil {
			log.Printf("[share] USB/IP connection error for %s: %v", busID, err)
		}
	}()

	// Forward outgoing data from the USB/IP server over whichever transport
	// the tunnel currently uses.
	go func() {
		buf := make([]byte, 65536)
		for {
			n, err := outReader.Read(buf)
			if err != nil {
				return
			}
			if err := tunnel.deliver(buf[:n]); err != nil {
				log.Printf("[share] tunnel %s send failed: %v", tunnelID, err)
				return
			}
		}
	}()

	// Wait in the background for the peer to connect directly.
	if accepted != nil {
		go sm.awaitDirect(tunnel, accepted)
	}

	// Send grant message, including where we can be reached directly. The
	// relay adds the public address it sees before passing this on.
	sm.client.SendJSON(map[string]interface{}{
		"type":          protocol.MsgDeviceGranted,
		"bus_id":        busID,
		"tunnel_id":     tunnelID,
		"request_id":    requestID,
		"dev_id":        dev.DevID(),
		"speed":         dev.Speed,
		"target_client": fromClient,
		"endpoints":     localEndpoints(sm.DirectPort()),
		"encrypted":     codec.encrypted(),
	})

	log.Printf("[share] device %s granted to %s (tunnel=%s, encrypted=%v)",
		busID, fromClient, tunnelID, codec.encrypted())

	// Broadcast updated device list
	sm.refreshDevices()
	sm.broadcastDeviceList()
}

func (sm *ShareManager) handleReleaseDevice(busID, fromClient string) {
	log.Printf("[share] device release: busID=%s from=%s", busID, fromClient)

	sm.mu.Lock()
	share, exists := sm.active[busID]
	if !exists {
		sm.mu.Unlock()
		return
	}

	// Close the tunnel input to signal HandleConnection to stop reading
	var tunnelDone <-chan struct{}
	if tunnel, ok := sm.tunnels[share.tunnelID]; ok {
		tunnel.in.Close()
		tunnel.closeDirect()
		tunnelDone = tunnel.done
		delete(sm.tunnels, share.tunnelID)
	}
	if sm.listener != nil {
		sm.listener.Forget(share.tunnelID)
	}

	server := share.server
	delete(sm.active, busID)
	sm.mu.Unlock()

	// Wait for HandleConnection goroutine to finish before detaching.
	// This ensures no more URBs are being submitted when we detach.
	if tunnelDone != nil {
		<-tunnelDone
		log.Printf("[share] HandleConnection goroutine finished for %s", busID)
	}

	// Now safe to detach - no more USB/IP protocol processing
	server.Detach()

	// Notify client
	sm.client.SendJSON(&protocol.DeviceReleased{
		Type:     protocol.MsgDeviceReleased,
		BusID:    busID,
		ClientID: sm.client.ID(),
	})

	log.Printf("[share] device %s released", busID)

	// Refresh device list
	sm.refreshDevices()
	sm.broadcastDeviceList()
}

func (sm *ShareManager) handleForceRelease(targetClient, fromClient, busID string) {
	if !sm.cfg.AllowForceDetach {
		log.Printf("[share] force-release denied for %s (not allowed by config)", busID)
		return
	}

	sm.mu.RLock()
	share, exists := sm.active[busID]
	sm.mu.RUnlock()

	if !exists {
		return
	}

	log.Printf("[share] force-releasing %s (requested by %s, was used by %s)", busID, protocol.ShortID(fromClient), protocol.ShortID(share.usedBy))
	sm.handleReleaseDevice(busID, share.usedBy)
}

func (sm *ShareManager) handleClientLeft(msg *protocol.ClientLeft) {
	sm.mu.RLock()
	var toRelease []string
	for busID, share := range sm.active {
		if share.usedBy == msg.ClientID {
			toRelease = append(toRelease, busID)
		}
	}
	sm.mu.RUnlock()

	for _, busID := range toRelease {
		log.Printf("[share] auto-releasing %s (client %s left)", busID, protocol.ShortID(msg.ClientID))
		sm.handleReleaseDevice(busID, msg.ClientID)
	}
}

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

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

	payload, err := tunnel.codec.decode(data)
	if err != nil {
		log.Printf("[share] tunnel %s: rejecting relayed frame: %v", tunnelID, err)
		go sm.handleReleaseDevice(tunnel.busID, "")
		return
	}

	// Buffered, non-blocking: this runs on the WebSocket read loop, which
	// must never stall on the USB side.
	if _, err := tunnel.in.Write(payload); err != nil {
		log.Printf("[share] tunnel %s input failed: %v — releasing %s", tunnelID, err, tunnel.busID)
		go sm.handleReleaseDevice(tunnel.busID, "")
	}
}

func (sm *ShareManager) cleanup() {
	sm.mu.Lock()
	defer sm.mu.Unlock()

	for busID, share := range sm.active {
		if tunnel, ok := sm.tunnels[share.tunnelID]; ok {
			tunnel.in.Close()
			tunnel.closeDirect()
		}
		share.server.Detach()
		log.Printf("[share] cleaned up device %s", busID)
	}

	sm.active = make(map[string]*activeShare)
	sm.tunnels = make(map[string]*shareTunnel)

	if sm.listener != nil {
		sm.listener.Close()
		sm.listener = nil
	}
}

// DeviceListForAPI returns device info formatted for the web API
func (sm *ShareManager) DeviceListForAPI() []map[string]interface{} {
	sm.mu.RLock()
	defer sm.mu.RUnlock()

	var result []map[string]interface{}
	for _, dev := range sm.devices {
		status := "available"
		usedBy := ""
		if share, ok := sm.active[dev.BusID]; ok {
			status = "in_use"
			usedBy = share.usedBy
		}

		result = append(result, map[string]interface{}{
			"bus_id":     dev.BusID,
			"vendor_id":  fmt.Sprintf("%04x", dev.VendorID),
			"product_id": fmt.Sprintf("%04x", dev.ProductID),
			"name":       dev.DisplayName(),
			"status":     status,
			"used_by":    usedBy,
			"speed":      dev.Speed,
		})
	}
	return result
}