From d3fbdde56f8609a0a9e61839e2cc9fe593455ef4 Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Wed, 26 Aug 2026 16:32:17 +0000 Subject: [PATCH 1/4] Add group assumption for OAuth authentication --- packages/auth/NEXT_CHANGELOG.md | 4 + packages/auth/README.md | 29 +++ packages/auth/src/credentials/m2m.ts | 8 +- packages/auth/src/oidc/tokensource.ts | 9 + .../tests/credentials/default/chain.test.ts | 53 +++- packages/auth/tests/credentials/m2m.test.ts | 201 ++++++++++++++- packages/auth/tests/oidc/tokensource.test.ts | 171 +++++++++++++ .../tests/group-assumption-client.test.ts | 232 ++++++++++++++++++ 8 files changed, 701 insertions(+), 6 deletions(-) create mode 100644 packages/examples/tests/group-assumption-client.test.ts diff --git a/packages/auth/NEXT_CHANGELOG.md b/packages/auth/NEXT_CHANGELOG.md index af1360e5f..044d70920 100644 --- a/packages/auth/NEXT_CHANGELOG.md +++ b/packages/auth/NEXT_CHANGELOG.md @@ -4,6 +4,10 @@ ### New Features and Improvements +- Added group role assumption for OAuth M2M and OIDC token exchange, including + default credential configuration through `DATABRICKS_GROUP_ID` and profile + `group_id`. + ### Bug Fixes ### Documentation diff --git a/packages/auth/README.md b/packages/auth/README.md index 6386341d2..4a5fc4b5c 100644 --- a/packages/auth/README.md +++ b/packages/auth/README.md @@ -5,3 +5,32 @@ > ## Beta > > **This SDK is in Beta and is supported for production use cases.** Interfaces might still change slightly before GA (e.g. name standardization and minor ergonomic tweaks). We are keen to hear feedback from early adopters — please [file issues](https://github.com/databricks/sdk-js/issues), and we will address them. + +## Group role assumption + +OAuth M2M credentials can request a token that assumes a group role: + +```ts +import {newM2mCredentials} from '@databricks/sdk-auth/credentials'; + +const credentials = newM2mCredentials({ + host: 'https://my-workspace.cloud.databricks.com', + clientId: process.env.DATABRICKS_CLIENT_ID ?? '', + clientSecret: process.env.DATABRICKS_CLIENT_SECRET ?? '', + groupId: 'group-123', +}); +``` + +OIDC token exchange accepts the same `groupId` option through +`newDatabricksOidcTokenProvider`. + +The default credential chain reads the group from `DATABRICKS_GROUP_ID` or +`group_id` in a Databricks configuration profile. With a group configured, it +uses OAuth M2M and skips authentication methods that cannot assume a group, +such as PAT and Databricks CLI authentication. Selecting one of those methods +explicitly with `auth_type` returns an error. + +Client options do not have a `groupId` field. A group supplied by a profile or +the environment affects only default credentials; it does not alter explicitly +provided credentials. The group is sent as `assume_group` only when obtaining +an OAuth token and is not added as a header to ordinary API requests. diff --git a/packages/auth/src/credentials/m2m.ts b/packages/auth/src/credentials/m2m.ts index f687c23e0..1904ad4ec 100644 --- a/packages/auth/src/credentials/m2m.ts +++ b/packages/auth/src/credentials/m2m.ts @@ -78,10 +78,14 @@ export function newM2mCredentials( ? options.scopes : DEFAULT_SCOPES; - const body = new URLSearchParams({ + const params = new URLSearchParams({ grant_type: 'client_credentials', scope: scopes.join(' '), - }).toString(); + }); + if (options.groupId !== undefined && options.groupId !== '') { + params.set('assume_group', options.groupId); + } + const body = params.toString(); // Client ID and secret are URL-encoded before Basic auth encoding to // avoid ambiguity with special characters in either value, matching the diff --git a/packages/auth/src/oidc/tokensource.ts b/packages/auth/src/oidc/tokensource.ts index e34b85172..9eaff0e8a 100644 --- a/packages/auth/src/oidc/tokensource.ts +++ b/packages/auth/src/oidc/tokensource.ts @@ -43,6 +43,12 @@ export interface DatabricksOidcTokenProviderConfig { */ host: string; + /** + * ID of the group whose role is assumed by the exchanged token. When + * omitted or empty, no group role is assumed. + */ + groupId?: string; + /** * TokenEndpointProvider returns the token endpoint for the Databricks OIDC * application. @@ -89,6 +95,9 @@ async function exchangeIdToken( params.set('subject_token_type', 'urn:ietf:params:oauth:token-type:jwt'); params.set('subject_token', idToken.value); params.set('grant_type', 'urn:ietf:params:oauth:grant-type:token-exchange'); + if (config.groupId !== undefined && config.groupId !== '') { + params.set('assume_group', config.groupId); + } const response = await fetch(endpoints.tokenEndpoint, { method: 'POST', diff --git a/packages/auth/tests/credentials/default/chain.test.ts b/packages/auth/tests/credentials/default/chain.test.ts index 0a69a0fbd..6ad0b23da 100644 --- a/packages/auth/tests/credentials/default/chain.test.ts +++ b/packages/auth/tests/credentials/default/chain.test.ts @@ -1,4 +1,4 @@ -import {describe, expect, it} from 'vitest'; +import {afterEach, describe, expect, it, vi} from 'vitest'; // Import Secret from the browser subpath so this test can run under both // the Node and browser runners (the default `/profiles` entry pulls in @@ -59,6 +59,11 @@ const loaderFor = describe('DefaultCredentials chain', () => { const selectedError = new Error('selected provider failed'); + + afterEach(() => { + vi.unstubAllGlobals(); + }); + const resolutionCases: { name: string; strategies: readonly Strategy[]; @@ -166,6 +171,52 @@ describe('DefaultCredentials chain', () => { expect(buildCount).toBe(1); }); + it('passes the configured group to the M2M token request', async () => { + const tokenEndpoint = `${HOST}/oidc/v1/token`; + const fetchMock = vi.fn((input, init) => { + const url = input instanceof Request ? input.url : input.toString(); + if (url === `${HOST}/.well-known/databricks-config`) { + return Promise.resolve( + new Response(JSON.stringify({oidc_endpoint: `${HOST}/oidc`}), { + status: 200, + }) + ); + } + if (url === `${HOST}/oidc/.well-known/oauth-authorization-server`) { + return Promise.resolve( + new Response(JSON.stringify({token_endpoint: tokenEndpoint}), { + status: 200, + }) + ); + } + const body = init?.body; + if (typeof body !== 'string') { + expect.fail('expected body to be a string'); + } + const params = new URLSearchParams(body); + expect(params.get('assume_group')).toBe('group-123'); + return Promise.resolve( + new Response(JSON.stringify({access_token: 'token'}), {status: 200}) + ); + }); + vi.stubGlobal('fetch', fetchMock); + const creds = new DefaultCredentials( + [patStrategy, m2mStrategy], + loaderFor({ + host: HOST, + groupId: 'group-123', + token: new Secret('ignored-pat'), + clientId: 'client-id', + clientSecret: new Secret('client-secret'), + }) + ); + + expect(await creds.authHeaders()).toEqual([ + {key: 'Authorization', value: 'Bearer token'}, + ]); + expect(fetchMock).toHaveBeenCalledTimes(3); + }); + it('invokes the profile loader exactly once', async () => { let loaderCalls = 0; const loader = (): Promise => { diff --git a/packages/auth/tests/credentials/m2m.test.ts b/packages/auth/tests/credentials/m2m.test.ts index 8253261ea..0b8c9df53 100644 --- a/packages/auth/tests/credentials/m2m.test.ts +++ b/packages/auth/tests/credentials/m2m.test.ts @@ -34,6 +34,14 @@ function textResponse(status: number, text: string): Response { return new Response(text, {status}); } +function requestParams(request: CapturedRequest): URLSearchParams { + const body = request.init?.body; + if (typeof body !== 'string') { + expect.fail('expected body to be a string'); + } + return new URLSearchParams(body); +} + const HOST = 'https://workspace.example'; const HOST_METADATA_URL = `${HOST}/.well-known/databricks-config`; const OIDC_ROOT = `${HOST}/oidc`; @@ -43,7 +51,7 @@ const TOKEN_ENDPOINT = `${OIDC_ROOT}/v1/token`; interface FetchStubs { hostMetadata?: () => Response; oauthServer?: () => Response; - token?: () => Response; + token?: (request: CapturedRequest) => Response; } function stubFetch(stubs: FetchStubs): { @@ -65,7 +73,7 @@ function stubFetch(stubs: FetchStubs): { return Promise.resolve((stubs.oauthServer ?? defaultOauthServer)()); } if (url === TOKEN_ENDPOINT && stubs.token !== undefined) { - return Promise.resolve(stubs.token()); + return Promise.resolve(stubs.token({url, init})); } return Promise.resolve(textResponse(599, `unexpected url: ${url}`)); }); @@ -88,6 +96,7 @@ describe('newM2mCredentials', () => { clientId?: string; clientSecret?: string; scopes?: string[]; + groupId?: string; tokenResponseBody: Record; want: { basicAuth: string; @@ -95,6 +104,7 @@ describe('newM2mCredentials', () => { tokenValue: string; tokenType: string | undefined; expiry: Date | undefined; + assumeGroup: string | null; }; }[] = [ { @@ -110,6 +120,7 @@ describe('newM2mCredentials', () => { tokenValue: 'test-token', tokenType: 'Bearer', expiry: new Date(NOW + 3600 * 1000), + assumeGroup: null, }, }, { @@ -121,6 +132,7 @@ describe('newM2mCredentials', () => { tokenValue: 'cde', tokenType: 'Some', expiry: undefined, + assumeGroup: null, }, }, { @@ -132,6 +144,7 @@ describe('newM2mCredentials', () => { tokenValue: 'no-type-token', tokenType: undefined, expiry: undefined, + assumeGroup: null, }, }, { @@ -145,6 +158,7 @@ describe('newM2mCredentials', () => { tokenValue: 't', tokenType: 'Bearer', expiry: undefined, + assumeGroup: null, }, }, { @@ -157,6 +171,7 @@ describe('newM2mCredentials', () => { tokenValue: 't', tokenType: 'Bearer', expiry: undefined, + assumeGroup: null, }, }, { @@ -169,6 +184,7 @@ describe('newM2mCredentials', () => { tokenValue: 't', tokenType: 'Bearer', expiry: undefined, + assumeGroup: null, }, }, { @@ -181,13 +197,47 @@ describe('newM2mCredentials', () => { tokenValue: 't', tokenType: 'Bearer', expiry: undefined, + assumeGroup: null, + }, + }, + { + name: 'group role is included in the token grant', + groupId: 'group-123', + tokenResponseBody: {token_type: 'Bearer', access_token: 't'}, + want: { + basicAuth: `Basic ${btoa('b:c')}`, + scope: 'all-apis', + tokenValue: 't', + tokenType: 'Bearer', + expiry: undefined, + assumeGroup: 'group-123', + }, + }, + { + name: 'empty group ID omits group role from the token grant', + groupId: '', + tokenResponseBody: {token_type: 'Bearer', access_token: 't'}, + want: { + basicAuth: `Basic ${btoa('b:c')}`, + scope: 'all-apis', + tokenValue: 't', + tokenType: 'Bearer', + expiry: undefined, + assumeGroup: null, }, }, ]; it.each(successCases)( '$name', - async ({clientId, clientSecret, scopes, tokenResponseBody, want}) => { + async ({ + clientId, + clientSecret, + scopes, + groupId, + tokenResponseBody, + want, + }) => { vi.setSystemTime(NOW); const {captured} = stubFetch({ token: () => jsonResponse(200, tokenResponseBody), @@ -198,6 +248,7 @@ describe('newM2mCredentials', () => { clientId: clientId ?? DEFAULT_CLIENT_ID, clientSecret: clientSecret ?? DEFAULT_CLIENT_SECRET, ...(scopes !== undefined && {scopes}), + ...(groupId !== undefined && {groupId}), }); expect(creds.name()).toBe('oauth-m2m'); const token = await creds.token(); @@ -229,6 +280,7 @@ describe('newM2mCredentials', () => { const params = new URLSearchParams(body); expect(params.get('grant_type')).toBe('client_credentials'); expect(params.get('scope')).toBe(want.scope); + expect(params.get('assume_group')).toBe(want.assumeGroup); } ); @@ -252,6 +304,80 @@ describe('newM2mCredentials', () => { ]); }); + it('retains the group role across repeated token grants', async () => { + const {captured} = stubFetch({ + token: () => jsonResponse(200, {access_token: 't'}), + }); + const creds = newM2mCredentials({ + host: HOST, + clientId: DEFAULT_CLIENT_ID, + clientSecret: DEFAULT_CLIENT_SECRET, + groupId: 'group-123', + }); + + await creds.token(); + await creds.token(); + + const tokenRequests = captured.filter(c => c.url === TOKEN_ENDPOINT); + expect(tokenRequests).toHaveLength(2); + for (const request of tokenRequests) { + expect(requestParams(request).get('assume_group')).toBe('group-123'); + } + }); + + it('keeps group roles isolated between credential providers', async () => { + const {captured} = stubFetch({ + token: request => { + const group = requestParams(request).get('assume_group') ?? 'normal'; + return jsonResponse(200, {access_token: `${group}-token`}); + }, + }); + const first = newM2mCredentials({ + host: HOST, + clientId: DEFAULT_CLIENT_ID, + clientSecret: DEFAULT_CLIENT_SECRET, + groupId: 'group-a', + }); + const second = newM2mCredentials({ + host: HOST, + clientId: DEFAULT_CLIENT_ID, + clientSecret: DEFAULT_CLIENT_SECRET, + groupId: 'group-b', + }); + + await expect(first.token()).resolves.toMatchObject({ + value: 'group-a-token', + }); + await expect(second.token()).resolves.toMatchObject({ + value: 'group-b-token', + }); + + const groups = captured + .filter(c => c.url === TOKEN_ENDPOINT) + .map(c => requestParams(c).get('assume_group')); + expect(groups).toEqual(['group-a', 'group-b']); + }); + + it('does not retry a rejected grouped token grant without the group', async () => { + const {captured} = stubFetch({token: () => textResponse(403, 'denied')}); + const creds = newM2mCredentials({ + host: HOST, + clientId: DEFAULT_CLIENT_ID, + clientSecret: DEFAULT_CLIENT_SECRET, + groupId: 'group-123', + }); + + await expect(creds.token()).rejects.toMatchObject({ + code: 'TOKEN_REQUEST_FAILED', + message: 'token request failed with status 403: denied', + }); + const tokenRequests = captured.filter(c => c.url === TOKEN_ENDPOINT); + expect(tokenRequests).toHaveLength(1); + expect(requestParams(tokenRequests[0]).get('assume_group')).toBe( + 'group-123' + ); + }); + it('retries discovery after a failed first call', async () => { let hostMetadataCalls = 0; const {captured} = stubFetch({ @@ -349,6 +475,75 @@ describe('newM2mCredentials', () => { } ); + const groupedEndpointCases: { + name: string; + host: string; + oidcRoot: string; + hostMetadata: Record; + }[] = [ + { + name: 'workspace', + host: 'https://workspace.example', + oidcRoot: 'https://workspace.example/oidc', + hostMetadata: {oidc_endpoint: 'https://workspace.example/oidc'}, + }, + { + name: 'account or unified', + host: 'https://accounts.example', + oidcRoot: 'https://accounts.example/oidc/accounts/account-id', + hostMetadata: { + oidc_endpoint: 'https://accounts.example/oidc/accounts/{account_id}', + account_id: 'account-id', + }, + }, + ]; + + it.each(groupedEndpointCases)( + 'sends group role to the $name token endpoint shape', + async ({host, oidcRoot, hostMetadata}) => { + const metadataUrl = `${host}/.well-known/databricks-config`; + const oauthServerUrl = `${oidcRoot}/.well-known/oauth-authorization-server`; + const tokenEndpoint = `${oidcRoot}/v1/token`; + const captured: CapturedRequest[] = []; + const fetchMock = vi.fn((input, init) => { + const request = {url: urlOf(input), init}; + captured.push(request); + if (request.url === metadataUrl) { + return Promise.resolve(jsonResponse(200, hostMetadata)); + } + if (request.url === oauthServerUrl) { + return Promise.resolve( + jsonResponse(200, {token_endpoint: tokenEndpoint}) + ); + } + if (request.url === tokenEndpoint) { + return Promise.resolve(jsonResponse(200, {access_token: 'token'})); + } + return Promise.resolve( + textResponse(599, `unexpected url: ${request.url}`) + ); + }); + vi.stubGlobal('fetch', fetchMock); + const credentials = newM2mCredentials({ + host, + clientId: DEFAULT_CLIENT_ID, + clientSecret: DEFAULT_CLIENT_SECRET, + groupId: 'group-123', + }); + + await expect(credentials.token()).resolves.toMatchObject({ + value: 'token', + }); + + expect(captured.map(request => request.url)).toEqual([ + metadataUrl, + oauthServerUrl, + tokenEndpoint, + ]); + expect(requestParams(captured[2]).get('assume_group')).toBe('group-123'); + } + ); + type ExpectedError = | {kind: 'm2m'; code: M2mCredentialsErrorCode; message: RegExp} | {kind: 'zod'}; diff --git a/packages/auth/tests/oidc/tokensource.test.ts b/packages/auth/tests/oidc/tokensource.test.ts index 5a27b4c21..cc1634705 100644 --- a/packages/auth/tests/oidc/tokensource.test.ts +++ b/packages/auth/tests/oidc/tokensource.test.ts @@ -51,6 +51,14 @@ function stubFetchText( return {captured, mock}; } +function requestParams(request: CapturedRequest): URLSearchParams { + const body = request.init?.body; + if (typeof body !== 'string') { + expect.fail('expected body to be a string'); + } + return new URLSearchParams(body); +} + const TOKEN_ENDPOINT = 'https://host.com/oidc/v1/token'; const ID_TOKEN = 'id-token-42'; @@ -134,8 +142,10 @@ describe('newDatabricksOidcTokenProvider', () => { clientId?: string; accountId?: string; audience?: string; + groupId?: string; wantAudience: string; wantClientIdInBody: boolean; + wantGroupId: string | null; }[] = [ { name: 'WIF workspace uses configured audience and sends client_id', @@ -143,6 +153,7 @@ describe('newDatabricksOidcTokenProvider', () => { audience: 'token-audience', wantAudience: 'token-audience', wantClientIdInBody: true, + wantGroupId: null, }, { name: 'WIF account uses configured audience and sends client_id', @@ -151,6 +162,7 @@ describe('newDatabricksOidcTokenProvider', () => { audience: 'token-audience', wantAudience: 'token-audience', wantClientIdInBody: true, + wantGroupId: null, }, { name: 'account default audience falls back to accountId', @@ -158,18 +170,39 @@ describe('newDatabricksOidcTokenProvider', () => { accountId: 'ac123', wantAudience: 'ac123', wantClientIdInBody: true, + wantGroupId: null, }, { name: 'workspace default audience falls back to the token endpoint', clientId: 'client-id', wantAudience: TOKEN_ENDPOINT, wantClientIdInBody: true, + wantGroupId: null, }, { name: 'account-wide federation omits client_id from the body', audience: 'token-audience', wantAudience: 'token-audience', wantClientIdInBody: false, + wantGroupId: null, + }, + { + name: 'WIF sends the requested group role', + clientId: 'client-id', + audience: 'token-audience', + groupId: 'group-123', + wantAudience: 'token-audience', + wantClientIdInBody: true, + wantGroupId: 'group-123', + }, + { + name: 'empty group ID omits group role from WIF', + clientId: 'client-id', + audience: 'token-audience', + groupId: '', + wantAudience: 'token-audience', + wantClientIdInBody: true, + wantGroupId: null, }, ]; @@ -179,8 +212,10 @@ describe('newDatabricksOidcTokenProvider', () => { clientId, accountId, audience, + groupId, wantAudience, wantClientIdInBody, + wantGroupId, }) => { vi.setSystemTime(NOW); const {captured} = stubFetchJson(200, { @@ -197,6 +232,7 @@ describe('newDatabricksOidcTokenProvider', () => { ...(clientId !== undefined && {clientId}), ...(accountId !== undefined && {accountId}), ...(audience !== undefined && {audience}), + ...(groupId !== undefined && {groupId}), }); const token = await ts.token(); @@ -236,6 +272,141 @@ describe('newDatabricksOidcTokenProvider', () => { expect(params.get('grant_type')).toBe( 'urn:ietf:params:oauth:grant-type:token-exchange' ); + expect(params.get('assume_group')).toBe(wantGroupId); + } + ); + + it('retains the group role across repeated token exchanges', async () => { + const {captured} = stubFetchJson(200, {access_token: 'token'}); + const {provider} = staticIdTokenProvider(ID_TOKEN); + const ts = newDatabricksOidcTokenProvider({ + host: 'http://host.com', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: provider, + audience: 'token-audience', + groupId: 'group-123', + }); + + await ts.token(); + await ts.token(); + + expect(captured).toHaveLength(2); + for (const request of captured) { + expect(requestParams(request).get('assume_group')).toBe('group-123'); + } + }); + + it('keeps group roles isolated between OIDC providers', async () => { + const captured: CapturedRequest[] = []; + const fetchMock = vi.fn((input, init) => { + const request = {url: urlOf(input), init}; + captured.push(request); + const group = requestParams(request).get('assume_group') ?? 'normal'; + return Promise.resolve( + new Response(JSON.stringify({access_token: `${group}-token`}), { + status: 200, + headers: {'Content-Type': 'application/json'}, + }) + ); + }); + vi.stubGlobal('fetch', fetchMock); + const first = newDatabricksOidcTokenProvider({ + host: 'http://host.com', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: staticIdTokenProvider(ID_TOKEN).provider, + audience: 'token-audience', + groupId: 'group-a', + }); + const second = newDatabricksOidcTokenProvider({ + host: 'http://host.com', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: staticIdTokenProvider(ID_TOKEN).provider, + audience: 'token-audience', + groupId: 'group-b', + }); + + await expect(first.token()).resolves.toMatchObject({ + value: 'group-a-token', + }); + await expect(second.token()).resolves.toMatchObject({ + value: 'group-b-token', + }); + + expect( + captured.map(request => requestParams(request).get('assume_group')) + ).toEqual(['group-a', 'group-b']); + }); + + it('does not retry a rejected grouped exchange without the group', async () => { + const {captured} = stubFetchText( + 400, + '{"error":"invalid_request","error_description":"assume_group is not supported"}' + ); + const {provider} = staticIdTokenProvider(ID_TOKEN); + const ts = newDatabricksOidcTokenProvider({ + host: 'http://host.com', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: provider, + audience: 'token-audience', + groupId: 'group-123', + }); + + let caught: unknown; + try { + await ts.token(); + } catch (error) { + caught = error; + } + expect(caught).toBeInstanceOf(Error); + expect((caught as Error).message).toContain('invalid_request'); + expect((caught as Error).message).toContain( + 'assume_group is not supported' + ); + expect(captured).toHaveLength(1); + expect(requestParams(captured[0]).get('assume_group')).toBe('group-123'); + }); + + const groupedEndpointCases: { + name: string; + tokenEndpoint: string; + clientId?: string; + }[] = [ + { + name: 'workspace WIF with client_id', + tokenEndpoint: 'https://host.com/oidc/v1/token', + clientId: 'client-id', + }, + { + name: 'account or unified WIF with client_id', + tokenEndpoint: 'https://host.com/oidc/accounts/account-id/v1/token', + clientId: 'client-id', + }, + { + name: 'account-wide federation without client_id', + tokenEndpoint: 'https://host.com/oidc/accounts/account-id/v1/token', + }, + ]; + + it.each(groupedEndpointCases)( + 'sends group role to the $name endpoint shape', + async ({tokenEndpoint, clientId}) => { + const {captured} = stubFetchJson(200, {access_token: 'token'}); + const ts = newDatabricksOidcTokenProvider({ + host: 'https://host.com', + tokenEndpointProvider: () => Promise.resolve({tokenEndpoint}), + idTokenProvider: staticIdTokenProvider(ID_TOKEN).provider, + audience: 'token-audience', + groupId: 'group-123', + ...(clientId !== undefined && {clientId}), + }); + + await ts.token(); + + expect(captured).toHaveLength(1); + expect(captured[0].url).toBe(tokenEndpoint); + const params = requestParams(captured[0]); + expect(params.get('assume_group')).toBe('group-123'); + expect(params.has('client_id')).toBe(clientId !== undefined); } ); diff --git a/packages/examples/tests/group-assumption-client.test.ts b/packages/examples/tests/group-assumption-client.test.ts new file mode 100644 index 000000000..3bc5dbccf --- /dev/null +++ b/packages/examples/tests/group-assumption-client.test.ts @@ -0,0 +1,232 @@ +import {mkdtempSync, writeFileSync} from 'node:fs'; +import {tmpdir} from 'node:os'; +import {join} from 'node:path'; + +import {afterEach, beforeEach, describe, expect, it, vi} from 'vitest'; + +import type {Credentials} from '@databricks/sdk-auth'; +import {newM2mCredentials} from '@databricks/sdk-auth/credentials'; +import type { + HttpClient, + HttpRequest, + HttpResponse, +} from '@databricks/sdk-core/http'; + +import {main} from '../src/profile-client'; + +const HOST = 'https://profile-host.cloud.databricks.com'; +const OIDC_ROOT = `${HOST}/oidc`; +const TOKEN_ENDPOINT = `${OIDC_ROOT}/v1/token`; +const ENV_KEYS = [ + 'HOME', + 'DATABRICKS_CONFIG_FILE', + 'DATABRICKS_CONFIG_PROFILE', + 'DATABRICKS_HOST', + 'DATABRICKS_TOKEN', + 'DATABRICKS_CLIENT_ID', + 'DATABRICKS_CLIENT_SECRET', + 'DATABRICKS_AUTH_TYPE', + 'DATABRICKS_GROUP_ID', +]; + +function jsonResponse(body: unknown): HttpResponse { + return { + statusCode: 200, + headers: new Headers({'content-type': 'application/json'}), + body: new Response(JSON.stringify(body)).body, + }; +} + +function urlOf(input: string | URL | Request): string { + if (typeof input === 'string') return input; + if (input instanceof URL) return input.href; + return input.url; +} + +function stubM2mFetch(): URLSearchParams[] { + const tokenForms: URLSearchParams[] = []; + const fetchMock = vi.fn((input, init) => { + const url = urlOf(input); + if (url === `${HOST}/.well-known/databricks-config`) { + return Promise.resolve( + new Response(JSON.stringify({oidc_endpoint: OIDC_ROOT}), {status: 200}) + ); + } + if (url === `${OIDC_ROOT}/.well-known/oauth-authorization-server`) { + return Promise.resolve( + new Response(JSON.stringify({token_endpoint: TOKEN_ENDPOINT}), { + status: 200, + }) + ); + } + if (url === TOKEN_ENDPOINT) { + const body = init?.body; + if (typeof body !== 'string') { + expect.fail('expected token request body to be a string'); + } + tokenForms.push(new URLSearchParams(body)); + return Promise.resolve( + new Response( + JSON.stringify({access_token: 'access-token', token_type: 'Bearer'}), + {status: 200} + ) + ); + } + return Promise.resolve( + new Response(`unexpected URL: ${url}`, {status: 599}) + ); + }); + vi.stubGlobal('fetch', fetchMock); + return tokenForms; +} + +function recordingHttpClient(): { + httpClient: HttpClient; + request: () => HttpRequest; +} { + let seen: HttpRequest | undefined; + return { + httpClient: { + send(request: HttpRequest): Promise { + seen = request; + return Promise.resolve(jsonResponse({model: 'demo-model'})); + }, + }, + request: (): HttpRequest => { + if (seen === undefined) { + throw new Error('the client did not send a request'); + } + return seen; + }, + }; +} + +function expectNoRoleHeaders(request: HttpRequest): void { + for (const header of [ + 'assume_group', + 'X-Databricks-Assume-Group', + 'X-Databricks-Role', + ]) { + expect(request.headers.has(header), `${header} should be absent`).toBe( + false + ); + } +} + +function expectGroupedClientRequest( + tokenForms: URLSearchParams[], + request: HttpRequest, + groupId: string +): void { + expect(tokenForms).toHaveLength(1); + expect([...tokenForms[0].entries()]).toEqual([ + ['grant_type', 'client_credentials'], + ['scope', 'all-apis'], + ['assume_group', groupId], + ]); + expect(request.headers.get('Authorization')).toBe('Bearer access-token'); + expectNoRoleHeaders(request); +} + +describe('generated client group assumption', () => { + const saved: Record = {}; + + beforeEach(() => { + for (const key of ENV_KEYS) { + saved[key] = process.env[key]; + Reflect.deleteProperty(process.env, key); + } + process.env.HOME = mkdtempSync(join(tmpdir(), 'group-client-home-')); + process.env.DATABRICKS_HOST = HOST; + }); + + afterEach(() => { + vi.unstubAllGlobals(); + for (const key of ENV_KEYS) { + if (saved[key] === undefined) { + Reflect.deleteProperty(process.env, key); + } else { + process.env[key] = saved[key]; + } + } + }); + + const defaultGroupCases: { + name: string; + configure: () => void; + }[] = [ + { + name: 'profile', + configure: (): void => { + const configFile = join( + mkdtempSync(join(tmpdir(), 'group-client-profile-')), + 'databrickscfg' + ); + writeFileSync( + configFile, + `[DEFAULT]\nhost = ${HOST}\nclient_id = client-id\nclient_secret = client-secret\ngroup_id = group-id\n` + ); + Reflect.deleteProperty(process.env, 'DATABRICKS_HOST'); + process.env.DATABRICKS_CONFIG_FILE = configFile; + process.env.DATABRICKS_CONFIG_PROFILE = 'DEFAULT'; + }, + }, + { + name: 'environment', + configure: (): void => { + process.env.DATABRICKS_CLIENT_ID = 'client-id'; + process.env.DATABRICKS_CLIENT_SECRET = 'client-secret'; + process.env.DATABRICKS_GROUP_ID = 'group-id'; + }, + }, + ]; + + it.each(defaultGroupCases)( + 'uses the $name group for default M2M without API role headers', + async ({configure}) => { + configure(); + const tokenForms = stubM2mFetch(); + const recorder = recordingHttpClient(); + + await main({httpClient: recorder.httpClient}); + + expectGroupedClientRequest(tokenForms, recorder.request(), 'group-id'); + } + ); + + it('uses an explicit M2M group without API role headers', async () => { + process.env.DATABRICKS_GROUP_ID = 'profile-group-must-not-win'; + const tokenForms = stubM2mFetch(); + const recorder = recordingHttpClient(); + const credentials = newM2mCredentials({ + host: HOST, + clientId: 'client-id', + clientSecret: 'client-secret', + groupId: 'explicit-group', + }); + + await main({host: HOST, httpClient: recorder.httpClient, credentials}); + + expectGroupedClientRequest( + tokenForms, + recorder.request(), + 'explicit-group' + ); + }); + + it('does not apply a profile group to explicit credentials', async () => { + process.env.DATABRICKS_GROUP_ID = 'profile-group'; + const recorder = recordingHttpClient(); + const credentials: Credentials = { + name: () => 'explicit', + authHeaders: () => + Promise.resolve([{key: 'Authorization', value: 'Bearer explicit'}]), + }; + + await main({httpClient: recorder.httpClient, credentials}); + + const request = recorder.request(); + expect(request.headers.get('Authorization')).toBe('Bearer explicit'); + expectNoRoleHeaders(request); + }); +}); From d61baaeb8936cbc9a130b9b9ab777c3b980cec2f Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Tue, 1 Sep 2026 15:18:15 +0000 Subject: [PATCH 2/4] Improve group assumption test readability --- packages/auth/tests/credentials/m2m.test.ts | 118 ++++++++----------- packages/auth/tests/oidc/tokensource.test.ts | 105 +++++++++-------- 2 files changed, 109 insertions(+), 114 deletions(-) diff --git a/packages/auth/tests/credentials/m2m.test.ts b/packages/auth/tests/credentials/m2m.test.ts index 0b8c9df53..9615c89af 100644 --- a/packages/auth/tests/credentials/m2m.test.ts +++ b/packages/auth/tests/credentials/m2m.test.ts @@ -42,6 +42,14 @@ function requestParams(request: CapturedRequest): URLSearchParams { return new URLSearchParams(body); } +function expectedM2mParams(groupId?: string): URLSearchParams { + return new URLSearchParams({ + grant_type: 'client_credentials', + scope: 'all-apis', + ...(groupId !== undefined && groupId !== '' && {assume_group: groupId}), + }); +} + const HOST = 'https://workspace.example'; const HOST_METADATA_URL = `${HOST}/.well-known/databricks-config`; const OIDC_ROOT = `${HOST}/oidc`; @@ -96,7 +104,6 @@ describe('newM2mCredentials', () => { clientId?: string; clientSecret?: string; scopes?: string[]; - groupId?: string; tokenResponseBody: Record; want: { basicAuth: string; @@ -104,7 +111,6 @@ describe('newM2mCredentials', () => { tokenValue: string; tokenType: string | undefined; expiry: Date | undefined; - assumeGroup: string | null; }; }[] = [ { @@ -120,7 +126,6 @@ describe('newM2mCredentials', () => { tokenValue: 'test-token', tokenType: 'Bearer', expiry: new Date(NOW + 3600 * 1000), - assumeGroup: null, }, }, { @@ -132,7 +137,6 @@ describe('newM2mCredentials', () => { tokenValue: 'cde', tokenType: 'Some', expiry: undefined, - assumeGroup: null, }, }, { @@ -144,7 +148,6 @@ describe('newM2mCredentials', () => { tokenValue: 'no-type-token', tokenType: undefined, expiry: undefined, - assumeGroup: null, }, }, { @@ -158,7 +161,6 @@ describe('newM2mCredentials', () => { tokenValue: 't', tokenType: 'Bearer', expiry: undefined, - assumeGroup: null, }, }, { @@ -171,7 +173,6 @@ describe('newM2mCredentials', () => { tokenValue: 't', tokenType: 'Bearer', expiry: undefined, - assumeGroup: null, }, }, { @@ -184,7 +185,6 @@ describe('newM2mCredentials', () => { tokenValue: 't', tokenType: 'Bearer', expiry: undefined, - assumeGroup: null, }, }, { @@ -197,47 +197,13 @@ describe('newM2mCredentials', () => { tokenValue: 't', tokenType: 'Bearer', expiry: undefined, - assumeGroup: null, - }, - }, - { - name: 'group role is included in the token grant', - groupId: 'group-123', - tokenResponseBody: {token_type: 'Bearer', access_token: 't'}, - want: { - basicAuth: `Basic ${btoa('b:c')}`, - scope: 'all-apis', - tokenValue: 't', - tokenType: 'Bearer', - expiry: undefined, - assumeGroup: 'group-123', - }, - }, - { - name: 'empty group ID omits group role from the token grant', - groupId: '', - tokenResponseBody: {token_type: 'Bearer', access_token: 't'}, - want: { - basicAuth: `Basic ${btoa('b:c')}`, - scope: 'all-apis', - tokenValue: 't', - tokenType: 'Bearer', - expiry: undefined, - assumeGroup: null, }, }, ]; it.each(successCases)( '$name', - async ({ - clientId, - clientSecret, - scopes, - groupId, - tokenResponseBody, - want, - }) => { + async ({clientId, clientSecret, scopes, tokenResponseBody, want}) => { vi.setSystemTime(NOW); const {captured} = stubFetch({ token: () => jsonResponse(200, tokenResponseBody), @@ -248,7 +214,6 @@ describe('newM2mCredentials', () => { clientId: clientId ?? DEFAULT_CLIENT_ID, clientSecret: clientSecret ?? DEFAULT_CLIENT_SECRET, ...(scopes !== undefined && {scopes}), - ...(groupId !== undefined && {groupId}), }); expect(creds.name()).toBe('oauth-m2m'); const token = await creds.token(); @@ -280,7 +245,49 @@ describe('newM2mCredentials', () => { const params = new URLSearchParams(body); expect(params.get('grant_type')).toBe('client_credentials'); expect(params.get('scope')).toBe(want.scope); - expect(params.get('assume_group')).toBe(want.assumeGroup); + } + ); + + const groupAssumptionCases: {name: string; groupId?: string}[] = [ + { + name: 'omits the group role when no group is configured' + }, + { + name: 'omits the group role when the group ID is empty', + groupId: '', + }, + { + name: 'requests group A', + groupId: 'group-a', + }, + { + name: 'requests group B', + groupId: 'group-b', + }, + ]; + + it.each(groupAssumptionCases)( + '$name on every token grant', + async ({groupId}) => { + const {captured} = stubFetch({ + token: () => jsonResponse(200, {access_token: 't'}), + }); + const creds = newM2mCredentials({ + host: HOST, + clientId: DEFAULT_CLIENT_ID, + clientSecret: DEFAULT_CLIENT_SECRET, + ...(groupId !== undefined && {groupId}), + }); + + await creds.token(); + await creds.token(); + + const tokenRequests = captured.filter(c => c.url === TOKEN_ENDPOINT); + expect(tokenRequests).toHaveLength(2); + const wantParams = [...expectedM2mParams(groupId).entries()]; + for (const request of tokenRequests) { + expect([...requestParams(request).entries()]).toStrictEqual(wantParams); + } } ); @@ -304,27 +311,6 @@ describe('newM2mCredentials', () => { ]); }); - it('retains the group role across repeated token grants', async () => { - const {captured} = stubFetch({ - token: () => jsonResponse(200, {access_token: 't'}), - }); - const creds = newM2mCredentials({ - host: HOST, - clientId: DEFAULT_CLIENT_ID, - clientSecret: DEFAULT_CLIENT_SECRET, - groupId: 'group-123', - }); - - await creds.token(); - await creds.token(); - - const tokenRequests = captured.filter(c => c.url === TOKEN_ENDPOINT); - expect(tokenRequests).toHaveLength(2); - for (const request of tokenRequests) { - expect(requestParams(request).get('assume_group')).toBe('group-123'); - } - }); - it('keeps group roles isolated between credential providers', async () => { const {captured} = stubFetch({ token: request => { diff --git a/packages/auth/tests/oidc/tokensource.test.ts b/packages/auth/tests/oidc/tokensource.test.ts index cc1634705..e6b1410f8 100644 --- a/packages/auth/tests/oidc/tokensource.test.ts +++ b/packages/auth/tests/oidc/tokensource.test.ts @@ -59,6 +59,20 @@ function requestParams(request: CapturedRequest): URLSearchParams { return new URLSearchParams(body); } +function expectedTokenExchangeParams( + clientId?: string, + groupId?: string +): URLSearchParams { + return new URLSearchParams({ + ...(clientId !== undefined && clientId !== '' && {client_id: clientId}), + scope: 'all-apis', + subject_token_type: 'urn:ietf:params:oauth:token-type:jwt', + subject_token: ID_TOKEN, + grant_type: 'urn:ietf:params:oauth:grant-type:token-exchange', + ...(groupId !== undefined && groupId !== '' && {assume_group: groupId}), + }); +} + const TOKEN_ENDPOINT = 'https://host.com/oidc/v1/token'; const ID_TOKEN = 'id-token-42'; @@ -142,10 +156,8 @@ describe('newDatabricksOidcTokenProvider', () => { clientId?: string; accountId?: string; audience?: string; - groupId?: string; wantAudience: string; wantClientIdInBody: boolean; - wantGroupId: string | null; }[] = [ { name: 'WIF workspace uses configured audience and sends client_id', @@ -153,7 +165,6 @@ describe('newDatabricksOidcTokenProvider', () => { audience: 'token-audience', wantAudience: 'token-audience', wantClientIdInBody: true, - wantGroupId: null, }, { name: 'WIF account uses configured audience and sends client_id', @@ -162,7 +173,6 @@ describe('newDatabricksOidcTokenProvider', () => { audience: 'token-audience', wantAudience: 'token-audience', wantClientIdInBody: true, - wantGroupId: null, }, { name: 'account default audience falls back to accountId', @@ -170,39 +180,18 @@ describe('newDatabricksOidcTokenProvider', () => { accountId: 'ac123', wantAudience: 'ac123', wantClientIdInBody: true, - wantGroupId: null, }, { name: 'workspace default audience falls back to the token endpoint', clientId: 'client-id', wantAudience: TOKEN_ENDPOINT, wantClientIdInBody: true, - wantGroupId: null, }, { name: 'account-wide federation omits client_id from the body', audience: 'token-audience', wantAudience: 'token-audience', wantClientIdInBody: false, - wantGroupId: null, - }, - { - name: 'WIF sends the requested group role', - clientId: 'client-id', - audience: 'token-audience', - groupId: 'group-123', - wantAudience: 'token-audience', - wantClientIdInBody: true, - wantGroupId: 'group-123', - }, - { - name: 'empty group ID omits group role from WIF', - clientId: 'client-id', - audience: 'token-audience', - groupId: '', - wantAudience: 'token-audience', - wantClientIdInBody: true, - wantGroupId: null, }, ]; @@ -212,10 +201,8 @@ describe('newDatabricksOidcTokenProvider', () => { clientId, accountId, audience, - groupId, wantAudience, wantClientIdInBody, - wantGroupId, }) => { vi.setSystemTime(NOW); const {captured} = stubFetchJson(200, { @@ -232,7 +219,6 @@ describe('newDatabricksOidcTokenProvider', () => { ...(clientId !== undefined && {clientId}), ...(accountId !== undefined && {accountId}), ...(audience !== undefined && {audience}), - ...(groupId !== undefined && {groupId}), }); const token = await ts.token(); @@ -272,29 +258,52 @@ describe('newDatabricksOidcTokenProvider', () => { expect(params.get('grant_type')).toBe( 'urn:ietf:params:oauth:grant-type:token-exchange' ); - expect(params.get('assume_group')).toBe(wantGroupId); } ); - it('retains the group role across repeated token exchanges', async () => { - const {captured} = stubFetchJson(200, {access_token: 'token'}); - const {provider} = staticIdTokenProvider(ID_TOKEN); - const ts = newDatabricksOidcTokenProvider({ - host: 'http://host.com', - tokenEndpointProvider: fixedEndpointProvider(), - idTokenProvider: provider, - audience: 'token-audience', - groupId: 'group-123', - }); + const groupAssumptionCases: {name: string; groupId?: string}[] = [ + { + name: 'omits the group role when no group is configured', + }, + { + name: 'omits the group role when the group ID is empty', + groupId: '', + }, + { + name: 'requests group A', + groupId: 'group-a', + }, + { + name: 'requests group B', + groupId: 'group-b', + }, + ]; + + it.each(groupAssumptionCases)( + '$name on every token exchange', + async ({groupId}) => { + const {captured} = stubFetchJson(200, {access_token: 'token'}); + const ts = newDatabricksOidcTokenProvider({ + host: 'http://host.com', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: staticIdTokenProvider(ID_TOKEN).provider, + clientId: 'client-id', + audience: 'token-audience', + ...(groupId !== undefined && {groupId}), + }); - await ts.token(); - await ts.token(); + await ts.token(); + await ts.token(); - expect(captured).toHaveLength(2); - for (const request of captured) { - expect(requestParams(request).get('assume_group')).toBe('group-123'); + expect(captured).toHaveLength(2); + const wantParams = [ + ...expectedTokenExchangeParams('client-id', groupId).entries(), + ]; + for (const request of captured) { + expect([...requestParams(request).entries()]).toStrictEqual(wantParams); + } } - }); + ); it('keeps group roles isolated between OIDC providers', async () => { const captured: CapturedRequest[] = []; @@ -404,9 +413,9 @@ describe('newDatabricksOidcTokenProvider', () => { expect(captured).toHaveLength(1); expect(captured[0].url).toBe(tokenEndpoint); - const params = requestParams(captured[0]); - expect(params.get('assume_group')).toBe('group-123'); - expect(params.has('client_id')).toBe(clientId !== undefined); + expect([...requestParams(captured[0]).entries()]).toStrictEqual([ + ...expectedTokenExchangeParams(clientId, 'group-123').entries(), + ]); } ); From 5507f36af408ab1769ec8af0ecb12bc4b18a79d1 Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Wed, 2 Sep 2026 07:49:01 +0000 Subject: [PATCH 3/4] Fix M2M test formatting --- packages/auth/tests/credentials/m2m.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/auth/tests/credentials/m2m.test.ts b/packages/auth/tests/credentials/m2m.test.ts index 9615c89af..1ba6b33db 100644 --- a/packages/auth/tests/credentials/m2m.test.ts +++ b/packages/auth/tests/credentials/m2m.test.ts @@ -250,7 +250,7 @@ describe('newM2mCredentials', () => { const groupAssumptionCases: {name: string; groupId?: string}[] = [ { - name: 'omits the group role when no group is configured' + name: 'omits the group role when no group is configured', }, { name: 'omits the group role when the group ID is empty', From 8c7315db8500ca4356e99f8c6f0c64750df305a6 Mon Sep 17 00:00:00 2001 From: Mihai Mitrea Date: Fri, 4 Sep 2026 07:47:56 +0000 Subject: [PATCH 4/4] Address group assumption review feedback --- packages/auth/README.md | 29 ---- packages/auth/tests/credentials/m2m.test.ts | 156 ++++++++++--------- packages/auth/tests/oidc/tokensource.test.ts | 140 ++++++++--------- 3 files changed, 149 insertions(+), 176 deletions(-) diff --git a/packages/auth/README.md b/packages/auth/README.md index 4a5fc4b5c..6386341d2 100644 --- a/packages/auth/README.md +++ b/packages/auth/README.md @@ -5,32 +5,3 @@ > ## Beta > > **This SDK is in Beta and is supported for production use cases.** Interfaces might still change slightly before GA (e.g. name standardization and minor ergonomic tweaks). We are keen to hear feedback from early adopters — please [file issues](https://github.com/databricks/sdk-js/issues), and we will address them. - -## Group role assumption - -OAuth M2M credentials can request a token that assumes a group role: - -```ts -import {newM2mCredentials} from '@databricks/sdk-auth/credentials'; - -const credentials = newM2mCredentials({ - host: 'https://my-workspace.cloud.databricks.com', - clientId: process.env.DATABRICKS_CLIENT_ID ?? '', - clientSecret: process.env.DATABRICKS_CLIENT_SECRET ?? '', - groupId: 'group-123', -}); -``` - -OIDC token exchange accepts the same `groupId` option through -`newDatabricksOidcTokenProvider`. - -The default credential chain reads the group from `DATABRICKS_GROUP_ID` or -`group_id` in a Databricks configuration profile. With a group configured, it -uses OAuth M2M and skips authentication methods that cannot assume a group, -such as PAT and Databricks CLI authentication. Selecting one of those methods -explicitly with `auth_type` returns an error. - -Client options do not have a `groupId` field. A group supplied by a profile or -the environment affects only default credentials; it does not alter explicitly -provided credentials. The group is sent as `assume_group` only when obtaining -an OAuth token and is not added as a header to ordinary API requests. diff --git a/packages/auth/tests/credentials/m2m.test.ts b/packages/auth/tests/credentials/m2m.test.ts index 1ba6b33db..1d3770a54 100644 --- a/packages/auth/tests/credentials/m2m.test.ts +++ b/packages/auth/tests/credentials/m2m.test.ts @@ -42,10 +42,13 @@ function requestParams(request: CapturedRequest): URLSearchParams { return new URLSearchParams(body); } -function expectedM2mParams(groupId?: string): URLSearchParams { +function expectedM2mParams( + groupId?: string, + scope = 'all-apis' +): URLSearchParams { return new URLSearchParams({ grant_type: 'client_credentials', - scope: 'all-apis', + scope, ...(groupId !== undefined && groupId !== '' && {assume_group: groupId}), }); } @@ -104,6 +107,8 @@ describe('newM2mCredentials', () => { clientId?: string; clientSecret?: string; scopes?: string[]; + groupId?: string; + tokenCalls?: number; tokenResponseBody: Record; want: { basicAuth: string; @@ -111,10 +116,12 @@ describe('newM2mCredentials', () => { tokenValue: string; tokenType: string | undefined; expiry: Date | undefined; + assumeGroup?: string; }; }[] = [ { - name: 'Bearer token with expiry and default scopes', + name: 'Bearer token with expiry, default scopes, and no group', + tokenCalls: 2, tokenResponseBody: { token_type: 'Bearer', access_token: 'test-token', @@ -128,6 +135,47 @@ describe('newM2mCredentials', () => { expiry: new Date(NOW + 3600 * 1000), }, }, + { + name: 'empty group ID omits group role assumption', + groupId: '', + tokenCalls: 2, + tokenResponseBody: {access_token: 't'}, + want: { + basicAuth: `Basic ${btoa('b:c')}`, + scope: 'all-apis', + tokenValue: 't', + tokenType: undefined, + expiry: undefined, + }, + }, + { + name: 'group A is sent on every token grant', + groupId: 'group-a', + tokenCalls: 2, + tokenResponseBody: {access_token: 't'}, + want: { + basicAuth: `Basic ${btoa('b:c')}`, + scope: 'all-apis', + tokenValue: 't', + tokenType: undefined, + expiry: undefined, + assumeGroup: 'group-a', + }, + }, + { + name: 'group B is sent on every token grant', + groupId: 'group-b', + tokenCalls: 2, + tokenResponseBody: {access_token: 't'}, + want: { + basicAuth: `Basic ${btoa('b:c')}`, + scope: 'all-apis', + tokenValue: 't', + tokenType: undefined, + expiry: undefined, + assumeGroup: 'group-b', + }, + }, { name: 'non-Bearer token_type is preserved', tokenResponseBody: {token_type: 'Some', access_token: 'cde'}, @@ -203,7 +251,15 @@ describe('newM2mCredentials', () => { it.each(successCases)( '$name', - async ({clientId, clientSecret, scopes, tokenResponseBody, want}) => { + async ({ + clientId, + clientSecret, + scopes, + groupId, + tokenCalls = 1, + tokenResponseBody, + want, + }) => { vi.setSystemTime(NOW); const {captured} = stubFetch({ token: () => jsonResponse(200, tokenResponseBody), @@ -214,79 +270,39 @@ describe('newM2mCredentials', () => { clientId: clientId ?? DEFAULT_CLIENT_ID, clientSecret: clientSecret ?? DEFAULT_CLIENT_SECRET, ...(scopes !== undefined && {scopes}), + ...(groupId !== undefined && {groupId}), }); expect(creds.name()).toBe('oauth-m2m'); - const token = await creds.token(); - - expect(token.value).toBe(want.tokenValue); - expect(token.type).toBe(want.tokenType); - expect(token.expiry).toEqual(want.expiry); + for (let call = 0; call < tokenCalls; call += 1) { + const token = await creds.token(); + expect(token.value).toBe(want.tokenValue); + expect(token.type).toBe(want.tokenType); + expect(token.expiry).toEqual(want.expiry); + } expect(captured.map(c => c.url)).toStrictEqual([ HOST_METADATA_URL, OAUTH_SERVER_URL, - TOKEN_ENDPOINT, + ...Array.from({length: tokenCalls}, () => TOKEN_ENDPOINT), ]); - const tokenRequest = captured[2]; - const init = tokenRequest.init; - if (init === undefined) { - expect.fail('expected fetch init to be provided'); - } - expect(init.method).toBe('POST'); - const headers = new Headers(init.headers); - expect(headers.get('Authorization')).toBe(want.basicAuth); - expect(headers.get('Content-Type')).toBe( - 'application/x-www-form-urlencoded' - ); - const body = init.body; - if (typeof body !== 'string') { - expect.fail('expected body to be a string'); - } - const params = new URLSearchParams(body); - expect(params.get('grant_type')).toBe('client_credentials'); - expect(params.get('scope')).toBe(want.scope); - } - ); - - const groupAssumptionCases: {name: string; groupId?: string}[] = [ - { - name: 'omits the group role when no group is configured', - }, - { - name: 'omits the group role when the group ID is empty', - groupId: '', - }, - { - name: 'requests group A', - groupId: 'group-a', - }, - { - name: 'requests group B', - groupId: 'group-b', - }, - ]; - - it.each(groupAssumptionCases)( - '$name on every token grant', - async ({groupId}) => { - const {captured} = stubFetch({ - token: () => jsonResponse(200, {access_token: 't'}), - }); - const creds = newM2mCredentials({ - host: HOST, - clientId: DEFAULT_CLIENT_ID, - clientSecret: DEFAULT_CLIENT_SECRET, - ...(groupId !== undefined && {groupId}), - }); - - await creds.token(); - await creds.token(); - - const tokenRequests = captured.filter(c => c.url === TOKEN_ENDPOINT); - expect(tokenRequests).toHaveLength(2); - const wantParams = [...expectedM2mParams(groupId).entries()]; - for (const request of tokenRequests) { - expect([...requestParams(request).entries()]).toStrictEqual(wantParams); + for (const tokenRequest of captured.slice(2)) { + const init = tokenRequest.init; + if (init === undefined) { + expect.fail('expected fetch init to be provided'); + } + expect(init.method).toBe('POST'); + const headers = new Headers(init.headers); + expect(headers.get('Authorization')).toBe(want.basicAuth); + expect(headers.get('Content-Type')).toBe( + 'application/x-www-form-urlencoded' + ); + expect([...requestParams(tokenRequest).entries()]).toStrictEqual([ + ...expectedM2mParams(groupId, want.scope).entries(), + ]); + expect(requestParams(tokenRequest).get('scope')).toBe(want.scope); + expect(requestParams(tokenRequest).get('assume_group')).toBe( + want.assumeGroup ?? null + ); } } ); diff --git a/packages/auth/tests/oidc/tokensource.test.ts b/packages/auth/tests/oidc/tokensource.test.ts index e6b1410f8..29cdfc5d4 100644 --- a/packages/auth/tests/oidc/tokensource.test.ts +++ b/packages/auth/tests/oidc/tokensource.test.ts @@ -156,13 +156,17 @@ describe('newDatabricksOidcTokenProvider', () => { clientId?: string; accountId?: string; audience?: string; + groupId?: string; + tokenCalls?: number; wantAudience: string; wantClientIdInBody: boolean; + wantAssumeGroup?: string; }[] = [ { - name: 'WIF workspace uses configured audience and sends client_id', + name: 'WIF workspace uses configured audience without a group', clientId: 'client-id', audience: 'token-audience', + tokenCalls: 2, wantAudience: 'token-audience', wantClientIdInBody: true, }, @@ -193,6 +197,35 @@ describe('newDatabricksOidcTokenProvider', () => { wantAudience: 'token-audience', wantClientIdInBody: false, }, + { + name: 'empty group ID omits group role assumption', + clientId: 'client-id', + audience: 'token-audience', + groupId: '', + tokenCalls: 2, + wantAudience: 'token-audience', + wantClientIdInBody: true, + }, + { + name: 'group A is sent on every token exchange', + clientId: 'client-id', + audience: 'token-audience', + groupId: 'group-a', + tokenCalls: 2, + wantAudience: 'token-audience', + wantClientIdInBody: true, + wantAssumeGroup: 'group-a', + }, + { + name: 'group B is sent on every token exchange', + clientId: 'client-id', + audience: 'token-audience', + groupId: 'group-b', + tokenCalls: 2, + wantAudience: 'token-audience', + wantClientIdInBody: true, + wantAssumeGroup: 'group-b', + }, ]; it.each(audienceCases)( @@ -201,8 +234,11 @@ describe('newDatabricksOidcTokenProvider', () => { clientId, accountId, audience, + groupId, + tokenCalls = 1, wantAudience, wantClientIdInBody, + wantAssumeGroup, }) => { vi.setSystemTime(NOW); const {captured} = stubFetchJson(200, { @@ -219,88 +255,38 @@ describe('newDatabricksOidcTokenProvider', () => { ...(clientId !== undefined && {clientId}), ...(accountId !== undefined && {accountId}), ...(audience !== undefined && {audience}), + ...(groupId !== undefined && {groupId}), }); - const token = await ts.token(); - expect(token.value).toBe('test-auth-token'); - expect(token.type).toBe('access-token'); - expect(token.expiry).toEqual(new Date(NOW + 3600 * 1000)); - - expect(audiences).toEqual([wantAudience]); - - expect(captured).toHaveLength(1); - const first = captured[0]; - expect(first.url).toBe(TOKEN_ENDPOINT); - const init = first.init; - if (init === undefined) { - expect.fail('expected fetch init to be provided'); + for (let call = 0; call < tokenCalls; call += 1) { + const token = await ts.token(); + expect(token.value).toBe('test-auth-token'); + expect(token.type).toBe('access-token'); + expect(token.expiry).toEqual(new Date(NOW + 3600 * 1000)); } - expect(init.method).toBe('POST'); - const headers = new Headers(init.headers); - expect(headers.get('Content-Type')).toBe( - 'application/x-www-form-urlencoded' - ); - const body = init.body; - if (typeof body !== 'string') { - expect.fail('expected body to be a string'); - } - const params = new URLSearchParams(body); - if (wantClientIdInBody) { - expect(params.get('client_id')).toBe(clientId); - } else { - expect(params.has('client_id')).toBe(false); - } - expect(params.get('scope')).toBe('all-apis'); - expect(params.get('subject_token_type')).toBe( - 'urn:ietf:params:oauth:token-type:jwt' - ); - expect(params.get('subject_token')).toBe(ID_TOKEN); - expect(params.get('grant_type')).toBe( - 'urn:ietf:params:oauth:grant-type:token-exchange' - ); - } - ); - - const groupAssumptionCases: {name: string; groupId?: string}[] = [ - { - name: 'omits the group role when no group is configured', - }, - { - name: 'omits the group role when the group ID is empty', - groupId: '', - }, - { - name: 'requests group A', - groupId: 'group-a', - }, - { - name: 'requests group B', - groupId: 'group-b', - }, - ]; - - it.each(groupAssumptionCases)( - '$name on every token exchange', - async ({groupId}) => { - const {captured} = stubFetchJson(200, {access_token: 'token'}); - const ts = newDatabricksOidcTokenProvider({ - host: 'http://host.com', - tokenEndpointProvider: fixedEndpointProvider(), - idTokenProvider: staticIdTokenProvider(ID_TOKEN).provider, - clientId: 'client-id', - audience: 'token-audience', - ...(groupId !== undefined && {groupId}), - }); - await ts.token(); - await ts.token(); + expect(audiences).toEqual( + Array.from({length: tokenCalls}, () => wantAudience) + ); - expect(captured).toHaveLength(2); - const wantParams = [ - ...expectedTokenExchangeParams('client-id', groupId).entries(), - ]; + expect(captured).toHaveLength(tokenCalls); for (const request of captured) { - expect([...requestParams(request).entries()]).toStrictEqual(wantParams); + expect(request.url).toBe(TOKEN_ENDPOINT); + const init = request.init; + if (init === undefined) { + expect.fail('expected fetch init to be provided'); + } + expect(init.method).toBe('POST'); + const headers = new Headers(init.headers); + expect(headers.get('Content-Type')).toBe( + 'application/x-www-form-urlencoded' + ); + const params = requestParams(request); + expect([...params.entries()]).toStrictEqual([ + ...expectedTokenExchangeParams(clientId, groupId).entries(), + ]); + expect(params.has('client_id')).toBe(wantClientIdInBody); + expect(params.get('assume_group')).toBe(wantAssumeGroup ?? null); } } );