fix: redirect on sso login (#9369)

determined-ai · May 28, 2024 · a0f2e33 · a0f2e33
1 parent 9abde37
commit a0f2e33
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 4 deletions.
diff --git a/webui/react/src/globalStorage.ts b/webui/react/src/globalStorage.ts
@@ -8,6 +8,7 @@ class GlobalStorage {
     this.storage = storage;
     this.keys = {
       authToken: 'auth-token',
+      landingRedirect: 'landing-redirect',
       serverAddress: 'server-address',
     };
   }
@@ -20,6 +21,10 @@ class GlobalStorage {
     return this.storage.get<string>(this.keys.serverAddress) || '';
   }
 
+  get landingRedirect() {
+    return this.storage.get<string>(this.keys.landingRedirect) || '';
+  }
+
   set authToken(token: string) {
     this.storage.set(this.keys.authToken, token);
   }
@@ -28,15 +33,27 @@ class GlobalStorage {
     this.storage.set(this.keys.serverAddress, address);
   }
 
+  set landingRedirect(address: string) {
+    this.storage.set(this.keys.landingRedirect, address);
+  }
+
   removeAuthToken() {
     this.storage.remove(this.keys.authToken);
   }
 
   removeServerAddress() {
     this.storage.remove(this.keys.serverAddress);
   }
+
+  removeLandingRedirect() {
+    this.storage.remove(this.keys.landingRedirect);
+  }
 }
 
 export const globalStorage = new GlobalStorage(
   new StorageManager({ basePath: 'global', store: window.localStorage }),
 );
+
+export const sessionStorage = new GlobalStorage(
+  new StorageManager({ basePath: 'session', store: window.sessionStorage }),
+);
diff --git a/webui/react/src/pages/SignIn.tsx b/webui/react/src/pages/SignIn.tsx
@@ -15,6 +15,7 @@ import Page from 'components/Page';
 import PageMessage from 'components/PageMessage';
 import useUI from 'components/ThemeProvider';
 import { handleRelayState, samlUrl } from 'ee/SamlAuth';
+import { sessionStorage } from 'globalStorage';
 import useAuthCheck from 'hooks/useAuthCheck';
 import usePolling from 'hooks/usePolling';
 import { defaultRoute, rbacDefaultRoute } from 'routes';
@@ -62,6 +63,10 @@ const SignIn: React.FC = () => {
    * the user to the most recent requested page.
    */
   useEffect(() => {
+    if (location.state != null) {
+      sessionStorage.landingRedirect = locationToPath(location.state) ?? '';
+    }
+
     if (isAuthenticated) {
       // Stop the spinner, prepping for user redirect.
       uiActions.hideSpinner();
@@ -76,10 +81,9 @@ const SignIn: React.FC = () => {
 
       // Reroute the authenticated user to the app.
       if (!queries.has('redirect')) {
-        routeToReactUrl(
-          locationToPath(location.state) ||
-            (rbacEnabled ? rbacDefaultRoute.path : defaultRoute.path),
-        );
+        const path = sessionStorage.landingRedirect;
+        sessionStorage.removeLandingRedirect();
+        routeToReactUrl(path || (rbacEnabled ? rbacDefaultRoute.path : defaultRoute.path));
       } else {
         routeAll(queries.get('redirect') || '');
       }