Skip to content

Commit

Permalink
feat(core,schemas): add denyAccess api context to custom jwt
Browse files Browse the repository at this point in the history
add denyAccess api context to custom jwt
  • Loading branch information
simeng-li committed Sep 2, 2024
1 parent a0807d7 commit 30d2026
Show file tree
Hide file tree
Showing 7 changed files with 175 additions and 8 deletions.
36 changes: 32 additions & 4 deletions packages/core/src/libraries/jwt-customizer.ts
Original file line number Diff line number Diff line change
@@ -1,14 +1,18 @@
import {
type CustomJwtErrorBody,
CustomJwtErrorCode,
LogtoJwtTokenKeyType,
jwtCustomizerUserContextGuard,
userInfoSelectFields,
type CustomJwtFetcher,
type JwtCustomizerType,
type JwtCustomizerUserContext,
type LogtoJwtTokenKey,
type JwtCustomizerApiContext,
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';

Expand All @@ -28,19 +32,43 @@ import {

import { type CloudConnectionLibrary } from './cloud-connection.js';

const apiContext: JwtCustomizerApiContext = Object.freeze({
denyAccess: (message?: string) => {
const error: CustomJwtErrorBody = {
code: CustomJwtErrorCode.AccessDenied,
message: message ?? 'Access denied',
};

throw new LocalVmError(error, 403);

Check warning on line 42 in packages/core/src/libraries/jwt-customizer.ts

View check run for this annotation

Codecov / codecov/patch

packages/core/src/libraries/jwt-customizer.ts#L37-L42

Added lines #L37 - L42 were not covered by tests
},
});

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

Check warning on line 56 in packages/core/src/libraries/jwt-customizer.ts

View workflow job for this annotation

GitHub Actions / ESLint Report Analysis

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

[no-warning-comments] Unexpected 'todo' comment: 'TODO: @simeng remove this when the dev...'.
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;
}

Check warning on line 70 in packages/core/src/libraries/jwt-customizer.ts

View check run for this annotation

Codecov / codecov/patch

packages/core/src/libraries/jwt-customizer.ts#L69-L70

Added lines #L69 - L70 were not covered by tests

// Assuming we only use zod for request body validation
if (error instanceof ZodError) {
const { errors } = error;
Expand Down
18 changes: 17 additions & 1 deletion packages/core/src/oidc/extra-token-claims.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';
Expand All @@ -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';

/**
Expand Down Expand Up @@ -196,11 +200,14 @@ export const getExtraTokenClaimsForJwtCustomization = async (
: jwtCustomizerLog.Type.AccessToken
}`
);

Check warning on line 203 in packages/core/src/oidc/extra-token-claims.ts

View check run for this annotation

Codecov / codecov/patch

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

Added line #L203 was not covered by tests
entry.append({
result: LogResult.Error,
error: { message: String(error) },
});

Check warning on line 208 in packages/core/src/oidc/extra-token-claims.ts

View check run for this annotation

Codecov / codecov/patch

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

Added line #L208 was not covered by tests
const { payload } = entry;

Check warning on line 210 in packages/core/src/oidc/extra-token-claims.ts

View check run for this annotation

Codecov / codecov/patch

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

Added line #L210 was not covered by tests
await queries.logs.insertLog({
id: generateStandardId(),
key: payload.key,
Expand All @@ -210,6 +217,15 @@ export const getExtraTokenClaimsForJwtCustomization = async (
token,
},
});

// If the error is an instance of `ResponseError`, we need to parse the error body to get the error code.
if (error instanceof ResponseError) {
const result = await trySafe(async () => parseCustomJwtResponseError(error));

if (result?.code === CustomJwtErrorCode.AccessDenied) {
throw new errors.AccessDenied(result.message);
}
}

Check warning on line 228 in packages/core/src/oidc/extra-token-claims.ts

View check run for this annotation

Codecov / codecov/patch

packages/core/src/oidc/extra-token-claims.ts#L220-L228

Added lines #L220 - L228 were not covered by tests
}
};
/* eslint-enable complexity */
9 changes: 7 additions & 2 deletions packages/core/src/routes/logto-config/jwt-customizer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,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';

Expand Down Expand Up @@ -249,8 +250,12 @@ 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);

throw new RequestError(
{ code: 'jwt_customizer.general', status: 422 },
{ message, code }
);
}

if (error instanceof ZodError) {
Expand Down
21 changes: 20 additions & 1 deletion packages/core/src/utils/custom-jwt/index.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,7 @@
import { LogtoJwtTokenKey, type JwtCustomizerType } from '@logto/schemas';
import { LogtoJwtTokenKey, type JwtCustomizerType, customJwtErrorBodyGuard } from '@logto/schemas';
import { type ResponseError } from '@withtyped/client';

import RequestError from '../../errors/RequestError/index.js';

import { type CustomJwtDeployRequestBody } from './types.js';

Expand All @@ -11,3 +14,19 @@ export const getJwtCustomizerScripts = (jwtCustomizers: Partial<JwtCustomizerTyp
Object.values(LogtoJwtTokenKey).map((key) => [key, { production: jwtCustomizers[key]?.script }])
) as CustomJwtDeployRequestBody;
};

export const parseCustomJwtResponseError = async (error: ResponseError) => {
const result = customJwtErrorBodyGuard.safeParse(await error.response.json());

if (!result.success) {
throw new RequestError(
{
code: 'jwt_customizer.general',
status: 500,
},
{ message: result.error.message }
);
}

Check warning on line 29 in packages/core/src/utils/custom-jwt/index.ts

View check run for this annotation

Codecov / codecov/patch

packages/core/src/utils/custom-jwt/index.ts#L22-L29

Added lines #L22 - L29 were not covered by tests

return result.data;
};
10 changes: 10 additions & 0 deletions packages/integration-tests/src/__mocks__/jwt-customizer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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'};
};`;
40 changes: 40 additions & 0 deletions packages/integration-tests/src/tests/api/logto-config.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ import {
clientCredentialsJwtCustomizerPayload,
accessTokenSampleScript,
clientCredentialsSampleScript,
accessTokenAccessDeniedSampleScript,
clientCredentialsAccessDeniedSampleScript,
} from '#src/__mocks__/jwt-customizer.js';
import {
deleteOidcKey,
Expand All @@ -26,6 +28,7 @@ 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,
Expand Down Expand Up @@ -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: 'oidc.access_denied',
status: 400,
}
);
}
);

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: 'oidc.access_denied',
status: 400,
}
);
}
);
});
49 changes: 49 additions & 0 deletions packages/schemas/src/types/logto-config/jwt-customizer.ts
Original file line number Diff line number Diff line change
Expand Up @@ -145,3 +145,52 @@ 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).optional(),
message: z.string(),
})
.catchall(z.unknown());

export type CustomJwtErrorBody = z.infer<typeof customJwtErrorBodyGuard>;

export type JwtCustomizerApiContext = {
/**
* Reject the the current token exchange request.
*
* @remarks
* By calling this function, the current token exchange request will be rejected,
* and a ODIC `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: JwtCustomizerApiContext;
};

0 comments on commit 30d2026

Please sign in to comment.