chore: improve sgid state validation (#386)

* chore: improve sgid state validation

* fix: add validation for redirectUrl after sgid oidc login

* fix: validate landingUrl clientside in first leg of sgid flow

* fix: revert wrong validation on oidc redirection

* fix: wrong type on redirectUrl

* fix: wrong type passed in sgid login mutation

* fix: update callbackUrlSchema to fail gracefully

Author: zhongliang02

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

@@ -5,19 +5,20 @@ import { z } from 'zod'
 
 import { safeSchemaJsonParse } from '~/utils/zod'
 import { HOME } 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 parsed = safeSchemaJsonParse(
       z.object({
-        landingUrl: z.string(),
+        landingUrl: callbackUrlSchema,
       }),
       String(router.query.state),
     )
     if (parsed.success) {
-      return parsed.data.landingUrl
+      return parsed.data.landingUrl.href
     }
     return HOME
   }, [router.query.state])

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

@@ -5,6 +5,7 @@ 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()
@@ -14,7 +15,7 @@ export const SgidLoginButton = (): JSX.Element | null => {
     },
   })
 
-  const landingUrl = getRedirectUrl(router.query)
+  const landingUrl = callbackUrlSchema.parse(getRedirectUrl(router.query)).href
 
   const handleSgidLogin = () => {
     return sgidLoginMutation.mutate({

src/schemas/url.ts

@@ -6,17 +6,27 @@ import { HOME } from '~/lib/routes'
 
 const baseUrl = new URL(getBaseUrl())
 
+const urlSchema = createUrlSchema({
+  baseOrigin: baseUrl.origin,
+  whitelist: {
+    protocols: ['http', 'https'],
+    hosts: [baseUrl.host],
+  },
+})
+
 export const callbackUrlSchema = z
   .string()
   .optional()
   .default(HOME)
-  .pipe(
-    createUrlSchema({
-      baseOrigin: baseUrl.origin,
-      whitelist: {
-        protocols: ['http', 'https'],
-        hosts: [baseUrl.host],
-      },
-    }),
-  )
+  .transform((val, ctx) => {
+    try {
+      return urlSchema.parse(val)
+    } catch {
+      ctx.addIssue({
+        code: z.ZodIssueCode.custom,
+        message: 'URL validation error',
+      })
+      return z.NEVER
+    }
+  })
   .catch(new URL(HOME, baseUrl.origin))

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

@@ -22,16 +22,12 @@ import {
 } from './sgid.utils'
 
 const sgidCallbackStateSchema = z.object({
-  landingUrl: z.string(),
+  landingUrl: callbackUrlSchema,
 })
 
 export const sgidRouter = router({
   login: publicProcedure
-    .input(
-      z.object({
-        landingUrl: callbackUrlSchema,
-      }),
-    )
+    .input(sgidCallbackStateSchema)
     .mutation(async ({ ctx, input: { landingUrl } }) => {
       if (!sgid) {
         throw new TRPCError({
@@ -158,7 +154,7 @@ export const sgidRouter = router({
         selectProfileStep: true,
         redirectUrl: appendWithRedirect(
           `${SIGN_IN}${SIGN_IN_SELECT_PROFILE_SUBROUTE}`,
-          parsedState.data.landingUrl,
+          parsedState.data.landingUrl.href,
         ),
       }
     }),