[BUG] Support for Token Exchange in OAuth flows causes NullPointerException

[agarwalvibhor1914]

Description-
Below support was added in version v1.0.6-oss-
Support for Token Exchange in OAuth flows where in third party tokens are exchanged for InHouse tokens.

We are using external token but when code is trying to exchange it with databricks token it causes NullPointerException in the logs. Because we dont have(and want) configurations in databricks to issue inHouse tokens.
Although driver continues with external token in the next step.

Cause-
DatabricksTokenFederationProvider -> referesh()
if (!isSameHost(claims.getIssuer(), this.config.getHost())) {
optionalToken = tryTokenExchange(accessToken, accessTokenType);
}
Optional tryTokenExchange(String accessToken, String accessTokenType) {
LOGGER.debug(
"Token tryTokenExchange(String accessToken, String accessTokenType = {})", accessTokenType);
try {
return Optional.of(exchangeToken(accessToken));
} catch (Exception e) {
LOGGER.error(e, "Token exchange failed, falling back to using external token");
return Optional.empty();
}
}

Because exchange with in-house token is not applicable in our case, exchangeToken throws exception.

java.lang.NullPointerException: Cannot invoke "String.trim()" because "s" is null
at com.databricks.internal.nimbusds.jose.JOSEObject.split(JOSEObject.java:218)
at com.databricks.internal.nimbusds.jwt.SignedJWT.parse(SignedJWT.java:129)
at com.databricks.jdbc.auth.DatabricksTokenFederationProvider.createToken(DatabricksTokenFederationProvider.java:166)
at com.databricks.jdbc.auth.DatabricksTokenFederationProvider.retrieveToken(DatabricksTokenFederationProvider.java:213)
at com.databricks.jdbc.auth.DatabricksTokenFederationProvider.exchangeToken(DatabricksTokenFederationProvider.java:191)
at com.databricks.jdbc.auth.DatabricksTokenFederationProvider.tryTokenExchange(DatabricksTokenFederationProvider.java:157)
at com.databricks.jdbc.auth.DatabricksTokenFederationProvider.refresh(DatabricksTokenFederationProvider.java:139)
at com.databricks.internal.sdk.core.oauth.RefreshableTokenSource.getToken(RefreshableTokenSource.java:84)
at com.databricks.jdbc.auth.DatabricksTokenFederationProvider.lambda$configure$0(DatabricksTokenFederationProvider.java:118)
at com.databricks.internal.sdk.core.DatabricksConfig.authenticate(DatabricksConfig.java:207)
at com.databricks.jdbc.dbclient.impl.thrift.DatabricksHttpTTransport.refreshHeadersIfRequired(DatabricksHttpTTransport.java:166)
at com.databricks.jdbc.dbclient.impl.thrift.DatabricksHttpTTransport.flush(DatabricksHttpTTransport.java:96)
at com.databricks.internal.apache.thrift.TServiceClient.sendBase(TServiceClient.java:75)
at com.databricks.internal.apache.thrift.TServiceClient.sendBase(TServiceClient.java:64)
at com.databricks.jdbc.model.client.thrift.generated.TCLIService$Client.send_OpenSession(TCLIService.java:213)
at com.databricks.jdbc.model.client.thrift.generated.TCLIService$Client.OpenSession(TCLIService.java:206)
at com.databricks.jdbc.dbclient.impl.thrift.DatabricksThriftAccessor.getThriftResponse(DatabricksThriftAccessor.java:103)
at com.databricks.jdbc.dbclient.impl.thrift.DatabricksThriftServiceClient.createSession(DatabricksThriftServiceClient.java:97)
at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke0(Native Method)
at java.base/jdk.internal.reflect.NativeMethodAccessorImpl.invoke(NativeMethodAccessorImpl.java:77)
at java.base/jdk.internal.reflect.DelegatingMethodAccessorImpl.invoke(DelegatingMethodAccessorImpl.java:43)
at java.base/java.lang.reflect.Method.invoke(Method.java:569)
at com.databricks.jdbc.telemetry.latency.DatabricksMetricsTimedProcessor$TimedInvocationHandler.invoke(DatabricksMetricsTimedProcessor.java:57)
at jdk.proxy2/jdk.proxy2.$Proxy9.createSession(Unknown Source)
at com.databricks.jdbc.api.impl.DatabricksSession.open(DatabricksSession.java:140)
at com.databricks.jdbc.api.impl.DatabricksConnection.open(DatabricksConnection.java:67)
at com.databricks.client.jdbc.Driver.connect(Driver.java:65)

To Reproduce
properties.put("Auth_AccessToken", authentication.accessToken());
properties.put("AuthMech", "11");
properties.put("Auth_Flow", "0");

var connection = driver.connect(url, properties);

Expected behaviour
When token can not be exchanged, no error logs should be produced. Driver should use external token silently.

Environment

• Java 21
• Driver Version 1.0.6-oss onwards

[jprakash-db]

@agarwalvibhor1914 I have a few questions

• Are you using JWT? It looks like it is JWT since you are reaching the token exchange step
• Can you share more granular logs complete logs, we have error logs for multiple steps so easier to pin point what the failure is
  Also the failure is a bit weird because every error within token exchange is self contained and does not propagate or hinder the application -

  databricks-jdbc/src/main/java/com/databricks/jdbc/auth/DatabricksTokenFederationProvider.java

  Line 157 in ff2751a

  return Optional.of(exchangeToken(accessToken));

I am unable to repro this error on my side , can you confirm is this something you are facing even in the latest versions of the oss driver such as 1.0.10-oss

[agarwalvibhor1914]

yes its JWT, and tried till version 1.0.9-oss.

When failed to get exchange token, then NPE is thrown over here, because accessToken is null, SignedJWT.parse(accessToken)-> JOSEObject.split(s)-> String t = s.trim().

databricks-jdbc/src/main/java/com/databricks/jdbc/auth/DatabricksTokenFederationProvider.java

Line 166 in ff2751a

SignedJWT signedJWT = SignedJWT.parse(accessToken);

This NPE gets a catch over here-

databricks-jdbc/src/main/java/com/databricks/jdbc/auth/DatabricksTokenFederationProvider.java

Line 159 in ff2751a

LOGGER.error(e, "Token exchange failed, falling back to using external token");

This is the error message and stacktrace we see in our logs, polluting log files and monitoring tools. Previously provided stack trace is complete.

As I mentioned before, the driver continue with the external JWT token.

[jprakash-db]

@agarwalvibhor1914 So the expectation is to be able to disable token federation all together , so that these logs don't show up in the first place. Will raise a PR to provide this capability

[agarwalvibhor1914]

@jprakash-db any planning for this fix of this issue?