From b45e34021d81a632a21af7d5831c0906d01f33dd Mon Sep 17 00:00:00 2001 From: "Ilya G." <2004gusak@gmail.com> Date: Fri, 13 Mar 2026 15:44:19 +0500 Subject: [PATCH] Allow reclaiming stale custom subdomains --- edge/__tests__/tunnel.test.ts | 45 +++++++++++++++++++ edge/src/tunnel.ts | 85 +++++++++++++++++++++++++++-------- 2 files changed, 112 insertions(+), 18 deletions(-) diff --git a/edge/__tests__/tunnel.test.ts b/edge/__tests__/tunnel.test.ts index 7ddaf2e..c459547 100644 --- a/edge/__tests__/tunnel.test.ts +++ b/edge/__tests__/tunnel.test.ts @@ -173,6 +173,51 @@ describe("Tunnel Durable Object", () => { ws.close(); }); + it("reclaims custom subdomain when prior tunnel is stale", async () => { + const testToken = "test-token-999"; + await env.DB.prepare( + "INSERT OR REPLACE INTO users (github_id, username, token) VALUES (?, ?, ?)" + ).bind("999", "staleuser", testToken).run(); + await env.DB.prepare( + "INSERT OR REPLACE INTO reserved_subdomains (subdomain, user_id) VALUES (?, ?)" + ).bind("myapp", "999").run(); + + const staleId = env.TUNNEL.newUniqueId(); + await env.DB.prepare( + "INSERT OR REPLACE INTO tunnels (subdomain, client_id) VALUES (?, ?)" + ).bind("myapp", staleId.toString()).run(); + + const stub = getStub(); + const response = await stub.fetch(`https://${DOMAIN}/_wormhole/register`, { + headers: { Upgrade: "websocket" }, + }); + + const ws = response.webSocket!; + ws.accept(); + + const messagePromise = new Promise((resolve) => { + ws.addEventListener("message", (event) => { + resolve(event.data as string); + }); + }); + + ws.send(JSON.stringify({ type: "register", protocol: "http", subdomain: "myapp", token: testToken })); + + const msg = await messagePromise; + const parsed = JSON.parse(msg); + + expect(parsed.type).toBe("registered"); + expect(parsed.subdomain).toBe("myapp"); + + const row = await env.DB.prepare( + "SELECT client_id FROM tunnels WHERE subdomain = ?" + ).bind("myapp").first<{ client_id: string }>(); + expect(row).not.toBeNull(); + expect(row!.client_id).not.toBe(staleId.toString()); + + ws.close(); + }); + it("rejects reserved system subdomains", async () => { const testToken = "test-token-12345"; await env.DB.prepare( diff --git a/edge/src/tunnel.ts b/edge/src/tunnel.ts index 4857ccc..8c27b25 100644 --- a/edge/src/tunnel.ts +++ b/edge/src/tunnel.ts @@ -22,6 +22,8 @@ interface HttpResponseMessage { const REQUEST_TIMEOUT_MS = 30_000; const MAX_SUBDOMAINS_PER_USER = 3; +const INTERNAL_STATUS_HEADER = "x-wormhole-internal"; +const INTERNAL_STATUS_VALUE = "status"; export class Tunnel implements DurableObject { private clientWs: WebSocket | null = null; @@ -38,6 +40,10 @@ export class Tunnel implements DurableObject { async fetch(request: Request): Promise { const url = new URL(request.url); + if (url.pathname === "/_wormhole/status" && request.headers.get(INTERNAL_STATUS_HEADER) === INTERNAL_STATUS_VALUE) { + return this.handleStatus(); + } + // Handle WebSocket upgrade for tunnel registration if (url.pathname === "/_wormhole/register") { return this.handleRegister(request); @@ -244,18 +250,7 @@ export class Tunnel implements DurableObject { async webSocketClose(ws: WebSocket, code: number, reason: string, wasClean: boolean): Promise { this.clientWs = null; - // Recover subdomain from storage if lost during hibernation - if (!this.subdomain) { - this.subdomain = (await this.state.storage.get("subdomain")) ?? null; - } - - // Remove subdomain from D1 and storage - if (this.subdomain) { - await this.env.DB.prepare("DELETE FROM tunnels WHERE subdomain = ?") - .bind(this.subdomain).run(); - await this.state.storage.delete("subdomain"); - this.subdomain = null; - } + await this.cleanupSubdomain(); // Reject all pending requests for (const [id, pending] of this.pendingRequests) { @@ -363,13 +358,19 @@ export class Tunnel implements DurableObject { // Check if subdomain is already in use by an active tunnel const active = await this.env.DB.prepare( "SELECT client_id FROM tunnels WHERE subdomain = ?" - ).bind(requested).first(); + ).bind(requested).first<{ client_id: string }>(); if (active) { - ws.send(JSON.stringify({ - type: "register_error", - error: `Subdomain "${requested}" is already in use.`, - })); - return; + const isActive = await this.isTunnelActive(active.client_id); + if (isActive) { + ws.send(JSON.stringify({ + type: "register_error", + error: `Subdomain "${requested}" is already in use.`, + })); + return; + } + await this.env.DB.prepare( + "DELETE FROM tunnels WHERE subdomain = ?" + ).bind(requested).run(); } // Auto-reserve on first use (if not already reserved by this user) @@ -440,6 +441,54 @@ export class Tunnel implements DurableObject { headers, })); } + + private isClientConnected(): boolean { + const ws = this.getClientWs(); + return !!ws && ws.readyState === WebSocket.OPEN; + } + + private async handleStatus(): Promise { + const connected = this.isClientConnected(); + if (!connected) { + await this.cleanupSubdomain(); + } + return new Response(JSON.stringify({ connected }), { + status: 200, + headers: { "Content-Type": "application/json" }, + }); + } + + private async cleanupSubdomain(): Promise { + if (!this.subdomain) { + this.subdomain = (await this.state.storage.get("subdomain")) ?? null; + } + + if (!this.subdomain) { + return; + } + + await this.env.DB.prepare("DELETE FROM tunnels WHERE subdomain = ?") + .bind(this.subdomain).run(); + await this.state.storage.delete("subdomain"); + this.subdomain = null; + } + + private async isTunnelActive(clientId: string): Promise { + try { + const id = this.env.TUNNEL.idFromString(clientId); + const stub = this.env.TUNNEL.get(id); + const response = await stub.fetch("https://internal/_wormhole/status", { + headers: { [INTERNAL_STATUS_HEADER]: INTERNAL_STATUS_VALUE }, + }); + if (!response.ok) { + return true; + } + const data = await response.json<{ connected?: boolean }>(); + return data.connected === true; + } catch { + return true; + } + } } function generateSubdomain(): string {