diff --git a/.gitignore b/.gitignore index 009cc736..23a79330 100644 --- a/.gitignore +++ b/.gitignore @@ -10,3 +10,6 @@ dist/ # Vitest browser + Playwright failure captures __screenshots__/ + +# Development Environment +.claude/scheduled_tasks.lock diff --git a/packages/auth/package.json b/packages/auth/package.json index 37f2747c..131b7352 100644 --- a/packages/auth/package.json +++ b/packages/auth/package.json @@ -17,6 +17,14 @@ "./credentials/browser": { "types": "./dist/credentials/index.browser.d.ts", "import": "./dist/credentials/index.browser.js" + }, + "./oidc": { + "types": "./dist/oidc/index.d.ts", + "import": "./dist/oidc/index.js" + }, + "./oidc/browser": { + "types": "./dist/oidc/index.browser.d.ts", + "import": "./dist/oidc/index.browser.js" } }, "files": [ diff --git a/packages/auth/src/oidc/env.ts b/packages/auth/src/oidc/env.ts new file mode 100644 index 00000000..c1e5d54b --- /dev/null +++ b/packages/auth/src/oidc/env.ts @@ -0,0 +1,31 @@ +/** + * Environment-variable OIDC ID token provider. Reads the ID token from a + * named environment variable. + * + * Node.js only. Not exported from the browser entry point. + */ + +import {env} from 'node:process'; + +import type {IDTokenProvider} from './oidc'; +import {idTokenProviderFn} from './oidc'; + +/** + * Returns an IDTokenProvider that reads the ID token from environment + * variable `name`. + * + * Note that the IDTokenProvider does not cache the token and will read the + * token from environment variable `name` each time. + * + * @param name - Name of the environment variable holding the ID token. + * @throws Error when the environment variable is unset or empty. + */ +export function newEnvIDTokenProvider(name: string): IDTokenProvider { + return idTokenProviderFn(() => { + const t = env[name]; + if (t === undefined || t === '') { + return Promise.reject(new Error(`missing env var "${name}"`)); + } + return Promise.resolve({value: t}); + }); +} diff --git a/packages/auth/src/oidc/file.ts b/packages/auth/src/oidc/file.ts new file mode 100644 index 00000000..0e669a99 --- /dev/null +++ b/packages/auth/src/oidc/file.ts @@ -0,0 +1,47 @@ +/** + * File-based OIDC ID token provider. Reads the ID token from a file on disk. + * + * Node.js only. Not exported from the browser entry point. + */ + +import {readFile} from 'node:fs/promises'; + +import type {IDTokenProvider} from './oidc'; +import {idTokenProviderFn} from './oidc'; + +/** + * Returns an IDTokenProvider that reads the ID token from a file. The file + * should contain a single line with the token. + * + * @param path - Filesystem path to the file containing the ID token. + * @throws Error when the path is empty, the file does not exist, or the file + * is empty. + */ +export function newFileTokenProvider(path: string): IDTokenProvider { + return idTokenProviderFn(async () => { + if (path === '') { + throw new Error('missing path'); + } + let content: string; + try { + content = await readFile(path, 'utf-8'); + } catch (e: unknown) { + if (isNodeErrorWithCode(e, 'ENOENT')) { + throw new Error(`file "${path}" does not exist`); + } + throw e; + } + if (content.length === 0) { + throw new Error(`file "${path}" is empty`); + } + return {value: content}; + }); +} + +function isNodeErrorWithCode(e: unknown, code: string): boolean { + if (!(e instanceof Error) || !('code' in e)) { + return false; + } + const errCode: unknown = e.code; + return typeof errCode === 'string' && errCode === code; +} diff --git a/packages/auth/src/oidc/index.browser.ts b/packages/auth/src/oidc/index.browser.ts new file mode 100644 index 00000000..a9d2a263 --- /dev/null +++ b/packages/auth/src/oidc/index.browser.ts @@ -0,0 +1,12 @@ +/** + * Browser entry point for OIDC providers. Excludes ID token providers that + * depend on Node.js-only APIs (`process.env`, the filesystem). + */ + +export type {IDToken, IDTokenProvider} from './oidc'; +export {idTokenProviderFn} from './oidc'; +export type { + DatabricksOIDCTokenProviderConfig, + OAuthAuthorizationServer, +} from './tokensource'; +export {newDatabricksOIDCTokenProvider} from './tokensource'; diff --git a/packages/auth/src/oidc/index.ts b/packages/auth/src/oidc/index.ts new file mode 100644 index 00000000..568cbeba --- /dev/null +++ b/packages/auth/src/oidc/index.ts @@ -0,0 +1,13 @@ +/** + * OIDC ID token providers and Databricks OIDC token-exchange provider. + */ + +export type {IDToken, IDTokenProvider} from './oidc'; +export {idTokenProviderFn} from './oidc'; +export {newEnvIDTokenProvider} from './env'; +export {newFileTokenProvider} from './file'; +export type { + DatabricksOIDCTokenProviderConfig, + OAuthAuthorizationServer, +} from './tokensource'; +export {newDatabricksOIDCTokenProvider} from './tokensource'; diff --git a/packages/auth/src/oidc/oidc.ts b/packages/auth/src/oidc/oidc.ts new file mode 100644 index 00000000..aea75f3c --- /dev/null +++ b/packages/auth/src/oidc/oidc.ts @@ -0,0 +1,32 @@ +/** + * Package oidc provides utilities for working with OIDC ID tokens. + * + * This package is experimental and subject to change. + */ + +/** + * IDToken represents an OIDC ID token that can be exchanged for a Databricks + * access token. + */ +export interface IDToken { + value: string; +} + +/** + * IDTokenProvider is anything that returns an IDToken given an audience. + */ +export interface IDTokenProvider { + idToken(audience: string): Promise; +} + +/** + * Adapter to allow the use of ordinary functions as IDTokenProvider. + * + * @example + * const provider = idTokenProviderFn(async () => ({ value: 'my-id-token' })); + */ +export function idTokenProviderFn( + fn: (audience: string) => Promise +): IDTokenProvider { + return {idToken: fn}; +} diff --git a/packages/auth/src/oidc/tokensource.ts b/packages/auth/src/oidc/tokensource.ts new file mode 100644 index 00000000..8dd3def8 --- /dev/null +++ b/packages/auth/src/oidc/tokensource.ts @@ -0,0 +1,133 @@ +/** + * Databricks OIDC token-exchange provider. Exchanges an OIDC ID token for a + * Databricks access token using the OAuth 2.0 token-exchange grant. + */ + +import {z} from 'zod'; + +import type {Token, TokenProvider} from '../auth'; +import {tokenProviderFn} from '../auth'; + +import type {IDTokenProvider} from './oidc'; + +/** + * OAuthAuthorizationServer describes the OAuth endpoints used to mint + * Databricks access tokens. + */ +export interface OAuthAuthorizationServer { + tokenEndpoint: string; +} + +/** + * DatabricksOIDCTokenProviderConfig is the configuration for a Databricks OIDC + * TokenProvider. + */ +export interface DatabricksOIDCTokenProviderConfig { + /** + * ClientID of the Databricks OIDC application. It corresponds to the + * Application ID of the Databricks Service Principal. + * + * This field is only required for Workload Identity Federation and should + * be empty for Account-wide token federation. + */ + clientId?: string; + + /** + * AccountID is the account ID of the Databricks Account. This field is + * only required for Account-wide token federation. + */ + accountId?: string; + + /** + * Host is the host of the Databricks account or workspace. + */ + host: string; + + /** + * TokenEndpointProvider returns the token endpoint for the Databricks OIDC + * application. + */ + tokenEndpointProvider: () => Promise; + + /** + * Audience is the audience of the Databricks OIDC application. + * This is only used for Workspace level tokens. + */ + audience?: string; + + /** + * IDTokenProvider returns the IDToken to be used for the token exchange. + */ + idTokenProvider: IDTokenProvider; +} + +/** + * Returns a new Databricks OIDC TokenProvider that exchanges an OIDC ID token + * for a Databricks access token using the OAuth 2.0 token-exchange grant. + */ +export function newDatabricksOIDCTokenProvider( + config: DatabricksOIDCTokenProviderConfig +): TokenProvider { + return tokenProviderFn(() => exchangeIdToken(config)); +} + +async function exchangeIdToken( + config: DatabricksOIDCTokenProviderConfig +): Promise { + if (config.host === '') { + throw new Error('missing Host'); + } + const endpoints = await config.tokenEndpointProvider(); + const audience = determineAudience(config, endpoints); + const idToken = await config.idTokenProvider.idToken(audience); + + const params = new URLSearchParams(); + if (config.clientId !== undefined && config.clientId !== '') { + params.set('client_id', config.clientId); + } + params.set('scope', 'all-apis'); + 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'); + + const response = await fetch(endpoints.tokenEndpoint, { + method: 'POST', + headers: {'Content-Type': 'application/x-www-form-urlencoded'}, + body: params.toString(), + }); + if (!response.ok) { + const text = await response.text(); + throw new Error( + `token request failed with status ${response.status.toString()}: ${text}` + ); + } + const parsed = tokenResponseSchema.parse(await response.json()); + const expiry = + parsed.expires_in !== undefined + ? new Date(Date.now() + parsed.expires_in * 1000) + : undefined; + return { + value: parsed.access_token, + ...(parsed.token_type !== undefined && {type: parsed.token_type}), + ...(expiry !== undefined && {expiry}), + }; +} + +function determineAudience( + config: DatabricksOIDCTokenProviderConfig, + endpoints: OAuthAuthorizationServer +): string { + if (config.audience !== undefined && config.audience !== '') { + return config.audience; + } + if (config.accountId !== undefined && config.accountId !== '') { + return config.accountId; + } + return endpoints.tokenEndpoint; +} + +const tokenResponseSchema = z.object({ + access_token: z.string(), + token_type: z.string().optional(), + expires_in: z.number().optional(), +}); diff --git a/packages/auth/tests/oidc/env.test.ts b/packages/auth/tests/oidc/env.test.ts new file mode 100644 index 00000000..93a129bd --- /dev/null +++ b/packages/auth/tests/oidc/env.test.ts @@ -0,0 +1,81 @@ +import {afterEach, describe, expect, it, vi} from 'vitest'; + +import {newEnvIDTokenProvider} from '../../src/oidc/env'; + +describe('newEnvIDTokenProvider', () => { + afterEach(() => { + vi.unstubAllEnvs(); + }); + + const successCases: { + name: string; + envName: string; + envValue: string; + want: string; + }[] = [ + { + name: 'success', + envName: 'OIDC_TEST_TOKEN_SUCCESS', + envValue: 'test-token-123', + want: 'test-token-123', + }, + { + name: 'different variable name', + envName: 'ANOTHER_OIDC_TOKEN', + envValue: 'another-token-456', + want: 'another-token-456', + }, + ]; + + it.each(successCases)('$name', async ({envName, envValue, want}) => { + vi.stubEnv(envName, envValue); + const provider = newEnvIDTokenProvider(envName); + const token = await provider.idToken('any-audience'); + expect(token.value).toBe(want); + }); + + it('does not cache and re-reads the environment variable each call', async () => { + const envName = 'OIDC_TEST_TOKEN_REREAD'; + const provider = newEnvIDTokenProvider(envName); + + vi.stubEnv(envName, 'first'); + expect((await provider.idToken('')).value).toBe('first'); + + vi.stubEnv(envName, 'second'); + expect((await provider.idToken('')).value).toBe('second'); + }); + + it('ignores the audience argument', async () => { + const envName = 'OIDC_TEST_TOKEN_AUDIENCE'; + vi.stubEnv(envName, 'tok'); + const provider = newEnvIDTokenProvider(envName); + expect((await provider.idToken('audience-a')).value).toBe('tok'); + expect((await provider.idToken('audience-b')).value).toBe('tok'); + }); + + const errorCases: { + name: string; + envName: string; + envValue?: string; + }[] = [ + { + name: 'missing env var', + envName: 'OIDC_TEST_TOKEN_MISSING', + }, + { + name: 'empty env var', + envName: 'OIDC_TEST_TOKEN_EMPTY', + envValue: '', + }, + ]; + + it.each(errorCases)('rejects on $name', async ({envName, envValue}) => { + if (envValue !== undefined) { + vi.stubEnv(envName, envValue); + } + const provider = newEnvIDTokenProvider(envName); + await expect(provider.idToken('')).rejects.toThrow( + `missing env var "${envName}"` + ); + }); +}); diff --git a/packages/auth/tests/oidc/file.test.ts b/packages/auth/tests/oidc/file.test.ts new file mode 100644 index 00000000..d6c8844d --- /dev/null +++ b/packages/auth/tests/oidc/file.test.ts @@ -0,0 +1,67 @@ +import {mkdtemp, rm, writeFile} from 'node:fs/promises'; +import {tmpdir} from 'node:os'; +import {join} from 'node:path'; + +import {afterEach, beforeEach, describe, expect, it} from 'vitest'; + +import {newFileTokenProvider} from '../../src/oidc/file'; + +describe('newFileTokenProvider', () => { + let tmpDir: string; + + beforeEach(async () => { + tmpDir = await mkdtemp(join(tmpdir(), 'oidc-file-test-')); + }); + + afterEach(async () => { + await rm(tmpDir, {recursive: true, force: true}); + }); + + it('reads the token from a file', async () => { + const path = join(tmpDir, 'token'); + await writeFile(path, 'content'); + const provider = newFileTokenProvider(path); + const token = await provider.idToken('any-audience'); + expect(token.value).toBe('content'); + }); + + it('preserves file contents verbatim including a trailing newline', async () => { + const path = join(tmpDir, 'token'); + await writeFile(path, 'token-with-newline\n'); + const provider = newFileTokenProvider(path); + const token = await provider.idToken(''); + expect(token.value).toBe('token-with-newline\n'); + }); + + it('does not cache and re-reads the file each call', async () => { + const path = join(tmpDir, 'token'); + await writeFile(path, 'first'); + const provider = newFileTokenProvider(path); + expect((await provider.idToken('')).value).toBe('first'); + + await writeFile(path, 'second'); + expect((await provider.idToken('')).value).toBe('second'); + }); + + it('rejects when the path is empty', async () => { + const provider = newFileTokenProvider(''); + await expect(provider.idToken('')).rejects.toThrow('missing path'); + }); + + it('rejects when the file does not exist', async () => { + const path = join(tmpDir, 'does-not-exist'); + const provider = newFileTokenProvider(path); + await expect(provider.idToken('')).rejects.toThrow( + `file "${path}" does not exist` + ); + }); + + it('rejects when the file is empty', async () => { + const path = join(tmpDir, 'token'); + await writeFile(path, ''); + const provider = newFileTokenProvider(path); + await expect(provider.idToken('')).rejects.toThrow( + `file "${path}" is empty` + ); + }); +}); diff --git a/packages/auth/tests/oidc/oidc.test.ts b/packages/auth/tests/oidc/oidc.test.ts new file mode 100644 index 00000000..8b3b5adb --- /dev/null +++ b/packages/auth/tests/oidc/oidc.test.ts @@ -0,0 +1,18 @@ +import {describe, expect, it} from 'vitest'; + +import {idTokenProviderFn} from '../../src/oidc/oidc'; + +describe('idTokenProviderFn', () => { + it('forwards the audience to the wrapped function', async () => { + const provider = idTokenProviderFn(audience => + Promise.resolve({value: `id-token-for-${audience}`}) + ); + const token = await provider.idToken('my-audience'); + expect(token.value).toBe('id-token-for-my-audience'); + }); + + it('propagates rejections from the wrapped function', async () => { + const provider = idTokenProviderFn(() => Promise.reject(new Error('boom'))); + await expect(provider.idToken('audience')).rejects.toThrow('boom'); + }); +}); diff --git a/packages/auth/tests/oidc/tokensource.test.ts b/packages/auth/tests/oidc/tokensource.test.ts new file mode 100644 index 00000000..5f7fe4b9 --- /dev/null +++ b/packages/auth/tests/oidc/tokensource.test.ts @@ -0,0 +1,253 @@ +import {afterEach, describe, expect, it, vi} from 'vitest'; +import {ZodError} from 'zod'; + +import {idTokenProviderFn} from '../../src/oidc/oidc'; +import type {OAuthAuthorizationServer} from '../../src/oidc/tokensource'; +import {newDatabricksOIDCTokenProvider} from '../../src/oidc/tokensource'; + +interface CapturedRequest { + url: string; + init: RequestInit | undefined; +} + +function urlOf(input: string | URL | Request): string { + if (typeof input === 'string') { + return input; + } + if (input instanceof URL) { + return input.href; + } + return input.url; +} + +function stubFetchJson( + status: number, + body: unknown +): {captured: CapturedRequest[]; mock: ReturnType} { + const captured: CapturedRequest[] = []; + const mock = vi.fn((input, init) => { + captured.push({url: urlOf(input), init}); + return Promise.resolve( + new Response(JSON.stringify(body), { + status, + headers: {'Content-Type': 'application/json'}, + }) + ); + }); + vi.stubGlobal('fetch', mock); + return {captured, mock}; +} + +function stubFetchText( + status: number, + text: string +): {captured: CapturedRequest[]; mock: ReturnType} { + const captured: CapturedRequest[] = []; + const mock = vi.fn((input, init) => { + captured.push({url: urlOf(input), init}); + return Promise.resolve(new Response(text, {status})); + }); + vi.stubGlobal('fetch', mock); + return {captured, mock}; +} + +const TOKEN_ENDPOINT = 'https://host.com/oidc/v1/token'; +const ID_TOKEN = 'id-token-42'; + +function fixedEndpointProvider(): () => Promise { + return () => Promise.resolve({tokenEndpoint: TOKEN_ENDPOINT}); +} + +function staticIdTokenProvider(value: string): { + provider: ReturnType; + audiences: string[]; +} { + const audiences: string[] = []; + const provider = idTokenProviderFn(audience => { + audiences.push(audience); + return Promise.resolve({value}); + }); + return {provider, audiences}; +} + +describe('newDatabricksOIDCTokenProvider', () => { + const NOW = 1_700_000_000_000; + + afterEach(() => { + vi.unstubAllGlobals(); + vi.useRealTimers(); + }); + + it('rejects with "missing Host" when host is empty', async () => { + const {provider} = staticIdTokenProvider(ID_TOKEN); + const ts = newDatabricksOIDCTokenProvider({ + host: '', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: provider, + }); + await expect(ts.token()).rejects.toThrow('missing Host'); + }); + + it('propagates errors from the IDTokenProvider', async () => { + const provider = idTokenProviderFn(() => + Promise.reject(new Error('error getting id token')) + ); + const ts = newDatabricksOIDCTokenProvider({ + host: 'http://host.com', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: provider, + }); + await expect(ts.token()).rejects.toThrow('error getting id token'); + }); + + it('rejects when the token endpoint returns a non-2xx response', async () => { + stubFetchText(500, 'Internal Server Error'); + const {provider} = staticIdTokenProvider(ID_TOKEN); + const ts = newDatabricksOIDCTokenProvider({ + host: 'http://host.com', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: provider, + audience: 'token-audience', + }); + await expect(ts.token()).rejects.toThrow( + /token request failed with status 500/ + ); + }); + + it('rejects when the token endpoint returns a malformed body', async () => { + stubFetchJson(200, {foo: 'bar'}); + const {provider} = staticIdTokenProvider(ID_TOKEN); + const ts = newDatabricksOIDCTokenProvider({ + host: 'http://host.com', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: provider, + audience: 'token-audience', + }); + await expect(ts.token()).rejects.toBeInstanceOf(ZodError); + }); + + const audienceCases: { + name: string; + clientId?: string; + accountId?: string; + audience?: string; + wantAudience: string; + wantClientIdInBody: boolean; + }[] = [ + { + name: 'WIF workspace uses configured audience and sends client_id', + clientId: 'client-id', + audience: 'token-audience', + wantAudience: 'token-audience', + wantClientIdInBody: true, + }, + { + name: 'WIF account uses configured audience and sends client_id', + clientId: 'client-id', + accountId: 'ac123', + audience: 'token-audience', + wantAudience: 'token-audience', + wantClientIdInBody: true, + }, + { + name: 'account default audience falls back to accountId', + clientId: 'client-id', + accountId: 'ac123', + wantAudience: 'ac123', + wantClientIdInBody: true, + }, + { + name: 'workspace default audience falls back to the token endpoint', + clientId: 'client-id', + wantAudience: TOKEN_ENDPOINT, + wantClientIdInBody: true, + }, + { + name: 'account-wide federation omits client_id from the body', + audience: 'token-audience', + wantAudience: 'token-audience', + wantClientIdInBody: false, + }, + ]; + + it.each(audienceCases)( + '$name', + async ({ + clientId, + accountId, + audience, + wantAudience, + wantClientIdInBody, + }) => { + vi.setSystemTime(NOW); + const {captured} = stubFetchJson(200, { + token_type: 'access-token', + access_token: 'test-auth-token', + expires_in: 3600, + }); + const {provider, audiences} = staticIdTokenProvider(ID_TOKEN); + + const ts = newDatabricksOIDCTokenProvider({ + host: 'http://host.com', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: provider, + ...(clientId !== undefined && {clientId}), + ...(accountId !== undefined && {accountId}), + ...(audience !== undefined && {audience}), + }); + + 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'); + } + 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' + ); + } + ); + + it('omits expiry when expires_in is not in the response', async () => { + stubFetchJson(200, {access_token: 'test-auth-token'}); + const {provider} = staticIdTokenProvider(ID_TOKEN); + const ts = newDatabricksOIDCTokenProvider({ + host: 'http://host.com', + tokenEndpointProvider: fixedEndpointProvider(), + idTokenProvider: provider, + audience: 'token-audience', + }); + const token = await ts.token(); + expect(token.value).toBe('test-auth-token'); + expect(token.type).toBeUndefined(); + expect(token.expiry).toBeUndefined(); + }); +}); diff --git a/packages/auth/vitest.config.browser.ts b/packages/auth/vitest.config.browser.ts index a4c22eac..67ff72bd 100644 --- a/packages/auth/vitest.config.browser.ts +++ b/packages/auth/vitest.config.browser.ts @@ -9,6 +9,10 @@ export default defineConfig({ headless: true, }, include: ['tests/**/*.test.ts'], - exclude: ['tests/credentials/u2m.test.ts'], + exclude: [ + 'tests/credentials/u2m.test.ts', + 'tests/oidc/env.test.ts', + 'tests/oidc/file.test.ts', + ], }, });