SNOW-618478: Support private key base64 in connection parameters (#1847)

Co-authored-by: Eric Simmerman <ericsimmerman@gmail.com>
Co-authored-by: Dawid Heyman <dawid.heyman@snowflake.com>
Co-authored-by: Waleed Fateem <72769898+sfc-gh-wfateem@users.noreply.github.com>
Co-authored-by: Mikołaj Kubik <mikolaj.kubik@snowflake.com>
Co-authored-by: Piotr Bulawa <piotr.bulawa@snowflake.com>

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

@@ -45,7 +45,8 @@ public class SFLoginInput {
   private OCSPMode ocspMode;
   private HttpClientSettingsKey httpClientKey;
   private String privateKeyFile;
-  private String privateKeyFilePwd;
+  private String privateKeyBase64;
+  private String privateKeyPwd;
   private String inFlightCtx; // Opaque string sent for Snowsight account activation
 
   private boolean disableConsoleLogin = true;
@@ -325,22 +326,37 @@ 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;
   }
 
-  SFLoginInput setPrivateKeyFilePwd(String privateKeyFilePwd) {
-    this.privateKeyFilePwd = privateKeyFilePwd;
+  SFLoginInput setPrivateKeyPwd(String privateKeyPwd) {
+    this.privateKeyPwd = privateKeyPwd;
     return this;
   }
 
   String getPrivateKeyFile() {
     return privateKeyFile;
   }
 
-  String getPrivateKeyFilePwd() {
-    return privateKeyFilePwd;
+  String getPrivateKeyPwd() {
+    return privateKeyPwd;
+  }
+
+  boolean isPrivateKeyProvided() {
+    return (getPrivateKey() != null
+        || getPrivateKeyFile() != null
+        || getPrivateKeyBase64() != null);
   }
 
   public String getApplication() {

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,7 +453,14 @@ 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:
+        case PRIVATE_KEY_PWD:
           if (propertyValue != null) {
             privateKeyPassword = (String) propertyValue;
           }
@@ -553,7 +561,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: {}, connection timeout: {}, socket timeout: {}, tracing: {}, private key file: {}, private key file pwd is {},"
+            + " query timeout: {}, connection timeout: {}, socket timeout: {}, tracing: {}, private key file: {}, private key pwd is {},"
             + " enable_diagnostics: {}, diagnostics_allowlist_path: {},"
             + " session parameters: client store temporary credential: {}, gzip disabled: {}, browser response timeout: {}",
         connectionPropertiesMap.get(SFSessionProperty.SERVER_URL),
@@ -585,7 +593,12 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
         connectionPropertiesMap.get(SFSessionProperty.TRACING),
         connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE),
         SFLoggerUtil.isVariableProvided(
-            (String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD)),
+            (String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_BASE64)),
+        SFLoggerUtil.isVariableProvided(
+            (String)
+                connectionPropertiesMap.getOrDefault(
+                    SFSessionProperty.PRIVATE_KEY_PWD,
+                    connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD))),
         connectionPropertiesMap.get(SFSessionProperty.ENABLE_DIAGNOSTICS),
         connectionPropertiesMap.get(SFSessionProperty.DIAGNOSTICS_ALLOWLIST_FILE),
         sessionParametersMap.get(CLIENT_STORE_TEMPORARY_CREDENTIAL),
@@ -643,8 +656,13 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
         .setSessionParameters(sessionParametersMap)
         .setPrivateKey((PrivateKey) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY))
         .setPrivateKeyFile((String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE))
-        .setPrivateKeyFilePwd(
-            (String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD))
+        .setPrivateKeyBase64(
+            (String) connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_BASE64))
+        .setPrivateKeyPwd(
+            (String)
+                connectionPropertiesMap.getOrDefault(
+                    SFSessionProperty.PRIVATE_KEY_PWD,
+                    connectionPropertiesMap.get(SFSessionProperty.PRIVATE_KEY_FILE_PWD)))
         .setApplication((String) connectionPropertiesMap.get(SFSessionProperty.APPLICATION))
         .setServiceName(getServiceName())
         .setOCSPMode(getOCSPMode())
@@ -765,7 +783,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,7 +54,15 @@ 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),
+  /**
+   * @deprecated Use {@link #PRIVATE_KEY_PWD} for clarity. The given password will be used to
+   *     decrypt the private key value independent of whether that value is supplied as a file or
+   *     base64 string
+   */
+  @Deprecated
   PRIVATE_KEY_FILE_PWD("private_key_file_pwd", false, String.class),
+  PRIVATE_KEY_PWD("private_key_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;
   }
@@ -422,7 +422,8 @@ private static SFLoginOutput newSession(
             new SessionUtilKeyPair(
                 loginInput.getPrivateKey(),
                 loginInput.getPrivateKeyFile(),
-                loginInput.getPrivateKeyFilePwd(),
+                loginInput.getPrivateKeyBase64(),
+                loginInput.getPrivateKeyPwd(),
                 loginInput.getAccountName(),
                 loginInput.getUserName());
 
@@ -677,7 +678,8 @@ private static SFLoginOutput newSession(
                     new SessionUtilKeyPair(
                         loginInput.getPrivateKey(),
                         loginInput.getPrivateKeyFile(),
-                        loginInput.getPrivateKeyFilePwd(),
+                        loginInput.getPrivateKeyBase64(),
+                        loginInput.getPrivateKeyPwd(),
                         loginInput.getAccountName(),
                         loginInput.getUserName());
 
@@ -1725,7 +1727,8 @@ public static void resetOCSPUrlIfNecessary(String serverUrl) throws IOException
    *
    * @param privateKey private key
    * @param privateKeyFile path to private key file
-   * @param privateKeyFilePwd password for private key file
+   * @param privateKeyBase64 base64 encoded content of the private key file
+   * @param privateKeyPwd password for private key file or base64 encoded private key
    * @param accountName account name
    * @param userName user name
    * @return JWT token
@@ -1734,16 +1737,41 @@ public static void resetOCSPUrlIfNecessary(String serverUrl) throws IOException
   public static String generateJWTToken(
       PrivateKey privateKey,
       String privateKeyFile,
-      String privateKeyFilePwd,
+      String privateKeyBase64,
+      String privateKeyPwd,
       String accountName,
       String userName)
       throws SFException {
     SessionUtilKeyPair s =
         new SessionUtilKeyPair(
-            privateKey, privateKeyFile, privateKeyFilePwd, accountName, userName);
+            privateKey, privateKeyFile, privateKeyBase64, privateKeyPwd, 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.
    *