SNOW-2204396: Add platform detection results to CLIENT_ENVIRONMENT (#2351)

Author: sfc-gh-dheyman

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 # Changelog
 - v3.27.1-SNAPSHOT
+    - Added platform detection on login to set PLATFORM metric in CLIENT_ENVIRONMENT
+  
 - v3.27.0
     - Added the `changelog.yml` GitHub workflow to ensure changelog is updated on release PRs.
     - Added HTTP 307 & 308 retries in case of internal IP redirects

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

@@ -51,6 +51,8 @@ public class SFLoginInput {
   private String privateKeyBase64;
   private String privateKeyPwd;
   private String inFlightCtx; // Opaque string sent for Snowsight account activation
+  private int platformDetectionTimeoutMs = 200; // Default 200ms timeout for platform detection
+  private boolean disablePlatformDetection = false; // Default false - platform detection enabled
 
   private SFOauthLoginInput oauthLoginInput;
 
@@ -637,4 +639,22 @@ public void setBrowserHandler(
       SessionUtilExternalBrowser.AuthExternalBrowserHandlers browserHandler) {
     this.browserHandler = browserHandler;
   }
+
+  public int getPlatformDetectionTimeoutMs() {
+    return platformDetectionTimeoutMs;
+  }
+
+  public SFLoginInput setPlatformDetectionTimeoutMs(int platformDetectionTimeoutMs) {
+    this.platformDetectionTimeoutMs = platformDetectionTimeoutMs;
+    return this;
+  }
+
+  public boolean isDisablePlatformDetection() {
+    return disablePlatformDetection;
+  }
+
+  public SFLoginInput setDisablePlatformDetection(boolean disablePlatformDetection) {
+    this.disablePlatformDetection = disablePlatformDetection;
+    return this;
+  }
 }

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

@@ -143,6 +143,8 @@ public class SFSession extends SFBaseSession {
    */
   private int retryTimeout = 300;
 
+  private int defaultPlatformDetectionTimeoutMs = 200;
+
   private boolean enableClientStoreTemporaryCredential = true;
   private boolean enableClientRequestMfaToken = true;
 
@@ -794,7 +796,16 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
                 : false)
         .setEnableClientStoreTemporaryCredential(enableClientStoreTemporaryCredential)
         .setEnableClientRequestMfaToken(enableClientRequestMfaToken)
-        .setBrowserResponseTimeout(browserResponseTimeout);
+        .setBrowserResponseTimeout(browserResponseTimeout)
+        .setPlatformDetectionTimeoutMs(
+            connectionPropertiesMap.get(SFSessionProperty.PLATFORM_DETECTION_TIMEOUT_MS) != null
+                ? (int) connectionPropertiesMap.get(SFSessionProperty.PLATFORM_DETECTION_TIMEOUT_MS)
+                : defaultPlatformDetectionTimeoutMs)
+        .setDisablePlatformDetection(
+            connectionPropertiesMap.get(SFSessionProperty.DISABLE_PLATFORM_DETECTION) != null
+                ? getBooleanValue(
+                    connectionPropertiesMap.get(SFSessionProperty.DISABLE_PLATFORM_DETECTION))
+                : false); // Default to false (platform detection enabled)
 
     logger.info(
         "Connecting to {} Snowflake domain",

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

@@ -153,7 +153,11 @@ public enum SFSessionProperty {
 
   CERT_REVOCATION_CHECK_MODE("CERT_REVOCATION_CHECK_MODE", false, String.class),
 
-  ALLOW_CERTIFICATES_WITHOUT_CRL_URL("ALLOW_CERTIFICATES_WITHOUT_CRL_URL", false, Boolean.class);
+  ALLOW_CERTIFICATES_WITHOUT_CRL_URL("ALLOW_CERTIFICATES_WITHOUT_CRL_URL", false, Boolean.class),
+
+  PLATFORM_DETECTION_TIMEOUT_MS("platformDetectionTimeoutMs", false, Integer.class),
+
+  DISABLE_PLATFORM_DETECTION("disablePlatformDetection", false, Boolean.class);
 
   // property key in string
   private String propertyKey;

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

@@ -20,6 +20,7 @@
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
+import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
@@ -54,6 +55,7 @@
 import net.snowflake.client.log.ArgSupplier;
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
+import net.snowflake.client.util.PlatformDetector;
 import net.snowflake.client.util.SecretDetector;
 import net.snowflake.client.util.Stopwatch;
 import net.snowflake.client.util.ThrowingFunction;
@@ -1146,6 +1148,23 @@ static Map<String, Object> createClientEnvironmentInfo(
 
     clientEnv.put("JDBC_JAR_NAME", SnowflakeDriver.getJdbcJarname());
 
+    // Add platform detection (if not disabled)
+    if (!loginInput.isDisablePlatformDetection()) {
+      try {
+        PlatformDetector platformDetector = new PlatformDetector();
+        AwsAttestationService awsAttestationService = new AwsAttestationService();
+        List<String> detectedPlatforms =
+            platformDetector.detectPlatforms(
+                loginInput.getPlatformDetectionTimeoutMs(), awsAttestationService);
+        clientEnv.put("PLATFORM", detectedPlatforms);
+      } catch (Exception e) {
+        logger.debug("Platform detection failed: {}", e.getMessage());
+        // Continue without platform information
+      }
+    } else {
+      logger.debug("Platform detection is disabled");
+    }
+
     // OAuth metrics data
     if (authenticatorType == AuthenticatorType.OAUTH
         && loginInput.getOriginalAuthenticator() != null) {

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

@@ -29,8 +29,13 @@ void initializeSignerRegion() {
     aws4Signer.setRegionName(getAWSRegion());
   }
 
-  AWSCredentials getAWSCredentials() {
-    return DefaultAWSCredentialsProviderChain.getInstance().getCredentials();
+  public AWSCredentials getAWSCredentials() {
+    try {
+      return DefaultAWSCredentialsProviderChain.getInstance().getCredentials();
+    } catch (Exception e) {
+      logger.debug("Failed to retrieve AWS credentials: {}", e.getMessage());
+      return null;
+    }
   }
 
   String getAWSRegion() {

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

@@ -0,0 +1,76 @@
+package net.snowflake.client.core.auth.wif;
+
+import com.amazonaws.auth.AWSCredentials;
+import java.io.IOException;
+import net.snowflake.client.core.HttpClientSettingsKey;
+import net.snowflake.client.core.HttpUtil;
+import net.snowflake.client.core.OCSPMode;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
+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;
+
+@SnowflakeJdbcInternalApi
+public class PlatformDetectionUtil {
+
+  private static final SFLogger logger = SFLoggerFactory.getLogger(PlatformDetectionUtil.class);
+
+  public static String performPlatformDetectionRequest(HttpRequestBase httpRequest, int timeoutMs)
+      throws SnowflakeSQLException, IOException {
+    return HttpUtil.executeGeneralRequestOmitSnowflakeHeaders(
+        httpRequest,
+        1,
+        timeoutMs / 1000,
+        timeoutMs,
+        0,
+        new HttpClientSettingsKey(OCSPMode.DISABLE_OCSP_CHECKS),
+        null);
+  }
+
+  public static boolean hasValidAwsIdentityForWif(AwsAttestationService attestationService) {
+    if (!hasValidAwsCredentials(attestationService)) {
+      logger.debug("No valid AWS credentials available for identity validation");
+      return false;
+    }
+    return true;
+  }
+
+  public static boolean hasValidAwsCredentials(AwsAttestationService attestationService) {
+    try {
+      AWSCredentials credentials = attestationService.getAWSCredentials();
+      return hasValidAwsCredentials(credentials);
+    } catch (Exception e) {
+      logger.debug("Failed to get AWS credentials from default provider chain: {}", e.getMessage());
+      return false;
+    }
+  }
+
+  public static boolean hasValidAwsCredentials(AWSCredentials awsCredentials) {
+    try {
+      if (awsCredentials == null) {
+        logger.debug("No AWS credentials found");
+        return false;
+      }
+
+      // Check if we have access key and secret key
+      String accessKey = awsCredentials.getAWSAccessKeyId();
+      String secretKey = awsCredentials.getAWSSecretKey();
+
+      if (accessKey == null
+          || accessKey.trim().isEmpty()
+          || secretKey == null
+          || secretKey.trim().isEmpty()) {
+        logger.debug("AWS credentials are incomplete");
+        return false;
+      }
+
+      logger.debug("Valid AWS credentials found");
+      return true;
+
+    } catch (Exception e) {
+      logger.debug("AWS credential check failed: {}", e.getMessage());
+      return false;
+    }
+  }
+}

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

@@ -8,23 +8,33 @@
 import net.snowflake.client.core.HttpUtil;
 import net.snowflake.client.core.SFException;
 import net.snowflake.client.core.SFLoginInput;
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
 import net.snowflake.client.jdbc.ErrorCode;
 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 {
+/**
+ * Utility class for Workload Identity Federation (WIF) specific operations. This class contains
+ * functions that are used exclusively within the WIF package.
+ */
+@SnowflakeJdbcInternalApi
+public class WorkloadIdentityUtil {
 
   private static final SFLogger logger = SFLoggerFactory.getLogger(WorkloadIdentityUtil.class);
 
-  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";
+  public static final String DEFAULT_METADATA_SERVICE_BASE_URL = "http://169.254.169.254";
+
+  public static final String SNOWFLAKE_AUDIENCE_HEADER_NAME = "X-Snowflake-Audience";
+  public static final String SNOWFLAKE_AUDIENCE = "snowflakecomputing.com";
 
-  static String performIdentityRequest(HttpRequestBase tokenRequest, SFLoginInput loginInput)
+  /**
+   * Performs an HTTP request for WIF identity token retrieval. This method is used by WIF
+   * authentication flows to communicate with cloud metadata services.
+   */
+  public static String performIdentityRequest(HttpRequestBase tokenRequest, SFLoginInput loginInput)
       throws SnowflakeSQLException, IOException {
     return HttpUtil.executeGeneralRequestOmitSnowflakeHeaders(
         tokenRequest,
@@ -36,7 +46,12 @@ static String performIdentityRequest(HttpRequestBase tokenRequest, SFLoginInput
         null);
   }
 
-  static SubjectAndIssuer extractClaimsWithoutVerifyingSignature(String token) throws SFException {
+  /**
+   * Extracts claims (subject and issuer) from a JWT token without verifying the signature. This is
+   * used in WIF flows where signature verification is handled elsewhere.
+   */
+  public static SubjectAndIssuer extractClaimsWithoutVerifyingSignature(String token)
+      throws SFException {
     Map<String, Object> claims = extractClaimsMap(token);
     if (claims == null) {
       throw new SFException(
@@ -65,11 +80,15 @@ private static Map<String, Object> extractClaimsMap(String token) {
     }
   }
 
-  static class SubjectAndIssuer {
+  /**
+   * Container class for JWT subject and issuer claims. Used in WIF authentication flows to pass
+   * extracted token claims.
+   */
+  public static class SubjectAndIssuer {
     private final String subject;
     private final String issuer;
 
-    SubjectAndIssuer(String subject, String issuer) {
+    public SubjectAndIssuer(String subject, String issuer) {
       this.issuer = issuer;
       this.subject = subject;
     }

src/main/java/net/snowflake/client/util/EnvironmentProvider.java

@@ -0,0 +1,14 @@
+package net.snowflake.client.util;
+
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
+
+/**
+ * Interface for providing environment variables to enable thread-safe testing. This abstraction
+ * allows dependency injection of environment variable access, making code testable with instance
+ * mocks (as opposed to static SnowflakeUtil) that work across threads.
+ */
+@SnowflakeJdbcInternalApi
+public interface EnvironmentProvider {
+
+  String getEnv(String name);
+}

src/main/java/net/snowflake/client/util/Platform.java

@@ -0,0 +1,34 @@
+package net.snowflake.client.util;
+
+import net.snowflake.client.core.SnowflakeJdbcInternalApi;
+
+/** Enum representing all detectable cloud platforms and identity providers. */
+@SnowflakeJdbcInternalApi
+public enum Platform {
+  IS_AWS_LAMBDA("is_aws_lambda"),
+  IS_AZURE_FUNCTION("is_azure_function"),
+  IS_GCE_CLOUD_RUN_SERVICE("is_gce_cloud_run_service"),
+  IS_GCE_CLOUD_RUN_JOB("is_gce_cloud_run_job"),
+  IS_GITHUB_ACTION("is_github_action"),
+  IS_EC2_INSTANCE("is_ec2_instance"),
+  HAS_AWS_IDENTITY_ARN_UNKNOWN("has_aws_identity_arn_unknown"),
+  IS_AZURE_VM("is_azure_vm"),
+  HAS_AZURE_MANAGED_IDENTITY("has_azure_managed_identity"),
+  IS_GCE_VM("is_gce_vm"),
+  HAS_GCP_IDENTITY("has_gcp_identity");
+
+  private final String value;
+
+  Platform(String value) {
+    this.value = value;
+  }
+
+  public String getValue() {
+    return value;
+  }
+
+  @Override
+  public String toString() {
+    return value;
+  }
+}