usb-server/internal/client/client.go

package client

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"net/url"
	"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/google/uuid"
	"github.com/gorilla/websocket"
)

const (
	// readTimeout is how long we tolerate silence from the relay. The relay
	// pings every 20s, and gorilla answers pings automatically, so exceeding
	// this means the connection is genuinely dead — including the case where
	// a NAT or proxy dropped it without sending a TCP reset.
	readTimeout = 60 * time.Second

	// pingInterval is how often we ping the relay ourselves, so that an idle
	// tunnel keeps NAT mappings alive from both directions.
	pingInterval = 20 * time.Second

	// writeTimeout bounds a single frame write.
	writeTimeout = 20 * time.Second

	// sendQueueDepth bounds outgoing backlog before we consider the link stuck.
	sendQueueDepth = 256

	// reconnectMin/reconnectMax bound the exponential backoff between
	// reconnect attempts, so a relay outage does not turn into a hot loop.
	reconnectMin = 1 * time.Second
	reconnectMax = 30 * time.Second
)

// outMsg is one queued outgoing WebSocket frame.
type outMsg struct {
	typ  int
	data []byte
}

// Client manages the connection to the relay server
type Client struct {
	cfg      *config.Config
	clientID string

	mu   sync.Mutex
	conn *websocket.Conn
	send chan outMsg
	dead chan struct{}

	// Callbacks for messages that only one manager can own.
	// In "both" mode the share manager takes the share-side ones and the use
	// manager the use-side ones, so they never collide.
	OnDeviceList     func(msg *protocol.DeviceList)     // use side
	OnDeviceGranted  func(msg *protocol.DeviceGranted)  // use side
	OnDeviceDenied   func(msg *protocol.DeviceDenied)   // use side
	OnDeviceReleased func(msg *protocol.DeviceReleased) // use side
	OnClientJoined   func(msg *protocol.ClientJoined)
	OnRequestDevice  func(targetClient, fromClient, busID, requestID string) // share side
	OnReleaseDevice  func(busID, fromClient string)                          // share side
	OnForceRelease   func(targetClient, fromClient, busID string)            // share side

	// Multicast callbacks. Both managers care about these, so they are lists
	// rather than single fields: in "both" mode a plain field would mean the
	// second manager to register silently unhooked the first.
	tunnelHandlers     []func(tunnelID string, data []byte)
	clientLeftHandlers []func(msg *protocol.ClientLeft)
	disconnectHandlers []func()
	handlerMu          sync.RWMutex

	// OnConnect fires once a registration has been sent successfully.
	OnConnect func()

	// secret derives per-tunnel keys and peer tokens. Nil when the config
	// carries only a group hash, in which case tunnels stay unencrypted and
	// direct connections are unavailable.
	secret *crypto.TunnelSecret

	// directPort is advertised to the relay so peers learn where to reach us.
	directPort int

	ctx    context.Context
	cancel context.CancelFunc
}

// NewClient creates a new client instance
func NewClient(cfg *config.Config) *Client {
	ctx, cancel := context.WithCancel(context.Background())

	c := &Client{
		cfg:      cfg,
		clientID: uuid.New().String(),
		ctx:      ctx,
		cancel:   cancel,
	}

	if cfg.HasTokens() {
		secret, err := crypto.DeriveTunnelSecret(cfg.Token1, cfg.Token2, cfg.Token3)
		if err != nil {
			log.Printf("[client] tunnel encryption unavailable: %v", err)
		} else {
			c.secret = secret
		}
	} else {
		log.Printf("[client] no tokens configured, only a group hash: " +
			"tunnels will not be encrypted and direct connections are unavailable")
	}

	return c
}

// TunnelSecret returns the group secret, or nil if it could not be derived.
func (c *Client) TunnelSecret() *crypto.TunnelSecret { return c.secret }

// SetDirectPort records the port peers should use to reach this client
// directly. It is announced with the next registration.
func (c *Client) SetDirectPort(port int) {
	c.mu.Lock()
	c.directPort = port
	c.mu.Unlock()
}

// ID returns the client ID
func (c *Client) ID() string {
	return c.clientID
}

// Config returns the client config
func (c *Client) Config() *config.Config {
	return c.cfg
}

// Context returns the client's lifetime context.
func (c *Client) Context() context.Context {
	return c.ctx
}

// relayURL normalises the configured relay address into a WebSocket URL.
func (c *Client) relayURL() (string, error) {
	u, err := url.Parse(c.cfg.RelayAddr)
	if err != nil {
		return "", fmt.Errorf("invalid relay address: %w", err)
	}

	switch u.Scheme {
	case "ws", "wss":
		// ok
	case "http":
		u.Scheme = "ws"
	case "https":
		u.Scheme = "wss"
	default:
		u.Scheme = "ws"
	}

	if u.Path == "" || u.Path == "/" {
		u.Path = "/ws"
	}

	return u.String(), nil
}

// Connect establishes connection to the relay server
func (c *Client) Connect() error {
	target, err := c.relayURL()
	if err != nil {
		return err
	}

	log.Printf("[client] connecting to %s", target)

	dialer := websocket.Dialer{
		HandshakeTimeout: 15 * time.Second,
		ReadBufferSize:   64 * 1024,
		WriteBufferSize:  64 * 1024,
	}

	conn, _, err := dialer.DialContext(c.ctx, target, nil)
	if err != nil {
		return fmt.Errorf("connecting to relay: %w", err)
	}

	conn.SetReadLimit(maxMessageSize)
	conn.SetReadDeadline(time.Now().Add(readTimeout))
	conn.SetPongHandler(func(string) error {
		conn.SetReadDeadline(time.Now().Add(readTimeout))
		return nil
	})

	c.mu.Lock()
	directPort := c.directPort
	c.mu.Unlock()

	reg := &protocol.Register{
		Type:           protocol.MsgRegister,
		Hash:           c.cfg.Hash,
		Mode:           c.cfg.Mode,
		ClientID:       c.clientID,
		Name:           c.cfg.Name,
		DirectPort:     directPort,
		LocalEndpoints: localEndpoints(directPort),
	}
	regData, err := json.Marshal(reg)
	if err != nil {
		conn.Close()
		return fmt.Errorf("encoding registration: %w", err)
	}

	// The registration is written directly because the write pump is not
	// running yet; every later write goes through the pump.
	conn.SetWriteDeadline(time.Now().Add(writeTimeout))
	if err := conn.WriteMessage(websocket.TextMessage, regData); err != nil {
		conn.Close()
		return fmt.Errorf("sending registration: %w", err)
	}

	c.mu.Lock()
	c.conn = conn
	c.send = make(chan outMsg, sendQueueDepth)
	c.dead = make(chan struct{})
	sendCh, deadCh := c.send, c.dead
	c.mu.Unlock()

	go c.writePump(conn, sendCh, deadCh)

	log.Printf("[client] registered as %s (mode=%s, name=%s)",
		protocol.ShortID(c.clientID), c.cfg.Mode, c.cfg.Name)

	if c.OnConnect != nil {
		c.OnConnect()
	}

	return nil
}

// maxMessageSize must match the relay's limit.
const maxMessageSize = 1024 * 1024

// writePump serialises all writes to the relay socket and sends keepalives.
func (c *Client) writePump(conn *websocket.Conn, send <-chan outMsg, dead <-chan struct{}) {
	ticker := time.NewTicker(pingInterval)
	defer ticker.Stop()
	defer conn.Close() // unblocks the read loop if we give up first

	for {
		select {
		case msg := <-send:
			conn.SetWriteDeadline(time.Now().Add(writeTimeout))
			if err := conn.WriteMessage(msg.typ, msg.data); err != nil {
				log.Printf("[client] write error: %v", err)
				return
			}

		case <-ticker.C:
			conn.SetWriteDeadline(time.Now().Add(writeTimeout))
			if err := conn.WriteMessage(websocket.PingMessage, nil); err != nil {
				return
			}

		case <-dead:
			return

		case <-c.ctx.Done():
			return
		}
	}
}

// RunReadLoop reads messages from the relay and dispatches them
func (c *Client) RunReadLoop() error {
	c.mu.Lock()
	conn := c.conn
	c.mu.Unlock()

	if conn == nil {
		return fmt.Errorf("not connected")
	}

	for {
		msgType, data, err := conn.ReadMessage()
		if err != nil {
			select {
			case <-c.ctx.Done():
				return nil
			default:
			}
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
				return fmt.Errorf("read error: %w", err)
			}
			return err
		}

		conn.SetReadDeadline(time.Now().Add(readTimeout))

		switch msgType {
		case websocket.TextMessage:
			c.handleTextMessage(data)
		case websocket.BinaryMessage:
			c.handleBinaryMessage(data)
		}
	}
}

// Run connects and runs the main loop with auto-reconnect
func (c *Client) Run() error {
	backoff := reconnectMin

	for {
		select {
		case <-c.ctx.Done():
			return nil
		default:
		}

		if err := c.Connect(); err != nil {
			log.Printf("[client] connection failed: %v, retrying in %s", err, backoff)
			if !c.sleep(backoff) {
				return nil
			}
			backoff = nextBackoff(backoff)
			continue
		}

		// Connected: reset the backoff so a later blip retries promptly.
		backoff = reconnectMin

		err := c.RunReadLoop()
		if err != nil {
			log.Printf("[client] disconnected: %v", err)
		} else {
			log.Printf("[client] disconnected")
		}

		c.teardown()

		// The relay dropped every tunnel involving us; local state that
		// still references one has to go too.
		c.fireDisconnect()

		select {
		case <-c.ctx.Done():
			return nil
		default:
		}

		log.Printf("[client] reconnecting in %s", backoff)
		if !c.sleep(backoff) {
			return nil
		}
		backoff = nextBackoff(backoff)
	}
}

// nextBackoff doubles the delay up to reconnectMax.
func nextBackoff(d time.Duration) time.Duration {
	d *= 2
	if d > reconnectMax {
		return reconnectMax
	}
	return d
}

// sleep waits for d, returning false if the client is shutting down.
func (c *Client) sleep(d time.Duration) bool {
	timer := time.NewTimer(d)
	defer timer.Stop()
	select {
	case <-timer.C:
		return true
	case <-c.ctx.Done():
		return false
	}
}

// teardown closes the current connection and stops its write pump.
func (c *Client) teardown() {
	c.mu.Lock()
	if c.dead != nil {
		close(c.dead)
		c.dead = nil
	}
	if c.conn != nil {
		c.conn.Close()
		c.conn = nil
	}
	c.send = nil
	c.mu.Unlock()
}

// Close shuts down the client
func (c *Client) Close() {
	c.cancel()
	c.teardown()
}

// AddTunnelHandler registers a handler for incoming tunnel frames.
// Handlers receive every frame and must ignore tunnel IDs they do not own.
func (c *Client) AddTunnelHandler(fn func(tunnelID string, data []byte)) {
	c.handlerMu.Lock()
	defer c.handlerMu.Unlock()
	c.tunnelHandlers = append(c.tunnelHandlers, fn)
}

// AddClientLeftHandler registers a handler for peer disconnects.
func (c *Client) AddClientLeftHandler(fn func(msg *protocol.ClientLeft)) {
	c.handlerMu.Lock()
	defer c.handlerMu.Unlock()
	c.clientLeftHandlers = append(c.clientLeftHandlers, fn)
}

// AddDisconnectHandler registers a handler that runs after the relay
// connection drops and before reconnecting.
//
// The relay forgets every tunnel when a client disconnects, so anything still
// attached locally now points at a tunnel that no longer exists. Handlers use
// this to tear that state down instead of leaving devices wedged until the
// process restarts.
func (c *Client) AddDisconnectHandler(fn func()) {
	c.handlerMu.Lock()
	defer c.handlerMu.Unlock()
	c.disconnectHandlers = append(c.disconnectHandlers, fn)
}

func (c *Client) fireDisconnect() {
	c.handlerMu.RLock()
	handlers := append([]func(){}, c.disconnectHandlers...)
	c.handlerMu.RUnlock()
	for _, fn := range handlers {
		fn()
	}
}

// Connected reports whether the client currently has a live relay connection.
func (c *Client) Connected() bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.conn != nil
}

// enqueue queues an outgoing frame. It never blocks on the socket; a full
// queue means the relay link is stuck, which is reported as an error so the
// caller can tear down whatever it was trying to send.
func (c *Client) enqueue(typ int, data []byte) error {
	c.mu.Lock()
	send, dead := c.send, c.dead
	c.mu.Unlock()

	if send == nil {
		return fmt.Errorf("not connected")
	}

	select {
	case send <- outMsg{typ: typ, data: data}:
		return nil
	case <-dead:
		return fmt.Errorf("connection closed")
	case <-c.ctx.Done():
		return fmt.Errorf("client shutting down")
	default:
		return fmt.Errorf("send queue full, relay link stalled")
	}
}

// SendJSON sends a JSON message to the relay
func (c *Client) SendJSON(v interface{}) error {
	data, err := json.Marshal(v)
	if err != nil {
		return fmt.Errorf("encoding message: %w", err)
	}
	return c.enqueue(websocket.TextMessage, data)
}

// SendBinary sends a binary message to the relay
func (c *Client) SendBinary(data []byte) error {
	return c.enqueue(websocket.BinaryMessage, data)
}

// SendTunnelData sends tunnel data with the tunnel ID prefix
func (c *Client) SendTunnelData(tunnelID string, data []byte) error {
	// Tunnel header: 16 bytes tunnel ID + payload
	msg := make([]byte, protocol.TunnelHeaderSize+len(data))
	copy(msg[:protocol.TunnelHeaderSize], tunnelID)
	copy(msg[protocol.TunnelHeaderSize:], data)
	return c.SendBinary(msg)
}

func (c *Client) handleTextMessage(data []byte) {
	var env protocol.Envelope
	if err := json.Unmarshal(data, &env); err != nil {
		return
	}

	switch env.Type {
	case protocol.MsgDeviceList:
		if c.OnDeviceList != nil {
			var msg protocol.DeviceList
			if json.Unmarshal(data, &msg) == nil {
				c.OnDeviceList(&msg)
			}
		}

	case protocol.MsgRequestDevice:
		if c.OnRequestDevice != nil {
			var msg struct {
				TargetClient string `json:"target_client"`
				FromClient   string `json:"from_client"`
				BusID        string `json:"bus_id"`
				RequestID    string `json:"request_id"`
			}
			if json.Unmarshal(data, &msg) == nil {
				c.OnRequestDevice(msg.TargetClient, msg.FromClient, msg.BusID, msg.RequestID)
			}
		}

	case protocol.MsgDeviceGranted:
		if c.OnDeviceGranted != nil {
			var msg protocol.DeviceGranted
			if json.Unmarshal(data, &msg) == nil {
				c.OnDeviceGranted(&msg)
			}
		}

	case protocol.MsgDeviceDenied:
		if c.OnDeviceDenied != nil {
			var msg protocol.DeviceDenied
			if json.Unmarshal(data, &msg) == nil {
				c.OnDeviceDenied(&msg)
			}
		}

	case protocol.MsgForceRelease:
		if c.OnForceRelease != nil {
			var msg struct {
				TargetClient string `json:"target_client"`
				FromClient   string `json:"from_client"`
				BusID        string `json:"bus_id"`
			}
			if json.Unmarshal(data, &msg) == nil {
				c.OnForceRelease(msg.TargetClient, msg.FromClient, msg.BusID)
			}
		}

	case protocol.MsgReleaseDevice:
		if c.OnReleaseDevice != nil {
			var msg struct {
				BusID      string `json:"bus_id"`
				FromClient string `json:"from_client"`
			}
			if json.Unmarshal(data, &msg) == nil {
				c.OnReleaseDevice(msg.BusID, msg.FromClient)
			}
		}

	case protocol.MsgDeviceReleased:
		if c.OnDeviceReleased != nil {
			var msg protocol.DeviceReleased
			if json.Unmarshal(data, &msg) == nil {
				c.OnDeviceReleased(&msg)
			}
		}

	case protocol.MsgClientJoined:
		if c.OnClientJoined != nil {
			var msg protocol.ClientJoined
			if json.Unmarshal(data, &msg) == nil {
				c.OnClientJoined(&msg)
			}
		}

	case protocol.MsgClientLeft:
		var msg protocol.ClientLeft
		if json.Unmarshal(data, &msg) == nil {
			c.handlerMu.RLock()
			handlers := append([]func(*protocol.ClientLeft){}, c.clientLeftHandlers...)
			c.handlerMu.RUnlock()
			for _, fn := range handlers {
				fn(&msg)
			}
		}

	case protocol.MsgPong:
		// ignore pong

	case protocol.MsgError:
		var msg protocol.ErrorMsg
		if json.Unmarshal(data, &msg) == nil {
			log.Printf("[client] error from relay: %s", msg.Message)
		}
	}
}

func (c *Client) handleBinaryMessage(data []byte) {
	if len(data) < protocol.TunnelHeaderSize {
		return
	}

	tunnelID := string(data[:protocol.TunnelHeaderSize])
	payload := data[protocol.TunnelHeaderSize:]

	c.handlerMu.RLock()
	handlers := c.tunnelHandlers
	c.handlerMu.RUnlock()

	for _, fn := range handlers {
		fn(tunnelID, payload)
	}
}