feat(tsconnect-worker): worker-side resource transfer registry

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

packages/tsconnect-worker/src/transfers.test.ts

@@ -0,0 +1,174 @@
+import { test, suite } from "node:test"
+import assert from "node:assert/strict"
+import { createTransferRegistry, type AnyResource } from "./transfers.js"
+import type { TransferToken } from "./protocol.js"
+
+// ── Helpers ───────────────────────────────────────────────────────────────────
+
+async function waitFor(predicate: () => boolean, timeoutMs = 2000): Promise<void> {
+  const start = Date.now()
+  while (!predicate()) {
+    if (Date.now() - start > timeoutMs) throw new Error("waitFor: condition not met within timeout")
+    await new Promise((r) => setImmediate(r))
+  }
+}
+
+function fakeConn(): {
+  obj: { localAddr: string; remoteAddr: string; closed: boolean; close(): void }
+} {
+  const obj = {
+    localAddr: "127.0.0.1:1",
+    remoteAddr: "127.0.0.1:2",
+    closed: false,
+    close() {
+      this.closed = true
+    },
+  }
+  return { obj }
+}
+
+function connResource(): { res: AnyResource; port1: MessagePort; obj: { closed: boolean } } {
+  const { obj } = fakeConn()
+  const { port1 } = new MessageChannel()
+  const res = { kind: "conn", obj, port1 } as unknown as AnyResource
+  return { res, port1, obj }
+}
+
+function entryWith(res: AnyResource, resourceId = 0): { resources: Map<number, AnyResource> } {
+  return { resources: new Map([[resourceId, res]]) }
+}
+
+// ── detach ────────────────────────────────────────────────────────────────────
+
+suite("detach", () => {
+  test("removes the resource from the entry and closes port1 without closing the socket", () => {
+    const registry = createTransferRegistry()
+    const { res, port1, obj } = connResource()
+    let port1Closed = false
+    const origClose = port1.close.bind(port1)
+    port1.close = () => {
+      port1Closed = true
+      origClose()
+    }
+    const entry = entryWith(res)
+
+    const { transferId, kind, meta } = registry.detach(entry, 0)
+
+    assert.equal(entry.resources.size, 0)
+    assert.equal(obj.closed, false)
+    assert.equal(port1Closed, true)
+    assert.equal(kind, "conn")
+    assert.deepEqual(meta, { kind: "conn", localAddr: "127.0.0.1:1", remoteAddr: "127.0.0.1:2" })
+    assert.ok(registry.hasTransfer(transferId))
+  })
+
+  test("throws when the resource id is unknown", () => {
+    const registry = createTransferRegistry()
+    assert.throws(() => registry.detach({ resources: new Map() }, 5), /resource not found/)
+  })
+})
+
+// ── claim ─────────────────────────────────────────────────────────────────────
+
+suite("claim", () => {
+  test("returns the resource and removes it from pending", () => {
+    const registry = createTransferRegistry()
+    const { res } = connResource()
+    const { transferId } = registry.detach(entryWith(res), 0)
+
+    const claimed = registry.claim(transferId)
+    assert.equal(claimed, res)
+    assert.equal(registry.hasTransfer(transferId), false)
+    assert.equal(registry.hasPending(), false)
+  })
+
+  test("throws on duplicate claim", () => {
+    const registry = createTransferRegistry()
+    const { res } = connResource()
+    const { transferId } = registry.detach(entryWith(res), 0)
+    registry.claim(transferId)
+    assert.throws(() => registry.claim(transferId), /transfer already claimed or expired/)
+  })
+
+  test("claim after TTL expiry throws and the resource was closed", async () => {
+    const registry = createTransferRegistry({ ttl: 10 })
+    const { res, obj } = connResource()
+    const { transferId } = registry.detach(entryWith(res), 0)
+
+    await waitFor(() => obj.closed)
+    assert.throws(() => registry.claim(transferId), /transfer already claimed or expired/)
+  })
+})
+
+// ── send / queue ──────────────────────────────────────────────────────────────
+
+suite("send", () => {
+  test("delivers immediately when the target key is registered", () => {
+    const registry = createTransferRegistry()
+    const delivered: Array<{ clientId: number; channel: string }> = []
+    registry.registerKey("B", 7, () => {})
+
+    const result = registry.send("A", "B", "chan", { x: 1 }, [], [], (clientId, msg) => {
+      delivered.push({ clientId, channel: msg.channel })
+    })
+
+    assert.deepEqual(result, { delivered: true })
+    assert.deepEqual(delivered, [{ clientId: 7, channel: "chan" }])
+  })
+
+  test("queues when the target key is not registered, then flushes on registerKey", () => {
+    const registry = createTransferRegistry()
+    const result = registry.send("A", "B", "chan", { x: 1 }, [], [], () => {
+      throw new Error("should not deliver")
+    })
+    assert.deepEqual(result, { delivered: false })
+
+    const flushed: string[] = []
+    registry.registerKey("B", 3, (msg) => flushed.push(msg.channel))
+    assert.deepEqual(flushed, ["chan"])
+  })
+
+  test("queue TTL expiry drops the envelope and closes referenced pending resources", async () => {
+    const registry = createTransferRegistry({ ttl: 10 })
+    const { res, obj } = connResource()
+    const { transferId } = registry.detach(entryWith(res), 0)
+    const token: TransferToken = {
+      transferId,
+      kind: "conn",
+      meta: { kind: "conn", localAddr: "a", remoteAddr: "b" },
+    }
+
+    registry.send("A", "B", "chan", null, [token], [], () => {
+      throw new Error("should not deliver")
+    })
+
+    await waitFor(() => obj.closed)
+    assert.equal(registry.hasTransfer(transferId), false)
+
+    const flushed: string[] = []
+    registry.registerKey("B", 1, (msg) => flushed.push(msg.channel))
+    assert.deepEqual(flushed, [])
+  })
+})
+
+// ── onEmpty ───────────────────────────────────────────────────────────────────
+
+suite("onEmpty", () => {
+  test("fires when the last pending transfer is claimed", () => {
+    let empties = 0
+    const registry = createTransferRegistry({ onEmpty: () => empties++ })
+    const { res } = connResource()
+    const { transferId } = registry.detach(entryWith(res), 0)
+    registry.claim(transferId)
+    assert.equal(empties, 1)
+  })
+
+  test("fires when the last pending transfer expires", async () => {
+    let empties = 0
+    const registry = createTransferRegistry({ ttl: 10, onEmpty: () => empties++ })
+    const { res, obj } = connResource()
+    registry.detach(entryWith(res), 0)
+    await waitFor(() => obj.closed)
+    await waitFor(() => empties === 1)
+  })
+})

packages/tsconnect-worker/src/transfers.ts

@@ -0,0 +1,171 @@
+import type { Conn, TCPListener, PacketConn } from "@webnet/tsconnect"
+import type { ResourceMeta, TransferToken, W2C } from "./protocol.js"
+
+export type AnyResource =
+  | { kind: "conn"; obj: Conn; port1: MessagePort }
+  | { kind: "listener"; obj: TCPListener; port1: MessagePort }
+  | { kind: "packetconn"; obj: PacketConn; port1: MessagePort }
+
+type TransferMessage = Extract<W2C, { type: "transfer" }>
+
+type Timer = ReturnType<typeof setTimeout>
+
+type QueuedEnvelope = {
+  channel: string
+  state: unknown
+  tokens: TransferToken[]
+  fromKey?: string
+  transferables: Transferable[]
+  timer: Timer
+}
+
+export type Deliver = (
+  clientId: number,
+  msg: TransferMessage,
+  transferables: Transferable[],
+) => void
+
+export type Flush = (msg: TransferMessage, transferables: Transferable[]) => void
+
+export type TransferRegistry = ReturnType<typeof createTransferRegistry>
+
+function metaOf(res: AnyResource): ResourceMeta {
+  if (res.kind === "conn")
+    return { kind: "conn", localAddr: res.obj.localAddr, remoteAddr: res.obj.remoteAddr }
+  if (res.kind === "listener") return { kind: "listener", addr: res.obj.addr }
+  return { kind: "packetconn", localAddr: res.obj.localAddr }
+}
+
+export function createTransferRegistry(opts: { ttl?: number; onEmpty?: () => void } = {}) {
+  const ttl = opts.ttl ?? 60_000
+  const pendingTransfers = new Map<string, { resource: AnyResource; timer: Timer }>()
+  const keyToClient = new Map<string, number>()
+  const transferQueue = new Map<string, QueuedEnvelope[]>()
+
+  function checkEmpty(): void {
+    if (pendingTransfers.size === 0) opts.onEmpty?.()
+  }
+
+  function closePending(transferId: string): void {
+    const pending = pendingTransfers.get(transferId)
+    if (!pending) return
+    clearTimeout(pending.timer)
+    pendingTransfers.delete(transferId)
+    try {
+      pending.resource.obj.close()
+    } catch {
+      /* already closed */
+    }
+  }
+
+  function detach(
+    entry: { resources: Map<number, AnyResource> },
+    resourceId: number,
+  ): { transferId: string; kind: AnyResource["kind"]; meta: ResourceMeta } {
+    const res = entry.resources.get(resourceId)
+    if (!res) throw new Error("resource not found")
+    entry.resources.delete(resourceId)
+    res.port1.onmessage = null
+    res.port1.close()
+
+    const transferId = crypto.randomUUID()
+    const timer = setTimeout(() => {
+      pendingTransfers.delete(transferId)
+      try {
+        res.obj.close()
+      } catch {
+        /* already closed */
+      }
+      checkEmpty()
+    }, ttl)
+    pendingTransfers.set(transferId, { resource: res, timer })
+    return { transferId, kind: res.kind, meta: metaOf(res) }
+  }
+
+  function claim(transferId: string): AnyResource {
+    const pending = pendingTransfers.get(transferId)
+    if (!pending) throw new Error("transfer already claimed or expired")
+    clearTimeout(pending.timer)
+    pendingTransfers.delete(transferId)
+    checkEmpty()
+    return pending.resource
+  }
+
+  function cancel(transferId: string): void {
+    closePending(transferId)
+    checkEmpty()
+  }
+
+  function send(
+    fromKey: string | undefined,
+    targetKey: string,
+    channel: string,
+    state: unknown,
+    tokens: TransferToken[],
+    transferables: Transferable[],
+    deliver: Deliver,
+  ): { delivered: boolean } {
+    const clientId = keyToClient.get(targetKey)
+    if (clientId !== undefined) {
+      deliver(clientId, { type: "transfer", channel, state, tokens, fromKey }, transferables)
+      return { delivered: true }
+    }
+
+    const envelope: QueuedEnvelope = {
+      channel,
+      state,
+      tokens,
+      fromKey,
+      transferables,
+      timer: setTimeout(() => {
+        const queue = transferQueue.get(targetKey)
+        if (queue) {
+          const idx = queue.indexOf(envelope)
+          if (idx >= 0) queue.splice(idx, 1)
+          if (queue.length === 0) transferQueue.delete(targetKey)
+        }
+        for (const token of tokens) closePending(token.transferId)
+        checkEmpty()
+      }, ttl),
+    }
+    const queue = transferQueue.get(targetKey) ?? []
+    queue.push(envelope)
+    transferQueue.set(targetKey, queue)
+    return { delivered: false }
+  }
+
+  function registerKey(key: string, clientId: number, flush: Flush): void {
+    keyToClient.set(key, clientId)
+    const queue = transferQueue.get(key)
+    if (!queue) return
+    transferQueue.delete(key)
+    for (const env of queue) {
+      clearTimeout(env.timer)
+      flush(
+        {
+          type: "transfer",
+          channel: env.channel,
+          state: env.state,
+          tokens: env.tokens,
+          fromKey: env.fromKey,
+        },
+        env.transferables,
+      )
+    }
+  }
+
+  function unregisterKey(key: string, clientId: number): void {
+    if (keyToClient.get(key) === clientId) keyToClient.delete(key)
+  }
+
+  return {
+    detach,
+    claim,
+    cancel,
+    send,
+    registerKey,
+    unregisterKey,
+    hasTransfer: (transferId: string) => pendingTransfers.has(transferId),
+    hasPending: () => pendingTransfers.size > 0,
+  }
+}

packages/tsconnect-worker/src/worker.ts

@@ -24,16 +24,14 @@ import type {
   FileOpsLimits,
 } from "./protocol.js"
 import { pumpStreamToPort, portToReadableStream } from "./protocol.js"
+import type { TransferToken } from "./protocol.js"
+import { createTransferRegistry } from "./transfers.js"
+import type { AnyResource } from "./transfers.js"
 
 const sw = self as unknown as SharedWorkerGlobalScope
 
 // ── State ────────────────────────────────────────────────────────────────────
 
-export type AnyResource =
-  | { kind: "conn"; obj: Conn; port1: MessagePort }
-  | { kind: "listener"; obj: TCPListener; port1: MessagePort }
-  | { kind: "packetconn"; obj: PacketConn; port1: MessagePort }
-
 type ClientEntry = {
   port: MessagePort
   resources: Map<number, AnyResource>
@@ -49,8 +47,10 @@ const clients = new Map<number, ClientEntry>()
 let nextClientId = 0
 let nextResourceId = 0
 
+const registry = createTransferRegistry({ onEmpty: () => maybeShutdownIpn() })
+
 // Clients that connected before init completed, with their lock name already extracted.
-const pendingClients: Array<{ port: MessagePort; lockName: string }> = []
+const pendingClients: Array<{ port: MessagePort; lockName: string; clientKey?: string }> = []
 let initState: "idle" | "pending" | "ready" | "failed" = "idle"
 let initError = ""
 
@@ -595,6 +595,78 @@ async function handleCall(
         ok(undefined)
         break
       }
+      case "detachResource": {
+        ok(registry.detach(entry, args[0] as number))
+        break
+      }
+      case "claimResource": {
+        const resource = registry.claim(args[0] as string)
+        if (resource.kind === "conn") {
+          const { port2, resourceId } = registerConnResource(entry, resource.obj)
+          const msg: W2C = {
+            type: "conn",
+            id,
+            localAddr: resource.obj.localAddr,
+            remoteAddr: resource.obj.remoteAddr,
+            port: port2,
+            resourceId,
+          }
+          send(port, msg, [port2])
+        } else if (resource.kind === "listener") {
+          const { port2, resourceId } = registerListenerResource(entry, resource.obj)
+          const msg: W2C = {
+            type: "listener",
+            id,
+            addr: resource.obj.addr,
+            port: port2,
+            resourceId,
+          }
+          send(port, msg, [port2])
+        } else {
+          const { port2, resourceId } = registerPacketConnResource(entry, resource.obj)
+          const msg: W2C = {
+            type: "packetconn",
+            id,
+            localAddr: resource.obj.localAddr,
+            port: port2,
+            resourceId,
+          }
+          send(port, msg, [port2])
+        }
+        break
+      }
+      case "transferSend": {
+        const [targetKey, channel, state, tokens, transferables] = args as [
+          string,
+          string,
+          unknown,
+          TransferToken[],
+          Transferable[] | undefined,
+        ]
+        const missing = tokens.find((t) => !registry.hasTransfer(t.transferId))
+        if (missing) {
+          fail(`unknown transfer: ${missing.transferId}`)
+          break
+        }
+        const { delivered } = registry.send(
+          entry.clientKey,
+          targetKey,
+          channel,
+          state,
+          tokens,
+          transferables ?? [],
+          (clientId, msg, xfer) => {
+            const target = clients.get(clientId)
+            if (target) send(target.port, msg, xfer)
+          },
+        )
+        ok({ delivered })
+        break
+      }
+      case "cancelTransfer":
+        registry.cancel(args[0] as string)
+        ok(undefined)
+        break
       default:
         fail(`unknown method: ${method}`)
     }
@@ -605,10 +677,19 @@ async function handleCall(
 
 // ── Client cleanup ────────────────────────────────────────────────────────────
 
+function maybeShutdownIpn(): void {
+  if (clients.size === 0 && !registry.hasPending() && ipn) {
+    ipn.shutdown().catch(() => {
+      /* best-effort */
+    })
+  }
+}
+
 function cleanupClient(clientId: number): void {
   const entry = clients.get(clientId)
   if (!entry) return
   clients.delete(clientId)
+  if (entry.clientKey !== undefined) registry.unregisterKey(entry.clientKey, clientId)
 
   for (const res of entry.resources.values()) {
     try {
@@ -629,21 +710,18 @@ function cleanupClient(clientId: number): void {
     }
   }
 
-  if (clients.size === 0 && ipn) {
-    ipn.shutdown().catch(() => {
-      /* best-effort */
-    })
-  }
+  maybeShutdownIpn()
 }
 
 // ── Client registration ───────────────────────────────────────────────────────
 
-function registerClient(port: MessagePort, lockName: string): void {
+function registerClient(port: MessagePort, lockName: string, clientKey?: string): void {
   const clientId = nextClientId++
   const entry: ClientEntry = {
     port,
     resources: new Map(),
     driveRegistered: false,
+    clientKey,
   }
   clients.set(clientId, entry)
 
@@ -659,6 +737,12 @@ function registerClient(port: MessagePort, lockName: string): void {
     return Promise.resolve()
   })
 
+  if (clientKey !== undefined) {
+    registry.registerKey(clientKey, clientId, (msg, transferables) =>
+      send(port, msg, transferables),
+    )
+  }
+
   send(port, { type: "ready" })
 
   port.onmessage = (e: MessageEvent) => {
@@ -733,7 +817,7 @@ sw.onconnect = (e: MessageEvent) => {
     }
 
     if (initState === "ready") {
-      registerClient(port, msg.lockName)
+      registerClient(port, msg.lockName, msg.clientKey)
       return
     }
 
@@ -742,12 +826,12 @@ sw.onconnect = (e: MessageEvent) => {
         send(port, { type: "initError", error: "init config required for first connection" })
         return
       }
-      pendingClients.push({ port, lockName: msg.lockName })
+      pendingClients.push({ port, lockName: msg.lockName, clientKey: msg.clientKey })
       const localConfig = msg.config
       try {
         await init(localConfig)
         for (const pending of pendingClients) {
-          registerClient(pending.port, pending.lockName)
+          registerClient(pending.port, pending.lockName, pending.clientKey)
         }
       } catch (err) {
         const error = errMsg(err)
@@ -758,7 +842,7 @@ sw.onconnect = (e: MessageEvent) => {
       pendingClients.length = 0
     } else {
       // initState === "pending": another tab is already initialising; queue.
-      pendingClients.push({ port, lockName: msg.lockName })
+      pendingClients.push({ port, lockName: msg.lockName, clientKey: msg.clientKey })
     }
   }