STCourses/packages/feature-ai/src/crypto.ts

/**
 * Secret box for secrets at rest (API keys, MCP env and headers): AES-GCM-256 with keys derived by HKDF-SHA-256
 * from STC_SECRET.
 *
 * - `v2.<base64 iv>.<base64 ciphertext+tag>` (v1.1): each purpose has its own HKDF key (`info` =
 *   `stcourses/<purpose>/v2`), and the AES-GCM additional data binds the token to the row it belongs to (for
 *   example the user id and provider of a BYOK key, or the MCP server id and field). A token copied to another
 *   row or purpose does not open, so the server cannot be used to decrypt one secret as another.
 * - `v1.<iv>.<ct>` (v1): one key for everything (info `aes-gcm-256/v1`), no additional data. Only the v1.1 data
 *   upgrade opens these, to re-seal them as v2 (`openLegacy`); `open` refuses them.
 */
export type SecretPurpose = "byok" | "server-api-key" | "mcp";

export interface SecretBox {
  seal(plain: string, purpose: SecretPurpose, aad: string): Promise<string>;
  open(token: string, purpose: SecretPurpose, aad: string): Promise<string>;
  /** v1 tokens only; for the data upgrade. */
  openLegacy(token: string): Promise<string>;
}

const enc = new TextEncoder();
const dec = new TextDecoder();
const b64 = (u: Uint8Array) => Buffer.from(u).toString("base64");
const unb64 = (s: string): Uint8Array<ArrayBuffer> => new Uint8Array(Buffer.from(s, "base64"));
const SALT = enc.encode("stcourses/byok");

export const isLegacyToken = (token: string): boolean => token.startsWith("v1.");

function deriveKey(secret: string, info: string): Promise<CryptoKey> {
  return (async () => {
    const ikm = await crypto.subtle.importKey("raw", enc.encode(secret), "HKDF", false, ["deriveKey"]);
    return crypto.subtle.deriveKey({ name: "HKDF", hash: "SHA-256", salt: SALT, info: enc.encode(info) }, ikm, { name: "AES-GCM", length: 256 }, false, ["encrypt", "decrypt"]);
  })();
}

function parseToken(token: string, version: "v1" | "v2"): { iv: Uint8Array<ArrayBuffer>; ct: Uint8Array<ArrayBuffer> } {
  const [v, iv, ct, extra] = token.split(".");
  if (v !== version || !iv || !ct || extra !== undefined) throw new Error("Unrecognised secret format");
  return { iv: unb64(iv), ct: unb64(ct) };
}

/** Additional data: purpose and row identity, length-prefixed so no two (purpose, aad) pairs encode alike. */
const additionalData = (purpose: SecretPurpose, aad: string) => enc.encode(`stc/${purpose.length}:${purpose}/${aad.length}:${aad}`);

export function createSecretBox(secret: string): SecretBox {
  const keys = new Map<string, Promise<CryptoKey>>();
  const key = (info: string) => {
    let k = keys.get(info);
    if (!k) keys.set(info, (k = deriveKey(secret, info)));
    return k;
  };
  return {
    async seal(plain, purpose, aad) {
      const iv = crypto.getRandomValues(new Uint8Array(12));
      const ct = new Uint8Array(await crypto.subtle.encrypt({ name: "AES-GCM", iv, additionalData: additionalData(purpose, aad) }, await key(`stcourses/${purpose}/v2`), enc.encode(plain)));
      return `v2.${b64(iv)}.${b64(ct)}`;
    },
    async open(token, purpose, aad) {
      const { iv, ct } = parseToken(token, "v2");
      const plain = await crypto.subtle.decrypt({ name: "AES-GCM", iv, additionalData: additionalData(purpose, aad) }, await key(`stcourses/${purpose}/v2`), ct);
      return dec.decode(plain);
    },
    async openLegacy(token) {
      const { iv, ct } = parseToken(token, "v1");
      return dec.decode(await crypto.subtle.decrypt({ name: "AES-GCM", iv }, await key("aes-gcm-256/v1"), ct));
    },
  };
}

/** The v1 token format, exactly as v1 wrote it. Only upgrade tests use it, to build v1 data. */
export async function legacyV1Encrypt(secret: string, plain: string): Promise<string> {
  const iv = crypto.getRandomValues(new Uint8Array(12));
  const ct = new Uint8Array(await crypto.subtle.encrypt({ name: "AES-GCM", iv }, await deriveKey(secret, "aes-gcm-256/v1"), enc.encode(plain)));
  return `v1.${b64(iv)}.${b64(ct)}`;
}

/** AAD for a user's own key: user id and provider kind (the user_api_keys row). */
export const byokAad = (userId: string, provider: string) => `${userId}\n${provider}`;
/** AAD for the admin's provider key (one per server). */
export const SERVER_KEY_AAD = "server";
/** AAD for an MCP server's env or headers: server id and field. */
export const mcpAad = (serverId: string, field: "env" | "headers") => `${serverId}\n${field}`;