Merge pull request #4846 from logto-io/gao-organization-toke-grant-tests

test(core): add organization token grant unit tests

packages/core/src/include.d/oidc-provider/lib-keep/helpers/validate_presence.d.ts

@@ -1,3 +1,4 @@
+// https://github.com/panva/node-oidc-provider/blob/cf2069cbb31a6a855876e95157372d25dde2511c/lib/helpers/validate_presence.js
 declare module 'oidc-provider/lib/helpers/validate_presence.js' {
   export default function validatePresence(
     ctx: KoaContextWithOIDC,

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

@@ -0,0 +1,308 @@
+import { UserScope } from '@logto/core-kit';
+import { type KoaContextWithOIDC, errors, type Adapter } from 'oidc-provider';
+import Sinon from 'sinon';
+
+import { createOidcContext } from '#src/test-utils/oidc-provider.js';
+import { MockTenant } from '#src/test-utils/tenant.js';
+
+import { buildHandler } from './organization-token.js';
+
+const { jest } = import.meta;
+
+// eslint-disable-next-line @typescript-eslint/no-empty-function
+const noop = async () => {};
+
+const mockHandler = (tenant = new MockTenant()) => {
+  return buildHandler(tenant.envSet, tenant.queries.organizations);
+};
+
+const clientId = 'some_client_id';
+const grantId = 'some_grant_id';
+const accountId = 'some_account_id';
+const requestScopes = ['foo', 'bar'];
+
+const mockAdapter: Adapter = {
+  upsert: jest.fn(),
+  find: jest.fn(),
+  findByUserCode: jest.fn(),
+  findByUid: jest.fn(),
+  consume: jest.fn(),
+  destroy: jest.fn(),
+  revokeByGrantId: jest.fn(),
+};
+
+type RefreshToken = InstanceType<KoaContextWithOIDC['oidc']['provider']['RefreshToken']>;
+type Grant = InstanceType<KoaContextWithOIDC['oidc']['provider']['Grant']>;
+type Client = InstanceType<KoaContextWithOIDC['oidc']['provider']['Client']>;
+
+// @ts-expect-error
+const validClient: Client = {
+  clientId,
+  grantTypeAllowed: jest.fn().mockResolvedValue(true),
+  clientAuthMethod: 'none',
+};
+
+const validRefreshToken: RefreshToken = {
+  kind: 'RefreshToken',
+  clientId,
+  grantId,
+  accountId,
+  consumed: undefined,
+  totalLifetime: jest.fn().mockReturnValue(1),
+  isSenderConstrained: jest.fn().mockReturnValue(false),
+  consume: jest.fn(),
+  iat: 0,
+  jti: '',
+  scope: [UserScope.Organizations, ...requestScopes].join(' '),
+  scopes: new Set([UserScope.Organizations, ...requestScopes]),
+  ttlPercentagePassed: jest.fn(),
+  isValid: false,
+  isExpired: false,
+  remainingTTL: 0,
+  expiration: 0,
+  save: jest.fn(),
+  adapter: mockAdapter,
+  destroy: jest.fn(),
+  emit: jest.fn(),
+};
+
+const stubRefreshToken = (ctx: KoaContextWithOIDC, overrides?: Partial<RefreshToken>) => {
+  return Sinon.stub(ctx.oidc.provider.RefreshToken, 'find').resolves({
+    ...validRefreshToken,
+    ...overrides,
+  });
+};
+
+const validOidcContext: Partial<KoaContextWithOIDC['oidc']> = {
+  requestParamScopes: new Set(requestScopes),
+  params: {
+    refresh_token: 'some_refresh_token',
+    organization_id: 'some_org_id',
+    scope: requestScopes.join(' '),
+  },
+  entities: {
+    RefreshToken: validRefreshToken,
+    Client: validClient,
+  },
+  client: validClient,
+};
+
+const validGrant: Grant = {
+  jti: '',
+  kind: '',
+  clientId,
+  accountId,
+  adapter: mockAdapter,
+  addOIDCScope: jest.fn(),
+  rejectOIDCScope: jest.fn(),
+  getOIDCScope: jest.fn(),
+  getOIDCScopeEncountered: jest.fn(),
+  getOIDCScopeFiltered: jest.fn(),
+  addOIDCClaims: jest.fn(),
+  rejectOIDCClaims: jest.fn(),
+  getOIDCClaims: jest.fn(),
+  getOIDCClaimsEncountered: jest.fn(),
+  getOIDCClaimsFiltered: jest.fn(),
+  addResourceScope: jest.fn(),
+  rejectResourceScope: jest.fn(),
+  getResourceScope: jest.fn(),
+  getResourceScopeEncountered: jest.fn(),
+  getResourceScopeFiltered: jest.fn(),
+  save: jest.fn(),
+  destroy: jest.fn(),
+  emit: jest.fn(),
+};
+
+const stubGrant = (
+  ctx: KoaContextWithOIDC,
+  overrides?: Partial<Grant> & Record<string, unknown>
+) => {
+  return Sinon.stub(ctx.oidc.provider.Grant, 'find').resolves({
+    ...validGrant,
+    ...overrides,
+  });
+};
+
+const stubAccount = (ctx: KoaContextWithOIDC, overrideAccountId = accountId) => {
+  return Sinon.stub(ctx.oidc.provider.Account, 'findAccount').resolves({
+    accountId: overrideAccountId,
+  });
+};
+
+const createPreparedContext = () => {
+  const ctx = createOidcContext(validOidcContext);
+  stubRefreshToken(ctx);
+  stubGrant(ctx);
+  stubAccount(ctx);
+  return ctx;
+};
+
+beforeAll(() => {
+  // `oidc-provider` will warn for dev interactions
+  Sinon.stub(console, 'warn');
+});
+
+afterAll(() => {
+  Sinon.restore();
+});
+
+describe('organization token grant', () => {
+  it('should throw when required parameters are missing', async () => {
+    const ctx = createOidcContext();
+    await expect(mockHandler()(ctx, noop)).rejects.toThrow(errors.InvalidRequest);
+  });
+
+  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 refresh token is not available', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('refresh token not found')
+    );
+  });
+
+  it('should throw when refresh token mismatch client id', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    stubRefreshToken(ctx, {
+      clientId: 'some_other_id',
+    });
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('client mismatch')
+    );
+  });
+
+  it('should throw when refresh token is expired', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    stubRefreshToken(ctx, {
+      isExpired: true,
+    });
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('refresh token is expired')
+    );
+  });
+
+  it('should throw when refresh token has no organization scope', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    stubRefreshToken(ctx, {
+      scopes: new Set(),
+    });
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InsufficientScope('refresh token missing required scope', UserScope.Organizations)
+    );
+  });
+
+  it('should throw when refresh token has no grant id or the grant cannot be found', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    const findRefreshToken = stubRefreshToken(ctx, {
+      grantId: undefined,
+    });
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('grantId not found')
+    );
+
+    findRefreshToken.resolves(validRefreshToken);
+    Sinon.stub(ctx.oidc.provider.Grant, 'find').resolves();
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('grant not found')
+    );
+  });
+
+  it('should throw when grant is expired', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    stubRefreshToken(ctx);
+    stubGrant(ctx, {
+      isExpired: true,
+    });
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('grant is expired')
+    );
+  });
+
+  it("should throw when grant's client id mismatch", async () => {
+    const ctx = createOidcContext(validOidcContext);
+    stubRefreshToken(ctx);
+    stubGrant(ctx, {
+      clientId: 'some_other_id',
+    });
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('client mismatch')
+    );
+  });
+
+  it('should throw when request scopes are not available in refresh token', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    stubRefreshToken(ctx, {
+      scope: UserScope.Organizations,
+      scopes: new Set([UserScope.Organizations]),
+    });
+    stubGrant(ctx);
+    await expect(mockHandler()(ctx, noop)).rejects.toThrow(errors.InvalidScope);
+  });
+
+  it('should throw when account cannot be found or account id mismatch', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    stubRefreshToken(ctx);
+    const stubbedGrant = stubGrant(ctx);
+    const stubFindAccount = Sinon.stub(ctx.oidc.provider.Account, 'findAccount').resolves();
+    await expect(mockHandler()(ctx, noop)).rejects.toThrow(errors.InvalidGrant);
+
+    stubbedGrant.resolves({ ...validGrant, accountId: 'some_other_id' });
+    stubFindAccount.resolves({ accountId });
+    await expect(mockHandler()(ctx, noop)).rejects.toMatchError(
+      new errors.InvalidGrant('accountId mismatch')
+    );
+  });
+
+  it('should throw when refresh token has been consumed', async () => {
+    const ctx = createOidcContext(validOidcContext);
+    stubRefreshToken(ctx, {
+      consumed: new Date(),
+    });
+    stubGrant(ctx);
+    stubAccount(ctx);
+    await expect(mockHandler()(ctx, noop)).rejects.toThrow(errors.InvalidGrant);
+  });
+
+  it('should throw if the user is not a member of the organization', async () => {
+    const ctx = createPreparedContext();
+    const tenant = new MockTenant();
+    Sinon.stub(tenant.queries.organizations.relations.users, 'exists').resolves(false);
+    await expect(mockHandler(tenant)(ctx, noop)).rejects.toThrow(errors.AccessDenied);
+  });
+
+  // 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();
+    const tenant = new MockTenant();
+
+    Sinon.stub(tenant.queries.organizations.relations.users, 'exists').resolves(true);
+    Sinon.stub(tenant.queries.organizations.relations.rolesUsers, 'getUserScopes').resolves([
+      { id: 'foo', name: 'foo' },
+      { id: 'bar', name: 'bar' },
+      { id: 'baz', name: 'baz' },
+    ]);
+
+    const entityStub = Sinon.stub(ctx.oidc, 'entity');
+    const noopStub = Sinon.stub().resolves();
+
+    await expect(mockHandler(tenant)(ctx, noopStub)).resolves.toBeUndefined();
+    expect(noopStub.callCount).toBe(1);
+
+    // 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,
+      scope: requestScopes.join(' '),
+      aud: 'urn:logto:organization:some_org_id',
+    });
+  });
+});

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

@@ -17,8 +17,6 @@
  * due to the lack of development type definitions in the `IdToken` class.
  */
 
-import assert from 'node:assert';
-
 import { UserScope, buildOrganizationUrn } from '@logto/core-kit';
 import { GrantType } from '@logto/schemas';
 import { isKeyInObject } from '@silverhand/essentials';
@@ -31,6 +29,7 @@ import instance from 'oidc-provider/lib/helpers/weak_cache.js';
 
 import { type EnvSet } from '#src/env-set/index.js';
 import type OrganizationQueries from '#src/queries/organizations.js';
+import assertThat from '#src/utils/assert-that.js';
 
 import { getSharedResourceServerData, reversedResourceAccessTokenTtl } from '../resource.js';
 
@@ -73,15 +72,15 @@ export const buildHandler: (
   queries: OrganizationQueries
   // eslint-disable-next-line complexity
 ) => Parameters<Provider['registerGrantType']>['1'] = (envSet, queries) => async (ctx, next) => {
-  validatePresence(ctx, ...requiredParameters);
-
   const providerInstance = instance(ctx.oidc.provider);
   const { rotateRefreshToken } = providerInstance.configuration();
   const { client, params, requestParamScopes, provider } = ctx.oidc;
   const { RefreshToken, Account, AccessToken, Grant } = provider;
 
-  assert(client, new InvalidClient('client must be available'));
-  assert(params, new InvalidGrant('parameters must be available'));
+  assertThat(params, new InvalidGrant('parameters must be available'));
+  validatePresence(ctx, ...requiredParameters);
+
+  assertThat(client, new InvalidClient('client must be available'));
 
   // @gao: I believe the presence of the param is validated by required parameters of this grant.
   // Add `String` to make TS happy.
@@ -115,7 +114,7 @@ export const buildHandler: (
   /* === End RFC 0001 === */
 
   if (!refreshToken.grantId) {
-    throw new InvalidGrant('grant id not found');
+    throw new InvalidGrant('grantId not found');
   }
 
   const grant = await Grant.find(refreshToken.grantId, {

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

@@ -1,6 +1,8 @@
-import Provider from 'oidc-provider';
+import Provider, { type KoaContextWithOIDC } from 'oidc-provider';
 import Sinon from 'sinon';
 
+import createMockContext from './jest-koa-mocks/create-mock-context.js';
+
 const { jest } = import.meta;
 
 export abstract class GrantMock {
@@ -46,3 +48,39 @@ export const createMockProvider = (
 
   return provider;
 };
+
+/**
+ * Create an empty OIDC context with minimal required properties and a mock provider.
+ *
+ * @param override - Override the default OIDC context properties.
+ */
+export const createOidcContext = (override?: Partial<KoaContextWithOIDC['oidc']>) => {
+  const issuer = 'https://mock-issuer.com';
+  const provider = new Provider(issuer);
+  const context: KoaContextWithOIDC = {
+    ...createMockContext(),
+    oidc: {
+      route: '',
+      cookies: {
+        get: jest.fn(),
+        set: jest.fn(),
+      },
+      params: {},
+      entities: {},
+      claims: {},
+      issuer,
+      provider,
+      entity: jest.fn(),
+      promptPending: jest.fn(),
+      requestParamClaims: new Set(),
+      requestParamScopes: new Set(),
+      prompts: new Set(),
+      acr: '',
+      amr: [],
+      getAccessToken: jest.fn(),
+      clientJwtAuthExpectedAudience: jest.fn(),
+      ...override,
+    },
+  };
+  return context;
+};