fix(egress): retry a gateway CONNECT once on connection failure and log each failed attempt (#9351)

* fix(egress): retry a gateway CONNECT once on connection failure and log each failed attempt

* test(egress): cover the 5xx CONNECT retry path

Author: waleedlatif1

apps/sim/lib/core/network/gateway.server.test.ts

@@ -45,6 +45,8 @@ let originRequests = 0
 let receivedHeaders: Record<string, unknown> = {}
 let receivedTarget = ''
 const proxiedPorts = new Set<number>()
+/** Each queued entry makes one CONNECT fail: `reset` drops the connection, a number answers with that status. */
+const connectFailures: Array<'reset' | number> = []
 const secureOrigin = httpsServer(async (req, res) => {
   res.setHeader('x-via-proxy', proxiedPorts.has(req.socket.remotePort ?? 0) ? 'yes' : 'no')
   originRequests++
@@ -70,6 +72,15 @@ proxy.on('connect', (req, socket, head) => {
     destination: req.url,
     sni: tlsSocket.servername,
   })
+  const failure = connectFailures.shift()
+  if (failure === 'reset') {
+    socket.destroy()
+    return
+  }
+  if (failure !== undefined) {
+    socket.end(`HTTP/1.1 ${failure} Gateway Failure\r\nContent-Length: 0\r\n\r\n`)
+    return
+  }
   if (!req.headers['proxy-authorization']?.match(/^Bearer (alpha|bravo|a{48})$/)) {
     socket.end('HTTP/1.1 407 Proxy Authentication Required\r\nContent-Length: 0\r\n\r\n')
     return
@@ -150,6 +161,11 @@ afterAll(async () => {
   rmSync(certificateDirectory, { recursive: true, force: true })
 })
 const url = () => 'https://origin.invalid'
+const gatewayFor = (organizationId: string): OutboundGateway => {
+  const gateway = state.gateways.get(organizationId)
+  if (!gateway) throw new Error(`No gateway for ${organizationId}`)
+  return gateway
+}
 const options = { profile: 'selfHostedService' as const }
 
 describe('organization gateways over real TLS CONNECT sockets', () => {
@@ -356,4 +372,42 @@ describe('organization gateways over real TLS CONNECT sockets', () => {
       await transport.dispatcher.destroy()
     }
   })
+
+  it.each(['reset', 503] as const)(
+    'retries a CONNECT that fails with %s before the tunnel opens and completes the request',
+    async (failure) => {
+      connectFailures.push(failure)
+      const start = admissions.length
+      try {
+        const response = await runWithOutboundOrganization('org_a', () =>
+          secureFetchWithPinnedIP(url(), '1.1.1.1', options)
+        )
+        expect(await response.text()).toBe('tls reached')
+        expect(admissions.length).toBe(start + 2)
+      } finally {
+        connectFailures.length = 0
+      }
+    }
+  )
+
+  it('does not retry a rejected proxy credential', async () => {
+    const start = admissions.length
+    await expect(
+      openGatewayTunnel({ ...gatewayFor('org_a'), token: 'denied' }, '1.1.1.1', 443)
+    ).rejects.toMatchObject({ code: 'GATEWAY_UNAVAILABLE' })
+    expect(admissions.length).toBe(start + 1)
+  })
+
+  it('gives up after two failed CONNECT attempts', async () => {
+    connectFailures.push('reset', 503)
+    const start = admissions.length
+    try {
+      await expect(openGatewayTunnel(gatewayFor('org_a'), '1.1.1.1', 443)).rejects.toMatchObject({
+        code: 'GATEWAY_UNAVAILABLE',
+      })
+      expect(admissions.length).toBe(start + 2)
+    } finally {
+      connectFailures.length = 0
+    }
+  })
 })

apps/sim/lib/core/network/gateway.server.ts

@@ -1,13 +1,18 @@
 import { isIP, Socket } from 'node:net'
 import { checkServerIdentity, connect as connectTls } from 'node:tls'
+import { createLogger } from '@sim/logger'
 import { evaluateAddress, STRICT_EGRESS_POLICY } from '@sim/security/egress'
+import { toError } from '@sim/utils/errors'
 import { Agent, Client } from 'undici/index.js'
 import type { OutboundGateway } from '@/lib/core/network/config.server'
 import { OutboundRoutingError } from '@/lib/core/network/routing'
 import type { EgressProfile } from '@/lib/core/security/egress/profiles'
 import { checkResolvedEgress, validateEgressUrl } from '@/lib/core/security/egress/validate'
 
 const CONNECT_TIMEOUT_MS = 10_000
+const MAX_CONNECT_ATTEMPTS = 2
+
+const logger = createLogger('OutboundGateway')
 
 /** The proxy CA never changes trust for the upstream service. */
 export async function secureOutboundTunnel(
@@ -64,7 +69,36 @@ export async function openGatewayTunnel(
   ) {
     throw new OutboundRoutingError('UNSUPPORTED_TRANSPORT')
   }
-  const authority = `${address}:443`
+  for (let attempt = 1; ; attempt++) {
+    const result = await attemptGatewayConnect(gateway, `${address}:443`)
+    if (result.socket) return result.socket
+    logger.warn('Gateway CONNECT failed', {
+      gatewayId: gateway.id,
+      organizationId: gateway.organizationId,
+      attempt,
+      statusCode: result.statusCode,
+      errorName: result.error?.name,
+      errorCode: getErrorCode(result.error),
+    })
+    if (!result.retryable || attempt >= MAX_CONNECT_ATTEMPTS) {
+      throw new OutboundRoutingError('GATEWAY_UNAVAILABLE')
+    }
+  }
+}
+
+type GatewayConnectResult =
+  | { socket: Socket; statusCode?: never; error?: never; retryable?: never }
+  | { socket?: never; statusCode?: number; error?: Error; retryable: boolean }
+
+/**
+ * A failure before the proxy answers 200 has sent no upstream bytes, so it is safe to retry.
+ * Credential or admission rejections (4xx) are not transient and never retry.
+ */
+async function attemptGatewayConnect(
+  gateway: OutboundGateway,
+  authority: string
+): Promise<GatewayConnectResult> {
+  /** A fresh client per attempt opens a new connection, which the NLB can route to another task. */
   const client = new Client(gateway.url, {
     allowH2: false,
     maxHeaderSize: 16_384,
@@ -77,19 +111,22 @@ export async function openGatewayTunnel(
       signal: AbortSignal.timeout(CONNECT_TIMEOUT_MS),
       headers: { host: authority, 'proxy-authorization': `Bearer ${gateway.token}` },
     })
-    if (statusCode !== 200 || !(socket instanceof Socket)) {
-      socket.destroy()
-      throw new OutboundRoutingError('GATEWAY_UNAVAILABLE')
-    }
-    return socket
-  } catch {
-    throw new OutboundRoutingError('GATEWAY_UNAVAILABLE')
+    if (statusCode === 200 && socket instanceof Socket) return { socket }
+    socket.destroy()
+    return { statusCode, retryable: statusCode >= 500 }
+  } catch (error) {
+    return { error: toError(error), retryable: true }
   } finally {
     /** An upgraded socket is detached from the dispatcher and now belongs to its caller. */
     await client.destroy()
   }
 }
 
+function getErrorCode(error: Error | undefined): string | undefined {
+  const code = error && 'code' in error ? error.code : undefined
+  return typeof code === 'string' ? code : undefined
+}
+
 /** A dispatcher belongs to one organization and gateway, never a process-wide default. */
 export function createGatewayDispatcher(
   gateway: OutboundGateway,