@bitkyc08/opencodex 2.7.42 → 2.7.43

package/src/oauth/kiro-credentials.ts

@@ -1,9 +1,16 @@
-import { existsSync, readFileSync } from "node:fs";
+import { randomUUID } from "node:crypto";
+import { chmodSync, closeSync, existsSync, fsyncSync, linkSync, openSync, readFileSync, renameSync, rmSync, writeFileSync } from "node:fs";
 import { homedir } from "node:os";
 import { isAbsolute, join } from "node:path";
 import { Database } from "bun:sqlite";
 
 const DEFAULT_EXPIRES_MS = 3600_000;
+const KIRO_CLI_RECOVERY_SUFFIX = ".opencodex-recovery";
+const KIRO_CLI_RECOVERY_HEADER_V1 = Buffer.from("opencodex-kiro-session-v1\n", "utf8");
+const KIRO_CLI_RECOVERY_HEADER = Buffer.from("opencodex-kiro-session-v2\n", "utf8");
+const KIRO_CLI_RECOVERY_PROCESS_INSTANCE = randomUUID();
+const KIRO_CLI_RECOVERY_PROCESS_INSTANCE_PATTERN = /^[A-Za-z0-9._-]{1,128}$/;
+const SQLITE_DATABASE_HEADER = Buffer.from("SQLite format 3\0", "binary");
 const KIRO_REGION_PATTERN = /^[a-z]{2}(?:-[a-z]+)+-\d$/;
 const CLIENT_ID_HASH_PATTERN = /^[A-Za-z0-9_-]{1,128}$/;
 const TOKEN_KEYS = [
@@ -49,6 +56,13 @@ export interface ImportedKiroCredential {
   clientSecret?: string;
 }
 
+/** Opaque backup retained on disk while a forced Kiro CLI login is pending persistence. */
+export interface KiroCliSessionSnapshot {
+  readonly path: string;
+  readonly database: Buffer;
+  readonly recoveryPath: string;
+}
+
 type JsonObject = Record<string, unknown>;
 
 function userHome(): string {
@@ -108,18 +122,30 @@ function jsonCredentialPaths(): string[] {
     .map(expandPath);
 }
 
-function sqliteEntries(): Array<{ location: KiroImportDiagnostic["location"]; path: string }> {
+function nativeKiroCliSessionEntries(): Array<{ location: "kiro-cli-data" | "kiro-cli-linux-data"; path: string }> {
   const home = userHome();
-  const configured = process.env.KIROCLI_DB_PATH?.trim() || process.env.KIRO_CLI_DB_FILE?.trim();
-  if (configured) return [{ location: "kiro-cli-db-env", path: expandPath(configured) }];
+  // Only the stores that `kiro-cli logout` / `kiro-cli login` themselves mutate. Import fallbacks
+  // (Amazon Q / SSO cache) and KIROCLI_DB_PATH selectors must not be snapshotted for rollback.
   return [
     { location: "kiro-cli-data", path: join(home, "Library", "Application Support", "kiro-cli", "data.sqlite3") },
     { location: "kiro-cli-linux-data", path: join(home, ".local", "share", "kiro-cli", "data.sqlite3") },
-    { location: "amazon-q-data", path: join(home, ".local", "share", "amazon-q", "data.sqlite3") },
-    { location: "kiro-sso-cache", path: join(home, ".kiro", "sso", "cache.db") },
   ];
 }
 
+function sqliteEntries(): Array<{ location: KiroImportDiagnostic["location"]; path: string }> {
+  const configured = process.env.KIROCLI_DB_PATH?.trim() || process.env.KIRO_CLI_DB_FILE?.trim();
+  if (configured) return [{ location: "kiro-cli-db-env", path: expandPath(configured) }];
+  return [
+    ...nativeKiroCliSessionEntries(),
+    { location: "amazon-q-data", path: join(userHome(), ".local", "share", "amazon-q", "data.sqlite3") },
+    { location: "kiro-sso-cache", path: join(userHome(), ".kiro", "sso", "cache.db") },
+  ];
+}
+
+function kiroCliImportSelectorConfigured(): boolean {
+  return Boolean(process.env.KIROCLI_DB_PATH?.trim() || process.env.KIRO_CLI_DB_FILE?.trim());
+}
+
 function selectTokenRow(db: Database): { value: string } | null | "ambiguous" | "selected_missing" {
   const rows = db.query("SELECT key, value FROM auth_kv WHERE key LIKE ? ORDER BY key ASC").all("%:token") as Array<{ key: string; value: string }>;
   const selectedKey = process.env.KIROCLI_TOKEN_KEY?.trim();
@@ -210,7 +236,16 @@ function readStateProfile(db: Database): { profileArn?: string; apiRegion?: stri
   }
 }
 
-function readSqliteCredentials(diagnostics: KiroImportDiagnostic[]): ImportedKiroCredential | undefined {
+interface LocatedKiroCliCredential {
+  credential: ImportedKiroCredential;
+  path: string;
+  database?: Buffer;
+}
+
+function readSqliteCredentials(
+  diagnostics: KiroImportDiagnostic[],
+  includeSnapshot = false,
+): LocatedKiroCliCredential | undefined {
   for (const { location, path } of sqliteEntries()) {
     if (!existsSync(path)) {
       diagnostics.push({ location, status: "missing" });
@@ -261,7 +296,13 @@ function readSqliteCredentials(diagnostics: KiroImportDiagnostic[]): ImportedKir
       const merged = { ...registrationData, ...tokenData, ...profile };
       const credential = credentialFromJson(merged, "sqlite");
       diagnostics.push({ location, status: credential ? "token_found" : "token_missing" });
-      if (credential) return credential;
+      if (credential) {
+        return {
+          credential,
+          path,
+          ...(includeSnapshot ? { database: db.serialize() } : {}),
+        };
+      }
     } catch (error) {
       if (error instanceof Error && error.message.includes("KIROCLI_TOKEN_KEY")) throw error;
       diagnostics.push({ location, status: "schema_mismatch" });
@@ -277,12 +318,12 @@ export function inspectKiroCredentialSources(): { credential: ImportedKiroCreden
   const json = readJsonCredentials(diagnostics);
   if (json) return { credential: json, diagnostics };
   const sqlite = readSqliteCredentials(diagnostics);
-  return { credential: sqlite ?? null, diagnostics };
+  return { credential: sqlite?.credential ?? null, diagnostics };
 }
 
 export function inspectKiroCliSqliteSources(): { credential: ImportedKiroCredential | null; diagnostics: KiroImportDiagnostic[] } {
   const diagnostics: KiroImportDiagnostic[] = [];
-  return { credential: readSqliteCredentials(diagnostics) ?? null, diagnostics };
+  return { credential: readSqliteCredentials(diagnostics)?.credential ?? null, diagnostics };
 }
 
 export function readImportedKiroCredential(): ImportedKiroCredential | null {
@@ -292,3 +333,274 @@ export function readImportedKiroCredential(): ImportedKiroCredential | null {
 export function readKiroCliSqliteCredential(): ImportedKiroCredential | null {
   return inspectKiroCliSqliteSources().credential;
 }
+
+/** Capture the complete active CLI database so a failed account switch can restore it exactly. */
+export function snapshotKiroCliSession(): KiroCliSessionSnapshot | null {
+  return inspectKiroCliSessionSnapshot().snapshot;
+}
+
+/**
+ * A store that exists but cannot be captured must never be logged out of: the recovery contract
+ * promises an exact restore, and without a snapshot a later failure would destroy the session for
+ * good. `missing`/`token_missing` are the only statuses that mean "there is nothing to lose".
+ */
+const KIRO_UNSNAPSHOTTABLE_SESSION_STATUSES: ReadonlySet<KiroDiagnosticStatus> = new Set([
+  "unreadable",
+  "schema_mismatch",
+  "invalid_json",
+  "token_ambiguous",
+  "token_key_missing",
+  "token_found",
+]);
+
+/**
+ * Capture the active CLI session and report why capture failed. `blocked` is true when a session
+ * store is present but could not be snapshotted (unreadable / schema-mismatched / ambiguous), so
+ * callers can abort before mutating it.
+ *
+ * Forced-login rollback must target the exact database `kiro-cli` mutates. Custom import selectors
+ * and lower-priority Amazon Q / SSO caches are never used here: snapshotting the wrong file would
+ * leave the real CLI session switched or lost after failure.
+ */
+export function inspectKiroCliSessionSnapshot(): {
+  snapshot: KiroCliSessionSnapshot | null;
+  diagnostics: KiroImportDiagnostic[];
+  blocked: boolean;
+} {
+  const diagnostics: KiroImportDiagnostic[] = [];
+  if (kiroCliImportSelectorConfigured()) {
+    diagnostics.push({ location: "kiro-cli-db-env", status: "token_ambiguous" });
+    return { snapshot: null, diagnostics, blocked: true };
+  }
+
+  for (const { location, path } of nativeKiroCliSessionEntries()) {
+    if (!existsSync(path)) {
+      diagnostics.push({ location, status: "missing" });
+      continue;
+    }
+
+    let db: Database | undefined;
+    try {
+      db = new Database(path, { readonly: true });
+      try { db.exec("PRAGMA busy_timeout = 5000"); } catch { /* read-only best effort */ }
+    } catch {
+      diagnostics.push({ location, status: "unreadable" });
+      return { snapshot: null, diagnostics, blocked: true };
+    }
+
+    try {
+      const row = selectTokenRow(db);
+      if (row === "ambiguous") {
+        diagnostics.push({ location, status: "token_ambiguous" });
+        return { snapshot: null, diagnostics, blocked: true };
+      }
+      if (row === "selected_missing") {
+        diagnostics.push({ location, status: "token_key_missing" });
+        return { snapshot: null, diagnostics, blocked: true };
+      }
+      if (!row) {
+        // An existing CLI database with no recognized token is still what logout mutates. Falling
+        // through to Amazon Q / another platform path would snapshot the wrong store.
+        diagnostics.push({ location, status: "token_missing" });
+        return { snapshot: null, diagnostics, blocked: true };
+      }
+      try {
+        JSON.parse(row.value);
+      } catch {
+        diagnostics.push({ location, status: "invalid_json" });
+        return { snapshot: null, diagnostics, blocked: true };
+      }
+      const database = db.serialize();
+      diagnostics.push({ location, status: "token_found" });
+      return {
+        snapshot: {
+          path,
+          database,
+          recoveryPath: `${path}${KIRO_CLI_RECOVERY_SUFFIX}`,
+        },
+        diagnostics,
+        blocked: false,
+      };
+    } catch (error) {
+      if (error instanceof Error && error.message.includes("KIROCLI_TOKEN_KEY")) {
+        diagnostics.push({ location, status: "token_key_missing" });
+        return { snapshot: null, diagnostics, blocked: true };
+      }
+      diagnostics.push({ location, status: "schema_mismatch" });
+      return { snapshot: null, diagnostics, blocked: true };
+    } finally {
+      db.close();
+    }
+  }
+
+  return {
+    snapshot: null,
+    diagnostics,
+    blocked: diagnostics.some(entry => KIRO_UNSNAPSHOTTABLE_SESSION_STATUSES.has(entry.status)),
+  };
+}
+
+/** Persist a complete, private recovery image before kiro-cli is allowed to mutate its store. */
+export function persistKiroCliSessionRecovery(snapshot: KiroCliSessionSnapshot): void {
+  if (existsSync(snapshot.recoveryPath)) {
+    throw new Error("Kiro CLI session recovery is already pending.");
+  }
+  const nonce = `${process.pid}.${Date.now()}.${Math.random().toString(16).slice(2)}`;
+  const staged = `${snapshot.recoveryPath}.${nonce}.tmp`;
+  try {
+    writeFileSync(staged, Buffer.concat([
+      KIRO_CLI_RECOVERY_HEADER,
+      Buffer.from(`${process.pid}\n`, "utf8"),
+      Buffer.from(`${KIRO_CLI_RECOVERY_PROCESS_INSTANCE}\n`, "utf8"),
+      snapshot.database,
+    ]), { flag: "wx", mode: 0o600 });
+    try { chmodSync(staged, 0o600); } catch { /* platform may ignore chmod */ }
+    const fd = openSync(staged, "r+");
+    try { fsyncSync(fd); } finally { closeSync(fd); }
+    // A hard link publishes the fsynced inode atomically and, unlike rename, can never replace
+    // another process's live recovery transaction if both raced past the initial existence check.
+    linkSync(staged, snapshot.recoveryPath);
+  } finally {
+    rmSync(staged, { force: true });
+  }
+}
+
+/** Remove recovery data only after the corresponding login transaction has settled. */
+export function discardKiroCliSessionRecovery(snapshot: KiroCliSessionSnapshot): void {
+  rmSync(snapshot.recoveryPath, { force: true });
+}
+
+interface ParsedKiroCliSessionRecovery {
+  ownerPid: number;
+  ownerProcessInstance?: string;
+  database: Buffer;
+}
+
+function parseKiroCliSessionRecovery(payload: Buffer): ParsedKiroCliSessionRecovery | null {
+  const isV2 = payload.subarray(0, KIRO_CLI_RECOVERY_HEADER.length).equals(KIRO_CLI_RECOVERY_HEADER);
+  const isV1 = payload.subarray(0, KIRO_CLI_RECOVERY_HEADER_V1.length).equals(KIRO_CLI_RECOVERY_HEADER_V1);
+  if (!isV2 && !isV1) return null;
+  const headerLength = isV2 ? KIRO_CLI_RECOVERY_HEADER.length : KIRO_CLI_RECOVERY_HEADER_V1.length;
+  const ownerEnd = payload.indexOf(0x0a, headerLength);
+  if (ownerEnd <= headerLength) return null;
+  const ownerPid = Number(payload.subarray(headerLength, ownerEnd).toString("utf8"));
+  let databaseStart = ownerEnd + 1;
+  let ownerProcessInstance: string | undefined;
+  if (isV2) {
+    const instanceEnd = payload.indexOf(0x0a, databaseStart);
+    if (instanceEnd <= databaseStart) return null;
+    ownerProcessInstance = payload.subarray(databaseStart, instanceEnd).toString("utf8");
+    if (!KIRO_CLI_RECOVERY_PROCESS_INSTANCE_PATTERN.test(ownerProcessInstance)) return null;
+    databaseStart = instanceEnd + 1;
+  }
+  const database = payload.subarray(databaseStart);
+  if (
+    !Number.isSafeInteger(ownerPid) || ownerPid <= 0 ||
+    !database.subarray(0, SQLITE_DATABASE_HEADER.length).equals(SQLITE_DATABASE_HEADER)
+  ) {
+    return null;
+  }
+  return { ownerPid, ...(ownerProcessInstance ? { ownerProcessInstance } : {}), database };
+}
+
+function isKiroRecoveryOwnerAlive(ownerPid: number, ownerProcessInstance?: string): boolean {
+  // A restarted supervisor may reuse the exact same PID (commonly PID 1). Only this process
+  // instance's nonce proves that a same-PID recovery transaction is still live.
+  if (ownerPid === process.pid) return ownerProcessInstance === KIRO_CLI_RECOVERY_PROCESS_INSTANCE;
+  try {
+    process.kill(ownerPid, 0);
+    return true;
+  } catch (error) {
+    if (error && typeof error === "object" && "code" in error) return error.code !== "ESRCH";
+    return true;
+  }
+}
+
+/** Restore a previously captured CLI database after every kiro-cli child process has exited. */
+export function restoreKiroCliSession(snapshot: KiroCliSessionSnapshot): void {
+  const nonce = `${process.pid}.${Date.now()}.${Math.random().toString(16).slice(2)}`;
+  const staged = `${snapshot.path}.ocx-restore.${nonce}.tmp`;
+  const displacedBase = `${snapshot.path}.ocx-restore.${nonce}.new`;
+  const displaced: Array<{ current: string; backup: string }> = [];
+  let published = false;
+  try {
+    writeFileSync(staged, snapshot.database, { mode: 0o600 });
+    try { chmodSync(staged, 0o600); } catch { /* platform may ignore chmod */ }
+    for (const suffix of ["", "-wal", "-shm", "-journal"]) {
+      const current = `${snapshot.path}${suffix}`;
+      if (!existsSync(current)) continue;
+      const backup = `${displacedBase}${suffix || ".db"}`;
+      renameSync(current, backup);
+      displaced.push({ current, backup });
+    }
+    renameSync(staged, snapshot.path);
+    published = true;
+  } catch (error) {
+    const recoveryErrors: unknown[] = [];
+    for (const { current, backup } of [...displaced].reverse()) {
+      if (!existsSync(backup) || existsSync(current)) continue;
+      try { renameSync(backup, current); } catch (recoveryError) { recoveryErrors.push(recoveryError); }
+    }
+    if (recoveryErrors.length > 0) throw new AggregateError([error, ...recoveryErrors], "Kiro CLI session restore failed during rollback.");
+    throw error;
+  } finally {
+    rmSync(staged, { force: true });
+    if (published) {
+      for (const { backup } of displaced) {
+        try { rmSync(backup, { force: true }); } catch { /* restored database is already published */ }
+      }
+    }
+  }
+}
+
+/** Restore a transaction abandoned by a crashed process before starting another forced login. */
+export function restoreStaleKiroCliSessionRecovery(): boolean {
+  for (const { path } of nativeKiroCliSessionEntries()) {
+    const recoveryPath = `${path}${KIRO_CLI_RECOVERY_SUFFIX}`;
+    if (!existsSync(recoveryPath)) continue;
+    const payload = readFileSync(recoveryPath);
+    const recovery = parseKiroCliSessionRecovery(payload);
+    if (!recovery) {
+      throw new Error(
+        `Kiro CLI session recovery data is invalid: ${recoveryPath}. Remove this file to continue.`,
+      );
+    }
+    if (isKiroRecoveryOwnerAlive(recovery.ownerPid, recovery.ownerProcessInstance)) {
+      throw new Error(
+        `Another Kiro CLI login transaction is still in progress (pid ${recovery.ownerPid}, ${recoveryPath}).`,
+      );
+    }
+    // Atomically claim the marker before restoring so a concurrent process cannot restore the
+    // same stale image (and later delete a newer recovery file published by the winner).
+    const claimedPath = `${recoveryPath}.claimed.${process.pid}.${Date.now()}`;
+    try {
+      renameSync(recoveryPath, claimedPath);
+    } catch {
+      continue;
+    }
+    try {
+      const claimed = readFileSync(claimedPath);
+      if (!claimed.equals(payload)) {
+        // Put the unexpected claim back so an operator / later process can inspect it.
+        try { renameSync(claimedPath, recoveryPath); } catch { /* keep claimed path for manual recovery */ }
+        throw new Error(
+          `Kiro CLI session recovery data changed while claiming ${recoveryPath}. Remove leftover claim files under the kiro-cli data directory and retry.`,
+        );
+      }
+      restoreKiroCliSession({
+        path,
+        database: recovery.database,
+        recoveryPath: claimedPath,
+      });
+      rmSync(claimedPath, { force: true });
+      return true;
+    } catch (error) {
+      // Preserve the only copy of the prior database when restore fails (disk full, permissions).
+      if (existsSync(claimedPath) && !existsSync(recoveryPath)) {
+        try { renameSync(claimedPath, recoveryPath); } catch { /* claimed file remains for manual recovery */ }
+      }
+      throw error;
+    }
+  }
+  return false;
+}