From db7fb818558f1d2e4b267081b5ba5f4f571dd0a6 Mon Sep 17 00:00:00 2001
From: Kersom <9053044+nixocio@users.noreply.github.com>
Date: Wed, 27 Oct 2021 16:37:54 -0400
Subject: [PATCH] Fix login redirect (#5386)

Allows the user to visit login page when the login redirect url is set.

Also, redirects to login page once logging out and there is session from
a SAML available.

See: https://github.com/ansible/awx/issues/11012
---
 awx/ui/src/App.js                 | 19 +++++++++++++---
 awx/ui/src/App.test.js            | 36 ++++++++++++++++++++++++++++++-
 awx/ui/src/contexts/Session.js    | 35 +++++++++++++++++++++++++++---
 awx/ui/src/screens/Login/Login.js | 15 ++-----------
 4 files changed, 85 insertions(+), 20 deletions(-)

diff --git a/awx/ui/src/App.js b/awx/ui/src/App.js
index ab0211bd38..674dec8b07 100644
--- a/awx/ui/src/App.js
+++ b/awx/ui/src/App.js
@@ -98,8 +98,13 @@ const AuthorizedRoutes = ({ routeConfig }) => {
   );
 };
 
-const ProtectedRoute = ({ children, ...rest }) => {
-  const { authRedirectTo, setAuthRedirectTo } = useSession();
+export function ProtectedRoute({ children, ...rest }) {
+  const {
+    authRedirectTo,
+    setAuthRedirectTo,
+    loginRedirectOverride,
+    isUserBeingLoggedOut,
+  } = useSession();
   const location = useLocation();
 
   useEffect(() => {
@@ -120,8 +125,16 @@ const ProtectedRoute = ({ children, ...rest }) => {
     );
   }
 
+  if (
+    loginRedirectOverride &&
+    !window.location.href.includes('/login') &&
+    !isUserBeingLoggedOut
+  ) {
+    window.location.replace(loginRedirectOverride);
+    return null;
+  }
   return <Redirect to="/login" />;
-};
+}
 
 function App() {
   const history = useHistory();
diff --git a/awx/ui/src/App.test.js b/awx/ui/src/App.test.js
index a8b842714a..edcf60ebb7 100644
--- a/awx/ui/src/App.test.js
+++ b/awx/ui/src/App.test.js
@@ -3,7 +3,7 @@ import { act } from 'react-dom/test-utils';
 import { RootAPI } from 'api';
 import * as SessionContext from 'contexts/Session';
 import { mountWithContexts } from '../testUtils/enzymeHelpers';
-import App from './App';
+import App, { ProtectedRoute } from './App';
 
 jest.mock('./api');
 
@@ -20,6 +20,8 @@ describe('<App />', () => {
     const contextValues = {
       setAuthRedirectTo: jest.fn(),
       isSessionExpired: false,
+      isUserBeingLoggedOut: false,
+      loginRedirectOverride: null,
     };
     jest
       .spyOn(SessionContext, 'useSession')
@@ -32,4 +34,36 @@ describe('<App />', () => {
     expect(wrapper.length).toBe(1);
     jest.clearAllMocks();
   });
+
+  test('redirect to login override', async () => {
+    const { location } = window;
+    delete window.location;
+    window.location = {
+      replace: jest.fn(),
+      href: '/',
+    };
+
+    expect(window.location.replace).not.toHaveBeenCalled();
+
+    const contextValues = {
+      setAuthRedirectTo: jest.fn(),
+      isSessionExpired: false,
+      isUserBeingLoggedOut: false,
+      loginRedirectOverride: '/sso/test',
+    };
+    jest
+      .spyOn(SessionContext, 'useSession')
+      .mockImplementation(() => contextValues);
+
+    await act(async () => {
+      mountWithContexts(
+        <ProtectedRoute>
+          <div>foo</div>
+        </ProtectedRoute>
+      );
+    });
+
+    expect(window.location.replace).toHaveBeenCalled();
+    window.location = location;
+  });
 });
diff --git a/awx/ui/src/contexts/Session.js b/awx/ui/src/contexts/Session.js
index f6db1e3be3..1f76826abf 100644
--- a/awx/ui/src/contexts/Session.js
+++ b/awx/ui/src/contexts/Session.js
@@ -5,10 +5,11 @@ import React, {
   useRef,
   useCallback,
 } from 'react';
-import { useHistory } from 'react-router-dom';
+import { useHistory, Redirect } from 'react-router-dom';
 import { DateTime } from 'luxon';
 import { RootAPI, MeAPI } from 'api';
 import { isAuthenticated } from 'util/auth';
+import useRequest from 'hooks/useRequest';
 import { SESSION_TIMEOUT_KEY } from '../constants';
 
 // The maximum supported timeout for setTimeout(), in milliseconds,
@@ -72,8 +73,31 @@ function SessionProvider({ children }) {
   const [sessionTimeout, setSessionTimeout] = useStorage(SESSION_TIMEOUT_KEY);
   const [sessionCountdown, setSessionCountdown] = useState(0);
   const [authRedirectTo, setAuthRedirectTo] = useState('/');
+  const [isUserBeingLoggedOut, setIsUserBeingLoggedOut] = useState(false);
+
+  const {
+    request: fetchLoginRedirectOverride,
+    result: { loginRedirectOverride },
+    isLoading,
+  } = useRequest(
+    useCallback(async () => {
+      const { data } = await RootAPI.read();
+      return {
+        loginRedirectOverride: data?.login_redirect_override,
+      };
+    }, []),
+    {
+      loginRedirectOverride: null,
+      isLoading: true,
+    }
+  );
+
+  useEffect(() => {
+    fetchLoginRedirectOverride();
+  }, [fetchLoginRedirectOverride]);
 
   const logout = useCallback(async () => {
+    setIsUserBeingLoggedOut(true);
     if (!isSessionExpired.current) {
       setAuthRedirectTo('/logout');
     }
@@ -82,14 +106,13 @@ function SessionProvider({ children }) {
     setSessionCountdown(0);
     clearTimeout(sessionTimeoutId.current);
     clearInterval(sessionIntervalId.current);
+    return <Redirect to="/login" />;
   }, [setSessionTimeout, setSessionCountdown]);
 
   useEffect(() => {
     if (!isAuthenticated(document.cookie)) {
-      history.replace('/login');
       return () => {};
     }
-
     const calcRemaining = () => {
       if (sessionTimeout) {
         return Math.max(
@@ -140,9 +163,15 @@ function SessionProvider({ children }) {
     clearInterval(sessionIntervalId.current);
   }, []);
 
+  if (isLoading) {
+    return null;
+  }
+
   return (
     <SessionContext.Provider
       value={{
+        isUserBeingLoggedOut,
+        loginRedirectOverride,
         authRedirectTo,
         handleSessionContinue,
         isSessionExpired,
diff --git a/awx/ui/src/screens/Login/Login.js b/awx/ui/src/screens/Login/Login.js
index bbd333524d..ede201c4e7 100644
--- a/awx/ui/src/screens/Login/Login.js
+++ b/awx/ui/src/screens/Login/Login.js
@@ -47,18 +47,12 @@ function AWXLogin({ alt, isAuthenticated }) {
     isLoading: isCustomLoginInfoLoading,
     error: customLoginInfoError,
     request: fetchCustomLoginInfo,
-    result: {
-      brandName,
-      logo,
-      loginInfo,
-      socialAuthOptions,
-      loginRedirectOverride,
-    },
+    result: { brandName, logo, loginInfo, socialAuthOptions },
   } = useRequest(
     useCallback(async () => {
       const [
         {
-          data: { custom_logo, custom_login_info, login_redirect_override },
+          data: { custom_logo, custom_login_info },
         },
         {
           data: { BRAND_NAME },
@@ -78,7 +72,6 @@ function AWXLogin({ alt, isAuthenticated }) {
         logo: logoSrc,
         loginInfo: custom_login_info,
         socialAuthOptions: authData,
-        loginRedirectOverride: login_redirect_override,
       };
     }, []),
     {
@@ -118,10 +111,6 @@ function AWXLogin({ alt, isAuthenticated }) {
   if (isCustomLoginInfoLoading) {
     return null;
   }
-  if (!isAuthenticated(document.cookie) && loginRedirectOverride) {
-    window.location.replace(loginRedirectOverride);
-    return null;
-  }
   if (isAuthenticated(document.cookie)) {
     return <Redirect to={authRedirectTo || '/'} />;
   }