Preserve wallet address on vault credential rotation (#8618)

0xFirekeeper · web-flow · commit 5ca4d7497ba2 · 2026-01-12T22:27:08.000+07:00
diff --git a/apps/dashboard/src/@/hooks/useApi.ts b/apps/dashboard/src/@/hooks/useApi.ts
@@ -330,20 +330,16 @@ export async function rotateSecretKeyClient(params: { project: Project }) {
     throw new Error(res.error);
   }
 
-  try {
-    // if the project has an encrypted vault admin key, rotate it as well
-    const service = params.project.services.find(
-      (service) => service.name === "engineCloud",
-    );
-    if (service?.encryptedAdminKey) {
-      await rotateVaultAccountAndAccessToken({
-        project: params.project,
-        projectSecretKey: res.data.data.secret,
-        projectSecretHash: res.data.data.secretHash,
-      });
-    }
-  } catch (error) {
-    console.error("Failed to rotate vault admin key", error);
+  // if the project has an encrypted vault admin key, rotate it as well
+  const service = params.project.services.find(
+    (service) => service.name === "engineCloud",
+  );
+  if (service?.encryptedAdminKey) {
+    await rotateVaultAccountAndAccessToken({
+      project: params.project,
+      projectSecretKey: res.data.data.secret,
+      projectSecretHash: res.data.data.secretHash,
+    });
   }
 
   return res.data;
diff --git a/apps/dashboard/src/app/(app)/team/[team_slug]/[project_slug]/(sidebar)/transactions/lib/vault.client.ts b/apps/dashboard/src/app/(app)/team/[team_slug]/[project_slug]/(sidebar)/transactions/lib/vault.client.ts
@@ -21,6 +21,41 @@ const SERVER_WALLET_ACCESS_TOKEN_PURPOSE =
 export const SERVER_WALLET_MANAGEMENT_ACCESS_TOKEN_PURPOSE =
   "Management Token for Dashboard";
 
+/**
+ * Retry a function with exponential backoff.
+ */
+async function withRetry<T>(
+  fn: () => Promise<T>,
+  options: { maxAttempts?: number; baseDelayMs?: number } = {},
+): Promise<T> {
+  const { maxAttempts = 3, baseDelayMs = 1000 } = options;
+  if (!Number.isInteger(maxAttempts) || maxAttempts < 1) {
+    throw new Error("maxAttempts must be at least 1");
+  }
+  if (!Number.isFinite(baseDelayMs) || baseDelayMs < 0) {
+    throw new Error("baseDelayMs must be a non-negative number");
+  }
+  let lastError: Error | undefined;
+
+  for (let attempt = 1; attempt <= maxAttempts; attempt++) {
+    try {
+      return await fn();
+    } catch (error) {
+      lastError = error instanceof Error ? error : new Error(String(error));
+      if (attempt < maxAttempts) {
+        // Exponential backoff with cap at 30s and jitter to prevent thundering herd
+        const delay = Math.min(baseDelayMs * 2 ** (attempt - 1), 30_000);
+        const jitter = Math.floor(Math.random() * Math.min(250, delay));
+        await new Promise<void>((resolve) =>
+          setTimeout(resolve, delay + jitter),
+        );
+      }
+    }
+  }
+
+  throw lastError ?? new Error("withRetry failed without capturing an error");
+}
+
 let vc: VaultClient | null = null;
 
 export async function initVaultClient() {
@@ -47,6 +82,20 @@ export async function rotateVaultAccountAndAccessToken(props: {
     throw new Error("No rotation code found");
   }
 
+  // IMPORTANT: Validate secret key BEFORE rotating to prevent bricking the project.
+  // If we rotate first and then secret key validation fails, the old rotation code
+  // is consumed but we can't save the new one, leaving the project unrecoverable.
+  if (props.projectSecretKey) {
+    const projectSecretKeyHash = await hashSecretKey(props.projectSecretKey);
+    const secretKeysHashed = [
+      ...props.project.secretKeys,
+      ...(props.projectSecretHash ? [{ hash: props.projectSecretHash }] : []),
+    ];
+    if (!secretKeysHashed.some((key) => key?.hash === projectSecretKeyHash)) {
+      throw new Error("Invalid project secret key");
+    }
+  }
+
   const rotateServiceAccountRes = await rotateServiceAccount({
     client: vaultClient,
     request: {
@@ -69,6 +118,9 @@ export async function rotateVaultAccountAndAccessToken(props: {
       vaultClient,
       adminKey,
       rotationCode,
+      // Skip wallet creation on rotation - preserve the existing project wallet
+      skipWalletCreation: true,
+      existingProjectWalletAddress: service?.projectWalletAddress ?? undefined,
     });
 
   return {
@@ -222,9 +274,19 @@ async function createAndEncryptVaultAccessTokens(props: {
   projectSecretHash?: string;
   adminKey: string;
   rotationCode: string;
+  skipWalletCreation?: boolean;
+  existingProjectWalletAddress?: string;
 }) {
-  const { project, projectSecretKey, vaultClient, adminKey, rotationCode } =
-    props;
+  const {
+    project,
+    projectSecretKey,
+    projectSecretHash,
+    vaultClient,
+    adminKey,
+    rotationCode,
+    skipWalletCreation,
+    existingProjectWalletAddress,
+  } = props;
 
   const [managementTokenResult, walletTokenResult] = await Promise.all([
     createManagementAccessToken({ project, adminKey, vaultClient }),
@@ -246,81 +308,78 @@ async function createAndEncryptVaultAccessTokens(props: {
   const managementToken = managementTokenResult.data;
   const walletToken = walletTokenResult.data;
 
-  // create a default project server wallet
-  const defaultProjectServerWallet = await createProjectServerWallet({
-    project,
-    managementAccessToken: managementToken.accessToken,
-    label: getProjectWalletLabel(project.name),
-  });
+  // CRITICAL: Save credentials IMMEDIATELY after creating tokens.
+  // This prevents a broken state if wallet creation or other operations fail.
+  // The rotationCode is consumed when rotating, so if we don't save the new one,
+  // the project becomes unrecoverable.
+  let encryptedAdminKey: string | null = null;
+  let encryptedWalletAccessToken: string | null = null;
 
   if (projectSecretKey) {
     // verify that the project secret key is valid
     const projectSecretKeyHash = await hashSecretKey(projectSecretKey);
     const secretKeysHashed = [
       ...project.secretKeys,
       // for newly rotated secret keys, we don't have the secret key in the project secret keys yet
-      ...(props.projectSecretHash ? [{ hash: props.projectSecretHash }] : []),
+      ...(projectSecretHash ? [{ hash: projectSecretHash }] : []),
     ];
     if (!secretKeysHashed.some((key) => key?.hash === projectSecretKeyHash)) {
       throw new Error("Invalid project secret key");
     }
 
     // encrypt admin key and wallet token with project secret key
-    const [encryptedAdminKey, encryptedWalletAccessToken] = await Promise.all([
+    [encryptedAdminKey, encryptedWalletAccessToken] = await Promise.all([
       encrypt(adminKey, projectSecretKey),
       encrypt(walletToken.accessToken, projectSecretKey),
     ]);
+  }
 
-    await updateProjectClient(
-      {
-        projectId: props.project.id,
-        teamId: props.project.teamId,
-      },
-      {
-        services: [
-          ...props.project.services.filter(
-            (service) => service.name !== "engineCloud",
-          ),
-          {
-            name: "engineCloud",
-            actions: [],
-            managementAccessToken: managementToken.accessToken,
-            maskedAdminKey: maskSecret(adminKey),
-            encryptedAdminKey,
-            encryptedWalletAccessToken,
-            rotationCode: rotationCode,
-            projectWalletAddress: defaultProjectServerWallet.address,
-          },
-        ],
-      },
-    );
-  } else {
-    // no secret key, only store the management token, remove any encrypted keys
-    await updateProjectClient(
-      {
-        projectId: props.project.id,
-        teamId: props.project.teamId,
-      },
-      {
-        services: [
-          ...props.project.services.filter(
-            (service) => service.name !== "engineCloud",
-          ),
-          {
-            name: "engineCloud",
-            actions: [],
-            managementAccessToken: managementToken.accessToken,
-            maskedAdminKey: maskSecret(adminKey),
-            encryptedAdminKey: null,
-            encryptedWalletAccessToken: null,
-            rotationCode: rotationCode,
-            projectWalletAddress: defaultProjectServerWallet.address,
-          },
-        ],
-      },
-    );
+  // For rotation, preserve existing wallet address. For new creation, create a default wallet.
+  let projectWalletAddress: string | null | undefined =
+    existingProjectWalletAddress ??
+    project.services.find((s) => s.name === "engineCloud")
+      ?.projectWalletAddress;
+
+  // Only create a new wallet if we don't have one (initial setup, not rotation)
+  if (!skipWalletCreation && !projectWalletAddress) {
+    const defaultProjectServerWallet = await createProjectServerWallet({
+      project,
+      managementAccessToken: managementToken.accessToken,
+      label: getProjectWalletLabel(project.name),
+    });
+    projectWalletAddress = defaultProjectServerWallet.address;
   }
 
+  // Save credentials with retry - this is critical because if rotation succeeded
+  // but this save fails, the new rotation code is lost and project becomes unrecoverable
+  await withRetry(
+    () =>
+      updateProjectClient(
+        {
+          projectId: project.id,
+          teamId: project.teamId,
+        },
+        {
+          services: [
+            ...project.services.filter(
+              (service) => service.name !== "engineCloud",
+            ),
+            {
+              name: "engineCloud",
+              actions: [],
+              managementAccessToken: managementToken.accessToken,
+              maskedAdminKey: maskSecret(adminKey),
+              encryptedAdminKey,
+              encryptedWalletAccessToken,
+              rotationCode: rotationCode,
+              projectWalletAddress: projectWalletAddress ?? null,
+            },
+          ],
+        },
+      ),
+    { maxAttempts: 3, baseDelayMs: 1000 },
+  );
+
   return {
     managementToken,
     walletToken,
diff --git a/apps/dashboard/src/app/(app)/team/[team_slug]/[project_slug]/(sidebar)/wallets/server-wallets/wallets/page.tsx b/apps/dashboard/src/app/(app)/team/[team_slug]/[project_slug]/(sidebar)/wallets/server-wallets/wallets/page.tsx
@@ -9,6 +9,7 @@ import { ServerWalletsTable } from "../../../transactions/components/server-wall
 import type { Wallet } from "../../../transactions/server-wallets/wallet-table/types";
 import { listSolanaAccounts } from "../../../transactions/solana-wallets/lib/vault.client";
 import type { SolanaWallet } from "../../../transactions/solana-wallets/wallet-table/types";
+import { VaultRecoveryCard } from "./vault-recovery-card.client";
 
 export const dynamic = "force-dynamic";
 
@@ -107,12 +108,10 @@ export default async function Page(props: {
   return (
     <div className="flex flex-col gap-10">
       {eoas.error ? (
-        <div className="rounded-xl border border-destructive/50 bg-destructive/10 p-4">
-          <p className="text-destructive font-semibold mb-2">
-            EVM Wallet Error
-          </p>
-          <p className="text-sm text-muted-foreground">{eoas.error.message}</p>
-        </div>
+        <VaultRecoveryCard
+          errorMessage={eoas.error.message}
+          project={project}
+        />
       ) : (
         <ServerWalletsTable
           client={client}
diff --git a/apps/dashboard/src/app/(app)/team/[team_slug]/[project_slug]/(sidebar)/wallets/server-wallets/wallets/vault-recovery-card.client.tsx b/apps/dashboard/src/app/(app)/team/[team_slug]/[project_slug]/(sidebar)/wallets/server-wallets/wallets/vault-recovery-card.client.tsx
