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

.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,42 +23,70 @@
   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:
    * 1. The org scopes query method requires `limit` and `offset` parameters. Other management API get
    * 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 @@
     };
 
     return jwtCustomizerUserContextGuard.parse(userContext);
-  };
+  }
 
   /**
    * This method is used to deploy the give JWT customizer scripts to the cloud worker service.
@@ -95,17 +128,24 @@
    * @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 @@
     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 @@
     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,13 +4,15 @@
   LogtoJwtTokenKeyType,
   LogResult,
   jwtCustomizer as jwtCustomizerLog,
+  type CustomJwtFetcher,
 } from '@logto/schemas';
 import { generateStandardId } from '@logto/shared';
 import { conditional, trySafe } from '@silverhand/essentials';
 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 @@
     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,37 +106,36 @@
         .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 &&
         token.accountId &&
         (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> },
-          },
-      search: {},
-    });
+          }),
+    };
+
+    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 @@
     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 @@
       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 @@
       }
     );
 
-    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.
   });
 });