mirror of
https://github.com/diegosouzapw/OmniRoute.git
synced 2026-08-07 15:52:52 +03:00
fix: enforce API key model restrictions and budget limits across all endpoints
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
This commit is contained in:
104
src/shared/utils/apiKeyPolicy.ts
Normal file
104
src/shared/utils/apiKeyPolicy.ts
Normal file
@@ -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 };
|
||||
}
|
||||
Reference in New Issue
Block a user