From cb8f5cf447b8f27daacb0ccac03005b2ddeb729c Mon Sep 17 00:00:00 2001 From: Vu Anh Phung Date: Fri, 13 Nov 2026 18:40:49 +0000 Subject: [PATCH] fix: send the token federation exchange through the connection agent FederationProvider posted the `/oidc/v1/token` exchange with node-fetch and no agent, so `proxy`, `checkServerCertificate`, `customCaCert` and `clientCert`/`clientKey` did not apply to it. Behind a required proxy the exchange failed and the driver silently fell back to the original token. Pass the connection's agent for the exchange URL, as OAuth already does. Co-authored-by: Isaac Signed-off-by: Vu Anh Phung --- CHANGELOG.md | 1 + lib/DBSQLClient.ts | 2 ++ .../auth/tokenProvider/FederationProvider.ts | 8 +++++++ tests/unit/DBSQLClient.test.ts | 24 +++++++++++++++++++ .../tokenProvider/FederationProvider.test.ts | 22 +++++++++++++++++ 5 files changed, 57 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index 34532b62..d1a38374 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/lib/DBSQLClient.ts b/lib/DBSQLClient.ts index 4256873a..6c5fa2de 100644 --- a/lib/DBSQLClient.ts +++ b/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), }); } diff --git a/lib/connection/auth/tokenProvider/FederationProvider.ts b/lib/connection/auth/tokenProvider/FederationProvider.ts index c680c35c..62adbd9a 100644 --- a/lib/connection/auth/tokenProvider/FederationProvider.ts +++ b/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,6 +80,8 @@ export default class FederationProvider implements ITokenProvider { private readonly returnOriginalTokenOnFailure: boolean; + private readonly getAgent?: (url: string) => Promise; + /** * Creates a new FederationProvider. * @param baseProvider - The underlying token provider @@ -86,6 +89,7 @@ export default class FederationProvider implements ITokenProvider { * @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, @@ -93,12 +97,14 @@ export default class FederationProvider implements ITokenProvider { options?: { clientId?: string; returnOriginalTokenOnFailure?: boolean; + getAgent?: (url: string) => Promise; }, ) { this.baseProvider = baseProvider; this.databricksHost = databricksHost; this.clientId = options?.clientId; this.returnOriginalTokenOnFailure = options?.returnOriginalTokenOnFailure ?? true; + this.getAgent = options?.getAgent; } async getToken(): Promise { @@ -162,12 +168,14 @@ export default class FederationProvider implements ITokenProvider { */ private async attemptTokenExchange(body: string): Promise { 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', }, diff --git a/tests/unit/DBSQLClient.test.ts b/tests/unit/DBSQLClient.test.ts index 386f268f..b8e007ba 100644 --- a/tests/unit/DBSQLClient.test.ts +++ b/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(); diff --git a/tests/unit/connection/auth/tokenProvider/FederationProvider.test.ts b/tests/unit/connection/auth/tokenProvider/FederationProvider.test.ts index fbf6a817..abb4c694 100644 --- a/tests/unit/connection/auth/tokenProvider/FederationProvider.test.ts +++ b/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, ReturnType>() + .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);