Enhance wallet session handling and authorization checks

- Updated the RootLayout component to include a delay before refetching the wallet session after authorization, ensuring cookies are set properly.
- Improved the wallet queries in the PageWallets and useUserWallets hooks to only execute when the user is authorized, preventing 403 errors.
- Added additional error handling for wallet address mismatches in the wallet router, enhancing security and user experience.
- Implemented loading states and delays for smoother transitions during data fetching.

Author: QSchlegel

src/components/common/overall-layout/layout.tsx

@@ -371,19 +371,35 @@ export default function RootLayout({
     // Don't refetch here - let the natural query refetch handle it if needed
   }, []);
 
-  const handleAuthModalAuthorized = useCallback(() => {
+  const handleAuthModalAuthorized = useCallback(async () => {
     setShowAuthModal(false);
     setCheckingSession(false);
     setHasCheckedSession(true); // Mark as checked so we don't check again
     // Show loading skeleton for smooth transition
     setShowPostAuthLoading(true);
-    // Refetch session after authorization to update state (but don't show modal again)
-    void refetchWalletSession();
+    
+    // Wait a moment for the cookie to be set by the browser, then refetch session
+    await new Promise(resolve => setTimeout(resolve, 200));
+    
+    // Refetch session to update state
+    await refetchWalletSession();
+    
+    // Invalidate wallet queries so they refetch with the new session
+    // Use a small delay to ensure cookie is available on subsequent requests
+    setTimeout(() => {
+      const userAddressForInvalidation = userAddress || address;
+      if (userAddressForInvalidation) {
+        void ctx.wallet.getUserWallets.invalidate({ address: userAddressForInvalidation });
+        void ctx.wallet.getUserNewWallets.invalidate({ address: userAddressForInvalidation });
+        void ctx.wallet.getUserNewWalletsNotOwner.invalidate({ address: userAddressForInvalidation });
+      }
+    }, 300);
+    
     // Hide loading after a brief delay to allow data to load
     setTimeout(() => {
       setShowPostAuthLoading(false);
-    }, 1000);
-  }, [refetchWalletSession]);
+    }, 1500);
+  }, [refetchWalletSession, ctx.wallet, userAddress, address]);
 
   // Memoize computed route values
   const isWalletPath = useMemo(() => router.pathname.includes("/wallets/[wallet]"), [router.pathname]);

src/components/pages/homepage/wallets/index.tsx

@@ -29,10 +29,21 @@ export default function PageWallets() {
   const [showArchived, setShowArchived] = useState(false);
   const userAddress = useUserStore((state) => state.userAddress);
 
+  // Check wallet session authorization before enabling queries
+  const { data: walletSession } = api.auth.getWalletSession.useQuery(
+    { address: userAddress ?? "" },
+    {
+      enabled: !!userAddress && userAddress.length > 0,
+      refetchOnWindowFocus: false,
+    },
+  );
+  const isAuthorized = walletSession?.authorized ?? false;
+
   const { data: newPendingWallets, isLoading: isLoadingNewWallets } = api.wallet.getUserNewWallets.useQuery(
     { address: userAddress! },
     {
-      enabled: userAddress !== undefined,
+      // Only enable query when user is authorized (prevents 403 errors)
+      enabled: userAddress !== undefined && isAuthorized,
       retry: (failureCount, error) => {
         // Don't retry on authorization errors (403)
         if (error && typeof error === "object") {
@@ -60,15 +71,24 @@ export default function PageWallets() {
     api.wallet.getUserNewWalletsNotOwner.useQuery(
       { address: userAddress! },
       {
-        enabled: userAddress !== undefined,
+        // Only enable query when user is authorized (prevents 403 errors)
+        enabled: userAddress !== undefined && isAuthorized,
         retry: (failureCount, error) => {
           // Don't retry on authorization errors (403)
           if (error && typeof error === "object") {
-            const err = error as { code?: string; message?: string; data?: { code?: string } };
+            const err = error as { 
+              code?: string; 
+              message?: string; 
+              data?: { code?: string; httpStatus?: number };
+              shape?: { code?: string; message?: string };
+            };
+            const errorMessage = err.message || err.shape?.message || "";
             const isAuthError =
               err.code === "FORBIDDEN" ||
               err.data?.code === "FORBIDDEN" ||
-              err.message?.includes("Address mismatch");
+              err.data?.httpStatus === 403 ||
+              err.shape?.code === "FORBIDDEN" ||
+              errorMessage.includes("Address mismatch");
             if (isAuthError) return false;
           }
           return failureCount < 1;

src/hooks/useUserWallets.ts

@@ -7,10 +7,22 @@ import { DbWalletWithLegacy } from "@/types/wallet";
 export default function useUserWallets() {
   const network = useSiteStore((state) => state.network);
   const userAddress = useUserStore((state) => state.userAddress);
+  
+  // Check wallet session authorization before enabling queries
+  const { data: walletSession } = api.auth.getWalletSession.useQuery(
+    { address: userAddress ?? "" },
+    {
+      enabled: !!userAddress && userAddress.length > 0,
+      refetchOnWindowFocus: false,
+    },
+  );
+  const isAuthorized = walletSession?.authorized ?? false;
+  
   const { data: wallets, isLoading } = api.wallet.getUserWallets.useQuery(
     { address: userAddress! },
     {
-      enabled: userAddress !== undefined,
+      // Only enable query when user is authorized (prevents 403 errors)
+      enabled: userAddress !== undefined && isAuthorized,
       staleTime: 1 * 60 * 1000, // 1 minute (user/wallet data)
       gcTime: 5 * 60 * 1000, // 5 minutes
       retry: (failureCount, error) => {

src/server/api/routers/wallets.ts

@@ -98,9 +98,12 @@ export const walletRouter = createTRPCRouter({
         : ctx.sessionAddress
           ? [ctx.sessionAddress]
           : [];
+      // If user has an active session, validate that the requested address is authorized
+      // Throw error if address doesn't match (security: prevent unauthorized access)
       if (addresses.length > 0 && !addresses.includes(input.address)) {
         throw new TRPCError({ code: "FORBIDDEN", message: "Address mismatch" });
       }
+      // Query wallets where the user is a signer
       return ctx.db.wallet.findMany({
         where: {
           signersAddresses: {
@@ -119,6 +122,8 @@ export const walletRouter = createTRPCRouter({
         : ctx.sessionAddress
           ? [ctx.sessionAddress]
           : [];
+      // If user has an active session, validate that the requested address is authorized
+      // Throw error if address doesn't match (security: prevent unauthorized access)
       if (addresses.length > 0 && !addresses.includes(input.address)) {
         throw new TRPCError({ code: "FORBIDDEN", message: "Address mismatch" });
       }
@@ -268,9 +273,12 @@ export const walletRouter = createTRPCRouter({
         : ctx.sessionAddress
           ? [ctx.sessionAddress]
           : [];
+      // If user has an active session, validate that the requested address is authorized
+      // Throw error if address doesn't match (security: prevent unauthorized access)
       if (addresses.length > 0 && !addresses.includes(input.address)) {
         throw new TRPCError({ code: "FORBIDDEN", message: "Address mismatch" });
       }
+      // Query new wallets owned by the user
       return ctx.db.newWallet.findMany({
         where: {
           ownerAddress: input.address,
@@ -287,9 +295,12 @@ export const walletRouter = createTRPCRouter({
         : ctx.sessionAddress
           ? [ctx.sessionAddress]
           : [];
+      // If user has an active session, validate that the requested address is authorized
+      // Throw error if address doesn't match (security: prevent unauthorized access)
       if (addresses.length > 0 && !addresses.includes(input.address)) {
         throw new TRPCError({ code: "FORBIDDEN", message: "Address mismatch" });
       }
+      // Query new wallets where user is a signer but not the owner
       return ctx.db.newWallet.findMany({
         where: {
           signersAddresses: {