diff --git a/CHANGELOG.md b/CHANGELOG.md index 34532b62..bc3c2603 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,7 @@ ## Unreleased +- Thrift Azure service-principal M2M (`oauthClientSecret` on an Azure host) now discovers the workspace's Entra tenant from the `/aad/auth` redirect when `azureTenantId` is not set, as the kernel backend and the Python connector do. Previously the client-credentials grant went to the `/organizations` authority, which Entra rejects for that grant, so the connection failed unless `azureTenantId` was passed. An explicit `azureTenantId` still skips discovery; U2M is unchanged. - 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/connection/auth/DatabricksOAuth/AzureTenantDiscovery.ts b/lib/connection/auth/DatabricksOAuth/AzureTenantDiscovery.ts new file mode 100644 index 00000000..33977aa7 --- /dev/null +++ b/lib/connection/auth/DatabricksOAuth/AzureTenantDiscovery.ts @@ -0,0 +1,74 @@ +import http from 'http'; +import fetch from 'node-fetch'; +import AuthenticationError from '../../../errors/AuthenticationError'; + +// Microsoft Entra login authorities (commercial, US Government, China). Matched as host suffixes, +// so a look-alike such as `login.microsoftonline.com.attacker.example` is rejected. +const MICROSOFT_LOGIN_HOST_SUFFIXES = ['.microsoftonline.com', '.microsoftonline.us', '.chinacloudapi.cn']; + +const DISCOVERY_TIMEOUT_MS = 30000; + +/** + * Returns the tenant id from the Entra login URL that a workspace's `/aad/auth` redirects to + * (`https://login.microsoftonline.com/{tenant}/oauth2/authorize?...`): its first path segment. + */ +export function tenantFromLoginUrl(location: string): string { + let url: URL; + try { + url = new URL(location); + } catch { + throw new AuthenticationError( + `Azure tenant discovery: /aad/auth redirected to an invalid URL; pass azureTenantId explicitly`, + ); + } + + const host = url.hostname.toLowerCase(); + if (!MICROSOFT_LOGIN_HOST_SUFFIXES.some((suffix) => host.endsWith(suffix))) { + throw new AuthenticationError( + `Azure tenant discovery: /aad/auth did not redirect to a Microsoft login host (got ${host}); pass azureTenantId explicitly`, + ); + } + + const tenant = url.pathname.split('/')[1]; + if (!tenant) { + throw new AuthenticationError( + 'Azure tenant discovery: could not read a tenant id from the /aad/auth redirect; pass azureTenantId explicitly', + ); + } + return tenant; +} + +/** + * Discovers the workspace's Entra tenant from the first redirect of `{host}/aad/auth`, as the + * Python connector and the SQL kernel do. + */ +export async function discoverAzureTenantId(host: string, agent?: http.Agent): Promise { + const base = (host.includes('://') ? host : `https://${host}`).replace(/\/+$/, ''); + let response; + try { + response = await fetch(`${base}/aad/auth`, { + method: 'GET', + redirect: 'manual', + agent, + timeout: DISCOVERY_TIMEOUT_MS, + }); + } catch (error) { + // Report only the code/type: a network error's message can carry request details. + const { code, type } = error as { code?: string; type?: string }; + throw new AuthenticationError( + `Azure tenant discovery: ${base}/aad/auth failed (${ + code ?? type ?? 'network error' + }); pass azureTenantId explicitly`, + ); + } + + // Drain the (small) body so the keep-alive agent can reuse the socket. + await response.text().catch(() => undefined); + const location = response.headers.get('location'); + if (response.status < 300 || response.status >= 400 || !location) { + throw new AuthenticationError( + `Azure tenant discovery: ${base}/aad/auth returned HTTP ${response.status} without a redirect; pass azureTenantId explicitly`, + ); + } + return tenantFromLoginUrl(location); +} diff --git a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts index 40ff80ae..58d15601 100644 --- a/lib/connection/auth/DatabricksOAuth/OAuthManager.ts +++ b/lib/connection/auth/DatabricksOAuth/OAuthManager.ts @@ -6,6 +6,7 @@ import OAuthToken from './OAuthToken'; import AuthorizationCode from './AuthorizationCode'; import { OAuthScope, OAuthScopes, defaultOAuthScopes, scopeDelimiter } from './OAuthScope'; import IClientContext from '../../../contracts/IClientContext'; +import { discoverAzureTenantId } from './AzureTenantDiscovery'; export enum OAuthFlow { U2M = 'U2M', @@ -289,10 +290,23 @@ export class AzureOAuthManager extends OAuthManager { public static datatricksAzureApp = '2ff814a6-3304-4ab8-85cb-cd0e6f879c1d'; + private discoveredTenantId?: string; + + protected async getClient(): Promise { + // Entra rejects the client-credentials grant on the multi-tenant `/organizations` authority, + // so M2M without an explicit tenant discovers the workspace's tenant first. + if (this.options.flow === OAuthFlow.M2M && !this.options.azureTenantId?.trim() && !this.discoveredTenantId) { + const connectionProvider = await this.context.getConnectionProvider(); + const agent = await connectionProvider.getAgent('https:'); + this.discoveredTenantId = await discoverAzureTenantId(this.options.host, agent); + } + return super.getClient(); + } + protected getOIDCConfigUrl(): string { // Use logical OR so empty / whitespace-only azureTenantId also falls back to /organizations/ // (`??` only substitutes for null/undefined, leaving `''` to produce a malformed `//v2.0/...` URL). - const tenantPath = this.options.azureTenantId?.trim() || 'organizations'; + const tenantPath = this.options.azureTenantId?.trim() || this.discoveredTenantId || 'organizations'; return `https://login.microsoftonline.com/${tenantPath}/v2.0/.well-known/openid-configuration`; } diff --git a/tests/unit/connection/auth/DatabricksOAuth/AzureTenantDiscovery.test.ts b/tests/unit/connection/auth/DatabricksOAuth/AzureTenantDiscovery.test.ts new file mode 100644 index 00000000..0c55e015 --- /dev/null +++ b/tests/unit/connection/auth/DatabricksOAuth/AzureTenantDiscovery.test.ts @@ -0,0 +1,137 @@ +import { expect, AssertionError } from 'chai'; +import http from 'http'; +import { AddressInfo } from 'net'; +import { + discoverAzureTenantId, + tenantFromLoginUrl, +} from '../../../../../lib/connection/auth/DatabricksOAuth/AzureTenantDiscovery'; +import AuthenticationError from '../../../../../lib/errors/AuthenticationError'; + +async function expectAuthError(promise: Promise, contains: string): Promise { + try { + await promise; + expect.fail('It should throw an error'); + } catch (error) { + if (error instanceof AssertionError) throw error; + expect(error).to.be.instanceOf(AuthenticationError); + expect((error as Error).message).to.contain(contains); + } +} + +describe('tenantFromLoginUrl', () => { + it('reads the tenant from the first path segment of the Entra login URL', () => { + expect(tenantFromLoginUrl('https://login.microsoftonline.com/tenant-123/oauth2/authorize?client_id=x')).to.equal( + 'tenant-123', + ); + }); + + it('accepts the sovereign Microsoft login hosts', () => { + for (const host of ['login.microsoftonline.us', 'login.chinacloudapi.cn']) { + expect(tenantFromLoginUrl(`https://${host}/tenant-abc/oauth2/authorize`)).to.equal('tenant-abc'); + } + }); + + it('rejects a redirect to a non-Microsoft or look-alike host', () => { + for (const location of [ + 'https://evil.example.com/tenant-abc/oauth2/authorize', + 'https://login.microsoftonline.com.attacker.example/tenant-abc/oauth2/authorize', + ]) { + expect(() => tenantFromLoginUrl(location)).to.throw(AuthenticationError, 'Microsoft login host'); + } + }); + + it('rejects a URL without a tenant segment or an invalid URL', () => { + expect(() => tenantFromLoginUrl('https://login.microsoftonline.com/')).to.throw(AuthenticationError, 'tenant id'); + expect(() => tenantFromLoginUrl('not a url')).to.throw(AuthenticationError, 'invalid URL'); + }); +}); + +describe('discoverAzureTenantId', () => { + let server: http.Server; + let base: string; + let respond: (req: http.IncomingMessage, res: http.ServerResponse) => void; + const paths: Array = []; + + before(async () => { + server = http.createServer((req, res) => { + paths.push(req.url); + respond(req, res); + }); + await new Promise((resolve) => { + server.listen(0, '127.0.0.1', resolve); + }); + base = `http://127.0.0.1:${(server.address() as AddressInfo).port}`; + }); + + after(async () => { + await new Promise((resolve) => { + server.close(() => resolve()); + }); + }); + + beforeEach(() => { + paths.length = 0; + }); + + it('reads the tenant from the first /aad/auth redirect without following it', async () => { + respond = (req, res) => { + res.writeHead(302, { Location: 'https://login.microsoftonline.com/tenant-xyz/oauth2/authorize?client_id=x' }); + res.end(); + }; + expect(await discoverAzureTenantId(`${base}/`)).to.equal('tenant-xyz'); + expect(paths).to.deep.equal(['/aad/auth']); + }); + + it('drains the redirect body so a keep-alive agent reuses the socket', async () => { + respond = (req, res) => { + res.writeHead(302, { Location: 'https://login.microsoftonline.com/tenant-xyz/oauth2/authorize' }); + // Larger than the stream buffers, so an undrained body would hold the only socket. + res.end('x'.repeat(1024 * 1024)); + }; + let connections = 0; + const onConnection = () => { + connections += 1; + }; + server.on('connection', onConnection); + const agent = new http.Agent({ keepAlive: true, maxSockets: 1 }); + try { + expect(await discoverAzureTenantId(base, agent)).to.equal('tenant-xyz'); + expect(await discoverAzureTenantId(base, agent)).to.equal('tenant-xyz'); + expect(connections).to.equal(1); + } finally { + server.off('connection', onConnection); + agent.destroy(); + } + }); + + it('fails when /aad/auth does not redirect', async () => { + respond = (req, res) => { + res.writeHead(200); + res.end(); + }; + await expectAuthError(discoverAzureTenantId(base), 'HTTP 200 without a redirect'); + }); + + it('wraps a network failure with the azureTenantId hint', async () => { + const closed = http.createServer(); + await new Promise((resolve) => { + closed.listen(0, '127.0.0.1', resolve); + }); + const { port } = closed.address() as AddressInfo; + await new Promise((resolve) => { + closed.close(() => resolve()); + }); + await expectAuthError( + discoverAzureTenantId(`http://127.0.0.1:${port}`), + 'failed (ECONNREFUSED); pass azureTenantId', + ); + }); + + it('fails when the redirect is not to a Microsoft login host', async () => { + respond = (req, res) => { + res.writeHead(302, { Location: 'https://idp.example.com/tenant-xyz/authorize' }); + res.end(); + }; + await expectAuthError(discoverAzureTenantId(base), 'Microsoft login host'); + }); +}); diff --git a/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.ts b/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.ts index 179f6e28..0d0edbcc 100644 --- a/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.ts +++ b/tests/unit/connection/auth/DatabricksOAuth/OAuthManager.test.ts @@ -1,10 +1,13 @@ import { AssertionError, expect } from 'chai'; +import https from 'https'; import sinon, { SinonStub } from 'sinon'; import { Issuer, BaseClient, TokenSet, GrantBody, IssuerMetadata, ClientMetadata, custom } from 'openid-client'; // Import the whole module once more - to stub some of its exports import openidClientLib from 'openid-client'; // Import the whole module to stub its default export import * as AuthorizationCodeModule from '../../../../../lib/connection/auth/DatabricksOAuth/AuthorizationCode'; +import * as AzureTenantDiscoveryModule from '../../../../../lib/connection/auth/DatabricksOAuth/AzureTenantDiscovery'; +import AuthenticationError from '../../../../../lib/errors/AuthenticationError'; import { AzureOAuthManager, @@ -155,9 +158,12 @@ class OpenIDClientStub implements BaseClient { afterEach(() => { (openidClientLib.Issuer as unknown as SinonStub).restore?.(); (AuthorizationCodeModule.default as unknown as SinonStub).restore?.(); + (AzureTenantDiscoveryModule.discoverAzureTenantId as unknown as SinonStub).restore?.(); }); function prepareTestInstances(options: Partial) { + // Azure M2M without a tenant discovers it over HTTP; keep these tests offline. + sinon.stub(AzureTenantDiscoveryModule, 'discoverAzureTenantId').resolves('test-tenant'); const oauthClient = sinon.spy(new OpenIDClientStub()); oauthClient.expectedClientId = options?.clientId; @@ -535,7 +541,7 @@ describe('AzureOAuthManager (tenant awareness)', () => { (mgr as unknown as Record T>)[name](...args); describe('getOIDCConfigUrl', () => { - it('falls back to /organizations/ when azureTenantId is not set (baseline-compatible)', () => { + it('uses /organizations/ when azureTenantId is not set and no tenant was discovered yet', () => { const mgr = makeAzure(); const url = call(mgr, 'getOIDCConfigUrl'); expect(url).to.equal('https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration'); @@ -548,7 +554,7 @@ describe('AzureOAuthManager (tenant awareness)', () => { expect(url).to.equal(`https://login.microsoftonline.com/${tenant}/v2.0/.well-known/openid-configuration`); }); - it('falls back to /organizations/ when azureTenantId is empty or whitespace', () => { + it('uses /organizations/ before discovery when azureTenantId is empty or whitespace', () => { for (const tenant of ['', ' ']) { const mgr = makeAzure({ azureTenantId: tenant }); const url = call(mgr, 'getOIDCConfigUrl'); @@ -594,6 +600,99 @@ describe('AzureOAuthManager (tenant awareness)', () => { }); }); +describe('AzureOAuthManager tenant discovery', () => { + let discoverStub: SinonStub; + let oidcConfigUrls: string[]; + + beforeEach(() => { + oidcConfigUrls = []; + const oauthClient = new OpenIDClientStub(); + const issuer: Issuer = { + Client: class extends OpenIDClientStub { + constructor(clientOptions: ClientMetadata) { + super(); + oauthClient.clientOptions = clientOptions; + return oauthClient; + } + }, + FAPI1Client: OpenIDClientStub, + metadata: { issuer: 'test' }, + [custom.http_options]: () => ({}), + discover: async () => issuer, + }; + sinon.stub(openidClientLib, 'Issuer').returns(issuer); + openidClientLib.Issuer.discover = async (url: string) => { + oidcConfigUrls.push(url); + return issuer; + }; + discoverStub = sinon.stub(AzureTenantDiscoveryModule, 'discoverAzureTenantId').resolves('discovered-tenant'); + }); + + afterEach(() => { + (openidClientLib.Issuer as unknown as SinonStub).restore?.(); + discoverStub.restore(); + }); + + const httpsAgent = new https.Agent(); + + function makeManager(overrides: Partial = {}) { + const context = new ClientContextStub(); + sinon + .stub(context.connectionProvider, 'getAgent') + .callsFake(async (url?: string | URL) => (url === 'https:' ? httpsAgent : undefined)); + return new AzureOAuthManager({ + host: 'adb-1234567890123456.1.azuredatabricks.net', + flow: OAuthFlow.M2M, + clientId: 'test_client_id', + clientSecret: 'test_client_secret', + context, + ...overrides, + }); + } + + it('discovers the tenant once for M2M without azureTenantId and uses it for the token request', async () => { + const mgr = makeManager(); + await mgr.getToken(); + await mgr.getToken(); + expect(discoverStub.calledOnce).to.be.true; + expect(discoverStub.firstCall.args[0]).to.equal('adb-1234567890123456.1.azuredatabricks.net'); + // The connection's https agent carries the proxy and TLS options. + expect(discoverStub.firstCall.args[1]).to.equal(httpsAgent); + expect(oidcConfigUrls).to.deep.equal([ + 'https://login.microsoftonline.com/discovered-tenant/v2.0/.well-known/openid-configuration', + ]); + }); + + it('uses an explicit azureTenantId without discovery', async () => { + await makeManager({ azureTenantId: 'explicit-tenant' }).getToken(); + expect(discoverStub.called).to.be.false; + expect(oidcConfigUrls).to.deep.equal([ + 'https://login.microsoftonline.com/explicit-tenant/v2.0/.well-known/openid-configuration', + ]); + }); + + it('keeps /organizations for U2M', async () => { + const mgr = makeManager({ flow: OAuthFlow.U2M, clientSecret: undefined }); + await (mgr as unknown as { getClient(): Promise }).getClient(); + expect(discoverStub.called).to.be.false; + expect(oidcConfigUrls).to.deep.equal([ + 'https://login.microsoftonline.com/organizations/v2.0/.well-known/openid-configuration', + ]); + }); + + it('surfaces a discovery failure instead of falling back to /organizations', async () => { + discoverStub.rejects(new AuthenticationError('Azure tenant discovery: no redirect')); + try { + await makeManager().getToken(); + expect.fail('It should throw an error'); + } catch (error) { + if (error instanceof AssertionError) throw error; + expect((error as Error).message).to.contain('Azure tenant discovery'); + } + expect(oidcConfigUrls).to.deep.equal([]); + }); +}); + describe('DatabricksOAuthManager.getScopes (AUTH-017)', () => { function makeManager(flow: OAuthFlow) { return new DatabricksOAuthManager({ host: 'example.cloud.databricks.com', flow, context: new ClientContextStub() });