From fd64578a170f6633659cc99b4c17f87f5981b34d Mon Sep 17 00:00:00 2001
From: ramvignesh-b <ramvignesh-b@github.com>
Date: Wed, 15 Apr 2026 16:57:03 +0530
Subject: [PATCH] fix: retain masterkey on non-logout error scenarios and
 refresh on db hit miss

---
 frontend/src/hooks/useAuth.test.ts |  4 +--
 frontend/src/hooks/useAuth.ts      | 39 +++++++++++++++++-----
 frontend/src/hooks/useLetters.tsx  |  9 +++--
 frontend/src/pages/Drawer.tsx      | 53 ++++++++++++++++++++++++++++--
 4 files changed, 89 insertions(+), 16 deletions(-)

diff --git a/frontend/src/hooks/useAuth.test.ts b/frontend/src/hooks/useAuth.test.ts
index b12fc0b..4d3d481 100644
--- a/frontend/src/hooks/useAuth.test.ts
+++ b/frontend/src/hooks/useAuth.test.ts
@@ -191,7 +191,7 @@ describe("initialize", () => {
     expect(useKeyStore.getState().masterKey).not.toBeNull();
   });
 
-  it("should clear both stores if the refresh attempt fails", async () => {
+  it("should preserve the master key even if the refresh attempt fails", async () => {
     server.use(
       http.post(
         `${API_URL}/api/auth/refresh/`,
@@ -206,6 +206,6 @@ describe("initialize", () => {
 
     expect(useAuthStore.getState().accessToken).toBeNull();
     expect(useAuthStore.getState().user).toBeNull();
-    expect(useKeyStore.getState().masterKey).toBeNull();
+    expect(useKeyStore.getState().masterKey).not.toBeNull();
   });
 });
diff --git a/frontend/src/hooks/useAuth.ts b/frontend/src/hooks/useAuth.ts
index 682b613..e915f28 100644
--- a/frontend/src/hooks/useAuth.ts
+++ b/frontend/src/hooks/useAuth.ts
@@ -4,6 +4,7 @@ import { endpoints } from "../config/endpoints";
 import type { UserProfile } from "../store/useAuthStore";
 import { useAuthStore } from "../store/useAuthStore";
 import { useKeyStore } from "../store/useKeyStore";
+import { CryptoUtils } from "../utils/crypto";
 import {
   clearMasterKey,
   loadMasterKey,
@@ -40,9 +41,17 @@ export const useAuth = () => {
   };
 
   const initialize = useCallback(async () => {
-    const { accessToken, user, setAuth, clearAuth, setInitializing } =
+    const { accessToken, user, setAuth, setInitializing } =
       useAuthStore.getState();
 
+    // Restore master key from IndexedDB
+    try {
+      const masterKey = await loadMasterKey();
+      if (masterKey) setMasterKey(masterKey);
+    } catch {
+      console.error("Master key restoration failed");
+    }
+
     // If session in memory, don't trigger refresh/me again
     if (accessToken && user) {
       setInitializing(false);
@@ -50,23 +59,34 @@ export const useAuth = () => {
     }
 
     try {
-      // try refresh
+      // try session refresh
       const { data: refreshData } = await publicApi.post(endpoints.REFRESH);
       const { data: userData } = await api.get(endpoints.ME, {
         headers: { Authorization: `Bearer ${refreshData.access}` },
       });
       setAuth(refreshData.access, userData);
-
-      // restore master key from IndexedDB
-      const masterKey = await loadMasterKey();
-      if (masterKey) setMasterKey(masterKey);
     } catch {
-      clearAuth();
-      setMasterKey(null);
-      await clearMasterKey();
+      // grace for temporary network errors
+    } finally {
+      setInitializing(false);
     }
   }, [setMasterKey]);
 
+  const unlock = async (password: string) => {
+    if (!user) return;
+
+    try {
+      const { masterKey } = await CryptoUtils.deriveKeyBundle(
+        password,
+        user.email,
+      );
+      await saveMasterKey(masterKey);
+      setMasterKey(masterKey);
+    } catch {
+      console.error("Master key restoration failed");
+    }
+  };
+
   return {
     isAuthenticated,
     user,
@@ -74,5 +94,6 @@ export const useAuth = () => {
     setAuthStore,
     logout,
     initialize,
+    unlock,
   };
 };
diff --git a/frontend/src/hooks/useLetters.tsx b/frontend/src/hooks/useLetters.tsx
index d0fab38..79f202e 100644
--- a/frontend/src/hooks/useLetters.tsx
+++ b/frontend/src/hooks/useLetters.tsx
@@ -56,11 +56,15 @@ async function decryptLetters(
 export function useLetters() {
   const [letters, setLetters] = useState<ProcessedLetter[]>([]);
   const [loading, setLoading] = useState(false);
+  const [isAuthRequired, setIsAuthRequired] = useState<boolean>(false);
   const { masterKey } = useKeyStore();
 
   useEffect(() => {
-    if (!masterKey) return;
-
+    if (!masterKey) {
+      setIsAuthRequired(true);
+      return;
+    }
+    setIsAuthRequired(false);
     setLoading(true);
     api
       .get(endpoints.LETTERS)
@@ -83,5 +87,6 @@ export function useLetters() {
     ...drawerItems,
     loading,
     refreshLetters: () => setLoading(true),
+    isAuthRequired,
   };
 }
diff --git a/frontend/src/pages/Drawer.tsx b/frontend/src/pages/Drawer.tsx
index 714b394..e08f873 100644
--- a/frontend/src/pages/Drawer.tsx
+++ b/frontend/src/pages/Drawer.tsx
@@ -1,4 +1,4 @@
-import { FeatherIcon } from "@phosphor-icons/react";
+import { FeatherIcon, LockKeyIcon } from "@phosphor-icons/react";
 import { useState } from "react";
 import { useNavigate } from "react-router-dom";
 import Logo from "../components/Logo";
@@ -9,10 +9,11 @@ import { useAuth } from "../hooks/useAuth";
 import { useLetters } from "../hooks/useLetters";
 
 export default function Drawer() {
-  const { user, logout } = useAuth();
+  const { user, logout, unlock } = useAuth();
+
   const [openSection, setOpenSection] = useState<string | null>();
   const navigate = useNavigate();
-  const { drafts, kept, sent, vault, loading } = useLetters();
+  const { drafts, kept, sent, vault, loading, isAuthRequired } = useLetters();
 
   if (!user) return null;
 
@@ -23,6 +24,52 @@ export default function Drawer() {
     <div className="min-h-screen w-full bg-base-100 text-base-content flex flex-col items-center py-12 px-5 pb-32 font-serif transition-colors">
       <div className="fixed inset-0 bg-[radial-gradient(circle_at_center,transparent_0%,rgba(0,0,0,0.5)_100%)] pointer-events-none z-0" />
 
+      {isAuthRequired && (
+        <div className="modal modal-open bg-base-100/20 backdrop-blur-md">
+          <div className="modal-box p-12 flex flex-col items-center">
+            <LockKeyIcon
+              size={48}
+              className="text-primary mx-auto mb-8 animate-pulse"
+            />
+            <h3 className="font-bold text-lg font-display text-primary">
+              Authentication Required
+            </h3>
+            <p className="py-4 font-sans">
+              We need your passkey to open your letters
+            </p>
+            <div className="divider w-1/2 mx-auto text-xs text-neutral-content/30 mt-0"></div>
+            <p className="text-xs text-neutral-content/30 font-mono italic">
+              P.S. We don't validate your input at the moment.
+            </p>
+            <div className="modal-action items-center gap-4">
+              <form
+                className="form-control w-full inline-flex"
+                onSubmit={async (e: React.SubmitEvent<HTMLFormElement>) => {
+                  e.preventDefault();
+                  const password = e.target.password.value;
+                  if (!password) return;
+                  unlock(password);
+                }}
+              >
+                <input
+                  name="password"
+                  required
+                  type="password"
+                  placeholder="password"
+                  className="font-sans validator input input-bordered rounded-r-none"
+                />
+                <div className="validator-message text-xs text-error"></div>
+                <button
+                  type="submit"
+                  className="btn btn-primary rounded-l-none"
+                >
+                  Unlock
+                </button>
+              </form>
+            </div>
+          </div>
+        </div>
+      )}
       <header className="text-center mb-12 z-10 animate-in fade-in slide-in-from-top-4 duration-1000">
         <Logo />
         <div className="font-sans text-xs tracking-[0.3em] uppercase text-base-content/40 mt-2">