curriculum-project-hub/hub/test/integration/local-kek-rotation.test.ts

import { mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { afterAll, beforeEach, describe, expect, it } from "vitest";
import { ProviderConnectionService, verifyStoredProviderEnvelopes } from "../../src/connections/providerConnections.js";
import {
  FeishuApplicationConnectionService,
  verifyStoredFeishuApplicationEnvelopes,
} from "../../src/connections/feishuApplicationConnections.js";
import { LocalSecretEnvelope, loadLocalSecretKeyringFile, type SecretEnvelopeV1 } from "../../src/security/secretEnvelope.js";
import { LOCAL_KEK_ROTATION_ADVISORY_LOCK, rotateLocalKek } from "../../src/security/localKekRotation.js";
import { DEFAULT_ORG_ID, prisma, resetDb } from "./helpers.js";

const temporaryDirectories: string[] = [];

describe("local KEK rotation", () => {
  beforeEach(async () => {
    await resetDb();
  });

  afterAll(async () => {
    await Promise.all(temporaryDirectories.splice(0).map((directory) => rm(directory, { recursive: true, force: true })));
    await prisma.$disconnect();
  });

  it("atomically retains the old KEK, rewraps DEKs, and authenticates every envelope", async () => {
    const oldKey = Buffer.alloc(32, "o");
    const oldSecrets = new LocalSecretEnvelope({
      activeKeyId: "local-old",
      keys: new Map([["local-old", oldKey]]),
    });
    await prisma.user.create({ data: { id: "rotation-admin", feishuOpenId: "ou_rotation", displayName: "Rotation" } });
    await prisma.organizationMembership.create({
      data: { organizationId: DEFAULT_ORG_ID, userId: "rotation-admin", role: "OWNER" },
    });
    await new ProviderConnectionService(prisma, oldSecrets, async () => {}).rotateByok({
      organizationId: DEFAULT_ORG_ID,
      providerId: "openrouter",
      actorUserId: "rotation-admin",
      baseUrl: "https://provider.example/api",
      authToken: "credential-stays-encrypted",
    });
    await new FeishuApplicationConnectionService(prisma, oldSecrets, async () => {}).rotateCustomerApplication({
      organizationId: DEFAULT_ORG_ID,
      actorUserId: "rotation-admin",
      appId: "cli_rotation",
      appSecret: "feishu-credential-stays-encrypted",
      botOpenId: "ou_rotation_bot",
    });
    const before = await prisma.providerCredentialVersion.findFirstOrThrow({
      include: { connection: true },
    });
    const beforeEnvelope = before.envelope as unknown as SecretEnvelopeV1;
    const beforeFeishu = await prisma.feishuApplicationCredentialVersion.findFirstOrThrow({
      include: { connection: true },
    });
    const beforeFeishuEnvelope = beforeFeishu.envelope as unknown as SecretEnvelopeV1;
    const directory = await mkdtemp(join(tmpdir(), "cph-kek-rotation-"));
    temporaryDirectories.push(directory);
    const keyringFile = join(directory, "keyring.json");
    await writeFile(keyringFile, JSON.stringify({
      version: 1,
      activeKeyId: "local-old",
      keys: { "local-old": oldKey.toString("base64") },
    }), { mode: 0o600 });
    const expectedOwner = {
      uid: process.getuid?.() ?? 0,
      gid: process.getgid?.() ?? 0,
    };
    if (expectedOwner.uid !== 0 || expectedOwner.gid !== 0) {
      await expect(rotateLocalKek({
        keyringFile,
        prisma,
        generatedKey: Buffer.alloc(32, "x"),
      })).rejects.toThrow("must be owned by 0:0");
    }

    let releaseLock!: () => void;
    const releaseLockPromise = new Promise<void>((resolve) => {
      releaseLock = resolve;
    });
    let lockAcquired!: () => void;
    const lockAcquiredPromise = new Promise<void>((resolve) => {
      lockAcquired = resolve;
    });
    const lockHolder = prisma.$transaction(async (tx) => {
      await tx.$queryRaw<Array<{ readonly locked: string }>>`
        SELECT pg_advisory_xact_lock(${LOCAL_KEK_ROTATION_ADVISORY_LOCK})::text AS "locked"
      `;
      lockAcquired();
      await releaseLockPromise;
    }, { timeout: 30_000 });
    await lockAcquiredPromise;
    try {
      await expect(rotateLocalKek({
        keyringFile,
        prisma,
        generatedKey: Buffer.alloc(32, "x"),
        expectedOwner,
      })).rejects.toThrow("rotation lock is already held");
    } finally {
      releaseLock();
      await lockHolder;
    }

    const result = await rotateLocalKek({
      keyringFile,
      prisma,
      now: new Date("2026-07-10T12:34:56.000Z"),
      generatedKey: Buffer.alloc(32, "n"),
      expectedOwner,
    });

    expect(result).toMatchObject({
      previousActiveKeyId: "local-old",
      activeKeyId: expect.stringMatching(/^local-20260710T123456Z-/),
      rewrapped: 2,
      verified: 2,
    });
    const after = await prisma.providerCredentialVersion.findUniqueOrThrow({ where: { id: before.id } });
    const afterEnvelope = after.envelope as unknown as SecretEnvelopeV1;
    expect(after.keyId).toBe(result.activeKeyId);
    expect(afterEnvelope.payload).toEqual(beforeEnvelope.payload);
    expect(afterEnvelope.wrappedDek).not.toEqual(beforeEnvelope.wrappedDek);
    const afterFeishu = await prisma.feishuApplicationCredentialVersion.findUniqueOrThrow({
      where: { id: beforeFeishu.id },
    });
    const afterFeishuEnvelope = afterFeishu.envelope as unknown as SecretEnvelopeV1;
    expect(afterFeishu.keyId).toBe(result.activeKeyId);
    expect(afterFeishuEnvelope.payload).toEqual(beforeFeishuEnvelope.payload);
    expect(afterFeishuEnvelope.wrappedDek).not.toEqual(beforeFeishuEnvelope.wrappedDek);
    const keyringSource = await readFile(keyringFile, "utf8");
    expect(keyringSource).toContain("local-old");
    expect(keyringSource).toContain(result.activeKeyId);
    expect(keyringSource).not.toContain("credential-stays-encrypted");
    expect(keyringSource).not.toContain("feishu-credential-stays-encrypted");

    const rotatedSecrets = new LocalSecretEnvelope(await loadLocalSecretKeyringFile(keyringFile));
    await expect(verifyStoredProviderEnvelopes(prisma, rotatedSecrets)).resolves.toBe(1);
    await expect(verifyStoredFeishuApplicationEnvelopes(prisma, rotatedSecrets)).resolves.toBe(1);
    expect(() => oldSecrets.decryptJson({
      purpose: "provider-connection",
      organizationId: before.connection.organizationId,
      connectionId: before.connection.id,
      secretVersionId: before.id,
    }, afterEnvelope)).toThrow("secret envelope key is unavailable");
    expect(() => oldSecrets.decryptJson({
      purpose: "feishu-application",
      organizationId: beforeFeishu.connection.organizationId,
      connectionId: beforeFeishu.connection.id,
      secretVersionId: beforeFeishu.id,
    }, afterFeishuEnvelope)).toThrow("secret envelope key is unavailable");
    await expect(prisma.auditEntry.count({
      where: { action: "provider_secret.kek_rewrapped", organizationId: DEFAULT_ORG_ID },
    })).resolves.toBe(1);
    await expect(prisma.auditEntry.count({
      where: { action: "feishu_application_secret.kek_rewrapped", organizationId: DEFAULT_ORG_ID },
    })).resolves.toBe(1);
  });
});