From cf1b6481cdebb9709f8d1c136c80731983c7d596 Mon Sep 17 00:00:00 2001
From: Piotr Fus <piotr.fus@snowflake.com>
Date: Mon, 9 Dec 2024 13:48:37 +0100
Subject: [PATCH] SNOW-1850888 Fix secure random initialization

---
 .../cloud/storage/EncryptionProvider.java     | 21 ++++---------------
 .../cloud/storage/GcmEncryptionProvider.java  | 21 +++++--------------
 2 files changed, 9 insertions(+), 33 deletions(-)

diff --git a/src/main/java/net/snowflake/client/jdbc/cloud/storage/EncryptionProvider.java b/src/main/java/net/snowflake/client/jdbc/cloud/storage/EncryptionProvider.java
index 7acb2fc4a..faa74ce86 100644
--- a/src/main/java/net/snowflake/client/jdbc/cloud/storage/EncryptionProvider.java
+++ b/src/main/java/net/snowflake/client/jdbc/cloud/storage/EncryptionProvider.java
@@ -41,7 +41,8 @@ public class EncryptionProvider {
   private static final String FILE_CIPHER = "AES/CBC/PKCS5Padding";
   private static final String KEY_CIPHER = "AES/ECB/PKCS5Padding";
   private static final int BUFFER_SIZE = 2 * 1024 * 1024; // 2 MB
-  private static SecureRandom secRnd;
+  private static ThreadLocal<SecureRandom> secRnd =
+      new ThreadLocal<>().withInitial(SecureRandom::new);
 
   /**
    * Decrypt a InputStream
@@ -165,11 +166,11 @@ public static CipherInputStream encrypt(
 
       // Create IV
       ivData = new byte[blockSize];
-      getSecRnd().nextBytes(ivData);
+      secRnd.get().nextBytes(ivData);
       final IvParameterSpec iv = new IvParameterSpec(ivData);
 
       // Create file key
-      getSecRnd().nextBytes(fileKeyBytes);
+      secRnd.get().nextBytes(fileKeyBytes);
       SecretKey fileKey = new SecretKeySpec(fileKeyBytes, 0, keySize, AES);
 
       // Init cipher
@@ -199,18 +200,4 @@ public static CipherInputStream encrypt(
 
     return cis;
   }
-
-  /*
-   * getSecRnd
-   * Gets a random number for encryption purposes.
-   */
-  private static synchronized SecureRandom getSecRnd()
-      throws NoSuchAlgorithmException, NoSuchProviderException {
-    if (secRnd == null) {
-      secRnd = SecureRandom.getInstance("SHA1PRNG");
-      byte[] bytes = new byte[10];
-      secRnd.nextBytes(bytes);
-    }
-    return secRnd;
-  }
 }
diff --git a/src/main/java/net/snowflake/client/jdbc/cloud/storage/GcmEncryptionProvider.java b/src/main/java/net/snowflake/client/jdbc/cloud/storage/GcmEncryptionProvider.java
index c3f53c0ea..b4a4682d8 100644
--- a/src/main/java/net/snowflake/client/jdbc/cloud/storage/GcmEncryptionProvider.java
+++ b/src/main/java/net/snowflake/client/jdbc/cloud/storage/GcmEncryptionProvider.java
@@ -36,21 +36,10 @@ class GcmEncryptionProvider {
   private static final String FILE_CIPHER = "AES/GCM/NoPadding";
   private static final String KEY_CIPHER = "AES/GCM/NoPadding";
   private static final int BUFFER_SIZE = 8 * 1024 * 1024; // 2 MB
-  private static final int blockSize;
-  private static final SecureRandom random;
+  private static final ThreadLocal<SecureRandom> random =
+      new ThreadLocal<>().withInitial(SecureRandom::new);
   private static final Base64.Decoder base64Decoder = Base64.getDecoder();
 
-  static {
-    try {
-      Cipher fileCipher = Cipher.getInstance(FILE_CIPHER);
-      blockSize = fileCipher.getBlockSize();
-
-      random = SecureRandom.getInstance("SHA1PRNG");
-    } catch (NoSuchAlgorithmException | NoSuchPaddingException e) {
-      throw new ExceptionInInitializerError(e);
-    }
-  }
-
   static InputStream encrypt(
       StorageObjectMetadata meta,
       long originalContentLength,
@@ -86,9 +75,9 @@ static InputStream encrypt(
 
   private static void initRandomIvsAndFileKey(
       byte[] dataIvData, byte[] fileKeyIvData, byte[] fileKeyBytes) {
-    random.nextBytes(dataIvData);
-    random.nextBytes(fileKeyIvData);
-    random.nextBytes(fileKeyBytes);
+    random.get().nextBytes(dataIvData);
+    random.get().nextBytes(fileKeyIvData);
+    random.get().nextBytes(fileKeyBytes);
   }
 
   private static byte[] encryptKey(byte[] kekBytes, byte[] keyBytes, byte[] keyIvData, byte[] aad)