Merge pull request #131 from ersintarhan/fix/api-key-model-restriction

Approved: Critical security fix for API key model restrictions. Minor improvements (error logging, type safety) will be applied in a follow-up commit.

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(

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(

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) {

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) {

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

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];

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) {

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(

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<ApiKeyPolicyResult> {
+  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 };
+}

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)