diff --git a/containers/api-proxy/guards/common-guard-checks.js b/containers/api-proxy/guards/common-guard-checks.js new file mode 100644 index 000000000..c9602944e --- /dev/null +++ b/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} 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 }; diff --git a/containers/api-proxy/proxy-request.js b/containers/api-proxy/proxy-request.js index 59411176d..3ee6da5e5 100644 --- a/containers/api-proxy/proxy-request.js +++ b/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; diff --git a/containers/api-proxy/server.websocket.test.js b/containers/api-proxy/server.websocket.test.js index 3b7c64649..eebc41c2a 100644 --- a/containers/api-proxy/server.websocket.test.js +++ b/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')); + }); +}); diff --git a/containers/api-proxy/websocket-proxy.js b/containers/api-proxy/websocket-proxy.js index 88d4699fe..f8d58759a 100644 --- a/containers/api-proxy/websocket-proxy.js +++ b/containers/api-proxy/websocket-proxy.js @@ -4,6 +4,53 @@ const http = require('http'); const tls = require('tls'); const { URL } = require('url'); const { computeTokenBudgetUsage } = require('./token-budget-log'); +const { buildCommonGuardChecks } = require('./guards/common-guard-checks'); + +/** Maps numeric status codes used by guards to HTTP/1.1 reason phrases. */ +const HTTP_STATUS_LINES = { + 400: '400 Bad Request', + 403: '403 Forbidden', + 429: '429 Too Many Requests', +}; + +/** + * Enforce all common security guards for a WebSocket upgrade request. + * Writes a raw HTTP error response to the socket and destroys it when any + * guard triggers, then returns true. Returns false when all guards pass. + * + * @param {object} ctx + * @param {import('net').Socket} ctx.socket + * @param {Function} ctx.logRequest + * @param {string} ctx.requestId + * @param {string} ctx.provider + * @param {object} guardDeps - Guard state getter and error-builder functions + * (same shape as the `deps` parameter of buildCommonGuardChecks). + * @returns {boolean} true when a guard blocked the request. + */ +function enforceWebSocketGuards({ socket, logRequest, requestId, provider }, guardDeps) { + // WebSocket upgrade requests have no JSON body, so model-specific guards + // receive null and are skipped (their getters return null for null models). + const guardChecks = buildCommonGuardChecks(guardDeps, null); + + for (const guard of guardChecks) { + if (!guard.isBlocked(guard.block)) continue; + + const block = guard.block; + logRequest('warn', guard.eventName, { + request_id: requestId, + provider, + ...guard.buildLogFields(block), + }); + + const statusLine = HTTP_STATUS_LINES[guard.statusCode] || String(guard.statusCode); + socket.write(`HTTP/1.1 ${statusLine}\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n`); + socket.write(JSON.stringify(guard.buildError(block))); + socket.destroy(); + return true; + } + + return false; +} function createProxyWebSocket({ limiter, @@ -23,8 +70,28 @@ function createProxyWebSocket({ buildPermissionDeniedLimitError, getAiCreditsBlockState, buildAiCreditsLimitError, + getModelMultiplierCapBlockState, + buildModelMultiplierCapError, + getRetiredModelBlockState, + buildRetiredModelError, + checkUnknownModelRejection, trackWebSocketTokenUsage, }) { + const guardDeps = { + getEffectiveTokenBlockState, + buildEffectiveTokenLimitError, + getMaxRunsBlockState, + buildMaxRunsExceededError, + getPermissionDeniedBlockState, + buildPermissionDeniedLimitError, + getAiCreditsBlockState, + buildAiCreditsLimitError, + getModelMultiplierCapBlockState, + buildModelMultiplierCapError, + getRetiredModelBlockState, + buildRetiredModelError, + checkUnknownModelRejection, + }; /** * Handle a WebSocket upgrade request by tunnelling through the Squid proxy. * @@ -65,60 +132,7 @@ function createProxyWebSocket({ const upstreamPath = buildUpstreamPath(req.url, targetHost, basePath); - const etBlock = getEffectiveTokenBlockState(); - if (etBlock && etBlock.maxExceeded) { - logRequest('warn', 'effective_tokens_limit_exceeded', { - request_id: requestId, - provider, - total_effective_tokens: etBlock.totalEffectiveTokens, - max_effective_tokens: etBlock.maxEffectiveTokens, - }); - socket.write('HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n'); - socket.write(JSON.stringify(buildEffectiveTokenLimitError(etBlock))); - socket.destroy(); - return; - } - - const mrBlock = getMaxRunsBlockState(); - if (mrBlock && mrBlock.maxExceeded) { - logRequest('warn', 'max_runs_exceeded', { - request_id: requestId, - provider, - invocation_count: mrBlock.invocationCount, - max_runs: mrBlock.maxRuns, - }); - socket.write('HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n'); - socket.write(JSON.stringify(buildMaxRunsExceededError(mrBlock))); - socket.destroy(); - return; - } - - const pdBlock = getPermissionDeniedBlockState(); - if (pdBlock && pdBlock.maxExceeded) { - logRequest('warn', 'permission_denied_limit_exceeded', { - request_id: requestId, - provider, - denied_count: pdBlock.deniedCount, - max_permission_denied: pdBlock.maxPermissionDenied, - }); - socket.write('HTTP/1.1 403 Forbidden\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n'); - socket.write(JSON.stringify(buildPermissionDeniedLimitError(pdBlock))); - socket.destroy(); - return; - } - - const aiCreditsBlock = getAiCreditsBlockState(); - if (aiCreditsBlock && aiCreditsBlock.maxExceeded) { - logRequest('warn', 'ai_credits_limit_exceeded', { - request_id: requestId, - provider, - total_ai_credits: aiCreditsBlock.totalAiCredits, - max_ai_credits: aiCreditsBlock.maxAiCredits, - hard_cap: aiCreditsBlock.hardCap === true, - }); - socket.write('HTTP/1.1 429 Too Many Requests\r\nContent-Type: application/json\r\nConnection: close\r\n\r\n'); - socket.write(JSON.stringify(buildAiCreditsLimitError(aiCreditsBlock))); - socket.destroy(); + if (enforceWebSocketGuards({ socket, logRequest, requestId, provider }, guardDeps)) { return; }