SNOW-1858529 Implement FLOE

src/main/java/net/snowflake/client/jdbc/cloud/storage/EncryptionProvider.java

@@ -41,8 +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 ThreadLocal<SecureRandom> secRnd =
-      new ThreadLocal<>().withInitial(SecureRandom::new);
+  private static final ThreadLocal<SecureRandom> SEC_RND =
+      ThreadLocal.withInitial(SecureRandom::new);
 
   /**
    * Decrypt a InputStream
@@ -70,7 +70,7 @@ public static InputStream decryptStream(
     byte[] kekBytes = Base64.getDecoder().decode(encMat.getQueryStageMasterKey());
     byte[] keyBytes = Base64.getDecoder().decode(keyBase64);
     byte[] ivBytes = Base64.getDecoder().decode(ivBase64);
-    SecretKey kek = new SecretKeySpec(kekBytes, 0, kekBytes.length, AES);
+    SecretKey kek = new SecretKeySpec(kekBytes, AES);
     Cipher keyCipher = Cipher.getInstance(KEY_CIPHER);
     keyCipher.init(Cipher.DECRYPT_MODE, kek);
     byte[] fileKeyBytes = keyCipher.doFinal(keyBytes);
@@ -98,7 +98,7 @@ public static void decrypt(
     // Decrypt file key
     {
       final Cipher keyCipher = Cipher.getInstance(KEY_CIPHER);
-      SecretKey kek = new SecretKeySpec(kekBytes, 0, kekBytes.length, AES);
+      SecretKey kek = new SecretKeySpec(kekBytes, AES);
       keyCipher.init(Cipher.DECRYPT_MODE, kek);
       byte[] fileKeyBytes = keyCipher.doFinal(keyBytes);
 
@@ -166,11 +166,11 @@ public static CipherInputStream encrypt(
 
       // Create IV
       ivData = new byte[blockSize];
-      secRnd.get().nextBytes(ivData);
+      SEC_RND.get().nextBytes(ivData);
       final IvParameterSpec iv = new IvParameterSpec(ivData);
 
       // Create file key
-      secRnd.get().nextBytes(fileKeyBytes);
+      SEC_RND.get().nextBytes(fileKeyBytes);
       SecretKey fileKey = new SecretKeySpec(fileKeyBytes, 0, keySize, AES);
 
       // Init cipher

src/main/java/net/snowflake/client/jdbc/cloud/storage/GcmEncryptionProvider.java

@@ -83,7 +83,7 @@ private static void initRandomIvsAndFileKey(
   private static byte[] encryptKey(byte[] kekBytes, byte[] keyBytes, byte[] keyIvData, byte[] aad)
       throws InvalidKeyException, InvalidAlgorithmParameterException, IllegalBlockSizeException,
           BadPaddingException, NoSuchPaddingException, NoSuchAlgorithmException {
-    SecretKey kek = new SecretKeySpec(kekBytes, 0, kekBytes.length, AES);
+    SecretKey kek = new SecretKeySpec(kekBytes, AES);
     GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_IN_BITS, keyIvData);
     Cipher keyCipher = Cipher.getInstance(KEY_CIPHER);
     keyCipher.init(Cipher.ENCRYPT_MODE, kek, gcmParameterSpec);
@@ -97,7 +97,7 @@ private static CipherInputStream encryptContent(
       InputStream src, byte[] keyBytes, byte[] dataIvBytes, byte[] aad)
       throws InvalidKeyException, InvalidAlgorithmParameterException, NoSuchPaddingException,
           NoSuchAlgorithmException {
-    SecretKey fileKey = new SecretKeySpec(keyBytes, 0, keyBytes.length, AES);
+    SecretKey fileKey = new SecretKeySpec(keyBytes, AES);
     GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_IN_BITS, dataIvBytes);
     Cipher fileCipher = Cipher.getInstance(FILE_CIPHER);
     fileCipher.init(Cipher.ENCRYPT_MODE, fileKey, gcmParameterSpec);
@@ -213,7 +213,7 @@ private static void decryptContentFromFile(
   private static byte[] decryptKey(byte[] kekBytes, byte[] ivBytes, byte[] keyBytes, byte[] aad)
       throws InvalidKeyException, InvalidAlgorithmParameterException, IllegalBlockSizeException,
           BadPaddingException, NoSuchPaddingException, NoSuchAlgorithmException {
-    SecretKey kek = new SecretKeySpec(kekBytes, 0, kekBytes.length, AES);
+    SecretKey kek = new SecretKeySpec(kekBytes, AES);
     GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_IN_BITS, ivBytes);
     Cipher keyCipher = Cipher.getInstance(KEY_CIPHER);
     keyCipher.init(Cipher.DECRYPT_MODE, kek, gcmParameterSpec);

src/main/java/net/snowflake/client/jdbc/cloud/storage/floe/Aead.java

@@ -0,0 +1,62 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.util.function.Supplier;
+import net.snowflake.client.jdbc.cloud.storage.floe.aead.AeadProvider;
+import net.snowflake.client.jdbc.cloud.storage.floe.aead.Gcm;
+
+public enum Aead {
+  AES_GCM_256((byte) 0, "AES", "AES/GCM/NoPadding", 32, 12, 16, () -> new Gcm(16));
+
+  private byte id;
+  private String jceKeyTypeName;
+  private String jceFullName;
+  private int keyLength;
+  private int ivLength;
+  private int authTagLength;
+  private Supplier<AeadProvider> aeadProvider;
+
+  Aead(
+      byte id,
+      String jceKeyTypeName,
+      String jceFullName,
+      int keyLength,
+      int ivLength,
+      int authTagLength,
+      Supplier<AeadProvider> aeadProvider) {
+    this.jceKeyTypeName = jceKeyTypeName;
+    this.jceFullName = jceFullName;
+    this.keyLength = keyLength;
+    this.id = id;
+    this.ivLength = ivLength;
+    this.authTagLength = authTagLength;
+    this.aeadProvider = aeadProvider;
+  }
+
+  byte getId() {
+    return id;
+  }
+
+  public String getJceKeyTypeName() {
+    return jceKeyTypeName;
+  }
+
+  String getJceFullName() {
+    return jceFullName;
+  }
+
+  int getKeyLength() {
+    return keyLength;
+  }
+
+  int getIvLength() {
+    return ivLength;
+  }
+
+  int getAuthTagLength() {
+    return authTagLength;
+  }
+
+  AeadProvider getAeadProvider() {
+    return aeadProvider.get();
+  }
+}

src/main/java/net/snowflake/client/jdbc/cloud/storage/floe/AeadAad.java

@@ -0,0 +1,26 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.nio.ByteBuffer;
+
+class AeadAad {
+  private final byte[] bytes;
+
+  private AeadAad(long segmentCounter, byte terminalityByte) {
+    ByteBuffer buf = ByteBuffer.allocate(9);
+    buf.putLong(segmentCounter);
+    buf.put(terminalityByte);
+    this.bytes = buf.array();
+  }
+
+  static AeadAad nonTerminal(long segmentCounter) {
+    return new AeadAad(segmentCounter, (byte) 0);
+  }
+
+  static AeadAad terminal(long segmentCounter) {
+    return new AeadAad(segmentCounter, (byte) 1);
+  }
+
+  byte[] getBytes() {
+    return bytes;
+  }
+}

src/main/java/net/snowflake/client/jdbc/cloud/storage/floe/AeadIv.java

@@ -0,0 +1,25 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.nio.ByteBuffer;
+
+class AeadIv {
+  private final byte[] bytes;
+
+  AeadIv(byte[] bytes) {
+    this.bytes = bytes;
+  }
+
+  static AeadIv generateRandom(FloeRandom floeRandom, int ivLength) {
+    return new AeadIv(floeRandom.ofLength(ivLength));
+  }
+
+  static AeadIv from(ByteBuffer buffer, int ivLength) {
+    byte[] bytes = new byte[ivLength];
+    buffer.get(bytes);
+    return new AeadIv(bytes);
+  }
+
+  byte[] getBytes() {
+    return bytes;
+  }
+}

src/main/java/net/snowflake/client/jdbc/cloud/storage/floe/AeadKey.java

@@ -0,0 +1,15 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import javax.crypto.SecretKey;
+
+class AeadKey {
+  private final SecretKey key;
+
+  AeadKey(SecretKey key) {
+    this.key = key;
+  }
+
+  SecretKey getKey() {
+    return key;
+  }
+}

src/main/java/net/snowflake/client/jdbc/cloud/storage/floe/BaseSegmentProcessor.java

@@ -0,0 +1,74 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.io.Closeable;
+import java.io.IOException;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
+
+abstract class BaseSegmentProcessor implements Closeable {
+  protected static final int NON_TERMINAL_SEGMENT_SIZE_MARKER = -1;
+  protected static final int headerTagLength = 32;
+
+  protected final FloeParameterSpec parameterSpec;
+  protected final FloeKey floeKey;
+  protected final FloeAad floeAad;
+
+  protected final KeyDerivator keyDerivator;
+
+  private AeadKey currentAeadKey;
+
+  private boolean isClosed;
+  private boolean completedExceptionally;
+
+  BaseSegmentProcessor(FloeParameterSpec parameterSpec, FloeKey floeKey, FloeAad floeAad) {
+    this.parameterSpec = parameterSpec;
+    this.floeKey = floeKey;
+    this.floeAad = floeAad;
+    this.keyDerivator = new KeyDerivator(parameterSpec);
+  }
+
+  protected AeadKey getKey(FloeKey floeKey, FloeIv floeIv, FloeAad floeAad, long segmentCounter) {
+    if (currentAeadKey == null || segmentCounter % parameterSpec.getKeyRotationModulo() == 0) {
+      // we don't need masking, because we derive a new key only when key rotation happens
+      currentAeadKey = deriveKey(floeKey, floeIv, floeAad, segmentCounter);
+    }
+    return currentAeadKey;
+  }
+
+  private AeadKey deriveKey(FloeKey floeKey, FloeIv floeIv, FloeAad floeAad, long segmentCounter) {
+    byte[] keyBytes =
+        keyDerivator.hkdfExpand(
+            floeKey,
+            floeIv,
+            floeAad,
+            new DekTagFloePurpose(segmentCounter),
+            parameterSpec.getAead().getKeyLength());
+    SecretKey key = new SecretKeySpec(keyBytes, parameterSpec.getAead().getJceKeyTypeName());
+    return new AeadKey(key);
+  }
+
+  protected void closeInternal() {
+    isClosed = true;
+  }
+
+  protected void markAsCompletedExceptionally() {
+    completedExceptionally = true;
+  }
+
+  protected void assertNotClosed() {
+    if (isClosed) {
+      throw new IllegalStateException("stream has already been closed");
+    }
+  }
+
+  @Override
+  public void close() throws IOException {
+    if (!isClosed && !completedExceptionally) {
+      throw new IllegalStateException("last segment was not processed");
+    }
+  }
+
+  protected boolean isClosed() {
+    return isClosed;
+  }
+}

src/main/java/net/snowflake/client/jdbc/cloud/storage/floe/Floe.java

@@ -0,0 +1,23 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import javax.crypto.SecretKey;
+
+public class Floe {
+  private final FloeParameterSpec parameterSpec;
+
+  private Floe(FloeParameterSpec parameterSpec) {
+    this.parameterSpec = parameterSpec;
+  }
+
+  public static Floe getInstance(FloeParameterSpec parameterSpec) {
+    return new Floe(parameterSpec);
+  }
+
+  public FloeEncryptor createEncryptor(SecretKey key, byte[] aad) {
+    return new FloeEncryptorImpl(parameterSpec, new FloeKey(key), new FloeAad(aad));
+  }
+
+  public FloeDecryptor createDecryptor(SecretKey key, byte[] aad, byte[] floeHeader) {
+    return new FloeDecryptorImpl(parameterSpec, new FloeKey(key), new FloeAad(aad), floeHeader);
+  }
+}

src/main/java/net/snowflake/client/jdbc/cloud/storage/floe/FloeAad.java

@@ -0,0 +1,16 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.util.Optional;
+
+class FloeAad {
+  private static final byte[] EMPTY_AAD = new byte[0];
+  private final byte[] aad;
+
+  FloeAad(byte[] aad) {
+    this.aad = Optional.ofNullable(aad).orElse(EMPTY_AAD);
+  }
+
+  byte[] getBytes() {
+    return aad;
+  }
+}

src/main/java/net/snowflake/client/jdbc/cloud/storage/floe/FloeDecryptor.java

@@ -0,0 +1,7 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+public interface FloeDecryptor extends AutoCloseable {
+  byte[] processSegment(byte[] ciphertext);
+
+  boolean isClosed();
+}