feat: replace redirect url pattern with route key pattern

src/components/AuthWrappers/EnforceLoginStatePageWrapper.tsx

@@ -1,45 +1,37 @@
-import { useMemo, type PropsWithChildren } from 'react'
+import { useEffect, type PropsWithChildren } from 'react'
 import { useRouter } from 'next/router'
 
 import { appendWithRedirect } from '~/utils/url'
 import { useLoginState } from '~/features/auth'
-import { SIGN_IN } from '~/lib/routes'
-import { callbackUrlSchema } from '~/schemas/url'
+import { HOME, SIGN_IN, type CallbackRoute } from '~/lib/routes'
 import { FullscreenSpinner } from '../FullscreenSpinner'
 
 interface EnforceLoginStatePageWrapperProps {
   /**
-   * Route to redirect to when user is not authenticated. Defaults to
-   * `SIGN_IN` route if not provided.
+   * Route to redirect to after user is authenticated. Defaults to
+   * `HOME` route if not provided.
    */
-  redirectTo?: string
+  redirectTo?: CallbackRoute
 }
 
-const Redirect = ({
-  redirectTo = SIGN_IN,
-}: EnforceLoginStatePageWrapperProps) => {
+const Redirect = ({ redirectTo = HOME }: EnforceLoginStatePageWrapperProps) => {
   const router = useRouter()
-  const redirectUrl = useMemo(() => {
-    if (typeof window === 'undefined') return '/'
-    const { pathname, search, hash } = window.location
-    return `${pathname}${search}${hash}`
-  }, [])
 
-  void router.replace(
-    callbackUrlSchema.parse(appendWithRedirect(redirectTo, redirectUrl)),
-  )
+  useEffect(() => {
+    void router.replace(appendWithRedirect(SIGN_IN, redirectTo))
+  }, [router, redirectTo])
 
   return <FullscreenSpinner />
 }
 
 /**
  * Page wrapper that renders children only if the login state localStorage flag has been set.
- * Otherwise, will redirect to the route passed into the `redirectTo` prop.
+ * Otherwise, will redirect to SIGN_IN, which will redirect to route passed into the `redirectTo` prop.
  *
  * @note 🚨 There is no authentication being performed by this component. This component is merely a wrapper that checks for the presence of the login flag in localStorage. This means that a user could add the flag and bypass the check. Any page children that require authentication should also perform authentication checks in that page itself!
  */
 export const EnforceLoginStatePageWrapper = ({
-  redirectTo = SIGN_IN,
+  redirectTo = HOME,
   children,
 }: PropsWithChildren<EnforceLoginStatePageWrapperProps>): React.ReactElement => {
   const { hasLoginStateFlag } = useLoginState()

src/components/AuthWrappers/PublicPageWrapper.tsx

@@ -1,13 +1,14 @@
-import { type PropsWithChildren } from 'react'
+import { useEffect, useState, type PropsWithChildren } from 'react'
 import { useRouter } from 'next/router'
 
-import { CALLBACK_URL_KEY } from '~/constants/params'
+import { getRedirectUrl } from '~/utils/url'
 import { useLoginState } from '~/features/auth'
+import { type CallbackRoute } from '~/lib/routes'
 import { callbackUrlSchema } from '~/schemas/url'
 import { FullscreenSpinner } from '../FullscreenSpinner'
 
 type PublicPageWrapperProps =
-  | { strict: true; redirectUrl?: string }
+  | { strict: true; redirectUrl?: CallbackRoute }
   | { strict: false }
 
 /**
@@ -23,14 +24,22 @@ export const PublicPageWrapper = ({
   const router = useRouter()
   const { hasLoginStateFlag } = useLoginState()
 
-  if (hasLoginStateFlag && rest.strict) {
-    if (rest.redirectUrl) {
-      void router.replace(callbackUrlSchema.parse(rest.redirectUrl))
+  const [isRedirecting, setIsRedirecting] = useState(true)
+
+  useEffect(() => {
+    if (hasLoginStateFlag && rest.strict) {
+      if (rest.redirectUrl) {
+        // must validate redirectUrl param
+        void router.replace(callbackUrlSchema.parse(rest.redirectUrl))
+      } else {
+        void router.replace(getRedirectUrl(router.query))
+      }
     } else {
-      void router.replace(
-        callbackUrlSchema.parse(router.query[CALLBACK_URL_KEY]),
-      )
+      setIsRedirecting(false)
     }
+  }, [hasLoginStateFlag, rest, router])
+
+  if (isRedirecting) {
     return <FullscreenSpinner />
   }
 

src/features/sign-in/components/EmailLogin/VerificationInput.tsx

@@ -17,12 +17,11 @@ import { Controller } from 'react-hook-form'
 import { useInterval } from 'usehooks-ts'
 
 import { trpc } from '~/utils/trpc'
-import { CALLBACK_URL_KEY } from '~/constants/params'
+import { getRedirectUrl } from '~/utils/url'
 import { useLoginState } from '~/features/auth'
 import { OTP_LENGTH } from '~/lib/auth'
 import { useZodForm } from '~/lib/form'
 import { emailVerifyOtpSchema } from '~/schemas/auth/email/sign-in'
-import { callbackUrlSchema } from '~/schemas/url'
 import { useSignInContext } from '../SignInContext'
 import { ResendOtpButton } from './ResendOtpButton'
 
@@ -61,7 +60,7 @@ export const VerificationInput = (): JSX.Element | null => {
       await utils.me.get.invalidate()
       // accessing router.query values returns decoded URI params automatically,
       // so there's no need to call decodeURIComponent manually when accessing the callback url.
-      await router.push(callbackUrlSchema.parse(router.query[CALLBACK_URL_KEY]))
+      await router.push(getRedirectUrl(router.query))
     },
     onError: (error) => {
       switch (error.message) {

src/features/sign-in/components/SgidCallback.tsx

@@ -2,9 +2,10 @@ import { useEffect } from 'react'
 import { useRouter } from 'next/router'
 
 import { trpc } from '~/utils/trpc'
+import { appendWithRedirect } from '~/utils/url'
 import { FullscreenSpinner } from '~/components/FullscreenSpinner'
 import { useLoginState } from '~/features/auth'
-import { callbackUrlSchema } from '~/schemas/url'
+import { SIGN_IN, SIGN_IN_SELECT_PROFILE } from '~/lib/routes'
 
 /**
  * This component is responsible for handling the callback from the SGID login.
@@ -19,27 +20,27 @@ export const SgidCallback = (): JSX.Element => {
     query: { code, state },
   } = router
 
-  const [{ redirectUrl, selectProfileStep }] =
-    trpc.auth.sgid.callback.useSuspenseQuery(
-      { code: String(code), state: String(state) },
-      { staleTime: Infinity },
-    )
+  const [response] = trpc.auth.sgid.callback.useSuspenseQuery(
+    { code: String(code), state: String(state) },
+    { staleTime: Infinity },
+  )
 
   useEffect(() => {
-    if (!selectProfileStep) {
-      setHasLoginStateFlag()
-      // eslint-disable-next-line @typescript-eslint/no-floating-promises
-      utils.me.get.invalidate()
+    if (!response.success) {
+      void router.replace(
+        `${SIGN_IN}?${new URLSearchParams({ error: response.reason })}`,
+      )
+    } else {
+      const { selectProfileStep, landingUrl } = response.data
+      if (!selectProfileStep) {
+        setHasLoginStateFlag()
+        void utils.me.get.invalidate()
+      }
+      void router.replace(
+        appendWithRedirect(SIGN_IN_SELECT_PROFILE, landingUrl),
+      )
     }
-    // eslint-disable-next-line @typescript-eslint/no-floating-promises
-    router.replace(callbackUrlSchema.parse(redirectUrl))
-  }, [
-    redirectUrl,
-    router,
-    selectProfileStep,
-    setHasLoginStateFlag,
-    utils.me.get,
-  ])
+  }, [response, router, setHasLoginStateFlag, utils.me.get])
 
   return <FullscreenSpinner />
 }

src/features/sign-in/components/SgidErrorFallback/SgidErrorFallback.tsx

@@ -4,21 +4,21 @@ import { type FallbackProps } from 'react-error-boundary'
 import { z } from 'zod'
 
 import { safeSchemaJsonParse } from '~/utils/zod'
-import { HOME } from '~/lib/routes'
+import { HOME, type CallbackRoute } from '~/lib/routes'
 import { callbackUrlSchema } from '~/schemas/url'
 import { SgidErrorModal } from './SgidErrorModal'
 
 export const SgidErrorFallback: ComponentType<FallbackProps> = ({ error }) => {
   const router = useRouter()
-  const redirectUrl = useMemo(() => {
+  const redirectUrl: CallbackRoute = useMemo(() => {
     const parsed = safeSchemaJsonParse(
       z.object({
         landingUrl: callbackUrlSchema,
       }),
       String(router.query.state),
     )
     if (parsed.success) {
-      return parsed.data.landingUrl.href
+      return parsed.data.landingUrl
     }
     return HOME
   }, [router.query.state])

src/features/sign-in/components/SgidErrorFallback/SgidErrorModal.tsx

@@ -14,11 +14,11 @@ import {
 
 import { appendWithRedirect } from '~/utils/url'
 import { SGID } from '~/lib/errors/auth.sgid'
-import { SIGN_IN } from '~/lib/routes'
+import { SIGN_IN, type CallbackRoute } from '~/lib/routes'
 
 interface SgidErrorModalProps {
   message: string
-  redirectUrl: string
+  redirectUrl: CallbackRoute
 }
 
 export const SgidErrorModal = ({

src/features/sign-in/components/SgidLogin/SgidLoginButton.tsx

@@ -5,7 +5,6 @@ import { Button } from '@opengovsg/design-system-react'
 import { trpc } from '~/utils/trpc'
 import { getRedirectUrl } from '~/utils/url'
 import { SingpassFullLogo } from '~/components/Svg/SingpassFullLogo'
-import { callbackUrlSchema } from '~/schemas/url'
 
 export const SgidLoginButton = (): JSX.Element | null => {
   const router = useRouter()
@@ -15,11 +14,9 @@ export const SgidLoginButton = (): JSX.Element | null => {
     },
   })
 
-  const landingUrl = callbackUrlSchema.parse(getRedirectUrl(router.query)).href
-
   const handleSgidLogin = () => {
     return sgidLoginMutation.mutate({
-      landingUrl,
+      landingUrl: getRedirectUrl(router.query),
     })
   }
 

src/features/sign-in/components/SgidProfileList/SgidProfileList.tsx

@@ -3,10 +3,9 @@ import { useRouter } from 'next/router'
 import { Divider, Stack } from '@chakra-ui/react'
 
 import { trpc } from '~/utils/trpc'
-import { CALLBACK_URL_KEY } from '~/constants/params'
+import { getRedirectUrl } from '~/utils/url'
 import { useLoginState } from '~/features/auth'
 import { withSuspense } from '~/hocs/withSuspense'
-import { callbackUrlSchema } from '~/schemas/url'
 import { SgidProfileItem } from './SgidProfileItem'
 import { SgidProfileListSkeleton } from './SgidProfileListSkeleton'
 
@@ -22,9 +21,7 @@ const SuspendableSgidProfileList = (): JSX.Element => {
     onSuccess: async () => {
       setHasLoginStateFlag()
       await utils.me.get.invalidate()
-      await router.replace(
-        callbackUrlSchema.parse(router.query[CALLBACK_URL_KEY]),
-      )
+      await router.replace(getRedirectUrl(router.query))
     },
   })
 

src/lib/routes.ts

@@ -1,6 +1,16 @@
 export const SIGN_IN = '/sign-in'
-export const SIGN_IN_SELECT_PROFILE_SUBROUTE = '/select-profile'
+export const SIGN_IN_SELECT_PROFILE = '/sign-in/select-profile'
 
 export const HOME = '/home'
 export const PROFILE = '/profile'
 export const SETTINGS_PROFILE = '/settings/profile'
+
+export const ALLOWED_CALLBACK_ROUTES = [
+  SIGN_IN,
+  SIGN_IN_SELECT_PROFILE,
+  HOME,
+  PROFILE,
+  SETTINGS_PROFILE,
+] as const
+
+export type CallbackRoute = (typeof ALLOWED_CALLBACK_ROUTES)[number]

src/schemas/url.ts

@@ -2,7 +2,7 @@ import { createUrlSchema } from '@opengovsg/starter-kitty-validators/url'
 import { z } from 'zod'
 
 import { getBaseUrl } from '~/utils/getBaseUrl'
-import { HOME } from '~/lib/routes'
+import { ALLOWED_CALLBACK_ROUTES, HOME } from '~/lib/routes'
 
 const baseUrl = new URL(getBaseUrl())
 
@@ -14,7 +14,10 @@ const urlSchema = createUrlSchema({
   },
 })
 
-export const callbackUrlSchema = z
+/**
+ * Zod schema for validating internal (same-app) URLs
+ */
+export const appUrlSchema = z
   .string()
   .optional()
   .default(HOME)
@@ -30,3 +33,8 @@ export const callbackUrlSchema = z
     }
   })
   .catch(new URL(HOME, baseUrl.origin))
+
+/**
+ * Zod schema for validating callbackUrls
+ */
+export const callbackUrlSchema = z.enum(ALLOWED_CALLBACK_ROUTES).catch(HOME)

src/server/modules/auth/sgid/sgid.router.ts

@@ -7,10 +7,12 @@ import set from 'lodash/set'
 import { z } from 'zod'
 
 import { trpcAssert } from '~/utils/trpcAssert'
-import { appendWithRedirect } from '~/utils/url'
-import { normaliseEmail, safeSchemaJsonParse } from '~/utils/zod'
+import {
+  createResponseSchema,
+  normaliseEmail,
+  safeSchemaJsonParse,
+} from '~/utils/zod'
 import { SGID } from '~/lib/errors/auth.sgid'
-import { SIGN_IN, SIGN_IN_SELECT_PROFILE_SUBROUTE } from '~/lib/routes'
 import { APP_SGID_SCOPE, sgid } from '~/lib/sgid'
 import { callbackUrlSchema } from '~/schemas/url'
 import { publicProcedure, router } from '~/server/trpc'
@@ -74,6 +76,14 @@ export const sgidRouter = router({
         code: z.string(),
       }),
     )
+    .output(
+      createResponseSchema(
+        z.object({
+          selectProfileStep: z.boolean(),
+          landingUrl: callbackUrlSchema,
+        }),
+      ),
+    )
     .query(async ({ ctx, input: { state, code } }) => {
       if (!sgid) {
         ctx.logger.error('SGID is not enabled')
@@ -107,15 +117,13 @@ export const sgidRouter = router({
 
       try {
         sgidUserInfo = await getUserInfo({ code, codeVerifier, nonce })
-      } catch (error) {
+      } catch {
         ctx.logger.warn({ state }, 'Unable to fetch user info from sgID')
         // Redirect back to sign in page with error.
         // TODO: Change this to throw an error instead, and then handle it in the client.
         return {
-          redirectUrl: `/sign-in?error=${
-            (error as Error).message ||
-            'Something went wrong whilst fetching SGID user info'
-          }`,
+          success: false,
+          reason: 'Something went wrong whilst fetching SGID user info',
         }
       }
 
@@ -151,11 +159,11 @@ export const sgidRouter = router({
       set(ctx.session, 'sgid.profiles', sgidProfileToStore.data)
       await ctx.session.save()
       return {
-        selectProfileStep: true,
-        redirectUrl: appendWithRedirect(
-          `${SIGN_IN}${SIGN_IN_SELECT_PROFILE_SUBROUTE}`,
-          parsedState.data.landingUrl.href,
-        ),
+        success: true,
+        data: {
+          selectProfileStep: true,
+          landingUrl: parsedState.data.landingUrl,
+        },
       }
     }),
   listStoredProfiles: publicProcedure.query(({ ctx }) => {