Merge pull request #5461 from logto-io/yemq-log-8282-add-POST-configs-jwt-customizer-API

feat(core): add PUT /configs/jwt-customizer API

packages/core/src/__mocks__/index.ts

@@ -159,7 +159,7 @@ export const mockCookieKeys: OidcConfigKey[] = [
   { id: 'cookie', value: 'bar', createdAt: 987_654_321 },
 ];
 
-export const mockLogtoConfigs: LogtoConfig[] = [
+const mockLogtoConfigs: LogtoConfig[] = [
   {
     tenantId: 'fake_tenant',
     key: LogtoOidcConfigKey.PrivateKeys,
@@ -172,6 +172,14 @@ export const mockLogtoConfigs: LogtoConfig[] = [
   },
 ];
 
+export const mockLogtoConfigRows = {
+  rows: mockLogtoConfigs,
+  rowCount: mockLogtoConfigs.length,
+  command: 'SELECT' as const,
+  fields: [],
+  notices: [],
+};
+
 export const mockPasscode: Passcode = {
   tenantId: 'fake_tenant',
   id: 'foo',
@@ -198,3 +206,19 @@ export const mockApplicationRole: ApplicationsRole = {
   applicationId: 'application_id',
   roleId: 'role_id',
 };
+
+export const mockJwtCustomizerConfigForAccessToken = {
+  tenantId: 'fake_tenant',
+  key: 'jwt.accessToken',
+  value: {
+    script: 'console.log("hello world");',
+    envVars: {
+      API_KEY: '<api-key>',
+    },
+    contextSample: {
+      user: {
+        username: 'user',
+      },
+    },
+  },
+};

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

@@ -34,6 +34,7 @@ const logtoConfigs: LogtoConfigLibrary = {
     resource: 'resource',
   }),
   getOidcConfigs: jest.fn(),
+  upsertJwtCustomizer: jest.fn(),
 };
 
 describe('getAccessToken()', () => {

packages/core/src/libraries/logto-config.ts

@@ -1,10 +1,11 @@
-import type { LogtoOidcConfigType } from '@logto/schemas';
 import {
   cloudApiIndicator,
   cloudConnectionDataGuard,
   logtoOidcConfigGuard,
   LogtoOidcConfigKey,
+  jwtCustomizerConfigGuard,
 } from '@logto/schemas';
+import type { LogtoOidcConfigType, LogtoJwtTokenKey } from '@logto/schemas';
 import chalk from 'chalk';
 import { z, ZodError } from 'zod';
 
@@ -14,7 +15,11 @@ import { consoleLog } from '#src/utils/console.js';
 export type LogtoConfigLibrary = ReturnType<typeof createLogtoConfigLibrary>;
 
 export const createLogtoConfigLibrary = ({
-  logtoConfigs: { getRowsByKeys, getCloudConnectionData: queryCloudConnectionData },
+  logtoConfigs: {
+    getRowsByKeys,
+    getCloudConnectionData: queryCloudConnectionData,
+    upsertJwtCustomizer: queryUpsertJwtCustomizer,
+  },
 }: Pick<Queries, 'logtoConfigs'>) => {
   const getOidcConfigs = async (): Promise<LogtoOidcConfigType> => {
     try {
@@ -59,5 +64,18 @@ export const createLogtoConfigLibrary = ({
     };
   };
 
-  return { getOidcConfigs, getCloudConnectionData };
+  // Can not narrow down the type of value if we utilize `buildInsertIntoWithPool` method.
+  const upsertJwtCustomizer = async <T extends LogtoJwtTokenKey>(
+    key: T,
+    value: z.infer<(typeof jwtCustomizerConfigGuard)[T]>
+  ) => {
+    const { value: rawValue } = await queryUpsertJwtCustomizer(key, value);
+
+    return {
+      key,
+      value: jwtCustomizerConfigGuard[key].parse(rawValue),
+    };
+  };
+
+  return { getOidcConfigs, getCloudConnectionData, upsertJwtCustomizer };
 };

packages/core/src/libraries/sign-in-experience/index.test.ts

@@ -57,6 +57,7 @@ const cloudConnection = createCloudConnectionLibrary({
     resource: 'resource',
   }),
   getOidcConfigs: jest.fn(),
+  upsertJwtCustomizer: jest.fn(),
 });
 
 const getLogtoConnectors = jest.spyOn(connectorLibrary, 'getLogtoConnectors');

packages/core/src/middleware/koa-guard.ts

@@ -92,16 +92,33 @@ export const isGuardMiddleware = <Type extends IMiddleware>(
 ): function_ is WithGuardConfig<Type> =>
   function_.name === 'guardMiddleware' && has(function_, 'config');
 
+/**
+ * Previous `tryParse` function's output type was `Output | undefined`.
+ * It can not properly infer the output type to be `Output` even if the guard is provided,
+ * which brings additional but unnecessary type checks.
+ */
+export const parse = <Output, Definition extends ZodTypeDef, Input>(
+  type: 'query' | 'body' | 'params' | 'files',
+  guard: ZodType<Output, Definition, Input>,
+  data: unknown
+) => {
+  try {
+    return guard.parse(data);
+  } catch (error: unknown) {
+    throw new RequestError({ code: 'guard.invalid_input', type }, error);
+  }
+};
+
 const tryParse = <Output, Definition extends ZodTypeDef, Input>(
   type: 'query' | 'body' | 'params' | 'files',
   guard: Optional<ZodType<Output, Definition, Input>>,
   data: unknown
 ) => {
-  try {
-    return guard?.parse(data);
-  } catch (error: unknown) {
-    throw new RequestError({ code: 'guard.invalid_input', type }, error);
+  if (!guard) {
+    return;
   }
+
+  return parse(type, guard, data);
 };
 
 export default function koaGuard<

packages/core/src/queries/logto-config.ts

@@ -1,14 +1,18 @@
-import type {
-  AdminConsoleData,
-  LogtoConfig,
-  LogtoConfigKey,
-  LogtoOidcConfigKey,
-  OidcConfigKey,
+import {
+  type jwtCustomizerConfigGuard,
+  LogtoTenantConfigKey,
+  LogtoConfigs,
+  type AdminConsoleData,
+  type LogtoConfig,
+  type LogtoConfigKey,
+  type LogtoOidcConfigKey,
+  type OidcConfigKey,
+  type LogtoJwtTokenKey,
 } from '@logto/schemas';
-import { LogtoTenantConfigKey, LogtoConfigs } from '@logto/schemas';
 import { convertToIdentifiers } from '@logto/shared';
 import type { CommonQueryMethods } from 'slonik';
 import { sql } from 'slonik';
+import type { z } from 'zod';
 
 const { table, fields } = convertToIdentifiers(LogtoConfigs);
 
@@ -47,11 +51,28 @@ export const createLogtoConfigQueries = (pool: CommonQueryMethods) => {
       returning *
     `);
 
+  // Can not narrow down the type of value if we utilize `buildInsertIntoWithPool` method.
+  const upsertJwtCustomizer = async <T extends LogtoJwtTokenKey>(
+    key: T,
+    value: z.infer<(typeof jwtCustomizerConfigGuard)[T]>
+  ) =>
+    pool.one<{ key: T; value: Record<string, string> }>(
+      sql`
+        insert into ${table} (${fields.key}, ${fields.value})
+          values (${key}, ${sql.jsonb(value)})
+          on conflict (${fields.tenantId}, ${fields.key}) do update set ${
+            fields.value
+          } = ${sql.jsonb(value)}
+          returning *
+      `
+    );
+
   return {
     getAdminConsoleConfig,
     updateAdminConsoleConfig,
     getCloudConnectionData,
     getRowsByKeys,
     updateOidcConfigsByKey,
+    upsertJwtCustomizer,
   };
 };

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

@@ -104,6 +104,52 @@
           }
         }
       }
+    },
+    "/api/configs/jwt-customizer/{tokenTypePath}": {
+      "put": {
+        "summary": "Create or update JWT customizer",
+        "description": "Create or update a JWT customizer for the given token type.",
+        "parameters": [
+          {
+            "in": "path",
+            "name": "tokenTypePath",
+            "description": "The token type to create a JWT customizer for."
+          }
+        ],
+        "requestBody": {
+          "content": {
+            "application/json": {
+              "schema": {
+                "properties": {
+                  "script": {
+                    "description": "The script of the JWT customizer."
+                  },
+                  "envVars": {
+                    "description": "The environment variables for the JWT customizer."
+                  },
+                  "contextSample": {
+                    "description": "The sample context for the JWT customizer script testing purpose."
+                  },
+                  "tokenSample": {
+                    "description": "The sample raw token payload for the JWT customizer script testing purpose."
+                  }
+                }
+              }
+            }
+          }
+        },
+        "responses": {
+          "200": {
+            "description": "The updated JWT customizer."
+          },
+          "201": {
+            "description": "The created JWT customizer."
+          },
+          "400": {
+            "description": "The request body is invalid."
+          }
+        }
+      }
     }
   }
 }

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

@@ -1,4 +1,4 @@
-import { LogtoOidcConfigKey, type AdminConsoleData } from '@logto/schemas';
+import { LogtoOidcConfigKey, type AdminConsoleData, LogtoJwtTokenKey } from '@logto/schemas';
 import { generateStandardId } from '@logto/shared';
 import { createMockUtils, pickDefault } from '@logto/shared/esm';
 import Sinon from 'sinon';
@@ -6,8 +6,9 @@ import Sinon from 'sinon';
 import {
   mockAdminConsoleData,
   mockCookieKeys,
-  mockLogtoConfigs,
   mockPrivateKeys,
+  mockLogtoConfigRows,
+  mockJwtCustomizerConfigForAccessToken,
 } from '#src/__mocks__/index.js';
 import { MockTenant } from '#src/test-utils/tenant.js';
 import { createRequester } from '#src/utils/test-utils.js';
@@ -52,13 +53,8 @@ const logtoConfigQueries = {
     },
   }),
   updateOidcConfigsByKey: jest.fn(),
-  getRowsByKeys: jest.fn(async () => ({
-    rows: mockLogtoConfigs,
-    rowCount: mockLogtoConfigs.length,
-    command: 'SELECT' as const,
-    fields: [],
-    notices: [],
-  })),
+  getRowsByKeys: jest.fn(async () => mockLogtoConfigRows),
+  // UpsertJwtCustomizer: jest.fn(),
 };
 
 const logtoConfigLibraries = {
@@ -66,6 +62,7 @@ const logtoConfigLibraries = {
     [LogtoOidcConfigKey.PrivateKeys]: mockPrivateKeys,
     [LogtoOidcConfigKey.CookieKeys]: mockCookieKeys,
   })),
+  upsertJwtCustomizer: jest.fn(),
 };
 
 const settingRoutes = await pickDefault(import('./logto-config.js'));
@@ -229,4 +226,44 @@ describe('configs routes', () => {
       [newPrivateKey2, newPrivateKey]
     );
   });
+
+  it('PUT /configs/jwt-customizer/:tokenType should add a record successfully', async () => {
+    logtoConfigQueries.getRowsByKeys.mockResolvedValueOnce({
+      ...mockLogtoConfigRows,
+      rows: [],
+      rowCount: 0,
+    });
+    logtoConfigLibraries.upsertJwtCustomizer.mockResolvedValueOnce(
+      mockJwtCustomizerConfigForAccessToken
+    );
+    const response = await routeRequester
+      .put(`/configs/jwt-customizer/access-token`)
+      .send(mockJwtCustomizerConfigForAccessToken.value);
+    expect(logtoConfigLibraries.upsertJwtCustomizer).toHaveBeenCalledWith(
+      LogtoJwtTokenKey.AccessToken,
+      mockJwtCustomizerConfigForAccessToken.value
+    );
+    expect(response.status).toEqual(201);
+    expect(response.body).toEqual(mockJwtCustomizerConfigForAccessToken.value);
+  });
+
+  it('PUT /configs/jwt-customizer/:tokenType should update a record successfully', async () => {
+    logtoConfigQueries.getRowsByKeys.mockResolvedValueOnce({
+      ...mockLogtoConfigRows,
+      rows: [mockJwtCustomizerConfigForAccessToken],
+      rowCount: 1,
+    });
+    logtoConfigLibraries.upsertJwtCustomizer.mockResolvedValueOnce(
+      mockJwtCustomizerConfigForAccessToken
+    );
+    const response = await routeRequester
+      .put('/configs/jwt-customizer/access-token')
+      .send(mockJwtCustomizerConfigForAccessToken.value);
+    expect(logtoConfigLibraries.upsertJwtCustomizer).toHaveBeenCalledWith(
+      LogtoJwtTokenKey.AccessToken,
+      mockJwtCustomizerConfigForAccessToken.value
+    );
+    expect(response.status).toEqual(200);
+    expect(response.body).toEqual(mockJwtCustomizerConfigForAccessToken.value);
+  });
 });

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

@@ -12,15 +12,23 @@ import {
   type OidcConfigKeysResponse,
   type OidcConfigKey,
   LogtoOidcConfigKeyType,
+  jwtCustomizerAccessTokenGuard,
+  jwtCustomizerClientCredentialsGuard,
+  LogtoJwtTokenKey,
 } from '@logto/schemas';
 import { z } from 'zod';
 
 import RequestError from '#src/errors/RequestError/index.js';
-import koaGuard from '#src/middleware/koa-guard.js';
+import koaGuard, { parse } from '#src/middleware/koa-guard.js';
 import { exportJWK } from '#src/utils/jwks.js';
 
 import type { AuthedRouter, RouterInitArgs } from './types.js';
 
+enum LogtoJwtTokenPath {
+  AccessToken = 'access-token',
+  ClientCredentials = 'client-credentials',
+}
+
 /**
  * Provide a simple API router key type and DB config key mapping
  */
@@ -29,6 +37,19 @@ const getOidcConfigKeyDatabaseColumnName = (key: LogtoOidcConfigKeyType): LogtoO
     ? LogtoOidcConfigKey.PrivateKeys
     : LogtoOidcConfigKey.CookieKeys;
 
+const getJwtTokenKeyAndBody = (tokenPath: LogtoJwtTokenPath, body: unknown) => {
+  if (tokenPath === LogtoJwtTokenPath.AccessToken) {
+    return {
+      key: LogtoJwtTokenKey.AccessToken,
+      body: parse('body', jwtCustomizerAccessTokenGuard, body),
+    };
+  }
+  return {
+    key: LogtoJwtTokenKey.ClientCredentials,
+    body: parse('body', jwtCustomizerClientCredentialsGuard, body),
+  };
+};
+
 /**
  * Remove actual values of the private keys from response.
  * @param type Logto config key DB column name. Values are either `oidc.privateKeys` or `oidc.cookieKeys`.
@@ -60,9 +81,9 @@ const getRedactedOidcKeyResponse = async (
 export default function logtoConfigRoutes<T extends AuthedRouter>(
   ...[router, { queries, logtoConfigs, invalidateCache }]: RouterInitArgs<T>
 ) {
-  const { getAdminConsoleConfig, updateAdminConsoleConfig, updateOidcConfigsByKey } =
+  const { getAdminConsoleConfig, getRowsByKeys, updateAdminConsoleConfig, updateOidcConfigsByKey } =
     queries.logtoConfigs;
-  const { getOidcConfigs } = logtoConfigs;
+  const { getOidcConfigs, upsertJwtCustomizer } = logtoConfigs;
 
   router.get(
     '/configs/admin-console',
@@ -182,4 +203,41 @@ export default function logtoConfigRoutes<T extends AuthedRouter>(
       return next();
     }
   );
+
+  router.put(
+    '/configs/jwt-customizer/:tokenTypePath',
+    koaGuard({
+      params: z.object({
+        tokenTypePath: z.nativeEnum(LogtoJwtTokenPath),
+      }),
+      /**
+       * Use `z.unknown()` to guard the request body as a JSON object, since the actual guard depends
+       * on the `tokenTypePath` and we can not get the value of `tokenTypePath` before parsing the request body,
+       * we will do more specific guard as long as we can get the value of `tokenTypePath`.
+       *
+       * Should specify `body` in koaGuard, otherwise the request body is not accessible even via `ctx.request.body`.
+       */
+      body: z.unknown(),
+      response: jwtCustomizerAccessTokenGuard.or(jwtCustomizerClientCredentialsGuard),
+      status: [200, 201, 400],
+    }),
+    async (ctx, next) => {
+      const {
+        params: { tokenTypePath },
+        body: rawBody,
+      } = ctx.guard;
+      const { key, body } = getJwtTokenKeyAndBody(tokenTypePath, rawBody);
+
+      const { rows } = await getRowsByKeys([key]);
+
+      const jwtCustomizer = await upsertJwtCustomizer(key, body);
+
+      if (rows.length === 0) {
+        ctx.status = 201;
+      }
+      ctx.body = jwtCustomizer.value;
+
+      return next();
+    }
+  );
 }

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

@@ -3,6 +3,8 @@ import {
   type AdminConsoleData,
   type OidcConfigKeysResponse,
   type LogtoOidcConfigKeyType,
+  type JwtCustomizerAccessToken,
+  type JwtCustomizerClientCredentials,
 } from '@logto/schemas';
 
 import { authedAdminApi } from './api.js';
@@ -30,3 +32,11 @@ export const rotateOidcKeys = async (
   authedAdminApi
     .post(`configs/oidc/${keyType}/rotate`, { json: { signingKeyAlgorithm } })
     .json<OidcConfigKeysResponse[]>();
+
+export const upsertJwtCustomizer = async (
+  keyTypePath: 'access-token' | 'client-credentials',
+  value: unknown
+) =>
+  authedAdminApi
+    .put(`configs/jwt-customizer/${keyTypePath}`, { json: value })
+    .json<JwtCustomizerAccessToken | JwtCustomizerClientCredentials>();

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

@@ -10,6 +10,7 @@ import {
   getOidcKeys,
   rotateOidcKeys,
   updateAdminConsoleConfig,
+  upsertJwtCustomizer,
 } from '#src/api/index.js';
 import { expectRejects } from '#src/helpers/index.js';
 
@@ -123,4 +124,48 @@ describe('admin console sign-in experience', () => {
     ]);
     expect(privateKeys2[1]?.id).toBe(privateKeys[0]?.id);
   });
+
+  it('should successfully add a new JWT customizer', async () => {
+    const accessTokenJwtCustomizerPayload = {
+      script: '',
+      envVars: {},
+      contextSample: {
+        user: {
+          username: 'test',
+          id: 'fake-id',
+        },
+      },
+    };
+    const clientCredentialsJwtCustomizerPayload = {
+      ...accessTokenJwtCustomizerPayload,
+      contextSample: {},
+    };
+
+    const accessToken = await upsertJwtCustomizer('access-token', accessTokenJwtCustomizerPayload);
+    expect(accessToken).toMatchObject(accessTokenJwtCustomizerPayload);
+    const newAccessTokenJwtCustomizerPayload = {
+      ...accessTokenJwtCustomizerPayload,
+      script: 'new script',
+    };
+    const updatedAccessToken = await upsertJwtCustomizer(
+      'access-token',
+      newAccessTokenJwtCustomizerPayload
+    );
+    expect(updatedAccessToken).toMatchObject(newAccessTokenJwtCustomizerPayload);
+
+    const clientCredentials = await upsertJwtCustomizer(
+      'client-credentials',
+      clientCredentialsJwtCustomizerPayload
+    );
+    expect(clientCredentials).toMatchObject(clientCredentialsJwtCustomizerPayload);
+    const newClientCredentialsJwtCustomizerPayload = {
+      ...clientCredentialsJwtCustomizerPayload,
+      script: 'new script client credentials',
+    };
+    const updatedClientCredentials = await upsertJwtCustomizer(
+      'client-credentials',
+      newClientCredentialsJwtCustomizerPayload
+    );
+    expect(updatedClientCredentials).toMatchObject(newClientCredentialsJwtCustomizerPayload);
+  });
 });