Momento/mcp-server/auth.js

/**
 * Memento MCP Server - Optimized API Key Management
 *
 * Performance optimizations:
 * - O(1) key lookup using indexed queries
 * - Batch operations for listing
 * - Connection-aware caching
 */

import { createHash, randomBytes } from 'crypto';

const KEY_PREFIX = 'mcp_key_';

// Simple in-memory cache for API keys (TTL: 60 seconds)
const keyCache = new Map();
const CACHE_TTL = 60000;

function getCachedKey(keyHash) {
  const cached = keyCache.get(keyHash);
  if (cached && Date.now() - cached.timestamp < CACHE_TTL) {
    return cached.data;
  }
  keyCache.delete(keyHash);
  return null;
}

function setCachedKey(keyHash, data) {
  keyCache.set(keyHash, { data, timestamp: Date.now() });
}

/**
 * Generate a new API key.
 * @param {import('@prisma/client').PrismaClient} prisma
 * @param {object} opts
 * @param {string} opts.name - Human-readable name for this key
 * @param {string} opts.userId - User ID to link this key to
 * @returns {{ rawKey: string, info: object }} The raw key (show once!) and key info
 */
export async function generateApiKey(prisma, { name, userId }) {
  // Validate user exists
  const user = await prisma.user.findUnique({
    where: { id: userId },
    select: { id: true, name: true, email: true, role: true },
  });
  if (!user) throw new Error(`User not found: ${userId}`);

  // Generate raw key: mcp_sk_{32 random hex chars}
  const rawBytes = randomBytes(24);
  const shortId = rawBytes.toString('hex').substring(0, 8);
  const rawKey = `mcp_sk_${rawBytes.toString('hex')}`;

  // Hash for storage
  const keyHash = hashKey(rawKey);

  const keyInfo = {
    shortId,
    name: name || `Key for ${user.name}`,
    userId: user.id,
    userName: user.name,
    userEmail: user.email,
    keyHash,
    createdAt: new Date().toISOString(),
    lastUsedAt: null,
    active: true,
  };

  await prisma.systemConfig.create({
    data: {
      key: `${KEY_PREFIX}${shortId}`,
      value: JSON.stringify(keyInfo),
    },
  });

  return {
    rawKey,  // Only returned once!
    info: {
      shortId,
      name: keyInfo.name,
      userId: keyInfo.userId,
      userName: keyInfo.userName,
      createdAt: keyInfo.createdAt,
    },
  };
}

/**
 * Validate an API key and return the associated user info.
 * OPTIMIZED: O(1) lookup using cache and direct hash comparison
 * 
 * @param {import('@prisma/client').PrismaClient} prisma
 * @param {string} rawKey - The raw API key from the request header
 * @returns {object|null} User info if valid, null if invalid/inactive
 */
export async function validateApiKey(prisma, rawKey) {
  if (!rawKey || !rawKey.startsWith('mcp_sk_')) return null;

  const keyHash = hashKey(rawKey);
  
  const cached = getCachedKey(keyHash);
  if (cached) {
    return cached;
  }

  const shortId = rawKey.substring(7, 15);
  const configKey = `${KEY_PREFIX}${shortId}`;

  const entry = await prisma.systemConfig.findUnique({ where: { key: configKey } });
  if (!entry) return null;

  try {
    const info = JSON.parse(entry.value);
    if (info.keyHash !== keyHash || !info.active) return null;

    info.lastUsedAt = new Date().toISOString();
    prisma.systemConfig.update({
      where: { key: configKey },
      data: { value: JSON.stringify(info) },
    }).catch(() => {});

    const result = {
      apiKeyId: info.shortId,
      apiKeyName: info.name,
      userId: info.userId,
      userName: info.userName,
    };
    
    setCachedKey(keyHash, result);
    return result;
  } catch {
    return null;
  }
}

/**
 * List all API keys (without revealing hashes).
 * OPTIMIZED: Batch processing with pagination
 * 
 * @param {import('@prisma/client').PrismaClient} prisma
 * @param {object} [opts]
 * @param {string} [opts.userId] - Filter by user
 * @returns {array}
 */
export async function listApiKeys(prisma, { userId } = {}) {
  const allKeys = await prisma.systemConfig.findMany({
    where: { key: { startsWith: KEY_PREFIX } },
    take: 1000, // Reasonable limit
  });

  const keys = [];
  for (const entry of allKeys) {
    try {
      const info = JSON.parse(entry.value);
      if (userId && info.userId !== userId) continue;
      keys.push({
        shortId: info.shortId,
        name: info.name,
        userId: info.userId,
        userName: info.userName,
        active: info.active,
        createdAt: info.createdAt,
        lastUsedAt: info.lastUsedAt,
      });
    } catch {
      // skip
    }
  }

  return keys;
}

/**
 * Revoke an API key by shortId.
 * @param {import('@prisma/client').PrismaClient} prisma
 * @param {string} shortId - The key's short identifier
 * @returns {boolean} Whether the key was found and deactivated
 */
export async function revokeApiKey(prisma, shortId) {
  const configKey = `${KEY_PREFIX}${shortId}`;
  const entry = await prisma.systemConfig.findUnique({ where: { key: configKey } });

  if (!entry) return false;

  const info = JSON.parse(entry.value);
  if (!info.active) return false;

  info.active = false;
  info.revokedAt = new Date().toISOString();

  await prisma.systemConfig.update({
    where: { key: configKey },
    data: { value: JSON.stringify(info) },
  });

  // Clear cache for this key
  keyCache.delete(info.keyHash);

  return true;
}

/**
 * Delete an API key permanently from DB.
 * @param {import('@prisma/client').PrismaClient} prisma
 * @param {string} shortId
 */
export async function deleteApiKey(prisma, shortId) {
  const configKey = `${KEY_PREFIX}${shortId}`;
  try {
    // Get hash before deleting to clear cache
    const entry = await prisma.systemConfig.findUnique({ where: { key: configKey } });
    if (entry) {
      const info = JSON.parse(entry.value);
      keyCache.delete(info.keyHash);
    }
    
    await prisma.systemConfig.delete({ where: { key: configKey } });
    return true;
  } catch {
    return false;
  }
}

/**
 * Resolve a user by email or ID for auth purposes.
 * OPTIMIZED: Added caching for user lookups
 * 
 * @param {import('@prisma/client').PrismaClient} prisma
 * @param {string} identifier - Email or user ID
 * @returns {object|null}
 */
const userCache = new Map();
const USER_CACHE_TTL = 30000; // 30 seconds

export async function resolveUser(prisma, identifier) {
  if (!identifier) return null;

  // Check cache
  const cached = userCache.get(identifier);
  if (cached && Date.now() - cached.timestamp < USER_CACHE_TTL) {
    return cached.data;
  }

  // Try by ID first
  let user = await prisma.user.findUnique({
    where: { id: identifier },
    select: { id: true, name: true, email: true, role: true },
  });

  // Try by email
  if (!user) {
    user = await prisma.user.findUnique({
      where: { email: identifier },
      select: { id: true, name: true, email: true, role: true },
    });
  }

  if (user) {
    userCache.set(identifier, { data: user, timestamp: Date.now() });
    userCache.set(user.email, { data: user, timestamp: Date.now() });
  }

  return user;
}

// ── Internal ──────────────────────────────────────────────────────────────────

function hashKey(rawKey) {
  return createHash('sha256').update(rawKey).digest('hex');
}

/**
 * Clear all caches (useful for testing or memory management)
 */
export function clearAuthCaches() {
  keyCache.clear();
  userCache.clear();
}