feat(core,console): enable custom JWT for OSS and can run script in local vm (#5794)

.changeset/itchy-eels-remain.md

@@ -0,0 +1,8 @@
+---
+"@logto/console": minor
+"@logto/core": minor
+---
+
+enable custom JWT feature for OSS version
+
+OSS version users can now use custom JWT feature to add custom claims to JWT access tokens payload (previously, this feature was only available to Logto Cloud).

packages/console/src/containers/ConsoleContent/Sidebar/hook.tsx

@@ -130,7 +130,6 @@ export const useSidebarMenuItems = (): {
         {
           Icon: JwtClaims,
           title: 'customize_jwt',
-          isHidden: !isCloud,
         },
         {
           Icon: Hook,

packages/console/src/hooks/use-console-routes/index.tsx

@@ -63,7 +63,7 @@ export const useConsoleRoutes = () => {
         },
         { path: 'signing-keys', element: <SigningKeys /> },
         isCloud && tenantSettings,
-        isCloud && customizeJwt
+        customizeJwt
       ),
     [tenantSettings]
   );

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

@@ -1,14 +1,19 @@
+import { runScriptFunctionInLocalVm, buildErrorResponse } from '@logto/core-kit/custom-jwt';
 import {
   userInfoSelectFields,
   jwtCustomizerUserContextGuard,
   type LogtoJwtTokenKey,
   type JwtCustomizerType,
   type JwtCustomizerUserContext,
+  type CustomJwtFetcher,
+  LogtoJwtTokenKeyType,
 } from '@logto/schemas';
 import { type ConsoleLog } from '@logto/shared';
 import { deduplicate, pick, pickState, assert } from '@silverhand/essentials';
 import deepmerge from 'deepmerge';
+import { z, ZodError } from 'zod';
 
+import { EnvSet } from '#src/env-set/index.js';
 import RequestError from '#src/errors/RequestError/index.js';
 import type { LogtoConfigLibrary } from '#src/libraries/logto-config.js';
 import { type ScopeLibrary } from '#src/libraries/scope.js';
@@ -18,26 +23,49 @@ import {
   getJwtCustomizerScripts,
   type CustomJwtDeployRequestBody,
 } from '#src/utils/custom-jwt/index.js';
+import { LocalVmError } from '#src/utils/custom-jwt/index.js';
 
 import { type CloudConnectionLibrary } from './cloud-connection.js';
 
-export const createJwtCustomizerLibrary = (
-  queries: Queries,
-  logtoConfigs: LogtoConfigLibrary,
-  cloudConnection: CloudConnectionLibrary,
-  userLibrary: UserLibrary,
-  scopeLibrary: ScopeLibrary
-) => {
-  const {
-    users: { findUserById },
-    rolesScopes: { findRolesScopesByRoleIds },
-    scopes: { findScopesByIds },
-    userSsoIdentities,
-    organizations: { relations },
-  } = queries;
-  const { findUserRoles } = userLibrary;
-  const { attachResourceToScopes } = scopeLibrary;
-  const { getJwtCustomizers } = logtoConfigs;
+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
+          ? pick(data, 'token', 'context', 'environmentVariables')
+          : pick(data, 'token', 'environmentVariables');
+      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) {
+      // Assuming we only use zod for request body validation
+      if (error instanceof ZodError) {
+        const { errors } = error;
+        throw new LocalVmError(
+          {
+            message: 'Invalid input',
+            errors,
+          },
+          400
+        );
+      }
+
+      throw new LocalVmError(
+        buildErrorResponse(error),
+        error instanceof SyntaxError || error instanceof TypeError ? 422 : 500
+      );
+    }
+  }
+
+  constructor(
+    private readonly queries: Queries,
+    private readonly logtoConfigs: LogtoConfigLibrary,
+    private readonly cloudConnection: CloudConnectionLibrary,
+    private readonly userLibrary: UserLibrary,
+    private readonly scopeLibrary: ScopeLibrary
+  ) {}
 
   /**
    * We does not include org roles' scopes for the following reason:
@@ -45,15 +73,20 @@ export const createJwtCustomizerLibrary = (
    * these APIs from console setup while this library method is a backend used method.
    * 2. Logto developers can get the org roles' id from this user context and hence query the org roles' scopes via management API.
    */
-  const getUserContext = async (userId: string): Promise<JwtCustomizerUserContext> => {
-    const user = await findUserById(userId);
-    const fullSsoIdentities = await userSsoIdentities.findUserSsoIdentitiesByUserId(userId);
-    const roles = await findUserRoles(userId);
-    const rolesScopes = await findRolesScopesByRoleIds(roles.map(({ id }) => id));
+  async getUserContext(userId: string): Promise<JwtCustomizerUserContext> {
+    const user = await this.queries.users.findUserById(userId);
+    const fullSsoIdentities = await this.queries.userSsoIdentities.findUserSsoIdentitiesByUserId(
+      userId
+    );
+    const roles = await this.userLibrary.findUserRoles(userId);
+    const rolesScopes = await this.queries.rolesScopes.findRolesScopesByRoleIds(
+      roles.map(({ id }) => id)
+    );
     const scopeIds = rolesScopes.map(({ scopeId }) => scopeId);
-    const scopes = await findScopesByIds(scopeIds);
-    const scopesWithResources = await attachResourceToScopes(scopes);
-    const organizationsWithRoles = await relations.users.getOrganizationsByUserId(userId);
+    const scopes = await this.queries.scopes.findScopesByIds(scopeIds);
+    const scopesWithResources = await this.scopeLibrary.attachResourceToScopes(scopes);
+    const organizationsWithRoles =
+      await this.queries.organizations.relations.users.getOrganizationsByUserId(userId);
     const userContext = {
       ...pick(user, ...userInfoSelectFields),
       ssoIdentities: fullSsoIdentities.map(pickState('issuer', 'identityId', 'detail')),
@@ -81,7 +114,7 @@ export const createJwtCustomizerLibrary = (
     };
 
     return jwtCustomizerUserContextGuard.parse(userContext);
-  };
+  }
 
   /**
    * This method is used to deploy the give JWT customizer scripts to the cloud worker service.
@@ -95,17 +128,24 @@ export const createJwtCustomizerLibrary = (
    * @params payload.value - JWT customizer value
    * @params payload.useCase - The use case of JWT customizer script, can be either `test` or `production`.
    */
-  const deployJwtCustomizerScript = async <T extends LogtoJwtTokenKey>(
+  async deployJwtCustomizerScript<T extends LogtoJwtTokenKey>(
     consoleLog: ConsoleLog,
     payload: {
       key: T;
       value: JwtCustomizerType[T];
       useCase: 'test' | 'production';
     }
-  ) => {
+  ) {
+    if (!EnvSet.values.isCloud) {
+      consoleLog.warn(
+        'Early terminate `deployJwtCustomizerScript` since we do not provide dedicated computing resource for OSS version.'
+      );
+      return;
+    }
+
     const [client, jwtCustomizers] = await Promise.all([
-      cloudConnection.getClient(),
-      getJwtCustomizers(consoleLog),
+      this.cloudConnection.getClient(),
+      this.logtoConfigs.getJwtCustomizers(consoleLog),
     ]);
 
     const customizerScriptsFromDatabase = getJwtCustomizerScripts(jwtCustomizers);
@@ -129,15 +169,19 @@ export const createJwtCustomizerLibrary = (
     await client.put(`/api/services/custom-jwt/worker`, {
       body: deepmerge(customizerScriptsFromDatabase, newCustomizerScripts),
     });
-  };
+  }
+
+  async undeployJwtCustomizerScript<T extends LogtoJwtTokenKey>(consoleLog: ConsoleLog, key: T) {
+    if (!EnvSet.values.isCloud) {
+      consoleLog.warn(
+        'Early terminate `undeployJwtCustomizerScript` since we do not deploy the script to dedicated computing resource for OSS version.'
+      );
+      return;
+    }
 
-  const undeployJwtCustomizerScript = async <T extends LogtoJwtTokenKey>(
-    consoleLog: ConsoleLog,
-    key: T
-  ) => {
     const [client, jwtCustomizers] = await Promise.all([
-      cloudConnection.getClient(),
-      getJwtCustomizers(consoleLog),
+      this.cloudConnection.getClient(),
+      this.logtoConfigs.getJwtCustomizers(consoleLog),
     ]);
 
     assert(jwtCustomizers[key], new RequestError({ code: 'entity.not_exists', key }));
@@ -160,10 +204,5 @@ export const createJwtCustomizerLibrary = (
     await client.put(`/api/services/custom-jwt/worker`, {
       body: deepmerge(customizerScriptsFromDatabase, newCustomizerScripts),
     });
-  };
-  return {
-    getUserContext,
-    deployJwtCustomizerScript,
-    undeployJwtCustomizerScript,
-  };
-};
+  }
+}

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

@@ -4,6 +4,7 @@ import {
   LogtoJwtTokenKeyType,
   LogResult,
   jwtCustomizer as jwtCustomizerLog,
+  type CustomJwtFetcher,
 } from '@logto/schemas';
 import { generateStandardId } from '@logto/shared';
 import { conditional, trySafe } from '@silverhand/essentials';
@@ -11,6 +12,7 @@ import { type KoaContextWithOIDC, type UnknownObject } from 'oidc-provider';
 
 import { EnvSet } from '#src/env-set/index.js';
 import { type CloudConnectionLibrary } from '#src/libraries/cloud-connection.js';
+import { JwtCustomizerLibrary } from '#src/libraries/jwt-customizer.js';
 import { type LogtoConfigLibrary } from '#src/libraries/logto-config.js';
 import { LogEntry } from '#src/middleware/koa-audit-log.js';
 import type Libraries from '#src/tenants/Libraries.js';
@@ -66,12 +68,6 @@ export const getExtraTokenClaimsForJwtCustomization = async (
     cloudConnection: CloudConnectionLibrary;
   }
 ): Promise<UnknownObject | undefined> => {
-  const { isCloud } = EnvSet.values;
-  // No cloud connection for OSS version, skip.
-  if (!isCloud) {
-    return;
-  }
-
   // Narrow down the token type to `AccessToken` and `ClientCredentials`.
   if (
     !(token instanceof ctx.oidc.provider.AccessToken) &&
@@ -110,14 +106,6 @@ export const getExtraTokenClaimsForJwtCustomization = async (
         .map((field) => [field, Reflect.get(token, field)])
     );
 
-    const client = await cloudConnection.getClient();
-
-    const commonPayload = {
-      script,
-      environmentVariables,
-      token: readOnlyToken,
-    };
-
     // We pass context to the cloud API only when it is a user's access token.
     const logtoUserInfo = conditional(
       !isTokenClientCredentials &&
@@ -125,22 +113,29 @@ export const getExtraTokenClaimsForJwtCustomization = async (
         (await libraries.jwtCustomizers.getUserContext(token.accountId))
     );
 
-    // `context` parameter is only eligible for user's access token for now.
-    return await client.post(`/api/services/custom-jwt`, {
-      body: isTokenClientCredentials
-        ? {
-            ...commonPayload,
-            tokenType: LogtoJwtTokenKeyType.ClientCredentials,
-          }
+    const payload: CustomJwtFetcher = {
+      script,
+      environmentVariables,
+      token: readOnlyToken,
+      ...(isTokenClientCredentials
+        ? { tokenType: LogtoJwtTokenKeyType.ClientCredentials }
         : {
-            ...commonPayload,
             tokenType: LogtoJwtTokenKeyType.AccessToken,
             // TODO (LOG-8555): the newly added `UserProfile` type includes undefined fields and can not be directly assigned to `Json` type. And the `undefined` fields should be removed by zod guard.
+            // `context` parameter is only eligible for user's access token for now.
             // eslint-disable-next-line no-restricted-syntax
             context: { user: logtoUserInfo as Record<string, Json> },
-          },
+          }),
+    };
+
+    if (EnvSet.values.isCloud) {
+      const client = await cloudConnection.getClient();
+      return await client.post(`/api/services/custom-jwt`, {
+        body: payload,
         search: {},
       });
+    }
+    return await JwtCustomizerLibrary.runScriptInLocalVm(payload);
   } catch (error: unknown) {
     const entry = new LogEntry(
       `${jwtCustomizerLog.prefix}.${

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

@@ -157,10 +157,9 @@ describe('configs JWT customizer routes', () => {
     expect(response.status).toEqual(204);
   });
 
-  it('POST /configs/jwt-customizer/test should return 200', async () => {
-    const cloudConnectionResponse = { success: true };
+  it('POST /configs/jwt-customizer/test should not call cloud connection client post', async () => {
     jest.spyOn(tenantContext.cloudConnection, 'getClient').mockResolvedValue(mockCloudClient);
-    jest.spyOn(mockCloudClient, 'post').mockResolvedValue(cloudConnectionResponse);
+    const clientPostSpy = jest.spyOn(mockCloudClient, 'post');
 
     const payload: JwtCustomizerTestRequestBody = {
       tokenType: LogtoJwtTokenKeyType.ClientCredentials,
@@ -169,7 +168,7 @@ describe('configs JWT customizer routes', () => {
       token: {},
     };
 
-    const response = await routeRequester.post('/configs/jwt-customizer/test').send(payload);
+    await routeRequester.post('/configs/jwt-customizer/test').send(payload);
 
     expect(tenantContext.libraries.jwtCustomizers.deployJwtCustomizerScript).toHaveBeenCalledWith(
       expect.any(ConsoleLog),
@@ -180,13 +179,8 @@ describe('configs JWT customizer routes', () => {
       }
     );
 
-    expect(mockCloudClient.post).toHaveBeenCalledWith('/api/services/custom-jwt', {
-      body: payload,
-      search: {
-        isTest: 'true',
-      },
-    });
+    expect(clientPostSpy).toHaveBeenCalledTimes(0);
 
-    expect(response.status).toEqual(200);
+    // TODO: Add the test on nested class static method.
   });
 });

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

@@ -13,6 +13,7 @@ import { ZodError, z } from 'zod';
 
 import { EnvSet } from '#src/env-set/index.js';
 import RequestError, { formatZodError } from '#src/errors/RequestError/index.js';
+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';
@@ -41,7 +42,6 @@ export default function logtoConfigJwtCustomizerRoutes<T extends ManagementApiRo
   const { getRowsByKeys, deleteJwtCustomizer } = queries.logtoConfigs;
   const { upsertJwtCustomizer, getJwtCustomizer, getJwtCustomizers, updateJwtCustomizer } =
     logtoConfigs;
-  const { deployJwtCustomizerScript, undeployJwtCustomizerScript } = libraries.jwtCustomizers;
 
   router.put(
     '/configs/jwt-customizer/:tokenTypePath',
@@ -79,7 +79,7 @@ export default function logtoConfigJwtCustomizerRoutes<T extends ManagementApiRo
 
       // Deploy first to avoid the case where the JWT customizer was saved to DB but not deployed successfully.
       if (!isIntegrationTest) {
-        await deployJwtCustomizerScript(getConsoleLogFromContext(ctx), {
+        await libraries.jwtCustomizers.deployJwtCustomizerScript(getConsoleLogFromContext(ctx), {
           key,
           value: body,
           useCase: 'production',
@@ -123,7 +123,7 @@ export default function logtoConfigJwtCustomizerRoutes<T extends ManagementApiRo
 
       // Deploy first to avoid the case where the JWT customizer was saved to DB but not deployed successfully.
       if (!isIntegrationTest) {
-        await deployJwtCustomizerScript(getConsoleLogFromContext(ctx), {
+        await libraries.jwtCustomizers.deployJwtCustomizerScript(getConsoleLogFromContext(ctx), {
           key,
           value: body,
           useCase: 'production',
@@ -195,7 +195,10 @@ export default function logtoConfigJwtCustomizerRoutes<T extends ManagementApiRo
 
       // Undeploy the script first to avoid the case where the JWT customizer was deleted from DB but worker script not updated successfully.
       if (!isIntegrationTest) {
-        await undeployJwtCustomizerScript(getConsoleLogFromContext(ctx), tokenKey);
+        await libraries.jwtCustomizers.undeployJwtCustomizerScript(
+          getConsoleLogFromContext(ctx),
+          tokenKey
+        );
       }
 
       await deleteJwtCustomizer(tokenKey);
@@ -204,10 +207,6 @@ export default function logtoConfigJwtCustomizerRoutes<T extends ManagementApiRo
     }
   );
 
-  if (!EnvSet.values.isCloud && !EnvSet.values.isUnitTest) {
-    return;
-  }
-
   router.post(
     '/configs/jwt-customizer/test',
     koaGuard({
@@ -220,7 +219,7 @@ export default function logtoConfigJwtCustomizerRoutes<T extends ManagementApiRo
       const { body } = ctx.guard;
 
       // Deploy the test script
-      await deployJwtCustomizerScript(getConsoleLogFromContext(ctx), {
+      await libraries.jwtCustomizers.deployJwtCustomizerScript(getConsoleLogFromContext(ctx), {
         key:
           body.tokenType === LogtoJwtTokenKeyType.AccessToken
             ? LogtoJwtTokenKey.AccessToken
@@ -229,13 +228,16 @@ export default function logtoConfigJwtCustomizerRoutes<T extends ManagementApiRo
         useCase: 'test',
       });
 
-      const client = await cloudConnection.getClient();
-
       try {
+        if (EnvSet.values.isCloud) {
+          const client = await cloudConnection.getClient();
           ctx.body = await client.post(`/api/services/custom-jwt`, {
             body,
             search: { isTest: 'true' },
           });
+        } else {
+          ctx.body = await JwtCustomizerLibrary.runScriptInLocalVm(body);
+        }
       } catch (error: unknown) {
         /**
          * All APIs should throw `RequestError` instead of `Error`.

packages/core/src/routes/logto-config/logto-config.openapi.json

@@ -246,7 +246,6 @@
     },
     "/api/configs/jwt-customizer/test": {
       "post": {
-        "tags": ["Cloud only"],
         "summary": "Test JWT customizer",
         "description": "Test the JWT customizer script with the given sample context and sample token payload.",
         "requestBody": {

packages/core/src/tenants/Libraries.ts

@@ -3,7 +3,7 @@ import { type CloudConnectionLibrary } from '#src/libraries/cloud-connection.js'
 import type { ConnectorLibrary } from '#src/libraries/connector.js';
 import { createDomainLibrary } from '#src/libraries/domain.js';
 import { createHookLibrary } from '#src/libraries/hook/index.js';
-import { createJwtCustomizerLibrary } from '#src/libraries/jwt-customizer.js';
+import { JwtCustomizerLibrary } from '#src/libraries/jwt-customizer.js';
 import type { LogtoConfigLibrary } from '#src/libraries/logto-config.js';
 import { OrganizationInvitationLibrary } from '#src/libraries/organization-invitation.js';
 import { createPasscodeLibrary } from '#src/libraries/passcode.js';
@@ -26,7 +26,7 @@ export default class Libraries {
   hooks = createHookLibrary(this.queries);
   scopes = createScopeLibrary(this.queries);
   socials = createSocialLibrary(this.queries, this.connectors);
-  jwtCustomizers = createJwtCustomizerLibrary(
+  jwtCustomizers = new JwtCustomizerLibrary(
     this.queries,
     this.logtoConfigs,
     this.cloudConnection,

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

@@ -1,2 +1,3 @@
 export * from './custom-jwt.js';
 export * from './types.js';
+export * from './local-vm-error.js';

packages/core/src/utils/custom-jwt/local-vm-error.ts

@@ -0,0 +1,17 @@
+import { ResponseError } from '@withtyped/client';
+
+// Extend the ResponseError from @withtyped/client, so that we can unify the error handling and display logic for both OSS version and Cloud version.
+export class LocalVmError extends ResponseError {
+  constructor(errorBody: Record<string, unknown>, statusCode: number) {
+    super(
+      new Response(
+        new Blob([JSON.stringify(errorBody)], {
+          type: 'application/json',
+        }),
+        {
+          status: statusCode,
+        }
+      )
+    );
+  }
+}

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

@@ -1,11 +1,38 @@
+import type { AccessTokenPayload, ClientCredentialsPayload } from '@logto/schemas';
+
+const standardTokenPayloadData = {
+  jti: 'f1d3d2d1-1f2d-3d4e-5d6f-7d8a9d0e1d2',
+  aud: 'http://localhost:3000/api/test',
+  scope: 'read write',
+  clientId: 'my_app',
+};
+
+export const accessTokenSample: AccessTokenPayload = {
+  ...standardTokenPayloadData,
+  accountId: 'uid_123',
+  grantId: 'grant_123',
+  gty: 'authorization_code',
+  kind: 'AccessToken',
+};
+
+export const clientCredentialsTokenSample: ClientCredentialsPayload = {
+  ...standardTokenPayloadData,
+  kind: 'ClientCredentials',
+};
+
 export const clientCredentialsJwtCustomizerPayload = {
   script: '',
-  environmentVariables: {},
+  environmentVariables: {
+    foo: 'bar',
+    API_KEY: '12345',
+  },
   contextSample: {},
+  tokenSample: clientCredentialsTokenSample,
 };
 
 export const accessTokenJwtCustomizerPayload = {
   ...clientCredentialsJwtCustomizerPayload,
+  tokenSample: accessTokenSample,
   contextSample: {
     user: {
       id: '123',
@@ -26,3 +53,11 @@ export const accessTokenJwtCustomizerPayload = {
     },
   },
 };
+
+export const accessTokenSampleScript = `const getCustomJwtClaims = async ({ token, context, environmentVariables }) => {
+  return { user_id: context?.user?.id ?? 'unknown' };
+}`;
+
+export const clientCredentialsSampleScript = `const getCustomJwtClaims = async ({ token, context, environmentVariables }) => {
+  return { ...environmentVariables };
+}`;

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

@@ -6,6 +6,8 @@ import {
   type AccessTokenJwtCustomizer,
   type ClientCredentialsJwtCustomizer,
   type JwtCustomizerConfigs,
+  type JwtCustomizerTestRequestBody,
+  type Json,
 } from '@logto/schemas';
 
 import { authedAdminApi } from './api.js';
@@ -60,3 +62,10 @@ export const updateJwtCustomizer = async (
   authedAdminApi
     .patch(`configs/jwt-customizer/${keyTypePath}`, { json: value })
     .json<AccessTokenJwtCustomizer | ClientCredentialsJwtCustomizer>();
+
+export const testJwtCustomizer = async (payload: JwtCustomizerTestRequestBody) =>
+  authedAdminApi
+    .post(`configs/jwt-customizer/test`, {
+      json: payload,
+    })
+    .json<Json>();

packages/integration-tests/src/tests/api/application/application.roles.test.ts

@@ -1,7 +1,12 @@
 import { ApplicationType, RoleType } from '@logto/schemas';
-import { generateStandardId } from '@logto/shared';
+import { generateStandardId, formUrlEncodedHeaders } from '@logto/shared';
 import { HTTPError } from 'ky';
 
+import {
+  clientCredentialsJwtCustomizerPayload,
+  clientCredentialsSampleScript,
+} from '#src/__mocks__/jwt-customizer.js';
+import { oidcApi } from '#src/api/api.js';
 import {
   createApplication,
   getApplicationRoles,
@@ -9,9 +14,14 @@ import {
   deleteRoleFromApplication,
   putRolesToApplication,
   getApplications,
+  createResource,
+  upsertJwtCustomizer,
+  deleteJwtCustomizer,
 } from '#src/api/index.js';
 import { createRole, assignApplicationsToRole } from '#src/api/role.js';
+import { createScope } from '#src/api/scope.js';
 import { expectRejects } from '#src/helpers/index.js';
+import { getAccessTokenPayload } from '#src/utils.js';
 
 describe('admin console application management (roles)', () => {
   it('should get empty list successfully', async () => {
@@ -146,4 +156,40 @@ describe('admin console application management (roles)', () => {
     expect(applications.find(({ name }) => name === 'test-m2m-app-002')).toBeFalsy();
     expect(applications.find(({ name }) => name === 'test-spa-app-002')).toBeTruthy();
   });
+
+  it('test m2m application client credentials grant type with custom JWT', async () => {
+    await upsertJwtCustomizer('client-credentials', {
+      ...clientCredentialsJwtCustomizerPayload,
+      script: clientCredentialsSampleScript,
+    });
+
+    const m2mApp = await createApplication(generateStandardId(), ApplicationType.MachineToMachine);
+    const resource = await createResource();
+    const createdScope = await createScope(resource.id);
+    const createdScope2 = await createScope(resource.id);
+    const role = await createRole({
+      type: RoleType.MachineToMachine,
+      scopeIds: [createdScope.id, createdScope2.id],
+    });
+    await assignApplicationsToRole([m2mApp.id], role.id);
+
+    const { access_token: accessToken } = await oidcApi
+      .post('token', {
+        headers: formUrlEncodedHeaders,
+        body: new URLSearchParams({
+          client_id: m2mApp.id,
+          client_secret: m2mApp.secret,
+          grant_type: 'client_credentials',
+          resource: resource.indicator,
+          scope: [createdScope.name, createdScope2.name].join(' '),
+        }),
+      })
+      .json<{ access_token: string }>();
+
+    const payload = getAccessTokenPayload(accessToken);
+    expect(payload).toHaveProperty('foo', 'bar');
+    expect(payload).toHaveProperty('API_KEY', '12345');
+
+    await deleteJwtCustomizer('client-credentials');
+  });
 });

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

@@ -3,11 +3,14 @@ import {
   type AdminConsoleData,
   LogtoOidcConfigKeyType,
   LogtoJwtTokenKey,
+  LogtoJwtTokenKeyType,
 } from '@logto/schemas';
 
 import {
   accessTokenJwtCustomizerPayload,
   clientCredentialsJwtCustomizerPayload,
+  accessTokenSampleScript,
+  clientCredentialsSampleScript,
 } from '#src/__mocks__/jwt-customizer.js';
 import {
   deleteOidcKey,
@@ -20,6 +23,7 @@ import {
   getJwtCustomizer,
   getJwtCustomizers,
   deleteJwtCustomizer,
+  testJwtCustomizer,
 } from '#src/api/index.js';
 import { expectRejects } from '#src/helpers/index.js';
 
@@ -241,4 +245,27 @@ describe('admin console sign-in experience', () => {
     await deleteJwtCustomizer('client-credentials');
     await expect(getJwtCustomizers()).resolves.toEqual([]);
   });
+
+  it('should successfully test an access token JWT customizer', async () => {
+    const testResult = await testJwtCustomizer({
+      tokenType: LogtoJwtTokenKeyType.AccessToken,
+      token: accessTokenJwtCustomizerPayload.tokenSample,
+      context: accessTokenJwtCustomizerPayload.contextSample,
+      script: accessTokenSampleScript,
+      environmentVariables: accessTokenJwtCustomizerPayload.environmentVariables,
+    });
+    expect(testResult).toMatchObject({
+      user_id: accessTokenJwtCustomizerPayload.contextSample.user.id,
+    });
+  });
+
+  it('should successfully test a client credentials JWT customizer', async () => {
+    const testResult = await testJwtCustomizer({
+      tokenType: LogtoJwtTokenKeyType.ClientCredentials,
+      token: clientCredentialsJwtCustomizerPayload.tokenSample,
+      script: clientCredentialsSampleScript,
+      environmentVariables: clientCredentialsJwtCustomizerPayload.environmentVariables,
+    });
+    expect(testResult).toMatchObject(clientCredentialsJwtCustomizerPayload.environmentVariables);
+  });
 });

packages/integration-tests/src/tests/api/oidc/get-access-token.test.ts

@@ -5,7 +5,18 @@ import { InteractionEvent, type Resource, RoleType } from '@logto/schemas';
 import { assert } from '@silverhand/essentials';
 import { createRemoteJWKSet, jwtVerify } from 'jose';
 
-import { createResource, deleteResource, deleteUser, putInteraction } from '#src/api/index.js';
+import {
+  accessTokenJwtCustomizerPayload,
+  accessTokenSampleScript,
+} from '#src/__mocks__/jwt-customizer.js';
+import {
+  createResource,
+  deleteJwtCustomizer,
+  deleteResource,
+  deleteUser,
+  putInteraction,
+  upsertJwtCustomizer,
+} from '#src/api/index.js';
 import { assignUsersToRole, createRole, deleteRole } from '#src/api/role.js';
 import { createScope, deleteScope } from '#src/api/scope.js';
 import MockClient, { defaultConfig } from '#src/client/index.js';
@@ -91,6 +102,11 @@ describe('get access token', () => {
   });
 
   it('can sign in and getAccessToken with guest user', async () => {
+    await upsertJwtCustomizer('access-token', {
+      ...accessTokenJwtCustomizerPayload,
+      script: accessTokenSampleScript,
+    });
+
     const client = new MockClient({
       resources: [testApiResourceInfo.indicator],
       scopes: testApiScopeNames,
@@ -108,6 +124,9 @@ describe('get access token', () => {
       'scope',
       testApiScopeNames.join(' ')
     );
+    expect(getAccessTokenPayload(accessToken)).toHaveProperty('user_id', guestUserId);
+
+    await deleteJwtCustomizer('access-token');
   });
 
   it('sign in and verify jwt', async () => {

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

@@ -1,12 +1,24 @@
+import { jsonObjectGuard } from '@logto/connector-kit';
 import { z } from 'zod';
 
 import { Organizations, Roles, UserSsoIdentities } from '../../db-entries/index.js';
-import { jsonObjectGuard, mfaFactorsGuard } from '../../foundations/index.js';
+import { mfaFactorsGuard } from '../../foundations/index.js';
 import { scopeResponseGuard } from '../scope.js';
 import { userInfoGuard } from '../user.js';
 
 import { accessTokenPayloadGuard, clientCredentialsPayloadGuard } from './oidc-provider.js';
 
+export const jwtCustomizerGuard = z.object({
+  script: z.string(),
+  environmentVariables: z.record(z.string()).optional(),
+  contextSample: jsonObjectGuard.optional(),
+});
+
+export enum LogtoJwtTokenKeyType {
+  AccessToken = 'access-token',
+  ClientCredentials = 'client-credentials',
+}
+
 export const jwtCustomizerUserContextGuard = userInfoGuard.extend({
   ssoIdentities: UserSsoIdentities.guard
     .pick({ issuer: true, identityId: true, detail: true })
@@ -32,12 +44,6 @@ export const jwtCustomizerUserContextGuard = userInfoGuard.extend({
 
 export type JwtCustomizerUserContext = z.infer<typeof jwtCustomizerUserContextGuard>;
 
-export const jwtCustomizerGuard = z.object({
-  script: z.string(),
-  environmentVariables: z.record(z.string()).optional(),
-  contextSample: jsonObjectGuard.optional(),
-});
-
 export const accessTokenJwtCustomizerGuard = jwtCustomizerGuard
   .extend({
     // Use partial token guard since users customization may not rely on all fields.
@@ -57,11 +63,6 @@ export const clientCredentialsJwtCustomizerGuard = jwtCustomizerGuard
 
 export type ClientCredentialsJwtCustomizer = z.infer<typeof clientCredentialsJwtCustomizerGuard>;
 
-export enum LogtoJwtTokenKeyType {
-  AccessToken = 'access-token',
-  ClientCredentials = 'client-credentials',
-}
-
 /**
  * This guard is for the core JWT customizer testing API request body guard.
  * Unlike the DB guard

packages/toolkit/core-kit/package.json

@@ -17,7 +17,11 @@
       "import": "./lib/index.js"
     },
     "./declaration": "./declaration/index.ts",
-    "./scss/*": "./scss/*.scss"
+    "./scss/*": "./scss/*.scss",
+    "./custom-jwt": {
+      "node": "./lib/custom-jwt/index.js",
+      "types": "./lib/custom-jwt/index.d.ts"
+    }
   },
   "types": "./lib/index.d.ts",
   "files": [

packages/toolkit/core-kit/src/custom-jwt/error-handling.ts

@@ -0,0 +1,10 @@
+import { types } from 'node:util';
+
+export const buildErrorResponse = (error: unknown) =>
+  /**
+   * Use `isNativeError` to check if the error is an instance of `Error`.
+   * If the error comes from `node:vm` module, then it will not be an instance of `Error` but can be captured by `isNativeError`.
+   */
+  types.isNativeError(error)
+    ? { message: error.message, stack: error.stack }
+    : { message: String(error) };

packages/toolkit/core-kit/src/custom-jwt/index.ts

@@ -0,0 +1,2 @@
+export * from './script-execution.js';
+export * from './error-handling.js';

packages/toolkit/core-kit/src/custom-jwt/script-execution.ts

@@ -0,0 +1,41 @@
+import { runInNewContext } from 'node:vm';
+
+/**
+ * This function is used to execute a named function in a customized code script in a local
+ * virtual machine with the given payload as input.
+ *
+ * @param script Custom code snippet.
+ * @param functionName The name of the function to be executed.
+ * @param payload The input payload for the function.
+ * @returns The result of the function execution.
+ */
+export const runScriptFunctionInLocalVm = async (
+  script: string,
+  functionName: string,
+  payload: unknown
+) => {
+  const globalContext = Object.freeze({
+    fetch: async (...args: Parameters<typeof fetch>) => fetch(...args),
+  });
+  const customFunction: unknown = runInNewContext(script + `;${functionName};`, globalContext);
+
+  if (typeof customFunction !== 'function') {
+    throw new TypeError(`The script does not have a function named \`${functionName}\``);
+  }
+
+  /**
+   * We can not use top-level await in `vm`, use the following implementation instead.
+   *
+   * Ref:
+   * 1. https://github.com/nodejs/node/issues/40898
+   * 2. https://github.com/n-riesco/ijavascript/issues/173#issuecomment-693924098
+   */
+  const result: unknown = await runInNewContext(
+    '(async () => customFunction(payload))();',
+    Object.freeze({ customFunction, payload }),
+    // Limit the execution time to 3 seconds, throws error if the script takes too long to execute.
+    { timeout: 3000 }
+  );
+
+  return result;
+};