From 69a981568f7f304ee36dda2915e1233c3d17534c Mon Sep 17 00:00:00 2001
From: Tim de Wolf <tpdewolf@gmail.com>
Date: Mon, 20 Mar 2023 15:41:15 +0100
Subject: [PATCH] fix: Set initial data with useEffect

---
 src/OidcJwtProvider.tsx | 21 ++++++++++++++++++++-
 src/store.ts            | 33 +++++++++++++--------------------
 2 files changed, 33 insertions(+), 21 deletions(-)

diff --git a/src/OidcJwtProvider.tsx b/src/OidcJwtProvider.tsx
index c5f40cb..30bf3b0 100644
--- a/src/OidcJwtProvider.tsx
+++ b/src/OidcJwtProvider.tsx
@@ -1,8 +1,15 @@
 import React, { createContext, useContext, useEffect, useRef, useState } from 'react';
 import { StoreApi, UseBoundStore, useStore } from 'zustand';
 
+import {
+  CSRF_TOKEN_STORAGE_KEY,
+  LOGGED_IN_TOKEN_STORAGE_KEY,
+  RETRY_LOGIN_STORAGE_KEY,
+  USER_INFO_TOKEN_STORAGE_KEY,
+} from './constants';
+import { Storage } from './storage';
 import { createOidcJwtClientStore } from './store';
-import { OidcJwtClientStore, OidcJwtProviderProps } from './types';
+import { AuthState, OidcJwtClientStore, OidcJwtProviderProps } from './types';
 import { removeTokenFromUrl } from './utils';
 
 const OidcJwtContext = createContext<UseBoundStore<StoreApi<OidcJwtClientStore>> | undefined>(
@@ -78,6 +85,18 @@ const OidcJwtProvider: React.FC<React.PropsWithChildren<OidcJwtProviderProps>> =
   const { client, removeTokenFromUrlFunction = removeTokenFromUrl, children } = props;
   const store = useRef(createOidcJwtClientStore(client, removeTokenFromUrlFunction)).current;
 
+  useEffect(() => {
+    const initialState: AuthState = {
+      userInfo: Storage.get(USER_INFO_TOKEN_STORAGE_KEY),
+      csrfToken: Storage.get(CSRF_TOKEN_STORAGE_KEY),
+      isLoggedIn: !!Storage.get(LOGGED_IN_TOKEN_STORAGE_KEY),
+      isInitialized: !client,
+      didRetryLogin: Storage.get(RETRY_LOGIN_STORAGE_KEY) === 1,
+    };
+
+    store.setState({ authState: initialState });
+  }, [client, store]);
+
   return (
     <OidcJwtContext.Provider value={store}>
       <OidcJwtInitializer {...props}>{children}</OidcJwtInitializer>
diff --git a/src/store.ts b/src/store.ts
index 507b712..7a2ef6e 100644
--- a/src/store.ts
+++ b/src/store.ts
@@ -1,40 +1,33 @@
 import { create } from 'zustand';
 
-import {
-  CSRF_TOKEN_STORAGE_KEY,
-  LOGGED_IN_TOKEN_STORAGE_KEY,
-  RETRY_LOGIN_STORAGE_KEY,
-  USER_INFO_TOKEN_STORAGE_KEY
-} from './constants';
-import { Storage } from './storage';
-import { OidcJwtClientOptions, OidcJwtClientStore } from './types';
+import { AuthState, OidcJwtClientOptions, OidcJwtClientStore } from './types';
 import { AuthService } from './utils/AuthService';
 
 function createOidcJwtClientStore(
   client: OidcJwtClientOptions | false,
-  removeTokenFromUrlFunction?: (url: string) => void
+  removeTokenFromUrlFunction?: (url: string) => void,
 ) {
-  return create<OidcJwtClientStore>(set => {
-    const initialState = {
-      userInfo: Storage.get(USER_INFO_TOKEN_STORAGE_KEY),
-      csrfToken: Storage.get(CSRF_TOKEN_STORAGE_KEY),
-      isLoggedIn: !!Storage.get(LOGGED_IN_TOKEN_STORAGE_KEY),
+  return create<OidcJwtClientStore>((set) => {
+    const initialState: AuthState = {
+      userInfo: undefined,
+      csrfToken: null,
+      isLoggedIn: false,
       isInitialized: !client,
-      didRetryLogin: Storage.get(RETRY_LOGIN_STORAGE_KEY) === 1,
+      didRetryLogin: false,
     };
 
     const service = client
       ? new AuthService({
-        client,
-        removeTokenFromUrlFunction,
-        state: initialState,
-      })
+          client,
+          removeTokenFromUrlFunction,
+          state: initialState,
+        })
       : null;
 
     return {
       service,
       authState: initialState,
-      setState: newState => set({ authState: newState }),
+      setState: (newState) => set({ authState: newState }),
     };
   });
 }