feat(core,schemas): token exchange grant (#6057)

packages/core/src/event-listeners/grant.ts

@@ -69,6 +69,7 @@ const grantTypeToExchangeByType: Record<GrantType, token.ExchangeByType> = {
   [GrantType.AuthorizationCode]: token.ExchangeByType.AuthorizationCode,
   [GrantType.RefreshToken]: token.ExchangeByType.RefreshToken,
   [GrantType.ClientCredentials]: token.ExchangeByType.ClientCredentials,
+  [GrantType.TokenExchange]: token.ExchangeByType.TokenExchange,
 };
 
 const getExchangeByType = (grantType: unknown): token.ExchangeByType => {

packages/core/src/oidc/grants/index.ts

@@ -7,6 +7,7 @@ import type Queries from '#src/tenants/Queries.js';
 
 import * as clientCredentials from './client-credentials.js';
 import * as refreshToken from './refresh-token.js';
+import * as tokenExchange from './token-exchange.js';
 
 export const registerGrants = (oidc: Provider, envSet: EnvSet, queries: Queries) => {
   const {
@@ -33,4 +34,9 @@ export const registerGrants = (oidc: Provider, envSet: EnvSet, queries: Queries)
     clientCredentials.buildHandler(envSet, queries),
     ...getParameterConfig(clientCredentials.parameters)
   );
+  oidc.registerGrantType(
+    GrantType.TokenExchange,
+    tokenExchange.buildHandler(envSet, queries),
+    ...getParameterConfig(tokenExchange.parameters)
+  );
 };

packages/core/src/oidc/grants/token-exchange.test.ts

@@ -0,0 +1,152 @@
+import { type SubjectToken } from '@logto/schemas';
+import { type KoaContextWithOIDC, errors } from 'oidc-provider';
+import Sinon from 'sinon';
+
+import { mockApplication } from '#src/__mocks__/index.js';
+import { createOidcContext } from '#src/test-utils/oidc-provider.js';
+import { MockTenant } from '#src/test-utils/tenant.js';
+
+import { buildHandler } from './token-exchange.js';
+
+const { jest } = import.meta;
+
+// eslint-disable-next-line @typescript-eslint/no-empty-function
+const noop = async () => {};
+const findSubjectToken = jest.fn();
+const updateSubjectTokenById = jest.fn();
+
+const mockTenant = new MockTenant(undefined, {
+  subjectTokens: {
+    findSubjectToken,
+    updateSubjectTokenById,
+  },
+  applications: {
+    findApplicationById: jest.fn().mockResolvedValue(mockApplication),
+  },
+});
+const mockHandler = (tenant = mockTenant) => {
+  return buildHandler(tenant.envSet, tenant.queries);
+};
+
+const clientId = 'some_client_id';
+const subjectTokenId = 'some_token_id';
+const accountId = 'some_account_id';
+
+type Client = InstanceType<KoaContextWithOIDC['oidc']['provider']['Client']>;
+
+// @ts-expect-error
+const validClient: Client = {
+  clientId,
+  grantTypeAllowed: jest.fn().mockResolvedValue(true),
+  clientAuthMethod: 'none',
+};
+
+const validSubjectToken: SubjectToken = {
+  id: subjectTokenId,
+  userId: accountId,
+  context: {},
+  expiresAt: Date.now() + 1000,
+  consumedAt: null,
+  tenantId: 'some_tenant_id',
+  creatorId: 'some_creator_id',
+  createdAt: Date.now(),
+};
+
+const validOidcContext: Partial<KoaContextWithOIDC['oidc']> = {
+  params: {
+    subject_token: 'some_subject_token',
+    subject_token_type: 'urn:ietf:params:oauth:token-type:access_token',
+  },
+  entities: {
+    Client: validClient,
+  },
+  client: validClient,
+};
+
+const createPreparedContext = () => {
+  const ctx = createOidcContext(validOidcContext);
+  return ctx;
+};
+
+beforeAll(() => {
+  // `oidc-provider` will warn for dev interactions
+  Sinon.stub(console, 'warn');
+});
+
+afterAll(() => {
+  Sinon.restore();
+});
+
+describe('token exchange', () => {
+  it('should throw when client is not available', async () => {
+    const ctx = createOidcContext({ ...validOidcContext, client: undefined });
+    await expect(mockHandler()(ctx, noop)).rejects.toThrow(errors.InvalidClient);
+  });
+
+  it('should throw when subject token type is incorrect', async () => {
+    const ctx = createOidcContext({
+      ...validOidcContext,
+      params: { ...validOidcContext.params, subject_token_type: 'invalid' },
+    });
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('unsupported subject token type')
+    );
+  });
+
+  it('should throw when subject token is not available', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('subject token not found')
+    );
+  });
+
+  it('should throw when subject token is expired', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    findSubjectToken.mockResolvedValueOnce({ ...validSubjectToken, expiresAt: Date.now() - 1000 });
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('subject token is expired')
+    );
+  });
+
+  it('should throw when subject token has been consumed', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    findSubjectToken.mockResolvedValueOnce({ ...validSubjectToken, consumedAt: Date.now() - 1000 });
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('subject token is already consumed')
+    );
+  });
+
+  it('should throw when account cannot be found', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    findSubjectToken.mockResolvedValueOnce(validSubjectToken);
+    Sinon.stub(ctx.oidc.provider.Account, 'findAccount').resolves();
+    await expect(mockHandler()(ctx, noop)).rejects.toThrow(errors.InvalidGrant);
+  });
+
+  // The handler returns void so we cannot check the return value, and it's also not
+  // straightforward to assert the token is issued correctly. Here we just do the sanity
+  // check and basic token validation. Comprehensive token validation should be done in
+  // integration tests.
+  it('should not explode when everything looks fine', async () => {
+    const ctx = createPreparedContext();
+    findSubjectToken.mockResolvedValueOnce(validSubjectToken);
+    Sinon.stub(ctx.oidc.provider.Account, 'findAccount').resolves({ accountId });
+
+    const entityStub = Sinon.stub(ctx.oidc, 'entity');
+    const noopStub = Sinon.stub().resolves();
+
+    await expect(mockHandler(mockTenant)(ctx, noopStub)).resolves.toBeUndefined();
+    expect(noopStub.callCount).toBe(1);
+    expect(updateSubjectTokenById).toHaveBeenCalled();
+
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment
+    const [key, value] = entityStub.lastCall.args;
+    expect(key).toBe('AccessToken');
+    expect(value).toMatchObject({
+      accountId,
+      clientId,
+      grantId: subjectTokenId,
+      gty: 'urn:ietf:params:oauth:grant-type:token-exchange',
+    });
+  });
+});

packages/core/src/oidc/grants/token-exchange.ts

@@ -0,0 +1,148 @@
+/**
+ * @overview This file implements the `token_exchange` grant type. The grant type is used to impersonate
+ *
+ * @see {@link https://github.com/logto-io/rfcs | Logto RFCs} for more information about RFC 0005.
+ */
+
+import { GrantType } from '@logto/schemas';
+import { trySafe } from '@silverhand/essentials';
+import type Provider from 'oidc-provider';
+import { errors } from 'oidc-provider';
+import resolveResource from 'oidc-provider/lib/helpers/resolve_resource.js';
+import validatePresence from 'oidc-provider/lib/helpers/validate_presence.js';
+import instance from 'oidc-provider/lib/helpers/weak_cache.js';
+
+import { type EnvSet } from '#src/env-set/index.js';
+import type Queries from '#src/tenants/Queries.js';
+import assertThat from '#src/utils/assert-that.js';
+
+const { InvalidClient, InvalidGrant } = errors;
+
+/**
+ * The valid parameters for the `urn:ietf:params:oauth:grant-type:token-exchange` grant type. Note the `resource` parameter is
+ * not included here since it should be handled per configuration when registering the grant type.
+ */
+export const parameters = Object.freeze([
+  'subject_token',
+  'subject_token_type',
+  'organization_id',
+  'scope',
+] as const);
+
+/**
+ * The required parameters for the grant type.
+ *
+ * @see {@link parameters} for the full list of valid parameters.
+ */
+const requiredParameters = Object.freeze([
+  'subject_token',
+  'subject_token_type',
+] as const) satisfies ReadonlyArray<(typeof parameters)[number]>;
+
+export const buildHandler: (
+  envSet: EnvSet,
+  queries: Queries
+) => Parameters<Provider['registerGrantType']>['1'] = (envSet, queries) => async (ctx, next) => {
+  const { client, params, requestParamScopes, provider } = ctx.oidc;
+  const { Account, AccessToken } = provider;
+  const {
+    subjectTokens: { findSubjectToken, updateSubjectTokenById },
+  } = queries;
+
+  assertThat(params, new InvalidGrant('parameters must be available'));
+  assertThat(client, new InvalidClient('client must be available'));
+  assertThat(
+    params.subject_token_type === 'urn:ietf:params:oauth:token-type:access_token',
+    new InvalidGrant('unsupported subject token type')
+  );
+
+  validatePresence(ctx, ...requiredParameters);
+
+  const providerInstance = instance(provider);
+  const {
+    features: { userinfo, resourceIndicators },
+  } = providerInstance.configuration();
+
+  const subjectToken = await trySafe(async () => findSubjectToken(String(params.subject_token)));
+  assertThat(subjectToken, new InvalidGrant('subject token not found'));
+  assertThat(subjectToken.expiresAt > Date.now(), new InvalidGrant('subject token is expired'));
+  assertThat(!subjectToken.consumedAt, new InvalidGrant('subject token is already consumed'));
+
+  const account = await Account.findAccount(ctx, subjectToken.userId);
+
+  if (!account) {
+    throw new InvalidGrant('refresh token invalid (referenced account not found)');
+  }
+
+  // TODO: (LOG-9501) Implement general security checks like dPop
+  ctx.oidc.entity('Account', account);
+
+  // TODO: (LOG-9140) Check organization permissions
+
+  const accessToken = new AccessToken({
+    accountId: account.accountId,
+    clientId: client.clientId,
+    gty: GrantType.TokenExchange,
+    client,
+    grantId: subjectToken.id, // There is no actual grant, so we use the subject token ID
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    scope: undefined!,
+  });
+
+  /* eslint-disable @silverhand/fp/no-mutation */
+
+  /** The scopes requested by the client. If not provided, use the scopes from the refresh token. */
+  const scope = requestParamScopes;
+  const resource = await resolveResource(
+    ctx,
+    {
+      // We don't restrict the resource indicators to the requested resource,
+      // because the subject token does not have a resource indicator.
+      // Use the params.resource to bypass the resource indicator check.
+      resourceIndicators: new Set([params.resource]),
+    },
+    { userinfo, resourceIndicators },
+    scope
+  );
+
+  if (resource) {
+    const resourceServerInfo = await resourceIndicators.getResourceServerInfo(
+      ctx,
+      resource,
+      client
+    );
+    // @ts-expect-error -- code from oidc-provider
+    // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-call
+    accessToken.resourceServer = new provider.ResourceServer(resource, resourceServerInfo);
+    // For access token scopes, there is no "grant" to check,
+    // filter the scopes based on the resource server's scopes
+    accessToken.scope = [...scope]
+      // @ts-expect-error -- code from oidc-provider
+      .filter(Set.prototype.has.bind(accessToken.resourceServer.scopes))
+      .join(' ');
+  } else {
+    // TODO: (LOG-9166) Check claims and scopes
+    accessToken.claims = ctx.oidc.claims;
+    accessToken.scope = Array.from(scope).join(' ');
+  }
+  // TODO: (LOG-9140) Handle organization token
+
+  /* eslint-enable @silverhand/fp/no-mutation */
+
+  ctx.oidc.entity('AccessToken', accessToken);
+  const accessTokenString = await accessToken.save();
+
+  // Consume the subject token
+  await updateSubjectTokenById(subjectToken.id, {
+    consumedAt: Date.now(),
+  });
+
+  ctx.body = {
+    access_token: accessTokenString,
+    expires_in: accessToken.expiration,
+    scope: accessToken.scope,
+    token_type: accessToken.tokenType,
+  };
+
+  await next();
+};

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

@@ -20,17 +20,17 @@ import {
 describe('getConstantClientMetadata()', () => {
   expect(getConstantClientMetadata(mockEnvSet, ApplicationType.SPA)).toEqual({
     application_type: 'web',
-    grant_types: [GrantType.AuthorizationCode, GrantType.RefreshToken],
+    grant_types: [GrantType.AuthorizationCode, GrantType.RefreshToken, GrantType.TokenExchange],
     token_endpoint_auth_method: 'none',
   });
   expect(getConstantClientMetadata(mockEnvSet, ApplicationType.Native)).toEqual({
     application_type: 'native',
-    grant_types: [GrantType.AuthorizationCode, GrantType.RefreshToken],
+    grant_types: [GrantType.AuthorizationCode, GrantType.RefreshToken, GrantType.TokenExchange],
     token_endpoint_auth_method: 'none',
   });
   expect(getConstantClientMetadata(mockEnvSet, ApplicationType.Traditional)).toEqual({
     application_type: 'web',
-    grant_types: [GrantType.AuthorizationCode, GrantType.RefreshToken],
+    grant_types: [GrantType.AuthorizationCode, GrantType.RefreshToken, GrantType.TokenExchange],
     token_endpoint_auth_method: 'client_secret_basic',
   });
   expect(getConstantClientMetadata(mockEnvSet, ApplicationType.MachineToMachine)).toEqual({

packages/core/src/oidc/utils.ts

@@ -39,7 +39,7 @@ export const getConstantClientMetadata = (
     grant_types:
       type === ApplicationType.MachineToMachine
         ? [GrantType.ClientCredentials]
-        : [GrantType.AuthorizationCode, GrantType.RefreshToken],
+        : [GrantType.AuthorizationCode, GrantType.RefreshToken, GrantType.TokenExchange],
     token_endpoint_auth_method: getTokenEndpointAuthMethod(),
     response_types: conditional(type === ApplicationType.MachineToMachine && []),
     // https://www.scottbrady91.com/jose/jwts-which-signing-algorithm-should-i-use

packages/core/src/queries/subject-token.ts

@@ -1,14 +1,28 @@
-import { SubjectTokens } from '@logto/schemas';
+import { type CreateSubjectToken, SubjectTokens } from '@logto/schemas';
 import type { CommonQueryMethods } from '@silverhand/slonik';
 
+import { buildFindEntityByIdWithPool } from '#src/database/find-entity-by-id.js';
 import { buildInsertIntoWithPool } from '#src/database/insert-into.js';
+import { buildUpdateWhereWithPool } from '#src/database/update-where.js';
+import { type OmitAutoSetFields } from '#src/utils/sql.js';
 
 export const createSubjectTokenQueries = (pool: CommonQueryMethods) => {
   const insertSubjectToken = buildInsertIntoWithPool(pool)(SubjectTokens, {
     returning: true,
   });
 
+  const findSubjectToken = buildFindEntityByIdWithPool(pool)(SubjectTokens);
+
+  const updateSubjectToken = buildUpdateWhereWithPool(pool)(SubjectTokens, true);
+
+  const updateSubjectTokenById = async (
+    id: string,
+    set: Partial<OmitAutoSetFields<CreateSubjectToken>>
+  ) => updateSubjectToken({ set, where: { id }, jsonbMode: 'merge' });
+
   return {
     insertSubjectToken,
+    findSubjectToken,
+    updateSubjectTokenById,
   };
 };