feat(deployment): re-wrap data keys onto the target key version

apps/api/src/app/console.ts

@@ -14,6 +14,7 @@ import { ExecutionContextService } from "@src/core/services/execution-context/ex
 import { TopUpDeploymentsController } from "@src/deployment/controllers/deployment/top-up-deployments.controller";
 import { GpuBotController } from "@src/deployment/controllers/gpu-bot/gpu-bot.controller";
 import { ProviderController } from "@src/provider/controllers/provider/provider.controller";
+import { DataKeyRotationController } from "@src/secret/controllers/data-key-rotation/data-key-rotation.controller";
 import { WorkloadAbuseController } from "@src/workload-abuse/controllers/workload-abuse.controller";
 import { APP_INITIALIZER, ON_APP_START } from "../core/providers/app-initializer";
 
@@ -122,6 +123,18 @@ program
     });
   });
 
+program
+  .command("rotate-data-keys")
+  .description("Re-wrap every user's data key onto the configured KMS key version, proving no stored secret changed")
+  .requiredOption("-t, --target-version <version>", "The key version to move data keys onto, which must be the one this console is configured to wrap under")
+  .option("-b, --batch-size <number>", "How many data keys are written per transaction", value => z.number({ coerce: true }).parse(value))
+  .option("-d, --dry-run", "Log what would be re-wrapped without opening or writing anything", false)
+  .action(async (options, command) => {
+    await executeCliHandler(command.name(), async () => {
+      return container.resolve(DataKeyRotationController).rotate(options);
+    });
+  });
+
 program
   .command("link-auth0-accounts")
   .description("Interactively link a secondary Auth0 account into a primary one (Auth0 only, no DB writes)")

apps/api/src/deployment/providers/kms.provider.ts

@@ -12,6 +12,7 @@ const CLOUD_PLATFORM_SCOPE = "https://www.googleapis.com/auth/cloud-platform";
 
 /** The Cloud KMS operations the console performs on the SDL secrets key, narrowed so they can be doubled in tests. */
 export interface SdlSecretsKmsClient {
+  getCryptoKeyVersion(request: { name: string }, options?: CallOptions): Promise<[protos.google.cloud.kms.v1.ICryptoKeyVersion, ...unknown[]]>;
   getPublicKey(request: { name: string }, options?: CallOptions): Promise<[protos.google.cloud.kms.v1.IPublicKey, ...unknown[]]>;
   asymmetricDecrypt(request: {
     name: string;
@@ -28,6 +29,7 @@ export interface SdlSecretsKmsClient {
  */
 export interface SdlSecretsKmsTarget {
   client: SdlSecretsKmsClient;
+  version: string;
   versionName: string;
   kid: string;
   resolveVersionName(kid: unknown): string | undefined;
@@ -53,6 +55,7 @@ export function createSdlSecretsKmsTarget(input: {
 
   return {
     client,
+    version,
     versionName: versionPath(version),
     kid: `${key}.v${version}`,
     resolveVersionName(kid) {

apps/api/src/secret/config/key-rotation.config.ts

@@ -0,0 +1,4 @@
+/** Cloud KMS answers a version's state as the enum's name over REST and as its ordinal over gRPC, and both reach here unchanged. */
+export const ENABLED_CRYPTO_KEY_VERSION_STATES: ReadonlySet<unknown> = new Set(["ENABLED", 1]);
+
+export const DEFAULT_KEY_ROTATION_BATCH_SIZE = 100;

apps/api/src/secret/controllers/data-key-rotation/data-key-rotation.controller.ts

@@ -0,0 +1,13 @@
+import { singleton } from "tsyringe";
+
+import type { RotateDataKeysOptions } from "@src/secret/services/data-key-rotation/data-key-rotation.service";
+import { DataKeyRotationService } from "@src/secret/services/data-key-rotation/data-key-rotation.service";
+
+@singleton()
+export class DataKeyRotationController {
+  constructor(private readonly dataKeyRotationService: DataKeyRotationService) {}
+
+  async rotate(options: RotateDataKeysOptions) {
+    return this.dataKeyRotationService.rotate(options);
+  }
+}

apps/api/src/secret/repositories/data-key/data-key.repository.ts

@@ -1,3 +1,4 @@
+import { and, asc, count, eq, gt, ne, sql } from "drizzle-orm";
 import { singleton } from "tsyringe";
 
 import { type ApiPgDatabase, type ApiPgTables, InjectPg, InjectPgTable } from "@src/core/providers";
@@ -51,4 +52,57 @@ export class DataKeyRepository extends BaseRepository<Table, DataKeyInput, DataK
   async countWrappedUnder(wrappedByKid: DataKeyOutput["wrappedByKid"]): Promise<number> {
     return this.count({ wrappedByKid });
   }
+
+  /** Where the fleet stands across every version at once, so an operator reads what is left behind without asking version by version. */
+  async countByWrappingVersion(): Promise<Record<string, number>> {
+    const rows = await this.cursor
+      .select({ wrappedByKid: this.table.wrappedByKid, total: count() })
+      .from(this.table)
+      .groupBy(this.table.wrappedByKid)
+      .orderBy(asc(this.table.wrappedByKid));
+
+    return Object.fromEntries(rows.map(row => [row.wrappedByKid, row.total]));
+  }
+
+  /**
+   * Keyset rather than offset, because every committed batch removes its own rows from this
+   * predicate: an offset walked forward over the shrinking remainder would step over as many rows
+   * as it moved. A row this run cannot move keeps its wrapping, so it is read again on the next
+   * pass — which is the whole of the rotation's resumability.
+   */
+  async *findNotWrappedUnderIteratively({ wrappedByKid, batchSize }: { wrappedByKid: string; batchSize: number }): AsyncGenerator<DataKeyOutput[]> {
+    let cursor: string | undefined;
+
+    while (true) {
+      const batch = await this.cursor
+        .select()
+        .from(this.table)
+        .where(and(ne(this.table.wrappedByKid, wrappedByKid), ...(cursor ? [gt(this.table.id, cursor)] : [])))
+        .orderBy(asc(this.table.id))
+        .limit(batchSize);
+
+      if (!batch.length) return;
+
+      yield this.toOutputList(batch);
+
+      if (batch.length < batchSize) return;
+
+      cursor = batch[batch.length - 1].id;
+    }
+  }
+
+  /**
+   * The wrapping this re-wrap opened is compared inside the statement's own WHERE, so a row another
+   * writer moved first is left as that writer wrote it rather than overwritten with a key wrapped
+   * from a reading that is no longer current.
+   */
+  async rewrapIfStillWrappedUnder(input: { id: string; wrappedUnder: string; wrappedKey: string; wrappedByKid: string }): Promise<boolean> {
+    const [rewrapped] = await this.cursor
+      .update(this.table)
+      .set({ wrappedKey: input.wrappedKey, wrappedByKid: input.wrappedByKid, updatedAt: sql`now()` })
+      .where(and(eq(this.table.id, input.id), eq(this.table.wrappedByKid, input.wrappedUnder)))
+      .returning({ id: this.table.id });
+
+    return !!rewrapped;
+  }
 }