Skip to content

Commit c2a8a8c

Browse files
hectorcast-dbclaude
hectorcast-db
and
claude
authored
Undo OIDC endpoint change for External Browser authentication in Azure (#658)
## What changes are proposed in this pull request? Undo OIDC endpoint change for External Browser authentication in Azure. Last PR fixed an issue on Databricks M2M but introduced a behavior change in External Browser authentication, which is not covered by tests. #657 This PR adds back the logic for endpoint selection, but only when External Browser authentication is used. ## How is this tested? N/A NO_CHANGELOG=true --------- Co-authored-by: Claude Sonnet 4.5 <noreply@anthropic.com>
1 parent 2662285 commit c2a8a8c

File tree

3 files changed

+186
-19
lines changed

3 files changed

+186
-19
lines changed

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

Lines changed: 18 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,13 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
5858
// Use the utility class to resolve client ID and client secret
5959
String clientId = OAuthClientUtils.resolveClientId(config);
6060
String clientSecret = OAuthClientUtils.resolveClientSecret(config);
61+
OpenIDConnectEndpoints oidcEndpoints = null;
62+
try {
63+
oidcEndpoints = OAuthClientUtils.resolveOidcEndpoints(config);
64+
} catch (IOException e) {
65+
LOGGER.error("Failed to resolve OIDC endpoints: {}", e.getMessage());
66+
return null;
67+
}
6168

6269
try {
6370
if (tokenCache == null) {
@@ -78,7 +85,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
7885
new SessionCredentialsTokenSource(
7986
cachedToken,
8087
config.getHttpClient(),
81-
config.getOidcEndpoints().getTokenEndpoint(),
88+
oidcEndpoints.getTokenEndpoint(),
8289
clientId,
8390
clientSecret,
8491
Optional.of(config.getEffectiveOAuthRedirectUrl()),
@@ -100,7 +107,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
100107

101108
// If no cached token or refresh failed, perform browser auth
102109
CachedTokenSource cachedTokenSource =
103-
performBrowserAuth(config, clientId, clientSecret, tokenCache);
110+
performBrowserAuth(config, clientId, clientSecret, tokenCache, oidcEndpoints);
104111
tokenCache.save(cachedTokenSource.getToken());
105112
return OAuthHeaderFactory.fromTokenSource(cachedTokenSource);
106113
} catch (IOException | DatabricksException e) {
@@ -109,7 +116,7 @@ public OAuthHeaderFactory configure(DatabricksConfig config) {
109116
}
110117
}
111118

112-
protected List<String> getScopes(DatabricksConfig config) {
119+
protected List<String> getScopes(DatabricksConfig config, OpenIDConnectEndpoints oidcEndpoints) {
113120
// Get user-provided scopes and add required default scopes.
114121
Set<String> scopes = new HashSet<>(config.getScopes());
115122
// Requesting a refresh token is most of the time the right thing to do from a
@@ -125,7 +132,11 @@ protected List<String> getScopes(DatabricksConfig config) {
125132
}
126133

127134
CachedTokenSource performBrowserAuth(
128-
DatabricksConfig config, String clientId, String clientSecret, TokenCache tokenCache)
135+
DatabricksConfig config,
136+
String clientId,
137+
String clientSecret,
138+
TokenCache tokenCache,
139+
OpenIDConnectEndpoints oidcEndpoints)
129140
throws IOException {
130141
LOGGER.debug("Performing browser authentication");
131142

@@ -138,8 +149,8 @@ CachedTokenSource performBrowserAuth(
138149
.withAccountId(config.getAccountId())
139150
.withRedirectUrl(config.getEffectiveOAuthRedirectUrl())
140151
.withBrowserTimeout(config.getOAuthBrowserAuthTimeout())
141-
.withScopes(getScopes(config))
142-
.withOpenIDConnectEndpoints(config.getOidcEndpoints())
152+
.withScopes(getScopes(config, oidcEndpoints))
153+
.withOpenIDConnectEndpoints(oidcEndpoints)
143154
.build();
144155
Consent consent = client.initiateConsent();
145156

@@ -151,7 +162,7 @@ CachedTokenSource performBrowserAuth(
151162
new SessionCredentialsTokenSource(
152163
token,
153164
config.getHttpClient(),
154-
config.getOidcEndpoints().getTokenEndpoint(),
165+
oidcEndpoints.getTokenEndpoint(),
155166
config.getClientId(),
156167
config.getClientSecret(),
157168
Optional.ofNullable(config.getEffectiveOAuthRedirectUrl()),

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

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,9 @@
11
package com.databricks.sdk.core.oauth;
22

33
import com.databricks.sdk.core.DatabricksConfig;
4+
import com.databricks.sdk.core.http.Request;
5+
import com.databricks.sdk.core.http.Response;
6+
import java.io.IOException;
47

58
/** Utility methods for OAuth client credentials resolution. */
69
public class OAuthClientUtils {
@@ -39,4 +42,33 @@ public static String resolveClientSecret(DatabricksConfig config) {
3942
}
4043
return null;
4144
}
45+
46+
/**
47+
* Resolves the OAuth OIDC endpoints from the configuration. Prioritizes regular OIDC endpoints,
48+
* then Azure OIDC endpoints. If the client ID and client secret are provided, the OIDC endpoints
49+
* are fetched from the discovery URL. If the Azure client ID and client secret are provided, the
50+
* OIDC endpoints are fetched from the Azure AD endpoint. If no client ID and client secret are
51+
* provided, the OIDC endpoints are fetched from the default OIDC endpoints.
52+
*
53+
* @param config The Databricks configuration
54+
* @return The resolved OIDC endpoints
55+
* @throws IOException if the OIDC endpoints cannot be resolved
56+
*/
57+
public static OpenIDConnectEndpoints resolveOidcEndpoints(DatabricksConfig config)
58+
throws IOException {
59+
if (config.getClientId() != null && config.getClientSecret() != null) {
60+
return config.getOidcEndpoints();
61+
} else if (config.getAzureClientId() != null && config.getAzureClientSecret() != null) {
62+
Request request = new Request("GET", config.getHost() + "/oidc/oauth2/v2.0/authorize");
63+
request.setRedirectionBehavior(false);
64+
Response resp = config.getHttpClient().execute(request);
65+
String realAuthUrl = resp.getFirstHeader("location");
66+
if (realAuthUrl == null) {
67+
return null;
68+
}
69+
return new OpenIDConnectEndpoints(
70+
realAuthUrl.replaceAll("/authorize", "/token"), realAuthUrl);
71+
}
72+
return config.getOidcEndpoints();
73+
}
4274
}

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

Lines changed: 136 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -303,7 +303,8 @@ void cacheWithValidRefreshableTokenTest() throws IOException {
303303
any(DatabricksConfig.class),
304304
any(String.class),
305305
any(String.class),
306-
any(TokenCache.class));
306+
any(TokenCache.class),
307+
any(OpenIDConnectEndpoints.class));
307308

308309
// Verify token was NOT saved back to cache (we're using the cached one as-is).
309310
Mockito.verify(mockTokenCache, Mockito.never()).save(any(Token.class));
@@ -363,7 +364,12 @@ void cacheWithValidNonRefreshableTokenTest() throws IOException {
363364

364365
// Verify performBrowserAuth was NOT called.
365366
Mockito.verify(provider, Mockito.never())
366-
.performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class));
367+
.performBrowserAuth(
368+
any(DatabricksConfig.class),
369+
any(),
370+
any(),
371+
any(TokenCache.class),
372+
any(OpenIDConnectEndpoints.class));
367373

368374
// Verify no token was saved (we're using the cached one as-is).
369375
Mockito.verify(mockTokenCache, Mockito.never()).save(any(Token.class));
@@ -430,7 +436,8 @@ void cacheWithInvalidAccessTokenValidRefreshTest() throws IOException {
430436
any(DatabricksConfig.class),
431437
any(String.class),
432438
any(String.class),
433-
any(TokenCache.class));
439+
any(TokenCache.class),
440+
any(OpenIDConnectEndpoints.class));
434441

435442
// Verify token was saved back to cache
436443
Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class));
@@ -508,7 +515,12 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException {
508515
Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache));
509516
Mockito.doReturn(cachedTokenSource)
510517
.when(provider)
511-
.performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class));
518+
.performBrowserAuth(
519+
any(DatabricksConfig.class),
520+
any(),
521+
any(),
522+
any(TokenCache.class),
523+
any(OpenIDConnectEndpoints.class));
512524

513525
// Spy on the config to inject the endpoints
514526
DatabricksConfig spyConfig = Mockito.spy(config);
@@ -527,7 +539,12 @@ void cacheWithInvalidAccessTokenRefreshFailingTest() throws IOException {
527539

528540
// Verify performBrowserAuth was called since refresh failed
529541
Mockito.verify(provider, Mockito.times(1))
530-
.performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class));
542+
.performBrowserAuth(
543+
any(DatabricksConfig.class),
544+
any(),
545+
any(),
546+
any(TokenCache.class),
547+
any(OpenIDConnectEndpoints.class));
531548

532549
// Verify token was saved after browser auth (for the new token)
533550
Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class));
@@ -572,17 +589,31 @@ void cacheWithInvalidTokensTest() throws IOException {
572589
new DatabricksConfig()
573590
.setAuthType("external-browser")
574591
.setHost("https://test.databricks.com")
575-
.setClientId("test-client-id");
592+
.setClientId("test-client-id")
593+
.setHttpClient(mockHttpClient);
576594

577595
// Create our provider and mock the browser auth method
578596
ExternalBrowserCredentialsProvider provider =
579597
Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache));
580598
Mockito.doReturn(cachedTokenSource)
581599
.when(provider)
582-
.performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class));
600+
.performBrowserAuth(
601+
any(DatabricksConfig.class),
602+
any(),
603+
any(),
604+
any(TokenCache.class),
605+
any(OpenIDConnectEndpoints.class));
606+
607+
// Spy on the config to inject the endpoints
608+
OpenIDConnectEndpoints endpoints =
609+
new OpenIDConnectEndpoints(
610+
"https://test.databricks.com/oidc/v1/token",
611+
"https://test.databricks.com/oidc/v1/authorize");
612+
DatabricksConfig spyConfig = Mockito.spy(config);
613+
Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints();
583614

584615
// Configure provider
585-
HeaderFactory headerFactory = provider.configure(config);
616+
HeaderFactory headerFactory = provider.configure(spyConfig);
586617
assertNotNull(headerFactory);
587618
// Verify headers contain the browser auth token (fallback)
588619
Map<String, String> headers = headerFactory.headers();
@@ -593,7 +624,12 @@ void cacheWithInvalidTokensTest() throws IOException {
593624

594625
// Verify performBrowserAuth was called since we had an invalid token
595626
Mockito.verify(provider, Mockito.times(1))
596-
.performBrowserAuth(any(DatabricksConfig.class), any(), any(), any(TokenCache.class));
627+
.performBrowserAuth(
628+
any(DatabricksConfig.class),
629+
any(),
630+
any(),
631+
any(TokenCache.class),
632+
any(OpenIDConnectEndpoints.class));
597633

598634
// Verify token was saved after browser auth (for the new token)
599635
Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class));
@@ -609,7 +645,7 @@ void doNotAddOfflineAccessScopeWhenDisableOauthRefreshTokenIsTrue() {
609645
.setScopes(Arrays.asList("my-test-scope"));
610646

611647
ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider();
612-
List<String> scopes = provider.getScopes(config);
648+
List<String> scopes = provider.getScopes(config, null);
613649

614650
assertEquals(1, scopes.size());
615651
assertTrue(scopes.contains("my-test-scope"));
@@ -625,7 +661,7 @@ void doNotRemoveUserProvidedScopesWhenDisableOauthRefreshTokenIsTrue() {
625661
.setScopes(Arrays.asList("my-test-scope", "offline_access"));
626662

627663
ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider();
628-
List<String> scopes = provider.getScopes(config);
664+
List<String> scopes = provider.getScopes(config, null);
629665

630666
assertEquals(2, scopes.size());
631667
assertTrue(scopes.contains("offline_access"));
@@ -641,10 +677,98 @@ void addOfflineAccessScopeWhenDisableOauthRefreshTokenIsFalse() {
641677
.setScopes(Arrays.asList("my-test-scope"));
642678

643679
ExternalBrowserCredentialsProvider provider = new ExternalBrowserCredentialsProvider();
644-
List<String> scopes = provider.getScopes(config);
680+
List<String> scopes = provider.getScopes(config, null);
645681

646682
assertEquals(2, scopes.size());
647683
assertTrue(scopes.contains("offline_access"));
648684
assertTrue(scopes.contains("my-test-scope"));
649685
}
686+
687+
@Test
688+
void externalBrowserAuthWithAzureClientIdTest() throws IOException {
689+
// Create mock HTTP client
690+
HttpClient mockHttpClient = Mockito.mock(HttpClient.class);
691+
692+
// Mock token cache
693+
TokenCache mockTokenCache = Mockito.mock(TokenCache.class);
694+
Mockito.doReturn(null).when(mockTokenCache).load();
695+
696+
// Create valid token for browser auth
697+
Token browserAuthToken =
698+
new Token(
699+
"azure_access_token", "Bearer", "azure_refresh_token", Instant.now().plusSeconds(3600));
700+
701+
// Create token source
702+
SessionCredentialsTokenSource browserAuthTokenSource =
703+
new SessionCredentialsTokenSource(
704+
browserAuthToken,
705+
mockHttpClient,
706+
"https://test.azuredatabricks.net/oidc/v1/token",
707+
"test-azure-client-id",
708+
null,
709+
Optional.empty(),
710+
Optional.empty());
711+
712+
CachedTokenSource cachedTokenSource =
713+
new CachedTokenSource.Builder(browserAuthTokenSource).setToken(browserAuthToken).build();
714+
715+
// Create Azure config with Azure client ID
716+
DatabricksConfig config =
717+
new DatabricksConfig()
718+
.setAuthType("external-browser")
719+
.setHost("https://test.azuredatabricks.net")
720+
.setAzureClientId("test-azure-client-id")
721+
.setHttpClient(mockHttpClient);
722+
723+
// Create provider and mock browser auth
724+
ExternalBrowserCredentialsProvider provider =
725+
Mockito.spy(new ExternalBrowserCredentialsProvider(mockTokenCache));
726+
Mockito.doReturn(cachedTokenSource)
727+
.when(provider)
728+
.performBrowserAuth(
729+
any(DatabricksConfig.class),
730+
any(),
731+
any(),
732+
any(TokenCache.class),
733+
any(OpenIDConnectEndpoints.class));
734+
735+
// Spy on config to inject OIDC endpoints
736+
OpenIDConnectEndpoints endpoints =
737+
new OpenIDConnectEndpoints(
738+
"https://test.azuredatabricks.net/oidc/v1/token",
739+
"https://test.azuredatabricks.net/oidc/v1/authorize");
740+
DatabricksConfig spyConfig = Mockito.spy(config);
741+
Mockito.doReturn(endpoints).when(spyConfig).getOidcEndpoints();
742+
743+
// Configure provider
744+
HeaderFactory headerFactory = provider.configure(spyConfig);
745+
assertNotNull(headerFactory);
746+
747+
// Verify headers contain the Azure token
748+
Map<String, String> headers = headerFactory.headers();
749+
assertEquals("Bearer azure_access_token", headers.get("Authorization"));
750+
751+
// Capture and verify the OpenIDConnectEndpoints passed to performBrowserAuth
752+
ArgumentCaptor<OpenIDConnectEndpoints> endpointsCaptor =
753+
ArgumentCaptor.forClass(OpenIDConnectEndpoints.class);
754+
Mockito.verify(provider, Mockito.times(1))
755+
.performBrowserAuth(
756+
any(DatabricksConfig.class),
757+
any(),
758+
any(),
759+
any(TokenCache.class),
760+
endpointsCaptor.capture());
761+
762+
// Verify the captured endpoints match what we expect for Azure
763+
OpenIDConnectEndpoints capturedEndpoints = endpointsCaptor.getValue();
764+
assertNotNull(capturedEndpoints);
765+
assertEquals(
766+
"https://test.azuredatabricks.net/oidc/v1/token", capturedEndpoints.getTokenEndpoint());
767+
assertEquals(
768+
"https://test.azuredatabricks.net/oidc/v1/authorize",
769+
capturedEndpoints.getAuthorizationEndpoint());
770+
771+
// Verify token was saved
772+
Mockito.verify(mockTokenCache, Mockito.times(1)).save(any(Token.class));
773+
}
650774
}

0 commit comments

Comments
 (0)