SNOW-618478: Support private key base64 in connection parameters

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

@@ -45,6 +45,7 @@ public class SFLoginInput {
   private OCSPMode ocspMode;
   private HttpClientSettingsKey httpClientKey;
   private String privateKeyFile;
+  private String privateKeyBase64;
   private String privateKeyFilePwd;
   private String inFlightCtx; // Opaque string sent for Snowsight account activation
 
@@ -325,6 +326,15 @@ SFLoginInput setPrivateKey(PrivateKey privateKey) {
     return this;
   }
 
+  String getPrivateKeyBase64() {
+    return privateKeyBase64;
+  }
+
+  SFLoginInput setPrivateKeyBase64(String privateKeyBase64) {
+    this.privateKeyBase64 = privateKeyBase64;
+    return this;
+  }
+
   SFLoginInput setPrivateKeyFile(String privateKeyFile) {
     this.privateKeyFile = privateKeyFile;
     return this;
@@ -343,6 +353,12 @@ String getPrivateKeyFilePwd() {
     return privateKeyFilePwd;
   }
 
+  boolean isPrivateKeyProvided() {
+    return (getPrivateKey() != null
+        || getPrivateKeyFile() != null
+        || getPrivateKeyBase64() != null);
+  }
+
   public String getApplication() {
     return application;
   }

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

@@ -86,6 +86,7 @@ public class SFSession extends SFBaseSession {
   private String idToken;
   private String mfaToken;
   private String privateKeyFileLocation;
+  private String privateKeyBase64;
   private String privateKeyPassword;
   private PrivateKey privateKey;
 
@@ -452,6 +453,12 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
           }
           break;
 
+        case PRIVATE_KEY_BASE64:
+          if (propertyValue != null) {
+            privateKeyBase64 = (String) propertyValue;
+          }
+          break;
+
         case PRIVATE_KEY_FILE_PWD:
           if (propertyValue != null) {
             privateKeyPassword = (String) propertyValue;
@@ -553,7 +560,7 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
             + " warehouse: {}, validate default parameters: {}, authenticator: {}, ocsp mode: {},"
             + " passcode in password: {}, passcode is {}, private key is {}, disable socks proxy: {},"
             + " application: {}, app id: {}, app version: {}, login timeout: {}, retry timeout: {}, network timeout: {},"
-            + " query timeout: {}, tracing: {}, private key file: {}, private key file pwd is {},"
+            + " query timeout: {}, tracing: {}, private key file: {}, private key base64 is {}, private key file pwd is {},"
             + " enable_diagnostics: {}, diagnostics_allowlist_path: {},"
             + " session parameters: client store temporary credential: {}, gzip disabled: {}, browser response timeout: {}",
         connectionPropertiesMap.get(SFSessionProperty.SERVER_URL),
@@ -582,6 +589,8 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
         connectionPropertiesMap.get(SFSessionProperty.QUERY_TIMEOUT),
         connectionPropertiesMap.get(SFSessionProperty.TRACING),
         connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE),
+        SFLoggerUtil.isVariableProvided(
+            (String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_BASE64)),
         SFLoggerUtil.isVariableProvided(
             (String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD)),
         connectionPropertiesMap.get(SFSessionProperty.ENABLE_DIAGNOSTICS),
@@ -631,6 +640,8 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
         .setSessionParameters(sessionParametersMap)
         .setPrivateKey((PrivateKey) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY))
         .setPrivateKeyFile((String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE))
+        .setPrivateKeyBase64(
+            (String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_BASE64))
         .setPrivateKeyFilePwd(
             (String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD))
         .setApplication((String) connectionPropertiesMap.get(SFSessionProperty.APPLICATION))
@@ -750,7 +761,10 @@ private boolean isSnowflakeAuthenticator() {
     Map<SFSessionProperty, Object> connectionPropertiesMap = getConnectionPropertiesMap();
     String authenticator = (String) connectionPropertiesMap.get(SFSessionProperty.AUTHENTICATOR);
     PrivateKey privateKey = (PrivateKey) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY);
-    return (authenticator == null && privateKey == null && privateKeyFileLocation == null)
+    return (authenticator == null
+            && privateKey == null
+            && privateKeyFileLocation == null
+            && privateKeyBase64 == null)
         || ClientAuthnDTO.AuthenticatorType.SNOWFLAKE.name().equalsIgnoreCase(authenticator);
   }
 

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

@@ -54,6 +54,7 @@ public enum SFSessionProperty {
   VALIDATE_DEFAULT_PARAMETERS("validateDefaultParameters", false, Boolean.class),
   INJECT_WAIT_IN_PUT("inject_wait_in_put", false, Integer.class),
   PRIVATE_KEY_FILE("private_key_file", false, String.class),
+  PRIVATE_KEY_BASE64("private_key_base64", false, String.class),
   PRIVATE_KEY_FILE_PWD("private_key_file_pwd", false, String.class),
   CLIENT_INFO("snowflakeClientInfo", false, String.class),
   ALLOW_UNDERSCORES_IN_HOST("allowUnderscoresInHost", false, Boolean.class),

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

@@ -240,7 +240,7 @@ private static ClientAuthnDTO.AuthenticatorType getAuthenticator(SFLoginInput lo
     // authenticator is null, then jdbc will decide authenticator depends on
     // if privateKey is specified or not. If yes, authenticator type will be
     // SNOWFLAKE_JWT, otherwise it will use SNOWFLAKE.
-    return (loginInput.getPrivateKey() != null || loginInput.getPrivateKeyFile() != null)
+    return loginInput.isPrivateKeyProvided()
         ? ClientAuthnDTO.AuthenticatorType.SNOWFLAKE_JWT
         : ClientAuthnDTO.AuthenticatorType.SNOWFLAKE;
   }
@@ -421,6 +421,7 @@ private static SFLoginOutput newSession(
             new SessionUtilKeyPair(
                 loginInput.getPrivateKey(),
                 loginInput.getPrivateKeyFile(),
+                loginInput.getPrivateKeyBase64(),
                 loginInput.getPrivateKeyFilePwd(),
                 loginInput.getAccountName(),
                 loginInput.getUserName());
@@ -676,6 +677,7 @@ private static SFLoginOutput newSession(
                     new SessionUtilKeyPair(
                         loginInput.getPrivateKey(),
                         loginInput.getPrivateKeyFile(),
+                        loginInput.getPrivateKeyBase64(),
                         loginInput.getPrivateKeyFilePwd(),
                         loginInput.getAccountName(),
                         loginInput.getUserName());
@@ -1723,6 +1725,7 @@ public static void resetOCSPUrlIfNecessary(String serverUrl) throws IOException
    *
    * @param privateKey private key
    * @param privateKeyFile path to private key file
+   * @param privateKeyBase64 base64 encoded content of the private key file
    * @param privateKeyFilePwd password for private key file
    * @param accountName account name
    * @param userName user name
@@ -1732,16 +1735,41 @@ public static void resetOCSPUrlIfNecessary(String serverUrl) throws IOException
   public static String generateJWTToken(
       PrivateKey privateKey,
       String privateKeyFile,
+      String privateKeyBase64,
       String privateKeyFilePwd,
       String accountName,
       String userName)
       throws SFException {
     SessionUtilKeyPair s =
         new SessionUtilKeyPair(
-            privateKey, privateKeyFile, privateKeyFilePwd, accountName, userName);
+            privateKey, privateKeyFile, privateKeyBase64, privateKeyFilePwd, accountName, userName);
     return s.issueJwtToken();
   }
 
+  /**
+   * Helper function to generate a JWT token. Use {@link #generateJWTToken(PrivateKey, String,
+   * String, String, String, String)}
+   *
+   * @param privateKey private key
+   * @param privateKeyFile path to private key file
+   * @param privateKeyFilePwd password for private key file
+   * @param accountName account name
+   * @param userName user name
+   * @return JWT token
+   * @throws SFException if Snowflake error occurs
+   */
+  @Deprecated
+  public static String generateJWTToken(
+      PrivateKey privateKey,
+      String privateKeyFile,
+      String privateKeyFilePwd,
+      String accountName,
+      String userName)
+      throws SFException {
+    return generateJWTToken(
+        privateKey, privateKeyFile, null, privateKeyFilePwd, accountName, userName);
+  }
+
   /**
    * Helper method to check if the request path is a login/auth request to use for retry strategy.
    *

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

@@ -13,9 +13,9 @@
 import com.nimbusds.jose.crypto.RSASSASigner;
 import com.nimbusds.jwt.JWTClaimsSet;
 import com.nimbusds.jwt.SignedJWT;
-import java.io.FileReader;
 import java.io.IOException;
 import java.io.StringReader;
+import java.nio.charset.StandardCharsets;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -80,6 +80,7 @@ class SessionUtilKeyPair {
   SessionUtilKeyPair(
       PrivateKey privateKey,
       String privateKeyFile,
+      String privateKeyBase64,
       String privateKeyFilePwd,
       String accountName,
       String userName)
@@ -100,18 +101,10 @@ class SessionUtilKeyPair {
       }
     }
 
-    // if there is both a file and a private key, there is a problem
-    if (!Strings.isNullOrEmpty(privateKeyFile) && privateKey != null) {
-      throw new SFException(
-          ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
-          "Cannot have both private key value and private key file.");
-    } else {
-      // if privateKeyFile has a value and privateKey is null
-      this.privateKey =
-          Strings.isNullOrEmpty(privateKeyFile)
-              ? privateKey
-              : extractPrivateKeyFromFile(privateKeyFile, privateKeyFilePwd);
-    }
+    ensurePrivateKeyProvidedInOnlyOneProperty(privateKey, privateKeyFile, privateKeyBase64);
+    this.privateKey =
+        buildPrivateKey(privateKey, privateKeyFile, privateKeyBase64, privateKeyFilePwd);
+
     // construct public key from raw bytes
     if (this.privateKey instanceof RSAPrivateCrtKey) {
       RSAPrivateCrtKey rsaPrivateCrtKey = (RSAPrivateCrtKey) this.privateKey;
@@ -130,6 +123,42 @@ class SessionUtilKeyPair {
     }
   }
 
+  private static void ensurePrivateKeyProvidedInOnlyOneProperty(
+      PrivateKey privateKey, String privateKeyFile, String privateKeyBase64) throws SFException {
+    if (!Strings.isNullOrEmpty(privateKeyFile) && privateKey != null) {
+      throw new SFException(
+          ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
+          "Cannot have both private key object and private key file.");
+    }
+    if (!Strings.isNullOrEmpty(privateKeyBase64) && !Strings.isNullOrEmpty(privateKeyFile)) {
+      throw new SFException(
+          ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
+          "Cannot have both private key file and private key base64 string value.");
+    }
+    if (!Strings.isNullOrEmpty(privateKeyBase64) && privateKey != null) {
+      throw new SFException(
+          ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
+          "Cannot have both private key object and private key base64 string value.");
+    }
+  }
+
+  private PrivateKey buildPrivateKey(
+      PrivateKey privateKey,
+      String privateKeyFile,
+      String privateKeyBase64,
+      String privateKeyFilePwd)
+      throws SFException {
+    if (!Strings.isNullOrEmpty(privateKeyBase64)) {
+      logger.trace("Reading private key from base64 string");
+      return extractPrivateKeyFromBase64(privateKeyBase64, privateKeyFilePwd);
+    }
+    if (!Strings.isNullOrEmpty(privateKeyFile)) {
+      logger.trace("Reading private key from file");
+      return extractPrivateKeyFromFile(privateKeyFile, privateKeyFilePwd);
+    }
+    return privateKey;
+  }
+
   private KeyFactory getKeyFactoryInstance() throws NoSuchAlgorithmException {
     if (isFipsMode) {
       return KeyFactory.getInstance("RSA", this.SecurityProvider);
@@ -148,33 +177,52 @@ private SecretKeyFactory getSecretKeyFactory(String algorithm) throws NoSuchAlgo
 
   private PrivateKey extractPrivateKeyFromFile(String privateKeyFile, String privateKeyFilePwd)
       throws SFException {
+
+    try {
+      Path privKeyPath = Paths.get(privateKeyFile);
+      FileUtil.logFileUsage(privKeyPath, "Extract private key from file", true);
+      byte[] bytes = Files.readAllBytes(privKeyPath);
+      return extractPrivateKeyFromBytes(bytes, privateKeyFilePwd);
+    } catch (IOException ie) {
+      logger.error("Could not read private key from file", ie);
+      throw new SFException(ie, ErrorCode.INVALID_PARAMETER_VALUE, ie.getCause());
+    }
+  }
+
+  private PrivateKey extractPrivateKeyFromBytes(byte[] privateKeyBytes, String privateKeyFilePwd)
+      throws SFException {
     if (isBouncyCastleProviderEnabled) {
       try {
-        return extractPrivateKeyWithBouncyCastle(privateKeyFile, privateKeyFilePwd);
+        return extractPrivateKeyWithBouncyCastle(privateKeyBytes, privateKeyFilePwd);
       } catch (IOException | PKCSException | OperatorCreationException e) {
         logger.error("Could not extract private key using Bouncy Castle provider", e);
         throw new SFException(e, ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY, e.getCause());
       }
     } else {
       try {
-        return extractPrivateKeyWithJdk(privateKeyFile, privateKeyFilePwd);
+        return extractPrivateKeyWithJdk(privateKeyBytes, privateKeyFilePwd);
       } catch (NoSuchAlgorithmException
           | InvalidKeySpecException
           | IOException
           | IllegalArgumentException
           | NullPointerException
           | InvalidKeyException e) {
         logger.error(
-            "Could not extract private key. Try setting the JVM argument: " + "-D{}" + "=TRUE",
+            "Could not extract private key using standard JDK. Try setting the JVM argument: "
+                + "-D{}"
+                + "=TRUE",
             SecurityUtil.ENABLE_BOUNCYCASTLE_PROVIDER_JVM);
-        throw new SFException(
-            e,
-            ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
-            privateKeyFile + ": " + e.getMessage());
+        throw new SFException(e, ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY, e.getMessage());
       }
     }
   }
 
+  private PrivateKey extractPrivateKeyFromBase64(String privateKeyBase64, String privateKeyFilePwd)
+      throws SFException {
+    byte[] decodedKey = Base64.decodeBase64(privateKeyBase64);
+    return extractPrivateKeyFromBytes(decodedKey, privateKeyFilePwd);
+  }
+
   public String issueJwtToken() throws SFException {
     JWTClaimsSet.Builder builder = new JWTClaimsSet.Builder();
     String sub = String.format(SUBJECT_FMT, this.accountName, this.userName);
@@ -232,13 +280,12 @@ public static int getTimeout() {
   }
 
   private PrivateKey extractPrivateKeyWithBouncyCastle(
-      String privateKeyFile, String privateKeyFilePwd)
+      byte[] privateKeyBytes, String privateKeyFilePwd)
       throws IOException, PKCSException, OperatorCreationException {
-    Path privKeyPath = Paths.get(privateKeyFile);
-    FileUtil.logFileUsage(
-        privKeyPath, "Extract private key from file using Bouncy Castle provider", true);
+    logger.trace("Extracting private key using Bouncy Castle provider");
     PrivateKeyInfo privateKeyInfo = null;
-    PEMParser pemParser = new PEMParser(new FileReader(privKeyPath.toFile()));
+    PEMParser pemParser =
+        new PEMParser(new StringReader(new String(privateKeyBytes, StandardCharsets.UTF_8)));
     Object pemObject = pemParser.readObject();
     if (pemObject instanceof PKCS8EncryptedPrivateKeyInfo) {
       // Handle the case where the private key is encrypted.
@@ -264,11 +311,10 @@ private PrivateKey extractPrivateKeyWithBouncyCastle(
     return converter.getPrivateKey(privateKeyInfo);
   }
 
-  private PrivateKey extractPrivateKeyWithJdk(String privateKeyFile, String privateKeyFilePwd)
+  private PrivateKey extractPrivateKeyWithJdk(byte[] privateKeyFileBytes, String privateKeyFilePwd)
       throws IOException, NoSuchAlgorithmException, InvalidKeySpecException, InvalidKeyException {
-    Path privKeyPath = Paths.get(privateKeyFile);
-    FileUtil.logFileUsage(privKeyPath, "Extract private key from file using Jdk", true);
-    String privateKeyContent = new String(Files.readAllBytes(privKeyPath));
+    logger.trace("Extracting private key using JDK");
+    String privateKeyContent = new String(privateKeyFileBytes, StandardCharsets.UTF_8);
     if (Strings.isNullOrEmpty(privateKeyFilePwd)) {
       // unencrypted private key file
       return generatePrivateKey(false, privateKeyContent, privateKeyFilePwd);

src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java

@@ -236,6 +236,14 @@ public void setPrivateKeyFile(String location, String password) {
     }
   }
 
+  public void setPrivateKeyBase64(String privateKeyBase64, String password) {
+    this.setAuthenticator(AUTHENTICATOR_SNOWFLAKE_JWT);
+    this.properties.put(SFSessionProperty.PRIVATE_KEY_BASE64.getPropertyKey(), privateKeyBase64);
+    if (!Strings.isNullOrEmpty(password)) {
+      this.properties.put(SFSessionProperty.PRIVATE_KEY_FILE_PWD.getPropertyKey(), password);
+    }
+  }
+
   public void setTracing(String tracing) {
     this.properties.put(SFSessionProperty.TRACING.getPropertyKey(), tracing);
   }