From 88ace8073be7e105b26728b02921bbab7907f362 Mon Sep 17 00:00:00 2001
From: Darcy Ye <darcyye@silverhand.io>
Date: Tue, 19 Mar 2024 17:26:38 +0800
Subject: [PATCH] refactor(core): use discriminate union for custom jwt test
 API

---
 packages/core/src/routes/logto-config.ts | 51 ++++++++++++++----------
 1 file changed, 30 insertions(+), 21 deletions(-)

diff --git a/packages/core/src/routes/logto-config.ts b/packages/core/src/routes/logto-config.ts
index fd10011b6f1..9c4f712fc3e 100644
--- a/packages/core/src/routes/logto-config.ts
+++ b/packages/core/src/routes/logto-config.ts
@@ -17,7 +17,7 @@ import {
   LogtoJwtTokenKey,
   LogtoJwtTokenPath,
   jsonObjectGuard,
-  customJwtFetcherGuard,
+  type CustomJwtFetcher,
 } from '@logto/schemas';
 import { z } from 'zod';
 
@@ -296,38 +296,47 @@ export default function logtoConfigRoutes<T extends AuthedRouter>(
   }
 
   router.post(
-    '/configs/jwt-customizer/:tokenTypePath/test',
+    '/configs/jwt-customizer/test',
     koaGuard({
-      params: z.object({
-        tokenTypePath: z.nativeEnum(LogtoJwtTokenPath),
-      }),
-      body: z.unknown(),
+      body: z.discriminatedUnion('tokenType', [
+        z.object({
+          tokenType: z.literal(LogtoJwtTokenKey.AccessToken),
+          payload: accessTokenJwtCustomizerGuard,
+        }),
+        z.object({
+          tokenType: z.literal(LogtoJwtTokenKey.ClientCredentials),
+          payload: clientCredentialsJwtCustomizerGuard,
+        }),
+      ]),
       response: jsonObjectGuard,
       /**
-       * 400 for cloud service zod error (data type does not match expectation, can be either request body or response body)
-       * 422 for cloud service syntax error
+       * Code 400 indicates Zod errors in cloud service (data type does not match expectation, can be either request body or response body).
+       * Code 422 indicates syntax errors in cloud service.
        */
       status: [200, 400, 422],
     }),
     async (ctx, next) => {
       const {
-        params: { tokenTypePath },
-        body: rawBody,
+        body: {
+          payload: { tokenSample, contextSample, ...rest },
+        },
       } = ctx.guard;
-      const {
-        body: { tokenSample, contextSample, ...rest },
-      } = getJwtTokenKeyAndBody(tokenTypePath, rawBody);
 
+      /**
+       * We have ensured the API request body via koa guard, manually cast the cloud service API call's
+       * `requestBody` type and let the cloud service API to throw if needed.
+       */
+      // eslint-disable-next-line no-restricted-syntax, @typescript-eslint/consistent-type-assertions
+      const requestBody = {
+        ...rest,
+        token: tokenSample,
+        context: contextSample,
+      } as CustomJwtFetcher;
       const client = await cloudConnection.getClient();
-      const testResult = await client.post(`/api/services/custom-jwt`, {
-        body: customJwtFetcherGuard.parse({
-          ...rest,
-          tokenSample,
-          contextSample,
-        }),
-      });
 
-      ctx.body = testResult;
+      ctx.body = await client.post(`/api/services/custom-jwt`, {
+        body: requestBody,
+      });
       return next();
     }
   );