@aikidosec/broker-client 1.0.5 → 1.0.7

package/README.md

@@ -60,6 +60,7 @@ Resource IDs are displayed in the Aikido UI when you register them.
 - `HTTP_PROXY` - Proxy server for HTTP requests (e.g., `http://proxy.company.local:8080`)
 - `HTTPS_PROXY` - Proxy server for HTTPS requests (e.g., `http://proxy.company.local:8080`)
 - `ALL_PROXY` - Universal proxy fallback for all protocols if protocol-specific proxy is not set
+- `BROKER_TARGET_URL` - Override the broker server URL (defaults to `https://broker.aikidobroker.com`)
 
 ## How It Works
 

package/app/client.js

@@ -10,23 +10,19 @@ import axios from 'axios';
 import { URL } from 'url';
 import { Address4, Address6 } from 'ip-address';
 import dns from 'native-dns';
-import fs from 'fs';
 import { ResourceManager } from './resourceManager.js';
 import { HttpsProxyAgent } from 'https-proxy-agent';
-
-// Configure logging
-const log = {
-  info: (msg) => console.log(`[INFO] ${new Date().toISOString()} - ${msg}`),
-  warn: (msg) => console.warn(`[WARN] ${new Date().toISOString()} - ${msg}`),
-  error: (msg) => console.error(`[ERROR] ${new Date().toISOString()} - ${msg}`),
-  debug: (msg) => console.log(`[DEBUG] ${new Date().toISOString()} - ${msg}`)
-};
+import { getClientId, setClientIdCache, getServerUrl, getClientSecret } from './config.js';
+import { STREAMING_THRESHOLD } from './streaming/state.js';
+import { registerStreamingHandlers } from './streaming/handlers.js';
+import { initStreamingResponse } from './streaming/initStreamingResponse.js';
+import { startStaleStreamCleanup, cleanupAllStreams } from './streaming/cleanup.js';
+import { log } from './log.js';
 
 // Broker Server Configuration
-const SERVER_URL = "https://broker.aikidobroker.com";
+const CLIENT_SECRET = getClientSecret();
+const SERVER_URL = getServerUrl();
 
-// Client Configuration (from environment)
-const CLIENT_SECRET = process.env.CLIENT_SECRET;
 const ALLOWED_SUBNETS = process.env.ALLOWED_INTERNAL_SUBNETS
   ? process.env.ALLOWED_INTERNAL_SUBNETS.split(',').map(s => s.trim()).filter(s => s)
   : [];
@@ -39,7 +35,7 @@ const DNS_SERVERS = process.env.DNS_SERVERS
 // Configure axios defaults
 const MAX_RESPONSE_SIZE = 100 * 1024 * 1024; // 100 MB
 const axiosConfig = {
-  timeout: 30000,
+  timeout: 300000,
   maxRedirects: 5,
   maxContentLength: MAX_RESPONSE_SIZE,
   maxBodyLength: MAX_RESPONSE_SIZE
@@ -51,31 +47,6 @@ const internalHttpClient = axios.create(axiosConfig);
 // Initialize ResourceManager
 const resourceManager = new ResourceManager();
 
-// Cache for client_id
-let _clientIdCache = null;
-
-/**
- * Get the client ID from cache or read from file.
- * Returns null if not registered yet.
- */
-function getClientId() {
-  if (_clientIdCache !== null) {
-    return _clientIdCache;
-  }
-  
-  const clientIdPath = '/config/client_id';
-  if (fs.existsSync(clientIdPath)) {
-    try {
-      _clientIdCache = fs.readFileSync(clientIdPath, 'utf8').trim();
-      return _clientIdCache;
-    } catch (e) {
-      log.error(`Failed to read client_id: ${e.message}`);
-    }
-  }
-  
-  return null;
-}
-
 /**
  * Resolve hostname using custom DNS servers if configured.
  * Falls back to system DNS if DNS_SERVERS not set.
@@ -133,6 +104,32 @@ async function resolveInternalHostname(hostname) {
   });
 }
 
+/**
+ * Get Content-Length via HEAD request to determine if streaming should be used.
+ * @param {string} url - URL to check
+ * @param {object} headers - Headers to send with the request
+ * @returns {Promise<number|null>} Content-Length in bytes, or null if unavailable
+ */
+async function getContentLengthViaHead(url, headers) {
+  try {
+    const headResponse = await internalHttpClient.head(url, {
+      headers,
+      validateStatus: () => true,
+      timeout: 5000
+    });
+    // Headers can be any case (Content-Length, content-length, CONTENT-LENGTH)
+    const contentLengthHeader = Object.keys(headResponse.headers)
+      .find(key => key.toLowerCase() === 'content-length');
+    const contentLength = contentLengthHeader 
+      ? parseInt(headResponse.headers[contentLengthHeader]) 
+      : NaN;
+    return isNaN(contentLength) ? null : contentLength;
+  } catch (e) {
+    log.warn(`HEAD request failed, will use buffered approach: ${e.message}`);
+    return null;
+  }
+}
+
 /**
  * Check if URL points to internal resource
  */
@@ -236,7 +233,9 @@ const socket = io(SERVER_URL, {
   randomizationFactor: 0.5,
   tryAllTransports: true, // if we don't, it won't try to fallback from websocket to polling
   autoConnect: false,  // Don't connect until after registration
-  withCredentials: true // make sure cookies work for sticky sessions
+  withCredentials: true, // make sure cookies work for sticky sessions
+  // Increase timeouts for heavy streaming workloads
+  pingTimeout: 60000,   // 60s (default 20s) - time to wait for pong before considering connection dead
 });
 
 // Socket.IO event handlers
@@ -256,8 +255,10 @@ socket.on('connect', async () => {
   }
 });
 
-socket.on('disconnect', () => {
-  log.warn("Disconnected from broker server");
+socket.on('disconnect', (reason) => {
+  log.warn(`Disconnected from broker server: ${reason}`);
+  // Clean up all active streams - they can't recover after reconnect
+  cleanupAllStreams();
 });
 
 socket.on('connect_error', (error) => {
@@ -284,6 +285,12 @@ socket.io.on('reconnect_failed', () => {
   log.error(`Socket.IO reconnection failed after all attempts`);
 });
 
+// Register streaming handlers (handle next chunk and abort stream on client disconnect)
+registerStreamingHandlers(socket);
+
+// Start cleanup interval for stale streams
+startStaleStreamCleanup();
+
 socket.on('forward_request', async (data, callback) => {
   /**
    * Receive request from broker server and forward to internal resource
@@ -335,29 +342,54 @@ socket.on('forward_request', async (data, callback) => {
       }
     }
     
-    // Forward the request to the internal resource
+    // Check Content-Length first with HEAD request to decide streaming vs buffered
+    let useStreaming = false;
+    let contentLength = null;
+    
+    if (method === 'GET') {
+      contentLength = await getContentLengthViaHead(resolvedUrl, headers);
+      if (contentLength !== null && contentLength > STREAMING_THRESHOLD) {
+        useStreaming = true;
+        log.info(`Large response detected (${(contentLength / (1024 * 1024)).toFixed(2)} MB) - using streaming`);
+      }
+    }
+    
+    // Make the request with appropriate response type
     const response = await internalHttpClient.request({
       method,
       url: resolvedUrl,
       headers,
       data: body,
       validateStatus: () => true, // Accept any status code
-      responseType: 'arraybuffer', // Get raw bytes, don't parse JSON
+      responseType: useStreaming ? 'stream' : 'arraybuffer',
     });
     
-    log.info(`Successfully forwarded request ${requestId} to ${targetUrl}, status: ${response.status}`);
-    
-    // Return response via acknowledgement
-    // Send body as base64 to preserve binary data byte-for-byte (critical for Docker registry digests)
-    const responseBody = response.data ? Buffer.from(response.data).toString('base64') : null;
-    
-    callback({
-      request_id: requestId,
-      status_code: response.status,
-      headers: response.headers,
-      body: responseBody,
-      version: 2
-    });
+    if (useStreaming) {
+      // Initialize streaming - server will pull chunks via get_next_chunk
+      initStreamingResponse({
+        requestId,
+        statusCode: response.status,
+        headers: response.headers,
+        stream: response.data,
+        callback
+      });
+    } else {
+      const responseSizeBytes = response.data ? response.data.length : 0;
+      const responseSizeMB = (responseSizeBytes / (1024 * 1024)).toFixed(2);
+      
+      log.info(`Successfully forwarded request ${requestId} to ${targetUrl}, status: ${response.status}, response size: ${responseSizeMB} MB`);
+      
+      // Send direct response
+      sendDirectResponse(
+        requestId,
+        response.status,
+        response.headers,
+        response.data,
+        callback
+      );
+      
+      log.info(`Response sent for request ${requestId}`);
+    }
     
   } catch (error) {
     log.error(`Error forwarding request ${requestId} to ${targetUrl}: ${error?.response?.status || error.message}`);
@@ -372,6 +404,25 @@ socket.on('forward_request', async (data, callback) => {
   }
 });
 
+/**
+ * Send a direct response (for small files under streaming threshold)
+ * @param {string} requestId - Request ID for tracking
+ * @param {number} statusCode - HTTP status code
+ * @param {object} headers - Response headers
+ * @param {Buffer|null} responseData - Raw response data (before base64)
+ * @param {function} callback - Socket.IO callback
+ */
+function sendDirectResponse(requestId, statusCode, headers, responseData, callback) {
+  const responseBody = responseData ? responseData.toString('base64') : null;
+  callback({
+    request_id: requestId,
+    status_code: statusCode,
+    headers: headers,
+    body: responseBody,
+    version: 2
+  });
+}
+
 function formatMessageBody(message) {
   return Buffer.from(message, 'utf-8').toString('base64');
 }
@@ -402,14 +453,7 @@ async function registerWithServer() {
         log.info(`✓ Successfully registered with broker server as ${clientId}`);
         
         // Save client_id to file and cache
-        _clientIdCache = clientId;
-        try {
-          fs.mkdirSync('/config', { recursive: true });
-          fs.writeFileSync('/config/client_id', clientId);
-          log.info("💾 Saved client_id to /config/client_id");
-        } catch (e) {
-          log.warn(`Could not save client_id file: ${e.message}`);
-        }
+        setClientIdCache(clientId);
 
         log.info("Waiting for server to propagate configuration...");
         await new Promise(resolve => setTimeout(resolve, 5000));
@@ -417,17 +461,7 @@ async function registerWithServer() {
       } else if (response.status === 409) {
         log.info("✓ Client already registered with server (this is OK)");
         // Try to extract client_id from response if available
-        try {
-          const clientId = response.data.client_id;
-          if (clientId) {
-            _clientIdCache = clientId;
-            fs.mkdirSync('/config', { recursive: true });
-            fs.writeFileSync('/config/client_id', clientId);
-            log.info("💾 Saved client_id to /config/client_id");
-          }
-        } catch (e) {
-          log.warn(`Could not save client_id file: ${e.message}`);
-        }
+        setClientIdCache(response.data.client_id);
         return;
       } else {
         log.warn(`Registration attempt ${attempt + 1} failed: ${response.status} - ${response.data}`);

package/app/config.js

@@ -0,0 +1,65 @@
+import fs from 'fs';
+
+// Client ID file path
+const CLIENT_ID_PATH = '/config/client_id';
+
+// Cache for client_id
+let _clientIdCache = null;
+
+// Client secret from environment
+const CLIENT_SECRET = process.env.CLIENT_SECRET;
+
+/**
+ * Get the server URL.
+ * Uses BROKER_TARGET_URL env var, defaults to https://broker.aikidobroker.com
+ */
+export function getServerUrl() {
+  return process.env.BROKER_TARGET_URL || 'https://broker.aikidobroker.com';
+}
+
+/**
+ * Get the client secret from environment.
+ */
+export function getClientSecret() {
+  return CLIENT_SECRET;
+}
+
+/**
+ * Get the client ID from cache or read from file.
+ * Returns null if not registered yet.
+ */
+export function getClientId() {
+  if (_clientIdCache !== null) {
+    return _clientIdCache;
+  }
+  
+  if (fs.existsSync(CLIENT_ID_PATH)) {
+    try {
+      _clientIdCache = fs.readFileSync(CLIENT_ID_PATH, 'utf8').trim();
+      return _clientIdCache;
+    } catch (e) {
+      console.error(`[ERROR] ${new Date().toISOString()} - Failed to read client_id: ${e.message}`);
+    }
+  }
+  
+  return null;
+}
+
+/**
+ * Set the client ID cache and persist to file.
+ * Does nothing if clientId is null/undefined.
+ */
+export function setClientIdCache(clientId) {
+  if (!clientId) {
+    return;
+  }
+  _clientIdCache = clientId;
+  try {
+    fs.mkdirSync('/config', { recursive: true });
+    fs.writeFileSync(CLIENT_ID_PATH, clientId);
+    console.log(`[INFO] ${new Date().toISOString()} - 💾 Saved client_id to ${CLIENT_ID_PATH}`);
+  } catch (e) {
+    console.error(`[ERROR] ${new Date().toISOString()} - Failed to write client_id: ${e.message}`);
+  }
+}
+

package/app/log.js

@@ -0,0 +1,10 @@
+/**
+ * Logging utility for streaming module
+ */
+
+export const log = {
+  info: (msg) => console.log(`[INFO] ${new Date().toISOString()} - ${msg}`),
+  warn: (msg) => console.warn(`[WARN] ${new Date().toISOString()} - ${msg}`),
+  error: (msg) => console.error(`[ERROR] ${new Date().toISOString()} - ${msg}`),
+  debug: (msg) => console.log(`[DEBUG] ${new Date().toISOString()} - ${msg}`)
+};

package/app/streaming/cleanup.js

@@ -0,0 +1,76 @@
+/**
+ * Stream cleanup utilities - manages cleanup of streaming state
+ */
+
+import { activeStreams, recalculateTotalBufferSize } from './state.js';
+import { log } from '../log.js';
+import { resumePausedStreams } from './flowControl.js';
+
+// Stale stream timeout (3 minutes)
+const STALE_STREAM_TIMEOUT_MS = 3 * 60 * 1000;
+const STALE_STREAM_INTERVAL_MS = 2 * 60 * 1000; // Check every 2 minutes
+
+// Cleanup interval reference
+let cleanupInterval = null;
+
+/**
+ * Start the stale stream cleanup interval
+ * Runs every 2 minutes to check for streams inactive for 3+ minutes
+ */
+export function startStaleStreamCleanup() {
+  if (cleanupInterval) return; // Already running
+  
+  cleanupInterval = setInterval(() => {
+    const now = Date.now();
+    
+    for (const [requestId, state] of activeStreams.entries()) {
+      const inactiveMs = now - state.lastActivity;
+      if (inactiveMs > STALE_STREAM_TIMEOUT_MS) {
+        try {
+          log.warn(`Cleaning up stale stream ${requestId} (inactive for ${(inactiveMs / 1000 / 60).toFixed(1)} minutes)`);
+          cleanupStream(requestId);
+        } catch (err) {
+          log.error(`Error cleaning up stale stream ${requestId}: ${err.message}`);
+        }
+      }
+    }
+
+  }, STALE_STREAM_INTERVAL_MS);
+  
+  log.info('Started stale stream cleanup interval');
+}
+
+/**
+ * Clean up a streaming request (e.g., on error or disconnect)
+ * @param {string} requestId - Request ID to clean up
+ */
+export function cleanupStream(requestId) {
+  const state = activeStreams.get(requestId);
+  if (state) {
+    // state.stream is the Node.js Readable stream from the axios HTTP response
+    // Destroying it releases the TCP connection and stops buffering data
+    state.stream.destroy();
+    activeStreams.delete(requestId);
+    
+    // Recalculate total buffer size to fix any drift
+    recalculateTotalBufferSize();
+    
+    // Resume any paused streams now that buffer space is freed
+    resumePausedStreams();
+  }
+}
+
+/**
+ * Clean up all active streams (e.g., on socket disconnect)
+ * Streams cannot recover after socket reconnect since server-side state is lost
+ */
+export function cleanupAllStreams() {
+  const count = activeStreams.size;
+  if (count === 0) return;
+  
+  for (const requestId of activeStreams.keys()) {
+    cleanupStream(requestId);
+  }
+  
+  log.info(`Cleaned up all ${count} streams after disconnect`);
+}

package/app/streaming/flowControl.js

@@ -0,0 +1,39 @@
+/**
+ * Stream flow control - manages pausing/resuming streams based on buffer limits
+ */
+
+import { activeStreams, totalBufferSize, GLOBAL_MAX_BUFFER_SIZE, PER_STREAM_MAX_BUFFER_SIZE } from './state.js';
+import { log } from '../log.js';
+
+/**
+ * Pause a stream if global or per-stream buffer limits are exceeded
+ * @param {string} requestId - Request ID
+ * @param {Readable} stream - The readable stream to pause
+ */
+export function pauseStreamIfOverLimit(requestId, stream) {
+  const state = activeStreams.get(requestId);
+  if (!state || state.paused) return;
+  
+  const shouldPause = totalBufferSize() >= GLOBAL_MAX_BUFFER_SIZE || state.buffer.length >= PER_STREAM_MAX_BUFFER_SIZE;
+  if (shouldPause) {
+    state.paused = true;
+    stream.pause();
+    log.info(`Stream paused for ${requestId}, stream buffer: ${(state.buffer.length / (1024 * 1024)).toFixed(2)} MB, global: ${(totalBufferSize() / (1024 * 1024)).toFixed(2)} MB`);
+  }
+}
+
+/**
+ * Resume any paused streams if both global and per-stream buffers are below thresholds
+ * Called after cleanup frees up buffer space
+ */
+export function resumePausedStreams() {
+  if (totalBufferSize() >= GLOBAL_MAX_BUFFER_SIZE) return;
+  
+  for (const [requestId, state] of activeStreams.entries()) {
+    if (state.paused && !state.complete && state.buffer.length < PER_STREAM_MAX_BUFFER_SIZE) {
+      state.paused = false;
+      state.stream.resume();
+      log.info(`Stream resumed for ${requestId} after cleanup, stream buffer: ${(state.buffer.length / (1024 * 1024)).toFixed(2)} MB, global: ${(totalBufferSize() / (1024 * 1024)).toFixed(2)} MB`);
+    }
+  }
+}

package/app/streaming/handlers.js

@@ -0,0 +1,150 @@
+/**
+ * Socket.IO event handlers for streaming
+ */
+
+import { activeStreams, totalBufferSize, subtractFromTotalBuffer, STREAM_CHUNK_SIZE, GLOBAL_MAX_BUFFER_SIZE, PER_STREAM_MAX_BUFFER_SIZE } from './state.js';
+import { cleanupStream } from './cleanup.js';
+import { log } from '../log.js';
+
+/**
+ * Register streaming event handlers on a Socket.IO socket
+ * @param {Socket} socket - Socket.IO socket instance
+ */
+export function registerStreamingHandlers(socket) {
+  socket.on('get_next_chunk', handleGetNextChunk);
+  socket.on('abort_stream', handleAbortStream);
+}
+
+/**
+ * Handle get_next_chunk request from server (pull-based streaming)
+ * Waits for data if buffer is empty but stream isn't complete
+ * @param {object} data - Request data containing request_id
+ * @param {function} callback - Socket.IO callback
+ */
+export async function handleGetNextChunk(data, callback) {
+  const requestId = data.request_id;
+  const state = activeStreams.get(requestId);
+  
+  if (!state) {
+    log.warn(`get_next_chunk: No active stream for request ${requestId}`);
+    callback({ error: 'no_stream', request_id: requestId });
+    return;
+  }
+  
+  state.lastActivity = Date.now();
+  
+  // Check for stream error
+  if (state.error) {
+    cleanupStreamWithError(requestId, state.error, callback);
+    return;
+  }
+  
+  // Wait for data if buffer doesn't have enough for a full chunk
+  await waitForData(state);
+  
+  // Check for error again after waiting
+  if (state.error) {
+    cleanupStreamWithError(requestId, state.error, callback);
+    return;
+  }
+  
+  // Extract chunk from buffer
+  const { chunkData, isComplete } = extractChunkFromBuffer(state, requestId);
+  
+  state.totalBytesSent += chunkData.length;
+  state.chunkIndex++;
+  
+  log.info(`get_next_chunk ${state.chunkIndex} for ${requestId}: ${chunkData.length} bytes, complete=${isComplete}, total=${(state.totalBytesSent / (1024 * 1024)).toFixed(2)} MB`);
+  
+  // Clean up if complete
+  if (isComplete) {
+    activeStreams.delete(requestId);
+    log.info(`Streaming complete for ${requestId}: ${state.chunkIndex} chunks, ${(state.totalBytesSent / (1024 * 1024)).toFixed(2)} MB total`);
+  }
+  
+  callback({
+    request_id: requestId,
+    data: chunkData.toString('base64'),
+    complete: isComplete,
+    chunk_index: state.chunkIndex
+  });
+}
+
+/**
+ * Wait for buffer to have enough data or stream to complete
+ */
+async function waitForData(state) {
+  const maxWaitMs = 60000;
+  const checkIntervalMs = 100;
+  let waited = 0;
+  
+  const hasEnoughData = () => state.buffer.length >= STREAM_CHUNK_SIZE;
+  const isStreamDone = () => state.complete || state.error;
+  const hasTimedOut = () => waited >= maxWaitMs;
+  
+  while (!hasEnoughData() && !isStreamDone() && !hasTimedOut()) {
+    await new Promise(resolve => setTimeout(resolve, checkIntervalMs));
+    waited += checkIntervalMs;
+  }
+}
+
+/**
+ * Extract a chunk from the buffer and manage buffer state
+ */
+function extractChunkFromBuffer(state, requestId) {
+  let chunkData;
+  let isComplete = false;
+  
+  if (state.buffer.length >= STREAM_CHUNK_SIZE) {
+    // Have enough data for a full chunk
+    chunkData = state.buffer.slice(0, STREAM_CHUNK_SIZE);
+    state.buffer = state.buffer.slice(STREAM_CHUNK_SIZE);
+    subtractFromTotalBuffer(chunkData.length);
+    
+    // Resume stream if it was paused and both limits are now satisfied
+    resumeStreamIfBelowBufferLimits(state);
+  } else if (state.complete) {
+    // Stream is done, send remaining buffer
+    chunkData = state.buffer;
+    subtractFromTotalBuffer(chunkData.length);
+    state.buffer = Buffer.alloc(0);
+    isComplete = true;
+  } else {
+    // Timeout waiting for data - send what we have
+    log.warn(`get_next_chunk: Timeout waiting for data for ${requestId}, sending ${state.buffer.length} bytes`);
+    chunkData = state.buffer;
+    subtractFromTotalBuffer(chunkData.length);
+    state.buffer = Buffer.alloc(0);
+  }
+  
+  return { chunkData, isComplete };
+}
+
+/**
+ * Resume a paused stream if buffer limits allow
+ */
+function resumeStreamIfBelowBufferLimits(state) {
+  const canResume = totalBufferSize() < GLOBAL_MAX_BUFFER_SIZE && state.buffer.length < PER_STREAM_MAX_BUFFER_SIZE;
+  if (state.paused && canResume) {
+    state.paused = false;
+    state.stream.resume();
+  }
+}
+
+/**
+ * Clean up stream and send error callback
+ */
+function cleanupStreamWithError(requestId, error, callback) {
+  activeStreams.delete(requestId);
+  callback({ error, request_id: requestId });
+}
+
+/**
+ * Handle abort_stream request from server (caller (aikido service) disconnected)
+ * @param {object} data - Request data containing request_id
+ */
+function handleAbortStream(data) {
+  const requestId = data.request_id;
+  log.warn(`Received abort_stream for ${requestId} - caller disconnected`);
+  cleanupStream(requestId);
+}

package/app/streaming/initStreamingResponse.js

@@ -0,0 +1,84 @@
+/**
+ * Initialize streaming response - stores stream state for pull-based chunk retrieval
+ */
+
+import { activeStreams, addToTotalBuffer } from './state.js';
+import { log } from '../log.js';
+import { pauseStreamIfOverLimit } from './flowControl.js';
+
+/**
+ * Initialize streaming response - stores stream state for pull-based chunk retrieval
+ * Server will call get_next_chunk to pull chunks
+ * @param {object} options
+ * @param {string} options.requestId - Request ID for tracking
+ * @param {number} options.statusCode - HTTP status code
+ * @param {object} options.headers - Response headers
+ * @param {ReadableStream} options.stream - Response stream from axios
+ * @param {function} options.callback - Socket.IO callback for initial response
+ */
+export function initStreamingResponse({ requestId, statusCode, headers, stream, callback }) {
+  log.info(`Initializing streaming response for request ${requestId}`);
+  
+  // Store stream state for pull-based retrieval
+  const streamState = {
+    stream,
+    buffer: Buffer.alloc(0),
+    complete: false,
+    error: null,
+    totalBytesSent: 0,
+    chunkIndex: 0,
+    paused: false,
+    lastActivity: Date.now()
+  };
+  
+  activeStreams.set(requestId, streamState);
+  
+  // Buffer data as it arrives from the internal resource
+  // Use bounded buffering with global and per-stream limits to control memory usage
+  stream.on('data', (chunk) => {
+    // Yield to event loop periodically to allow ping/pong processing
+    // This prevents stream data from starving Socket.IO heartbeats
+    setImmediate(() => {
+      const state = activeStreams.get(requestId);
+      if (!state) {
+        // Stream was cleaned up, ignore late data
+        return;
+      }
+      
+      state.buffer = Buffer.concat([state.buffer, chunk]);
+      addToTotalBuffer(chunk.length);
+      
+      // Pause stream if buffer limits exceeded (backpressure)
+      pauseStreamIfOverLimit(requestId, stream);
+    });
+  });
+  
+  stream.on('end', () => {
+    const state = activeStreams.get(requestId);
+    if (state) {
+      state.complete = true;
+      log.info(`Stream ended for request ${requestId}, ${(state.buffer.length / (1024 * 1024)).toFixed(2)} MB remaining in buffer`);
+    }
+  });
+  
+  stream.on('error', (error) => {
+    const state = activeStreams.get(requestId);
+    if (state) {
+      state.error = error.message;
+      state.complete = true;
+      log.error(`Stream error for request ${requestId}: ${error.message}`);
+    }
+  });
+  
+  // Send initial response - server will start pulling chunks
+  callback({
+    request_id: requestId,
+    status_code: statusCode,
+    headers: headers,
+    body: null,
+    version: 2,
+    streaming: true
+  });
+  
+  log.info(`Streaming initialized for request ${requestId}, server will pull chunks`);
+}

package/app/streaming/state.js

@@ -0,0 +1,82 @@
+/**
+ * Shared state for streaming - buffer tracking and active streams
+ */
+
+// Chunk size for streaming responses (10MB before base64 encoding)
+export const STREAM_CHUNK_SIZE = 10 * 1024 * 1024;
+
+// Global maximum buffer size across ALL streams
+export const GLOBAL_MAX_BUFFER_SIZE = STREAM_CHUNK_SIZE * 30; // 300MB
+
+// Per-stream maximum buffer size
+export const PER_STREAM_MAX_BUFFER_SIZE = STREAM_CHUNK_SIZE * 1.5; // 15MB
+
+// Threshold for using streaming vs direct response
+export const STREAMING_THRESHOLD = STREAM_CHUNK_SIZE;
+
+// Active streams for pull-based streaming (server pulls chunks from client)
+// Key: request_id, Value: { stream, buffer, complete, error, totalBytesSent, chunkIndex, paused, lastActivity }
+export const activeStreams = new Map();
+
+// Track total buffer size across all streams
+let _totalBufferSize = 0;
+
+/**
+ * Get the current total buffer size across all streams
+ * @returns {number} Total bytes buffered
+ */
+export function totalBufferSize() {
+  return _totalBufferSize;
+}
+
+/**
+ * Set the total buffer size (used for reset)
+ * @param {number} value - New value
+ */
+export function setTotalBufferSize(value) {
+  _totalBufferSize = value;
+}
+
+/**
+ * Add to the total buffer size
+ * @param {number} bytes - Bytes to add
+ */
+export function addToTotalBuffer(bytes) {
+  _totalBufferSize += bytes;
+}
+
+/**
+ * Subtract from the total buffer size
+ * @param {number} bytes - Bytes to subtract
+ */
+export function subtractFromTotalBuffer(bytes) {
+  _totalBufferSize -= bytes;
+}
+
+/**
+ * Recalculate total buffer size from all active streams
+ * Used to fix any drift in the counter
+ * @returns {number} Actual total bytes buffered
+ */
+export function recalculateTotalBufferSize() {
+  let actual = 0;
+  for (const [, state] of activeStreams.entries()) {
+    if (state.buffer) {
+      actual += state.buffer.length;
+    }
+  }
+  
+  if (actual !== _totalBufferSize) {
+    _totalBufferSize = actual;
+  }
+  
+  return actual;
+}
+
+/**
+ * Get count of active streams (for monitoring)
+ * @returns {number} Number of active streams
+ */
+export function getActiveStreamCount() {
+  return activeStreams.size;
+}

package/app/streamingHandler.js

@@ -0,0 +1,15 @@
+/**
+ * Streaming Handler - handles large file streaming with pull-based chunk retrieval
+ * Server pulls chunks via get_next_chunk calls
+ * 
+ * Uses bounded buffering to limit memory usage - pauses stream when buffer is full,
+ * resumes when buffer is drained.
+ * 
+ * This file re-exports from the streaming/ module for backwards compatibility.
+ */
+
+export { STREAM_CHUNK_SIZE, STREAMING_THRESHOLD, totalBufferSize as getTotalBufferSize, getActiveStreamCount } from './streaming/state.js';
+export { registerStreamingHandlers } from './streaming/handlers.js';
+export { initStreamingResponse } from './streaming/initStreamingResponse.js';
+export { startStaleStreamCleanup, cleanupStream, cleanupAllStreams } from './streaming/cleanup.js';
+

package/package.json

@@ -1,12 +1,14 @@
 {
   "name": "@aikidosec/broker-client",
-  "version": "1.0.5",
+  "version": "1.0.7",
   "description": "Aikido Broker Client - Runs in customer network to forward requests to internal resources",
   "main": "app/client.js",
   "type": "module",
   "scripts": {
     "start": "node app/client.js",
-    "dev": "node --watch app/client.js"
+    "dev": "node --watch app/client.js",
+    "test": "node --experimental-vm-modules node_modules/jest/bin/jest.js",
+    "test:watch": "node --experimental-vm-modules node_modules/jest/bin/jest.js --watch"
   },
   "keywords": [
     "broker",
@@ -38,9 +40,19 @@
     "native-dns": "0.7.0",
     "socket.io-client": "4.8.1"
   },
+  "devDependencies": {
+    "jest": "30.2.0"
+  },
   "files": [
     "app/",
     "README.md",
     "LICENSE"
-  ]
+  ],
+  "jest": {
+    "testEnvironment": "node",
+    "testMatch": [
+      "**/tests/**/*.test.js"
+    ],
+    "transform": {}
+  }
 }