//go:build linux

// Package bridge accepts USB devices handed in from another process.
//
// It exists for hosts where this process cannot open USB devices itself.
// Android is the case that motivated it: apps there have no access to
// /dev/bus/usb, and must ask the framework, which shows a permission dialog
// and returns an already-open file descriptor. A small app-side shim obtains
// that descriptor plus the device's raw descriptors and passes both here over
// a Unix socket, using SCM_RIGHTS to transfer the descriptor itself.
//
// Nothing about this is Android-specific though: any supervising process can
// use it to hand devices to an unprivileged client.
package bridge

import (
	"encoding/json"
	"fmt"
	"log"
	"net"
	"os"
	"path/filepath"
	"sync"

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

// maxRequestSize caps one request. Descriptor blobs are a few hundred bytes;
// this leaves plenty of room while bounding what a caller can make us buffer.
const maxRequestSize = 64 * 1024

// Request is one device handover, sent as a single JSON message with the
// device's file descriptor attached as SCM_RIGHTS ancillary data.
type Request struct {
	// Action is "add" or "remove".
	Action string `json:"action"`

	// BusID identifies the device within this client, e.g. "1-2". It must be
	// stable for as long as the device is shared: it is what peers request.
	BusID string `json:"bus_id"`

	// Descriptors is the raw descriptor blob, base64 encoded by encoding/json:
	// the device descriptor followed by all configuration descriptors. On
	// Android this is UsbDeviceConnection.getRawDescriptors().
	Descriptors []byte `json:"descriptors,omitempty"`

	BusNum       uint32 `json:"bus_num,omitempty"`
	DevNum       uint32 `json:"dev_num,omitempty"`
	Speed        uint32 `json:"speed,omitempty"`
	ConfigValue  uint8  `json:"config_value,omitempty"`
	Manufacturer string `json:"manufacturer,omitempty"`
	Product      string `json:"product,omitempty"`
	Serial       string `json:"serial,omitempty"`
}

// Response reports the outcome of a request.
type Response struct {
	OK    bool   `json:"ok"`
	Error string `json:"error,omitempty"`
}

// Server listens for device handovers on a Unix socket.
type Server struct {
	listener net.Listener
	path     string

	// OnChange fires after a device is added or removed, so the share manager
	// can refresh and announce its list without waiting for the next poll.
	OnChange func()

	mu     sync.Mutex
	closed bool
}

// Listen starts a bridge server on the given Unix socket path.
//
// The socket is created with 0600 permissions: whoever can write to it can
// make this client share arbitrary USB devices.
func Listen(path string) (*Server, error) {
	if path == "" {
		return nil, fmt.Errorf("socket path is required")
	}

	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return nil, fmt.Errorf("creating socket directory: %w", err)
	}

	// A leftover socket from a previous run would make Listen fail.
	if info, err := os.Stat(path); err == nil && info.Mode()&os.ModeSocket != 0 {
		os.Remove(path)
	}

	ln, err := net.Listen("unix", path)
	if err != nil {
		return nil, fmt.Errorf("listening on %s: %w", path, err)
	}

	if err := os.Chmod(path, 0600); err != nil {
		ln.Close()
		return nil, fmt.Errorf("securing socket: %w", err)
	}

	s := &Server{listener: ln, path: path}
	go s.acceptLoop()

	log.Printf("[bridge] listening on %s for device handovers", path)
	return s, nil
}

// Close stops the server and removes the socket.
func (s *Server) Close() error {
	s.mu.Lock()
	s.closed = true
	s.mu.Unlock()

	err := s.listener.Close()
	os.Remove(s.path)
	usb.ReleaseAdoptedFDs()
	return err
}

func (s *Server) acceptLoop() {
	for {
		conn, err := s.listener.Accept()
		if err != nil {
			s.mu.Lock()
			closed := s.closed
			s.mu.Unlock()
			if closed {
				return
			}
			log.Printf("[bridge] accept error: %v", err)
			return
		}

		go s.handleConn(conn.(*net.UnixConn))
	}
}

// handleConn processes requests on one connection until it closes.
func (s *Server) handleConn(conn *net.UnixConn) {
	defer conn.Close()

	for {
		req, fd, err := readRequest(conn)
		if err != nil {
			// A clean disconnect is the normal way a session ends.
			return
		}

		resp := s.apply(req, fd)
		if err := writeResponse(conn, resp); err != nil {
			return
		}
	}
}

// apply carries out one request, taking ownership of fd.
func (s *Server) apply(req *Request, fd int) Response {
	closeFD := func() {
		if fd >= 0 {
			unix.Close(fd)
		}
	}

	switch req.Action {
	case "add":
		if req.BusID == "" {
			closeFD()
			return Response{Error: "bus_id is required"}
		}
		if fd < 0 {
			return Response{Error: "no file descriptor was attached; " +
				"send the open device descriptor as SCM_RIGHTS ancillary data"}
		}
		if len(req.Descriptors) == 0 {
			closeFD()
			return Response{Error: "descriptors are required: this process cannot read them itself"}
		}

		meta := usb.ExternalDeviceMeta{
			BusNum:       req.BusNum,
			DevNum:       req.DevNum,
			Speed:        req.Speed,
			ConfigValue:  req.ConfigValue,
			Manufacturer: req.Manufacturer,
			Product:      req.Product,
			Serial:       req.Serial,
		}
		if err := usb.RegisterExternalDevice(req.BusID, req.Descriptors, meta); err != nil {
			closeFD()
			return Response{Error: err.Error()}
		}

		// Register the descriptor only after the device parsed cleanly, so a
		// rejected request leaves nothing behind.
		if err := usb.AdoptDeviceFD(req.BusID, fd); err != nil {
			usb.UnregisterExternalDevice(req.BusID)
			closeFD()
			return Response{Error: err.Error()}
		}

		log.Printf("[bridge] device %s registered from outside (%s %s)",
			req.BusID, req.Manufacturer, req.Product)
		s.notify()
		return Response{OK: true}

	case "remove":
		closeFD()
		if req.BusID == "" {
			return Response{Error: "bus_id is required"}
		}
		usb.UnregisterExternalDevice(req.BusID)
		log.Printf("[bridge] device %s withdrawn", req.BusID)
		s.notify()
		return Response{OK: true}

	default:
		closeFD()
		return Response{Error: fmt.Sprintf("unknown action %q (expected add or remove)", req.Action)}
	}
}

func (s *Server) notify() {
	if s.OnChange != nil {
		s.OnChange()
	}
}

// readRequest reads one JSON message plus an optional attached descriptor.
// It returns fd = -1 when no descriptor was sent.
func readRequest(conn *net.UnixConn) (*Request, int, error) {
	buf := make([]byte, maxRequestSize)
	oob := make([]byte, unix.CmsgSpace(4)) // room for exactly one descriptor

	n, oobn, _, _, err := conn.ReadMsgUnix(buf, oob)
	if err != nil {
		return nil, -1, err
	}
	if n == 0 {
		return nil, -1, fmt.Errorf("empty request")
	}

	fd := extractFD(oob[:oobn])

	var req Request
	if err := json.Unmarshal(buf[:n], &req); err != nil {
		if fd >= 0 {
			unix.Close(fd)
		}
		return nil, -1, fmt.Errorf("parsing request: %w", err)
	}

	return &req, fd, nil
}

// extractFD pulls a single descriptor out of ancillary data.
// Any extra descriptors are closed rather than leaked.
func extractFD(oob []byte) int {
	if len(oob) == 0 {
		return -1
	}

	msgs, err := unix.ParseSocketControlMessage(oob)
	if err != nil {
		return -1
	}

	result := -1
	for _, msg := range msgs {
		fds, err := unix.ParseUnixRights(&msg)
		if err != nil {
			continue
		}
		for _, fd := range fds {
			if result == -1 {
				result = fd
			} else {
				unix.Close(fd)
			}
		}
	}
	return result
}

func writeResponse(conn *net.UnixConn, resp Response) error {
	data, err := json.Marshal(resp)
	if err != nil {
		return err
	}
	_, err = conn.Write(data)
	return err
}