/** * db/apiKeys.js — API key management. */ import { v4 as uuidv4 } from "uuid"; import { getDbInstance, rowToCamel } from "./core"; import { backupDbFile } from "./backup"; // ──────────────── Performance Optimizations ──────────────── // Schema check memoization - only run once let _schemaChecked = false; // LRU cache for API key validation (valid keys only) const _keyValidationCache = new Map(); const _keyMetadataCache = new Map(); const CACHE_TTL = 60 * 1000; // 1 minute TTL const MAX_CACHE_SIZE = 1000; // Compiled regex cache for wildcard patterns const _regexCache = new Map(); // Cache for model permission checks const _modelPermissionCache = new Map(); // Prepared statements cache let _stmtGetAllKeys: any = null; let _stmtGetKeyById: any = null; let _stmtValidateKey: any = null; let _stmtGetKeyMetadata: any = null; let _stmtInsertKey: any = null; let _stmtUpdatePermissions: any = null; let _stmtDeleteKey: any = null; /** * Clear all caches (called on key create/update/delete) */ function invalidateCaches() { _keyValidationCache.clear(); _keyMetadataCache.clear(); _modelPermissionCache.clear(); } /** * LRU eviction for cache */ function evictIfNeeded(cache: Map) { if (cache.size > MAX_CACHE_SIZE) { // Remove oldest 20% of entries const entriesToRemove = Math.floor(MAX_CACHE_SIZE * 0.2); let i = 0; for (const key of cache.keys()) { if (i++ >= entriesToRemove) break; cache.delete(key); } } } /** * Get or compile regex for wildcard pattern */ function getWildcardRegex(pattern: string): RegExp { let regex = _regexCache.get(pattern); if (!regex) { const regexStr = pattern.replace(/\*/g, ".*"); regex = new RegExp(`^${regexStr}$`); _regexCache.set(pattern, regex); // Prevent unbounded growth if (_regexCache.size > 100) { const firstKey = _regexCache.keys().next().value; if (firstKey) _regexCache.delete(firstKey); } } return regex; } // Ensure the allowed_models column exists (memoized) function ensureAllowedModelsColumn(db) { if (_schemaChecked) return; try { const columns = db.prepare("PRAGMA table_info(api_keys)").all(); const columnNames = new Set(columns.map((column) => column.name)); if (!columnNames.has("allowed_models")) { db.exec("ALTER TABLE api_keys ADD COLUMN allowed_models TEXT"); console.log("[DB] Added api_keys.allowed_models column"); } _schemaChecked = true; } catch (error) { console.warn("[DB] Failed to verify api_keys schema:", error.message); } } /** * Initialize prepared statements (lazy initialization) */ function getPreparedStatements(db: any) { if (!_stmtGetAllKeys) { _stmtGetAllKeys = db.prepare("SELECT * FROM api_keys ORDER BY created_at"); _stmtGetKeyById = db.prepare("SELECT * FROM api_keys WHERE id = ?"); _stmtValidateKey = db.prepare("SELECT 1 FROM api_keys WHERE key = ?"); _stmtGetKeyMetadata = db.prepare( "SELECT id, name, machine_id, allowed_models FROM api_keys WHERE key = ?" ); _stmtInsertKey = db.prepare( "INSERT INTO api_keys (id, name, key, machine_id, allowed_models, created_at) VALUES (?, ?, ?, ?, ?, ?)" ); _stmtUpdatePermissions = db.prepare("UPDATE api_keys SET allowed_models = ? WHERE id = ?"); _stmtDeleteKey = db.prepare("DELETE FROM api_keys WHERE id = ?"); } return { getAllKeys: _stmtGetAllKeys, getKeyById: _stmtGetKeyById, validateKey: _stmtValidateKey, getKeyMetadata: _stmtGetKeyMetadata, insertKey: _stmtInsertKey, updatePermissions: _stmtUpdatePermissions, deleteKey: _stmtDeleteKey, }; } export async function getApiKeys() { const db = getDbInstance(); ensureAllowedModelsColumn(db); const stmt = getPreparedStatements(db); const rows = stmt.getAllKeys.all() as Record[]; return rows.map((row) => { const camelRow = rowToCamel(row) as Record; // Parse allowed_models from JSON string to array camelRow.allowedModels = parseAllowedModels(camelRow.allowedModels); return camelRow; }); } export async function getApiKeyById(id: string) { const db = getDbInstance(); ensureAllowedModelsColumn(db); const stmt = getPreparedStatements(db); const row = stmt.getKeyById.get(id) as Record | undefined; if (!row) return null; const camelRow = rowToCamel(row) as Record; camelRow.allowedModels = parseAllowedModels(camelRow.allowedModels); return camelRow; } /** * Helper function to safely parse allowed_models JSON */ function parseAllowedModels(value: any): string[] { if (!value || typeof value !== "string" || value.trim() === "") { return []; } try { const parsed = JSON.parse(value); return Array.isArray(parsed) ? parsed : []; } catch { return []; } } export async function createApiKey(name, machineId) { if (!machineId) { throw new Error("machineId is required"); } const db = getDbInstance(); ensureAllowedModelsColumn(db); const now = new Date().toISOString(); const { generateApiKeyWithMachine } = await import("@/shared/utils/apiKey"); const result = generateApiKeyWithMachine(machineId); const apiKey = { id: uuidv4(), name: name, key: result.key, machineId: machineId, allowedModels: [], // Empty array means all models allowed createdAt: now, }; const stmt = getPreparedStatements(db); stmt.insertKey.run(apiKey.id, apiKey.name, apiKey.key, apiKey.machineId, "[]", apiKey.createdAt); backupDbFile("pre-write"); return apiKey; } export async function updateApiKeyPermissions(id, allowedModels) { const db = getDbInstance(); ensureAllowedModelsColumn(db); // allowedModels should be an array of model IDs (strings) // Empty array means all models are allowed const modelsJson = JSON.stringify(allowedModels || []); const stmt = getPreparedStatements(db); const result = stmt.updatePermissions.run(modelsJson, id); if (result.changes === 0) return false; // Invalidate caches since permissions changed invalidateCaches(); backupDbFile("pre-write"); return true; } export async function deleteApiKey(id) { const db = getDbInstance(); const stmt = getPreparedStatements(db); const result = stmt.deleteKey.run(id); if (result.changes === 0) return false; // Invalidate caches since a key was removed invalidateCaches(); backupDbFile("pre-write"); return true; } /** * Validate API key with caching for performance * Cached valid keys reduce DB hits on every request */ export async function validateApiKey(key) { if (!key || typeof key !== "string") return false; const now = Date.now(); // Check cache first const cached = _keyValidationCache.get(key); if (cached && now - cached.timestamp < CACHE_TTL) { return cached.valid; } const db = getDbInstance(); const stmt = getPreparedStatements(db); const row = stmt.validateKey.get(key); const valid = !!row; // Only cache valid keys to prevent cache pollution if (valid) { evictIfNeeded(_keyValidationCache); _keyValidationCache.set(key, { valid: true, timestamp: now }); } return valid; } /** * Get API key metadata with caching for performance */ export async function getApiKeyMetadata(key) { if (!key || typeof key !== "string") return null; const now = Date.now(); // Check cache first const cached = _keyMetadataCache.get(key); if (cached && now - cached.timestamp < CACHE_TTL) { return cached.metadata; } const db = getDbInstance(); ensureAllowedModelsColumn(db); const stmt = getPreparedStatements(db); const row = stmt.getKeyMetadata.get(key) as Record | undefined; if (!row) return null; const metadata = { id: row.id, name: row.name, machineId: row.machine_id, allowedModels: parseAllowedModels(row.allowed_models), }; // Cache the result evictIfNeeded(_keyMetadataCache); _keyMetadataCache.set(key, { metadata, timestamp: now }); return metadata; } /** * Check if a model is allowed for a given API key * @param {string} key - The API key * @param {string} modelId - The model ID to check * @returns {boolean} - true if allowed, false if not */ export async function isModelAllowedForKey(key, modelId) { // If no key provided, allow (request may be using different auth method like JWT) // If no modelId provided, deny (invalid request) if (!key) return true; if (!modelId) return false; // Create cache key const cacheKey = `${key}:${modelId}`; const now = Date.now(); // Check permission cache const cached = _modelPermissionCache.get(cacheKey); if (cached && now - cached.timestamp < CACHE_TTL) { return cached.allowed; } const metadata = await getApiKeyMetadata(key); // SECURITY: Key not found in database = deny access (invalid/non-existent key) if (!metadata) return false; const { allowedModels } = metadata; // Empty array means all models allowed if (!allowedModels || allowedModels.length === 0) { return true; } let allowed = false; // Check if model matches any allowed pattern // Support exact match and prefix match (e.g., "openai/*" allows all OpenAI models) for (const pattern of allowedModels) { if (pattern === modelId) { allowed = true; break; } if (pattern.endsWith("/*")) { const prefix = pattern.slice(0, -2); // Remove "/*" if (modelId.startsWith(prefix + "/") || modelId.startsWith(prefix)) { allowed = true; break; } } // Support wildcard patterns using cached regex if (pattern.includes("*")) { const regex = getWildcardRegex(pattern); if (regex.test(modelId)) { allowed = true; break; } } } // Cache the result evictIfNeeded(_modelPermissionCache); _modelPermissionCache.set(cacheKey, { allowed, timestamp: now }); return allowed; } /** * Clear all caches (exported for testing/debugging) */ export function clearApiKeyCaches() { invalidateCaches(); _modelPermissionCache.clear(); _regexCache.clear(); }