From 0c0a56d4de4eb9eb26d5f539828dc464cea0ec7e Mon Sep 17 00:00:00 2001 From: Ersin Tarhan Date: Wed, 25 Feb 2026 06:43:55 +0300 Subject: [PATCH] fix: enforce API key model restrictions and budget limits across all endpoints MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit isModelAllowedForKey() existed in src/lib/db/apiKeys.ts but was never called anywhere. API keys with allowedModels restrictions could access any model through any endpoint. Changes: - Add shared enforceApiKeyPolicy() middleware (model restriction + budget) - Wire it into chat handler (replacing inline budget-only check) - Wire it into all /v1/* endpoints: embeddings, images/generations, audio/speech, audio/transcriptions, moderations, rerank - Wire it into provider-specific endpoints: /v1/providers/[provider]/embeddings, /v1/providers/[provider]/images/generations The middleware checks: 1. Model restriction — if key has allowedModels, verify the model is permitted 2. Budget limit — if key has budget configured, verify it hasn't been exceeded Fixes #130 --- src/app/api/v1/audio/speech/route.ts | 5 + src/app/api/v1/audio/transcriptions/route.ts | 5 + src/app/api/v1/embeddings/route.ts | 5 + src/app/api/v1/images/generations/route.ts | 5 + src/app/api/v1/moderations/route.ts | 6 + .../providers/[provider]/embeddings/route.ts | 5 + .../[provider]/images/generations/route.ts | 5 + src/app/api/v1/rerank/route.ts | 5 + src/shared/utils/apiKeyPolicy.ts | 104 ++++++++++++++++++ src/sse/handlers/chat.ts | 31 ++---- 10 files changed, 155 insertions(+), 21 deletions(-) create mode 100644 src/shared/utils/apiKeyPolicy.ts diff --git a/src/app/api/v1/audio/speech/route.ts b/src/app/api/v1/audio/speech/route.ts index 985edb7143..c16ba8cbfa 100644 --- a/src/app/api/v1/audio/speech/route.ts +++ b/src/app/api/v1/audio/speech/route.ts @@ -4,6 +4,7 @@ import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/serv import { parseSpeechModel } from "@omniroute/open-sse/config/audioRegistry.ts"; import { errorResponse } from "@omniroute/open-sse/utils/error.ts"; import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts"; +import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; /** * Handle CORS preflight @@ -41,6 +42,10 @@ export async function POST(request) { return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing model"); } + // Enforce API key policies (model restrictions + budget limits) + const policy = await enforceApiKeyPolicy(request, body.model); + if (policy.rejection) return policy.rejection; + const { provider } = parseSpeechModel(body.model); if (!provider) { return errorResponse( diff --git a/src/app/api/v1/audio/transcriptions/route.ts b/src/app/api/v1/audio/transcriptions/route.ts index 83adf67cfc..adfeb00fff 100644 --- a/src/app/api/v1/audio/transcriptions/route.ts +++ b/src/app/api/v1/audio/transcriptions/route.ts @@ -4,6 +4,7 @@ import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/serv import { parseTranscriptionModel } from "@omniroute/open-sse/config/audioRegistry.ts"; import { errorResponse } from "@omniroute/open-sse/utils/error.ts"; import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts"; +import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; /** * Handle CORS preflight @@ -43,6 +44,10 @@ export async function POST(request) { return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing model"); } + // Enforce API key policies (model restrictions + budget limits) + const policy = await enforceApiKeyPolicy(request, model as string); + if (policy.rejection) return policy.rejection; + const { provider } = parseTranscriptionModel(model); if (!provider) { return errorResponse( diff --git a/src/app/api/v1/embeddings/route.ts b/src/app/api/v1/embeddings/route.ts index 3c3bf53989..bd21987075 100644 --- a/src/app/api/v1/embeddings/route.ts +++ b/src/app/api/v1/embeddings/route.ts @@ -9,6 +9,7 @@ import { errorResponse } from "@omniroute/open-sse/utils/error.ts"; import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts"; import * as log from "@/sse/utils/logger"; import { toJsonErrorPayload } from "@/shared/utils/upstreamError"; +import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; /** * Handle CORS preflight @@ -78,6 +79,10 @@ export async function POST(request) { return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing input"); } + // Enforce API key policies (model restrictions + budget limits) + const policy = await enforceApiKeyPolicy(request, body.model); + if (policy.rejection) return policy.rejection; + // Parse model to get provider const { provider } = parseEmbeddingModel(body.model); if (!provider) { diff --git a/src/app/api/v1/images/generations/route.ts b/src/app/api/v1/images/generations/route.ts index f56ab70ddf..d95bd090f6 100644 --- a/src/app/api/v1/images/generations/route.ts +++ b/src/app/api/v1/images/generations/route.ts @@ -6,6 +6,7 @@ import { errorResponse } from "@omniroute/open-sse/utils/error.ts"; import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts"; import * as log from "@/sse/utils/logger"; import { toJsonErrorPayload } from "@/shared/utils/upstreamError"; +import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; /** * Handle CORS preflight @@ -75,6 +76,10 @@ export async function POST(request) { return errorResponse(HTTP_STATUS.BAD_REQUEST, "Invalid prompt: expected a non-empty string"); } + // Enforce API key policies (model restrictions + budget limits) + const policy = await enforceApiKeyPolicy(request, body.model); + if (policy.rejection) return policy.rejection; + // Parse model to get provider const { provider } = parseImageModel(body.model); if (!provider) { diff --git a/src/app/api/v1/moderations/route.ts b/src/app/api/v1/moderations/route.ts index 8a5651308e..753fa37ebb 100644 --- a/src/app/api/v1/moderations/route.ts +++ b/src/app/api/v1/moderations/route.ts @@ -4,6 +4,7 @@ import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/serv import { parseModerationModel } from "@omniroute/open-sse/config/moderationRegistry.ts"; import { errorResponse } from "@omniroute/open-sse/utils/error.ts"; import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts"; +import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; /** * Handle CORS preflight @@ -38,6 +39,11 @@ export async function POST(request) { } const model = body.model || "omni-moderation-latest"; + + // Enforce API key policies (model restrictions + budget limits) + const policy = await enforceApiKeyPolicy(request, model); + if (policy.rejection) return policy.rejection; + const { provider } = parseModerationModel(model); // Default to openai if no provider prefix diff --git a/src/app/api/v1/providers/[provider]/embeddings/route.ts b/src/app/api/v1/providers/[provider]/embeddings/route.ts index b10bf11d33..c4fe8f0ff4 100644 --- a/src/app/api/v1/providers/[provider]/embeddings/route.ts +++ b/src/app/api/v1/providers/[provider]/embeddings/route.ts @@ -5,6 +5,7 @@ import { getRegistryEntry } from "@omniroute/open-sse/config/providerRegistry.ts import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/services/auth"; import { handleEmbedding } from "@omniroute/open-sse/handlers/embeddings.ts"; import * as log from "@/sse/utils/logger"; +import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; /** * Handle CORS preflight @@ -53,6 +54,10 @@ export async function POST(request, { params }) { body.model = `${providerAlias}/${body.model}`; } + // Enforce API key policies (model restrictions + budget limits) + const policy = await enforceApiKeyPolicy(request, body.model); + if (policy.rejection) return policy.rejection; + // Validate provider match if (body.model) { const prefix = body.model.split("/")[0]; diff --git a/src/app/api/v1/providers/[provider]/images/generations/route.ts b/src/app/api/v1/providers/[provider]/images/generations/route.ts index ace00aba09..8eaf3cb6cf 100644 --- a/src/app/api/v1/providers/[provider]/images/generations/route.ts +++ b/src/app/api/v1/providers/[provider]/images/generations/route.ts @@ -6,6 +6,7 @@ import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/serv import { getImageProvider } from "@omniroute/open-sse/config/imageRegistry.ts"; import * as log from "@/sse/utils/logger"; import { toJsonErrorPayload } from "@/shared/utils/upstreamError"; +import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; /** * Handle CORS preflight @@ -60,6 +61,10 @@ export async function POST(request, { params }) { body.model = `${rawProvider}/${body.model}`; } + // Enforce API key policies (model restrictions + budget limits) + const policy = await enforceApiKeyPolicy(request, body.model); + if (policy.rejection) return policy.rejection; + // Validate provider match const modelProvider = body.model.split("/")[0]; if (modelProvider !== rawProvider) { diff --git a/src/app/api/v1/rerank/route.ts b/src/app/api/v1/rerank/route.ts index 5ed078a23c..6a4b128842 100644 --- a/src/app/api/v1/rerank/route.ts +++ b/src/app/api/v1/rerank/route.ts @@ -4,6 +4,7 @@ import { getProviderCredentials, extractApiKey, isValidApiKey } from "@/sse/serv import { parseRerankModel } from "@omniroute/open-sse/config/rerankRegistry.ts"; import { errorResponse } from "@omniroute/open-sse/utils/error.ts"; import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts"; +import { enforceApiKeyPolicy } from "@/shared/utils/apiKeyPolicy"; /** * Handle CORS preflight @@ -44,6 +45,10 @@ export async function POST(request) { return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing model"); } + // Enforce API key policies (model restrictions + budget limits) + const policy = await enforceApiKeyPolicy(request, body.model); + if (policy.rejection) return policy.rejection; + const { provider } = parseRerankModel(body.model); if (!provider) { return errorResponse( diff --git a/src/shared/utils/apiKeyPolicy.ts b/src/shared/utils/apiKeyPolicy.ts new file mode 100644 index 0000000000..9a314d8ac1 --- /dev/null +++ b/src/shared/utils/apiKeyPolicy.ts @@ -0,0 +1,104 @@ +/** + * API Key Policy Enforcement — Shared middleware for all /v1/* endpoints. + * + * Enforces API key policies: model restrictions and budget limits. + * Should be called after API key authentication in every endpoint that + * accepts a model parameter. + * + * @module shared/utils/apiKeyPolicy + */ + +import { extractApiKey } from "@/sse/services/auth"; +import { getApiKeyMetadata, isModelAllowedForKey } from "@/lib/localDb"; +import { checkBudget } from "@/domain/costRules"; +import { errorResponse } from "@omniroute/open-sse/utils/error.ts"; +import { HTTP_STATUS } from "@omniroute/open-sse/config/constants.ts"; + +export interface ApiKeyPolicyResult { + /** API key string (null if no key provided) */ + apiKey: string | null; + /** Metadata from DB (null if no key or key not found) */ + apiKeyInfo: any | null; + /** If set, the request should be rejected with this Response */ + rejection: Response | null; +} + +/** + * Enforce API key policies for a request. + * + * Checks: + * 1. Model restriction — if the key has `allowedModels`, verify the requested model is permitted + * 2. Budget limit — if the key has a budget configured, verify it hasn't been exceeded + * + * @param request - The incoming HTTP request + * @param modelStr - The model ID from the request body + * @returns ApiKeyPolicyResult with apiKey, metadata, and optional rejection response + * + * @example + * ```ts + * const policy = await enforceApiKeyPolicy(request, body.model); + * if (policy.rejection) return policy.rejection; + * // proceed with request, optionally use policy.apiKeyInfo + * ``` + */ +export async function enforceApiKeyPolicy( + request: Request, + modelStr: string | null +): Promise { + const apiKey = extractApiKey(request); + + // No API key = local mode, skip policy checks + if (!apiKey) { + return { apiKey: null, apiKeyInfo: null, rejection: null }; + } + + // Fetch key metadata (includes allowedModels) + let apiKeyInfo: any = null; + try { + apiKeyInfo = await getApiKeyMetadata(apiKey); + } catch { + // If metadata fetch fails, don't block — degrade gracefully + return { apiKey, apiKeyInfo: null, rejection: null }; + } + + // Key not found in DB — skip policy (auth layer handles validation) + if (!apiKeyInfo) { + return { apiKey, apiKeyInfo: null, rejection: null }; + } + + // ── Check 1: Model restriction ── + if (modelStr && apiKeyInfo.allowedModels && apiKeyInfo.allowedModels.length > 0) { + const allowed = await isModelAllowedForKey(apiKey, modelStr); + if (!allowed) { + return { + apiKey, + apiKeyInfo, + rejection: errorResponse( + HTTP_STATUS.FORBIDDEN, + `Model "${modelStr}" is not allowed for this API key` + ), + }; + } + } + + // ── Check 2: Budget limit ── + if (apiKeyInfo.id) { + try { + const budgetOk = checkBudget(apiKeyInfo.id); + if (!budgetOk.allowed) { + return { + apiKey, + apiKeyInfo, + rejection: errorResponse( + HTTP_STATUS.RATE_LIMITED, + budgetOk.reason || "Budget limit exceeded" + ), + }; + } + } catch { + // Budget check is best-effort — don't block on errors + } + } + + return { apiKey, apiKeyInfo, rejection: null }; +} diff --git a/src/sse/handlers/chat.ts b/src/sse/handlers/chat.ts index 878c87c9dd..d2b83cdd3b 100644 --- a/src/sse/handlers/chat.ts +++ b/src/sse/handlers/chat.ts @@ -23,7 +23,7 @@ import { } from "@omniroute/open-sse/utils/proxyFetch.ts"; import * as log from "../utils/logger"; import { updateProviderCredentials, checkAndRefreshToken } from "../services/tokenRefresh"; -import { getSettings, getCombos, getApiKeyMetadata } from "@/lib/localDb"; +import { getSettings, getCombos } from "@/lib/localDb"; import { resolveProxyForConnection } from "@/lib/localDb"; import { logProxyEvent } from "../../lib/proxyLogger"; import { logTranslationEvent } from "../../lib/translatorEvents"; @@ -34,8 +34,9 @@ import { getCircuitBreaker, CircuitBreakerOpenError } from "../../shared/utils/c import { isModelAvailable, setModelUnavailable } from "../../domain/modelAvailability"; import { RequestTelemetry, recordTelemetry } from "../../shared/utils/requestTelemetry"; import { generateRequestId } from "../../shared/utils/requestId"; -import { checkBudget, recordCost } from "../../domain/costRules"; +import { recordCost } from "../../domain/costRules"; import { logAuditEvent } from "../../lib/compliance/index"; +import { enforceApiKeyPolicy } from "../../shared/utils/apiKeyPolicy"; /** * Handle chat completion request @@ -97,15 +98,8 @@ export async function handleChat(request: any, clientRawRequest: any = null) { // Log API key (masked) const authHeader = request.headers.get("Authorization"); const apiKey = extractApiKey(request); - let apiKeyInfo = null; if (authHeader && apiKey) { - const masked = log.maskKey(apiKey); - log.debug("AUTH", `API Key: ${masked}`); - try { - apiKeyInfo = await getApiKeyMetadata(apiKey); - } catch { - apiKeyInfo = null; - } + log.debug("AUTH", `API Key: ${log.maskKey(apiKey)}`); } else { log.debug("AUTH", "No API key provided (local mode)"); } @@ -129,19 +123,14 @@ export async function handleChat(request: any, clientRawRequest: any = null) { return errorResponse(HTTP_STATUS.BAD_REQUEST, "Missing model"); } - // Pipeline: Budget check (if API key has budget limits) + // Pipeline: API key policy enforcement (model restrictions + budget limits) telemetry.startPhase("policy"); - if (apiKeyInfo?.id) { - try { - const budgetOk = checkBudget(apiKeyInfo.id); - if (!budgetOk.allowed) { - log.warn("BUDGET", `API key ${apiKeyInfo.id} exceeded budget: ${budgetOk.reason}`); - return errorResponse(429, budgetOk.reason || "Budget limit exceeded"); - } - } catch { - // Budget check is best-effort — don't block on errors - } + const policy = await enforceApiKeyPolicy(request, modelStr); + if (policy.rejection) { + log.warn("POLICY", `API key policy rejected: ${modelStr} (key=${policy.apiKeyInfo?.id || "unknown"})`); + return policy.rejection; } + const apiKeyInfo = policy.apiKeyInfo; telemetry.endPhase(); // Check if model is a combo (has multiple models with fallback)