Zustand Persisted Store Issue with Encryption and Decryption Not Hydrating Correctly

[Rutvik-Sanathra]

Hi @dai-shi, @dbritto-dev, I really appreciate your hard work on Zustand!,
However, I'm facing an issue with Zustand's persisted store when applying encryption and decryption. The store doesn't hydrate correctly after a page refresh, even though the data is stored in localStorage. Any help would be greatly appreciated!

CODE

import { create } from 'zustand';
import { persist, createJSONStorage } from 'zustand/middleware';
import { encryptData, decryptData } from '@/utils/encryption';

const encryptedStorage = {
  getItem: (name: string) => {
    const encryptedValue = localStorage.getItem(name);
    return encryptedValue ? decryptData(encryptedValue) : null;
  },
  setItem: (name: string, value: string) => {
    const encryptedValue = encryptData(value);
    localStorage.setItem(name, encryptedValue);
  },
  removeItem: (name: string) => {
    localStorage.removeItem(name);
  },
};

export type UserType = {
  token: string;
  first_name: string;
  last_name: string;
  role: string;
  email: string;
  employee_code: string;
  mobile: string;
};

interface StoreState {
  user: UserType | null;
  hydrated: boolean;
  setUser: (user: UserType | null) => void;
  resetUser: () => void;
  setHydrated: (state: boolean) => void;
}

const useStore = create<StoreState>()(
  persist(
    (set) => ({
      user: null,
      hydrated: false,
      setUser: (user) => set({ user }),
      resetUser: () => set({ user: null }),
      setHydrated: (state) => set({ hydrated: state }),
    }),
    {
      name: 'user',
      storage: createJSONStorage(() => encryptedStorage),
      onRehydrateStorage:  (state) => {
        if (state) {
          state.setHydrated(true);
        }
      },
    }
  )
);

export const useUser = () => useStore((state) => state.user);
export const useSetUser = () => useStore((state) => state.setUser);
export const useResetUser = () => useStore((state) => state.resetUser);
export const useHydrated = () => useStore((state) => state.hydrated);

AUTH GUARD FILE

import { useUser, useHydrated } from '@/store';
import { Navigate, Outlet, useLocation } from 'react-router';

const roleBasedRedirects: Record<string, string> = {
  'SUPER ADMIN': '/super-admin',
  ADMIN: '/admin',
  BSO: '/bso',
};

const AuthGuard = () => {
  const user = useUser();
  const isHydrated = useHydrated();
  const location = useLocation();

  console.log(isHydrated,"isHydrated") // always getting false
  
  if (!isHydrated) {
    return <p>Loading...</p>
  }

  console.log(user, 'User after hydration');

  if (!user?.token) return <Navigate to="/auth/login/sso" replace />;

  const defaultRedirect = roleBasedRedirects[user?.role?.toUpperCase()];
  if (!location.pathname.startsWith(defaultRedirect)) {
    return <Navigate to={defaultRedirect} replace />;
  }

  return <Outlet />;
};

export default AuthGuard;