tunnel closes if upstream is not connected

Author: remorses

cloudflare-tunnel/package.json

@@ -6,7 +6,7 @@
   "repository": "https://github.com/remorses/",
   "scripts": {
     "deployment": "pnpm tsc && wrangler deploy --env preview",
-    "deployment:prod": "pnpm tsc && wrangler deploy && pnpm deployment",
+    "deployment:prod": "sudo pnpm tsc && wrangler deploy && pnpm deployment",
     "dev": "wrangler dev",
     "test": "vitest --run"
   },

cloudflare-tunnel/src/tunnel.test.ts

@@ -209,4 +209,30 @@ describe.concurrent('Tunnel WebSocket', () => {
     client1.close()
     client2.close()
   })
+
+  test('client connection fails when no upstream is connected', async () => {
+    const tunnelId = getTunnelId()
+
+    // Try to connect client without upstream
+    const client = new WebSocket(`${WS_URL}/downstream?id=${tunnelId}`)
+
+    // Track both open and close events
+    const eventPromise = new Promise<{ code: number; reason: string }>((resolve, reject) => {
+      client.on('open', () => {
+        // Connection opened, wait for close
+      })
+      client.on('close', (code, reason) => {
+        resolve({ code, reason: reason.toString() })
+      })
+      client.on('error', (error) => {
+        reject(error)
+      })
+    })
+
+    // Wait for the connection to be closed
+    const closeEvent = await eventPromise
+
+    expect(closeEvent.code).toBe(4008)
+    expect(closeEvent.reason).toBe('No upstream available')
+  })
 })

cloudflare-tunnel/src/tunnel.ts

@@ -82,6 +82,17 @@ export class Tunnel {
       this.ctx.acceptWebSocket(server)
       server.serializeAttachment({ role: 'up' } satisfies Attachment)
     } else {
+      // Check if upstream exists before accepting downstream
+      const existingUp = this.getUpstream()
+      if (!existingUp) {
+        // Accept the WebSocket connection
+        server.accept()
+        // Immediately close it
+        server.close(4008, 'No upstream available')
+        // Return successful WebSocket upgrade response
+        // The client will receive the close event after connection
+        return addCors(new Response(null, { status: 101, webSocket: client }))
+      }
       // Accept with hibernation and tag as downstream
       this.ctx.acceptWebSocket(server)
       server.serializeAttachment({ role: 'down' } satisfies Attachment)

cloudflare-tunnel/wrangler.jsonc

@@ -12,20 +12,24 @@
   "routes": [
     {
       "pattern": "holocron.so/_tunnel/*",
-      "zone_name": "holocron.so"
+      "zone_name": "holocron.so",
     },
     {
       "pattern": "fumabase.com/_tunnel/*",
-      "zone_name": "fumabase.com"
+      "zone_name": "fumabase.com",
     },
     {
       "pattern": "unframer.co/_tunnel/*",
-      "zone_name": "unframer.co"
-    }
+      "zone_name": "unframer.co",
+    },
+    {
+      "pattern": "mcp.unframer.co/_tunnel/*",
+      "zone_name": "unframer.co",
+    },
   ],
 
   "durable_objects": {
-    "bindings": [{ "name": "TUNNEL_DO", "class_name": "Tunnel" }]
+    "bindings": [{ "name": "TUNNEL_DO", "class_name": "Tunnel" }],
   },
 
   "migrations": [{ "tag": "v1", "new_classes": ["Tunnel"] }],
@@ -35,16 +39,24 @@
       "routes": [
         {
           "pattern": "preview.fumabase.com/_tunnel/*",
-          "zone_name": "fumabase.com"
+          "zone_name": "fumabase.com",
         },
         {
           "pattern": "preview.holocron.so/_tunnel/*",
-          "zone_name": "holocron.so"
-        }
+          "zone_name": "holocron.so",
+        },
+        {
+          "pattern": "mcp.preview.unframer.co/_tunnel/*",
+          "zone_name": "unframer.co",
+        },
+        {
+          "pattern": "preview.unframer.co/_tunnel/*",
+          "zone_name": "unframer.co",
+        },
       ],
       "durable_objects": {
-        "bindings": [{ "name": "TUNNEL_DO", "class_name": "Tunnel" }]
-      }
-    }
-  }
+        "bindings": [{ "name": "TUNNEL_DO", "class_name": "Tunnel" }],
+      },
+    },
+  },
 }

docs-website/src/lib/docs-state.tsx

@@ -25,6 +25,8 @@ export type ChatHistory = {
 
 export type DrawerState = 'closed' | 'minimized' | 'open'
 
+export type WebSocketConnectionState = 'CONNECTED' | 'UPSTREAM_NOT_CONNECTED' | 'NOT_CONNECTED'
+
 export type PersistentDocsState = {
   chatId: string
   drawerState: DrawerState
@@ -34,7 +36,7 @@ export type PersistentDocsState = {
 export type DocsState = {
   // tree?: PageTree.Root
   toc?: TOCItemType[]
-  websocketServerPreviewConnected?: boolean
+  websocketConnectionState?: WebSocketConnectionState
   currentSlug?: string
   // docsJson?: DocsJsonType
   filesInDraft: FilesInDraft
@@ -54,6 +56,7 @@ const defaultState: DocsState = {
   filesInDraft: {},
   deletedPages: [],
   previewMode: 'preview',
+  websocketConnectionState: 'NOT_CONNECTED',
 }
 
 const defaultPersistentState: PersistentDocsState = {

docs-website/src/routes/_catchall-client.tsx

@@ -64,7 +64,7 @@ const openapiPath = `/api-reference`
 
 const allowedOrigins = [env.NEXT_PUBLIC_URL!.replace(/\/$/, ''), 'http://localhost:7664']
 
-let onFirstStateMessage = () => {}
+let onFirstStateMessage = () => { }
 const firstStateReceived = new Promise<void>((resolve) => {
   onFirstStateMessage = resolve
 })
@@ -170,43 +170,80 @@ async function websocketIdHandling(websocketId: string) {
   if (globalThis.websocketHandlingDone) return
   globalThis.websocketHandlingDone = true
 
-  console.log('connecting over preview websocketId', websocketId)
-  const websocketUrl = `wss://${WEBSITE_DOMAIN}/_tunnel/client?id=${websocketId}`
-  const ws = new WebSocket(websocketUrl)
-  ws.onopen = () => {
-    useDocsState.setState({
-      websocketServerPreviewConnected: true,
-    })
-    ws.send(JSON.stringify({ type: 'ready' }))
-  }
-  ws.onclose = () => {
-    useDocsState.setState({
-      websocketServerPreviewConnected: false,
-    })
+  let ws: WebSocket | null = null
+  let reconnectTimeout: NodeJS.Timeout | null = null
+  let pingInterval: NodeJS.Timeout | null = null
+
+  const clearTimers = () => {
+    if (reconnectTimeout) {
+      clearTimeout(reconnectTimeout)
+      reconnectTimeout = null
+    }
+    if (pingInterval) {
+      clearInterval(pingInterval)
+      pingInterval = null
+    }
   }
-  ws.onmessage = async (event) => {
-    let data: IframeRpcMessage
-    try {
-      data = JSON.parse(event.data)
-    } catch {
-      console.error(`websocket sent invalid json`, event.data)
-      return
+
+  const connect = () => {
+    console.log('connecting over preview websocketId', websocketId)
+    const websocketUrl = `wss://${WEBSITE_DOMAIN}/_tunnel/downstream?id=${websocketId}`
+    ws = new WebSocket(websocketUrl)
+
+    ws.onopen = () => {
+      clearTimers()
+      useDocsState.setState({
+        websocketConnectionState: 'CONNECTED',
+      })
+      ws!.send(JSON.stringify({ type: 'ready' }))
+
+      // Start ping interval
+      pingInterval = setInterval(() => {
+        if (ws && ws.readyState === WebSocket.OPEN) {
+          ws.send(JSON.stringify({ type: 'ping' }))
+        }
+      }, 5 * 1000)
     }
-    const { id, revalidate, state: partialState } = data || {}
-    if (partialState) {
-      await setDocsStateForMessage(partialState)
+
+    ws.onclose = (event) => {
+      clearTimers()
+
+      // Check if upstream is not connected based on close code
+      if (event.code === 4008) {
+        useDocsState.setState({
+          websocketConnectionState: 'UPSTREAM_NOT_CONNECTED',
+        })
+      } else {
+        useDocsState.setState({
+          websocketConnectionState: 'NOT_CONNECTED',
+        })
+      }
+
+      // Always attempt to reconnect after 5 seconds
+      reconnectTimeout = setTimeout(() => connect(), 5000)
     }
-    if (revalidate) {
-      await revalidator?.revalidate()
+
+    ws.onmessage = async (event) => {
+      let data: IframeRpcMessage
+      try {
+        data = JSON.parse(event.data)
+      } catch {
+        console.error(`websocket sent invalid json`, event.data)
+        return
+      }
+      const { id, revalidate, state: partialState } = data || {}
+      if (partialState) {
+        await setDocsStateForMessage(partialState)
+      }
+      if (revalidate) {
+        await revalidator?.revalidate()
+      }
+      ws!.send(JSON.stringify({ id } satisfies IframeRpcMessage))
     }
-    ws.send(JSON.stringify({ id } satisfies IframeRpcMessage))
   }
-  // ping interval
-  setInterval(() => {
-    if (ws.readyState === WebSocket.OPEN) {
-      ws.send(JSON.stringify({ type: 'ping' }))
-    }
-  }, 1000)
+
+  // Start initial connection
+  connect()
 }
 
 if (typeof window !== 'undefined') {
@@ -293,11 +330,11 @@ export function CSSVariables({ docsJson }: { docsJson: DocsJsonType }) {
   const toCssBlock = (obj: Record<string, string> | undefined) =>
     obj
       ? Object.entries(obj)
-          .map(([key, value]) => {
-            const cssVar = key.startsWith('--') ? key : `--${key}`
-            return `${cssVar}: ${value} !important;`
-          })
-          .join('\n  ')
+        .map(([key, value]) => {
+          const cssVar = key.startsWith('--') ? key : `--${key}`
+          return `${cssVar}: ${value} !important;`
+        })
+        .join('\n  ')
       : ''
 
   // Don't render if both empty
@@ -552,7 +589,7 @@ function DocsLayoutWrapper({ children, docsJson }: { children: React.ReactNode;
 }
 
 const noop = (callback: () => void) => {
-  return () => {}
+  return () => { }
 }
 
 function PreviewBanner({ websocketId }: { websocketId?: string }) {
@@ -564,7 +601,7 @@ function PreviewBanner({ websocketId }: { websocketId?: string }) {
     globalNavigate(url.pathname + url.search)
   }
 
-  const websocketServerPreviewConnected = useDocsState((state) => state.websocketServerPreviewConnected)
+  const websocketConnectionState = useDocsState((state) => state.websocketConnectionState)
 
   const shouldShow = useSyncExternalStore(
     noop,
@@ -575,19 +612,41 @@ function PreviewBanner({ websocketId }: { websocketId?: string }) {
     return null
   }
 
+  const getStatusColor = () => {
+    switch (websocketConnectionState) {
+      case 'CONNECTED':
+        return 'bg-green-500'
+      case 'UPSTREAM_NOT_CONNECTED':
+        return 'bg-yellow-500'
+      case 'NOT_CONNECTED':
+      default:
+        return 'bg-red-500'
+    }
+  }
+
+  const getStatusMessage = () => {
+    switch (websocketConnectionState) {
+      case 'CONNECTED':
+        return 'Connected to local preview. Added content will be highlighted green'
+      case 'UPSTREAM_NOT_CONNECTED':
+        return 'Waiting for upstream server. Please run the preview server'
+      case 'NOT_CONNECTED':
+      default:
+        return 'Server disconnected. Please restart the preview server'
+    }
+  }
+
   return (
     <Banner className='sticky top-0 z-50 bg-fd-muted text-fd-accent-foreground isolate px-4 py-1 flex items-center justify-between'>
       <div className='flex items-center gap-2'>
         <div
           className={cn(
             'w-2 h-2 rounded-full animate-pulse',
-            websocketServerPreviewConnected ? 'bg-green-500' : 'bg-red-500',
+            getStatusColor(),
           )}
         ></div>
         <span className='font-medium text-sm'>
-          {websocketServerPreviewConnected
-            ? 'Connected to local preview. Added content will be highlighted green'
-            : 'Server disconnected. Please restart the preview server'}
+          {getStatusMessage()}
         </span>
       </div>
       <button