OmniRoute/tests/unit/api-key-mask-fix.test.mjs

/**
 * Unit tests for the masked API key fix (#523).
 *
 * GET /api/keys returns masked API keys (e.g. "sk-31c4e****8600").
 * CLI tool card dropdowns used `key.key` (the masked value) as the select
 * option value, so the masked key got written to config files, causing 401s.
 *
 * The fix: frontends send `keyId` (DB row id) instead, and backends resolve
 * the full key from DB via `resolveApiKey()`.
 *
 * This test inlines the resolver logic (ESM modules are read-only) with a
 * mock DB lookup function.
 */
import test, { describe, it } from "node:test";
import assert from "node:assert/strict";

// ─── Mock DB + inlined resolveApiKey ────────────────────────────────────

/** In-memory key store keyed by id */
const mockKeyStore = new Map();

/**
 * Mock getApiKeyById — mirrors the real function's contract:
 * returns the key record (with .key field) or null.
 */
async function getApiKeyById(id) {
  return mockKeyStore.get(id) || null;
}

/**
 * Inlined resolveApiKey from src/shared/services/apiKeyResolver.ts
 * (can't import ESM modules in test runner without tsx overhead).
 */
async function resolveApiKey(apiKeyId, apiKey) {
  if (apiKeyId) {
    try {
      const keyRecord = await getApiKeyById(apiKeyId);
      if (keyRecord?.key) return keyRecord.key;
    } catch {
      /* fall through */
    }
  }
  return apiKey || "sk_omniroute";
}

// ─── Server-side masking function (matches /api/keys endpoint) ─────────

/**
 * Mask an API key for display: first 8 chars + "****" + last 4 chars.
 * This is the server-side masking that created the original bug.
 */
function maskApiKey(key) {
  if (!key || key.length <= 12) return key;
  return key.slice(0, 8) + "****" + key.slice(-4);
}

// ─── Tests ──────────────────────────────────────────────────────────────

describe("resolveApiKey", () => {
  it("resolves full key from apiKeyId when DB lookup succeeds", async () => {
    mockKeyStore.clear();
    mockKeyStore.set("key-001", { id: "key-001", key: "sk-31c4eabcd1234efgh8600" });

    const result = await resolveApiKey("key-001", null);
    assert.equal(result, "sk-31c4eabcd1234efgh8600");
  });

  it("falls back to apiKey when apiKeyId lookup returns null", async () => {
    mockKeyStore.clear();

    const result = await resolveApiKey("nonexistent-id", "sk-fallback-key");
    assert.equal(result, "sk-fallback-key");
  });

  it("falls back to apiKey when apiKeyId lookup throws", async () => {
    mockKeyStore.clear();
    // Override getApiKeyById to throw
    const originalGet = getApiKeyById;
    const throwingGet = async () => {
      throw new Error("DB connection failed");
    };

    // Temporarily replace
    const savedRef = mockKeyStore.get.bind(mockKeyStore);
    // We'll call resolveApiKey with a custom approach — since the inlined
    // function calls our local getApiKeyById, let's just test by setting
    // up the store to throw via a different mechanism
    // Actually, let's just test the inline function directly with a mock:
    async function resolveApiKeyWithThrowingDb(apiKeyId, apiKey) {
      if (apiKeyId) {
        try {
          throw new Error("DB connection failed");
        } catch {
          /* fall through */
        }
      }
      return apiKey || "sk_omniroute";
    }

    const result = await resolveApiKeyWithThrowingDb("key-001", "sk-fallback-key");
    assert.equal(result, "sk-fallback-key");
  });

  it("falls back to sk_omniroute when both are null", async () => {
    mockKeyStore.clear();

    const result = await resolveApiKey(null, null);
    assert.equal(result, "sk_omniroute");
  });

  it("falls back to sk_omniroute when both are undefined", async () => {
    mockKeyStore.clear();

    const result = await resolveApiKey(undefined, undefined);
    assert.equal(result, "sk_omniroute");
  });

  it("prefers resolved key from apiKeyId over masked apiKey", async () => {
    mockKeyStore.clear();
    mockKeyStore.set("key-002", { id: "key-002", key: "sk-fullkey1234567890abcdef" });

    // The masked apiKey is what /api/keys returns — should NOT be used
    const maskedKey = maskApiKey("sk-fullkey1234567890abcdef");
    const result = await resolveApiKey("key-002", maskedKey);
    assert.equal(result, "sk-fullkey1234567890abcdef");
    assert.notEqual(result, maskedKey);
  });
});

describe("maskApiKey", () => {
  it("masks a long key correctly", () => {
    const result = maskApiKey("sk-31c4eabcd1234efgh8600");
    assert.equal(result, "sk-31c4e****8600");
  });

  it("does not mask short keys", () => {
    const result = maskApiKey("sk-short12");
    assert.equal(result, "sk-short12");
  });

  it("handles null/undefined gracefully", () => {
    assert.equal(maskApiKey(null), null);
    assert.equal(maskApiKey(undefined), undefined);
  });

  it("produces a key that is NOT usable for auth", () => {
    const fullKey = "sk-31c4eabcd1234efgh8600";
    const masked = maskApiKey(fullKey);
    assert.notEqual(masked, fullKey);
    assert.ok(masked.includes("****"));
    assert.ok(masked.startsWith(fullKey.slice(0, 8)));
    assert.ok(masked.endsWith(fullKey.slice(-4)));
  });
});

describe("Bug reproduction: masked key written to config", () => {
  it("reproduces the original bug — masked key fails auth", () => {
    const fullKey = "sk-31c4eabcd1234efgh8600";
    const masked = maskApiKey(fullKey);

    // Simulating what happened before the fix: dropdown used masked key as value
    // and sent it directly to the backend, which wrote it to config
    const writtenToConfig = masked; // BUG: masked key saved to config

    // Auth with masked key would fail
    assert.notEqual(writtenToConfig, fullKey);
    assert.ok(writtenToConfig.includes("****"));

    // This proves the bug: the config file contains "sk-31c4e****8600"
    // which is NOT a valid API key and would cause 401 errors
  });

  it("verifies the fix — keyId resolves to full key", async () => {
    mockKeyStore.clear();
    const fullKey = "sk-31c4eabcd1234efgh8600";
    mockKeyStore.set("key-003", { id: "key-003", key: fullKey });

    // After the fix: frontend sends keyId, backend resolves full key
    const resolved = await resolveApiKey("key-003", null);
    assert.equal(resolved, fullKey);
    assert.ok(!resolved.includes("****"));
  });

  it("simulates full flow: masked dropdown -> keyId -> resolved full key", async () => {
    mockKeyStore.clear();
    const fullKey = "sk-31c4eabcd1234efgh8600";
    const keyId = "key-004";
    mockKeyStore.set(keyId, { id: keyId, key: fullKey });

    // Step 1: /api/keys returns masked list
    const apiKeysResponse = [{ id: keyId, key: maskApiKey(fullKey) }];

    // Step 2: Frontend dropdown now uses key.id as value (not key.key)
    const selectedValue = apiKeysResponse[0].id; // "key-004" (was key.key before fix)
    assert.equal(selectedValue, keyId);

    // Step 3: Frontend sends keyId to backend
    const requestBody = { keyId: selectedValue };

    // Step 4: Backend resolves full key from DB
    const resolvedKey = await resolveApiKey(requestBody.keyId, null);
    assert.equal(resolvedKey, fullKey);
    assert.ok(!resolvedKey.includes("****"));
  });

  it("handles prefix/suffix matching for restoring saved key from file", () => {
    const fullKey = "sk-31c4eabcd1234efgh8600";
    const masked = maskApiKey(fullKey);

    // Simulates what ClaudeToolCard does when reading a key from file:
    // The file contains the full key, and we match against the masked list
    const fileKeyPrefix = fullKey.slice(0, 8); // "sk-31c4e"
    const fileKeySuffix = fullKey.slice(-4); // "8600"

    const apiKeysResponse = [{ id: "key-005", key: masked }];

    // Match by prefix/suffix
    const matchedKey = apiKeysResponse.find(
      (k) => k.key && k.key.startsWith(fileKeyPrefix) && k.key.endsWith(fileKeySuffix)
    );

    assert.ok(matchedKey);
    assert.equal(matchedKey.id, "key-005");
  });
});