diff --git a/electron-agent/agent/TunnelAgent.ts b/electron-agent/agent/TunnelAgent.ts index 16ffde1..e0f8d3a 100644 --- a/electron-agent/agent/TunnelAgent.ts +++ b/electron-agent/agent/TunnelAgent.ts @@ -18,13 +18,6 @@ interface RelayRequest body: string; // base64 } -interface RelayResponse -{ - reqId: string; - status: number; - headers: Record; - body: string; // base64 -} export class TunnelAgent { @@ -84,6 +77,8 @@ export class TunnelAgent ); } + private _send( msg: object ): void { this.ws?.send( JSON.stringify( msg ) ); } + private _forward( data: Buffer ): void { let req: RelayRequest; @@ -100,40 +95,32 @@ export class TunnelAgent headers: { ...req.headers, 'content-length': String( bodyBuf.length ) }, }; - const chunks: Buffer[] = []; - const localReq = http.request( opts, localRes => { - localRes.on( 'data', ( chunk: Buffer ) => chunks.push( Buffer.from( chunk ) ) ); - localRes.on( 'end', () => + const respHeaders: Record = {}; + for ( const [ k, v ] of Object.entries( localRes.headers ) ) { - const respHeaders: Record = {}; - for ( const [ k, v ] of Object.entries( localRes.headers ) ) - { - if ( typeof v === 'string' ) respHeaders[ k ] = v; - else if ( Array.isArray( v ) ) respHeaders[ k ] = v.join( ', ' ); - } + if ( typeof v === 'string' ) respHeaders[ k ] = v; + else if ( Array.isArray( v ) ) respHeaders[ k ] = v.join( ', ' ); + } - const resp: RelayResponse = { - reqId: req.reqId, - status: localRes.statusCode ?? 200, - headers: respHeaders, - body: Buffer.concat( chunks ).toString( 'base64' ), - }; - this.ws?.send( JSON.stringify( resp ) ); - } ); + // Send headers immediately — starts streaming + this._send( { type: 'res_start', reqId: req.reqId, status: localRes.statusCode ?? 200, headers: respHeaders } ); + + localRes.on( 'data', ( chunk: Buffer ) => + this._send( { type: 'res_data', reqId: req.reqId, chunk: Buffer.from( chunk ).toString( 'base64' ) } ) + ); + + localRes.on( 'end', () => + this._send( { type: 'res_end', reqId: req.reqId } ) + ); } ); localReq.on( 'error', err => { - const errResp: RelayResponse = { - reqId: req.reqId, - status: 502, - headers: { 'content-type': 'application/json' }, - body: Buffer.from( JSON.stringify( { error: 'Local service error', detail: err.message } ) ) - .toString( 'base64' ), - }; - this.ws?.send( JSON.stringify( errResp ) ); + this._send( { type: 'res_start', reqId: req.reqId, status: 502, headers: { 'content-type': 'application/json' } } ); + this._send( { type: 'res_data', reqId: req.reqId, chunk: Buffer.from( JSON.stringify( { error: err.message } ) ).toString( 'base64' ) } ); + this._send( { type: 'res_end', reqId: req.reqId } ); } ); if ( bodyBuf.length > 0 ) localReq.write( bodyBuf ); diff --git a/source/server/relay/pending.ts b/source/server/relay/pending.ts index 15fc13b..1554840 100644 --- a/source/server/relay/pending.ts +++ b/source/server/relay/pending.ts @@ -1,27 +1,63 @@ export interface RelayResponse { - reqId: string; - status: number; + reqId: string; + status: number; headers: Record; - body: string; // base64-encoded + body: string; // base64 — legacy single-shot } -const callbacks = new Map void>(); +export type AgentMsg = + | { type: 'res_start'; reqId: string; status: number; headers: Record } + | { type: 'res_data'; reqId: string; chunk: string } // base64 + | { type: 'res_end'; reqId: string } + | RelayResponse; // legacy: no type field -export function addPending( reqId: string, cb: ( r: RelayResponse ) => void ): void +export interface StreamHandler { - callbacks.set( reqId, cb ); + onStart: ( status: number, headers: Record ) => void; + onData: ( chunk: Buffer ) => void; + onEnd: () => void; } -export function resolvePending( resp: RelayResponse ): void +const pending = new Map(); + +export function addPending( reqId: string, handler: StreamHandler ): void { - const cb = callbacks.get( resp.reqId ); - if ( !cb ) return; - callbacks.delete( resp.reqId ); - cb( resp ); + pending.set( reqId, handler ); } export function removePending( reqId: string ): void { - callbacks.delete( reqId ); + pending.delete( reqId ); +} + +export function dispatchAgentMsg( msg: AgentMsg ): void +{ + const handler = pending.get( msg.reqId ); + if ( !handler ) return; + + if ( 'type' in msg ) + { + if ( msg.type === 'res_start' ) + { + handler.onStart( msg.status, msg.headers ); + } + else if ( msg.type === 'res_data' ) + { + handler.onData( Buffer.from( msg.chunk, 'base64' ) ); + } + else if ( msg.type === 'res_end' ) + { + pending.delete( msg.reqId ); + handler.onEnd(); + } + } + else + { + // Legacy single-shot RelayResponse — forward as stream + handler.onStart( msg.status, msg.headers ); + handler.onData( Buffer.from( msg.body, 'base64' ) ); + pending.delete( msg.reqId ); + handler.onEnd(); + } } diff --git a/source/server/routes/agent.ts b/source/server/routes/agent.ts index 9ba2a38..3fd482c 100644 --- a/source/server/routes/agent.ts +++ b/source/server/routes/agent.ts @@ -4,7 +4,7 @@ import jwt from 'jsonwebtoken'; import { AuthPayload } from '../middleware/requireAuth'; import { getTunnel } from '../db'; import { registry } from '../relay/TunnelRegistry'; -import { resolvePending } from '../relay/pending'; +import { dispatchAgentMsg } from '../relay/pending'; export function handleAgentUpgrade( req: IncomingMessage, ws: WebSocket ): void { @@ -36,7 +36,7 @@ export function handleAgentUpgrade( req: IncomingMessage, ws: WebSocket ): void ws.on( 'message', ( data ) => { - try { resolvePending( JSON.parse( data.toString() ) ); } + try { dispatchAgentMsg( JSON.parse( data.toString() ) ); } catch { /* ignore malformed messages */ } } ); diff --git a/source/server/routes/proxy.ts b/source/server/routes/proxy.ts index 54c7dc7..4b7620d 100644 --- a/source/server/routes/proxy.ts +++ b/source/server/routes/proxy.ts @@ -4,15 +4,15 @@ import crypto from 'crypto'; import { getTunnel, TunnelConfig } from '../db'; import { registry } from '../relay/TunnelRegistry'; import { AuthPayload, extractBearer } from '../middleware/requireAuth'; -import { addPending, removePending, RelayResponse } from '../relay/pending'; +import { addPending, removePending } from '../relay/pending'; export interface RelayRequest { - reqId: string; - method: string; - path: string; + reqId: string; + method: string; + path: string; headers: Record; - body: string; // base64-encoded + body: string; // base64-encoded } const router = Router(); @@ -48,10 +48,10 @@ router.all( [ '/:tunnelId', '/:tunnelId/*' ], async ( req, res ) => if ( !ws ) { res.status( 503 ).json( { error: 'Tunnel agent not connected' } ); return; } // Strip the /:tunnelId prefix to get the path to forward - const prefix = '/' + req.params.tunnelId; + const prefix = '/' + req.params.tunnelId; const forwardPath = req.url.slice( prefix.length ) || '/'; - // Buffer the raw request body (no JSON middleware on /t routes) + // Buffer the raw request body const chunks: Buffer[] = []; await new Promise( ( resolve, reject ) => { @@ -74,39 +74,41 @@ router.all( [ '/:tunnelId', '/:tunnelId/*' ], async ( req, res ) => const relayReq: RelayRequest = { reqId, - method: req.method, - path: forwardPath, + method: req.method, + path: forwardPath, headers, - body: bodyBuf.toString( 'base64' ), + body: bodyBuf.toString( 'base64' ), }; - const TIMEOUT_MS = 30_000; - const response = await new Promise( ( resolve ) => + // 2-minute timeout for the first response chunk — reset once streaming starts + await new Promise( ( resolve ) => { const timer = setTimeout( () => { removePending( reqId ); - resolve( null ); - }, TIMEOUT_MS ); + if ( !res.headersSent ) res.status( 504 ).json( { error: 'Agent timed out' } ); + resolve(); + }, 120_000 ); - addPending( reqId, ( resp ) => + addPending( reqId, { - clearTimeout( timer ); - resolve( resp ); + onStart: ( status, respHeaders ) => + { + clearTimeout( timer ); // streaming started — no more timeout + res.status( status ); + for ( const [ key, val ] of Object.entries( respHeaders ) ) + { + if ( key.toLowerCase() === 'transfer-encoding' ) continue; + res.setHeader( key, val ); + } + res.flushHeaders(); + }, + onData: ( chunk ) => res.write( chunk ), + onEnd: () => { res.end(); resolve(); }, } ); ws.send( JSON.stringify( relayReq ) ); } ); - - if ( !response ) { res.status( 504 ).json( { error: 'Agent timed out' } ); return; } - - res.status( response.status ); - for ( const [ key, val ] of Object.entries( response.headers ) ) - { - if ( key.toLowerCase() === 'transfer-encoding' ) continue; - res.setHeader( key, val ); - } - res.end( Buffer.from( response.body, 'base64' ) ); } ); export default router;