refactor(api-proxy): deduplicate guard enforcement between HTTP and WebSocket paths, fix 3 missing WebSocket guards

containers/api-proxy/guards/common-guard-checks.js

@@ -0,0 +1,141 @@
+'use strict';
+
+/**
+ * Builds the array of guard descriptors shared by the HTTP and WebSocket proxy
+ * paths.  Centralising the list here ensures that every guard added to the HTTP
+ * path is automatically enforced for WebSocket connections too.
+ *
+ * Each descriptor has:
+ *   - block          — current block-state returned by the guard's getter
+ *   - isBlocked      — predicate: returns true when the request should be halted
+ *   - statusCode     — HTTP status code to send when blocked
+ *   - eventName      — event name used for log / telemetry
+ *   - buildError     — function(block) → error body object (JSON-serialisable)
+ *   - buildLogFields — function(block) → extra log-fields object
+ *
+ * Model-specific guards (model_multiplier_cap, retired_model,
+ * unknown_model_ai_credits) are only included when `model` is truthy (non-null
+ * and non-empty string).  For WebSocket upgrade requests there is no JSON
+ * request body, so callers should pass null and the model guards are skipped.
+ *
+ * @param {object} deps - Guard state getter and error-builder functions.
+ * @param {Function} deps.getEffectiveTokenBlockState
+ * @param {Function} deps.buildEffectiveTokenLimitError
+ * @param {Function} deps.getMaxRunsBlockState
+ * @param {Function} deps.buildMaxRunsExceededError
+ * @param {Function} deps.getPermissionDeniedBlockState
+ * @param {Function} deps.buildPermissionDeniedLimitError
+ * @param {Function} deps.getAiCreditsBlockState
+ * @param {Function} deps.buildAiCreditsLimitError
+ * @param {Function} deps.getModelMultiplierCapBlockState
+ * @param {Function} deps.buildModelMultiplierCapError
+ * @param {Function} deps.getRetiredModelBlockState
+ * @param {Function} deps.buildRetiredModelError
+ * @param {Function} deps.checkUnknownModelRejection
+ * @param {string|null} model - Model name extracted from the request, or null
+ *   to skip model-specific guards.
+ * @returns {Array<object>} Array of guard descriptor objects.
+ */
+function buildCommonGuardChecks(deps, model) {
+  const {
+    getEffectiveTokenBlockState,
+    buildEffectiveTokenLimitError,
+    getMaxRunsBlockState,
+    buildMaxRunsExceededError,
+    getPermissionDeniedBlockState,
+    buildPermissionDeniedLimitError,
+    getAiCreditsBlockState,
+    buildAiCreditsLimitError,
+    getModelMultiplierCapBlockState,
+    buildModelMultiplierCapError,
+    getRetiredModelBlockState,
+    buildRetiredModelError,
+    checkUnknownModelRejection,
+  } = deps;
+
+  return [
+    {
+      block: getEffectiveTokenBlockState(),
+      isBlocked: block => block && block.maxExceeded,
+      statusCode: 429,
+      eventName: 'effective_tokens_limit_exceeded',
+      buildError: buildEffectiveTokenLimitError,
+      buildLogFields: block => ({
+        total_effective_tokens: block.totalEffectiveTokens,
+        max_effective_tokens: block.maxEffectiveTokens,
+      }),
+    },
+    {
+      block: getMaxRunsBlockState(),
+      isBlocked: block => block && block.maxExceeded,
+      statusCode: 429,
+      eventName: 'max_runs_exceeded',
+      buildError: buildMaxRunsExceededError,
+      buildLogFields: block => ({
+        invocation_count: block.invocationCount,
+        max_runs: block.maxRuns,
+      }),
+    },
+    {
+      block: getPermissionDeniedBlockState(),
+      isBlocked: block => block && block.maxExceeded,
+      statusCode: 403,
+      eventName: 'permission_denied_limit_exceeded',
+      buildError: buildPermissionDeniedLimitError,
+      buildLogFields: block => ({
+        denied_count: block.deniedCount,
+        max_permission_denied: block.maxPermissionDenied,
+      }),
+    },
+    {
+      block: getAiCreditsBlockState(),
+      isBlocked: block => block && block.maxExceeded,
+      statusCode: 429,
+      eventName: 'ai_credits_limit_exceeded',
+      buildError: buildAiCreditsLimitError,
+      buildLogFields: block => ({
+        total_ai_credits: block.totalAiCredits,
+        max_ai_credits: block.maxAiCredits,
+        hard_cap: block.hardCap === true,
+      }),
+    },
+    // Model-specific guards — only active when a model was identified in the request.
+    ...(model ? [
+      {
+        block: getModelMultiplierCapBlockState(model),
+        isBlocked: block => !!block,
+        statusCode: 400,
+        eventName: 'model_multiplier_cap_exceeded',
+        buildError: buildModelMultiplierCapError,
+        buildLogFields: block => ({
+          model: block.model,
+          model_multiplier: block.multiplier,
+          max_model_multiplier: block.maxModelMultiplier,
+        }),
+      },
+      {
+        block: getRetiredModelBlockState(model),
+        isBlocked: block => !!block,
+        statusCode: 400,
+        eventName: 'retired_model',
+        buildError: buildRetiredModelError,
+        buildLogFields: block => ({
+          model: block.model,
+          suggestion: block.suggestion,
+        }),
+      },
+      {
+        block: checkUnknownModelRejection(model),
+        isBlocked: block => !!block,
+        statusCode: 400,
+        eventName: 'unknown_model_ai_credits',
+        buildError: block => block.error,
+        buildLogFields: block => ({
+          model: block.model,
+        }),
+      },
+    ] : []),
+  ];
+}
+
+module.exports = { buildCommonGuardChecks };

containers/api-proxy/proxy-request.js

@@ -66,6 +66,7 @@ const {
   buildRetiredModelError,
 } = require('./guards/retired-model-guard');
 const { writeBlockedRequestDiag } = require('./blocked-request-diagnostics');
+const { buildCommonGuardChecks } = require('./guards/common-guard-checks');
 
 // ── Optional token tracker (graceful degradation when not bundled) ────────────
 let trackTokenUsage;
@@ -228,6 +229,11 @@ const proxyWebSocket = createProxyWebSocket({
   buildPermissionDeniedLimitError,
   getAiCreditsBlockState,
   buildAiCreditsLimitError,
+  getModelMultiplierCapBlockState,
+  buildModelMultiplierCapError,
+  getRetiredModelBlockState,
+  buildRetiredModelError,
+  checkUnknownModelRejection,
   trackWebSocketTokenUsage,
 });
 
@@ -402,92 +408,23 @@ function sendGuardBlockedResponse(block, {
 
 function enforceGuards({ body, provider, req, res, requestId, startTime, span, inboundBytes }) {
   const checkModelMultiplier = req.method === 'POST' || req.method === 'PUT' || req.method === 'PATCH';
-  const guardChecks = [
-    {
-      block: getEffectiveTokenBlockState(),
-      isBlocked: block => block && block.maxExceeded,
-      statusCode: 429,
-      eventName: 'effective_tokens_limit_exceeded',
-      buildError: buildEffectiveTokenLimitError,
-      buildLogFields: block => ({
-        total_effective_tokens: block.totalEffectiveTokens,
-        max_effective_tokens: block.maxEffectiveTokens,
-      }),
-    },
-    {
-      block: getMaxRunsBlockState(),
-      isBlocked: block => block && block.maxExceeded,
-      statusCode: 429,
-      eventName: 'max_runs_exceeded',
-      buildError: buildMaxRunsExceededError,
-      buildLogFields: block => ({
-        invocation_count: block.invocationCount,
-        max_runs: block.maxRuns,
-      }),
-    },
-    {
-      block: getPermissionDeniedBlockState(),
-      isBlocked: block => block && block.maxExceeded,
-      statusCode: 403,
-      eventName: 'permission_denied_limit_exceeded',
-      buildError: buildPermissionDeniedLimitError,
-      buildLogFields: block => ({
-        denied_count: block.deniedCount,
-        max_permission_denied: block.maxPermissionDenied,
-      }),
-    },
-    {
-      block: getAiCreditsBlockState(),
-      isBlocked: block => block && block.maxExceeded,
-      statusCode: 429,
-      eventName: 'ai_credits_limit_exceeded',
-      buildError: buildAiCreditsLimitError,
-      buildLogFields: block => ({
-        total_ai_credits: block.totalAiCredits,
-        max_ai_credits: block.maxAiCredits,
-        hard_cap: block.hardCap === true,
-      }),
-    },
-    ...(checkModelMultiplier
-      ? [{
-        block: getModelMultiplierCapBlockState(extractModelFromBody(body)),
-        isBlocked: block => !!block,
-        statusCode: 400,
-        eventName: 'model_multiplier_cap_exceeded',
-        buildError: buildModelMultiplierCapError,
-        buildLogFields: block => ({
-          model: block.model,
-          model_multiplier: block.multiplier,
-          max_model_multiplier: block.maxModelMultiplier,
-        }),
-      }]
-      : []),
-    ...(checkModelMultiplier
-      ? [{
-        block: getRetiredModelBlockState(extractModelFromBody(body)),
-        isBlocked: block => !!block,
-        statusCode: 400,
-        eventName: 'retired_model',
-        buildError: buildRetiredModelError,
-        buildLogFields: block => ({
-          model: block.model,
-          suggestion: block.suggestion,
-        }),
-      }]
-      : []),
-    ...(checkModelMultiplier
-      ? [{
-        block: checkUnknownModelRejection(extractModelFromBody(body)),
-        isBlocked: block => !!block,
-        statusCode: 400,
-        eventName: 'unknown_model_ai_credits',
-        buildError: block => block.error,
-        buildLogFields: block => ({
-          model: block.model,
-        }),
-      }]
-      : []),
-  ];
+  const model = checkModelMultiplier ? extractModelFromBody(body) : null;
+
+  const guardChecks = buildCommonGuardChecks({
+    getEffectiveTokenBlockState,
+    buildEffectiveTokenLimitError,
+    getMaxRunsBlockState,
+    buildMaxRunsExceededError,
+    getPermissionDeniedBlockState,
+    buildPermissionDeniedLimitError,
+    getAiCreditsBlockState,
+    buildAiCreditsLimitError,
+    getModelMultiplierCapBlockState,
+    buildModelMultiplierCapError,
+    getRetiredModelBlockState,
+    buildRetiredModelError,
+    checkUnknownModelRejection,
+  }, model);
 
   for (const guard of guardChecks) {
     if (!guard.isBlocked(guard.block)) continue;

containers/api-proxy/server.websocket.test.js

@@ -276,3 +276,109 @@ describe('proxyWebSocket', () => {
     });
   });
 });
+
+// ── Security guard tests ──────────────────────────────────────────────────────
+//
+// These tests verify that common (non-model-specific) security guards are
+// enforced on the WebSocket upgrade path using the shared buildCommonGuardChecks
+// factory.  Model-specific guards (model_multiplier_cap, retired_model,
+// unknown_model_ai_credits) are intentionally skipped because WebSocket
+// upgrades pass model=null (no JSON body to extract a model from).
+// Guards are triggered by directly calling their apply functions (same
+// technique used in guards/*.test.js unit tests).
+
+describe('proxyWebSocket security guards', () => {
+  let wsProxy;
+  let applyMaxRunsInvocation, resetMaxRunsGuardForTests;
+  let applyEffectiveTokenUsage, resetEffectiveTokenGuardForTests;
+  let applyPermissionDenied, resetPermissionDeniedGuardForTests;
+  let applyAiCreditsUsage, resetAiCreditsGuardForTests;
+
+  beforeAll(() => {
+    jest.resetModules();
+    wsProxy = require('./server').proxyWebSocket;
+    ({ applyMaxRunsInvocation, resetMaxRunsGuardForTests } = require('./guards/max-runs-guard'));
+    ({ applyEffectiveTokenUsage, resetEffectiveTokenGuardForTests } = require('./guards/effective-token-guard'));
+    ({ applyPermissionDenied, resetPermissionDeniedGuardForTests } = require('./guards/max-permission-denied-guard'));
+    ({ applyAiCreditsUsage, resetAiCreditsGuardForTests } = require('./guards/ai-credits-guard'));
+  });
+
+  afterAll(() => {
+    jest.resetModules();
+  });
+
+  afterEach(() => {
+    delete process.env.AWF_MAX_RUNS;
+    delete process.env.AWF_MAX_EFFECTIVE_TOKENS;
+    delete process.env.AWF_MAX_PERMISSION_DENIED;
+    delete process.env.AWF_MAX_AI_CREDITS;
+    resetMaxRunsGuardForTests();
+    resetEffectiveTokenGuardForTests();
+    resetPermissionDeniedGuardForTests();
+    resetAiCreditsGuardForTests();
+    jest.restoreAllMocks();
+  });
+
+  it('blocks with 429 when max-runs limit is exceeded', () => {
+    process.env.AWF_MAX_RUNS = '1';
+    applyMaxRunsInvocation(); // consume the single allowed run
+
+    const socket = makeMockSocket();
+    wsProxy(makeUpgradeReq(), socket, Buffer.alloc(0), 'api.openai.com', {}, 'openai');
+
+    expect(socket.write).toHaveBeenCalledWith(expect.stringContaining('HTTP/1.1 429 Too Many Requests'));
+    expect(socket.write).toHaveBeenCalledWith(expect.stringContaining('"max_runs_exceeded"'));
+    expect(socket.destroy).toHaveBeenCalled();
+  });
+
+  it('blocks with 429 when effective-token limit is exceeded', () => {
+    process.env.AWF_MAX_EFFECTIVE_TOKENS = '1';
+    applyEffectiveTokenUsage({ output_tokens: 5 }, 'gpt-4o'); // exceeds cap of 1
+
+    const socket = makeMockSocket();
+    wsProxy(makeUpgradeReq(), socket, Buffer.alloc(0), 'api.openai.com', {}, 'openai');
+
+    expect(socket.write).toHaveBeenCalledWith(expect.stringContaining('HTTP/1.1 429 Too Many Requests'));
+    expect(socket.write).toHaveBeenCalledWith(expect.stringContaining('"effective_tokens_limit_exceeded"'));
+    expect(socket.destroy).toHaveBeenCalled();
+  });
+
+  it('blocks with 403 when permission-denied limit is exceeded', () => {
+    process.env.AWF_MAX_PERMISSION_DENIED = '1';
+    applyPermissionDenied(); // consume the single allowed denial
+
+    const socket = makeMockSocket();
+    wsProxy(makeUpgradeReq(), socket, Buffer.alloc(0), 'api.openai.com', {}, 'openai');
+
+    expect(socket.write).toHaveBeenCalledWith(expect.stringContaining('HTTP/1.1 403 Forbidden'));
+    expect(socket.write).toHaveBeenCalledWith(expect.stringContaining('"permission_denied_limit_exceeded"'));
+    expect(socket.destroy).toHaveBeenCalled();
+  });
+
+  it('blocks with 429 when ai-credits limit is exceeded', () => {
+    process.env.AWF_MAX_AI_CREDITS = '0.000001'; // tiny cap — any real usage will exceed it
+    applyAiCreditsUsage({ input_tokens: 1_000_000, output_tokens: 1_000_000 }, 'gpt-4o');
+
+    const socket = makeMockSocket();
+    wsProxy(makeUpgradeReq(), socket, Buffer.alloc(0), 'api.openai.com', {}, 'openai');
+
+    expect(socket.write).toHaveBeenCalledWith(expect.stringContaining('HTTP/1.1 429 Too Many Requests'));
+    expect(socket.write).toHaveBeenCalledWith(expect.stringContaining('"ai_credits_limit_exceeded"'));
+    expect(socket.destroy).toHaveBeenCalled();
+  });
+
+  it('allows the upgrade when no guards are triggered', () => {
+    // No guard env vars set and no usage applied — all guards pass.
+    // Without HTTPS_PROXY the upgrade will fail with 502, but the key point is
+    // that it gets past the guard checks (no 429/403 is written).
+    const socket = makeMockSocket();
+    wsProxy(makeUpgradeReq(), socket, Buffer.alloc(0), 'api.openai.com', {}, 'openai');
+
+    const guardStatuses = ['HTTP/1.1 429', 'HTTP/1.1 403 Forbidden'];
+    for (const status of guardStatuses) {
+      expect(socket.write).not.toHaveBeenCalledWith(expect.stringContaining(status));
+    }
+    // The 502 from missing HTTPS_PROXY confirms we got past the guards.
+    expect(socket.write).toHaveBeenCalledWith(expect.stringContaining('HTTP/1.1 502 Bad Gateway'));
+  });
+});