SNOW-896818 Decrypt Private Keys With BouncyCastle

Added the JVM argument -Dnet.snowflake.jdbc.enableBouncyCastle to allow the JDBC
driver to leverage the BouncyCastle provider to decrypt private keys instead of
using the default security providers in the JDK. This helps work around a
limitation with being able to decrypt private keys with PBES2 parameters.
Ref: https://bugs.openjdk.org/browse/JDK-8228481

Author: sfc-gh-wfateem

FIPS/src/test/java/net/snowflake/client/jdbc/ConnectionFipsIT.java

@@ -23,6 +23,7 @@
 import net.snowflake.client.ConditionalIgnoreRule;
 import net.snowflake.client.RunningOnGithubActions;
 import net.snowflake.client.category.TestCategoryFips;
+import net.snowflake.client.core.SecurityUtil;
 import org.apache.commons.codec.binary.Base64;
 import org.bouncycastle.crypto.CryptoServicesRegistrar;
 import org.bouncycastle.crypto.fips.FipsStatus;
@@ -161,7 +162,7 @@ public static void setup() throws Exception {
     }
 
     // attempts an SSL connection to Google
-    //connectToGoogle();
+    // connectToGoogle();
   }
 
   @AfterClass
@@ -205,9 +206,10 @@ public static void teardown() throws Exception {
           JAVA_SYSTEM_PROPERTY_SSL_TRUSTSTORE_TYPE,
           JAVA_SYSTEM_PROPERTY_SSL_TRUSTSTORE_TYPE_ORIGINAL_VALUE);
     }
+    System.clearProperty(SecurityUtil.ENABLE_BOUNCYCASTLE_PROVIDER_JVM);
 
     // attempts an SSL connection to Google
-    //connectToGoogle();
+    // connectToGoogle();
   }
 
   @Test
@@ -319,6 +321,22 @@ public void connectWithFipsAndPut() throws Exception {
     }
   }
 
+  /** Added in > 3.15.1 */
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubActions.class)
+  public void connectWithFipsKeyPairWithBouncyCastle() throws Exception {
+    System.setProperty(SecurityUtil.ENABLE_BOUNCYCASTLE_PROVIDER_JVM, "true");
+    connectWithFipsKeyPair();
+  }
+
+  /** Added in > 3.15.1 */
+  @Test
+  @ConditionalIgnoreRule.ConditionalIgnore(condition = RunningOnGithubActions.class)
+  public void testConnectUsingKeyPairWithBouncyCastle() throws Exception {
+    System.setProperty(SecurityUtil.ENABLE_BOUNCYCASTLE_PROVIDER_JVM, "true");
+    testConnectUsingKeyPair();
+  }
+
   private static void connectToGoogle() throws Exception {
     URL url = new URL("https://www.google.com/");
     HttpsURLConnection con = (HttpsURLConnection) url.openConnection();

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

@@ -142,6 +142,8 @@ public abstract class SFBaseSession {
 
   private Map<String, Object> commonParameters;
 
+  private boolean isJdbcArrowTreatDecimalAsInt = true;
+
   protected SFBaseSession(SFConnectionHandler sfConnectionHandler) {
     this.sfConnectionHandler = sfConnectionHandler;
   }
@@ -270,6 +272,14 @@ public void setJdbcTreatDecimalAsInt(boolean jdbcTreatDecimalAsInt) {
     isJdbcTreatDecimalAsInt = jdbcTreatDecimalAsInt;
   }
 
+  public boolean isJdbcArrowTreatDecimalAsInt() {
+    return isJdbcArrowTreatDecimalAsInt;
+  }
+
+  public void setJdbcArrowTreatDecimalAsInt(boolean jdbcArrowTreatDecimalAsInt) {
+    isJdbcArrowTreatDecimalAsInt = jdbcArrowTreatDecimalAsInt;
+  }
+
   public String getServerUrl() {
     if (connectionPropertiesMap.containsKey(SFSessionProperty.SERVER_URL)) {
       return (String) connectionPropertiesMap.get(SFSessionProperty.SERVER_URL);

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

@@ -87,8 +87,6 @@ public class SFSession extends SFBaseSession {
 
   private SFClientConfig sfClientConfig;
 
-  private SecurityUtil securityUtil;
-
   /**
    * Amount of seconds a user is willing to tolerate for establishing the connection with database.
    * In our case, it means the first login request to get authorization token.
@@ -147,7 +145,6 @@ public SFSession() {
 
   public SFSession(DefaultSFConnectionHandler sfConnectionHandler) {
     super(sfConnectionHandler);
-    securityUtil = new SecurityUtil();
   }
 
   /**
@@ -478,6 +475,12 @@ public void addSFSessionProperty(String propertyName, Object propertyValue) thro
           }
           break;
 
+        case JDBC_ARROW_TREAT_DECIMAL_AS_INT:
+          if (propertyValue != null) {
+            setJdbcArrowTreatDecimalAsInt(getBooleanValue(propertyValue));
+          }
+          break;
+
         default:
           break;
       }
@@ -616,7 +619,6 @@ public synchronized void open() throws SFException, SnowflakeSQLException {
       TelemetryService.disable();
     }
 
-    securityUtil.addBouncyCastleProvider();
     // propagate OCSP mode to SFTrustManager. Note OCSP setting is global on JVM.
     HttpUtil.initHttpClient(httpClientSettingsKey, null);
     SFLoginOutput loginOutput =

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

@@ -82,7 +82,7 @@ public enum SFSessionProperty {
 
   DISABLE_GCS_DEFAULT_CREDENTIALS("disableGcsDefaultCredentials", false, Boolean.class),
 
-  ENABLE_BOUNCY_CASTLE("enableBouncyCastle", false, Boolean.class);
+  JDBC_ARROW_TREAT_DECIMAL_AS_INT("JDBC_ARROW_TREAT_DECIMAL_AS_INT", false, Boolean.class);
 
   // property key in string
   private String propertyKey;

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

@@ -6,20 +6,22 @@
 import net.snowflake.client.log.SFLogger;
 import net.snowflake.client.log.SFLoggerFactory;
 
+@SnowflakeJdbcInternalApi
 public class SecurityUtil {
 
   private static final SFLogger LOGGER = SFLoggerFactory.getLogger(SecurityUtil.class);
 
-  /** provider name */
-  private final String BOUNCY_CASTLE_PROVIDER = "BC";
-
   /** provider name for FIPS */
-  private final String BOUNCY_CASTLE_FIPS_PROVIDER = "BCFIPS";
+  public static final String BOUNCY_CASTLE_FIPS_PROVIDER = "BCFIPS";
 
+  public static final String BOUNCY_CASTLE_PROVIDER = "BC";
   private static final String DEFAULT_SECURITY_PROVIDER_NAME =
       "org.bouncycastle.jce.provider.BouncyCastleProvider";
 
-  public void addBouncyCastleProvider() {
+  public static final String ENABLE_BOUNCYCASTLE_PROVIDER_JVM =
+      "net.snowflake.jdbc.enableBouncyCastle";
+
+  public static void addBouncyCastleProvider() {
     // Add Bouncy Castle to the list of security providers. This is required to
     // verify the signature on OCSP response and attached certificates.
     // It is also required to decrypt password protected private keys.
@@ -31,7 +33,7 @@ public void addBouncyCastleProvider() {
     }
   }
 
-  public Provider instantiateSecurityProvider() {
+  private static Provider instantiateSecurityProvider() {
 
     try {
       Class klass = Class.forName(DEFAULT_SECURITY_PROVIDER_NAME);
@@ -50,7 +52,7 @@ public Provider instantiateSecurityProvider() {
                   + "import BouncyCastleFipsProvider in the application.",
               DEFAULT_SECURITY_PROVIDER_NAME, ex.getMessage());
       LOGGER.error(errMsg, true);
-      throw new RuntimeException(errMsg);
+      throw new RuntimeException(errMsg, ex);
     }
   }
 }

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

@@ -39,6 +39,7 @@
 import net.snowflake.client.log.SFLoggerFactory;
 import org.apache.commons.codec.binary.Base64;
 import org.bouncycastle.asn1.pkcs.PrivateKeyInfo;
+import org.bouncycastle.openssl.PEMKeyPair;
 import org.bouncycastle.openssl.PEMParser;
 import org.bouncycastle.openssl.jcajce.JcaPEMKeyConverter;
 import org.bouncycastle.openssl.jcajce.JceOpenSSLPKCS8DecryptorProviderBuilder;
@@ -67,21 +68,13 @@ class SessionUtilKeyPair {
 
   private Provider SecurityProvider = null;
 
-  private SecretKeyFactory secretKeyFactory = null;
-
   private static final String ISSUER_FMT = "%s.%s.%s";
 
   private static final String SUBJECT_FMT = "%s.%s";
 
   private static final int JWT_DEFAULT_AUTH_TIMEOUT = 10;
 
-  /** provider name */
-  private static final String BOUNCY_CASTLE_PROVIDER = "BC";
-
-  /** provider name for FIPS */
-  private static final String BOUNCY_CASTLE_FIPS_PROVIDER = "BCFIPS";
-
-  private boolean ENABLE_BOUNCYCASTLE_PROVIDER = true;
+  private boolean isBouncyCastleProviderEnabled = false;
 
   SessionUtilKeyPair(
       PrivateKey privateKey,
@@ -92,10 +85,14 @@ class SessionUtilKeyPair {
       throws SFException {
     this.userName = userName.toUpperCase();
     this.accountName = accountName.toUpperCase();
-
+    String enableBouncyCastleJvm =
+        System.getProperty(SecurityUtil.ENABLE_BOUNCYCASTLE_PROVIDER_JVM);
+    if (enableBouncyCastleJvm != null) {
+      isBouncyCastleProviderEnabled = enableBouncyCastleJvm.equalsIgnoreCase("true");
+    }
     // check if in FIPS mode
     for (Provider p : Security.getProviders()) {
-      if (BOUNCY_CASTLE_FIPS_PROVIDER.equals(p.getName())) {
+      if (SecurityUtil.BOUNCY_CASTLE_FIPS_PROVIDER.equals(p.getName())) {
         this.isFipsMode = true;
         this.SecurityProvider = p;
         break;
@@ -151,11 +148,11 @@ private SecretKeyFactory getSecretKeyFactory(String algorithm) throws NoSuchAlgo
   private PrivateKey extractPrivateKeyFromFile(String privateKeyFile, String privateKeyFilePwd)
       throws SFException {
 
-    if (ENABLE_BOUNCYCASTLE_PROVIDER) {
+    if (isBouncyCastleProviderEnabled) {
       try {
         return extractPrivateKeyWithBouncyCastle(privateKeyFile, privateKeyFilePwd);
       } catch (IOException | PKCSException | OperatorCreationException e) {
-        logger.error("Could not extract private key using Bouncy Castle provider");
+        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 {
@@ -168,7 +165,8 @@ private PrivateKey extractPrivateKeyFromFile(String privateKeyFile, String priva
           | NullPointerException
           | InvalidKeyException e) {
         logger.error(
-            "Could not extract private key. Try setting " + ENABLE_BOUNCYCASTLE_PROVIDER + "=TRUE");
+            "Could not extract private key. Try setting the JVM argument: " + "-D{}" + "=TRUE",
+            SecurityUtil.ENABLE_BOUNCYCASTLE_PROVIDER_JVM);
         throw new SFException(
             e,
             ErrorCode.INVALID_OR_UNSUPPORTED_PRIVATE_KEY,
@@ -246,12 +244,20 @@ private PrivateKey extractPrivateKeyWithBouncyCastle(
       InputDecryptorProvider pkcs8Prov =
           new JceOpenSSLPKCS8DecryptorProviderBuilder().build(privateKeyFilePwd.toCharArray());
       privateKeyInfo = encryptedPrivateKeyInfo.decryptPrivateKeyInfo(pkcs8Prov);
+    } else if (pemObject instanceof PEMKeyPair) {
+      // PKCS#1 private key
+      privateKeyInfo = ((PEMKeyPair) pemObject).getPrivateKeyInfo();
     } else if (pemObject instanceof PrivateKeyInfo) {
       // Handle the case where the private key is unencrypted.
       privateKeyInfo = (PrivateKeyInfo) pemObject;
     }
     pemParser.close();
-    JcaPEMKeyConverter converter = new JcaPEMKeyConverter().setProvider(BOUNCY_CASTLE_PROVIDER);
+    JcaPEMKeyConverter converter =
+        new JcaPEMKeyConverter()
+            .setProvider(
+                isFipsMode
+                    ? SecurityUtil.BOUNCY_CASTLE_FIPS_PROVIDER
+                    : SecurityUtil.BOUNCY_CASTLE_PROVIDER);
     return converter.getPrivateKey(privateKeyInfo);
   }
 
@@ -260,23 +266,36 @@ private PrivateKey extractPrivateKeyWithJdk(String privateKeyFile, String privat
     String privateKeyContent = new String(Files.readAllBytes(Paths.get(privateKeyFile)));
     if (Strings.isNullOrEmpty(privateKeyFilePwd)) {
       // unencrypted private key file
-      PemReader pr = new PemReader(new StringReader(privateKeyContent));
-      byte[] decoded = pr.readPemObject().getContent();
-      pr.close();
-      PKCS8EncodedKeySpec encodedKeySpec = new PKCS8EncodedKeySpec(decoded);
-      KeyFactory keyFactory = getKeyFactoryInstance();
-      return keyFactory.generatePrivate(encodedKeySpec);
+      return generatePrivateKey(false, privateKeyContent, privateKeyFilePwd);
     } else {
       // encrypted private key file
-      PemReader pr = new PemReader(new StringReader(privateKeyContent));
-      byte[] decoded = pr.readPemObject().getContent();
-      pr.close();
-      EncryptedPrivateKeyInfo pkInfo = new EncryptedPrivateKeyInfo(decoded);
-      PBEKeySpec keySpec = new PBEKeySpec(privateKeyFilePwd.toCharArray());
-      SecretKeyFactory pbeKeyFactory = this.getSecretKeyFactory(pkInfo.getAlgName());
-      PKCS8EncodedKeySpec encodedKeySpec = pkInfo.getKeySpec(pbeKeyFactory.generateSecret(keySpec));
-      KeyFactory keyFactory = getKeyFactoryInstance();
-      return keyFactory.generatePrivate(encodedKeySpec);
+      return generatePrivateKey(true, privateKeyContent, privateKeyFilePwd);
+    }
+  }
+
+  private PrivateKey generatePrivateKey(
+      boolean isEncrypted, String privateKeyContent, String privateKeyFilePwd)
+      throws IOException, NoSuchAlgorithmException, InvalidKeySpecException, InvalidKeyException {
+    if (isEncrypted) {
+      try (PemReader pr = new PemReader(new StringReader(privateKeyContent))) {
+        byte[] decoded = pr.readPemObject().getContent();
+        pr.close();
+        EncryptedPrivateKeyInfo pkInfo = new EncryptedPrivateKeyInfo(decoded);
+        PBEKeySpec keySpec = new PBEKeySpec(privateKeyFilePwd.toCharArray());
+        SecretKeyFactory pbeKeyFactory = this.getSecretKeyFactory(pkInfo.getAlgName());
+        PKCS8EncodedKeySpec encodedKeySpec =
+            pkInfo.getKeySpec(pbeKeyFactory.generateSecret(keySpec));
+        KeyFactory keyFactory = getKeyFactoryInstance();
+        return keyFactory.generatePrivate(encodedKeySpec);
+      }
+    } else {
+      try (PemReader pr = new PemReader(new StringReader(privateKeyContent))) {
+        byte[] decoded = pr.readPemObject().getContent();
+        pr.close();
+        PKCS8EncodedKeySpec encodedKeySpec = new PKCS8EncodedKeySpec(decoded);
+        KeyFactory keyFactory = getKeyFactoryInstance();
+        return keyFactory.generatePrivate(encodedKeySpec);
+      }
     }
   }
 }

src/main/java/net/snowflake/client/core/arrow/AbstractArrowVectorConverter.java

@@ -38,6 +38,8 @@ abstract class AbstractArrowVectorConverter implements ArrowVectorConverter {
 
   protected TimeZone sessionTimeZone;
 
+  private boolean shouldTreatDecimalAsInt;
+
   /** Field names of the struct vectors used by timestamp */
   public static final String FIELD_NAME_EPOCH = "epoch"; // seconds since epoch
 
@@ -53,6 +55,11 @@ abstract class AbstractArrowVectorConverter implements ArrowVectorConverter {
     this.valueVector = valueVector;
     this.columnIndex = vectorIndex + 1;
     this.context = context;
+    this.shouldTreatDecimalAsInt =
+        context == null
+            || context.getSession() == null
+            || context.getSession().isJdbcArrowTreatDecimalAsInt()
+            || context.getSession().isJdbcTreatDecimalAsInt();
   }
 
   @Override
@@ -146,6 +153,10 @@ public BigDecimal toBigDecimal(int index) throws SFException {
         ErrorCode.INVALID_VALUE_CONVERT, logicalTypeStr, SnowflakeUtil.BIG_DECIMAL_STR, "");
   }
 
+  boolean shouldTreatDecimalAsInt() {
+    return shouldTreatDecimalAsInt;
+  }
+
   @Override
   public void setTreatNTZAsUTC(boolean isUTC) {
     this.treatNTZasUTC = isUTC;

src/main/java/net/snowflake/client/core/arrow/BigIntToFixedConverter.java

@@ -137,9 +137,10 @@ public BigDecimal toBigDecimal(int index) {
   public Object toObject(int index) throws SFException {
     if (bigIntVector.isNull(index)) {
       return null;
-    } else {
-      return getLong(index);
+    } else if (!shouldTreatDecimalAsInt()) {
+      return BigDecimal.valueOf(getLong(index), sfScale);
     }
+    return getLong(index);
   }
 
   @Override

src/main/java/net/snowflake/client/core/arrow/IntToFixedConverter.java

@@ -105,7 +105,12 @@ public BigDecimal toBigDecimal(int index) throws SFException {
 
   @Override
   public Object toObject(int index) throws SFException {
-    return isNull(index) ? null : (long) getInt(index);
+    if (isNull(index)) {
+      return null;
+    } else if (!shouldTreatDecimalAsInt()) {
+      return BigDecimal.valueOf((long) getInt(index), sfScale);
+    }
+    return (long) getInt(index);
   }
 
   @Override

src/main/java/net/snowflake/client/core/arrow/SmallIntToFixedConverter.java

@@ -99,6 +99,8 @@ public double toDouble(int index) throws SFException {
   public Object toObject(int index) throws SFException {
     if (isNull(index)) {
       return null;
+    } else if (!shouldTreatDecimalAsInt()) {
+      return BigDecimal.valueOf((long) getShort(index), sfScale);
     }
     return (long) getShort(index);
   }

src/main/java/net/snowflake/client/core/arrow/TinyIntToFixedConverter.java

@@ -90,7 +90,12 @@ public BigDecimal toBigDecimal(int index) throws SFException {
 
   @Override
   public Object toObject(int index) throws SFException {
-    return isNull(index) ? null : (long) toByte(index);
+    if (isNull(index)) {
+      return null;
+    } else if (!shouldTreatDecimalAsInt()) {
+      return BigDecimal.valueOf((long) getByte(index), sfScale);
+    }
+    return (long) toByte(index);
   }
 
   @Override