feat(auth): allow skipping forced recovery key download (#900)

* feat(auth): allow skipping forced recovery key download

* refactor: move from session storage to cookie

app/client/modules/auth/routes/download-recovery-key.tsx

@@ -9,16 +9,21 @@ import { Input } from "~/client/components/ui/input";
 import { Label } from "~/client/components/ui/label";
 import { downloadResticPasswordMutation } from "~/client/api-client/@tanstack/react-query.gen";
 import { parseError } from "~/client/lib/errors";
+import {
+	RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_MAX_AGE,
+	RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_NAME,
+} from "~/lib/recovery-key-skip";
 import { useNavigate } from "@tanstack/react-router";
 
 const RECOVERY_KEY_CREDENTIAL_REQUIRED_MESSAGE =
-	"Downloading the recovery key requires a local credential password. Ask an operator to run `bun run cli reset-password` for your user, then sign in with that password and try again.";
+	"Downloading the recovery key requires a local credential password. Ask an operator to run `docker exec -it zerobyte bun run cli reset-password` for your user, then sign in with that password and try again.";
 
 type Props = {
 	hasCredentialPassword: boolean;
+	userId: string | null;
 };
 
-export function DownloadRecoveryKeyPage({ hasCredentialPassword }: Props) {
+export function DownloadRecoveryKeyPage({ hasCredentialPassword, userId }: Props) {
 	const navigate = useNavigate();
 	const [password, setPassword] = useState("");
 	const [blockedMessage, setBlockedMessage] = useState<string | null>(null);
@@ -56,11 +61,14 @@ export function DownloadRecoveryKeyPage({ hasCredentialPassword }: Props) {
 		}
 
 		setBlockedMessage(null);
-		downloadResticPassword.mutate({
-			body: {
-				password,
-			},
-		});
+		downloadResticPassword.mutate({ body: { password } });
+	};
+
+	const handleSkip = () => {
+		if (!userId) return;
+
+		document.cookie = `${RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_NAME}=${userId}; path=/; max-age=${RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_MAX_AGE}`;
+		void navigate({ to: "/volumes", replace: true });
 	};
 
 	return (
@@ -114,6 +122,15 @@ export function DownloadRecoveryKeyPage({ hasCredentialPassword }: Props) {
 							Download Recovery Key
 						</Button>
 					)}
+					<Button
+						type="button"
+						variant="ghost"
+						onClick={handleSkip}
+						disabled={downloadResticPassword.isPending}
+						className="w-full"
+					>
+						Skip
+					</Button>
 				</div>
 			</form>
 		</AuthLayout>

app/client/modules/auth/routes/login.tsx

@@ -10,6 +10,7 @@ import { InputOTP, InputOTPGroup, InputOTPSeparator, InputOTPSlot } from "~/clie
 import { Label } from "~/client/components/ui/label";
 import { authClient } from "~/client/lib/auth-client";
 import { logger } from "~/client/lib/logger";
+import { RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_NAME } from "~/lib/recovery-key-skip";
 import { decodeLoginError, getLoginErrorDescription } from "~/client/lib/sso-errors";
 import { ResetPasswordDialog } from "../components/reset-password-dialog";
 import { useNavigate } from "@tanstack/react-router";
@@ -32,6 +33,12 @@ type LoginPageProps = {
 	error?: string;
 };
 
+function hasSkippedRecoveryKeyDownload(userId: string) {
+	return document.cookie
+		.split(";")
+		.some((cookie) => cookie.trim() === `${RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_NAME}=${userId}`);
+}
+
 export function LoginPage({ error }: LoginPageProps = {}) {
 	const navigate = useNavigate();
 	const [showResetDialog, setShowResetDialog] = useState(false);
@@ -77,7 +84,7 @@ export function LoginPage({ error }: LoginPageProps = {}) {
 		}
 
 		const d = await authClient.getSession();
-		if (data.user && !d.data?.user.hasDownloadedResticPassword) {
+		if (data.user && !d.data?.user.hasDownloadedResticPassword && !hasSkippedRecoveryKeyDownload(data.user.id)) {
 			void navigate({ to: "/download-recovery-key" });
 		} else {
 			void navigate({ to: "/volumes" });
@@ -113,7 +120,11 @@ export function LoginPage({ error }: LoginPageProps = {}) {
 		if (data) {
 			toast.success("Login successful");
 			const session = await authClient.getSession();
-			if (session.data?.user && !session.data.user.hasDownloadedResticPassword) {
+			if (
+				session.data?.user &&
+				!session.data.user.hasDownloadedResticPassword &&
+				!hasSkippedRecoveryKeyDownload(session.data.user.id)
+			) {
 				void navigate({ to: "/download-recovery-key" });
 			} else {
 				void navigate({ to: "/volumes" });
@@ -130,7 +141,10 @@ export function LoginPage({ error }: LoginPageProps = {}) {
 
 	if (requires2FA) {
 		return (
-			<AuthLayout title="Two-Factor Authentication" description="Enter the 6-digit code from your authenticator app">
+			<AuthLayout
+				title="Two-Factor Authentication"
+				description="Enter the 6-digit code from your authenticator app"
+			>
 				<div className="space-y-6">
 					<div className="space-y-4 flex flex-col items-center">
 						<Label htmlFor="totp-code">Authentication code</Label>
@@ -199,7 +213,11 @@ export function LoginPage({ error }: LoginPageProps = {}) {
 		<AuthLayout title="Login to your account" description="Enter your credentials below to login to your account">
 			<Form {...form}>
 				<form onSubmit={form.handleSubmit(onSubmit)} className="space-y-4">
-					<div className={cn("rounded-md border border-destructive/50 p-3 text-sm", { hidden: !errorDescription })}>
+					<div
+						className={cn("rounded-md border border-destructive/50 p-3 text-sm", {
+							hidden: !errorDescription,
+						})}
+					>
 						{errorDescription}
 					</div>
 					<FormField

app/client/modules/settings/routes/settings.tsx

@@ -49,7 +49,7 @@ import { useOrganizationContext } from "~/client/hooks/use-org-context";
 import { cn } from "~/client/lib/utils";
 
 const RECOVERY_KEY_CREDENTIAL_REQUIRED_MESSAGE =
-	"Downloading the recovery key requires a local credential password. Ask an operator to run `bun run cli reset-password` for your user, then sign in with that password and try again.";
+	"Downloading the recovery key requires a local credential password. Ask an operator to run `docker exec -it zerobyte bun run cli reset-password` for your user, then sign in with that password and try again.";
 
 type Props = {
 	appContext: AppContext;

app/context.ts

@@ -15,4 +15,5 @@ export type AppContext = {
 	user: User | null;
 	hasUsers: boolean;
 	sidebarOpen: boolean;
+	hasSkippedRecoveryKeyDownload: boolean;
 };

app/lib/recovery-key-skip.ts

@@ -0,0 +1,2 @@
+export const RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_NAME = "zerobyte_recovery_key_download_skipped";
+export const RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_MAX_AGE = 60 * 60;

app/middleware/auth.ts

@@ -1,9 +1,10 @@
 import { createMiddleware } from "@tanstack/react-start";
 import { redirect } from "@tanstack/react-router";
 import { auth } from "~/server/lib/auth";
-import { getRequestHeaders } from "@tanstack/react-start/server";
+import { getCookie, getRequestHeaders } from "@tanstack/react-start/server";
 import { authService } from "~/server/modules/auth/auth.service";
 import { isAuthRoute } from "~/lib/auth-routes";
+import { RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_NAME } from "~/lib/recovery-key-skip";
 
 export const authMiddleware = createMiddleware().server(async ({ next, request }) => {
 	const headers = getRequestHeaders();
@@ -20,7 +21,13 @@ export const authMiddleware = createMiddleware().server(async ({ next, request }
 	}
 
 	if (session?.user?.id) {
-		if (!session.user.hasDownloadedResticPassword && pathname !== "/download-recovery-key") {
+		const hasSkippedRecoveryKeyDownload = getCookie(RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_NAME) === session.user.id;
+
+		if (
+			!session.user.hasDownloadedResticPassword &&
+			!hasSkippedRecoveryKeyDownload &&
+			pathname !== "/download-recovery-key"
+		) {
 			throw redirect({ to: "/download-recovery-key" });
 		}
 

app/routes/(auth)/download-recovery-key.tsx

@@ -11,6 +11,7 @@ const getRecoveryKeyUserState = createServerFn({ method: "GET" }).handler(async
 
 	return {
 		hasCredentialPassword: session?.user ? await userHasCredentialPassword(session.user.id) : false,
+		userId: session?.user.id ?? null,
 	};
 });
 
@@ -30,7 +31,7 @@ export const Route = createFileRoute("/(auth)/download-recovery-key")({
 });
 
 function RouteComponent() {
-	const { hasCredentialPassword } = Route.useLoaderData();
+	const { hasCredentialPassword, userId } = Route.useLoaderData();
 
-	return <DownloadRecoveryKeyPage hasCredentialPassword={hasCredentialPassword} />;
+	return <DownloadRecoveryKeyPage hasCredentialPassword={hasCredentialPassword} userId={userId} />;
 }

app/routes/(dashboard)/route.tsx

@@ -9,6 +9,7 @@ import { getOrganizationContext } from "~/server/lib/functions/organization-cont
 import { getServerConstants } from "~/server/lib/functions/server-constants";
 import { userHasCredentialPassword } from "~/server/modules/auth/helpers";
 import { authService } from "~/server/modules/auth/auth.service";
+import { RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_NAME } from "~/lib/recovery-key-skip";
 
 export const fetchUser = createServerFn({ method: "GET" }).handler(async () => {
 	const headers = getRequestHeaders();
@@ -18,11 +19,14 @@ export const fetchUser = createServerFn({ method: "GET" }).handler(async () => {
 
 	const sidebarCookie = getCookie(SIDEBAR_COOKIE_NAME);
 	const sidebarOpen = !sidebarCookie ? true : sidebarCookie === "true";
+	const hasSkippedRecoveryKeyDownload =
+		!!session?.user && getCookie(RECOVERY_KEY_DOWNLOAD_SKIPPED_COOKIE_NAME) === session.user.id;
 
 	return {
 		user: session?.user ? { ...session.user, hasCredentialPassword } : null,
 		hasUsers,
 		sidebarOpen,
+		hasSkippedRecoveryKeyDownload,
 	};
 });
 
@@ -45,7 +49,11 @@ export const Route = createFileRoute("/(dashboard)")({
 			}),
 		]);
 
-		if (authContext.user && !authContext.user.hasDownloadedResticPassword) {
+		if (
+			authContext.user &&
+			!authContext.user.hasDownloadedResticPassword &&
+			!authContext.hasSkippedRecoveryKeyDownload
+		) {
 			throw redirect({ to: "/download-recovery-key" });
 		}