Skip to content

Commit 093a612

Browse files
hectorcast-dbhectorcast-db
authored
Separate Databricks OIDC endpoint logic into a separate method (#659)
## What changes are proposed in this pull request? Separate Databricks OIDC endpoint logic into a separate method. The `getOidcEndpoints()` property was incorrectly returning Azure OIDC endpoints when `ARM_CLIENT_ID` was set, even for Databricks OAuth flows (like `oauth-m2m`) that don't use Azure authentication. This was changed on a previous PR to fix M2M authentication: #657 Given that some users may be depending on this behavior, this PR returns the previous behavior to this method and instead creates a dedicated `getDatabricksOidcEndpoints()` method. ## How is this tested? Run unit tests and integration tests with appropriate environment variables set.
1 parent c2a8a8c commit 093a612

File tree

8 files changed

+65
-22
lines changed

8 files changed

+65
-22
lines changed

‎NEXT_CHANGELOG.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
### New Features and Improvements
66

77
### Bug Fixes
8+
* Fixed Databricks M2M OAuth to correctly use Databricks OIDC endpoints instead of incorrectly using Azure endpoints when `ARM_CLIENT_ID` is set. Added new `getDatabricksOidcEndpoints` method that returns only Databricks OIDC endpoints, and updated all Databricks OAuth flows to use it. The old `getOidcEndpoints` property is deprecated but maintained for backward compatibility.
89

910
### Security Vulnerabilities
1011

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/DatabricksConfig.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -758,7 +758,46 @@ public ClientType getClientType() {
758758
}
759759
}
760760

761+
/**
762+
* @deprecated Use {@link #getDatabricksOidcEndpoints()} instead. This method incorrectly returns
763+
* Azure OIDC endpoints when azure_client_id is set, even for Databricks OAuth flows that
764+
* don't use Azure authentication. This caused bugs where Databricks M2M OAuth would fail when
765+
* ARM_CLIENT_ID was set for other purposes. Use instead: - getDatabricksOidcEndpoints(): For
766+
* Databricks OAuth (oauth-m2m, external-browser, etc.). -
767+
* getAzureEntraIdWorkspaceEndpoints(): For Azure Entra ID OIDC endpoints.
768+
* @return The OIDC endpoints. This method dinamically returns the OIDC endpoints based on the
769+
* config.
770+
*/
771+
@Deprecated
761772
public OpenIDConnectEndpoints getOidcEndpoints() throws IOException {
773+
if (isAzure() && getAzureClientId() != null) {
774+
return getAzureEntraIdWorkspaceEndpoints();
775+
}
776+
return getDatabricksOidcEndpoints();
777+
}
778+
779+
/**
780+
* @return The Azure Entra ID OIDC endpoints.
781+
*/
782+
public OpenIDConnectEndpoints getAzureEntraIdWorkspaceEndpoints() throws IOException {
783+
if (isAzure() && getAzureClientId() != null) {
784+
Request request = new Request("GET", getHost() + "/oidc/oauth2/v2.0/authorize");
785+
request.setRedirectionBehavior(false);
786+
Response resp = getHttpClient().execute(request);
787+
String realAuthUrl = resp.getFirstHeader("location");
788+
if (realAuthUrl == null) {
789+
return null;
790+
}
791+
return new OpenIDConnectEndpoints(
792+
realAuthUrl.replaceAll("/authorize", "/token"), realAuthUrl);
793+
}
794+
return null;
795+
}
796+
797+
/**
798+
* @return The Databricks OIDC endpoints.
799+
*/
800+
public OpenIDConnectEndpoints getDatabricksOidcEndpoints() throws IOException {
762801
if (discoveryUrl == null) {
763802
return fetchDefaultOidcEndpoints();
764803
}

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/DefaultCredentialsProvider.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -101,7 +101,7 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
101101
// This would also need to be updated to support unified hosts.
102102
OpenIDConnectEndpoints endpoints = null;
103103
try {
104-
endpoints = config.getOidcEndpoints();
104+
endpoints = config.getDatabricksOidcEndpoints();
105105
} catch (Exception e) {
106106
LOG.warn("Failed to get OpenID Connect endpoints", e);
107107
}

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthClientUtils.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ public static String resolveClientSecret(DatabricksConfig config) {
5757
public static OpenIDConnectEndpoints resolveOidcEndpoints(DatabricksConfig config)
5858
throws IOException {
5959
if (config.getClientId() != null && config.getClientSecret() != null) {
60-
return config.getOidcEndpoints();
60+
return config.getDatabricksOidcEndpoints();
6161
} else if (config.getAzureClientId() != null && config.getAzureClientSecret() != null) {
6262
Request request = new Request("GET", config.getHost() + "/oidc/oauth2/v2.0/authorize");
6363
request.setRedirectionBehavior(false);
@@ -69,6 +69,6 @@ public static OpenIDConnectEndpoints resolveOidcEndpoints(DatabricksConfig confi
6969
return new OpenIDConnectEndpoints(
7070
realAuthUrl.replaceAll("/authorize", "/token"), realAuthUrl);
7171
}
72-
return config.getOidcEndpoints();
72+
return config.getDatabricksOidcEndpoints();
7373
}
7474
}

‎databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/OAuthM2MServicePrincipalCredentialsProvider.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
2727
// TODO: Azure returns 404 for UC workspace after redirecting to
2828
// https://login.microsoftonline.com/{cfg.azure_tenant_id}/.well-known/oauth-authorization-server
2929
try {
30-
OpenIDConnectEndpoints jsonResponse = config.getOidcEndpoints();
30+
OpenIDConnectEndpoints jsonResponse = config.getDatabricksOidcEndpoints();
3131
ClientCredentials clientCredentials =
3232
new ClientCredentials.Builder()
3333
.withHttpClient(config.getHttpClient())

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/DatabricksConfigTest.java‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ public void testWorkspaceLevelOidcEndpointsWithAccountId() throws IOException {
108108
c.resolve(
109109
new Environment(new HashMap<>(), new ArrayList<String>(), System.getProperty("os.name")));
110110
assertEquals(
111-
c.getOidcEndpoints().getAuthorizationEndpoint(),
111+
c.getDatabricksOidcEndpoints().getAuthorizationEndpoint(),
112112
"https://test-workspace.cloud.databricks.com/oidc/v1/authorize");
113113
}
114114
}
@@ -128,7 +128,7 @@ public void testWorkspaceLevelOidcEndpointsRetries() throws IOException {
128128
c.resolve(
129129
new Environment(new HashMap<>(), new ArrayList<String>(), System.getProperty("os.name")));
130130
assertEquals(
131-
c.getOidcEndpoints().getAuthorizationEndpoint(),
131+
c.getDatabricksOidcEndpoints().getAuthorizationEndpoint(),
132132
"https://test-workspace.cloud.databricks.com/oidc/v1/authorize");
133133
}
134134
}
@@ -139,7 +139,7 @@ public void testAccountLevelOidcEndpoints() throws IOException {
139139
new DatabricksConfig()
140140
.setHost("https://accounts.cloud.databricks.com")
141141
.setAccountId("1234567890")
142-
.getOidcEndpoints()
142+
.getDatabricksOidcEndpoints()
143143
.getAuthorizationEndpoint(),
144144
"https://accounts.cloud.databricks.com/oidc/accounts/1234567890/v1/authorize");
145145
}
@@ -163,7 +163,7 @@ public void testDiscoveryEndpoint() throws IOException {
163163
.setHost(server.getUrl())
164164
.setDiscoveryUrl(discoveryUrl)
165165
.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build())
166-
.getOidcEndpoints();
166+
.getDatabricksOidcEndpoints();
167167

168168
assertEquals(
169169
oidcEndpoints.getAuthorizationEndpoint(), "https://test.auth.endpoint/oidc/v1/authorize");

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/UnifiedHostTest.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -119,7 +119,7 @@ public void testOidcEndpointsForUnifiedHost() throws IOException {
119119
.setExperimentalIsUnifiedHost(true)
120120
.setAccountId("test-account-123");
121121

122-
OpenIDConnectEndpoints endpoints = config.getOidcEndpoints();
122+
OpenIDConnectEndpoints endpoints = config.getDatabricksOidcEndpoints();
123123

124124
assertEquals(
125125
"https://unified.databricks.com/oidc/accounts/test-account-123/v1/authorize",
@@ -138,7 +138,7 @@ public void testOidcEndpointsForUnifiedHostMissingAccountId() {
138138
// No account ID set
139139

140140
DatabricksException exception =
141-
assertThrows(DatabricksException.class, () -> config.getOidcEndpoints());
141+
assertThrows(DatabricksException.class, () -> config.getDatabricksOidcEndpoints());
142142
assertTrue(exception.getMessage().contains("account_id is required"));
143143
}
144144

‎databricks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/ExternalBrowserCredentialsProviderTest.java‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -44,15 +44,16 @@ void clientAndConsentTest() throws IOException {
4444
.setHttpClient(new CommonsHttpClient.Builder().withTimeoutSeconds(30).build());
4545
config.resolve();
4646

47-
assertEquals("tokenEndPointFromServer", config.getOidcEndpoints().getTokenEndpoint());
47+
assertEquals(
48+
"tokenEndPointFromServer", config.getDatabricksOidcEndpoints().getTokenEndpoint());
4849

4950
OAuthClient testClient =
5051
new OAuthClient.Builder()
5152
.withHttpClient(config.getHttpClient())
5253
.withClientId(config.getClientId())
5354
.withClientSecret(config.getClientSecret())
5455
.withHost(config.getHost())
55-
.withOpenIDConnectEndpoints(config.getOidcEndpoints())
56+
.withOpenIDConnectEndpoints(config.getDatabricksOidcEndpoints())
5657
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
5758
.withScopes(config.getScopes())
5859
.build();
@@ -92,15 +93,16 @@ void clientAndConsentTestWithCustomRedirectUrl() throws IOException {
9293
.setScopes(Arrays.asList("sql"));
9394
config.resolve();
9495

95-
assertEquals("tokenEndPointFromServer", config.getOidcEndpoints().getTokenEndpoint());
96+
assertEquals(
97+
"tokenEndPointFromServer", config.getDatabricksOidcEndpoints().getTokenEndpoint());
9698

9799
OAuthClient testClient =
98100
new OAuthClient.Builder()
99101
.withHttpClient(config.getHttpClient())
100102
.withClientId(config.getClientId())
101103
.withClientSecret(config.getClientSecret())
102104
.withHost(config.getHost())
103-
.withOpenIDConnectEndpoints(config.getOidcEndpoints())
105+
.withOpenIDConnectEndpoints(config.getDatabricksOidcEndpoints())
104106
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
105107
.withScopes(config.getScopes())
106108
.build();
@@ -129,8 +131,9 @@ void openIDConnectEndPointsTestAccounts() throws IOException {
129131
config.resolve();
130132

131133
String prefix = "https://accounts.cloud.databricks.com/oidc/accounts/" + config.getAccountId();
132-
assertEquals(prefix + "/v1/token", config.getOidcEndpoints().getTokenEndpoint());
133-
assertEquals(prefix + "/v1/authorize", config.getOidcEndpoints().getAuthorizationEndpoint());
134+
assertEquals(prefix + "/v1/token", config.getDatabricksOidcEndpoints().getTokenEndpoint());
135+
assertEquals(
136+
prefix + "/v1/authorize", config.getDatabricksOidcEndpoints().getAuthorizationEndpoint());
134137
}
135138

136139
@Test
@@ -278,7 +281,7 @@ void cacheWithValidRefreshableTokenTest() throws IOException {
278281

279282
// Spy on the config to inject the endpoints.
280283
DatabricksConfig spyConfig = Mockito.spy(config);
281-
Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints();
284+
Mockito.doReturn(endpoints).when(spyConfig).getDatabricksOidcEndpoints();
282285

283286
// Configure provider.
284287
HeaderFactory headerFactory = provider.configure(spyConfig);
@@ -343,7 +346,7 @@ void cacheWithValidNonRefreshableTokenTest() throws IOException {
343346

344347
// Spy on the config to inject the endpoints.
345348
DatabricksConfig spyConfig = Mockito.spy(config);
346-
Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints();
349+
Mockito.doReturn(endpoints).when(spyConfig).getDatabricksOidcEndpoints();
347350

348351
// Configure provider.
349352
HeaderFactory headerFactory = provider.configure(spyConfig);
@@ -415,7 +418,7 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException {
415418

416419
// Spy on the config to inject the endpoints
417420
DatabricksConfig spyConfig = Mockito.spy(config);
418-
Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints();
421+
Mockito.doReturn(endpoints).when(spyConfig).getDatabricksOidcEndpoints();
419422

420423
// Configure provider
421424
HeaderFactory headerFactory = provider.configure(spyConfig);
@@ -524,7 +527,7 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException {
524527

525528
// Spy on the config to inject the endpoints
526529
DatabricksConfig spyConfig = Mockito.spy(config);
527-
Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints();
530+
Mockito.doReturn(endpoints).when(spyConfig).getDatabricksOidcEndpoints();
528531

529532
// Configure provider
530533
HeaderFactory headerFactory = provider.configure(spyConfig);
@@ -610,7 +613,7 @@ void cacheWithInvalidTokensTest() throws IOException {
610613
"https://test.databricks.com/oidc/v1/token",
611614
"https://test.databricks.com/oidc/v1/authorize");
612615
DatabricksConfig spyConfig = Mockito.spy(config);
613-
Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints();
616+
Mockito.doReturn(endpoints).when(spyConfig).getDatabricksOidcEndpoints();
614617

615618
// Configure provider
616619
HeaderFactory headerFactory = provider.configure(spyConfig);
@@ -738,7 +741,7 @@ void externalBrowserAuthWithAzureClientIdTest() throws IOException {
738741
"https://test.azuredatabricks.net/oidc/v1/token",
739742
"https://test.azuredatabricks.net/oidc/v1/authorize");
740743
DatabricksConfig spyConfig = Mockito.spy(config);
741-
Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints();
744+
Mockito.doReturn(endpoints).when(spyConfig).getDatabricksOidcEndpoints();
742745

743746
// Configure provider
744747
HeaderFactory headerFactory = provider.configure(spyConfig);

0 commit comments

Comments
 (0)