Merge branch 'master' into SNOW-1016470-increase-code-coverage-to-at-least-90

Author: sfc-gh-ext-simba-vb

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

@@ -62,6 +62,7 @@ public class SFLoginInput {
   // Workload Identity Federation
   private String workloadIdentityProvider;
   private WorkloadIdentityAttestation workloadIdentityAttestation;
+  private String workloadIdentityEntraResource;
 
   // OAuth
   private int redirectUriPort = -1;
@@ -612,4 +613,13 @@ public void setWorkloadIdentityAttestation(
   public WorkloadIdentityAttestation getWorkloadIdentityAttestation() {
     return workloadIdentityAttestation;
   }
+
+  public String getWorkloadIdentityEntraResource() {
+    return this.workloadIdentityEntraResource;
+  }
+
+  public SFLoginInput setWorkloadIdentityEntraResource(String workloadIdentityEntraResource) {
+    this.workloadIdentityEntraResource = workloadIdentityEntraResource;
+    return this;
+  }
 }

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

@@ -723,6 +723,9 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
         .setOauthLoginInput(oauthLoginInput)
         .setWorkloadIdentityProvider(
             (String) connectionPropertiesMap.get(SFSessionProperty.WORKLOAD_IDENTITY_PROVIDER))
+        .setWorkloadIdentityEntraResource(
+            (String)
+                connectionPropertiesMap.get(SFSessionProperty.WORKLOAD_IDENTITY_ENTRA_RESOURCE))
         .setPrivateKeyBase64(
             (String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_BASE64))
         .setPrivateKeyPwd(

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

@@ -27,6 +27,7 @@ public enum SFSessionProperty {
   OAUTH_AUTHORIZATION_URL("oauthAuthorizationUrl", false, String.class),
   OAUTH_TOKEN_REQUEST_URL("oauthTokenRequestUrl", false, String.class),
   WORKLOAD_IDENTITY_PROVIDER("workloadIdentityProvider", false, String.class),
+  WORKLOAD_IDENTITY_ENTRA_RESOURCE("workloadIdentityEntraResource", false, String.class),
   WAREHOUSE("warehouse", false, String.class),
   LOGIN_TIMEOUT("loginTimeout", false, Integer.class),
   NETWORK_TIMEOUT("networkTimeout", false, Integer.class),

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

@@ -34,6 +34,7 @@
 import net.snowflake.client.core.auth.oauth.TokenResponseDTO;
 import net.snowflake.client.core.auth.wif.AwsAttestationService;
 import net.snowflake.client.core.auth.wif.AwsIdentityAttestationCreator;
+import net.snowflake.client.core.auth.wif.AzureAttestationService;
 import net.snowflake.client.core.auth.wif.AzureIdentityAttestationCreator;
 import net.snowflake.client.core.auth.wif.GcpIdentityAttestationCreator;
 import net.snowflake.client.core.auth.wif.OidcIdentityAttestationCreator;
@@ -378,7 +379,7 @@ private static WorkloadIdentityAttestation getWorkloadIdentityAttestation(SFLogi
         new WorkloadIdentityAttestationProvider(
             new AwsIdentityAttestationCreator(new AwsAttestationService()),
             new GcpIdentityAttestationCreator(loginInput),
-            new AzureIdentityAttestationCreator(),
+            new AzureIdentityAttestationCreator(new AzureAttestationService(), loginInput),
             new OidcIdentityAttestationCreator(loginInput.getToken()));
     return attestationProvider.getAttestation(loginInput.getWorkloadIdentityProvider());
   }

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

@@ -0,0 +1,38 @@
+package net.snowflake.client.core.auth.wif;
+
+import net.snowflake.client.core.SFLoginInput;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
+import net.snowflake.client.jdbc.SnowflakeUtil;
+import net.snowflake.client.log.SFLogger;
+import net.snowflake.client.log.SFLoggerFactory;
+import org.apache.http.client.methods.HttpRequestBase;
+
+@SnowflakeJdbcInternalApi
+public class AzureAttestationService {
+
+  private static final SFLogger logger = SFLoggerFactory.getLogger(AzureAttestationService.class);
+
+  // Expected to be set in Azure Functions environment
+  String getIdentityEndpoint() {
+    return SnowflakeUtil.systemGetEnv("IDENTITY_ENDPOINT");
+  }
+
+  // Expected to be set in Azure Functions environment
+  String getIdentityHeader() {
+    return SnowflakeUtil.systemGetEnv("IDENTITY_HEADER");
+  }
+
+  // Expected to be set in Azure Functions environment
+  String getClientId() {
+    return SnowflakeUtil.systemGetEnv("MANAGED_IDENTITY_CLIENT_ID");
+  }
+
+  String fetchTokenFromMetadataService(HttpRequestBase tokenRequest, SFLoginInput loginInput) {
+    try {
+      return WorkloadIdentityUtil.performIdentityRequest(tokenRequest, loginInput);
+    } catch (Exception e) {
+      logger.debug("Azure metadata server request was not successful: {}", e);
+      return null;
+    }
+  }
+}

src/main/java/net/snowflake/client/core/auth/wif/AzureIdentityAttestationCreator.java

@@ -1,17 +1,127 @@
 package net.snowflake.client.core.auth.wif;
 
+import static net.snowflake.client.core.auth.wif.WorkloadIdentityUtil.DEFAULT_METADATA_SERVICE_BASE_URL;
+import static net.snowflake.client.core.auth.wif.WorkloadIdentityUtil.SubjectAndIssuer;
+import static net.snowflake.client.core.auth.wif.WorkloadIdentityUtil.extractClaimsWithoutVerifyingSignature;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import com.google.common.base.Strings;
+import net.snowflake.client.core.SFLoginInput;
 import net.snowflake.client.core.SnowflakeJdbcInternalApi;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
+import org.apache.http.client.methods.HttpGet;
 
 @SnowflakeJdbcInternalApi
 public class AzureIdentityAttestationCreator implements WorkloadIdentityAttestationCreator {
 
   private static final SFLogger logger =
       SFLoggerFactory.getLogger(AzureIdentityAttestationCreator.class);
+  public static final ObjectMapper objectMapper = new ObjectMapper();
+
+  private static final String EXPECTED_AZURE_TOKEN_ISSUER_PREFIX = "https://sts.windows.net/";
+  private static final String DEFAULT_WORKLOAD_IDENTITY_ENTRA_RESOURCE =
+      "api://fd3f753b-eed3-462c-b6a7-a4b5bb650aad";
+
+  private final AzureAttestationService azureAttestationService;
+  private final SFLoginInput loginInput;
+  private final String workloadIdentityEntraResource;
+  private final String azureMetadataServiceBaseUrl;
+
+  public AzureIdentityAttestationCreator(
+      AzureAttestationService azureAttestationService, SFLoginInput loginInput) {
+    this.azureAttestationService = azureAttestationService;
+    this.azureMetadataServiceBaseUrl = DEFAULT_METADATA_SERVICE_BASE_URL;
+    this.loginInput = loginInput;
+    this.workloadIdentityEntraResource = getEntraResource(loginInput);
+  }
+
+  /** Only for testing purpose */
+  public AzureIdentityAttestationCreator(
+      AzureAttestationService azureAttestationService,
+      SFLoginInput loginInput,
+      String azureMetadataServiceBaseUrl) {
+    this.azureAttestationService = azureAttestationService;
+    this.azureMetadataServiceBaseUrl = azureMetadataServiceBaseUrl;
+    this.loginInput = loginInput;
+    this.workloadIdentityEntraResource = getEntraResource(loginInput);
+  }
 
   @Override
   public WorkloadIdentityAttestation createAttestation() {
-    throw new RuntimeException("Azure Workload Identity not supported");
+    logger.debug("Creating Azure identity attestation...");
+    String identityEndpoint = azureAttestationService.getIdentityEndpoint();
+    HttpGet request;
+    if (Strings.isNullOrEmpty(identityEndpoint)) {
+      request = createAzureVMIdentityRequest();
+    } else {
+      String identityHeader = azureAttestationService.getIdentityHeader();
+      if (Strings.isNullOrEmpty(identityHeader)) {
+        logger.warn("Managed identity is not enabled on this Azure function.");
+        return null;
+      }
+      request =
+          createAzureFunctionsIdentityRequest(
+              identityEndpoint, identityHeader, azureAttestationService.getClientId());
+    }
+    String tokenJson = azureAttestationService.fetchTokenFromMetadataService(request, loginInput);
+    if (tokenJson == null) {
+      logger.debug("Could not fetch Azure token.");
+      return null;
+    }
+    String token = extractTokenFromJson(tokenJson);
+    if (token == null) {
+      logger.error("No access token found in Azure response.");
+      return null;
+    }
+    SubjectAndIssuer claims = extractClaimsWithoutVerifyingSignature(token);
+    if (claims == null) {
+      logger.error("Could not extract claims from token");
+      return null;
+    }
+    if (!claims.getIssuer().startsWith(EXPECTED_AZURE_TOKEN_ISSUER_PREFIX)) {
+      logger.error("Unexpected Azure token issuer: {}", claims.getIssuer());
+      return null;
+    }
+    return new WorkloadIdentityAttestation(
+        WorkloadIdentityProviderType.AZURE, token, claims.toMap());
+  }
+
+  private String getEntraResource(SFLoginInput loginInput) {
+    if (!Strings.isNullOrEmpty(loginInput.getWorkloadIdentityEntraResource())) {
+      return loginInput.getWorkloadIdentityEntraResource();
+    } else {
+      return DEFAULT_WORKLOAD_IDENTITY_ENTRA_RESOURCE;
+    }
+  }
+
+  private String extractTokenFromJson(String tokenJson) {
+    try {
+      JsonNode jsonNode = objectMapper.readTree(tokenJson);
+      return jsonNode.get("access_token").asText();
+    } catch (Exception e) {
+      logger.error("Unable to extract token from Azure metadata response: {}", e.getMessage());
+      return null;
+    }
+  }
+
+  private HttpGet createAzureFunctionsIdentityRequest(
+      String identityEndpoint, String identityHeader, String managedIdentityClientId) {
+    String queryParams = "api-version=2019-08-01&resource=" + workloadIdentityEntraResource;
+    if (managedIdentityClientId != null) {
+      queryParams += "&client_id=" + managedIdentityClientId;
+    }
+    HttpGet request = new HttpGet(String.format("%s?%s", identityEndpoint, queryParams));
+    request.addHeader("X-IDENTITY-HEADER", identityHeader);
+    return request;
+  }
+
+  private HttpGet createAzureVMIdentityRequest() {
+    String urlWithoutQueryString = azureMetadataServiceBaseUrl + "/metadata/identity/oauth2/token?";
+    String queryParams = "api-version=2018-02-01&resource=" + workloadIdentityEntraResource;
+    HttpGet request = new HttpGet(urlWithoutQueryString + queryParams);
+    request.setHeader("Metadata", "True");
+    return request;
   }
 }

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

@@ -1,7 +1,9 @@
 package net.snowflake.client.core.auth.wif;
 
+import static net.snowflake.client.core.auth.wif.WorkloadIdentityUtil.DEFAULT_METADATA_SERVICE_BASE_URL;
+import static net.snowflake.client.core.auth.wif.WorkloadIdentityUtil.performIdentityRequest;
+
 import java.util.Collections;
-import net.snowflake.client.core.HttpUtil;
 import net.snowflake.client.core.SFLoginInput;
 import net.snowflake.client.core.SnowflakeJdbcInternalApi;
 import net.snowflake.client.log.SFLogger;
@@ -25,7 +27,7 @@ public class GcpIdentityAttestationCreator implements WorkloadIdentityAttestatio
 
   public GcpIdentityAttestationCreator(SFLoginInput loginInput) {
     this.loginInput = loginInput;
-    gcpMetadataServiceBaseUrl = DEFAULT_GCP_METADATA_SERVICE_BASE_URL;
+    gcpMetadataServiceBaseUrl = DEFAULT_METADATA_SERVICE_BASE_URL;
   }
 
   /** Only for testing purpose */
@@ -36,6 +38,7 @@ public GcpIdentityAttestationCreator(SFLoginInput loginInput) {
 
   @Override
   public WorkloadIdentityAttestation createAttestation() {
+    logger.debug("Creating GCP identity attestation...");
     String token = fetchTokenFromMetadataService();
     if (token == null) {
       logger.debug("No GCP token was found.");
@@ -71,15 +74,9 @@ private String fetchTokenFromMetadataService() {
     HttpGet tokenRequest = new HttpGet(uri);
     tokenRequest.setHeader(METADATA_FLAVOR_HEADER_NAME, METADATA_FLAVOR);
     try {
-      return HttpUtil.executeGeneralRequestOmitRequestGuid(
-          tokenRequest,
-          loginInput.getLoginTimeout(),
-          3, // 3s timeout
-          loginInput.getSocketTimeoutInMillis(),
-          0,
-          loginInput.getHttpClientSettingsKey());
+      return performIdentityRequest(tokenRequest, loginInput);
     } catch (Exception e) {
-      logger.debug("GCP metadata server request was not successful: " + e.getMessage());
+      logger.debug("GCP metadata server request was not successful: {}" + e);
       return null;
     }
   }

src/main/java/net/snowflake/client/core/auth/wif/OidcIdentityAttestationCreator.java

@@ -18,6 +18,7 @@ public OidcIdentityAttestationCreator(String token) {
 
   @Override
   public WorkloadIdentityAttestation createAttestation() {
+    logger.debug("Creating OIDC identity attestation...");
     if (token == null) {
       logger.debug("No OIDC token was specified");
       return null;

src/main/java/net/snowflake/client/core/auth/wif/WorkloadIdentityAttestation.java

@@ -30,4 +30,17 @@ public String getCredential() {
   public Map<String, String> getUserIdentifierComponents() {
     return userIdentifierComponents;
   }
+
+  @Override
+  public String toString() {
+    return "WorkloadIdentityAttestation{"
+        + "provider="
+        + provider
+        + ", credential='"
+        + credential
+        + '\''
+        + ", userIdentifierComponents="
+        + userIdentifierComponents
+        + '}';
+  }
 }

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

@@ -2,10 +2,15 @@
 
 import com.nimbusds.jwt.JWT;
 import com.nimbusds.jwt.JWTParser;
+import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
+import net.snowflake.client.core.HttpUtil;
+import net.snowflake.client.core.SFLoginInput;
+import net.snowflake.client.jdbc.SnowflakeSQLException;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
+import org.apache.http.client.methods.HttpRequestBase;
 
 class WorkloadIdentityUtil {
 
@@ -14,6 +19,20 @@ class WorkloadIdentityUtil {
   static final String SNOWFLAKE_AUDIENCE_HEADER_NAME = "X-Snowflake-Audience";
   static final String SNOWFLAKE_AUDIENCE = "snowflakecomputing.com";
 
+  // Address commonly used by AWS, Azure & GCP to host instance metadata service
+  static final String DEFAULT_METADATA_SERVICE_BASE_URL = "http://169.254.169.254";
+
+  static String performIdentityRequest(HttpRequestBase tokenRequest, SFLoginInput loginInput)
+      throws SnowflakeSQLException, IOException {
+    return HttpUtil.executeGeneralRequestOmitRequestGuid(
+        tokenRequest,
+        loginInput.getLoginTimeout(),
+        3, // 3s timeout
+        loginInput.getSocketTimeoutInMillis(),
+        0,
+        loginInput.getHttpClientSettingsKey());
+  }
+
   static SubjectAndIssuer extractClaimsWithoutVerifyingSignature(String token) {
     Map<String, Object> claims = extractClaimsMap(token);
     if (claims == null) {
@@ -38,7 +57,7 @@ private static Map<String, Object> extractClaimsMap(String token) {
       JWT jwt = JWTParser.parse(token);
       return jwt.getJWTClaimsSet().getClaims();
     } catch (Exception e) {
-      logger.debug("Unable to extract JWT claims from token", e);
+      logger.error("Unable to extract JWT claims from token: {}", e);
       return null;
     }
   }