Prevent multiple OIDC authn redirects on Lookout UI (#4196)

This fixes race conditions in the Lookout UI which cause repeated and unnecessary redirects when authenticating the user using OIDC. We also add loading and error screens which are displayed during the sign-in process. On the happy path, the loading screen is displayed very briefly only when (re-)authentication is required.
armadaproject · Feb 6, 2025 · f575442 · f575442
1 parent 6e52825
commit f575442
Show file tree

Hide file tree

Showing 22 changed files with 352 additions and 226 deletions.
diff --git a/internal/lookout/ui/src/App.tsx b/internal/lookout/ui/src/App.tsx
@@ -1,15 +1,16 @@
-import { Dispatch, ReactNode, SetStateAction, useEffect, useState } from "react"
+import { useEffect } from "react"
 
 import { CssBaseline, styled, ThemeProvider } from "@mui/material"
 import { QueryClient, QueryClientProvider } from "@tanstack/react-query"
 import { SnackbarProvider } from "notistack"
-import { UserManager, WebStorageStateStore, UserManagerSettings, User } from "oidc-client-ts"
-import { BrowserRouter, Navigate, Route, Routes, useNavigate } from "react-router-dom"
+import { UserManager, WebStorageStateStore, UserManagerSettings } from "oidc-client-ts"
+import { BrowserRouter, Navigate, Route, Routes } from "react-router-dom"
 
 import NavBar from "./components/NavBar"
 import JobSetsContainer from "./containers/JobSetsContainer"
 import { JobsTableContainer } from "./containers/lookoutV2/JobsTableContainer"
-import { UserManagerContext, useUserManager } from "./oidc"
+import { OidcAuthProvider } from "./oidcAuth"
+import { OIDC_REDIRECT_PATHNAME } from "./oidcAuth/OidcAuthProvider"
 import { Services, ServicesProvider } from "./services/context"
 import { theme } from "./theme/theme"
 import { CommandSpec, OidcConfig, withRouter } from "./utils"
@@ -43,71 +44,11 @@ type AppProps = {
   commandSpecs: CommandSpec[]
 }
 
-// Handling authentication on page opening
-
-interface AuthWrapperProps {
-  children: ReactNode
-  userManager: UserManager | undefined
-  isAuthenticated: boolean
-}
-
-interface OidcCallbackProps {
-  setIsAuthenticated: Dispatch<SetStateAction<boolean>>
-}
-
-export const UserManagerProvider = UserManagerContext.Provider
-
-function OidcCallback({ setIsAuthenticated }: OidcCallbackProps): JSX.Element {
-  const navigate = useNavigate()
-  const userManager = useUserManager()
-  const [error, setError] = useState<string | undefined>()
-
-  useEffect(() => {
-    if (!userManager) return
-    userManager
-      .signinRedirectCallback()
-      .then(() => {
-        setIsAuthenticated(true)
-        navigate("/")
-      })
-      .catch((e) => {
-        setError(`${e}`)
-        console.error(e)
-      })
-  }, [navigate, userManager, setIsAuthenticated])
-
-  if (error) return <p>Something went wrong; more details are available in the console.</p>
-  return <p>Authenticating...</p>
-}
-
-function AuthWrapper({ children, userManager, isAuthenticated }: AuthWrapperProps) {
-  useEffect(() => {
-    if (!userManager || isAuthenticated) return // Skip if userManager is not available or user is already authenticated
-
-    const handleAuthentication = async () => {
-      try {
-        const user = await userManager.getUser()
-        if (!user || user.expired) {
-          await userManager.signinRedirect()
-        }
-      } catch (error) {
-        console.error("Error during authentication:", error)
-      }
-    }
-
-    ;(async () => {
-      await handleAuthentication()
-    })()
-  }, [userManager, isAuthenticated])
-
-  return <>{children}</>
-}
-
 export function createUserManager(config: OidcConfig): UserManager {
   const userManagerSettings: UserManagerSettings = {
     authority: config.authority,
     client_id: config.clientId,
-    redirect_uri: `${window.location.origin}/oidc`,
+    redirect_uri: `${window.location.origin}${OIDC_REDIRECT_PATHNAME}`,
     scope: config.scope,
     userStore: new WebStorageStateStore({ store: window.localStorage }),
     loadUserInfo: true,
@@ -121,90 +62,68 @@ export function createUserManager(config: OidcConfig): UserManager {
 const V2Redirect = withRouter(({ router }) => <Navigate to={{ ...router.location, pathname: "/" }} />)
 
 export function App(props: AppProps): JSX.Element {
-  const [userManager, setUserManager] = useState<UserManager | undefined>(undefined)
-  const [isAuthenticated, setIsAuthenticated] = useState(false)
-  const [username, setUsername] = useState<string | undefined>(undefined)
-
-  useEffect(() => {
-    if (!userManager && props.oidcConfig) {
-      const userManagerInstance = createUserManager(props.oidcConfig)
-      setUserManager(userManagerInstance)
-
-      userManagerInstance.getUser().then((user: User | null) => {
-        if (user) {
-          setUsername(user.profile.sub)
-        }
-      })
-    }
-  }, [props.oidcConfig])
-
   useEffect(() => {
     if (props.customTitle) {
       document.title = `${props.customTitle} - Armada Lookout`
     }
   }, [props.customTitle])
 
-  const result = (
+  return (
     <ThemeProvider theme={theme} defaultMode="light">
       <CssBaseline />
       <SnackbarProvider anchorOrigin={{ horizontal: "right", vertical: "bottom" }} autoHideDuration={8000} maxSnack={3}>
         <QueryClientProvider client={queryClient}>
-          <BrowserRouter>
-            <UserManagerProvider value={userManager}>
-              <AuthWrapper userManager={userManager} isAuthenticated={isAuthenticated}>
-                <ServicesProvider services={props.services}>
-                  <AppContainer>
-                    <NavBar customTitle={props.customTitle} username={username} />
-                    <AppContent>
-                      <Routes>
-                        <Route
-                          path="/"
-                          element={
-                            <JobsTableContainer
-                              getJobsService={props.services.v2GetJobsService}
-                              groupJobsService={props.services.v2GroupJobsService}
-                              updateJobsService={props.services.v2UpdateJobsService}
-                              runInfoService={props.services.v2RunInfoService}
-                              jobSpecService={props.services.v2JobSpecService}
-                              cordonService={props.services.v2CordonService}
-                              debug={props.debugEnabled}
-                              autoRefreshMs={props.jobsAutoRefreshMs}
-                              commandSpecs={props.commandSpecs}
-                            />
-                          }
-                        />
-                        <Route
-                          path="/job-sets"
-                          element={
-                            <JobSetsContainer
-                              v2GroupJobsService={props.services.v2GroupJobsService}
-                              v2UpdateJobSetsService={props.services.v2UpdateJobSetsService}
-                              jobSetsAutoRefreshMs={props.jobSetsAutoRefreshMs}
-                            />
-                          }
-                        />
-                        <Route path="/oidc" element={<OidcCallback setIsAuthenticated={setIsAuthenticated} />} />
-                        <Route path="/v2" element={<V2Redirect />} />
-                        <Route
-                          path="*"
-                          element={
-                            // This wildcard route ensures that users who follow old
-                            // links to /job-sets or /jobs see something other than
-                            // a blank page.
-                            <Navigate to="/" />
-                          }
-                        />
-                      </Routes>
-                    </AppContent>
-                  </AppContainer>
-                </ServicesProvider>
-              </AuthWrapper>
-            </UserManagerProvider>
-          </BrowserRouter>
+          <OidcAuthProvider oidcConfig={props.oidcConfig}>
+            <BrowserRouter>
+              <ServicesProvider services={props.services}>
+                <AppContainer>
+                  <NavBar customTitle={props.customTitle} />
+                  <AppContent>
+                    <Routes>
+                      <Route
+                        path="/"
+                        element={
+                          <JobsTableContainer
+                            getJobsService={props.services.v2GetJobsService}
+                            groupJobsService={props.services.v2GroupJobsService}
+                            updateJobsService={props.services.v2UpdateJobsService}
+                            runInfoService={props.services.v2RunInfoService}
+                            jobSpecService={props.services.v2JobSpecService}
+                            cordonService={props.services.v2CordonService}
+                            debug={props.debugEnabled}
+                            autoRefreshMs={props.jobsAutoRefreshMs}
+                            commandSpecs={props.commandSpecs}
+                          />
+                        }
+                      />
+                      <Route
+                        path="/job-sets"
+                        element={
+                          <JobSetsContainer
+                            v2GroupJobsService={props.services.v2GroupJobsService}
+                            v2UpdateJobSetsService={props.services.v2UpdateJobSetsService}
+                            jobSetsAutoRefreshMs={props.jobSetsAutoRefreshMs}
+                          />
+                        }
+                      />
+                      <Route path="/v2" element={<V2Redirect />} />
+                      <Route
+                        path="*"
+                        element={
+                          // This wildcard route ensures that users who follow old
+                          // links to /job-sets or /jobs see something other than
+                          // a blank page.
+                          <Navigate to="/" />
+                        }
+                      />
+                    </Routes>
+                  </AppContent>
+                </AppContainer>
+              </ServicesProvider>
+            </BrowserRouter>
+          </OidcAuthProvider>
         </QueryClientProvider>
       </SnackbarProvider>
     </ThemeProvider>
   )
-
-  return result
 }
diff --git a/internal/lookout/ui/src/components/NavBar.tsx b/internal/lookout/ui/src/components/NavBar.tsx
@@ -5,6 +5,7 @@ import { AppBar, IconButton, Tab, Tabs, Toolbar, Typography } from "@mui/materia
 import { Link } from "react-router-dom"
 
 import { SettingsDialog } from "./SettingsDialog"
+import { useUsername } from "../oidcAuth"
 import { Router, withRouter } from "../utils"
 
 import "./NavBar.css"
@@ -46,13 +47,13 @@ function locationFromIndex(pages: Page[], index: number): string {
 interface NavBarProps {
   customTitle: string
   router: Router
-  username?: string
 }
 
-function NavBar({ customTitle, router, username }: NavBarProps) {
+function NavBar({ customTitle, router }: NavBarProps) {
   const currentLocation = router.location.pathname
   const currentValue = locationMap.has(currentLocation) ? locationMap.get(currentLocation) : 0
   const [settingsOpen, setSettingsOpen] = useState(false)
+  const username = useUsername()
 
   return (
     <>

diff --git a/internal/lookout/ui/src/components/lookoutV2/CancelDialog.tsx b/internal/lookout/ui/src/components/lookoutV2/CancelDialog.tsx
@@ -10,7 +10,7 @@ import dialogStyles from "./DialogStyles.module.css"
 import { JobStatusTable } from "./JobStatusTable"
 import { useCustomSnackbar } from "../../hooks/useCustomSnackbar"
 import { isTerminatedJobState, Job, JobFilter, JobId } from "../../models/lookoutV2Models"
-import { getAccessToken, useUserManager } from "../../oidc"
+import { useGetAccessToken } from "../../oidcAuth"
 import { IGetJobsService } from "../../services/lookoutV2/GetJobsService"
 import { UpdateJobsService } from "../../services/lookoutV2/UpdateJobsService"
 import { pl, waitMillis, PlatformCancelReason } from "../../utils"
@@ -41,7 +41,7 @@ export const CancelDialog = ({
   const [isPlatformCancel, setIsPlatformCancel] = useState(false)
   const openSnackbar = useCustomSnackbar()
 
-  const userManager = useUserManager()
+  const getAccessToken = useGetAccessToken()
 
   // Actions
   const fetchSelectedJobs = useCallback(async () => {
@@ -67,7 +67,7 @@ export const CancelDialog = ({
     setIsCancelling(true)
 
     const reason = isPlatformCancel ? PlatformCancelReason : ""
-    const accessToken = userManager && (await getAccessToken(userManager))
+    const accessToken = await getAccessToken()
     const response = await updateJobsService.cancelJobs(cancellableJobs, reason, accessToken)
 
     if (response.failedJobIds.length === 0) {

diff --git a/internal/lookout/ui/src/components/lookoutV2/ReprioritiseDialog.tsx b/internal/lookout/ui/src/components/lookoutV2/ReprioritiseDialog.tsx
@@ -18,7 +18,7 @@ import dialogStyles from "./DialogStyles.module.css"
 import { JobStatusTable } from "./JobStatusTable"
 import { useCustomSnackbar } from "../../hooks/useCustomSnackbar"
 import { isTerminatedJobState, Job, JobFilter, JobId } from "../../models/lookoutV2Models"
-import { getAccessToken, useUserManager } from "../../oidc"
+import { useGetAccessToken } from "../../oidcAuth"
 import { IGetJobsService } from "../../services/lookoutV2/GetJobsService"
 import { UpdateJobsService } from "../../services/lookoutV2/UpdateJobsService"
 import { pl, waitMillis } from "../../utils"
@@ -51,7 +51,7 @@ export const ReprioritiseDialog = ({
   const [hasAttemptedReprioritise, setHasAttemptedReprioritise] = useState(false)
   const openSnackbar = useCustomSnackbar()
 
-  const userManager = useUserManager()
+  const getAccessToken = useGetAccessToken()
 
   // Actions
   const fetchSelectedJobs = useCallback(async () => {
@@ -79,7 +79,7 @@ export const ReprioritiseDialog = ({
 
     setIsReprioritising(true)
 
-    const accessToken = userManager && (await getAccessToken(userManager))
+    const accessToken = await getAccessToken()
     const response = await updateJobsService.reprioritiseJobs(reprioritisableJobs, newPriority, accessToken)
 
     if (response.failedJobIds.length === 0) {

diff --git a/internal/lookout/ui/src/components/lookoutV2/sidebar/SidebarTabJobResult.tsx b/internal/lookout/ui/src/components/lookoutV2/sidebar/SidebarTabJobResult.tsx
@@ -21,7 +21,7 @@ import { NoRunsAlert } from "./NoRunsAlert"
 import { SidebarTabHeading, SidebarTabSubheading } from "./sidebarTabContentComponents"
 import { useCustomSnackbar } from "../../../hooks/useCustomSnackbar"
 import { Job, JobRun, JobState } from "../../../models/lookoutV2Models"
-import { getAccessToken, useUserManager } from "../../../oidc"
+import { useGetAccessToken } from "../../../oidcAuth"
 import { ICordonService } from "../../../services/lookoutV2/CordonService"
 import { IGetJobInfoService } from "../../../services/lookoutV2/GetJobInfoService"
 import { IGetRunInfoService } from "../../../services/lookoutV2/GetRunInfoService"
@@ -216,11 +216,11 @@ export const SidebarTabJobResult = ({
     setOpen(false)
   }
 
-  const userManager = useUserManager()
+  const getAccessToken = useGetAccessToken()
 
   const cordon = async (cluster: string, node: string) => {
     try {
-      const accessToken = userManager && (await getAccessToken(userManager))
+      const accessToken = await getAccessToken()
       await cordonService.cordonNode(cluster, node, accessToken)
       openSnackbar("Successfully cordoned node " + node, "success")
     } catch (e) {

diff --git a/internal/lookout/ui/src/containers/CancelJobSetsDialog.tsx b/internal/lookout/ui/src/containers/CancelJobSetsDialog.tsx
@@ -4,7 +4,7 @@ import { Dialog, DialogContent, DialogTitle } from "@mui/material"
 
 import CancelJobSets from "../components/job-sets/cancel-job-sets/CancelJobSets"
 import CancelJobSetsOutcome from "../components/job-sets/cancel-job-sets/CancelJobSetsOutcome"
-import { getAccessToken, useUserManager } from "../oidc"
+import { useGetAccessToken } from "../oidcAuth"
 import { ApiJobState } from "../openapi/armada"
 import { JobSet } from "../services/JobService"
 import { CancelJobSetsResponse, UpdateJobSetsService } from "../services/lookoutV2/UpdateJobSetsService"
@@ -52,7 +52,7 @@ export default function CancelJobSetsDialog(props: CancelJobSetsDialogProps) {
 
   const statesToCancel = getStatesToCancel(includeQueued, includeRunning)
 
-  const userManager = useUserManager()
+  const getAccessToken = useGetAccessToken()
 
   async function cancelJobSets() {
     if (requestStatus === "Loading") {
@@ -61,7 +61,7 @@ export default function CancelJobSetsDialog(props: CancelJobSetsDialogProps) {
 
     setRequestStatus("Loading")
     const reason = isPlatformCancel ? PlatformCancelReason : ""
-    const accessToken = userManager && (await getAccessToken(userManager))
+    const accessToken = await getAccessToken()
     const cancelJobSetsResponse = await props.updateJobSetsService.cancelJobSets(
       props.queue,
       jobSetsToCancel,

diff --git a/internal/lookout/ui/src/containers/ReprioritizeJobSetsDialog.tsx b/internal/lookout/ui/src/containers/ReprioritizeJobSetsDialog.tsx
@@ -4,7 +4,7 @@ import { Dialog, DialogContent, DialogTitle } from "@mui/material"
 
 import ReprioritizeJobSets from "../components/job-sets/reprioritize-job-sets/ReprioritizeJobSets"
 import ReprioritizeJobSetsOutcome from "../components/job-sets/reprioritize-job-sets/ReprioritizeJobSetsOutcome"
-import { getAccessToken, useUserManager } from "../oidc"
+import { useGetAccessToken } from "../oidcAuth"
 import { JobSet } from "../services/JobService"
 import { ReprioritizeJobSetsResponse, UpdateJobSetsService } from "../services/lookoutV2/UpdateJobSetsService"
 import { ApiResult, priorityIsValid, RequestStatus } from "../utils"
@@ -37,15 +37,15 @@ export default function ReprioritizeJobSetsDialog(props: ReprioritizeJobSetsDial
 
   const jobSetsToReprioritize = getReprioritizeableJobSets(props.selectedJobSets)
 
-  const userManager = useUserManager()
+  const getAccessToken = useGetAccessToken()
 
   async function reprioritizeJobSets() {
     if (requestStatus == "Loading" || !priorityIsValid(priority)) {
       return
     }
 
     setRequestStatus("Loading")
-    const accessToken = userManager && (await getAccessToken(userManager))
+    const accessToken = await getAccessToken()
     const reprioritizeJobSetsResponse = await props.updateJobSetsService.reprioritizeJobSets(
       props.queue,
       jobSetsToReprioritize,