env and file oidc

.gitignore

@@ -10,3 +10,6 @@ dist/
 
 # Vitest browser + Playwright failure captures
 __screenshots__/
+
+# Development Environment
+.claude/scheduled_tasks.lock

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": [

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});
+  });
+}

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;
+}

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';

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';

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<IDToken>;
+}
+
+/**
+ * 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<IDToken>
+): IDTokenProvider {
+  return {idToken: fn};
+}

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<OAuthAuthorizationServer>;
+
+  /**
+   * 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<Token> {
+  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(),
+});

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}"`
+    );
+  });
+});