fix: send the token federation exchange through the connection agent

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- The Thrift token federation exchange (`enableTokenFederation`) now goes through the connection's agent, so `proxy`, `checkServerCertificate`, `customCaCert` and `clientCert`/`clientKey` apply to it. Previously it connected directly with Node's default agent, so behind a required proxy the exchange failed and the driver fell back to the original token.
 - Thrift CloudFetch downloads are now bounded end to end by `retriesTimeout` (default 15 minutes), including every attempt and retry backoff. Previously a stalled download, where the response headers or body never arrived, blocked the fetch forever. A download still running when the deadline expires is aborted and fails with `RetryError` (`TIMEOUT_EXCEEDED`, message `CloudFetch result download failed: Retry timeout exceeded`). (#520)
 - Thrift connections made through `proxy` now apply `checkServerCertificate`, `customCaCert` and `clientCert`/`clientKey`. Previously the proxied path ignored them and verified against Node's process-wide trust store, so a certificate outside the configured trust set could be accepted. **Compatibility:** proxied connections now verify unless `checkServerCertificate: false` is set; as on direct connections, `NODE_TLS_REJECT_UNAUTHORIZED=0` no longer turns verification off. For an `https` proxy these options also apply to the proxy's own certificate. (#485)
 - Honor the client `useCloudFetch: false` default on the kernel backend so complete results can be fetched without cloud-storage downloads. (#574)

lib/DBSQLClient.ts

@@ -365,6 +365,8 @@ export default class DBSQLClient extends EventEmitter implements IDBSQLClient, I
     if (enableFederation) {
       wrapped = new FederationProvider(wrapped, host, {
         clientId: federationClientId,
+        // The connection's agent carries `proxy` and the TLS options.
+        getAgent: async (url) => (await this.getConnectionProvider()).getAgent(url),
       });
     }
 

lib/connection/auth/tokenProvider/FederationProvider.ts

@@ -1,3 +1,4 @@
+import http from 'http';
 import nodeFetch from 'node-fetch';
 import ITokenProvider from './ITokenProvider';
 import Token from './Token';
@@ -79,26 +80,31 @@ export default class FederationProvider implements ITokenProvider {
 
   private readonly returnOriginalTokenOnFailure: boolean;
 
+  private readonly getAgent?: (url: string) => Promise<http.Agent | undefined>;
+
   /**
    * Creates a new FederationProvider.
    * @param baseProvider - The underlying token provider
    * @param databricksHost - The Databricks workspace host URL
    * @param options - Optional configuration
    * @param options.clientId - Client ID for M2M/service principal federation
    * @param options.returnOriginalTokenOnFailure - Return original token if exchange fails (default: true)
+   * @param options.getAgent - Agent for the exchange request, so the connection's proxy and TLS options apply
    */
   constructor(
     baseProvider: ITokenProvider,
     databricksHost: string,
     options?: {
       clientId?: string;
       returnOriginalTokenOnFailure?: boolean;
+      getAgent?: (url: string) => Promise<http.Agent | undefined>;
     },
   ) {
     this.baseProvider = baseProvider;
     this.databricksHost = databricksHost;
     this.clientId = options?.clientId;
     this.returnOriginalTokenOnFailure = options?.returnOriginalTokenOnFailure ?? true;
+    this.getAgent = options?.getAgent;
   }
 
   async getToken(): Promise<Token> {
@@ -162,12 +168,14 @@ export default class FederationProvider implements ITokenProvider {
    */
   private async attemptTokenExchange(body: string): Promise<Token> {
     const url = this.buildExchangeUrl();
+    const agent = await this.getAgent?.(url);
     const controller = new AbortController();
     const timeoutId = setTimeout(() => controller.abort(), REQUEST_TIMEOUT_MS);
 
     try {
       const response = await fetchImpl(url, {
         method: 'POST',
+        agent,
         headers: {
           'Content-Type': 'application/x-www-form-urlencoded',
         },

tests/unit/DBSQLClient.test.ts

@@ -1,6 +1,7 @@
 import { expect, AssertionError } from 'chai';
 import sinon from 'sinon';
 import fs from 'fs';
+import http from 'http';
 import { Response } from 'node-fetch';
 import DBSQLClient, { ThriftLibrary } from '../../lib/DBSQLClient';
 import DBSQLSession from '../../lib/DBSQLSession';
@@ -9,6 +10,7 @@ import KernelBackend from '../../lib/kernel/KernelBackend';
 import * as KernelNativeLoader from '../../lib/kernel/KernelNativeLoader';
 
 import PlainHttpAuthentication from '../../lib/connection/auth/PlainHttpAuthentication';
+import FederationProvider from '../../lib/connection/auth/tokenProvider/FederationProvider';
 import DatabricksOAuth from '../../lib/connection/auth/DatabricksOAuth';
 import { DatabricksOAuthManager, AzureOAuthManager } from '../../lib/connection/auth/DatabricksOAuth/OAuthManager';
 
@@ -674,6 +676,28 @@ describe('DBSQLClient[Symbol.asyncDispose]', () => {
 });
 
 describe('DBSQLClient.createAuthProvider', () => {
+  it('should send the token federation exchange through the connection agent', async () => {
+    const client = new DBSQLClient();
+    const connectionProvider = new ConnectionProviderStub();
+    const agent = new http.Agent();
+    const getAgent = sinon.stub(connectionProvider, 'getAgent').resolves(agent);
+    client['connectionProvider'] = connectionProvider;
+
+    const provider = client['createAuthProvider']({
+      ...connectOptions,
+      authType: 'static-token',
+      staticToken: 'token',
+      enableTokenFederation: true,
+    });
+
+    const federation = (provider as unknown as { tokenProvider: FederationProvider }).tokenProvider;
+    expect(federation).to.be.instanceOf(FederationProvider);
+    const url = 'https://dbc-XXXXXXXX-XXXX.cloud.databricks.com/oidc/v1/token';
+    expect(await federation['getAgent']!(url)).to.equal(agent);
+    expect(getAgent.calledOnce).to.be.true;
+    expect((getAgent.firstCall.args as unknown[])[0]).to.equal(url);
+  });
+
   it('should use access token auth method', () => {
     const client = new DBSQLClient();
 

tests/unit/connection/auth/tokenProvider/FederationProvider.test.ts

@@ -1,4 +1,5 @@
 import { expect } from 'chai';
+import http from 'http';
 import sinon from 'sinon';
 import type nodeFetch from 'node-fetch';
 import FederationProvider, {
@@ -132,6 +133,27 @@ describe('FederationProvider', () => {
       expect(passedSignal.aborted).to.be.false;
     });
 
+    it('should send the exchange through the agent from getAgent', async () => {
+      const foreignJwt = createJWT({ iss: 'https://idp.example.com' });
+      const agent = new http.Agent();
+      const getAgent = sinon.stub().resolves(agent);
+      const federationProvider = new FederationProvider(
+        new MockTokenProvider(foreignJwt),
+        'my-workspace.cloud.databricks.com',
+        { getAgent },
+      );
+
+      const fetchStub = sinon
+        .stub<Parameters<typeof nodeFetch>, ReturnType<typeof nodeFetch>>()
+        .resolves(buildFakeResponse({ ok: true, body: { access_token: 'exchanged-databricks-token' } }));
+      setFederationFetchForTest(fetchStub as unknown as typeof nodeFetch);
+
+      await federationProvider.getToken();
+
+      expect(getAgent.calledOnceWithExactly('https://my-workspace.cloud.databricks.com/oidc/v1/token')).to.be.true;
+      expect(fetchStub.firstCall.args[1]!.agent).to.equal(agent);
+    });
+
     it('should propagate abort from the controller to the signal observed by fetch', async () => {
       const foreignJwt = createJWT({ iss: 'https://idp.example.com' });
       const baseProvider = new MockTokenProvider(foreignJwt);