From 731410747588424bf21450da88fcd5ad58c8dc6a Mon Sep 17 00:00:00 2001
From: Darcy Ye <darcyye@silverhand.io>
Date: Mon, 27 Jan 2025 00:19:13 +0800
Subject: [PATCH] refactor: refactor API query guard to fit swagger flow

---
 .../src/routes/saml-application/anonymous.ts  | 49 ++++++++++++++-----
 1 file changed, 38 insertions(+), 11 deletions(-)

diff --git a/packages/core/src/routes/saml-application/anonymous.ts b/packages/core/src/routes/saml-application/anonymous.ts
index 201b7acac19..0a0a394e415 100644
--- a/packages/core/src/routes/saml-application/anonymous.ts
+++ b/packages/core/src/routes/saml-application/anonymous.ts
@@ -16,17 +16,15 @@ import { generateAutoSubmitForm } from '#src/saml-application/SamlApplication/ut
 import assertThat from '#src/utils/assert-that.js';
 import { getConsoleLogFromContext } from '#src/utils/console.js';
 
-const samlApplicationSignInCallbackQueryParametersGuard = z.union([
-  z.object({
+const samlApplicationSignInCallbackQueryParametersGuard = z
+  .object({
     code: z.string(),
-    state: z.string().optional(),
-    redirectUri: z.string().optional(),
-  }),
-  z.object({
+    state: z.string(),
+    redirectUri: z.string(),
     error: z.string(),
-    error_description: z.string().optional(),
-  }),
-]);
+    error_description: z.string(),
+  })
+  .partial();
 
 export default function samlApplicationAnonymousRoutes<T extends AnonymousRouter>(
   ...[router, { id: tenantId, libraries, queries, envSet }]: RouterInitArgs<T>
@@ -70,6 +68,7 @@ export default function samlApplicationAnonymousRoutes<T extends AnonymousRouter
       status: [200, 400, 404],
     }),
     koaAuditLog(queries),
+    // eslint-disable-next-line complexity
     async (ctx, next) => {
       const consoleLog = getConsoleLogFromContext(ctx);
       const {
@@ -84,11 +83,39 @@ export default function samlApplicationAnonymousRoutes<T extends AnonymousRouter
         applicationId: id,
       });
 
+      // Validate query parameters
+      if (!query.code && !query.error) {
+        throw new RequestError({
+          code: 'guard.invalid_input',
+          message: 'Either code or error must be present',
+          type: 'query',
+        });
+      }
+
+      // eslint-disable-next-line @typescript-eslint/prefer-nullish-coalescing
+      if (query.code && (query.error || query.error_description)) {
+        throw new RequestError({
+          code: 'guard.invalid_input',
+          type: 'query',
+          message: 'Cannot have both code and error fields',
+        });
+      }
+
+      // eslint-disable-next-line @typescript-eslint/prefer-nullish-coalescing
+      if (query.error && (query.code || query.state || query.redirectUri)) {
+        throw new RequestError({
+          code: 'guard.invalid_input',
+          type: 'query',
+          message: 'When error is present, only error_description is allowed',
+        });
+      }
+
       // Handle error in query parameters
-      if ('error' in query) {
+      if (query.error) {
         throw new RequestError({
-          code: 'oidc.invalid_request',
+          code: 'guard.invalid_input',
           message: query.error_description,
+          type: 'query',
         });
       }