Fix google auth (#3096)

* Fix: authenticator for google sso

openmetadata-ui/src/main/resources/ui/src/App.tsx

@@ -23,7 +23,7 @@ const App: FunctionComponent = () => {
       <div className="content-wrapper" data-testid="content-wrapper">
         <ToastContextProvider>
           <Router>
-            <AuthProvider>
+            <AuthProvider childComponentType={AppRouter}>
               <AppRouter />
             </AuthProvider>
           </Router>

openmetadata-ui/src/main/resources/ui/src/auth-provider/AuthProvider.tsx

@@ -20,6 +20,7 @@ import { isEmpty, isNil } from 'lodash';
 import { observer } from 'mobx-react';
 import { UserPermissions } from 'Models';
 import React, {
+  ComponentType,
   createContext,
   ReactNode,
   useContext,
@@ -29,8 +30,9 @@ import React, {
 } from 'react';
 import { useHistory, useLocation } from 'react-router-dom';
 import appState from '../AppState';
-import GoogleAuthenticator from '../authenticators/GoogleAuthenticator';
+// import GoogleAuthenticator from '../authenticators/GoogleAuthenticator';
 import MsalAuthenticator from '../authenticators/MsalAuthenticator';
+import OidcAuthenticator from '../authenticators/OidcAuthenticator';
 import OktaAuthenticator from '../authenticators/OktaAuthenticator';
 import axiosClient from '../axiosAPIs';
 import {
@@ -53,6 +55,7 @@ import useToastContext from '../hooks/useToastContext';
 import {
   getAuthConfig,
   getNameFromEmail,
+  getUserManagerConfig,
   isProtectedRoute,
   isTourRoute,
   msalInstance,
@@ -64,13 +67,17 @@ import { AuthenticatorRef, OidcUser } from './AuthProvider.interface';
 import OktaAuthProvider from './okta-auth-provider';
 
 interface AuthProviderProps {
+  childComponentType: ComponentType;
   children: ReactNode;
 }
 
 const cookieStorage = new CookieStorage();
 const userAPIQueryFields = 'profile,teams,roles';
 
-export const AuthProvider = ({ children }: AuthProviderProps) => {
+export const AuthProvider = ({
+  childComponentType,
+  children,
+}: AuthProviderProps) => {
   const location = useLocation();
   const history = useHistory();
   const showToast = useToastContext();
@@ -187,6 +194,7 @@ export const AuthProvider = ({ children }: AuthProviderProps) => {
   };
 
   const handleSuccessfulLogin = (user: OidcUser) => {
+    setLoading(true);
     getUserByName(getNameFromEmail(user.profile.email), userAPIQueryFields)
       .then((res: AxiosResponse) => {
         if (res.data) {
@@ -207,6 +215,9 @@ export const AuthProvider = ({ children }: AuthProviderProps) => {
           setIsSigningIn(true);
           history.push(ROUTES.SIGNUP);
         }
+      })
+      .finally(() => {
+        setLoading(false);
       });
   };
 
@@ -217,7 +228,8 @@ export const AuthProvider = ({ children }: AuthProviderProps) => {
   const getAuthenticatedUser = (config: Record<string, string | boolean>) => {
     switch (config?.provider) {
       case AuthTypes.OKTA:
-      case AuthTypes.AZURE: {
+      case AuthTypes.AZURE:
+      case AuthTypes.GOOGLE: {
         getLoggedInUserDetails();
 
         break;
@@ -324,13 +336,25 @@ export const AuthProvider = ({ children }: AuthProviderProps) => {
         );
       }
       case AuthTypes.GOOGLE: {
-        return (
+        return authConfig ? (
-          <GoogleAuthenticator
+          // <GoogleAuthenticator
+          //   ref={authenticatorRef}
+          //   onLoginSuccess={handleSuccessfulLogin}
+          //   onLogoutSuccess={handleSuccessfulLogout}>
+          //   {children}
+          // </GoogleAuthenticator>
+          <OidcAuthenticator
+            childComponentType={childComponentType}
             ref={authenticatorRef}
+            userConfig={getUserManagerConfig({
+              ...(authConfig as Record<string, string>),
+            })}
             onLoginSuccess={handleSuccessfulLogin}
             onLogoutSuccess={handleSuccessfulLogout}>
             {children}
-          </GoogleAuthenticator>
+          </OidcAuthenticator>
+        ) : (
+          <Loader />
         );
       }
       case AuthTypes.AZURE: {

openmetadata-ui/src/main/resources/ui/src/authenticators/OidcAuthenticator.tsx

@@ -0,0 +1,164 @@
+/*
+ *  Copyright 2021 Collate
+ *  Licensed under the Apache License, Version 2.0 (the "License");
+ *  you may not use this file except in compliance with the License.
+ *  You may obtain a copy of the License at
+ *  http://www.apache.org/licenses/LICENSE-2.0
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+import { isEmpty } from 'lodash';
+import { UserManager, WebStorageStateStore } from 'oidc-client';
+import React, {
+  ComponentType,
+  forwardRef,
+  Fragment,
+  ReactNode,
+  useImperativeHandle,
+  useMemo,
+} from 'react';
+import { Callback, makeAuthenticator, makeUserManager } from 'react-oidc';
+import { Redirect, Route, Switch } from 'react-router-dom';
+import AppState from '../AppState';
+import { useAuthContext } from '../auth-provider/AuthProvider';
+import {
+  AuthenticatorRef,
+  OidcUser,
+} from '../auth-provider/AuthProvider.interface';
+import Appbar from '../components/app-bar/Appbar';
+import Loader from '../components/Loader/Loader';
+import { oidcTokenKey, ROUTES } from '../constants/constants';
+import SigninPage from '../pages/login';
+import PageNotFound from '../pages/page-not-found';
+
+interface Props {
+  childComponentType: ComponentType;
+  children: ReactNode;
+  userConfig: Record<string, string | boolean | WebStorageStateStore>;
+  onLoginSuccess: (user: OidcUser) => void;
+  onLogoutSuccess: () => void;
+}
+
+const getAuthenticator = (type: ComponentType, userManager: UserManager) => {
+  return makeAuthenticator({
+    userManager: userManager,
+    signinArgs: {
+      app: 'openmetadata',
+    },
+  })(type);
+};
+
+const OidcAuthenticator = forwardRef<AuthenticatorRef, Props>(
+  (
+    {
+      childComponentType,
+      children,
+      userConfig,
+      onLoginSuccess,
+      onLogoutSuccess,
+    }: Props,
+    ref
+  ) => {
+    const {
+      loading,
+      isAuthenticated,
+      setIsAuthenticated,
+      isAuthDisabled,
+      isSigningIn,
+      setIsSigningIn,
+      setLoadingIndicator,
+    } = useAuthContext();
+    const { userDetails, newUser } = AppState;
+    const userManager = useMemo(
+      () => makeUserManager(userConfig),
+      [userConfig]
+    );
+
+    const login = () => {
+      setIsSigningIn(true);
+    };
+
+    const logout = () => {
+      setLoadingIndicator(true);
+      setIsAuthenticated(false);
+      localStorage.removeItem(
+        `oidc.user:${userConfig.authority}:${userConfig.client_id}`
+      );
+      onLogoutSuccess();
+    };
+
+    useImperativeHandle(ref, () => ({
+      invokeLogin() {
+        login();
+      },
+      invokeLogout() {
+        logout();
+      },
+    }));
+
+    const AppWithAuth = getAuthenticator(childComponentType, userManager);
+
+    return (
+      <Fragment>
+        {!loading ? (
+          <>
+            <Appbar />
+            <Switch>
+              <Route exact path={ROUTES.HOME}>
+                {!isAuthDisabled && !isAuthenticated && !isSigningIn ? (
+                  <Redirect to={ROUTES.SIGNIN} />
+                ) : (
+                  <Redirect to={ROUTES.MY_DATA} />
+                )}
+              </Route>
+              <Route exact component={PageNotFound} path={ROUTES.NOT_FOUND} />
+              {!isSigningIn ? (
+                <Route exact component={SigninPage} path={ROUTES.SIGNIN} />
+              ) : null}
+              <Route
+                path={ROUTES.CALLBACK}
+                render={() => (
+                  <>
+                    <Callback
+                      userManager={userManager}
+                      onSuccess={(user) => {
+                        localStorage.setItem(oidcTokenKey, user.id_token);
+                        setIsAuthenticated(true);
+                        onLoginSuccess(user as OidcUser);
+                      }}
+                    />
+                    <Loader />
+                  </>
+                )}
+              />
+              {isAuthenticated || isAuthDisabled ? (
+                <Fragment>{children}</Fragment>
+              ) : !isSigningIn && isEmpty(userDetails) && isEmpty(newUser) ? (
+                <Redirect to={ROUTES.SIGNIN} />
+              ) : (
+                <AppWithAuth />
+              )}
+            </Switch>
+            {/* TODO: Uncomment below lines to show Welcome modal on Sign-up */}
+            {/* {isAuthenticatedRoute && isFirstTimeUser ? (
+            <FirstTimeUserModal
+              onCancel={() => handleFirstTourModal(true)}
+              onSave={() => handleFirstTourModal(false)}
+            />
+          ) : null} */}
+          </>
+        ) : (
+          <Loader />
+        )}
+      </Fragment>
+    );
+  }
+);
+
+OidcAuthenticator.displayName = 'OidcAuthenticator';
+
+export default OidcAuthenticator;

openmetadata-ui/src/main/resources/ui/src/pages/login/index.tsx

@@ -34,7 +34,7 @@ const SigninPage = () => {
 
   const getSignInButton = (): JSX.Element => {
     let btnComponent: JSX.Element;
-    switch (authConfig.provider) {
+    switch (authConfig?.provider) {
       case AuthTypes.GOOGLE: {
         btnComponent = (
           <LoginButton

openmetadata-ui/src/main/resources/ui/src/router/AppRouter.tsx

@@ -17,12 +17,14 @@ import { useAuthContext } from '../auth-provider/AuthProvider';
 import Appbar from '../components/app-bar/Appbar';
 import Loader from '../components/Loader/Loader';
 import { ROUTES } from '../constants/constants';
+import { AuthTypes } from '../enums/signin.enum';
 import SigninPage from '../pages/login';
 import PageNotFound from '../pages/page-not-found';
 import AuthenticatedAppRouter from './AuthenticatedAppRouter';
 
 const AppRouter = () => {
   const {
+    authConfig,
     isAuthDisabled,
     isAuthenticated,
     loading,
@@ -35,28 +37,34 @@ const AppRouter = () => {
     <Loader />
   ) : (
     <>
-      <Appbar />
+      {authConfig?.provider === AuthTypes.GOOGLE ? (
-      <Switch>
+        <AuthenticatedAppRouter />
-        <Route exact path={ROUTES.HOME}>
+      ) : (
-          {!isAuthDisabled && !isAuthenticated && !isSigningIn ? (
+        <>
-            <Redirect to={ROUTES.SIGNIN} />
+          <Appbar />
-          ) : (
+          <Switch>
-            <Redirect to={ROUTES.MY_DATA} />
+            <Route exact path={ROUTES.HOME}>
-          )}
+              {!isAuthDisabled && !isAuthenticated && !isSigningIn ? (
-        </Route>
+                <Redirect to={ROUTES.SIGNIN} />
-        {!isSigningIn ? (
+              ) : (
-          <Route exact component={SigninPage} path={ROUTES.SIGNIN} />
+                <Redirect to={ROUTES.MY_DATA} />
-        ) : null}
+              )}
-        {callbackComponent ? (
+            </Route>
-          <Route component={callbackComponent} path={ROUTES.CALLBACK} />
+            {!isSigningIn ? (
-        ) : null}
+              <Route exact component={SigninPage} path={ROUTES.SIGNIN} />
-        <Route exact component={PageNotFound} path={ROUTES.NOT_FOUND} />
+            ) : null}
-        {isAuthDisabled || isAuthenticated ? (
+            {callbackComponent ? (
-          <AuthenticatedAppRouter />
+              <Route component={callbackComponent} path={ROUTES.CALLBACK} />
-        ) : (
+            ) : null}
-          <Redirect to={ROUTES.SIGNIN} />
+            <Route exact component={PageNotFound} path={ROUTES.NOT_FOUND} />
-        )}
+            {isAuthDisabled || isAuthenticated ? (
-      </Switch>
+              <AuthenticatedAppRouter />
+            ) : (
+              <Redirect to={ROUTES.SIGNIN} />
+            )}
+          </Switch>
+        </>
+      )}
     </>
   );
 };

openmetadata-ui/src/main/resources/ui/src/utils/AuthProvider.util.ts

@@ -18,15 +18,12 @@ import {
   PopupRequest,
   PublicClientApplication,
 } from '@azure/msal-browser';
-import { CookieStorage } from 'cookie-storage';
 import { isNil } from 'lodash';
 import { WebStorageStateStore } from 'oidc-client';
 import { ROUTES } from '../constants/constants';
 import { AuthTypes } from '../enums/signin.enum';
 import { isDev } from '../utils/EnvironmentUtils';
 
-const cookieStorage = new CookieStorage();
-
 export let msalInstance: IPublicClientApplication;
 
 export const getOidcExpiry = () => {
@@ -52,7 +49,8 @@ export const getUserManagerConfig = (
       ? callbackUrl
       : `${window.location.origin}/callback`,
     scope: 'openid email profile',
-    userStore: new WebStorageStateStore({ store: cookieStorage }),
+    // userStore: new WebStorageStateStore({ store: cookieStorage }),
+    userStore: new WebStorageStateStore({ store: localStorage }),
   };
 };
 
@@ -82,8 +80,14 @@ export const getAuthConfig = (
       break;
     case AuthTypes.GOOGLE:
       {
+        // config = {
+        //   clientId,
+        //   provider,
+        // };
         config = {
+          authority,
           clientId,
+          callbackUrl,
           provider,
         };
       }