[Dashboard] Guard the vault rotation dialog against losing new keys

apps/dashboard/src/app/(app)/team/[team_slug]/[project_slug]/(sidebar)/vault/components/rotate-admin-key.client.tsx

@@ -9,7 +9,7 @@ import {
   LogOutIcon,
   RefreshCcwIcon,
 } from "lucide-react";
-import { useState } from "react";
+import { useEffect, useState } from "react";
 import { toast } from "sonner";
 import { rotateVaultServiceAccount } from "@/actions/vault";
 import type { Project } from "@/api/project/projects";
@@ -38,6 +38,8 @@ export default function RotateAdminKeyButton(props: {
   const [modalOpen, setModalOpen] = useState(false);
   const [keysConfirmed, setKeysConfirmed] = useState(false);
   const [keysDownloaded, setKeysDownloaded] = useState(false);
+  const [adminKeyCopied, setAdminKeyCopied] = useState(false);
+  const [accessTokenCopied, setAccessTokenCopied] = useState(false);
   const [stayManaged, setStayManaged] = useState(props.isManagedVault);
   const [secretKeyInput, setSecretKeyInput] = useState("");
   const router = useDashboardRouter();
@@ -61,22 +63,61 @@ export default function RotateAdminKeyButton(props: {
     },
   });
 
-  // Only an ejected vault returns the new keys; a managed vault keeps them.
-  const rotatedKeys =
-    rotateAdminKeyMutation.data?.adminKey &&
-    rotateAdminKeyMutation.data.walletAccessToken
+  const rotationResult = rotateAdminKeyMutation.data;
+
+  // An ejected vault keeps no server-side copy of its credentials, so this
+  // response is the only time they exist outside the enclave. A managed vault
+  // re-seals them with the project secret key and returns only the mask.
+  const ejectedKeys =
+    rotationResult &&
+    !rotationResult.isManagedVault &&
+    rotationResult.adminKey &&
+    rotationResult.walletAccessToken
       ? {
-          adminKey: rotateAdminKeyMutation.data.adminKey,
-          walletAccessToken: rotateAdminKeyMutation.data.walletAccessToken,
+          adminKey: rotationResult.adminKey,
+          walletAccessToken: rotationResult.walletAccessToken,
         }
       : undefined;
 
+  // The rotation went through but the response is missing values an ejected
+  // vault must return. Nothing can re-derive them, so this is a dead end.
+  const ejectedKeysMissing =
+    !!rotationResult && !rotationResult.isManagedVault && !ejectedKeys;
+
+  // Closing discards the only copy of the new credentials: while the request is
+  // in flight the response has nowhere to land, and once it lands the user has
+  // not stored it yet. Both states must hold the dialog open.
+  // The confirm checkbox is the only way out of this dialog, and closing it
+  // discards the response. Require the keys to actually leave the screen
+  // first: a download, or a copy of both values.
+  const keysCaptured =
+    keysDownloaded || (adminKeyCopied && accessTokenCopied);
+
+  const closeBlocked =
+    rotateAdminKeyMutation.isPending || (!!ejectedKeys && !keysConfirmed);
+
+  // eslint-disable-next-line no-restricted-syntax
+  useEffect(() => {
+    if (!closeBlocked) {
+      return;
+    }
+
+    const warnBeforeUnload = (event: BeforeUnloadEvent) => {
+      event.preventDefault();
+      // Required by Safari and older Chrome to show the prompt.
+      event.returnValue = "";
+    };
+
+    window.addEventListener("beforeunload", warnBeforeUnload);
+    return () => window.removeEventListener("beforeunload", warnBeforeUnload);
+  }, [closeBlocked]);
+
   const handleDownloadKeys = () => {
-    if (!rotatedKeys) {
+    if (!ejectedKeys) {
       return;
     }
 
-    const fileContent = `Project:\n${props.project.name} (${props.project.publishableKey})\n\nVault Admin Key:\n${rotatedKeys.adminKey}\n\nVault Access Token:\n${rotatedKeys.walletAccessToken}\n`;
+    const fileContent = `Project:\n${props.project.name} (${props.project.publishableKey})\n\nVault Admin Key:\n${ejectedKeys.adminKey}\n\nVault Access Token:\n${ejectedKeys.walletAccessToken}\n`;
     const blob = new Blob([fileContent], { type: "text/plain;charset=utf-8" });
     const url = URL.createObjectURL(blob);
     const link = document.createElement("a");
@@ -95,7 +136,7 @@ export default function RotateAdminKeyButton(props: {
   };
 
   const handleCloseModal = () => {
-    if (rotatedKeys && !keysConfirmed) {
+    if (closeBlocked) {
       return;
     }
 
@@ -129,10 +170,33 @@ export default function RotateAdminKeyButton(props: {
         Rotate Admin Key
       </Button>
 
-      <Dialog modal={true} onOpenChange={handleCloseModal} open={modalOpen}>
+      <Dialog
+        modal={true}
+        onOpenChange={(open) => {
+          if (!open) {
+            handleCloseModal();
+          }
+        }}
+        open={modalOpen}
+      >
         <DialogContent
           className="overflow-hidden p-0"
-          dialogCloseClassName={cn(rotatedKeys && !keysConfirmed && "hidden")}
+          dialogCloseClassName={cn(closeBlocked && "hidden")}
+          onEscapeKeyDown={(event) => {
+            if (closeBlocked) {
+              event.preventDefault();
+            }
+          }}
+          onInteractOutside={(event) => {
+            if (closeBlocked) {
+              event.preventDefault();
+            }
+          }}
+          onPointerDownOutside={(event) => {
+            if (closeBlocked) {
+              event.preventDefault();
+            }
+          }}
         >
           {rotateAdminKeyMutation.isPending ? (
             <>
@@ -142,11 +206,11 @@ export default function RotateAdminKeyButton(props: {
               <div className="flex flex-col items-center justify-center gap-4 p-10">
                 <Spinner className="size-8" />
                 <p className="text-muted-foreground text-xs">
-                  This may take a few seconds.
+                  This may take a few seconds. Keep this dialog open.
                 </p>
               </div>
             </>
-          ) : rotatedKeys ? (
+          ) : ejectedKeys ? (
             <div>
               <DialogHeader className="p-6">
                 <DialogTitle>New Vault Keys</DialogTitle>
@@ -162,8 +226,9 @@ export default function RotateAdminKeyButton(props: {
                       <CopyTextButton
                         className="!h-auto w-full justify-between bg-background px-3 py-3 font-mono text-xs"
                         copyIconPosition="right"
-                        textToCopy={rotatedKeys.adminKey}
-                        textToShow={maskSecret(rotatedKeys.adminKey)}
+                        onClick={() => setAdminKeyCopied(true)}
+                        textToCopy={ejectedKeys.adminKey}
+                        textToShow={maskSecret(ejectedKeys.adminKey)}
                         tooltip="Copy Admin Key"
                       />
                       <p className="text-muted-foreground text-xs">
@@ -180,8 +245,9 @@ export default function RotateAdminKeyButton(props: {
                       <CopyTextButton
                         className="!h-auto w-full justify-between bg-background px-3 py-3 font-mono text-xs"
                         copyIconPosition="right"
-                        textToCopy={rotatedKeys.walletAccessToken}
-                        textToShow={maskSecret(rotatedKeys.walletAccessToken)}
+                        onClick={() => setAccessTokenCopied(true)}
+                        textToCopy={ejectedKeys.walletAccessToken}
+                        textToShow={maskSecret(ejectedKeys.walletAccessToken)}
                         tooltip="Copy Vault Access Token"
                       />
                       <p className="text-muted-foreground text-xs">
@@ -218,10 +284,16 @@ export default function RotateAdminKeyButton(props: {
                   <CheckboxWithLabel className="text-foreground">
                     <Checkbox
                       checked={keysConfirmed}
+                      disabled={!keysCaptured}
                       onCheckedChange={(v) => setKeysConfirmed(!!v)}
                     />
                     I confirm that I've securely stored these keys
                   </CheckboxWithLabel>
+                  {!keysCaptured && (
+                    <p className="mt-2 text-muted-foreground text-xs">
+                      Download the keys, or copy both values, to continue.
+                    </p>
+                  )}
                 </Alert>
               </div>
 
@@ -235,7 +307,47 @@ export default function RotateAdminKeyButton(props: {
                 </Button>
               </div>
             </div>
-          ) : rotateAdminKeyMutation.data ? (
+          ) : ejectedKeysMissing ? (
+            <div>
+              <DialogHeader className="p-6">
+                <DialogTitle>Rotation completed without new keys</DialogTitle>
+              </DialogHeader>
+
+              <div className="space-y-6 p-6 pt-0">
+                <Alert variant="destructive">
+                  <CircleAlertIcon className="size-5" />
+                  <AlertTitle>New keys were not returned</AlertTitle>
+                  <AlertDescription>
+                    The rotation went through, so your previous admin key and
+                    access tokens no longer work, but the replacements did not
+                    reach this page and cannot be shown again. Your vault and
+                    server wallets are unaffected. Contact support to have new
+                    credentials issued for this project.
+                  </AlertDescription>
+                </Alert>
+
+                {rotationResult?.maskedAdminKey && (
+                  <div>
+                    <h3 className="mb-2 font-medium text-sm">
+                      New Vault Admin Key
+                    </h3>
+                    <div className="flex w-full items-center rounded-lg border bg-background px-3 py-3 font-mono text-xs">
+                      {rotationResult.maskedAdminKey}
+                    </div>
+                    <p className="mt-2 text-muted-foreground text-xs">
+                      Give support this masked key to identify the rotation.
+                    </p>
+                  </div>
+                )}
+              </div>
+
+              <div className="flex justify-end gap-3 border-t bg-card px-6 py-4">
+                <Button onClick={handleCloseModal} variant="primary">
+                  Close
+                </Button>
+              </div>
+            </div>
+          ) : rotationResult?.isManagedVault ? (
             <div>
               <DialogHeader className="p-6">
                 <DialogTitle>Admin Key Rotated</DialogTitle>
@@ -248,7 +360,7 @@ export default function RotateAdminKeyButton(props: {
                   </h3>
                   <div className="flex flex-col gap-2">
                     <div className="flex w-full items-center rounded-lg border bg-background px-3 py-3 font-mono text-xs">
-                      {rotateAdminKeyMutation.data.maskedAdminKey}
+                      {rotationResult.maskedAdminKey}
                     </div>
                     <p className="text-muted-foreground text-xs">
                       Your admin key and wallet access token were re-encrypted
@@ -344,7 +456,7 @@ export default function RotateAdminKeyButton(props: {
                     <AlertDescription>
                       {willStayManaged
                         ? "This will invalidate your current admin key and all existing access tokens. Your stored credentials will be re-encrypted automatically, so your server wallets keep working."
-                        : "This action will invalidate your current admin key and all existing access tokens. You will need to update your backend to use these new access tokens."}
+                        : "This action will invalidate your current admin key and all existing access tokens. You will need to update your backend to use these new access tokens. The new keys are shown once and cannot be recovered, so keep this dialog open until you have stored them."}
                     </AlertDescription>
                   </Alert>
                 </div>