Skip to content

Commit 78501a5

Browse files
emmyzhou-dbparthban-db
emmyzhou-db
and
parthban-db
authored
Add support for OIDC ID token authentication using an environment variables and files (#445)
## What changes are proposed in this pull request? Add Environment Variable and File-based ID Token Sources ## Key Changes - Added `EnvVarIDTokenSource` to read ID tokens from environment variables - Added `FileIDTokenSource` to read ID tokens from files - Both implement the `IDTokenSource` interface for OIDC authentication - Added envvar-oidc and file-oidc to the list of auth types ## How is this tested? - Unit tests for both classes NO_CHANGELOG=true --------- Co-authored-by: Parth Bansal <[email protected]>
1 parent 8aad2be commit 78501a5

File tree

7 files changed

+375
-0
lines changed

7 files changed

+375
-0
lines changed

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,14 @@ public class DatabricksConfig {
148148
@ConfigAttribute(env = "TOKEN_AUDIENCE")
149149
private String tokenAudience;
150150

151+
/** Path to the file containing an OIDC ID token. */
152+
@ConfigAttribute(env = "DATABRICKS_OIDC_TOKEN_FILEPATH", auth = "file-oidc")
153+
private String oidcTokenFilepath;
154+
155+
/** Environment variable name that contains an OIDC ID token. */
156+
@ConfigAttribute(env = "DATABRICKS_OIDC_TOKEN_ENV", auth = "env-oidc")
157+
private String oidcTokenEnv;
158+
151159
public Environment getEnv() {
152160
return env;
153161
}
@@ -528,6 +536,24 @@ public DatabricksConfig setTokenAudience(String tokenAudience) {
528536
return this;
529537
}
530538

539+
public String getOidcTokenFilepath() {
540+
return oidcTokenFilepath;
541+
}
542+
543+
public DatabricksConfig setOidcTokenFilepath(String oidcTokenFilepath) {
544+
this.oidcTokenFilepath = oidcTokenFilepath;
545+
return this;
546+
}
547+
548+
public String getOidcTokenEnv() {
549+
return oidcTokenEnv;
550+
}
551+
552+
public DatabricksConfig setOidcTokenEnv(String oidcTokenEnv) {
553+
this.oidcTokenEnv = oidcTokenEnv;
554+
return this;
555+
}
556+
531557
public boolean isAzure() {
532558
if (azureWorkspaceResourceId != null) {
533559
return true;

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

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

33
import com.databricks.sdk.core.oauth.*;
4+
import com.google.common.base.Strings;
45
import java.util.ArrayList;
56
import java.util.List;
67
import org.slf4j.Logger;
@@ -111,6 +112,20 @@ private void addOIDCCredentialsProviders(DatabricksConfig config) {
111112
}
112113

113114
List<NamedIDTokenSource> namedIdTokenSources = new ArrayList<>();
115+
namedIdTokenSources.add(
116+
new NamedIDTokenSource(
117+
"env-oidc",
118+
new EnvVarIDTokenSource(
119+
// Use configured environment variable name if set, otherwise default to
120+
// DATABRICKS_OIDC_TOKEN
121+
Strings.isNullOrEmpty(config.getOidcTokenEnv())
122+
? "DATABRICKS_OIDC_TOKEN"
123+
: config.getOidcTokenEnv(),
124+
config.getEnv())));
125+
126+
namedIdTokenSources.add(
127+
new NamedIDTokenSource("file-oidc", new FileIDTokenSource(config.getOidcTokenFilepath())));
128+
114129
namedIdTokenSources.add(
115130
new NamedIDTokenSource(
116131
"github-oidc",

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

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
package com.databricks.sdk.core.oauth;
2+
3+
import com.databricks.sdk.core.DatabricksException;
4+
import com.databricks.sdk.core.utils.Environment;
5+
import com.google.common.base.Strings;
6+
7+
/** Implementation of {@link IDTokenSource} that reads the ID token from an environment variable. */
8+
public class EnvVarIDTokenSource implements IDTokenSource {
9+
/* The name of the environment variable to read the ID token from. */
10+
private final String envVarName;
11+
/* The environment to read variables from. */
12+
private final Environment env;
13+
14+
/**
15+
* Creates a new EnvVarIDTokenSource that reads from the specified environment variable.
16+
*
17+
* @param envVarName The name of the environment variable to read the ID token from.
18+
* @param env The environment to read variables from.
19+
*/
20+
public EnvVarIDTokenSource(String envVarName, Environment env) {
21+
this.envVarName = envVarName;
22+
this.env = env;
23+
}
24+
25+
/**
26+
* Retrieves an ID Token from the environment variable.
27+
*
28+
* @param audience The intended recipient of the ID Token (unused in this implementation).
29+
* @return An {@link IDToken} containing the token value from the environment variable.
30+
* @throws IllegalArgumentException if the environment variable name is null or empty.
31+
* @throws DatabricksException if the environment variable is not set or is empty.
32+
*/
33+
@Override
34+
public IDToken getIDToken(String audience) {
35+
if (Strings.isNullOrEmpty(envVarName)) {
36+
throw new IllegalArgumentException("Environment variable name cannot be null or empty");
37+
}
38+
39+
try {
40+
String token = env.get(envVarName);
41+
return new IDToken(token);
42+
} catch (IllegalArgumentException e) {
43+
throw new DatabricksException(
44+
String.format("Received empty ID token from environment variable %s", envVarName), e);
45+
}
46+
}
47+
}

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

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
package com.databricks.sdk.core.oauth;
2+
3+
import com.databricks.sdk.core.DatabricksException;
4+
import com.google.common.base.Strings;
5+
import java.io.IOException;
6+
import java.nio.charset.StandardCharsets;
7+
import java.nio.file.Files;
8+
import java.nio.file.InvalidPathException;
9+
import java.nio.file.Path;
10+
import java.nio.file.Paths;
11+
import java.util.List;
12+
import java.util.stream.Collectors;
13+
14+
/**
15+
* Implementation of {@link IDTokenSource} that reads the ID token from a file. The token is read
16+
* using UTF-8 encoding and any leading/trailing whitespace is trimmed. The file should contain
17+
* exactly one non-empty line with the token value. Files with multiple non-empty lines or only
18+
* empty lines will result in an error.
19+
*
20+
* @see IDTokenSource
21+
*/
22+
public class FileIDTokenSource implements IDTokenSource {
23+
/* The path to the file containing the ID token. */
24+
private final String filePath;
25+
26+
/**
27+
* Creates a new FileIDTokenSource that reads from the specified file.
28+
*
29+
* @param filePath Path to the file containing the ID token. The file should contain a single line
30+
* with the token value.
31+
* @throws IllegalArgumentException if the file path is null or empty.
32+
*/
33+
public FileIDTokenSource(String filePath) {
34+
this.filePath = filePath;
35+
}
36+
37+
/**
38+
* Retrieves an ID Token from the file. The file is read using UTF-8 encoding and the first line
39+
* is used as the token value. Any leading or trailing whitespace in the token is trimmed.
40+
*
41+
* @param audience The intended recipient of the ID Token. This parameter is not used in this
42+
* implementation as the token is read directly from the file.
43+
* @return An {@link IDToken} containing the token value from the file.
44+
* @throws IllegalArgumentException if the file path is null or empty.
45+
* @throws DatabricksException in the following cases:
46+
* <ul>
47+
* <li>If the file path is invalid or malformed
48+
* <li>If the file does not exist
49+
* <li>If there are security permission issues accessing the file
50+
* <li>If the file is empty or contains only whitespace
51+
* <li>If the file cannot be read due to I/O errors
52+
* <li>If the token format in the file is invalid
53+
* </ul>
54+
*/
55+
@Override
56+
public IDToken getIDToken(String audience) {
57+
if (Strings.isNullOrEmpty(filePath)) {
58+
throw new IllegalArgumentException("File path cannot be null or empty");
59+
}
60+
61+
Path path;
62+
try {
63+
path = Paths.get(filePath);
64+
} catch (InvalidPathException e) {
65+
throw new DatabricksException("Invalid file path: " + filePath, e);
66+
}
67+
68+
boolean isFileExists;
69+
try {
70+
isFileExists = Files.exists(path);
71+
} catch (SecurityException e) {
72+
throw new DatabricksException(
73+
String.format(
74+
"Security permission denied when checking if file %s exists: %s",
75+
filePath, e.getMessage()),
76+
e);
77+
}
78+
79+
if (!isFileExists) {
80+
throw new DatabricksException(String.format("File %s does not exist", filePath));
81+
}
82+
83+
List<String> rawLines;
84+
try {
85+
rawLines = Files.readAllLines(path, StandardCharsets.UTF_8);
86+
} catch (IOException e) {
87+
throw new DatabricksException(
88+
String.format("Failed to read ID token from file %s: %s", filePath, e.getMessage()), e);
89+
} catch (SecurityException e) {
90+
throw new DatabricksException(
91+
String.format(
92+
"Security permission denied when reading file %s: %s", filePath, e.getMessage()),
93+
e);
94+
}
95+
96+
// Filter out empty lines
97+
List<String> nonEmptyLines =
98+
rawLines.stream()
99+
.map(String::trim)
100+
.filter(line -> !line.isEmpty())
101+
.collect(Collectors.toList());
102+
103+
if (nonEmptyLines.isEmpty()) {
104+
throw new DatabricksException(String.format("File %s contains only empty lines", filePath));
105+
}
106+
107+
if (nonEmptyLines.size() > 1) {
108+
throw new DatabricksException(
109+
String.format(
110+
"The token should be a single line but the file %s contains %d non-empty lines",
111+
filePath, nonEmptyLines.size()));
112+
}
113+
114+
String token = nonEmptyLines.get(0);
115+
116+
try {
117+
return new IDToken(token);
118+
} catch (IllegalArgumentException e) {
119+
throw new DatabricksException(
120+
String.format("Received empty ID token from file %s", filePath));
121+
}
122+
}
123+
}

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ public class IDToken {
1212
* Constructs an IDToken with a value.
1313
*
1414
* @param value The ID Token string.
15+
* @throws IllegalArgumentException if the token value is null or empty.
1516
*/
1617
public IDToken(String value) {
1718
if (value == null || value.isEmpty()) {

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

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
package com.databricks.sdk.core.oauth;
2+
3+
import static org.junit.jupiter.api.Assertions.*;
4+
5+
import com.databricks.sdk.core.DatabricksException;
6+
import com.databricks.sdk.core.utils.Environment;
7+
import java.util.HashMap;
8+
import java.util.Map;
9+
import java.util.stream.Stream;
10+
import org.junit.jupiter.params.ParameterizedTest;
11+
import org.junit.jupiter.params.provider.Arguments;
12+
import org.junit.jupiter.params.provider.MethodSource;
13+
14+
/** Tests for EnvVarIDTokenSource. */
15+
public class EnvVarIDTokenSourceTest {
16+
private static final String TEST_ENV_VAR_NAME = "TEST_ID_TOKEN";
17+
private static final String TEST_TOKEN = "test-id-token";
18+
private static final String TEST_AUDIENCE = "test-audience";
19+
20+
private Environment createTestEnvironment(Map<String, String> envVars) {
21+
return new Environment(envVars, new String[0], "test");
22+
}
23+
24+
private static Stream<Arguments> provideTestCases() {
25+
return Stream.of(
26+
// Test case: Success case
27+
Arguments.of(
28+
"Success case",
29+
TEST_ENV_VAR_NAME,
30+
createEnvVars(TEST_ENV_VAR_NAME, TEST_TOKEN),
31+
TEST_TOKEN,
32+
null),
33+
// Test case: Null environment variable name
34+
Arguments.of(
35+
"Null environment variable name",
36+
null,
37+
new HashMap<>(),
38+
null,
39+
IllegalArgumentException.class),
40+
// Test case: Empty environment variable name
41+
Arguments.of(
42+
"Empty environment variable name",
43+
"",
44+
new HashMap<>(),
45+
null,
46+
IllegalArgumentException.class),
47+
// Test case: Missing environment variable
48+
Arguments.of(
49+
"Missing environment variable",
50+
TEST_ENV_VAR_NAME,
51+
new HashMap<>(),
52+
null,
53+
DatabricksException.class),
54+
// Test case: Empty token value
55+
Arguments.of(
56+
"Empty token value",
57+
TEST_ENV_VAR_NAME,
58+
createEnvVars(TEST_ENV_VAR_NAME, ""),
59+
null,
60+
DatabricksException.class));
61+
}
62+
63+
private static Map<String, String> createEnvVars(String key, String value) {
64+
Map<String, String> envVars = new HashMap<>();
65+
envVars.put(key, value);
66+
return envVars;
67+
}
68+
69+
@ParameterizedTest(name = "{0}")
70+
@MethodSource("provideTestCases")
71+
void testGetIDToken(
72+
String testName,
73+
String envVarName,
74+
Map<String, String> envVars,
75+
String expectedToken,
76+
Class<? extends Exception> expectedException) {
77+
Environment env = envVars != null ? createTestEnvironment(envVars) : null;
78+
EnvVarIDTokenSource source = new EnvVarIDTokenSource(envVarName, env);
79+
80+
if (expectedException != null) {
81+
assertThrows(expectedException, () -> source.getIDToken(TEST_AUDIENCE));
82+
} else {
83+
IDToken token = source.getIDToken(TEST_AUDIENCE);
84+
assertNotNull(token);
85+
assertEquals(expectedToken, token.getValue());
86+
}
87+
}
88+
}

0 commit comments

Comments
 (0)