From c9452337e14611b782c80ad1588c68b0da5288c2 Mon Sep 17 00:00:00 2001 From: Anso Date: Fri, 8 May 2026 09:02:59 -0400 Subject: [PATCH] test(pilot): in-process simulation of mid-tunnel disconnects (#987) * test(pilot): in-process simulation of mid-tunnel disconnects The PR #979 hardening added per-stream idle timers, drain handling, paused-request maps, and a tcpAwaitingDrain set on the bridge. None of those cleanup paths had an end-to-end test driving them through a real disconnect. A future refactor that inverts the order of clearIdleTimer + streams.delete (or that misses one of the aux maps) would leak per-tunnel memory in a way unit tests cannot catch. Spin up a real http.Server with attachUpgrade and a real ws.WebSocket client (same self-contained pattern as pilot-tunnel-integration.test.ts). Four cases: - HTTP request mid-flight: open a real loopback HTTP request, wait for the bridge to forward http_req to the agent, terminate the agent WS, assert the loopback request resolves with 502 (the teardownStream path) and the manager marks the node Offline. - TCP stream with bytes outstanding: openTcpStream from the bridge, ack open from the test agent, write some bytes, terminate the tunnel, assert the TcpStream emits 'close' and the bridge is gone from the manager. - Reconnect with the long-lived token after a clean close does NOT bump tunnels_replaced (no live tunnel to replace). - Reconnect WHILE the prior tunnel is still live DOES bump tunnels_replaced. Inverse case proves the counter wiring is correct in both directions. Self-contained: no shared helper modules, no dependencies on other in-flight pilot test files. Can land independently of PRs #982, #983, or #985. * test(pilot): address PR E code-review on disconnect tests Two small findings: - Defensive: attach a no-op 'error' listener on the TcpStream in test #2 after the 'open' event has fired. teardownStream only emits 'error' for unaccepted streams today, so the listener is not exercised; the listener exists so a future refactor that delays accepted=true past 'open' cannot crash the worker on an unhandled 'error' event. - Removed a stale cross-file reference in the header comment that pointed at pilot-tunnel-integration.test.ts; that file lives in a separate in-flight branch and does not exist on main. No behavior change. All 4 disconnect tests pass. * chore(test): drop unused PROTOCOL_VERSION import The import was leftover from an earlier draft and ESLint flagged it as no-unused-vars (error level for symbols not prefixed with _), failing the lint job in CI. --- .../__tests__/pilot-tunnel-disconnect.test.ts | 339 ++++++++++++++++++ 1 file changed, 339 insertions(+) create mode 100644 backend/src/__tests__/pilot-tunnel-disconnect.test.ts diff --git a/backend/src/__tests__/pilot-tunnel-disconnect.test.ts b/backend/src/__tests__/pilot-tunnel-disconnect.test.ts new file mode 100644 index 00000000..be70deb3 --- /dev/null +++ b/backend/src/__tests__/pilot-tunnel-disconnect.test.ts @@ -0,0 +1,339 @@ +/** + * In-process simulation of mid-tunnel disconnects. + * + * The PR #979 hardening pass added per-stream idle timers, drain handling, + * paused-request maps, and a tcpAwaitingDrain set on the bridge. None of + * those cleanup paths had an end-to-end test that drives them through a + * real disconnect; a future refactor that inverts the order of + * clearIdleTimer + streams.delete (or that misses one of the aux maps) + * could leak per-tunnel memory in a way the unit tests cannot catch. + * + * This file spins up a real http.Server with attachUpgrade and a real + * ws.WebSocket client, then exercises three disconnect scenarios plus a + * reconnect-counter case. Self-contained so it can land independently of + * any other in-flight pilot work. + */ +import { describe, it, expect, beforeAll, afterAll, afterEach, vi } from 'vitest'; +import http from 'http'; +import crypto from 'crypto'; +import jwt from 'jsonwebtoken'; +import { WebSocket, WebSocketServer } from 'ws'; +import { setupTestDb, cleanupTestDb } from './helpers/setupTestDb'; +import { attachUpgrade } from '../websocket/upgradeHandler'; +import { + BinaryFrameType, + decodeBinaryFrame, + decodeJsonFrame, + encodeJsonFrame, + type JsonFrame, +} from '../pilot/protocol'; +import { PilotTunnelManager } from '../services/PilotTunnelManager'; + +let tmpDir: string; +let DatabaseService: typeof import('../services/DatabaseService').DatabaseService; + +let server: http.Server; +let port: number; +let pilotTunnelWss: WebSocketServer; +let mainWss: WebSocketServer; +let nodeId: number; + +beforeAll(async () => { + tmpDir = await setupTestDb(); + ({ DatabaseService } = await import('../services/DatabaseService')); + + server = http.createServer(); + mainWss = new WebSocketServer({ noServer: true }); + pilotTunnelWss = new WebSocketServer({ noServer: true }); + attachUpgrade(server, { wss: mainWss, pilotTunnelWss }); + + await new Promise((resolve, reject) => { + server.once('error', reject); + server.listen(0, '127.0.0.1', () => { + const addr = server.address(); + if (!addr || typeof addr === 'string') { + reject(new Error('listen returned unexpected address')); + return; + } + port = addr.port; + resolve(); + }); + }); + + nodeId = DatabaseService.getInstance().addNode({ + name: `pilot-disconnect-${Date.now()}`, + type: 'remote', + mode: 'pilot_agent', + compose_dir: '/tmp/x', + is_default: false, + api_url: '', + api_token: '', + }); +}); + +afterAll(async () => { + const mgr = PilotTunnelManager.getInstance(); + mgr.closeTunnel(nodeId); + mgr.removeAllListeners('tunnel-up'); + mgr.removeAllListeners('tunnel-down'); + pilotTunnelWss.close(); + mainWss.close(); + await new Promise((resolve) => server.close(() => resolve())); + cleanupTestDb(tmpDir); +}); + +afterEach(() => { + PilotTunnelManager.getInstance().closeTunnel(nodeId); +}); + +function mintEnrollToken(): string { + const db = DatabaseService.getInstance(); + const jwtSecret = db.getGlobalSettings().auth_jwt_secret; + if (!jwtSecret) throw new Error('test DB has no auth_jwt_secret'); + const ttlSeconds = 15 * 60; + const token = jwt.sign( + { scope: 'pilot_enroll', nodeId, enrollNonce: crypto.randomUUID() }, + jwtSecret, + { expiresIn: ttlSeconds }, + ); + const hash = crypto.createHash('sha256').update(token).digest('hex'); + db.createPilotEnrollment(nodeId, hash, Date.now() + ttlSeconds * 1000); + return token; +} + +/** + * Open a ws client that decodes both JSON and binary frames, with a queue + * + waiter so tests can await individual frames in arrival order. Mirrors + * pilot-tunnel-integration.test.ts. + */ +async function openTunnel(token: string): Promise<{ + ws: WebSocket; + nextJsonFrame: () => Promise; + onBinaryFrame: (cb: (type: BinaryFrameType, streamId: number, payload: Buffer) => void) => void; + closed: Promise; +}> { + const ws = new WebSocket(`ws://127.0.0.1:${port}/api/pilot/tunnel`, { + headers: { + Authorization: `Bearer ${token}`, + 'x-sencho-agent-version': 'disconnect-test/1.0', + }, + }); + + const queue: JsonFrame[] = []; + const waiters: Array<(frame: JsonFrame) => void> = []; + let binaryHandler: ((t: BinaryFrameType, s: number, p: Buffer) => void) | null = null; + + ws.on('message', (data, isBinary) => { + try { + if (isBinary) { + const buf = Buffer.isBuffer(data) ? data : Buffer.from(data as ArrayBuffer); + const frame = decodeBinaryFrame(buf); + binaryHandler?.(frame.type, frame.streamId, frame.payload); + } else { + const frame = decodeJsonFrame(data.toString()); + const waiter = waiters.shift(); + if (waiter) waiter(frame); + else queue.push(frame); + } + } catch { + // ignore malformed frames in tests + } + }); + + const closed = new Promise((resolve) => ws.once('close', () => resolve())); + + await new Promise((resolve, reject) => { + ws.once('open', () => resolve()); + ws.once('error', reject); + }); + + return { + ws, + nextJsonFrame: () => + new Promise((resolve, reject) => { + if (queue.length) { resolve(queue.shift()!); return; } + const t = setTimeout(() => reject(new Error('timed out waiting for frame')), 5000); + waiters.push((f) => { clearTimeout(t); resolve(f); }); + }), + onBinaryFrame: (cb) => { binaryHandler = cb; }, + closed, + }; +} + +async function enrollAndCaptureLongLivedToken(): Promise { + const enrollToken = mintEnrollToken(); + const tunnel = await openTunnel(enrollToken); + await tunnel.nextJsonFrame(); // hello + const ack = await tunnel.nextJsonFrame(); + if (ack.t !== 'ctrl' || ack.op !== 'enroll_ack' || typeof ack.payload?.token !== 'string') { + throw new Error('expected enroll_ack with token'); + } + const longLivedToken = ack.payload.token; + tunnel.ws.close(); + await tunnel.closed; + await vi.waitFor( + () => expect(PilotTunnelManager.getInstance().hasActiveTunnel(nodeId)).toBe(false), + { timeout: 2000 }, + ); + return longLivedToken; +} + +describe('pilot tunnel mid-disconnect cleanup (in-process integration)', () => { + it('completes a loopback HTTP request with 502 when the tunnel dies mid-flight', async () => { + const enrollToken = mintEnrollToken(); + const tunnel = await openTunnel(enrollToken); + await tunnel.nextJsonFrame(); // hello + await tunnel.nextJsonFrame(); // enroll_ack + + const loopbackUrl = PilotTunnelManager.getInstance().getLoopbackUrl(nodeId); + if (!loopbackUrl) throw new Error('loopback URL missing'); + + // Issue a request against the loopback. The bridge sends an http_req + // frame to the agent (us). We deliberately do NOT respond; the + // request stays parked in the bridge's stream map. + const url = new URL(loopbackUrl); + const responsePromise = new Promise((resolve, reject) => { + const req = http.request({ + host: url.hostname, + port: Number(url.port), + method: 'GET', + path: '/api/health', + }, (res) => { + resolve(res.statusCode || 0); + res.resume(); + }); + req.on('error', reject); + req.end(); + }); + + // Wait for the http_req frame to arrive on the agent side, confirming + // the bridge has a live stream. Then kill the tunnel from the agent + // side, simulating an agent crash. + await tunnel.nextJsonFrame(); // http_req + tunnel.ws.terminate(); + await tunnel.closed; + + // The parked loopback request must resolve (not hang) with 502 from + // the bridge's teardownStream path, and the manager must drop the + // tunnel so the next caller sees offline. + const status = await responsePromise; + expect(status).toBe(502); + + await vi.waitFor( + () => expect(PilotTunnelManager.getInstance().hasActiveTunnel(nodeId)).toBe(false), + { timeout: 2000 }, + ); + + // updateNodeStatus('offline') ran via the bridge.once('closed') hook. + const dbNode = DatabaseService.getInstance().getNode(nodeId); + expect(dbNode?.status).toBe('offline'); + }, 10_000); + + it('closes a TCP stream cleanly when the tunnel dies with bytes outstanding', async () => { + const enrollToken = mintEnrollToken(); + const tunnel = await openTunnel(enrollToken); + await tunnel.nextJsonFrame(); // hello + await tunnel.nextJsonFrame(); // enroll_ack + + const bridge = PilotTunnelManager.getInstance().getBridge(nodeId); + if (!bridge) throw new Error('bridge missing'); + + // openTcpStream sends a tcp_open frame to the agent. We capture it, + // then ack with tcp_open_ack { ok: true } so the stream transitions + // to accepted. Without the ack the bridge keeps the stream parked + // and write() returns false silently. + const streamHandle = bridge.openTcpStream({ stack: 'demo', service: 'svc', port: 80 }); + if (!streamHandle) throw new Error('openTcpStream returned null'); + + const tcpOpen = await tunnel.nextJsonFrame(); + if (tcpOpen.t !== 'tcp_open') throw new Error('expected tcp_open'); + tunnel.ws.send(encodeJsonFrame({ t: 'tcp_open_ack', s: tcpOpen.s, ok: true })); + + // Wait for the handle's 'open' event to fire so write() will actually + // serialize bytes through the bridge. + await new Promise((resolve) => streamHandle.once('open', () => resolve())); + + // Defensive: 'error' fires from teardownStream only when the stream + // is not yet accepted, which is not the case here. Attach a no-op + // listener so a future refactor that delays accepted=true past 'open' + // does not crash the worker on an unhandled 'error' event. + streamHandle.on('error', () => { /* defensive */ }); + + // Listen for the stream's 'close' event before yanking the tunnel. + const closedEvent = new Promise((resolve) => streamHandle.once('close', () => resolve())); + + // Push some bytes, then kill the tunnel from the agent side. The + // bridge's teardownStream path must emit 'close' on the TcpStream + // and clear the streams + tcpAwaitingDrain maps. + streamHandle.write(Buffer.from('hello, world')); + tunnel.ws.terminate(); + + await closedEvent; + await vi.waitFor( + () => expect(PilotTunnelManager.getInstance().hasActiveTunnel(nodeId)).toBe(false), + { timeout: 2000 }, + ); + + // openTcpStream after the tunnel is gone must return null (not crash). + const afterClose = PilotTunnelManager.getInstance().getBridge(nodeId); + expect(afterClose).toBeNull(); + }, 10_000); + + it('reconnect with the long-lived token does NOT bump tunnels_replaced', async () => { + const longLivedToken = await enrollAndCaptureLongLivedToken(); + const before = PilotTunnelManager.getInstance().getMetricsSnapshot().counters; + + const tunnel = await openTunnel(longLivedToken); + await tunnel.nextJsonFrame(); // hello + + await vi.waitFor( + () => expect(PilotTunnelManager.getInstance().hasActiveTunnel(nodeId)).toBe(true), + { timeout: 2000 }, + ); + + const after = PilotTunnelManager.getInstance().getMetricsSnapshot().counters; + + // tunnels_total bumps on every successful registration. + expect(after.tunnels_total).toBe(before.tunnels_total + 1); + // tunnels_replaced should NOT bump: the prior bridge was closed cleanly + // and removed from the manager before this connect, so this is a + // fresh registration, not a replacement of a still-live tunnel. + expect(after.tunnels_replaced).toBe(before.tunnels_replaced); + + tunnel.ws.close(); + await tunnel.closed; + }, 10_000); + + it('reconnecting BEFORE the prior tunnel is removed DOES bump tunnels_replaced', async () => { + // Inverse of the previous test: prove the counter-bump path is wired + // by forcing a true split-brain. Open one tunnel, do NOT close it, + // then open a second tunnel for the same node. The manager closes + // the first bridge to make room, which is the documented "replaced" + // path. + const longLivedToken = await enrollAndCaptureLongLivedToken(); + + const first = await openTunnel(longLivedToken); + await first.nextJsonFrame(); // hello + await vi.waitFor( + () => expect(PilotTunnelManager.getInstance().hasActiveTunnel(nodeId)).toBe(true), + { timeout: 2000 }, + ); + + const before = PilotTunnelManager.getInstance().getMetricsSnapshot().counters; + + const second = await openTunnel(longLivedToken); + await second.nextJsonFrame(); // hello + + const after = PilotTunnelManager.getInstance().getMetricsSnapshot().counters; + expect(after.tunnels_replaced).toBe(before.tunnels_replaced + 1); + expect(after.tunnels_total).toBe(before.tunnels_total + 1); + + // The first tunnel's WS should have been closed by the manager + // (PilotCloseCode.Replaced = 4000). Do not assert on the close code + // here because the test runner may see the close before the code + // surfaces; just await the close to keep cleanup tidy. + await first.closed; + second.ws.close(); + await second.closed; + }, 10_000); +});