Skip to content

Commit aa748cc

Browse files
sfc-gh-rkowalskisfc-gh-rkowalski
authored
Merge branch 'master' into SNOW-1950032-bugfix
2 parents 6df4923 + cabd6c3 commit aa748cc

36 files changed

Lines changed: 1196 additions & 123 deletions

‎Jenkinsfile‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ timestamps {
5656
string(name: 'parent_build_number', value: env.BUILD_NUMBER),
5757
string(name: 'timeout_value', value: '420'),
5858
string(name: 'PR_Key', value: scmInfo.GIT_BRANCH.substring(3)),
59-
string(name: 'svn_revision', value: 'bptp-built')
59+
string(name: 'svn_revision', value: 'bptp-stable')
6060
]]
6161
}
6262

‎src/main/java/net/snowflake/client/core/SFLoginInput.java‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ public class SFLoginInput {
6262
// Workload Identity Federation
6363
private String workloadIdentityProvider;
6464
private WorkloadIdentityAttestation workloadIdentityAttestation;
65+
private String workloadIdentityEntraResource;
6566

6667
// OAuth
6768
private int redirectUriPort = -1;
@@ -612,4 +613,13 @@ public void setWorkloadIdentityAttestation(
612613
public WorkloadIdentityAttestation getWorkloadIdentityAttestation() {
613614
return workloadIdentityAttestation;
614615
}
616+
617+
public String getWorkloadIdentityEntraResource() {
618+
return this.workloadIdentityEntraResource;
619+
}
620+
621+
public SFLoginInput setWorkloadIdentityEntraResource(String workloadIdentityEntraResource) {
622+
this.workloadIdentityEntraResource = workloadIdentityEntraResource;
623+
return this;
624+
}
615625
}

‎src/main/java/net/snowflake/client/core/SFSession.java‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -723,6 +723,9 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
723723
.setOauthLoginInput(oauthLoginInput)
724724
.setWorkloadIdentityProvider(
725725
(String) connectionPropertiesMap.get(SFSessionProperty.WORKLOAD_IDENTITY_PROVIDER))
726+
.setWorkloadIdentityEntraResource(
727+
(String)
728+
connectionPropertiesMap.get(SFSessionProperty.WORKLOAD_IDENTITY_ENTRA_RESOURCE))
726729
.setPrivateKeyBase64(
727730
(String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_BASE64))
728731
.setPrivateKeyPwd(

‎src/main/java/net/snowflake/client/core/SFSessionProperty.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@ public enum SFSessionProperty {
2727
OAUTH_AUTHORIZATION_URL("oauthAuthorizationUrl", false, String.class),
2828
OAUTH_TOKEN_REQUEST_URL("oauthTokenRequestUrl", false, String.class),
2929
WORKLOAD_IDENTITY_PROVIDER("workloadIdentityProvider", false, String.class),
30+
WORKLOAD_IDENTITY_ENTRA_RESOURCE("workloadIdentityEntraResource", false, String.class),
3031
WAREHOUSE("warehouse", false, String.class),
3132
LOGIN_TIMEOUT("loginTimeout", false, Integer.class),
3233
NETWORK_TIMEOUT("networkTimeout", false, Integer.class),

‎src/main/java/net/snowflake/client/core/SFTrustManager.java‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -165,8 +165,8 @@ public class SFTrustManager extends X509ExtendedTrustManager {
165165
private static final int DEFAULT_OCSP_RESPONDER_CONNECTION_TIMEOUT = 10000;
166166
/** Default OCSP Cache server host name prefix */
167167
private static final String DEFAULT_OCSP_CACHE_HOST_PREFIX = "http://ocsp.snowflakecomputing.";
168-
/** Default OCSP Cache server host name */
169-
private static final String DEFAULT_OCSP_CACHE_HOST = DEFAULT_OCSP_CACHE_HOST_PREFIX + "com";
168+
/** Default domain for OCSP cache host */
169+
private static final String DEFAULT_OCSP_CACHE_HOST_DOMAIN = "com";
170170

171171
/** OCSP response file cache directory */
172172
private static final FileCacheManager fileCacheManager;
@@ -329,7 +329,7 @@ static void resetOCSPResponseCacherServerURL(String ocspCacheServerUrl) throws I
329329
}
330330
}
331331

332-
private static void setOCSPResponseCacheServerURL(String topLevelDomain) {
332+
static void setOCSPResponseCacheServerURL(String serverURL) {
333333
String ocspCacheUrl = systemGetProperty(SF_OCSP_RESPONSE_CACHE_SERVER_URL);
334334
if (ocspCacheUrl != null) {
335335
SF_OCSP_RESPONSE_CACHE_SERVER_URL_VALUE = ocspCacheUrl;
@@ -345,6 +345,14 @@ private static void setOCSPResponseCacheServerURL(String topLevelDomain) {
345345
true);
346346
}
347347
if (SF_OCSP_RESPONSE_CACHE_SERVER_URL_VALUE == null) {
348+
String topLevelDomain = DEFAULT_OCSP_CACHE_HOST_DOMAIN;
349+
try {
350+
URL url = new URL(serverURL);
351+
int domainIndex = url.getHost().lastIndexOf(".") + 1;
352+
topLevelDomain = url.getHost().substring(domainIndex);
353+
} catch (Exception e) {
354+
logger.debug("Exception while setting top level domain (for OCSP)", e);
355+
}
348356
SF_OCSP_RESPONSE_CACHE_SERVER_URL_VALUE =
349357
String.format("%s%s/%s", DEFAULT_OCSP_CACHE_HOST_PREFIX, topLevelDomain, CACHE_FILE_NAME);
350358
}
@@ -772,8 +780,6 @@ void validateRevocationStatus(X509Certificate[] chain, String peerHost)
772780
ocspCacheServer.resetOCSPResponseCacheServer(peerHost);
773781
}
774782

775-
String topLevelDomain = peerHost.substring(peerHost.lastIndexOf(".") + 1);
776-
setOCSPResponseCacheServerURL(topLevelDomain);
777783
boolean isCached = isCached(pairIssuerSubjectList);
778784
if (useOCSPResponseCacheServer() && !isCached) {
779785
if (!ocspCacheServer.new_endpoint_enabled) {

‎src/main/java/net/snowflake/client/core/SessionUtil.java‎

Lines changed: 24 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package net.snowflake.client.core;
22

33
import static net.snowflake.client.core.SFTrustManager.resetOCSPResponseCacherServerURL;
4+
import static net.snowflake.client.core.SFTrustManager.setOCSPResponseCacheServerURL;
45
import static net.snowflake.client.jdbc.SnowflakeUtil.isNullOrEmpty;
56
import static net.snowflake.client.jdbc.SnowflakeUtil.systemGetEnv;
67
import static net.snowflake.client.jdbc.SnowflakeUtil.systemGetProperty;
@@ -31,8 +32,9 @@
3132
import net.snowflake.client.core.auth.oauth.OAuthAccessTokenForRefreshTokenProvider;
3233
import net.snowflake.client.core.auth.oauth.OAuthAccessTokenProviderFactory;
3334
import net.snowflake.client.core.auth.oauth.TokenResponseDTO;
34-
import net.snowflake.client.core.auth.wif.AWSAttestationService;
35+
import net.snowflake.client.core.auth.wif.AwsAttestationService;
3536
import net.snowflake.client.core.auth.wif.AwsIdentityAttestationCreator;
37+
import net.snowflake.client.core.auth.wif.AzureAttestationService;
3638
import net.snowflake.client.core.auth.wif.AzureIdentityAttestationCreator;
3739
import net.snowflake.client.core.auth.wif.GcpIdentityAttestationCreator;
3840
import net.snowflake.client.core.auth.wif.OidcIdentityAttestationCreator;
@@ -339,19 +341,12 @@ static SFLoginOutput openSession(
339341
}
340342

341343
if (authenticator.equals(AuthenticatorType.WORKLOAD_IDENTITY)) {
342-
WorkloadIdentityAttestationProvider attestationProvider =
343-
new WorkloadIdentityAttestationProvider(
344-
new AwsIdentityAttestationCreator(new AWSAttestationService()),
345-
new GcpIdentityAttestationCreator(loginInput),
346-
new AzureIdentityAttestationCreator(),
347-
new OidcIdentityAttestationCreator());
348-
WorkloadIdentityAttestation attestation =
349-
attestationProvider.getAttestation(loginInput.getWorkloadIdentityProvider());
344+
WorkloadIdentityAttestation attestation = getWorkloadIdentityAttestation(loginInput);
350345
if (attestation != null) {
351346
loginInput.setWorkloadIdentityAttestation(attestation);
352347
} else {
353348
throw new SFException(
354-
ErrorCode.WORKFLOW_IDENTITY_FLOW_ERROR,
349+
ErrorCode.WORKLOAD_IDENTITY_FLOW_ERROR,
355350
"Unable to obtain workload identity attestation. Make sure that correct workload identity provider has been set and that Snowflake-JDBC driver runs on supported environment.");
356351
}
357352
}
@@ -378,6 +373,17 @@ static SFLoginOutput openSession(
378373
}
379374
}
380375

376+
private static WorkloadIdentityAttestation getWorkloadIdentityAttestation(SFLoginInput loginInput)
377+
throws SFException {
378+
WorkloadIdentityAttestationProvider attestationProvider =
379+
new WorkloadIdentityAttestationProvider(
380+
new AwsIdentityAttestationCreator(new AwsAttestationService()),
381+
new GcpIdentityAttestationCreator(loginInput),
382+
new AzureIdentityAttestationCreator(new AzureAttestationService(), loginInput),
383+
new OidcIdentityAttestationCreator(loginInput.getToken()));
384+
return attestationProvider.getAttestation(loginInput.getWorkloadIdentityProvider());
385+
}
386+
381387
static void checkIfExperimentalAuthnEnabled(AuthenticatorType authenticator) throws SFException {
382388
if (authenticator.equals(AuthenticatorType.PROGRAMMATIC_ACCESS_TOKEN)
383389
|| authenticator.equals(AuthenticatorType.OAUTH_CLIENT_CREDENTIALS)
@@ -510,6 +516,13 @@ static SFLoginOutput newSession(
510516
Map<SFSessionProperty, Object> connectionPropertiesMap,
511517
String tracingLevel)
512518
throws SFException, SnowflakeSQLException {
519+
try {
520+
// Adjust OCSP cache server if it is private link
521+
resetOCSPUrlIfNecessary(loginInput.getServerUrl());
522+
} catch (IOException ex) {
523+
throw new SFException(ex, ErrorCode.IO_ERROR, "unexpected URL syntax exception");
524+
}
525+
513526
Stopwatch stopwatch = new Stopwatch();
514527
stopwatch.start();
515528
// build URL for login request
@@ -608,13 +621,6 @@ static SFLoginOutput newSession(
608621
throw new SFException(ex, ErrorCode.INTERNAL_ERROR, "unexpected URI syntax exception:1");
609622
}
610623

611-
try {
612-
// Adjust OCSP cache server if it is private link
613-
resetOCSPUrlIfNecessary(loginInput.getServerUrl());
614-
} catch (IOException ex) {
615-
throw new SFException(ex, ErrorCode.IO_ERROR, "unexpected URL syntax exception");
616-
}
617-
618624
HttpPost postRequest = null;
619625

620626
try {
@@ -1927,6 +1933,7 @@ enum TokenRequestType {
19271933
* @throws IOException If exception encountered
19281934
*/
19291935
public static void resetOCSPUrlIfNecessary(String serverUrl) throws IOException {
1936+
setOCSPResponseCacheServerURL(serverUrl);
19301937
if (PrivateLinkDetector.isPrivateLink(serverUrl)) {
19311938
// Privatelink uses special OCSP Cache server
19321939
URL url = new URL(serverUrl);

‎src/main/java/net/snowflake/client/core/auth/AuthenticatorType.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ public enum AuthenticatorType {
5656
PROGRAMMATIC_ACCESS_TOKEN,
5757

5858
/*
59-
* Authenticator to support existing authentication by existing AWS/GCP/Azure workload identity
59+
* Authenticator to support existing authentication by existing AWS/GCP/Azure/OIDC workload identity
6060
*/
6161
WORKLOAD_IDENTITY
6262
}

src/main/java/net/snowflake/client/core/auth/wif/AWSAttestationService.java renamed to src/main/java/net/snowflake/client/core/auth/wif/AwsAttestationService.java

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,17 +12,21 @@
1212
import net.snowflake.client.core.SnowflakeJdbcInternalApi;
1313
import net.snowflake.client.jdbc.EnvironmentVariables;
1414
import net.snowflake.client.jdbc.SnowflakeUtil;
15+
import net.snowflake.client.log.SFLogger;
16+
import net.snowflake.client.log.SFLoggerFactory;
1517

1618
@SnowflakeJdbcInternalApi
17-
public class AWSAttestationService {
19+
public class AwsAttestationService {
20+
21+
public static final SFLogger logger = SFLoggerFactory.getLogger(AwsAttestationService.class);
1822

1923
private static final String SECURE_TOKEN_SERVICE_NAME = "sts";
2024
private static boolean regionInitialized = false;
2125
private static String region;
2226

2327
private final AWS4Signer aws4Signer;
2428

25-
public AWSAttestationService() {
29+
public AwsAttestationService() {
2630
aws4Signer = new AWS4Signer();
2731
aws4Signer.setServiceName(SECURE_TOKEN_SERVICE_NAME);
2832
}
@@ -32,12 +36,18 @@ AWSCredentials getAWSCredentials() {
3236
}
3337

3438
String getAWSRegion() {
35-
if (!regionInitialized) {
36-
String envRegion = SnowflakeUtil.systemGetEnv(EnvironmentVariables.AWS_REGION.getName());
37-
region = envRegion != null ? envRegion : new InstanceMetadataRegionProvider().getRegion();
39+
try {
40+
if (!regionInitialized) {
41+
String envRegion = SnowflakeUtil.systemGetEnv(EnvironmentVariables.AWS_REGION.getName());
42+
region = envRegion != null ? envRegion : new InstanceMetadataRegionProvider().getRegion();
43+
}
44+
return region;
45+
} catch (Exception e) {
46+
logger.debug("Could not get AWS region", e);
47+
return null;
48+
} finally {
3849
regionInitialized = true;
3950
}
40-
return region;
4151
}
4252

4353
String getArn() {

‎src/main/java/net/snowflake/client/core/auth/wif/AwsIdentityAttestationCreator.java‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,9 +19,9 @@ public class AwsIdentityAttestationCreator implements WorkloadIdentityAttestatio
1919
private static final SFLogger logger =
2020
SFLoggerFactory.getLogger(AwsIdentityAttestationCreator.class);
2121

22-
private final AWSAttestationService attestationService;
22+
private final AwsAttestationService attestationService;
2323

24-
public AwsIdentityAttestationCreator(AWSAttestationService attestationService) {
24+
public AwsIdentityAttestationCreator(AwsAttestationService attestationService) {
2525
this.attestationService = attestationService;
2626
}
2727

‎src/main/java/net/snowflake/client/core/auth/wif/AzureAttestationService.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package net.snowflake.client.core.auth.wif;
2+
3+
import net.snowflake.client.core.SFLoginInput;
4+
import net.snowflake.client.core.SnowflakeJdbcInternalApi;
5+
import net.snowflake.client.jdbc.SnowflakeUtil;
6+
import net.snowflake.client.log.SFLogger;
7+
import net.snowflake.client.log.SFLoggerFactory;
8+
import org.apache.http.client.methods.HttpRequestBase;
9+
10+
@SnowflakeJdbcInternalApi
11+
public class AzureAttestationService {
12+
13+
private static final SFLogger logger = SFLoggerFactory.getLogger(AzureAttestationService.class);
14+
15+
// Expected to be set in Azure Functions environment
16+
String getIdentityEndpoint() {
17+
return SnowflakeUtil.systemGetEnv("IDENTITY_ENDPOINT");
18+
}
19+
20+
// Expected to be set in Azure Functions environment
21+
String getIdentityHeader() {
22+
return SnowflakeUtil.systemGetEnv("IDENTITY_HEADER");
23+
}
24+
25+
// Expected to be set in Azure Functions environment
26+
String getClientId() {
27+
return SnowflakeUtil.systemGetEnv("MANAGED_IDENTITY_CLIENT_ID");
28+
}
29+
30+
String fetchTokenFromMetadataService(HttpRequestBase tokenRequest, SFLoginInput loginInput) {
31+
try {
32+
return WorkloadIdentityUtil.performIdentityRequest(tokenRequest, loginInput);
33+
} catch (Exception e) {
34+
logger.debug("Azure metadata server request was not successful: {}", e);
35+
return null;
36+
}
37+
}
38+
}

0 commit comments

Comments
 (0)