Fix HID transfers, harden the tunnel, add E2E crypto and direct peers

The HID failure came down to the endpoint type map being indexed by endpoint number without the direction bit. A composite device can have endpoint 1 as both interrupt IN (0x81) and bulk OUT (0x01); the last one read won, so interrupt URBs were submitted as bulk and the kernel rejected them. The device attached and stayed silent. Endpoint data now comes from the raw descriptors read from /dev/bus/usb rather than sysfs, which only ever exposes the active alternate setting — a webcam's isochronous endpoints are invisible there because they only exist after SET_INTERFACE. Two sysfs parsing bugs fell out of that too: the numeric endpoint attributes are hex without a prefix (wMaxPacketSize "0040" was read as 40, not 64), and bInterval was never read at all. Reliability: three places could freeze the whole process. The share path fed io.Pipe from the WebSocket read loop, so one slow USB transfer stalled every tunnel and the keepalives with them. The relay wrote to client sockets while holding the hub lock, so one peer that stopped reading blocked routing and registration for everyone. Control transfers ran inline in the protocol loop behind a 5s timeout. Also fixed: a use-after- free where a discarded URB's memory could be collected while the kernel still owned it, a reap loop that spun at 100% CPU on ioctl errors, a missing attach timeout, a double close(done) panic, and Hash[:8] in the relay's log line, which let a client with a short hash take the server down. Adds mode "both", so one client can offer and consume devices at once. The tunnel and client-left callbacks became multicast for it: as plain fields the second manager to register silently unhooked the first. Tunnel traffic is now AES-256-GCM end to end, on the relay path as well as directly. The key is derived from the three tokens, not from the group hash — the relay is told the hash, so a key derived from it would protect nothing from the one party in the middle. Group IDs are unchanged, so existing setups keep working; only clients configured without the tokens drop to unencrypted, relay-only operation. Peers now try to connect directly, with the relay supplying the public address neither side can determine for itself. Candidates are raced because an unreachable address hangs until timeout rather than refusing. Falling back to the relay is not an error. Platform reach: cross-compiled targets for ARM, MIPS and RISC-V (the Linux client needed no code changes — usbdevfs is not architecture specific), multi-arch Docker images, an Android bridge that accepts devices over SCM_RIGHTS because apps cannot open /dev/bus/usb, and macOS builds via system_profiler enumeration. Adds a Windows KMDF filter driver under driver/windows with its Go side. UNTESTED: it has never been compiled or run, needs the WDK to build and an EV certificate to distribute. Treat it as a starting point. Adds "usb-client diag": says per machine whether sharing and using are possible, what stands in the way, and what fixes it. Reports can be uploaded to a relay to get them off machines that are awkward to copy from. 96 tests, all green under -race. Builds for linux, windows and darwin on amd64 and arm64. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
2026-08-11 22:02:04 +02:00
parent 54178dce75
commit 9ed473a965
95 changed files with 12181 additions and 892 deletions
@@ -0,0 +1,153 @@
+package relay
+
+import (
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+)
+
+// Diagnostics drop-off.
+//
+// Getting a report off an awkward machine — a headless NAS, a Windows box in
+// the middle of driver debugging — is otherwise a matter of copying thousands
+// of lines by hand. The relay is already reachable from every client, so it
+// makes a convenient place to leave one.
+//
+// Reports are held in memory only, capped in size and count, and expire. The
+// relay is not a storage service, and treating it like one is how it would
+// become one.
+const (
+	// maxDiagReports bounds how many are kept; the oldest is dropped first.
+	maxDiagReports = 32
+
+	// maxDiagSize bounds one report.
+	maxDiagSize = 4 << 20 // 4 MB
+
+	// diagTTL is how long a report survives. Long enough to fetch and read,
+	// short enough that machine details do not linger.
+	diagTTL = 24 * time.Hour
+)
+
+// RetentionNote describes the retention policy for the client to print.
+const RetentionNote = "24 hours"
+
+type diagReport struct {
+	data     []byte
+	stored   time.Time
+	fetched  int
+	remoteIP string
+}
+
+type diagStore struct {
+	mu      sync.Mutex
+	reports map[string]*diagReport
+}
+
+func newDiagStore() *diagStore {
+	return &diagStore{reports: make(map[string]*diagReport)}
+}
+
+// put stores a report, evicting the oldest if the store is full.
+func (s *diagStore) put(id string, data []byte, remoteIP string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.expireLocked()
+
+	if len(s.reports) >= maxDiagReports {
+		var oldestID string
+		var oldest time.Time
+		for id, report := range s.reports {
+			if oldestID == "" || report.stored.Before(oldest) {
+				oldestID, oldest = id, report.stored
+			}
+		}
+		delete(s.reports, oldestID)
+	}
+
+	s.reports[id] = &diagReport{
+		data:     data,
+		stored:   time.Now(),
+		remoteIP: remoteIP,
+	}
+}
+
+func (s *diagStore) get(id string) ([]byte, bool) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.expireLocked()
+
+	report, ok := s.reports[id]
+	if !ok {
+		return nil, false
+	}
+	report.fetched++
+	return report.data, true
+}
+
+// expireLocked drops reports past their TTL. Callers must hold the lock.
+func (s *diagStore) expireLocked() {
+	cutoff := time.Now().Add(-diagTTL)
+	for id, report := range s.reports {
+		if report.stored.Before(cutoff) {
+			delete(s.reports, id)
+		}
+	}
+}
+
+// handleDiag serves the diagnostics endpoint: PUT to store, GET to retrieve.
+func (s *Server) handleDiag(w http.ResponseWriter, r *http.Request) {
+	id := strings.TrimPrefix(r.URL.Path, "/diag/")
+	if id == "" || strings.Contains(id, "/") {
+		http.Error(w, "report ID required: /diag/<id>", http.StatusBadRequest)
+		return
+	}
+
+	switch r.Method {
+	case http.MethodPut, http.MethodPost:
+		s.storeDiag(w, r, id)
+	case http.MethodGet:
+		s.fetchDiag(w, id)
+	default:
+		http.Error(w, "use PUT to store and GET to retrieve", http.StatusMethodNotAllowed)
+	}
+}
+
+func (s *Server) storeDiag(w http.ResponseWriter, r *http.Request, id string) {
+	// LimitReader rather than trusting Content-Length: a client can lie about
+	// that, and this endpoint takes uploads from anyone who can reach it.
+	data, err := io.ReadAll(io.LimitReader(r.Body, maxDiagSize+1))
+	if err != nil {
+		http.Error(w, "could not read the report", http.StatusBadRequest)
+		return
+	}
+	if len(data) > maxDiagSize {
+		http.Error(w, fmt.Sprintf("report exceeds the %d byte limit", maxDiagSize),
+			http.StatusRequestEntityTooLarge)
+		return
+	}
+	if len(data) == 0 {
+		http.Error(w, "empty report", http.StatusBadRequest)
+		return
+	}
+
+	s.diag.put(id, data, clientIP(r))
+
+	w.WriteHeader(http.StatusCreated)
+	fmt.Fprintf(w, "stored as %s, kept for %s\n", id, RetentionNote)
+}
+
+func (s *Server) fetchDiag(w http.ResponseWriter, id string) {
+	data, ok := s.diag.get(id)
+	if !ok {
+		http.Error(w, "no such report (wrong ID, or it expired)", http.StatusNotFound)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+	w.Write(data)
+}
@@ -0,0 +1,233 @@
+package relay
+
+import (
+	"bytes"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+)
+
+func newDiagServer() *Server {
+	return &Server{hub: NewHub(), diag: newDiagStore()}
+}
+
+func TestDiagStoreAndFetch(t *testing.T) {
+	s := newDiagServer()
+	body := []byte(`{"tool":{"os":"windows"}}`)
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodPut, "/diag/report-1", bytes.NewReader(body))
+	s.handleDiag(rec, req)
+
+	if rec.Code != http.StatusCreated {
+		t.Fatalf("store returned %d, want %d", rec.Code, http.StatusCreated)
+	}
+
+	rec = httptest.NewRecorder()
+	req = httptest.NewRequest(http.MethodGet, "/diag/report-1", nil)
+	s.handleDiag(rec, req)
+
+	if rec.Code != http.StatusOK {
+		t.Fatalf("fetch returned %d, want 200", rec.Code)
+	}
+	if !bytes.Equal(rec.Body.Bytes(), body) {
+		t.Errorf("fetched %q, want %q", rec.Body.String(), body)
+	}
+	if ct := rec.Header().Get("Content-Type"); ct != "application/json" {
+		t.Errorf("content type %q, want application/json", ct)
+	}
+}
+
+func TestDiagMissingReportIs404(t *testing.T) {
+	s := newDiagServer()
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodGet, "/diag/nope", nil)
+	s.handleDiag(rec, req)
+
+	if rec.Code != http.StatusNotFound {
+		t.Errorf("got %d, want 404", rec.Code)
+	}
+}
+
+func TestDiagRejectsBadPaths(t *testing.T) {
+	s := newDiagServer()
+
+	for _, path := range []string{"/diag/", "/diag/a/b"} {
+		rec := httptest.NewRecorder()
+		req := httptest.NewRequest(http.MethodGet, path, nil)
+		s.handleDiag(rec, req)
+
+		if rec.Code != http.StatusBadRequest {
+			t.Errorf("%s returned %d, want 400", path, rec.Code)
+		}
+	}
+}
+
+// The endpoint takes uploads from anyone who can reach the relay, so it must
+// bound what one caller can make it hold.
+func TestDiagRejectsOversizedReport(t *testing.T) {
+	s := newDiagServer()
+
+	huge := bytes.Repeat([]byte("x"), maxDiagSize+100)
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodPut, "/diag/big", bytes.NewReader(huge))
+	s.handleDiag(rec, req)
+
+	if rec.Code != http.StatusRequestEntityTooLarge {
+		t.Errorf("got %d, want 413", rec.Code)
+	}
+
+	// And it must not have been stored anyway.
+	if _, ok := s.diag.get("big"); ok {
+		t.Error("an oversized report was stored")
+	}
+}
+
+// A lying Content-Length must not get past the limit either.
+func TestDiagLimitIgnoresContentLength(t *testing.T) {
+	s := newDiagServer()
+
+	huge := bytes.Repeat([]byte("x"), maxDiagSize+100)
+	req := httptest.NewRequest(http.MethodPut, "/diag/liar", bytes.NewReader(huge))
+	req.ContentLength = 10 // claims to be small
+
+	rec := httptest.NewRecorder()
+	s.handleDiag(rec, req)
+
+	if rec.Code == http.StatusCreated {
+		t.Error("an oversized body was accepted because it claimed to be small")
+	}
+}
+
+func TestDiagRejectsEmptyReport(t *testing.T) {
+	s := newDiagServer()
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodPut, "/diag/empty", bytes.NewReader(nil))
+	s.handleDiag(rec, req)
+
+	if rec.Code != http.StatusBadRequest {
+		t.Errorf("got %d, want 400", rec.Code)
+	}
+}
+
+func TestDiagRejectsOtherMethods(t *testing.T) {
+	s := newDiagServer()
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodDelete, "/diag/x", nil)
+	s.handleDiag(rec, req)
+
+	if rec.Code != http.StatusMethodNotAllowed {
+		t.Errorf("got %d, want 405", rec.Code)
+	}
+}
+
+// The store is bounded, so a stream of uploads cannot grow it without limit.
+func TestDiagEvictsOldestWhenFull(t *testing.T) {
+	store := newDiagStore()
+
+	for i := 0; i < maxDiagReports+5; i++ {
+		store.put(fmt.Sprintf("report-%d", i), []byte("{}"), "127.0.0.1")
+		// Ordering by timestamp needs the timestamps to differ.
+		time.Sleep(time.Millisecond)
+	}
+
+	store.mu.Lock()
+	count := len(store.reports)
+	store.mu.Unlock()
+
+	if count > maxDiagReports {
+		t.Errorf("store holds %d reports, limit is %d", count, maxDiagReports)
+	}
+
+	// The newest must have survived; the very first must not have.
+	if _, ok := store.get(fmt.Sprintf("report-%d", maxDiagReports+4)); !ok {
+		t.Error("the most recent report was evicted")
+	}
+	if _, ok := store.get("report-0"); ok {
+		t.Error("the oldest report survived eviction")
+	}
+}
+
+func TestDiagExpiresOldReports(t *testing.T) {
+	store := newDiagStore()
+
+	store.put("old", []byte("{}"), "127.0.0.1")
+
+	// Backdate it past the TTL.
+	store.mu.Lock()
+	store.reports["old"].stored = time.Now().Add(-diagTTL - time.Minute)
+	store.mu.Unlock()
+
+	if _, ok := store.get("old"); ok {
+		t.Error("a report past its TTL was still served")
+	}
+}
+
+func TestClientIPPrefersForwardedHeader(t *testing.T) {
+	tests := []struct {
+		name  string
+		setup func(*http.Request)
+		want  string
+	}{
+		{
+			name:  "remote address",
+			setup: func(r *http.Request) { r.RemoteAddr = "203.0.113.7:12345" },
+			want:  "203.0.113.7",
+		},
+		{
+			name: "forwarded header wins",
+			setup: func(r *http.Request) {
+				r.RemoteAddr = "10.0.0.1:12345"
+				r.Header.Set("X-Forwarded-For", "203.0.113.7, 10.0.0.2")
+			},
+			want: "203.0.113.7",
+		},
+		{
+			name: "real ip header",
+			setup: func(r *http.Request) {
+				r.RemoteAddr = "10.0.0.1:12345"
+				r.Header.Set("X-Real-IP", "203.0.113.9")
+			},
+			want: "203.0.113.9",
+		},
+		{
+			name: "garbage header falls back",
+			setup: func(r *http.Request) {
+				r.RemoteAddr = "203.0.113.7:12345"
+				r.Header.Set("X-Forwarded-For", "not-an-ip")
+			},
+			want: "203.0.113.7",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			req := httptest.NewRequest(http.MethodGet, "/", nil)
+			req.Header = http.Header{}
+			tt.setup(req)
+
+			if got := clientIP(req); got != tt.want {
+				t.Errorf("clientIP() = %q, want %q", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestRetentionNoteIsMentionedOnStore(t *testing.T) {
+	s := newDiagServer()
+
+	rec := httptest.NewRecorder()
+	req := httptest.NewRequest(http.MethodPut, "/diag/x", strings.NewReader("{}"))
+	s.handleDiag(rec, req)
+
+	if !strings.Contains(rec.Body.String(), RetentionNote) {
+		t.Errorf("the response does not say how long the report is kept: %q", rec.Body.String())
+	}
+}
@@ -3,36 +3,101 @@ package relay
 import (
 	"encoding/json"
 	"log"
+	"net"
+	"strconv"
 	"sync"

 	"github.com/duffy/usb-server/internal/protocol"
 	"github.com/gorilla/websocket"
 )

+// sendQueueDepth bounds per-client outgoing backlog. A client that falls this
+// far behind is not going to catch up, and buffering more would let one stuck
+// peer consume the relay's memory.
+const sendQueueDepth = 256
+
+// outMsg is one queued WebSocket frame.
+type outMsg struct {
+	typ  int // websocket.TextMessage or websocket.BinaryMessage
+	data []byte
+}
+
 // Client represents a connected WebSocket client
 type Client struct {
 	ID   string
 	Hash string
-	Mode string // "share" or "use"
+	Mode string // "share", "use" or "both"
 	Name string
 	Conn *websocket.Conn
-	Send chan []byte // buffered channel for outgoing messages

-	mu sync.Mutex
+	// DirectPort is the port this client accepts direct tunnel connections on,
+	// 0 if it accepts none.
+	DirectPort int
+
+	// PublicIP is the source address the relay sees this client connect from.
+	// Peers cannot determine their own public address, so the relay supplies
+	// it when passing on a grant — that is the whole reason it is involved in
+	// setting up connections that then bypass it.
+	PublicIP string
+
+	// Send carries outgoing frames to this client's write pump. All writes go
+	// through it: writing to the socket directly from another client's read
+	// loop would block that peer — and, because the hub held its lock across
+	// the write, every other client with it.
+	Send chan outMsg
+
+	closeOnce sync.Once
+	dead      chan struct{}
 }

-// WriteJSON sends a JSON message to the client
-func (c *Client) WriteJSON(v interface{}) error {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	return c.Conn.WriteJSON(v)
+// newClient creates a client with its outgoing queue ready.
+func newClient(id, hash, mode, name string, conn *websocket.Conn) *Client {
+	return &Client{
+		ID:   id,
+		Hash: hash,
+		Mode: mode,
+		Name: name,
+		Conn: conn,
+		Send: make(chan outMsg, sendQueueDepth),
+		dead: make(chan struct{}),
+	}
 }

-// WriteBinary sends a binary message to the client
-func (c *Client) WriteBinary(data []byte) error {
-	c.mu.Lock()
-	defer c.mu.Unlock()
-	return c.Conn.WriteMessage(websocket.BinaryMessage, data)
+// enqueue queues a frame without blocking.
+// It reports false when the client's queue is full or it is already gone; the
+// caller should treat that as a disconnect rather than retrying.
+func (c *Client) enqueue(typ int, data []byte) bool {
+	select {
+	case <-c.dead:
+		return false
+	default:
+	}
+
+	select {
+	case c.Send <- outMsg{typ: typ, data: data}:
+		return true
+	case <-c.dead:
+		return false
+	default:
+		log.Printf("[hub] send queue full for %s (%s), dropping client",
+			protocol.ShortID(c.ID), c.Name)
+		c.kill()
+		return false
+	}
+}
+
+// kill marks the client dead and wakes its write pump. Idempotent.
+func (c *Client) kill() {
+	c.closeOnce.Do(func() { close(c.dead) })
+}
+
+// enqueueJSON marshals and queues a JSON control message.
+func (c *Client) enqueueJSON(v interface{}) bool {
+	data, err := json.Marshal(v)
+	if err != nil {
+		return false
+	}
+	return c.enqueue(websocket.TextMessage, data)
 }

 // Hub manages all connected clients and routes messages between them
@@ -58,39 +123,80 @@ func NewHub() *Hub {
 	}
 }

+// peers returns a snapshot of the clients in a hash group, excluding one ID.
+//
+// Taking a snapshot and releasing the lock before doing anything with the
+// clients is deliberate: holding the hub lock across a send is what let a
+// single slow peer stall registration and routing for everyone.
+func (h *Hub) peers(hash, excludeID string) []*Client {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	group := h.groups[hash]
+	result := make([]*Client, 0, len(group))
+	for _, c := range group {
+		if c.ID != excludeID {
+			result = append(result, c)
+		}
+	}
+	return result
+}
+
+// peer looks up a single client in a hash group.
+func (h *Hub) peer(hash, clientID string) *Client {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	group := h.groups[hash]
+	if group == nil {
+		return nil
+	}
+	return group[clientID]
+}
+
 // Register adds a client to its hash group
 func (h *Hub) Register(client *Client) {
 	h.mu.Lock()
-	defer h.mu.Unlock()
-
 	if h.groups[client.Hash] == nil {
 		h.groups[client.Hash] = make(map[string]*Client)
 	}
+	// A reconnecting client reuses its ID; drop the stale entry so its
+	// write pump exits instead of lingering with a dead socket.
+	if old, exists := h.groups[client.Hash][client.ID]; exists && old != client {
+		old.kill()
+	}
 	h.groups[client.Hash][client.ID] = client
+	h.mu.Unlock()

-	log.Printf("[hub] client registered: id=%s hash=%s..%s mode=%s name=%s",
-		client.ID, client.Hash[:8], client.Hash[len(client.Hash)-4:], client.Mode, client.Name)
+	log.Printf("[hub] client registered: id=%s hash=%s mode=%s name=%s",
+		protocol.ShortID(client.ID), protocol.ShortID(client.Hash), client.Mode, client.Name)

 	// Notify other clients in the group
-	h.broadcastToGroup(client.Hash, client.ID, &protocol.ClientJoined{
+	joined := &protocol.ClientJoined{
 		Type:     protocol.MsgClientJoined,
 		ClientID: client.ID,
 		Mode:     client.Mode,
 		Name:     client.Name,
-	})
+	}
+	for _, peer := range h.peers(client.Hash, client.ID) {
+		peer.enqueueJSON(joined)
+	}
 }

 // Unregister removes a client and cleans up its tunnels
 func (h *Hub) Unregister(client *Client) {
 	h.mu.Lock()
-	defer h.mu.Unlock()
-
 	group := h.groups[client.Hash]
 	if group == nil {
+		h.mu.Unlock()
 		return
 	}

-	delete(group, client.ID)
+	// Only remove this exact client: a reconnect may already have installed a
+	// newer connection under the same ID.
+	if group[client.ID] == client {
+		delete(group, client.ID)
+	}
 	if len(group) == 0 {
 		delete(h.groups, client.Hash)
 	}
@@ -101,21 +207,26 @@ func (h *Hub) Unregister(client *Client) {
 			delete(h.tunnels, tid)
 		}
 	}
+	h.mu.Unlock()

-	log.Printf("[hub] client unregistered: id=%s name=%s", client.ID, client.Name)
+	client.kill()

-	// Notify others
-	h.broadcastToGroup(client.Hash, client.ID, &protocol.ClientLeft{
+	log.Printf("[hub] client unregistered: id=%s name=%s", protocol.ShortID(client.ID), client.Name)
+
+	left := &protocol.ClientLeft{
 		Type:     protocol.MsgClientLeft,
 		ClientID: client.ID,
-	})
+	}
+	for _, peer := range h.peers(client.Hash, client.ID) {
+		peer.enqueueJSON(left)
+	}
 }

 // HandleTextMessage processes a JSON control message
 func (h *Hub) HandleTextMessage(sender *Client, data []byte) {
 	var env protocol.Envelope
 	if err := json.Unmarshal(data, &env); err != nil {
-		log.Printf("[hub] invalid message from %s: %v", sender.ID, err)
+		log.Printf("[hub] invalid message from %s: %v", protocol.ShortID(sender.ID), err)
 		return
 	}

@@ -135,9 +246,9 @@ func (h *Hub) HandleTextMessage(sender *Client, data []byte) {
 	case protocol.MsgDeviceReleased:
 		h.handleDeviceReleased(sender, data)
 	case protocol.MsgPing:
-		sender.WriteJSON(&protocol.Pong{Type: protocol.MsgPong})
+		sender.enqueueJSON(&protocol.Pong{Type: protocol.MsgPong})
 	default:
-		log.Printf("[hub] unknown message type from %s: %s", sender.ID, env.Type)
+		log.Printf("[hub] unknown message type from %s: %s", protocol.ShortID(sender.ID), env.Type)
 	}
 }

@@ -159,40 +270,32 @@ func (h *Hub) HandleBinaryMessage(sender *Client, data []byte) {

 	// Forward to the other end of the tunnel
 	var targetID string
-	if sender.ID == tunnel.ShareClient {
+	switch sender.ID {
+	case tunnel.ShareClient:
 		targetID = tunnel.UseClient
-	} else if sender.ID == tunnel.UseClient {
+	case tunnel.UseClient:
 		targetID = tunnel.ShareClient
-	} else {
+	default:
 		return
 	}

-	h.mu.RLock()
-	group := h.groups[sender.Hash]
-	if group != nil {
-		if target := group[targetID]; target != nil {
-			target.WriteBinary(data)
-		}
+	if target := h.peer(sender.Hash, targetID); target != nil {
+		target.enqueue(websocket.BinaryMessage, data)
 	}
-	h.mu.RUnlock()
 }

-// handleDeviceList broadcasts device list from share client to all use clients
+// handleDeviceList broadcasts a device list to every client in the group that
+// can consume devices.
 func (h *Hub) handleDeviceList(sender *Client, data []byte) {
-	if sender.Mode != protocol.ModeShare {
+	if !protocol.CanShare(sender.Mode) {
 		return
 	}

-	h.mu.RLock()
-	group := h.groups[sender.Hash]
-	for _, client := range group {
-		if client.ID != sender.ID && client.Mode == protocol.ModeUse {
-			client.mu.Lock()
-			client.Conn.WriteMessage(websocket.TextMessage, data)
-			client.mu.Unlock()
+	for _, client := range h.peers(sender.Hash, sender.ID) {
+		if protocol.CanUse(client.Mode) {
+			client.enqueue(websocket.TextMessage, data)
 		}
 	}
-	h.mu.RUnlock()
 }

 // handleRequestDevice forwards a device request to the target share client
@@ -202,22 +305,19 @@ func (h *Hub) handleRequestDevice(sender *Client, data []byte) {
 		return
 	}

-	h.mu.RLock()
-	group := h.groups[sender.Hash]
-	if group != nil {
-		if target := group[msg.TargetClient]; target != nil && target.Mode == protocol.ModeShare {
-			// Add the sender's ID so the share client knows who's requesting
-			enriched := map[string]interface{}{
-				"type":          protocol.MsgRequestDevice,
-				"target_client": msg.TargetClient,
-				"bus_id":        msg.BusID,
-				"request_id":    msg.RequestID,
-				"from_client":   sender.ID,
-			}
-			target.WriteJSON(enriched)
-		}
+	target := h.peer(sender.Hash, msg.TargetClient)
+	if target == nil || !protocol.CanShare(target.Mode) {
+		return
 	}
-	h.mu.RUnlock()
+
+	// Add the sender's ID so the share client knows who's requesting
+	target.enqueueJSON(map[string]interface{}{
+		"type":          protocol.MsgRequestDevice,
+		"target_client": msg.TargetClient,
+		"bus_id":        msg.BusID,
+		"request_id":    msg.RequestID,
+		"from_client":   sender.ID,
+	})
 }

 // handleDeviceGranted registers the tunnel and forwards to the requesting client
@@ -229,8 +329,10 @@ func (h *Hub) handleDeviceGranted(sender *Client, data []byte) {
 	if err := json.Unmarshal(data, &granted); err != nil {
 		return
 	}
+	if granted.TunnelID == "" {
+		return
+	}

-	// Register tunnel
 	h.mu.Lock()
 	h.tunnels[granted.TunnelID] = &Tunnel{
 		ID:          granted.TunnelID,
@@ -241,19 +343,44 @@ func (h *Hub) handleDeviceGranted(sender *Client, data []byte) {
 	h.mu.Unlock()

 	log.Printf("[hub] tunnel created: %s (share=%s, use=%s, device=%s)",
-		granted.TunnelID, sender.ID, granted.TargetClient, granted.BusID)
+		granted.TunnelID, protocol.ShortID(sender.ID), protocol.ShortID(granted.TargetClient), granted.BusID)

-	// Forward to use client
-	h.mu.RLock()
-	group := h.groups[sender.Hash]
-	if group != nil {
-		if target := group[granted.TargetClient]; target != nil {
-			target.mu.Lock()
-			target.Conn.WriteMessage(websocket.TextMessage, data)
-			target.mu.Unlock()
+	target := h.peer(sender.Hash, granted.TargetClient)
+	if target == nil {
+		return
+	}
+
+	// Add the address we see the granting client at. It cannot know its own
+	// public address, and this is what lets the two peers connect directly
+	// across NAT and take their USB traffic off this relay entirely.
+	out := data
+	if extra := publicEndpoint(sender); extra != "" {
+		granted.Endpoints = appendUnique(granted.Endpoints, extra)
+		if reencoded, err := json.Marshal(granted); err == nil {
+			out = reencoded
 		}
 	}
-	h.mu.RUnlock()
+
+	target.enqueue(websocket.TextMessage, out)
+}
+
+// publicEndpoint builds the host:port at which a client's direct listener
+// should be reachable from outside, or "" if it accepts no direct connections.
+func publicEndpoint(c *Client) string {
+	if c.DirectPort == 0 || c.PublicIP == "" {
+		return ""
+	}
+	return net.JoinHostPort(c.PublicIP, strconv.Itoa(c.DirectPort))
+}
+
+// appendUnique adds an entry unless it is already present.
+func appendUnique(list []string, item string) []string {
+	for _, existing := range list {
+		if existing == item {
+			return list
+		}
+	}
+	return append(list, item)
 }

 // handleDeviceDenied forwards denial to the requesting client
@@ -266,16 +393,9 @@ func (h *Hub) handleDeviceDenied(sender *Client, data []byte) {
 		return
 	}

-	h.mu.RLock()
-	group := h.groups[sender.Hash]
-	if group != nil {
-		if target := group[denied.TargetClient]; target != nil {
-			target.mu.Lock()
-			target.Conn.WriteMessage(websocket.TextMessage, data)
-			target.mu.Unlock()
-		}
+	if target := h.peer(sender.Hash, denied.TargetClient); target != nil {
+		target.enqueue(websocket.TextMessage, data)
 	}
-	h.mu.RUnlock()
 }

 // handleReleaseDevice forwards a release to the share client
@@ -296,21 +416,14 @@ func (h *Hub) handleReleaseDevice(sender *Client, data []byte) {
 	}
 	h.mu.Unlock()

-	// Forward to share client
-	h.mu.RLock()
-	group := h.groups[sender.Hash]
-	if group != nil {
-		if target := group[msg.TargetClient]; target != nil {
-			enriched := map[string]interface{}{
-				"type":         protocol.MsgReleaseDevice,
-				"target_client": msg.TargetClient,
-				"bus_id":       msg.BusID,
-				"from_client":  sender.ID,
-			}
-			target.WriteJSON(enriched)
-		}
+	if target := h.peer(sender.Hash, msg.TargetClient); target != nil {
+		target.enqueueJSON(map[string]interface{}{
+			"type":          protocol.MsgReleaseDevice,
+			"target_client": msg.TargetClient,
+			"bus_id":        msg.BusID,
+			"from_client":   sender.ID,
+		})
 	}
-	h.mu.RUnlock()
 }

 // handleForceRelease forwards a force-release request to the target share client
@@ -331,43 +444,36 @@ func (h *Hub) handleForceRelease(sender *Client, data []byte) {
 	}
 	h.mu.Unlock()

-	// Forward to share client
-	h.mu.RLock()
-	group := h.groups[sender.Hash]
-	if group != nil {
-		if target := group[msg.TargetClient]; target != nil && target.Mode == protocol.ModeShare {
-			enriched := map[string]interface{}{
-				"type":          protocol.MsgForceRelease,
-				"target_client": msg.TargetClient,
-				"bus_id":        msg.BusID,
-				"from_client":   sender.ID,
-			}
-			target.WriteJSON(enriched)
-		}
+	target := h.peer(sender.Hash, msg.TargetClient)
+	if target == nil || !protocol.CanShare(target.Mode) {
+		return
 	}
-	h.mu.RUnlock()
+
+	target.enqueueJSON(map[string]interface{}{
+		"type":          protocol.MsgForceRelease,
+		"target_client": msg.TargetClient,
+		"bus_id":        msg.BusID,
+		"from_client":   sender.ID,
+	})
 }

 // handleDeviceReleased broadcasts device released notification
 func (h *Hub) handleDeviceReleased(sender *Client, data []byte) {
-	h.mu.RLock()
-	group := h.groups[sender.Hash]
-	for _, client := range group {
-		if client.ID != sender.ID && client.Mode == protocol.ModeUse {
-			client.mu.Lock()
-			client.Conn.WriteMessage(websocket.TextMessage, data)
-			client.mu.Unlock()
+	for _, client := range h.peers(sender.Hash, sender.ID) {
+		if protocol.CanUse(client.Mode) {
+			client.enqueue(websocket.TextMessage, data)
 		}
 	}
-	h.mu.RUnlock()
 }

-// broadcastToGroup sends a message to all clients in a hash group except the sender
-func (h *Hub) broadcastToGroup(hash, excludeID string, msg interface{}) {
-	group := h.groups[hash]
-	for _, client := range group {
-		if client.ID != excludeID {
-			client.WriteJSON(msg)
-		}
+// GroupStats reports the number of clients per hash group, for diagnostics.
+func (h *Hub) GroupStats() map[string]int {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	stats := make(map[string]int, len(h.groups))
+	for hash, group := range h.groups {
+		stats[protocol.ShortID(hash)] = len(group)
 	}
+	return stats
 }
@@ -0,0 +1,465 @@
+package relay
+
+import (
+	"encoding/json"
+	"testing"
+	"time"
+
+	"github.com/duffy/usb-server/internal/protocol"
+	"github.com/gorilla/websocket"
+)
+
+// newTestClient builds a client without a socket. Nothing in the routing path
+// touches Conn — only the write pump does, and these tests read Send directly.
+func newTestClient(id, hash, mode string) *Client {
+	return newClient(id, hash, mode, "test-"+id, nil)
+}
+
+// drain collects everything queued for a client without blocking.
+func drain(c *Client) []outMsg {
+	var msgs []outMsg
+	for {
+		select {
+		case m := <-c.Send:
+			msgs = append(msgs, m)
+		default:
+			return msgs
+		}
+	}
+}
+
+// typeOf extracts the "type" field of a queued JSON control message.
+func typeOf(t *testing.T, m outMsg) string {
+	t.Helper()
+	var env protocol.Envelope
+	if err := json.Unmarshal(m.data, &env); err != nil {
+		t.Fatalf("queued message is not JSON: %v", err)
+	}
+	return env.Type
+}
+
+// countType drains a client and reports how many messages of one type it got.
+// Counting by type rather than total keeps these assertions independent of the
+// client_joined notifications registration produces.
+func countType(t *testing.T, c *Client, msgType string) int {
+	t.Helper()
+	n := 0
+	for _, m := range drain(c) {
+		if typeOf(t, m) == msgType {
+			n++
+		}
+	}
+	return n
+}
+
+// registerAll registers every client, then drains them, so that no client is
+// left holding join notifications from a peer that registered after it.
+func registerAll(h *Hub, clients ...*Client) {
+	for _, c := range clients {
+		h.Register(c)
+	}
+	for _, c := range clients {
+		drain(c)
+	}
+}
+
+func TestDeviceListReachesUseAndBothButNotShare(t *testing.T) {
+	h := NewHub()
+
+	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
+	user := newTestClient("user", "grp", protocol.ModeUse)
+	both := newTestClient("both", "grp", protocol.ModeBoth)
+	otherSharer := newTestClient("sharer2", "grp", protocol.ModeShare)
+
+	registerAll(h, sharer, user, both, otherSharer)
+
+	list, _ := json.Marshal(&protocol.DeviceList{
+		Type:     protocol.MsgDeviceList,
+		ClientID: sharer.ID,
+		Devices:  []protocol.USBDevice{{BusID: "1-1"}},
+	})
+	h.HandleTextMessage(sharer, list)
+
+	if got := countType(t, user, protocol.MsgDeviceList); got != 1 {
+		t.Errorf("use client received %d device lists, want 1", got)
+	}
+	if got := countType(t, both, protocol.MsgDeviceList); got != 1 {
+		t.Errorf("both client received %d device lists, want 1", got)
+	}
+	if got := countType(t, otherSharer, protocol.MsgDeviceList); got != 0 {
+		t.Errorf("share-only client received %d device lists, want 0", got)
+	}
+	if got := countType(t, sharer, protocol.MsgDeviceList); got != 0 {
+		t.Errorf("sender received %d copies of its own list, want 0", got)
+	}
+}
+
+// A "both" client must be able to offer devices, which means its device list
+// has to be routed like any share client's.
+func TestBothClientCanShare(t *testing.T) {
+	h := NewHub()
+
+	both := newTestClient("both", "grp", protocol.ModeBoth)
+	user := newTestClient("user", "grp", protocol.ModeUse)
+	registerAll(h, both, user)
+
+	list, _ := json.Marshal(&protocol.DeviceList{
+		Type:     protocol.MsgDeviceList,
+		ClientID: both.ID,
+		Devices:  []protocol.USBDevice{{BusID: "2-1"}},
+	})
+	h.HandleTextMessage(both, list)
+
+	if got := countType(t, user, protocol.MsgDeviceList); got != 1 {
+		t.Fatalf("use client received %d lists from a both-mode sharer, want 1", got)
+	}
+}
+
+func TestRequestDeviceReachesShareCapableTargetsOnly(t *testing.T) {
+	h := NewHub()
+
+	requester := newTestClient("req", "grp", protocol.ModeUse)
+	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
+	useOnly := newTestClient("useonly", "grp", protocol.ModeUse)
+
+	registerAll(h, requester, sharer, useOnly)
+
+	req, _ := json.Marshal(&protocol.RequestDevice{
+		Type:         protocol.MsgRequestDevice,
+		TargetClient: sharer.ID,
+		BusID:        "1-1",
+		RequestID:    "r1",
+	})
+	h.HandleTextMessage(requester, req)
+
+	var msgs []outMsg
+	for _, m := range drain(sharer) {
+		if typeOf(t, m) == protocol.MsgRequestDevice {
+			msgs = append(msgs, m)
+		}
+	}
+	if len(msgs) != 1 {
+		t.Fatalf("share client received %d requests, want 1", len(msgs))
+	}
+
+	// The relay must stamp in who is asking; the share side needs it to reply.
+	var got map[string]interface{}
+	json.Unmarshal(msgs[0].data, &got)
+	if got["from_client"] != requester.ID {
+		t.Errorf("from_client = %v, want %q", got["from_client"], requester.ID)
+	}
+
+	// A use-only client is not a valid target.
+	req2, _ := json.Marshal(&protocol.RequestDevice{
+		Type:         protocol.MsgRequestDevice,
+		TargetClient: useOnly.ID,
+		BusID:        "1-1",
+		RequestID:    "r2",
+	})
+	h.HandleTextMessage(requester, req2)
+
+	if got := countType(t, useOnly, protocol.MsgRequestDevice); got != 0 {
+		t.Errorf("use-only client received %d device requests, want 0", got)
+	}
+}
+
+func TestGroupsAreIsolatedByHash(t *testing.T) {
+	h := NewHub()
+
+	a := newTestClient("a", "hash-a", protocol.ModeShare)
+	b := newTestClient("b", "hash-b", protocol.ModeUse)
+	registerAll(h, a, b)
+
+	list, _ := json.Marshal(&protocol.DeviceList{
+		Type: protocol.MsgDeviceList, ClientID: a.ID,
+	})
+	h.HandleTextMessage(a, list)
+
+	if got := len(drain(b)); got != 0 {
+		t.Errorf("client in another hash group received %d messages, want 0", got)
+	}
+}
+
+func TestTunnelForwardsBothWays(t *testing.T) {
+	h := NewHub()
+
+	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
+	user := newTestClient("user", "grp", protocol.ModeUse)
+	registerAll(h, sharer, user)
+
+	tunnelID := "0123456789abcdef" // exactly TunnelHeaderSize
+	granted, _ := json.Marshal(map[string]interface{}{
+		"type":          protocol.MsgDeviceGranted,
+		"bus_id":        "1-1",
+		"tunnel_id":     tunnelID,
+		"request_id":    "r1",
+		"target_client": user.ID,
+	})
+	h.HandleTextMessage(sharer, granted)
+
+	if msgs := drain(user); len(msgs) != 1 || typeOf(t, msgs[0]) != protocol.MsgDeviceGranted {
+		t.Fatalf("grant was not forwarded to the use client: %v", msgs)
+	}
+
+	// use -> share
+	frame := append([]byte(tunnelID), 0xAA, 0xBB)
+	h.HandleBinaryMessage(user, frame)
+	msgs := drain(sharer)
+	if len(msgs) != 1 {
+		t.Fatalf("share client received %d tunnel frames, want 1", len(msgs))
+	}
+	if msgs[0].typ != websocket.BinaryMessage {
+		t.Errorf("tunnel frame sent as type %d, want binary", msgs[0].typ)
+	}
+
+	// share -> use
+	h.HandleBinaryMessage(sharer, frame)
+	if got := len(drain(user)); got != 1 {
+		t.Errorf("use client received %d tunnel frames, want 1", got)
+	}
+}
+
+func TestTunnelFramesForUnknownTunnelAreDropped(t *testing.T) {
+	h := NewHub()
+
+	a := newTestClient("a", "grp", protocol.ModeShare)
+	b := newTestClient("b", "grp", protocol.ModeUse)
+	registerAll(h, a, b)
+
+	h.HandleBinaryMessage(a, append([]byte("nonexistenttunnl"), 0x01))
+
+	if got := len(drain(b)); got != 0 {
+		t.Errorf("frame for an unknown tunnel was forwarded (%d messages)", got)
+	}
+}
+
+func TestUnregisterNotifiesPeersAndDropsTunnels(t *testing.T) {
+	h := NewHub()
+
+	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
+	user := newTestClient("user", "grp", protocol.ModeUse)
+	registerAll(h, sharer, user)
+
+	tunnelID := "0123456789abcdef"
+	granted, _ := json.Marshal(map[string]interface{}{
+		"type": protocol.MsgDeviceGranted, "bus_id": "1-1",
+		"tunnel_id": tunnelID, "target_client": user.ID,
+	})
+	h.HandleTextMessage(sharer, granted)
+	drain(user)
+
+	h.Unregister(sharer)
+
+	msgs := drain(user)
+	if len(msgs) != 1 || typeOf(t, msgs[0]) != protocol.MsgClientLeft {
+		t.Fatalf("peer was not told about the disconnect: %v", msgs)
+	}
+
+	h.mu.RLock()
+	_, stillThere := h.tunnels[tunnelID]
+	h.mu.RUnlock()
+	if stillThere {
+		t.Error("tunnel survived the share client leaving")
+	}
+}
+
+// Registration must not panic on short or empty identifiers: the relay
+// truncated hashes for logging, so a client with a 3-character hash used to
+// take the whole server down.
+func TestRegisterSurvivesShortIdentifiers(t *testing.T) {
+	h := NewHub()
+
+	for _, c := range []*Client{
+		newTestClient("", "", protocol.ModeUse),
+		newTestClient("x", "ab", protocol.ModeShare),
+		newTestClient("y", "abc", protocol.ModeBoth),
+	} {
+		h.Register(c)
+		h.Unregister(c)
+	}
+}
+
+// A client that stops draining must be dropped rather than allowed to consume
+// unbounded memory or block the peer producing the traffic.
+func TestFullSendQueueDropsClient(t *testing.T) {
+	h := NewHub()
+
+	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
+	slow := newTestClient("slow", "grp", protocol.ModeUse)
+	registerAll(h, sharer, slow)
+
+	list, _ := json.Marshal(&protocol.DeviceList{
+		Type: protocol.MsgDeviceList, ClientID: sharer.ID,
+	})
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		for i := 0; i < sendQueueDepth+50; i++ {
+			h.HandleTextMessage(sharer, list)
+		}
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(5 * time.Second):
+		t.Fatal("routing blocked on a client that never reads")
+	}
+
+	select {
+	case <-slow.dead:
+	default:
+		t.Error("client with a full queue was not dropped")
+	}
+}
+
+// Reconnecting with the same ID must retire the stale entry, not leave two.
+func TestReRegisterReplacesStaleClient(t *testing.T) {
+	h := NewHub()
+
+	first := newTestClient("dup", "grp", protocol.ModeUse)
+	h.Register(first)
+
+	second := newTestClient("dup", "grp", protocol.ModeUse)
+	h.Register(second)
+
+	select {
+	case <-first.dead:
+	default:
+		t.Error("stale connection was not killed on re-registration")
+	}
+
+	if got := h.GroupStats()[protocol.ShortID("grp")]; got != 1 {
+		t.Errorf("group holds %d clients, want 1", got)
+	}
+}
+
+func TestValidModeAndCapabilities(t *testing.T) {
+	tests := []struct {
+		mode                    string
+		valid, canShare, canUse bool
+	}{
+		{protocol.ModeShare, true, true, false},
+		{protocol.ModeUse, true, false, true},
+		{protocol.ModeBoth, true, true, true},
+		{"", false, false, false},
+		{"admin", false, false, false},
+	}
+
+	for _, tt := range tests {
+		if got := protocol.ValidMode(tt.mode); got != tt.valid {
+			t.Errorf("ValidMode(%q) = %v, want %v", tt.mode, got, tt.valid)
+		}
+		if got := protocol.CanShare(tt.mode); got != tt.canShare {
+			t.Errorf("CanShare(%q) = %v, want %v", tt.mode, got, tt.canShare)
+		}
+		if got := protocol.CanUse(tt.mode); got != tt.canUse {
+			t.Errorf("CanUse(%q) = %v, want %v", tt.mode, got, tt.canUse)
+		}
+	}
+}
+
+// The relay is the only party that knows a client's public address, so it
+// must add it to a grant. Without this, two peers behind NAT could never find
+// each other and every tunnel would stay relayed.
+func TestGrantGetsPublicEndpointAppended(t *testing.T) {
+	h := NewHub()
+
+	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
+	sharer.DirectPort = 41000
+	sharer.PublicIP = "203.0.113.7"
+	user := newTestClient("user", "grp", protocol.ModeUse)
+	registerAll(h, sharer, user)
+
+	granted, _ := json.Marshal(map[string]interface{}{
+		"type":          protocol.MsgDeviceGranted,
+		"bus_id":        "1-1",
+		"tunnel_id":     "0123456789abcdef",
+		"target_client": user.ID,
+		"endpoints":     []string{"192.168.1.5:41000"},
+		"encrypted":     true,
+	})
+	h.HandleTextMessage(sharer, granted)
+
+	msgs := drain(user)
+	if len(msgs) != 1 {
+		t.Fatalf("use client received %d messages, want 1", len(msgs))
+	}
+
+	var got protocol.DeviceGranted
+	if err := json.Unmarshal(msgs[0].data, &got); err != nil {
+		t.Fatalf("decoding forwarded grant: %v", err)
+	}
+
+	want := "203.0.113.7:41000"
+	var found, keptLocal bool
+	for _, ep := range got.Endpoints {
+		if ep == want {
+			found = true
+		}
+		if ep == "192.168.1.5:41000" {
+			keptLocal = true
+		}
+	}
+	if !found {
+		t.Errorf("endpoints %v do not include the public address %q", got.Endpoints, want)
+	}
+	if !keptLocal {
+		t.Errorf("endpoints %v lost the sharer's own local address", got.Endpoints)
+	}
+	if !got.Encrypted {
+		t.Error("the encrypted flag did not survive re-encoding")
+	}
+}
+
+// A client that accepts no direct connections must not have a bogus endpoint
+// invented for it.
+func TestGrantWithoutDirectPortIsUnchanged(t *testing.T) {
+	h := NewHub()
+
+	sharer := newTestClient("sharer", "grp", protocol.ModeShare)
+	sharer.PublicIP = "203.0.113.7" // reachable, but no listener
+	user := newTestClient("user", "grp", protocol.ModeUse)
+	registerAll(h, sharer, user)
+
+	granted, _ := json.Marshal(map[string]interface{}{
+		"type": protocol.MsgDeviceGranted, "bus_id": "1-1",
+		"tunnel_id": "0123456789abcdef", "target_client": user.ID,
+	})
+	h.HandleTextMessage(sharer, granted)
+
+	msgs := drain(user)
+	if len(msgs) != 1 {
+		t.Fatalf("use client received %d messages, want 1", len(msgs))
+	}
+
+	var got protocol.DeviceGranted
+	json.Unmarshal(msgs[0].data, &got)
+	if len(got.Endpoints) != 0 {
+		t.Errorf("endpoints = %v, want none for a client with no direct port", got.Endpoints)
+	}
+}
+
+func TestPublicEndpointRequiresBothParts(t *testing.T) {
+	tests := []struct {
+		name string
+		port int
+		ip   string
+		want string
+	}{
+		{"both present", 41000, "203.0.113.7", "203.0.113.7:41000"},
+		{"no port", 0, "203.0.113.7", ""},
+		{"no ip", 41000, "", ""},
+		{"neither", 0, "", ""},
+		{"ipv6", 41000, "2001:db8::1", "[2001:db8::1]:41000"},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			c := &Client{DirectPort: tt.port, PublicIP: tt.ip}
+			if got := publicEndpoint(c); got != tt.want {
+				t.Errorf("publicEndpoint() = %q, want %q", got, tt.want)
+			}
+		})
+	}
+}
@@ -3,13 +3,32 @@ package relay
 import (
 	"encoding/json"
 	"log"
+	"net"
 	"net/http"
+	"strings"
 	"time"

 	"github.com/duffy/usb-server/internal/protocol"
 	"github.com/gorilla/websocket"
 )

+const (
+	// readTimeout is how long a client may stay silent before we drop it.
+	// It must exceed pingInterval so that keepalive pongs refresh it.
+	readTimeout = 60 * time.Second
+
+	// pingInterval is how often the relay pings each client.
+	pingInterval = 20 * time.Second
+
+	// writeTimeout bounds a single frame write. Without it, a peer that has
+	// stopped reading would pin its write pump forever.
+	writeTimeout = 20 * time.Second
+
+	// maxMessageSize caps an inbound frame. Tunnel frames are at most 64 KB
+	// of USB payload plus the tunnel header; 1 MB leaves ample headroom.
+	maxMessageSize = 1024 * 1024
+)
+
 var upgrader = websocket.Upgrader{
 	ReadBufferSize:  64 * 1024,
 	WriteBufferSize: 64 * 1024,
@@ -22,6 +41,7 @@ var upgrader = websocket.Upgrader{
 type Server struct {
 	hub  *Hub
 	addr string
+	diag *diagStore
 }

 // NewServer creates a new relay server
@@ -29,6 +49,7 @@ func NewServer(addr string) *Server {
 	return &Server{
 		hub:  NewHub(),
 		addr: addr,
+		diag: newDiagStore(),
 	}
 }

@@ -37,9 +58,21 @@ func (s *Server) Run() error {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/ws", s.handleWebSocket)
 	mux.HandleFunc("/health", s.handleHealth)
+	mux.HandleFunc("/diag/", s.handleDiag)
+
+	// Timeouts bound how long a stuck client can hold a connection. The
+	// WebSocket route needs no write timeout — those connections are
+	// long-lived by design — so it is left to the per-message deadlines the
+	// write pump sets.
+	server := &http.Server{
+		Addr:              s.addr,
+		Handler:           mux,
+		ReadHeaderTimeout: 15 * time.Second,
+		IdleTimeout:       120 * time.Second,
+	}

 	log.Printf("[relay] starting on %s", s.addr)
-	return http.ListenAndServe(s.addr, mux)
+	return server.ListenAndServe()
 }

 func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
@@ -57,10 +90,10 @@ func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 	defer conn.Close()

 	// Set read limits and deadlines
-	conn.SetReadLimit(1024 * 1024) // 1MB max message
-	conn.SetReadDeadline(time.Now().Add(60 * time.Second))
+	conn.SetReadLimit(maxMessageSize)
+	conn.SetReadDeadline(time.Now().Add(readTimeout))
 	conn.SetPongHandler(func(string) error {
-		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
+		conn.SetReadDeadline(time.Now().Add(readTimeout))
 		return nil
 	})

@@ -78,55 +111,34 @@ func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 		return
 	}

-	if reg.Hash == "" || reg.ClientID == "" || (reg.Mode != protocol.ModeShare && reg.Mode != protocol.ModeUse) {
-		conn.WriteJSON(&protocol.ErrorMsg{Type: protocol.MsgError, Message: "missing required fields"})
+	if reg.Hash == "" || reg.ClientID == "" || !protocol.ValidMode(reg.Mode) {
+		conn.WriteJSON(&protocol.ErrorMsg{Type: protocol.MsgError, Message: "missing or invalid registration fields"})
 		return
 	}

-	client := &Client{
-		ID:   reg.ClientID,
-		Hash: reg.Hash,
-		Mode: reg.Mode,
-		Name: reg.Name,
-		Conn: conn,
-		Send: make(chan []byte, 256),
-	}
+	client := newClient(reg.ClientID, reg.Hash, reg.Mode, reg.Name, conn)
+	client.DirectPort = reg.DirectPort
+	client.PublicIP = clientIP(r)

 	s.hub.Register(client)
 	defer s.hub.Unregister(client)

-	// Start ping ticker
-	done := make(chan struct{})
-	go func() {
-		ticker := time.NewTicker(30 * time.Second)
-		defer ticker.Stop()
-		for {
-			select {
-			case <-ticker.C:
-				client.mu.Lock()
-				err := conn.WriteMessage(websocket.PingMessage, nil)
-				client.mu.Unlock()
-				if err != nil {
-					return
-				}
-			case <-done:
-				return
-			}
-		}
-	}()
-	defer close(done)
+	// The write pump owns the socket's write side: every frame for this
+	// client, plus keepalive pings, goes through it. Nothing else may write,
+	// which is what keeps one unresponsive peer from blocking the hub.
+	go s.writePump(client)

 	// Read loop
 	for {
 		msgType, data, err := conn.ReadMessage()
 		if err != nil {
 			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseNormalClosure) {
-				log.Printf("[relay] read error from %s: %v", client.ID, err)
+				log.Printf("[relay] read error from %s: %v", protocol.ShortID(client.ID), err)
 			}
 			break
 		}

-		conn.SetReadDeadline(time.Now().Add(60 * time.Second))
+		conn.SetReadDeadline(time.Now().Add(readTimeout))

 		switch msgType {
 		case websocket.TextMessage:
@@ -135,4 +147,67 @@ func (s *Server) handleWebSocket(w http.ResponseWriter, r *http.Request) {
 			s.hub.HandleBinaryMessage(client, data)
 		}
 	}
+
+	client.kill()
+}
+
+// clientIP determines the address a client connects from, which is passed on
+// to its peers so they can reach it directly.
+//
+// X-Forwarded-For is honoured because relays are commonly deployed behind a
+// reverse proxy, where RemoteAddr would otherwise be the proxy itself. Only
+// the first entry is used: later ones are supplied by upstream hops and are
+// not trustworthy. A wrong value here costs a failed direct attempt and a
+// fallback to relaying, never a security property — the peer still has to
+// prove group membership in the handshake.
+func clientIP(r *http.Request) string {
+	if fwd := r.Header.Get("X-Forwarded-For"); fwd != "" {
+		first := strings.TrimSpace(strings.Split(fwd, ",")[0])
+		if ip := net.ParseIP(first); ip != nil {
+			return ip.String()
+		}
+	}
+	if real := strings.TrimSpace(r.Header.Get("X-Real-IP")); real != "" {
+		if ip := net.ParseIP(real); ip != nil {
+			return ip.String()
+		}
+	}
+
+	host, _, err := net.SplitHostPort(r.RemoteAddr)
+	if err != nil {
+		return ""
+	}
+	if ip := net.ParseIP(host); ip != nil {
+		return ip.String()
+	}
+	return ""
+}
+
+// writePump serialises all writes to one client's socket.
+func (s *Server) writePump(client *Client) {
+	ticker := time.NewTicker(pingInterval)
+	defer ticker.Stop()
+	defer client.Conn.Close() // unblocks the read loop when we give up
+
+	for {
+		select {
+		case msg := <-client.Send:
+			client.Conn.SetWriteDeadline(time.Now().Add(writeTimeout))
+			if err := client.Conn.WriteMessage(msg.typ, msg.data); err != nil {
+				log.Printf("[relay] write error to %s: %v", protocol.ShortID(client.ID), err)
+				client.kill()
+				return
+			}
+
+		case <-ticker.C:
+			client.Conn.SetWriteDeadline(time.Now().Add(writeTimeout))
+			if err := client.Conn.WriteMessage(websocket.PingMessage, nil); err != nil {
+				client.kill()
+				return
+			}
+
+		case <-client.dead:
+			return
+		}
+	}
 }