usb-server/internal/usbip/server.go

//go:build linux

package usbip

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"
	"unsafe"

	"github.com/duffy/usb-server/internal/usb"
	"golang.org/x/sys/unix"
)

// Server handles USB/IP protocol on the share side.
// It manages a single USB device and forwards URBs between
// the USB/IP client (via tunnel) and the physical device (via usbdevfs).
type Server struct {
	device      *usb.Device
	handle      *usb.DeviceHandle
	mu          sync.Mutex
	pendingURBs map[uint32]*pendingURB // seqnum -> pending URB
	closed      bool
	epTypes     map[uint8]uint8 // endpoint number (1-15) -> usbdevfs URB type
}

type pendingURB struct {
	seqNum       uint32
	devID        uint32
	direction    uint32
	endpoint     uint32
	buffer       []byte
	urbPtr       unsafe.Pointer // pointer to submitted usbdevfs_urb
	isISO        bool
	numPackets   int32
	isoMem       []byte   // keeps ISO URB+descriptors memory alive for GC
	packetLens   []uint32 // original request lengths per ISO packet (for offset computation)
}

// NewServer creates a USB/IP server for a specific device
func NewServer(dev *usb.Device) *Server {
	return &Server{
		device:      dev,
		pendingURBs: make(map[uint32]*pendingURB),
		epTypes:     make(map[uint8]uint8),
	}
}

// Attach opens the device, disconnects the kernel driver, and claims all interfaces
func (s *Server) Attach() error {
	handle, err := usb.OpenDevice(s.device.DevPath, s.device.BusID)
	if err != nil {
		return fmt.Errorf("opening device: %w", err)
	}
	s.handle = handle

	// Build endpoint type map from device info (must be done before driver detach)
	s.buildEndpointTypeMap()

	// For each interface: disconnect kernel driver and claim it
	for _, iface := range s.device.Interfaces {
		ifnum := uint32(iface.Number)

		// Try atomic disconnect+claim first (USBDEVFS_DISCONNECT_CLAIM, Linux 3.7+)
		err := handle.DisconnectClaimInterface(ifnum)
		if err == nil {
			log.Printf("[usbip-server] interface %d: disconnect+claim OK", ifnum)
			continue
		}

		// Fallback: disconnect driver per interface, then claim
		if disconnErr := handle.DisconnectDriverForInterface(ifnum); disconnErr != nil {
			log.Printf("[usbip-server] interface %d: disconnect warning: %v", ifnum, disconnErr)
		}

		if claimErr := handle.ClaimInterface(ifnum); claimErr != nil {
			log.Printf("[usbip-server] error: could not claim interface %d: %v", ifnum, claimErr)
			// This is a critical error - clean up and fail
			for _, prev := range s.device.Interfaces {
				if uint32(prev.Number) < ifnum {
					handle.ReleaseInterface(uint32(prev.Number))
				}
			}
			handle.ConnectDriver()
			handle.Close()
			s.handle = nil
			return fmt.Errorf("claiming interface %d: %w", ifnum, claimErr)
		}
		log.Printf("[usbip-server] interface %d: fallback disconnect+claim OK", ifnum)
	}

	return nil
}

// Detach releases all interfaces, closes the device, and rebinds kernel drivers.
func (s *Server) Detach() {
	s.mu.Lock()
	s.closed = true
	s.mu.Unlock()

	if s.handle == nil {
		return
	}

	// 1. Discard all pending URBs to clean up device state
	s.mu.Lock()
	for seqNum, pending := range s.pendingURBs {
		if pending.urbPtr != nil {
			unix.Syscall(unix.SYS_IOCTL, uintptr(s.handle.Fd()),
				0x8000550B, // USBDEVFS_DISCARDURB
				uintptr(pending.urbPtr))
		}
		delete(s.pendingURBs, seqNum)
	}
	s.mu.Unlock()

	// 2. Release all claimed interfaces
	for _, iface := range s.device.Interfaces {
		if err := s.handle.ReleaseInterface(uint32(iface.Number)); err != nil {
			log.Printf("[usbip-server] release interface %d: %v", iface.Number, err)
		}
	}

	// 3. Close the device file descriptor.
	// The kernel auto-cancels remaining URBs on close.
	s.handle.Close()
	s.handle = nil

	// 4. Force kernel driver re-binding via sysfs authorized toggle.
	// After USBDEVFS_DISCONNECT_CLAIM, the kernel sets privileges_dropped=true.
	// This means closing the fd does NOT auto-rebind drivers.
	// Also USBDEVFS_RESET after ReleaseInterface doesn't rebind because
	// the kernel sets needs_binding=false on release.
	// The reliable solution: toggle authorized 0->1 which forces complete
	// re-enumeration and driver binding.
	s.rebindDrivers()
}

// rebindDrivers forces the kernel to re-bind drivers to the device
// by toggling the sysfs authorized attribute.
func (s *Server) rebindDrivers() {
	authPath := filepath.Join(s.device.SysPath, "authorized")

	// Deauthorize: kernel disconnects device, unbinds all drivers
	if err := os.WriteFile(authPath, []byte("0"), 0644); err != nil {
		log.Printf("[usbip-server] sysfs deauthorize failed: %v, trying fallback", err)
		s.rebindDriversFallback()
		return
	}

	// Brief delay for the kernel to process the deauthorization
	time.Sleep(100 * time.Millisecond)

	// Re-authorize: kernel re-enumerates device, binds drivers
	if err := os.WriteFile(authPath, []byte("1"), 0644); err != nil {
		log.Printf("[usbip-server] sysfs re-authorize failed: %v", err)
		return
	}

	log.Printf("[usbip-server] device re-authorized, kernel drivers re-bound")
}

// rebindDriversFallback tries alternative methods to rebind drivers
func (s *Server) rebindDriversFallback() {
	// Try writing bus_id to each original driver's bind file
	for _, iface := range s.device.Interfaces {
		if iface.Driver == "" {
			continue
		}
		ifaceName := fmt.Sprintf("%s:%d.%d", s.device.BusID, s.device.ConfigValue, iface.Number)
		bindPath := filepath.Join("/sys/bus/usb/drivers", iface.Driver, "bind")
		if err := os.WriteFile(bindPath, []byte(ifaceName), 0644); err != nil {
			log.Printf("[usbip-server] bind %s to %s failed: %v", ifaceName, iface.Driver, err)
		} else {
			log.Printf("[usbip-server] re-bound %s to driver %s", ifaceName, iface.Driver)
		}
	}
}

// buildEndpointTypeMap builds the endpoint number -> URB type map from device descriptors
func (s *Server) buildEndpointTypeMap() {
	for _, iface := range s.device.Interfaces {
		for _, ep := range iface.Endpoints {
			epNum := ep.Address & 0x0F
			// Map USB descriptor transfer type to usbdevfs URB type
			var urbType uint8
			switch ep.TransferType {
			case usb.TransferTypeControl:
				urbType = 2
			case usb.TransferTypeIsochronous:
				urbType = 0
			case usb.TransferTypeBulk:
				urbType = 3
			case usb.TransferTypeInterrupt:
				urbType = 1
			default:
				urbType = 3 // default bulk
			}
			s.epTypes[epNum] = urbType
			typeNames := map[uint8]string{0: "ISO", 1: "interrupt", 2: "control", 3: "bulk"}
			log.Printf("[usbip-server] endpoint %d (0x%02x): %s", epNum, ep.Address, typeNames[urbType])
		}
	}
}

// getURBType returns the usbdevfs URB type for an endpoint number
func (s *Server) getURBType(endpoint uint8) uint8 {
	if endpoint == 0 {
		return 2 // control
	}
	if t, ok := s.epTypes[endpoint]; ok {
		return t
	}
	return 3 // default: bulk
}

// BuildDeviceDescriptor creates a USB/IP device descriptor from our device info
func (s *Server) BuildDeviceDescriptor() DeviceDescriptor {
	var desc DeviceDescriptor
	SetPath(&desc.Path, s.device.SysPath)
	SetBusID(&desc.BusID, s.device.BusID)
	desc.BusNum = s.device.BusNum
	desc.DevNum = s.device.DevNum
	desc.Speed = s.device.Speed
	desc.IDVendor = s.device.VendorID
	desc.IDProduct = s.device.ProductID
	desc.BcdDevice = s.device.BcdDevice
	desc.BDeviceClass = s.device.DeviceClass
	desc.BDeviceSubClass = s.device.DeviceSubClass
	desc.BDeviceProtocol = s.device.DeviceProtocol
	desc.BConfigurationValue = s.device.ConfigValue
	desc.BNumConfigurations = s.device.NumConfigs
	desc.BNumInterfaces = uint8(len(s.device.Interfaces))
	return desc
}

// BuildInterfaceDescriptors creates USB/IP interface descriptors
func (s *Server) BuildInterfaceDescriptors() []InterfaceDescriptor {
	var descs []InterfaceDescriptor
	for _, iface := range s.device.Interfaces {
		descs = append(descs, InterfaceDescriptor{
			BInterfaceClass:    iface.Class,
			BInterfaceSubClass: iface.SubClass,
			BInterfaceProtocol: iface.Protocol,
		})
	}
	return descs
}

// HandleConnection processes USB/IP protocol on a bidirectional stream.
// It reads USB/IP requests from the reader, processes them, and writes responses to the writer.
// This is the main loop for handling a connected USB/IP client.
func (s *Server) HandleConnection(r io.Reader, w io.Writer) error {
	// Start the URB reaper goroutine
	retChan := make(chan []byte, 64)
	done := make(chan struct{})
	defer close(done)

	go s.reapLoop(retChan, done)

	// Forward completed URBs to the writer
	go func() {
		for {
			select {
			case data, ok := <-retChan:
				if !ok {
					return
				}
				if _, err := w.Write(data); err != nil {
					return
				}
			case <-done:
				return
			}
		}
	}()

	// Read and process incoming USB/IP messages
	for {
		// Read the URB header (20 bytes basic + 28 bytes specific = 48 total)
		hdr, err := ReadURBHeader(r)
		if err != nil {
			if err == io.EOF {
				return nil
			}
			return fmt.Errorf("reading URB header: %w", err)
		}

		switch hdr.Command {
		case CmdSubmit:
			if err := s.handleCmdSubmit(r, hdr, retChan); err != nil {
				return fmt.Errorf("handling CMD_SUBMIT: %w", err)
			}
		case CmdUnlink:
			if err := s.handleCmdUnlink(r, hdr, retChan); err != nil {
				return fmt.Errorf("handling CMD_UNLINK: %w", err)
			}
		default:
			return fmt.Errorf("unknown URB command: 0x%08x", hdr.Command)
		}
	}
}

func (s *Server) handleCmdSubmit(r io.Reader, hdr *URBHeader, retChan chan<- []byte) error {
	body, err := ReadCmdSubmit(r)
	if err != nil {
		return err
	}

	// Read transfer buffer for OUT direction
	var transferBuf []byte
	if hdr.Direction == DirOut && body.TransferBufferLen > 0 {
		transferBuf = make([]byte, body.TransferBufferLen)
		if _, err := io.ReadFull(r, transferBuf); err != nil {
			return fmt.Errorf("reading transfer buffer: %w", err)
		}
	}

	// Read ISO packet descriptors if present
	var isoDescs []ISOPacketDescriptor
	numPackets := int32(0)
	if body.NumberOfPackets != 0xFFFFFFFF && body.NumberOfPackets > 0 {
		numPackets = int32(body.NumberOfPackets)
		isoDescs = make([]ISOPacketDescriptor, numPackets)
		if err := binary.Read(r, binary.BigEndian, &isoDescs); err != nil {
			return fmt.Errorf("reading ISO descriptors: %w", err)
		}
	}

	endpoint := uint8(hdr.Endpoint)
	urbType := s.getURBType(endpoint)

	// Handle control transfers specially (endpoint 0)
	if endpoint == 0 && hdr.Direction == DirIn {
		buf := make([]byte, body.TransferBufferLen)
		n, err := s.handle.ControlTransfer(
			body.Setup[0], body.Setup[1],
			binary.LittleEndian.Uint16(body.Setup[2:4]),
			binary.LittleEndian.Uint16(body.Setup[4:6]),
			binary.LittleEndian.Uint16(body.Setup[6:8]),
			5000, buf,
		)
		var status int32
		if err != nil {
			status = -32 // -EPIPE
			n = 0
		}
		resp, err := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, status, buf[:n])
		if err != nil {
			return err
		}
		retChan <- resp
		return nil
	}

	if endpoint == 0 && hdr.Direction == DirOut {
		bmRequestType := body.Setup[0]
		bRequest := body.Setup[1]
		wValue := binary.LittleEndian.Uint16(body.Setup[2:4])
		wIndex := binary.LittleEndian.Uint16(body.Setup[4:6])

		var status int32

		// Intercept standard USB requests that require special usbdevfs ioctls.
		// Raw control transfers via USBDEVFS_CONTROL don't update kernel state.
		switch {
		case bmRequestType == 0x01 && bRequest == 0x0B:
			// SET_INTERFACE (Standard, Interface recipient)
			// MUST use USBDEVFS_SETINTERFACE so the kernel updates endpoint state
			// and allocates bandwidth for ISO endpoints (critical for webcams).
			if err := s.handle.SetInterface(uint32(wIndex), uint32(wValue)); err != nil {
				log.Printf("[usbip-server] SET_INTERFACE(iface=%d, alt=%d) failed: %v", wIndex, wValue, err)
				status = -32 // -EPIPE
			} else {
				log.Printf("[usbip-server] SET_INTERFACE(iface=%d, alt=%d) OK", wIndex, wValue)
			}

		case bmRequestType == 0x00 && bRequest == 0x09:
			// SET_CONFIGURATION (Standard, Device recipient)
			if err := s.handle.SetConfiguration(uint32(wValue)); err != nil {
				log.Printf("[usbip-server] SET_CONFIGURATION(%d) failed: %v", wValue, err)
				status = -32
			} else {
				log.Printf("[usbip-server] SET_CONFIGURATION(%d) OK", wValue)
			}

		case bmRequestType == 0x02 && bRequest == 0x01 && wValue == 0x0000:
			// CLEAR_FEATURE(ENDPOINT_HALT) (Standard, Endpoint recipient)
			if err := s.handle.ClearHalt(uint32(wIndex)); err != nil {
				log.Printf("[usbip-server] CLEAR_HALT(ep=0x%02x) failed: %v", wIndex, err)
				status = -32
			}

		default:
			// Generic OUT control transfer
			buf := transferBuf
			if buf == nil {
				buf = make([]byte, 0)
			}
			_, err := s.handle.ControlTransfer(
				bmRequestType, bRequest, wValue, wIndex,
				binary.LittleEndian.Uint16(body.Setup[6:8]),
				5000, buf,
			)
			if err != nil {
				status = -32 // -EPIPE
			}
		}

		resp, err := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, status, nil)
		if err != nil {
			return err
		}
		retChan <- resp
		return nil
	}

	ep := endpoint
	if hdr.Direction == DirIn {
		ep |= 0x80
	}

	// Handle isochronous transfers.
	// Trust the USB/IP NumberOfPackets field rather than our endpoint type map,
	// because the map is built at enumeration time (alternate setting 0) and
	// webcams only activate ISO endpoints after SET_INTERFACE to alt > 0.
	if numPackets > 0 {
		return s.handleISOSubmit(hdr, body, transferBuf, isoDescs, numPackets, ep, retChan)
	}

	// Bulk and interrupt transfers
	var buf []byte
	if hdr.Direction == DirIn {
		buf = make([]byte, body.TransferBufferLen)
	} else {
		buf = transferBuf
	}

	urb, err := s.handle.SubmitURB(&usb.SubmitURBParams{
		Type:        urbType,
		Endpoint:    ep,
		Flags:       0,
		Buffer:      buf,
		UserContext:  uintptr(hdr.SeqNum),
	})
	if err != nil {
		resp, _ := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, -32, nil)
		retChan <- resp
		return nil
	}

	s.mu.Lock()
	s.pendingURBs[hdr.SeqNum] = &pendingURB{
		seqNum:    hdr.SeqNum,
		devID:     hdr.DevID,
		direction: hdr.Direction,
		endpoint:  hdr.Endpoint,
		buffer:    buf,
		urbPtr:    unsafe.Pointer(urb),
	}
	s.mu.Unlock()

	return nil
}

// handleISOSubmit handles isochronous URB submission
func (s *Server) handleISOSubmit(hdr *URBHeader, body *CmdSubmitBody, transferBuf []byte,
	isoDescs []ISOPacketDescriptor, numPackets int32, ep uint8, retChan chan<- []byte) error {

	// Collect packet lengths and compute total buffer size
	packetLens := make([]uint32, numPackets)
	var totalBufLen uint32
	for i := int32(0); i < numPackets; i++ {
		packetLens[i] = isoDescs[i].Length
		totalBufLen += isoDescs[i].Length
	}

	// Prepare buffer
	var buf []byte
	if hdr.Direction == DirIn {
		buf = make([]byte, totalBufLen)
	} else {
		// For OUT: the incoming data is packed (compact), expand to sequential layout
		buf = make([]byte, totalBufLen)
		if transferBuf != nil {
			srcOff := uint32(0)
			dstOff := uint32(0)
			for i := int32(0); i < numPackets; i++ {
				pktLen := packetLens[i]
				if srcOff+pktLen <= uint32(len(transferBuf)) {
					copy(buf[dstOff:dstOff+pktLen], transferBuf[srcOff:srcOff+pktLen])
				}
				srcOff += pktLen
				dstOff += pktLen
			}
		}
	}

	// Submit ISO URB
	urb, isoMem, err := s.handle.SubmitISOURB(&usb.SubmitISOURBParams{
		Endpoint:        ep,
		Flags:           0x02, // URB_ISO_ASAP
		Buffer:          buf,
		NumberOfPackets:  numPackets,
		PacketLengths:   packetLens,
		UserContext:      uintptr(hdr.SeqNum),
	})
	if err != nil {
		// Submit failed - send error response
		resp, _ := BuildRetSubmit(hdr.SeqNum, hdr.DevID, hdr.Direction, hdr.Endpoint, -32, nil)
		retChan <- resp
		return nil
	}

	s.mu.Lock()
	s.pendingURBs[hdr.SeqNum] = &pendingURB{
		seqNum:     hdr.SeqNum,
		devID:      hdr.DevID,
		direction:  hdr.Direction,
		endpoint:   hdr.Endpoint,
		buffer:     buf,
		urbPtr:     unsafe.Pointer(urb),
		isISO:      true,
		numPackets: numPackets,
		isoMem:     isoMem,
		packetLens: packetLens,
	}
	s.mu.Unlock()

	return nil
}

func (s *Server) handleCmdUnlink(r io.Reader, hdr *URBHeader, retChan chan<- []byte) error {
	body, err := ReadCmdUnlink(r)
	if err != nil {
		return err
	}

	s.mu.Lock()
	pending, exists := s.pendingURBs[body.UnlinkSeqNum]
	if exists {
		delete(s.pendingURBs, body.UnlinkSeqNum)
	}
	s.mu.Unlock()

	var status int32
	if exists && pending.urbPtr != nil {
		// Try to discard the URB via proper ioctl
		_, _, errno := unix.Syscall(unix.SYS_IOCTL, uintptr(s.handle.Fd()),
			uintptr(0x8000550B), // USBDEVFS_DISCARDURB
			uintptr(pending.urbPtr))
		if errno == 0 {
			status = -104 // -ECONNRESET
		}
	}

	resp, err := BuildRetUnlink(hdr.SeqNum, hdr.DevID, status)
	if err != nil {
		return err
	}
	retChan <- resp

	return nil
}

// reapLoop continuously reaps completed URBs and sends responses
func (s *Server) reapLoop(retChan chan<- []byte, done <-chan struct{}) {
	for {
		select {
		case <-done:
			return
		default:
		}

		s.mu.Lock()
		if s.closed || s.handle == nil {
			s.mu.Unlock()
			return
		}
		// Save handle reference under lock to prevent nil deref race
		handle := s.handle
		s.mu.Unlock()

		urbInfo, err := handle.ReapURBInfo()
		if err != nil {
			select {
			case <-done:
				return
			default:
				continue
			}
		}

		seqNum := uint32(urbInfo.UserContext)

		s.mu.Lock()
		pending, exists := s.pendingURBs[seqNum]
		if exists {
			delete(s.pendingURBs, seqNum)
		}
		s.mu.Unlock()

		if !exists {
			continue
		}

		var resp []byte
		if pending.isISO {
			resp, err = s.buildISOResponse(urbInfo, pending)
		} else {
			var data []byte
			if pending.direction == DirIn && urbInfo.ActualLength > 0 {
				data = pending.buffer[:urbInfo.ActualLength]
			}
			resp, err = BuildRetSubmit(
				pending.seqNum,
				pending.devID,
				pending.direction,
				pending.endpoint,
				urbInfo.Status,
				data,
			)
		}
		if err != nil {
			continue
		}

		select {
		case retChan <- resp:
		case <-done:
			return
		}
	}
}

// buildISOResponse builds a RET_SUBMIT for a completed ISO URB
func (s *Server) buildISOResponse(urbInfo *usb.ReapedURBInfo, pending *pendingURB) ([]byte, error) {
	// Read ISO packet results from the URB memory
	isoResults := usb.ReadISOResults(pending.isoMem, pending.numPackets)

	// Build USB/IP ISO descriptors and pack transfer data
	var usbipDescs []ISOPacketDescriptor
	var packedData []byte
	bufOffset := uint32(0) // offset in our sequential buffer

	for i := int32(0); i < pending.numPackets; i++ {
		pktLen := pending.packetLens[i]
		actualLen := isoResults[i].ActualLength
		status := isoResults[i].Status

		usbipDescs = append(usbipDescs, ISOPacketDescriptor{
			Offset:       bufOffset,
			Length:        pktLen,
			ActualLength: actualLen,
			Status:       status,
		})

		// Pack actual data (compact, no gaps) for IN direction
		if pending.direction == DirIn && actualLen > 0 {
			end := bufOffset + actualLen
			if end > uint32(len(pending.buffer)) {
				end = uint32(len(pending.buffer))
			}
			packedData = append(packedData, pending.buffer[bufOffset:end]...)
		}

		bufOffset += pktLen
	}

	return BuildRetSubmitISO(
		pending.seqNum,
		pending.devID,
		pending.direction,
		pending.endpoint,
		urbInfo.Status,
		packedData,
		uint32(urbInfo.StartFrame),
		pending.numPackets,
		urbInfo.ErrorCount,
		usbipDescs,
	)
}

// HandleDevlistRequest handles an OP_REQ_DEVLIST for this device
func (s *Server) HandleDevlistRequest() ([]byte, error) {
	desc := s.BuildDeviceDescriptor()
	ifaceDescs := s.BuildInterfaceDescriptors()
	return BuildDevlistReply([]DeviceDescriptor{desc}, [][]InterfaceDescriptor{ifaceDescs})
}

// HandleImportRequest handles an OP_REQ_IMPORT for this device
func (s *Server) HandleImportRequest(requestedBusID string) ([]byte, error) {
	if requestedBusID != s.device.BusID {
		return BuildImportReply(1, nil) // device not found
	}
	desc := s.BuildDeviceDescriptor()
	return BuildImportReply(0, &desc)
}

// ReadManagementRequest reads and dispatches a management phase message.
// Returns the response bytes and whether we should transition to transfer phase.
func (s *Server) ReadManagementRequest(r io.Reader) (response []byte, startTransfer bool, err error) {
	hdr, err := ReadOpHeader(r)
	if err != nil {
		return nil, false, err
	}

	switch hdr.Command {
	case OpReqDevlist:
		resp, err := s.HandleDevlistRequest()
		return resp, false, err

	case OpReqImport:
		var busID [32]byte
		if _, err := io.ReadFull(r, busID[:]); err != nil {
			return nil, false, err
		}
		reqBusID := GetBusID(busID)
		resp, err := s.HandleImportRequest(reqBusID)
		if err != nil {
			return nil, false, err
		}

		// Check if import was successful (status in response)
		var checkBuf bytes.Buffer
		checkBuf.Write(resp)
		checkHdr, _ := ReadOpHeader(&checkBuf)
		if checkHdr != nil && checkHdr.Status == 0 {
			return resp, true, nil // successful import -> transfer phase
		}
		return resp, false, nil

	default:
		return nil, false, fmt.Errorf("unknown management command: 0x%04x", hdr.Command)
	}
}