sing-box/service/usbip/server.go

//go:build linux || (darwin && cgo)

package usbip

import (
	"context"
	"fmt"
	"net"
	"sync"
	"time"

	"github.com/sagernet/sing-box/adapter"
	boxService "github.com/sagernet/sing-box/adapter/service"
	"github.com/sagernet/sing-box/common/listener"
	C "github.com/sagernet/sing-box/constant"
	"github.com/sagernet/sing-box/log"
	"github.com/sagernet/sing-box/option"
	"github.com/sagernet/sing/common"
	E "github.com/sagernet/sing/common/exceptions"
	N "github.com/sagernet/sing/common/network"
)

// ServerService is the unified USB/IP server. It owns the wire-protocol
// surface (devlist, import, import-ext, control channel) and delegates
// all mutable state to an exportLedger. Platform-specific device
// acquisition is delegated to an ExportHost.
type ServerService struct {
	boxService.Adapter
	ctx      context.Context
	cancel   context.CancelFunc
	logger   log.ContextLogger
	listener *listener.Listener
	matches  []option.USBIPDeviceMatch
	host     ExportHost
	ledger   *exportLedger

	reconcileAccess sync.Mutex
}

// NewServerService constructs a ServerService for the running platform.
func NewServerService(ctx context.Context, logger log.ContextLogger, tag string, options option.USBIPServerServiceOptions) (adapter.Service, error) {
	for i, m := range options.Devices {
		if m.IsZero() {
			return nil, E.New("devices[", i, "]: at least one of busid/vendor_id/product_id/serial is required")
		}
	}
	if options.ListenPort == 0 {
		options.ListenPort = DefaultPort
	}
	host, err := newPlatformExportHost(logger, options.Devices)
	if err != nil {
		return nil, err
	}
	ctx, cancel := context.WithCancel(ctx)
	return &ServerService{
		Adapter: boxService.NewAdapter(C.TypeUSBIPServer, tag),
		ctx:     ctx,
		cancel:  cancel,
		logger:  logger,
		matches: options.Devices,
		host:    host,
		ledger:  newExportLedger(logger, importLeaseTTL, time.Now),
		listener: listener.New(listener.Options{
			Context: ctx,
			Logger:  logger,
			Network: []string{N.NetworkTCP},
			Listen:  options.ListenOptions,
		}),
	}, nil
}

func (s *ServerService) Start(stage adapter.StartStage) error {
	if stage != adapter.StartStateStart {
		return nil
	}
	err := s.host.Start(s.ctx)
	if err != nil {
		return err
	}
	err = s.reconcileAndBroadcast(false)
	if err != nil {
		_ = s.host.Close()
		return err
	}
	tcpListener, err := s.listener.ListenTCP()
	if err != nil {
		_ = s.host.Close()
		return err
	}
	go s.acceptLoop(tcpListener)
	go s.eventLoop()
	return nil
}

func (s *ServerService) Close() error {
	if s.cancel != nil {
		s.cancel()
	}
	for _, conn := range s.ledger.CloseAllSubscribers() {
		_ = conn.Close()
	}
	err := common.Close(common.PtrOrNil(s.listener))
	s.reconcileAccess.Lock()
	defer s.reconcileAccess.Unlock()
	_ = s.host.Close()
	s.ledger.ResetForClose()
	return err
}

func (s *ServerService) eventLoop() {
	events, err := s.host.Events(s.ctx)
	if err != nil {
		s.logger.Warn("subscribe topology events: ", err)
		return
	}
	if events == nil {
		return
	}
	for {
		select {
		case <-s.ctx.Done():
			return
		case _, ok := <-events:
			if !ok {
				return
			}
		}
		err = s.reconcileAndBroadcast(true)
		if err != nil {
			s.logger.Warn("reconcile exports: ", err)
		}
	}
}

func (s *ServerService) reconcileAndBroadcast(notify bool) error {
	s.reconcileAccess.Lock()
	defer s.reconcileAccess.Unlock()
	if s.ctx != nil && s.ctx.Err() != nil {
		return nil
	}
	snapshot, released, _, err := s.host.Reconcile(s.ctx, s.ledger.IsBusy)
	if err != nil {
		return err
	}
	s.ledger.ApplyHostSnapshot(snapshot, released)
	if notify {
		s.ledger.BroadcastIfChanged(s.ctx)
	} else {
		s.ledger.SeedBroadcastState(s.ctx)
	}
	return nil
}

func (s *ServerService) handleStandardConn(conn net.Conn, header OpHeader) {
	closeConn := true
	defer func() {
		if closeConn {
			_ = conn.Close()
		}
	}()
	switch header.Code {
	case OpReqDevList:
		s.handleDevList(conn)
	case OpReqImport:
		closeConn = !s.handleImport(conn)
	case OpReqImportExt:
		closeConn = !s.handleImportExt(conn)
	default:
		s.logger.Debug(fmt.Sprintf("unknown opcode 0x%04x", header.Code))
	}
}

func (s *ServerService) handleControlConn(conn net.Conn) {
	defer conn.Close()
	helloMessage, err := readControlMessage(conn)
	if err != nil {
		s.logger.Debug("read control hello: ", err)
		return
	}
	hello := helloMessage.Frame
	if hello.Type != controlFrameHello {
		s.logger.Debug("unexpected control frame ", hello.Type, " before hello")
		return
	}
	if hello.Version != controlProtocolVersion {
		s.logger.Debug("unsupported control version ", hello.Version)
		return
	}
	if hello.Capabilities&controlRequiredCapabilities != controlRequiredCapabilities {
		s.logger.Debug("missing control capabilities 0x", hello.Capabilities)
		return
	}
	capabilities := negotiatedControlCapabilities(hello.Capabilities)
	sub, seq := s.ledger.Subscribe(s.ctx, conn, capabilities)
	defer s.ledger.Unsubscribe(sub)
	err = writeControlAckWithCapabilities(conn, seq, capabilities)
	if err != nil {
		s.logger.Debug("write control ack: ", err)
		return
	}
	readDone := make(chan struct{})
	go s.readControlConn(sub, readDone)
	for {
		select {
		case <-s.ctx.Done():
			return
		case <-readDone:
			return
		case message := <-sub.send:
			err = writeControlMessage(conn, message.Frame, message.Payload)
			if err != nil {
				s.logger.Debug("write control frame: ", err)
				return
			}
		}
	}
}

func (s *ServerService) handleDevList(conn net.Conn) {
	entries := s.buildDevListEntries()
	err := WriteOpRepDevList(conn, entries)
	if err != nil {
		s.logger.Debug("write devlist: ", err)
	}
}

func (s *ServerService) buildDevListEntries() []DeviceEntry {
	exports := s.ledger.AvailableExports()
	if len(exports) == 0 {
		return nil
	}
	entries := make([]DeviceEntry, 0, len(exports))
	for _, export := range exports {
		snapshot := export.Snapshot(s.ctx, false)
		if snapshot.Err != nil || snapshot.State != deviceStateAvailable {
			continue
		}
		entries = append(entries, snapshot.Entry)
	}
	return entries
}

func (s *ServerService) handleImport(conn net.Conn) bool {
	busid, err := ReadOpReqImportBody(conn)
	if err != nil {
		s.logger.Debug("read import body: ", err)
		return false
	}
	return s.handleImportBusID(conn, busid, false)
}

func (s *ServerService) handleImportExt(conn net.Conn) bool {
	request, err := ReadOpReqImportExtBody(conn)
	if err != nil {
		s.logger.Debug("read import-ext body: ", err)
		return false
	}
	if !s.ledger.ConsumeLease(request) {
		s.logger.Info("import-ext rejected (invalid lease): ", request.BusID)
		_ = WriteOpRepImportExt(conn, OpStatusError, nil)
		return false
	}
	return s.handleImportBusID(conn, request.BusID, true)
}

func (s *ServerService) handleImportBusID(conn net.Conn, busid string, extended bool) bool {
	writeReply := WriteOpRepImport
	if extended {
		writeReply = WriteOpRepImportExt
	}
	export, ok, reason := s.ledger.TryReserveForImport(s.ctx, busid)
	if !ok {
		s.logger.Info("import rejected (", busid, ": ", reason, ")")
		_ = writeReply(conn, OpStatusError, nil)
		return false
	}
	info, err := export.DeviceInfo(s.ctx)
	if err != nil {
		s.ledger.ReleaseImport(s.ctx, busid, false)
		s.logger.Warn("refresh ", busid, ": ", err)
		_ = writeReply(conn, OpStatusError, nil)
		return false
	}
	session, err := export.NewServerDataSession(s.ctx, conn)
	if err != nil {
		s.ledger.ReleaseImport(s.ctx, busid, false)
		s.logger.Warn("open data session ", busid, ": ", err)
		_ = writeReply(conn, OpStatusError, nil)
		return false
	}
	s.ledger.ConfirmImport(s.ctx)
	err = writeReply(conn, OpStatusOK, &info)
	if err != nil {
		s.logger.Warn("reply import ", busid, ": ", err)
		_ = session.Close()
		<-session.Done()
		released, _ := s.host.FinishImport(s.ctx, busid)
		s.ledger.ReleaseImport(s.ctx, busid, released)
		return false
	}
	s.logger.Info("attached ", busid, " to remote ", conn.RemoteAddr())
	go s.waitImportDone(busid, session)
	return true
}

func (s *ServerService) waitImportDone(busid string, session DataSession) {
	<-session.Done()
	released, err := s.host.FinishImport(s.ctx, busid)
	if err != nil {
		s.logger.Debug("finish import ", busid, ": ", err)
	}
	s.ledger.ReleaseImport(s.ctx, busid, released)
}