feat(core): add denyAccess api context to customJwt script (#6532)

* feat(schemas,core): add denyAccess api conext to custom jwt

add denyAccess api context to the custom jwt

* fix(test): fix integration test

fix integration test

* chore(schemas): update type name

update api context type name

* chore(schemas): fix typo

fix typo

* feat(core): add dev feature guard

add dev feature guard

packages/core/src/libraries/jwt-customizer.ts

@@ -1,4 +1,6 @@
 import {
+  type CustomJwtErrorBody,
+  CustomJwtErrorCode,
   LogtoJwtTokenKeyType,
   jwtCustomizerUserContextGuard,
   userInfoSelectFields,
@@ -6,9 +8,11 @@ import {
   type JwtCustomizerType,
   type JwtCustomizerUserContext,
   type LogtoJwtTokenKey,
+  type CustomJwtApiContext,
+  type CustomJwtScriptPayload,
 } from '@logto/schemas';
 import { type ConsoleLog } from '@logto/shared';
-import { assert, deduplicate, pick, pickState } from '@silverhand/essentials';
+import { assert, conditional, deduplicate, pick, pickState } from '@silverhand/essentials';
 import deepmerge from 'deepmerge';
 import { ZodError, z } from 'zod';
 
@@ -28,19 +32,49 @@ import {
 
 import { type CloudConnectionLibrary } from './cloud-connection.js';
 
+const apiContext: CustomJwtApiContext = Object.freeze({
+  denyAccess: (message = 'Access denied') => {
+    const error: CustomJwtErrorBody = {
+      code: CustomJwtErrorCode.AccessDenied,
+      message,
+    };
+
+    throw new LocalVmError(
+      {
+        message,
+        error,
+      },
+      403
+    );
+  },
+});
+
 export class JwtCustomizerLibrary {
   // Convert errors to WithTyped client response error to share the error handling logic.
   static async runScriptInLocalVm(data: CustomJwtFetcher) {
     try {
-      const payload =
-        data.tokenType === LogtoJwtTokenKeyType.AccessToken
+      // @ts-expect-error -- remove this when the dev feature is ready
+      const payload: CustomJwtScriptPayload = {
+        ...(data.tokenType === LogtoJwtTokenKeyType.AccessToken
           ? pick(data, 'token', 'context', 'environmentVariables')
-          : pick(data, 'token', 'environmentVariables');
+          : pick(data, 'token', 'environmentVariables')),
+        ...conditional(
+          // TODO: @simeng remove this when the dev feature is ready
+          EnvSet.values.isDevFeaturesEnabled && {
+            api: apiContext,
+          }
+        ),
+      };
+
       const result = await runScriptFunctionInLocalVm(data.script, 'getCustomJwtClaims', payload);
 
       // If the `result` is not a record, we cannot merge it to the existing token payload.
       return z.record(z.unknown()).parse(result);
     } catch (error: unknown) {
+      if (error instanceof LocalVmError) {
+        throw error;
+      }
+
       // Assuming we only use zod for request body validation
       if (error instanceof ZodError) {
         const { errors } = error;

packages/core/src/oidc/extra-token-claims.ts

@@ -6,10 +6,12 @@ import {
   jwtCustomizer as jwtCustomizerLog,
   type CustomJwtFetcher,
   GrantType,
+  CustomJwtErrorCode,
 } from '@logto/schemas';
 import { generateStandardId } from '@logto/shared';
 import { conditional, trySafe } from '@silverhand/essentials';
-import { type KoaContextWithOIDC, type UnknownObject } from 'oidc-provider';
+import { ResponseError } from '@withtyped/client';
+import { errors, type KoaContextWithOIDC, type UnknownObject } from 'oidc-provider';
 import { z } from 'zod';
 
 import { EnvSet } from '#src/env-set/index.js';
@@ -20,6 +22,8 @@ import { LogEntry } from '#src/middleware/koa-audit-log.js';
 import type Libraries from '#src/tenants/Libraries.js';
 import type Queries from '#src/tenants/Queries.js';
 
+import { parseCustomJwtResponseError } from '../utils/custom-jwt/index.js';
+
 import { tokenExchangeActGuard } from './grants/token-exchange/types.js';
 
 /**
@@ -196,11 +200,14 @@ export const getExtraTokenClaimsForJwtCustomization = async (
           : jwtCustomizerLog.Type.AccessToken
       }`
     );
+
     entry.append({
       result: LogResult.Error,
       error: { message: String(error) },
     });
+
     const { payload } = entry;
+
     await queries.logs.insertLog({
       id: generateStandardId(),
       key: payload.key,
@@ -210,6 +217,20 @@ export const getExtraTokenClaimsForJwtCustomization = async (
         token,
       },
     });
+
+    // TODO: @simeng remove this once the feature is ready
+    if (!EnvSet.values.isDevFeaturesEnabled) {
+      return;
+    }
+
+    // If the error is an instance of `ResponseError`, we need to parse the customJwtError body to get the error code.
+    if (error instanceof ResponseError) {
+      const customJwtError = await trySafe(async () => parseCustomJwtResponseError(error));
+
+      if (customJwtError?.code === CustomJwtErrorCode.AccessDenied) {
+        throw new errors.AccessDenied(customJwtError.message);
+      }
+    }
   }
 };
 /* eslint-enable complexity */

packages/core/src/routes/logto-config/jwt-customizer.ts

@@ -1,4 +1,5 @@
 import {
+  CustomJwtErrorCode,
   LogtoJwtTokenKey,
   LogtoJwtTokenKeyType,
   accessTokenJwtCustomizerGuard,
@@ -18,6 +19,7 @@ import { JwtCustomizerLibrary } from '#src/libraries/jwt-customizer.js';
 import koaGuard, { parse } from '#src/middleware/koa-guard.js';
 import { koaQuotaGuard } from '#src/middleware/koa-quota-guard.js';
 import { getConsoleLogFromContext } from '#src/utils/console.js';
+import { parseCustomJwtResponseError } from '#src/utils/custom-jwt/index.js';
 
 import type { ManagementApiRouter, RouterInitArgs } from '../types.js';
 
@@ -249,8 +251,11 @@ export default function logtoConfigJwtCustomizerRoutes<T extends ManagementApiRo
          * format of `RequestError`, we manually transform it here to keep the error format consistent.
          */
         if (error instanceof ResponseError) {
-          const { message } = z.object({ message: z.string() }).parse(await error.response.json());
-          throw new RequestError({ code: 'jwt_customizer.general', status: 422 }, { message });
+          const { code, message } = await parseCustomJwtResponseError(error);
+
+          const status = code === CustomJwtErrorCode.AccessDenied ? 403 : 422;
+
+          throw new RequestError({ code: 'jwt_customizer.general', status }, { message, code });
         }
 
         if (error instanceof ZodError) {

packages/core/src/utils/custom-jwt/index.ts

@@ -1,4 +1,13 @@
-import { LogtoJwtTokenKey, type JwtCustomizerType } from '@logto/schemas';
+import {
+  LogtoJwtTokenKey,
+  type JwtCustomizerType,
+  customJwtErrorBodyGuard,
+  type CustomJwtErrorBody,
+} from '@logto/schemas';
+import { type ResponseError } from '@withtyped/client';
+import { z } from 'zod';
+
+import RequestError from '../../errors/RequestError/index.js';
 
 import { type CustomJwtDeployRequestBody } from './types.js';
 
@@ -11,3 +20,52 @@ export const getJwtCustomizerScripts = (jwtCustomizers: Partial<JwtCustomizerTyp
     Object.values(LogtoJwtTokenKey).map((key) => [key, { production: jwtCustomizers[key]?.script }])
   ) as CustomJwtDeployRequestBody;
 };
+
+/**
+ * Parse the withtyped ResponseError body.
+ * @see {@link https://github.com/withtyped/withtyped/blob/master/packages/server/src/index.ts} handleError method
+ */
+const errorResponseGuard = z.object({
+  message: z.string(),
+  error: z.unknown().optional(),
+});
+
+/**
+ * Parse the CustomJwtErrorBody from the response.
+ *
+ * @param {ResponseError} error The response error, should always be the cloud service ResponseError type.
+ * @returns {Promise<CustomJwtErrorBody>} The parsed CustomJwtErrorBody.
+ * @throws {RequestError} error if the response is not a standard CustomJwtError.
+ */
+export const parseCustomJwtResponseError = async (
+  error: ResponseError
+): Promise<CustomJwtErrorBody> => {
+  const errorResponse = errorResponseGuard.safeParse(await error.response.json());
+
+  // Can not parse the ResponseError body.
+  if (!errorResponse.success) {
+    throw new RequestError(
+      {
+        code: 'jwt_customizer.general',
+        status: 500,
+      },
+      { message: error.message }
+    );
+  }
+
+  const errorResponseData = errorResponse.data;
+  const errorBody = customJwtErrorBodyGuard.safeParse(errorResponseData.error);
+
+  // Not a standard CustomJwtError body.
+  if (!errorBody.success) {
+    throw new RequestError(
+      {
+        code: 'jwt_customizer.general',
+        status: 422,
+      },
+      { message: errorResponseData.message }
+    );
+  }
+
+  return errorBody.data;
+};

packages/integration-tests/src/__mocks__/jwt-customizer.ts

@@ -58,6 +58,16 @@ export const accessTokenSampleScript = `const getCustomJwtClaims = async ({ toke
   return { user_id: context?.user?.id ?? 'unknown', hasPassword: context?.user?.hasPassword };
 };`;
 
+export const accessTokenAccessDeniedSampleScript = `const getCustomJwtClaims = async ({ token, context, environmentVariables, api }) => {
+  api.denyAccess('You are not allowed to access this resource');
+  return { test: 'foo'};
+};`;
+
 export const clientCredentialsSampleScript = `const getCustomJwtClaims = async ({ token, context, environmentVariables }) => {
   return { ...environmentVariables };
 }`;
+
+export const clientCredentialsAccessDeniedSampleScript = `const getCustomJwtClaims = async ({ token, context, environmentVariables, api }) => {
+  api.denyAccess('You are not allowed to access this resource');
+  return { test: 'foo'};
+};`;

packages/integration-tests/src/tests/api/logto-config.test.ts

@@ -11,6 +11,8 @@ import {
   clientCredentialsJwtCustomizerPayload,
   accessTokenSampleScript,
   clientCredentialsSampleScript,
+  accessTokenAccessDeniedSampleScript,
+  clientCredentialsAccessDeniedSampleScript,
 } from '#src/__mocks__/jwt-customizer.js';
 import {
   deleteOidcKey,
@@ -26,13 +28,14 @@ import {
   testJwtCustomizer,
 } from '#src/api/index.js';
 import { expectRejects } from '#src/helpers/index.js';
+import { devFeatureTest } from '#src/utils.js';
 
 const defaultAdminConsoleConfig: AdminConsoleData = {
   signInExperienceCustomized: false,
   organizationCreated: false,
 };
 
-describe('admin console sign-in experience', () => {
+describe('logto config', () => {
   it('should get admin console config successfully', async () => {
     const adminConsoleConfig = await getAdminConsoleConfig();
 
@@ -268,4 +271,41 @@ describe('admin console sign-in experience', () => {
     });
     expect(testResult).toMatchObject(clientCredentialsJwtCustomizerPayload.environmentVariables);
   });
+
+  devFeatureTest.it(
+    'should throw access denied error when calling the denyAccess api in the script',
+    async () => {
+      await expectRejects(
+        testJwtCustomizer({
+          tokenType: LogtoJwtTokenKeyType.AccessToken,
+          token: accessTokenJwtCustomizerPayload.tokenSample,
+          context: accessTokenJwtCustomizerPayload.contextSample,
+          script: accessTokenAccessDeniedSampleScript,
+          environmentVariables: accessTokenJwtCustomizerPayload.environmentVariables,
+        }),
+        {
+          code: 'jwt_customizer.general',
+          status: 403,
+        }
+      );
+    }
+  );
+
+  devFeatureTest.it(
+    'should throw access denied error when calling the denyAccess api in the script',
+    async () => {
+      await expectRejects(
+        testJwtCustomizer({
+          tokenType: LogtoJwtTokenKeyType.ClientCredentials,
+          token: clientCredentialsJwtCustomizerPayload.tokenSample,
+          script: clientCredentialsAccessDeniedSampleScript,
+          environmentVariables: clientCredentialsJwtCustomizerPayload.environmentVariables,
+        }),
+        {
+          code: 'jwt_customizer.general',
+          status: 403,
+        }
+      );
+    }
+  );
 });

packages/schemas/src/types/logto-config/jwt-customizer.ts

@@ -145,3 +145,50 @@ export const customJwtFetcherGuard = z.discriminatedUnion('tokenType', [
 ]);
 
 export type CustomJwtFetcher = z.infer<typeof customJwtFetcherGuard>;
+
+export enum CustomJwtErrorCode {
+  /**
+   * The `AccessDenied` error explicitly thrown
+   * by calling the `api.denyAccess` function in the custom JWT script.
+   */
+  AccessDenied = 'AccessDenied',
+  /** General JWT customizer error,
+   * this is the fallback custom jwt error code
+   * for any internal error thrown by the JWT customizer (localVM, azure function, or CF worker).
+   */
+  General = 'General',
+}
+
+export const customJwtErrorBodyGuard = z.object({
+  code: z.nativeEnum(CustomJwtErrorCode),
+  message: z.string(),
+});
+
+export type CustomJwtErrorBody = z.infer<typeof customJwtErrorBodyGuard>;
+
+export type CustomJwtApiContext = {
+  /**
+   * Reject the the current token exchange request.
+   *
+   * @remarks
+   * By calling this function, the current token exchange request will be rejected,
+   * and a OIDC `AccessDenied` error will be thrown to the client with the given message.
+   *
+   * @param message The message to be shown to the user.
+   * @throws {ResponseError} with `CustomJwtErrorBody`
+   */
+  denyAccess: (message?: string) => never;
+};
+
+/**
+ * The payload type for the custom JWT script.
+ *
+ * @remarks
+ * We use this type to guard the input payload for the custom JWT script.
+ */
+export type CustomJwtScriptPayload = {
+  token: Record<string, unknown>;
+  context?: Record<string, unknown>;
+  environmentVariables?: Record<string, string>;
+  api: CustomJwtApiContext;
+};