SNOW-1858529 Implement header handling in FLOE

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

@@ -41,8 +41,7 @@ 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 ThreadLocal<SecureRandom> secRnd = ThreadLocal.withInitial(SecureRandom::new);
 
   /**
    * Decrypt a InputStream

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

@@ -37,7 +37,7 @@ class GcmEncryptionProvider {
   private static final String KEY_CIPHER = "AES/GCM/NoPadding";
   private static final int BUFFER_SIZE = 8 * 1024 * 1024; // 2 MB
   private static final ThreadLocal<SecureRandom> random =
-      new ThreadLocal<>().withInitial(SecureRandom::new);
+      ThreadLocal.withInitial(SecureRandom::new);
   private static final Base64.Decoder base64Decoder = Base64.getDecoder();
 
   static InputStream encrypt(

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

@@ -0,0 +1,16 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+public enum Aead {
+  AES_GCM_128((byte) 0),
+  AES_GCM_256((byte) 1);
+
+  private byte id;
+
+  Aead(byte id) {
+    this.id = id;
+  }
+
+  byte getId() {
+    return id;
+  }
+}

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,15 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.util.Optional;
+
+class FloeAad {
+  private final byte[] aad;
+
+  FloeAad(byte[] aad) {
+    this.aad = Optional.ofNullable(aad).orElse(new byte[0]);
+  }
+
+  byte[] getBytes() {
+    return aad;
+  }
+}

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

@@ -0,0 +1,18 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+abstract class FloeBase {
+  protected static final int headerTagLength = 32;
+
+  protected final FloeParameterSpec parameterSpec;
+  protected final FloeKey floeKey;
+  protected final FloeAad floeAad;
+
+  protected final FloeKdf floeKdf;
+
+  FloeBase(FloeParameterSpec parameterSpec, FloeKey floeKey, FloeAad floeAad) {
+    this.parameterSpec = parameterSpec;
+    this.floeKey = floeKey;
+    this.floeAad = floeAad;
+    this.floeKdf = new FloeKdf(parameterSpec);
+  }
+}

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

@@ -0,0 +1,3 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+public interface FloeDecryptor {}

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

@@ -0,0 +1,40 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+
+public class FloeDecryptorImpl extends FloeBase implements FloeDecryptor {
+  FloeDecryptorImpl(
+      FloeParameterSpec parameterSpec, FloeKey floeKey, FloeAad floeAad, byte[] floeHeaderAsBytes) {
+    super(parameterSpec, floeKey, floeAad);
+    validate(floeHeaderAsBytes);
+  }
+
+  public void validate(byte[] floeHeaderAsBytes) {
+    byte[] encodedParams = parameterSpec.paramEncode();
+    if (floeHeaderAsBytes.length
+        != encodedParams.length + parameterSpec.getFloeIvLength().getLength() + headerTagLength) {
+      throw new IllegalArgumentException("invalid header length");
+    }
+    ByteBuffer floeHeader = ByteBuffer.wrap(floeHeaderAsBytes);
+
+    byte[] encodedParamsFromHeader = new byte[10];
+    floeHeader.get(encodedParamsFromHeader, 0, encodedParamsFromHeader.length);
+    if (!Arrays.equals(encodedParams, encodedParamsFromHeader)) {
+      throw new IllegalArgumentException("invalid parameters header");
+    }
+
+    byte[] floeIvBytes = new byte[parameterSpec.getFloeIvLength().getLength()];
+    floeHeader.get(floeIvBytes, 0, floeIvBytes.length);
+    FloeIv floeIv = new FloeIv(floeIvBytes);
+
+    byte[] headerTagFromHeader = new byte[headerTagLength];
+    floeHeader.get(headerTagFromHeader, 0, headerTagFromHeader.length);
+
+    byte[] headerTag =
+        floeKdf.hkdfExpand(floeKey, floeIv, floeAad, FloePurpose.HEADER_TAG, headerTagLength);
+    if (!Arrays.equals(headerTag, headerTagFromHeader)) {
+      throw new IllegalArgumentException("invalid header tag");
+    }
+  }
+}

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

@@ -0,0 +1,5 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+public interface FloeEncryptor {
+  byte[] getHeader();
+}

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

@@ -0,0 +1,38 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.nio.ByteBuffer;
+
+class FloeEncryptorImpl extends FloeBase implements FloeEncryptor {
+  private final FloeIv floeIv;
+
+  private final byte[] header;
+
+  FloeEncryptorImpl(FloeParameterSpec parameterSpec, FloeKey floeKey, FloeAad floeAad) {
+    super(parameterSpec, floeKey, floeAad);
+    this.floeIv =
+        FloeIv.generateRandom(parameterSpec.getFloeRandom(), parameterSpec.getFloeIvLength());
+    this.header = buildHeader();
+  }
+
+  private byte[] buildHeader() {
+    byte[] parametersEncoded = parameterSpec.paramEncode();
+    byte[] floeIvBytes = floeIv.getBytes();
+    byte[] headerTag =
+        floeKdf.hkdfExpand(floeKey, floeIv, floeAad, FloePurpose.HEADER_TAG, headerTagLength);
+
+    ByteBuffer result =
+        ByteBuffer.allocate(parametersEncoded.length + floeIvBytes.length + headerTag.length);
+    result.put(parametersEncoded);
+    result.put(floeIvBytes);
+    result.put(headerTag);
+    if (result.hasRemaining()) {
+      throw new IllegalArgumentException("Header is too long");
+    }
+    return result.array();
+  }
+
+  @Override
+  public byte[] getHeader() {
+    return header;
+  }
+}

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

@@ -0,0 +1,21 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+class FloeIv {
+  private final byte[] bytes;
+
+  FloeIv(byte[] bytes) {
+    this.bytes = bytes;
+  }
+
+  static FloeIv generateRandom(FloeRandom floeRandom, FloeIvLength floeIvLength) {
+    return new FloeIv(floeRandom.ofLength(floeIvLength.getLength()));
+  }
+
+  byte[] getBytes() {
+    return bytes;
+  }
+
+  int lengthInBytes() {
+    return bytes.length;
+  }
+}

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

@@ -0,0 +1,13 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+public class FloeIvLength {
+  private final int length;
+
+  public FloeIvLength(int length) {
+    this.length = length;
+  }
+
+  public int getLength() {
+    return length;
+  }
+}

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

@@ -0,0 +1,47 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.nio.ByteBuffer;
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.util.Arrays;
+import javax.crypto.Mac;
+
+class FloeKdf {
+  private final FloeParameterSpec parameterSpec;
+
+  FloeKdf(FloeParameterSpec parameterSpec) {
+    this.parameterSpec = parameterSpec;
+  }
+
+  byte[] hkdfExpand(
+      FloeKey floeKey, FloeIv floeIv, FloeAad floeAad, FloePurpose purpose, int length) {
+    byte[] encodedParams = parameterSpec.paramEncode();
+    ByteBuffer info =
+        ByteBuffer.allocate(
+            encodedParams.length
+                + floeIv.getBytes().length
+                + purpose.getBytes().length
+                + floeAad.getBytes().length);
+    info.put(encodedParams);
+    info.put(floeIv.getBytes());
+    info.put(purpose.getBytes());
+    info.put(floeAad.getBytes());
+    return jceHkdfExpand(parameterSpec.getHash(), floeKey, info.array(), length);
+  }
+
+  private byte[] jceHkdfExpand(Hash hash, FloeKey prk, byte[] info, int len) {
+    try {
+      Mac mac = Mac.getInstance(hash.getJceName());
+      mac.init(prk.getKey());
+      mac.update(info);
+      mac.update((byte) 1);
+      byte[] bytes = mac.doFinal();
+      if (bytes.length != len) {
+        return Arrays.copyOf(bytes, len);
+      }
+      return bytes;
+    } catch (NoSuchAlgorithmException | InvalidKeyException e) {
+      throw new RuntimeException(e);
+    }
+  }
+}

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

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

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

@@ -0,0 +1,51 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.nio.ByteBuffer;
+import java.nio.ByteOrder;
+
+public class FloeParameterSpec {
+  private final Aead aead;
+  private final Hash hash;
+  private final int encryptedSegmentLength;
+  private final FloeIvLength floeIvLength;
+  private final FloeRandom floeRandom;
+
+  public FloeParameterSpec(Aead aead, Hash hash, int encryptedSegmentLength, int floeIvLength) {
+    this(
+        aead, hash, encryptedSegmentLength, new FloeIvLength(floeIvLength), new SecureFloeRandom());
+  }
+
+  FloeParameterSpec(
+      Aead aead,
+      Hash hash,
+      int encryptedSegmentLength,
+      FloeIvLength floeIvLength,
+      FloeRandom floeRandom) {
+    this.aead = aead;
+    this.hash = hash;
+    this.encryptedSegmentLength = encryptedSegmentLength;
+    this.floeIvLength = floeIvLength;
+    this.floeRandom = floeRandom;
+  }
+
+  byte[] paramEncode() {
+    ByteBuffer result = ByteBuffer.allocate(10).order(ByteOrder.BIG_ENDIAN);
+    result.put(aead.getId());
+    result.put(hash.getId());
+    result.putInt(encryptedSegmentLength);
+    result.putInt(floeIvLength.getLength());
+    return result.array();
+  }
+
+  public Hash getHash() {
+    return hash;
+  }
+
+  public FloeIvLength getFloeIvLength() {
+    return floeIvLength;
+  }
+
+  FloeRandom getFloeRandom() {
+    return floeRandom;
+  }
+}

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

@@ -0,0 +1,17 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.nio.charset.StandardCharsets;
+
+public enum FloePurpose {
+  HEADER_TAG("HEADER_TAG:".getBytes(StandardCharsets.UTF_8));
+
+  private final byte[] bytes;
+
+  FloePurpose(byte[] bytes) {
+    this.bytes = bytes;
+  }
+
+  public byte[] getBytes() {
+    return bytes;
+  }
+}

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

@@ -0,0 +1,5 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+interface FloeRandom {
+  byte[] ofLength(int length);
+}

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

@@ -0,0 +1,21 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+public enum Hash {
+  SHA384((byte) 0, "HmacSHA384");
+
+  private byte id;
+  private final String jceName;
+
+  Hash(byte id, String jceName) {
+    this.id = id;
+    this.jceName = jceName;
+  }
+
+  byte getId() {
+    return id;
+  }
+
+  public String getJceName() {
+    return jceName;
+  }
+}

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

@@ -0,0 +1,15 @@
+package net.snowflake.client.jdbc.cloud.storage.floe;
+
+import java.security.SecureRandom;
+
+class SecureFloeRandom implements FloeRandom {
+  private static final ThreadLocal<SecureRandom> random =
+      ThreadLocal.withInitial(SecureRandom::new);
+
+  @Override
+  public byte[] ofLength(int length) {
+    byte[] bytes = new byte[length];
+    random.get().nextBytes(bytes);
+    return bytes;
+  }
+}