feat(core): add the support on custom JWT feature

packages/core/src/libraries/cloud-connection.ts

@@ -27,17 +27,27 @@ const accessTokenResponseGuard = z.object({
 /**
  * The scope here can be empty and still work, because the cloud API requests made using this client do not rely on scope verification.
  * The `CloudScope.SendEmail` is added for now because it needs to call the cloud email service API.
+ * The `CloudScope.FetchCustomJwt` is added for now because it needs to call the cloud custom JWT service API.
  */
 const scopes: string[] = [CloudScope.SendEmail, CloudScope.FetchCustomJwt];
 const accessTokenExpirationMargin = 60;
 
 /** The library for connecting to Logto Cloud service. */
 export class CloudConnectionLibrary {
+  private _isAuthenticated = false;
   private client?: Client<typeof router>;
   private accessTokenCache?: { expiresAt: number; accessToken: string };
 
   constructor(private readonly logtoConfigs: LogtoConfigLibrary) {}
 
+  get isAuthenticated() {
+    return this._isAuthenticated;
+  }
+
+  private set isAuthenticated(value: boolean) {
+    this._isAuthenticated = value;
+  }
+
   public getCloudConnectionData = async (): Promise<CloudConnection> => {
     const { getCloudConnectionData: getCloudServiceM2mCredentials } = this.logtoConfigs;
     const credentials = await getCloudServiceM2mCredentials();
@@ -66,6 +76,8 @@ export class CloudConnectionLibrary {
       if (expiresAt > Date.now() / 1000 + accessTokenExpirationMargin) {
         return accessToken;
       }
+      // Set the cloud connection to not authenticated if the access token is expired.
+      this.isAuthenticated = false;
     }
 
     const { tokenEndpoint, appId, appSecret, resource } = await this.getCloudConnectionData();
@@ -93,6 +105,8 @@ export class CloudConnectionLibrary {
       expiresAt: Date.now() / 1000 + result.data.expires_in,
       accessToken: result.data.access_token,
     };
+    // Set the cloud connection to `authenticated` if the access token is valid.
+    this.isAuthenticated = true;
 
     return result.data.access_token;
   };

packages/core/src/oidc/init.test.ts

@@ -5,8 +5,10 @@ import initOidc from './init.js';
 
 describe('oidc provider init', () => {
   it('init should not throw', async () => {
-    const { queries, libraries } = new MockTenant();
+    const { queries, libraries, logtoConfigs, cloudConnection } = new MockTenant();
 
-    expect(() => initOidc(mockEnvSet, queries, libraries)).not.toThrow();
+    expect(() =>
+      initOidc(mockEnvSet, queries, libraries, logtoConfigs, cloudConnection)
+    ).not.toThrow();
   });
 });

packages/core/src/oidc/init.ts

@@ -1,3 +1,4 @@
+/* eslint-disable max-lines */
 /* istanbul ignore file */
 
 import assert from 'node:assert';
@@ -12,16 +13,21 @@ import {
   inSeconds,
   logtoCookieKey,
   type LogtoUiCookie,
+  LogtoJwtTokenKey,
+  type JsonObject,
 } from '@logto/schemas';
-import { conditional, tryThat } from '@silverhand/essentials';
+import { conditional, trySafe, tryThat } from '@silverhand/essentials';
+import { got } from 'got';
 import i18next from 'i18next';
 import koaBody from 'koa-body';
 import Provider, { errors } from 'oidc-provider';
 import snakecaseKeys from 'snakecase-keys';
 
-import { type EnvSet } from '#src/env-set/index.js';
+import { EnvSet } from '#src/env-set/index.js';
 import RequestError from '#src/errors/RequestError/index.js';
 import { addOidcEventListeners } from '#src/event-listeners/index.js';
+import { type CloudConnectionLibrary } from '#src/libraries/cloud-connection.js';
+import { type LogtoConfigLibrary } from '#src/libraries/logto-config.js';
 import koaAuditLog from '#src/middleware/koa-audit-log.js';
 import koaBodyEtag from '#src/middleware/koa-body-etag.js';
 import postgresAdapter from '#src/oidc/adapter.js';
@@ -45,7 +51,13 @@ import { OIDCExtraParametersKey, InteractionMode } from './type.js';
 // Temporarily removed 'EdDSA' since it's not supported by browser yet
 const supportedSigningAlgs = Object.freeze(['RS256', 'PS256', 'ES256', 'ES384', 'ES512'] as const);
 
-export default function initOidc(envSet: EnvSet, queries: Queries, libraries: Libraries): Provider {
+export default function initOidc(
+  envSet: EnvSet,
+  queries: Queries,
+  libraries: Libraries,
+  logtoConfigs: LogtoConfigLibrary,
+  cloudConnection: CloudConnectionLibrary
+): Provider {
   const {
     resources: { findDefaultResource },
     users: { findUserById },
@@ -198,6 +210,64 @@ export default function initOidc(envSet: EnvSet, queries: Queries, libraries: Li
       },
     },
     extraParams: [OIDCExtraParametersKey.InteractionMode],
+    extraTokenClaims: async (ctx, token) => {
+      if (!EnvSet.values.isDevFeaturesEnabled) {
+        return;
+      }
+
+      /**
+       * The execution on this function relies on the existence of authenticated cloud connection client.
+       *
+       * The process that cloud connection get access token also includes this function (`extraTokenClaims`
+       * is a function that will always be executed during the process of generating an access token), it
+       * could trigger infinite loop if we do not terminal the process early.
+       */
+      if (!cloudConnection.isAuthenticated) {
+        return;
+      }
+
+      const isTokenClientCredentials = token instanceof ctx.oidc.provider.ClientCredentials;
+
+      const {
+        value: { script, envVars },
+      } = (await trySafe(
+        logtoConfigs.getJwtCustomizer(
+          isTokenClientCredentials
+            ? LogtoJwtTokenKey.ClientCredentials
+            : LogtoJwtTokenKey.AccessToken
+        )
+      )) ?? { value: {} };
+
+      if (script) {
+        // Wait for cloud API to be ready and we can use cloud connection client to request the API.
+        const accessToken = await cloudConnection.getAccessToken();
+        const { endpoint: cloudApiEndpoint } = await cloudConnection.getCloudConnectionData();
+
+        // 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))
+        );
+        const result =
+          (await trySafe(
+            got
+              .post(`${cloudApiEndpoint}/services/custom-jwt`, {
+                headers: {
+                  Authorization: `Bearer ${accessToken}`,
+                },
+                json: {
+                  script,
+                  envVars,
+                  token,
+                  ...conditional(logtoUserInfo && { context: { user: logtoUserInfo } }),
+                },
+              })
+              .json<JsonObject>()
+          )) ?? {};
+        return result;
+      }
+    },
     extraClientMetadata: {
       properties: Object.values(CustomClientMetadataKey),
       validator: (_, key, value) => {
@@ -351,3 +421,4 @@ export default function initOidc(envSet: EnvSet, queries: Queries, libraries: Li
 
   return oidc;
 }
+/* eslint-enable max-lines */

packages/core/src/tenants/Tenant.ts

@@ -87,7 +87,7 @@ export default class Tenant implements TenantContext {
     app.use(koaSecurityHeaders(mountedApps, id));
 
     // Mount OIDC
-    const provider = initOidc(envSet, queries, libraries);
+    const provider = initOidc(envSet, queries, libraries, logtoConfigs, cloudConnection);
     app.use(mount('/oidc', provider.app));
 
     const tenantContext: TenantContext = {