Skip to content

Commit 823211d

Browse files
sfc-gh-dheymansfc-gh-dheyman
authored
SNOW-1902245: Add support for GCP Workload Identity Federation (#2144)
1 parent 3ec16d3 commit 823211d

14 files changed

Lines changed: 406 additions & 12 deletions

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

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -659,6 +659,32 @@ public static String executeGeneralRequest(
659659
null);
660660
}
661661

662+
@SnowflakeJdbcInternalApi
663+
public static String executeGeneralRequestOmitRequestGuid(
664+
HttpRequestBase httpRequest,
665+
int retryTimeout,
666+
int authTimeout,
667+
int socketTimeout,
668+
int retryCount,
669+
HttpClientSettingsKey ocspAndProxyAndGzipKey)
670+
throws SnowflakeSQLException, IOException {
671+
return executeRequestInternal(
672+
httpRequest,
673+
retryTimeout,
674+
authTimeout,
675+
socketTimeout,
676+
retryCount,
677+
0,
678+
null,
679+
false,
680+
false,
681+
false,
682+
false,
683+
getHttpClient(ocspAndProxyAndGzipKey),
684+
new ExecTimeTelemetryData(),
685+
null);
686+
}
687+
662688
/**
663689
* Executes an HTTP request for Snowflake.
664690
*

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

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -261,7 +261,8 @@ public int getSocketTimeoutInMillis() {
261261
return (int) socketTimeout.toMillis();
262262
}
263263

264-
SFLoginInput setSocketTimeout(Duration socketTimeout) {
264+
@SnowflakeJdbcInternalApi
265+
public SFLoginInput setSocketTimeout(Duration socketTimeout) {
265266
this.socketTimeout = socketTimeout;
266267
return this;
267268
}
@@ -461,7 +462,8 @@ public HttpClientSettingsKey getHttpClientSettingsKey() {
461462
return httpClientKey;
462463
}
463464

464-
SFLoginInput setHttpClientSettingsKey(HttpClientSettingsKey key) {
465+
@SnowflakeJdbcInternalApi
466+
public SFLoginInput setHttpClientSettingsKey(HttpClientSettingsKey key) {
465467
this.httpClientKey = key;
466468
return this;
467469
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -342,7 +342,7 @@ static SFLoginOutput openSession(
342342
WorkloadIdentityAttestationProvider attestationProvider =
343343
new WorkloadIdentityAttestationProvider(
344344
new AwsIdentityAttestationCreator(new AWSAttestationService()),
345-
new GcpIdentityAttestationCreator(),
345+
new GcpIdentityAttestationCreator(loginInput),
346346
new AzureIdentityAttestationCreator(),
347347
new OidcIdentityAttestationCreator());
348348
WorkloadIdentityAttestation attestation =

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

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

22-
private static final String SNOWFLAKE_AUDIENCE_HEADER_NAME = "X-Snowflake-Audience";
23-
private static final String SNOWFLAKE_AUDIENCE = "snowflakecomputing.com";
24-
2522
private final AWSAttestationService attestationService;
2623

2724
public AwsIdentityAttestationCreator(AWSAttestationService attestationService) {
@@ -63,7 +60,9 @@ private Request<Void> createStsRequest(String hostname) {
6360
URI.create(
6461
String.format("https://%s/?Action=GetCallerIdentity&Version=2011-06-15", hostname)));
6562
request.addHeader("Host", hostname);
66-
request.addHeader(SNOWFLAKE_AUDIENCE_HEADER_NAME, SNOWFLAKE_AUDIENCE);
63+
request.addHeader(
64+
WorkloadIdentityUtil.SNOWFLAKE_AUDIENCE_HEADER_NAME,
65+
WorkloadIdentityUtil.SNOWFLAKE_AUDIENCE);
6766
return request;
6867
}
6968

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

Lines changed: 86 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,19 +1,101 @@
11
package net.snowflake.client.core.auth.wif;
22

3-
import net.snowflake.client.core.SFException;
3+
import com.nimbusds.jwt.JWT;
4+
import com.nimbusds.jwt.JWTParser;
5+
import java.util.Collections;
6+
import java.util.Map;
7+
import net.snowflake.client.core.HttpUtil;
8+
import net.snowflake.client.core.SFLoginInput;
49
import net.snowflake.client.core.SnowflakeJdbcInternalApi;
5-
import net.snowflake.client.jdbc.ErrorCode;
610
import net.snowflake.client.log.SFLogger;
711
import net.snowflake.client.log.SFLoggerFactory;
12+
import org.apache.http.client.methods.HttpGet;
813

914
@SnowflakeJdbcInternalApi
1015
public class GcpIdentityAttestationCreator implements WorkloadIdentityAttestationCreator {
1116

17+
private static final String METADATA_FLAVOR_HEADER_NAME = "Metadata-Flavor";
18+
private static final String METADATA_FLAVOR = "Google";
19+
private static final String EXPECTED_GCP_TOKEN_ISSUER = "https://accounts.google.com";
20+
private static final String DEFAULT_GCP_METADATA_SERVICE_BASE_URL = "http://169.254.169.254";
21+
22+
private final String gcpMetadataServiceBaseUrl;
23+
1224
private static final SFLogger logger =
1325
SFLoggerFactory.getLogger(GcpIdentityAttestationCreator.class);
1426

27+
private final SFLoginInput loginInput;
28+
29+
public GcpIdentityAttestationCreator(SFLoginInput loginInput) {
30+
this.loginInput = loginInput;
31+
gcpMetadataServiceBaseUrl = DEFAULT_GCP_METADATA_SERVICE_BASE_URL;
32+
}
33+
34+
/** Only for testing purpose */
35+
GcpIdentityAttestationCreator(SFLoginInput loginInput, String gcpBaseUrl) {
36+
this.loginInput = loginInput;
37+
this.gcpMetadataServiceBaseUrl = gcpBaseUrl;
38+
}
39+
1540
@Override
16-
public WorkloadIdentityAttestation createAttestation() throws SFException {
17-
throw new SFException(ErrorCode.FEATURE_UNSUPPORTED, "GCP Workload Identity not supported");
41+
public WorkloadIdentityAttestation createAttestation() {
42+
String token = fetchTokenFromMetadataService();
43+
if (token == null) {
44+
logger.debug("No GCP token was found.");
45+
return null;
46+
}
47+
// if the token has been returned, we can assume that we're on GCP environment
48+
Map<String, Object> claims = extractClaims(token);
49+
if (claims == null) {
50+
logger.error("Failed to parse JWT and extract claims");
51+
return null;
52+
}
53+
String issuer = (String) claims.get("iss");
54+
if (issuer == null) {
55+
logger.error("Missing issuer claim in GCP token");
56+
return null;
57+
}
58+
String subject = (String) claims.get("sub");
59+
if (subject == null) {
60+
logger.error("Missing sub claim in GCP token");
61+
return null;
62+
}
63+
if (!issuer.equals(EXPECTED_GCP_TOKEN_ISSUER)) {
64+
logger.error("Unexpected GCP token issuer:" + issuer);
65+
return null;
66+
}
67+
return new WorkloadIdentityAttestation(
68+
WorkloadIdentityProviderType.GCP, token, Collections.singletonMap("sub", subject));
69+
}
70+
71+
private Map<String, Object> extractClaims(String token) {
72+
try {
73+
JWT jwt = JWTParser.parse(token);
74+
return jwt.getJWTClaimsSet().getClaims();
75+
} catch (Exception e) {
76+
logger.debug("Unable to extract JWT claims from token", e);
77+
return null;
78+
}
79+
}
80+
81+
private String fetchTokenFromMetadataService() {
82+
String uri =
83+
gcpMetadataServiceBaseUrl
84+
+ "/computeMetadata/v1/instance/service-accounts/default/identity?audience="
85+
+ WorkloadIdentityUtil.SNOWFLAKE_AUDIENCE;
86+
HttpGet tokenRequest = new HttpGet(uri);
87+
tokenRequest.setHeader(METADATA_FLAVOR_HEADER_NAME, METADATA_FLAVOR);
88+
try {
89+
return HttpUtil.executeGeneralRequestOmitRequestGuid(
90+
tokenRequest,
91+
loginInput.getLoginTimeout(),
92+
3, // 3s timeout
93+
loginInput.getSocketTimeoutInMillis(),
94+
0,
95+
loginInput.getHttpClientSettingsKey());
96+
} catch (Exception e) {
97+
logger.debug("GCP metadata server request was not successful: " + e.getMessage());
98+
return null;
99+
}
18100
}
19101
}

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

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,6 @@
1+
package net.snowflake.client.core.auth.wif;
2+
3+
class WorkloadIdentityUtil {
4+
static final String SNOWFLAKE_AUDIENCE_HEADER_NAME = "X-Snowflake-Audience";
5+
static final String SNOWFLAKE_AUDIENCE = "snowflakecomputing.com";
6+
}

‎src/test/java/net/snowflake/client/core/auth/wif/GcpIdentityAttestationCreatorLatestIT.java‎

Lines changed: 136 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,136 @@
1+
package net.snowflake.client.core.auth.wif;
2+
3+
import static org.junit.jupiter.api.Assertions.assertEquals;
4+
import static org.junit.jupiter.api.Assertions.assertNotNull;
5+
import static org.junit.jupiter.api.Assertions.assertNull;
6+
7+
import java.time.Duration;
8+
import net.snowflake.client.category.TestTags;
9+
import net.snowflake.client.core.HttpClientSettingsKey;
10+
import net.snowflake.client.core.OCSPMode;
11+
import net.snowflake.client.core.SFLoginInput;
12+
import net.snowflake.client.jdbc.BaseWiremockTest;
13+
import org.junit.jupiter.api.Tag;
14+
import org.junit.jupiter.api.Test;
15+
16+
@Tag(TestTags.AUTHENTICATION)
17+
class GcpIdentityAttestationCreatorLatestIT extends BaseWiremockTest {
18+
19+
private static final String SCENARIOS_BASE_DIR = MAPPINGS_BASE_DIR + "/wif/gcp";
20+
21+
/*
22+
* {
23+
* "iss": "https://accounts.google.com",
24+
* "iat": 1743692017,
25+
* "exp": 1775228014,
26+
* "aud": "www.example.com",
27+
* "sub": "some-subject"
28+
* }
29+
*/
30+
private static final String SUCCESSFUL_FLOW_SCENARIO_MAPPINGS =
31+
SCENARIOS_BASE_DIR + "/successful_flow.json";
32+
33+
/*
34+
* {
35+
* "iss": "https://not.google.com",
36+
* "iat": 1743761213,
37+
* "exp": 1743764813,
38+
* "aud": "www.example.com",
39+
* "sub": "some-subject"
40+
* }
41+
*/
42+
private static final String INVALID_ISSUER_SCENARIO_MAPPINGS =
43+
SCENARIOS_BASE_DIR + "/invalid_issuer_claim.json";
44+
45+
/*
46+
* {
47+
* "sub": "some-subject",
48+
* "iat": 1743761213,
49+
* "exp": 1743764813,
50+
* "aud": "www.example.com"
51+
* }
52+
*/
53+
private static final String MISSING_ISSUER_SCENARIO_MAPPINGS =
54+
SCENARIOS_BASE_DIR + "/missing_issuer_claim.json";
55+
56+
/*
57+
* {
58+
* "iss": "https://accounts.google.com",
59+
* "iat": 1743761213,
60+
* "exp": 1743764813,
61+
* "aud": "www.example.com"
62+
* }
63+
*/
64+
private static final String MISSING_SUB_SCENARIO_MAPPINGS =
65+
SCENARIOS_BASE_DIR + "/missing_sub_claim.json";
66+
67+
// token equal to "unparsable.token"
68+
private static final String TOKEN_PARSE_ERROR_SCENARIO_MAPPINGS =
69+
SCENARIOS_BASE_DIR + "/unparsable_token.json";
70+
71+
// 400 Bad Request
72+
private static final String HTTP_ERROR_MAPPINGS = SCENARIOS_BASE_DIR + "/http_error.json";
73+
74+
@Test
75+
public void successfulFlowScenario() {
76+
importMappingFromResources(SUCCESSFUL_FLOW_SCENARIO_MAPPINGS);
77+
SFLoginInput loginInput = createLoginInputStub();
78+
79+
GcpIdentityAttestationCreator attestationCreator =
80+
new GcpIdentityAttestationCreator(loginInput, getBaseUrl());
81+
WorkloadIdentityAttestation attestation = attestationCreator.createAttestation();
82+
assertNotNull(attestation);
83+
assertEquals(WorkloadIdentityProviderType.GCP, attestation.getProvider());
84+
assertEquals("some-subject", attestation.getUserIdentifiedComponents().get("sub"));
85+
assertNotNull(attestation.getCredential());
86+
}
87+
88+
@Test
89+
public void invalidIssuerScenario() {
90+
importMappingFromResources(INVALID_ISSUER_SCENARIO_MAPPINGS);
91+
creatAttestationAndAssertNull();
92+
}
93+
94+
@Test
95+
public void missingIssuerScenario() {
96+
importMappingFromResources(MISSING_ISSUER_SCENARIO_MAPPINGS);
97+
creatAttestationAndAssertNull();
98+
}
99+
100+
@Test
101+
public void missingSubScenario() {
102+
importMappingFromResources(MISSING_SUB_SCENARIO_MAPPINGS);
103+
creatAttestationAndAssertNull();
104+
}
105+
106+
@Test
107+
public void unparsableTokenScenario() {
108+
importMappingFromResources(TOKEN_PARSE_ERROR_SCENARIO_MAPPINGS);
109+
creatAttestationAndAssertNull();
110+
}
111+
112+
@Test
113+
public void httpErrorScenario() {
114+
importMappingFromResources(HTTP_ERROR_MAPPINGS);
115+
creatAttestationAndAssertNull();
116+
}
117+
118+
private void creatAttestationAndAssertNull() {
119+
SFLoginInput loginInput = createLoginInputStub();
120+
GcpIdentityAttestationCreator attestationCreator =
121+
new GcpIdentityAttestationCreator(loginInput, getBaseUrl());
122+
WorkloadIdentityAttestation attestation = attestationCreator.createAttestation();
123+
assertNull(attestation);
124+
}
125+
126+
private String getBaseUrl() {
127+
return String.format("http://%s:%d/", WIREMOCK_HOST, wiremockHttpPort);
128+
}
129+
130+
private SFLoginInput createLoginInputStub() {
131+
SFLoginInput loginInputStub = new SFLoginInput();
132+
loginInputStub.setSocketTimeout(Duration.ofMinutes(5));
133+
loginInputStub.setHttpClientSettingsKey(new HttpClientSettingsKey(OCSPMode.FAIL_OPEN));
134+
return loginInputStub;
135+
}
136+
}

‎src/test/java/net/snowflake/client/core/auth/wif/WorkloadIdentityAttestationProviderTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ public void shouldCreateProperAttestationCreatorByType() throws SFException {
3232
WorkloadIdentityAttestationProvider provider =
3333
new WorkloadIdentityAttestationProvider(
3434
new AwsIdentityAttestationCreator(null),
35-
new GcpIdentityAttestationCreator(),
35+
new GcpIdentityAttestationCreator(null),
3636
new AzureIdentityAttestationCreator(),
3737
new OidcIdentityAttestationCreator());
3838
WorkloadIdentityAttestationCreator attestationCreator =

‎src/test/resources/wiremock/mappings/wif/gcp/http_error.json‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
{
2+
"mappings": [
3+
{
4+
"request": {
5+
"urlPattern": "/computeMetadata/v1/instance/service-accounts/default/identity.*",
6+
"queryParameters": {
7+
"audience": {
8+
"equalTo": "snowflakecomputing.com"
9+
}
10+
},
11+
"method": "GET",
12+
"headers": {
13+
"Metadata-Flavor": {
14+
"equalTo": "Google"
15+
}
16+
}
17+
},
18+
"response": {
19+
"status": 400
20+
}
21+
}
22+
]
23+
}

‎src/test/resources/wiremock/mappings/wif/gcp/invalid_issuer_claim.json‎

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
{
2+
"mappings": [
3+
{
4+
"request": {
5+
"urlPattern": "/computeMetadata/v1/instance/service-accounts/default/identity.*",
6+
"queryParameters": {
7+
"audience": {
8+
"equalTo": "snowflakecomputing.com"
9+
}
10+
},
11+
"method": "GET",
12+
"headers": {
13+
"Metadata-Flavor": {
14+
"equalTo": "Google"
15+
}
16+
}
17+
},
18+
"response": {
19+
"status": 200,
20+
"body": "eyJ0eXAiOiJhdCtqd3QiLCJhbGciOiJFUzI1NiIsImtpZCI6ImU2M2I5NzA1OTRiY2NmZTAxMDlkOTg4OWM2MDk3OWEwIn0.eyJpc3MiOiJodHRwczovL25vdC5nb29nbGUuY29tIiwiaWF0IjoxNzQzNzYxMjEzLCJleHAiOjE3NDM3NjQ4MTMsImF1ZCI6Ind3dy5leGFtcGxlLmNvbSIsInN1YiI6InNvbWUtc3ViamVjdCJ9.lQjch0MLams6AprSVOdZfVnHxqySB1sgDPW8sV839QPFUutMPj3SNkxGvrkIJj2QPDeCh02bdAzXe4N75tLORg"
21+
}
22+
}
23+
]
24+
}

0 commit comments

Comments
 (0)