fix: discover the Azure tenant for Thrift service-principal M2M

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)

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

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<BaseClient> {
+    // 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`;
   }
 

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<unknown>, contains: string): Promise<void> {
+  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<string | undefined> = [];
+
+  before(async () => {
+    server = http.createServer((req, res) => {
+      paths.push(req.url);
+      respond(req, res);
+    });
+    await new Promise<void>((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<void>((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<void>((resolve) => {
+      closed.listen(0, '127.0.0.1', resolve);
+    });
+    const { port } = closed.address() as AddressInfo;
+    await new Promise<void>((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');
+  });
+});

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<OAuthManagerOptions>) {
+    // 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<string, (...a: unknown[]) => 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<string>(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<string>(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<OpenIDClientStub> = {
+      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<OAuthManagerOptions> = {}) {
+    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<unknown> }).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() });