Support accepting a private key as a base64 encoded string with the n…

…ew session property key "private_key_base64" Rename private key password variable for clarity Add deprecation to private key file password
Qualytics · Aug 10, 2024 · dd04f7d · dd04f7d
1 parent e2a092d
commit dd04f7d
Show file tree

Hide file tree

Showing 9 changed files with 368 additions and 53 deletions.
diff --git a/src/main/java/net/snowflake/client/core/SFLoginInput.java b/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;
   }
 
+  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 getPrivateKeyBase64() {
+    return privateKeyBase64;
+  }
+
+  String getPrivateKeyPwd() {
+    return privateKeyPwd;
+  }
+
+  boolean isPrivateKeyProvided() {
+    return (getPrivateKey() != null
+        || getPrivateKeyFile() != null
+        || getPrivateKeyBase64() != null);
   }
 
   public String getApplication() {

diff --git a/src/main/java/net/snowflake/client/core/SFSession.java b/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;
           }
@@ -583,7 +591,7 @@ 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.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),
@@ -631,8 +639,9 @@ 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())
@@ -750,7 +759,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);
   }
 

diff --git a/src/main/java/net/snowflake/client/core/SFSessionProperty.java b/src/main/java/net/snowflake/client/core/SFSessionProperty.java
@@ -54,7 +54,14 @@ 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),
+  /**
+   * @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_BASE64("private_key_base64", 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),
 

diff --git a/src/main/java/net/snowflake/client/core/SessionUtil.java b/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,7 +421,8 @@ private static SFLoginOutput newSession(
             new SessionUtilKeyPair(
                 loginInput.getPrivateKey(),
                 loginInput.getPrivateKeyFile(),
-                loginInput.getPrivateKeyFilePwd(),
+                loginInput.getPrivateKeyBase64(),
+                loginInput.getPrivateKeyPwd(),
                 loginInput.getAccountName(),
                 loginInput.getUserName());
 
@@ -676,7 +677,8 @@ private static SFLoginOutput newSession(
                     new SessionUtilKeyPair(
                         loginInput.getPrivateKey(),
                         loginInput.getPrivateKeyFile(),
-                        loginInput.getPrivateKeyFilePwd(),
+                        loginInput.getPrivateKeyBase64(),
+                        loginInput.getPrivateKeyPwd(),
                         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,13 +1735,39 @@ 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
+   *
+   * @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 {
+    SessionUtilKeyPair s =
+        new SessionUtilKeyPair(
+            privateKey, privateKeyFile, null, privateKeyFilePwd, accountName, userName);
     return s.issueJwtToken();
   }
 

diff --git a/src/main/java/net/snowflake/client/core/SessionUtilKeyPair.java b/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,17 +101,30 @@ class SessionUtilKeyPair {
       }
     }
 
-    // if there is both a file and a private key, there is a problem
+    // Ensure that we only received one of: privateKey, privateKeyFile, or privateKeyBase64
     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.");
+          "Cannot have both private key object and private key file.");
+    } else if (!Strings.isNullOrEmpty(privateKeyBase64) && privateKey != null) {
+      throw new SFException(
+          ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
+          "Cannot have both private key object and private key string value.");
+    } else 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 string value.");
     } else {
-      // if privateKeyFile has a value and privateKey is null
-      this.privateKey =
-          Strings.isNullOrEmpty(privateKeyFile)
-              ? privateKey
-              : extractPrivateKeyFromFile(privateKeyFile, privateKeyFilePwd);
+      if (!Strings.isNullOrEmpty(privateKeyBase64)) {
+        // privateKeyBase64 has a value and other options for passing private key are null
+        this.privateKey = extractPrivateKeyFromBase64(privateKeyBase64, privateKeyFilePwd);
+      } else {
+        // either extract from privateKeyFile or use the passed object
+        this.privateKey =
+            Strings.isNullOrEmpty(privateKeyFile)
+                ? privateKey
+                : extractPrivateKeyFromFile(privateKeyFile, privateKeyFilePwd);
+      }
     }
     // construct public key from raw bytes
     if (this.privateKey instanceof RSAPrivateCrtKey) {
@@ -148,33 +162,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 privateKeyBytesPwd)
+      throws SFException {
     if (isBouncyCastleProviderEnabled) {
       try {
-        return extractPrivateKeyWithBouncyCastle(privateKeyFile, privateKeyFilePwd);
+        return extractPrivateKeyWithBouncyCastle(privateKeyBytes, privateKeyBytesPwd);
       } 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, privateKeyBytesPwd);
       } 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 privateKeyBytesPwd)
+      throws SFException {
+    byte[] decodedKey = Base64.decodeBase64(privateKeyBase64);
+    return extractPrivateKeyFromBytes(decodedKey, privateKeyBytesPwd);
+  }
+
   public String issueJwtToken() throws SFException {
     JWTClaimsSet.Builder builder = new JWTClaimsSet.Builder();
     String sub = String.format(SUBJECT_FMT, this.accountName, this.userName);
@@ -232,13 +265,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);
+
     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 +296,9 @@ 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));
+    String privateKeyContent = new String(privateKeyFileBytes, StandardCharsets.UTF_8);
     if (Strings.isNullOrEmpty(privateKeyFilePwd)) {
       // unencrypted private key file
       return generatePrivateKey(false, privateKeyContent, privateKeyFilePwd);

diff --git a/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java b/src/main/java/net/snowflake/client/jdbc/SnowflakeBasicDataSource.java
@@ -232,7 +232,15 @@ public void setPrivateKeyFile(String location, String password) {
     this.setAuthenticator(AUTHENTICATOR_SNOWFLAKE_JWT);
     this.properties.put(SFSessionProperty.PRIVATE_KEY_FILE.getPropertyKey(), location);
     if (!Strings.isNullOrEmpty(password)) {
-      this.properties.put(SFSessionProperty.PRIVATE_KEY_FILE_PWD.getPropertyKey(), password);
+      this.properties.put(SFSessionProperty.PRIVATE_KEY_PWD.getPropertyKey(), 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_PWD.getPropertyKey(), password);
     }
   }