Skip to content

Commit 2634bdd

Browse files
committed
SNOW-1858529 Implement header handling in FLOE
1 parent 871df20 commit 2634bdd

22 files changed

+493
-3
lines changed

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

+1-2
Original file line numberDiff line numberDiff line change
@@ -41,8 +41,7 @@ public class EncryptionProvider {
4141
private static final String FILE_CIPHER = "AES/CBC/PKCS5Padding";
4242
private static final String KEY_CIPHER = "AES/ECB/PKCS5Padding";
4343
private static final int BUFFER_SIZE = 2 * 1024 * 1024; // 2 MB
44-
private static ThreadLocal<SecureRandom> secRnd =
45-
new ThreadLocal<>().withInitial(SecureRandom::new);
44+
private static ThreadLocal<SecureRandom> secRnd = ThreadLocal.withInitial(SecureRandom::new);
4645

4746
/**
4847
* Decrypt a InputStream

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

+1-1
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ class GcmEncryptionProvider {
3737
private static final String KEY_CIPHER = "AES/GCM/NoPadding";
3838
private static final int BUFFER_SIZE = 8 * 1024 * 1024; // 2 MB
3939
private static final ThreadLocal<SecureRandom> random =
40-
new ThreadLocal<>().withInitial(SecureRandom::new);
40+
ThreadLocal.withInitial(SecureRandom::new);
4141
private static final Base64.Decoder base64Decoder = Base64.getDecoder();
4242

4343
static InputStream encrypt(
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
public enum Aead {
4+
AES_GCM_128((byte) 0),
5+
AES_GCM_256((byte) 1);
6+
7+
private byte id;
8+
9+
Aead(byte id) {
10+
this.id = id;
11+
}
12+
13+
byte getId() {
14+
return id;
15+
}
16+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import javax.crypto.SecretKey;
4+
5+
public class Floe {
6+
private final FloeParameterSpec parameterSpec;
7+
8+
private Floe(FloeParameterSpec parameterSpec) {
9+
this.parameterSpec = parameterSpec;
10+
}
11+
12+
public static Floe getInstance(FloeParameterSpec parameterSpec) {
13+
return new Floe(parameterSpec);
14+
}
15+
16+
public FloeEncryptor createEncryptor(SecretKey key, byte[] aad) {
17+
return new FloeEncryptorImpl(parameterSpec, new FloeKey(key), new FloeAad(aad));
18+
}
19+
20+
public FloeDecryptor createDecryptor(SecretKey key, byte[] aad, byte[] floeHeader) {
21+
return new FloeDecryptorImpl(parameterSpec, new FloeKey(key), new FloeAad(aad), floeHeader);
22+
}
23+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import java.util.Optional;
4+
5+
class FloeAad {
6+
private final byte[] aad;
7+
8+
FloeAad(byte[] aad) {
9+
this.aad = Optional.ofNullable(aad).orElse(new byte[0]);
10+
}
11+
12+
byte[] getBytes() {
13+
return aad;
14+
}
15+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,18 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
abstract class FloeBase {
4+
protected static final int headerTagLength = 32;
5+
6+
protected final FloeParameterSpec parameterSpec;
7+
protected final FloeKey floeKey;
8+
protected final FloeAad floeAad;
9+
10+
protected final FloeKdf floeKdf;
11+
12+
FloeBase(FloeParameterSpec parameterSpec, FloeKey floeKey, FloeAad floeAad) {
13+
this.parameterSpec = parameterSpec;
14+
this.floeKey = floeKey;
15+
this.floeAad = floeAad;
16+
this.floeKdf = new FloeKdf(parameterSpec);
17+
}
18+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
public interface FloeDecryptor {}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import java.nio.ByteBuffer;
4+
import java.util.Arrays;
5+
6+
public class FloeDecryptorImpl extends FloeBase implements FloeDecryptor {
7+
FloeDecryptorImpl(
8+
FloeParameterSpec parameterSpec, FloeKey floeKey, FloeAad floeAad, byte[] floeHeaderAsBytes) {
9+
super(parameterSpec, floeKey, floeAad);
10+
validate(floeHeaderAsBytes);
11+
}
12+
13+
public void validate(byte[] floeHeaderAsBytes) {
14+
byte[] encodedParams = parameterSpec.paramEncode();
15+
if (floeHeaderAsBytes.length
16+
!= encodedParams.length + parameterSpec.getFloeIvLength().getLength() + headerTagLength) {
17+
throw new IllegalArgumentException("invalid header length");
18+
}
19+
ByteBuffer floeHeader = ByteBuffer.wrap(floeHeaderAsBytes);
20+
21+
byte[] encodedParamsFromHeader = new byte[10];
22+
floeHeader.get(encodedParamsFromHeader, 0, encodedParamsFromHeader.length);
23+
if (!Arrays.equals(encodedParams, encodedParamsFromHeader)) {
24+
throw new IllegalArgumentException("invalid parameters header");
25+
}
26+
27+
byte[] floeIvBytes = new byte[parameterSpec.getFloeIvLength().getLength()];
28+
floeHeader.get(floeIvBytes, 0, floeIvBytes.length);
29+
FloeIv floeIv = new FloeIv(floeIvBytes);
30+
31+
byte[] headerTagFromHeader = new byte[headerTagLength];
32+
floeHeader.get(headerTagFromHeader, 0, headerTagFromHeader.length);
33+
34+
byte[] headerTag =
35+
floeKdf.hkdfExpand(floeKey, floeIv, floeAad, FloePurpose.HEADER_TAG, headerTagLength);
36+
if (!Arrays.equals(headerTag, headerTagFromHeader)) {
37+
throw new IllegalArgumentException("invalid header tag");
38+
}
39+
}
40+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
public interface FloeEncryptor {
4+
byte[] getHeader();
5+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import java.nio.ByteBuffer;
4+
5+
class FloeEncryptorImpl extends FloeBase implements FloeEncryptor {
6+
private final FloeIv floeIv;
7+
8+
private final byte[] header;
9+
10+
FloeEncryptorImpl(FloeParameterSpec parameterSpec, FloeKey floeKey, FloeAad floeAad) {
11+
super(parameterSpec, floeKey, floeAad);
12+
this.floeIv =
13+
FloeIv.generateRandom(parameterSpec.getFloeRandom(), parameterSpec.getFloeIvLength());
14+
this.header = buildHeader();
15+
}
16+
17+
private byte[] buildHeader() {
18+
byte[] parametersEncoded = parameterSpec.paramEncode();
19+
byte[] floeIvBytes = floeIv.getBytes();
20+
byte[] headerTag =
21+
floeKdf.hkdfExpand(floeKey, floeIv, floeAad, FloePurpose.HEADER_TAG, headerTagLength);
22+
23+
ByteBuffer result =
24+
ByteBuffer.allocate(parametersEncoded.length + floeIvBytes.length + headerTag.length);
25+
result.put(parametersEncoded);
26+
result.put(floeIvBytes);
27+
result.put(headerTag);
28+
if (result.hasRemaining()) {
29+
throw new IllegalArgumentException("Header is too long");
30+
}
31+
return result.array();
32+
}
33+
34+
@Override
35+
public byte[] getHeader() {
36+
return header;
37+
}
38+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
class FloeIv {
4+
private final byte[] bytes;
5+
6+
FloeIv(byte[] bytes) {
7+
this.bytes = bytes;
8+
}
9+
10+
static FloeIv generateRandom(FloeRandom floeRandom, FloeIvLength floeIvLength) {
11+
return new FloeIv(floeRandom.ofLength(floeIvLength.getLength()));
12+
}
13+
14+
byte[] getBytes() {
15+
return bytes;
16+
}
17+
18+
int lengthInBytes() {
19+
return bytes.length;
20+
}
21+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,13 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
public class FloeIvLength {
4+
private final int length;
5+
6+
public FloeIvLength(int length) {
7+
this.length = length;
8+
}
9+
10+
public int getLength() {
11+
return length;
12+
}
13+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import java.nio.ByteBuffer;
4+
import java.security.InvalidKeyException;
5+
import java.security.NoSuchAlgorithmException;
6+
import java.util.Arrays;
7+
import javax.crypto.Mac;
8+
9+
class FloeKdf {
10+
private final FloeParameterSpec parameterSpec;
11+
12+
FloeKdf(FloeParameterSpec parameterSpec) {
13+
this.parameterSpec = parameterSpec;
14+
}
15+
16+
byte[] hkdfExpand(
17+
FloeKey floeKey, FloeIv floeIv, FloeAad floeAad, FloePurpose purpose, int length) {
18+
byte[] encodedParams = parameterSpec.paramEncode();
19+
ByteBuffer info =
20+
ByteBuffer.allocate(
21+
encodedParams.length
22+
+ floeIv.getBytes().length
23+
+ purpose.getBytes().length
24+
+ floeAad.getBytes().length);
25+
info.put(encodedParams);
26+
info.put(floeIv.getBytes());
27+
info.put(purpose.getBytes());
28+
info.put(floeAad.getBytes());
29+
return jceHkdfExpand(parameterSpec.getHash(), floeKey, info.array(), length);
30+
}
31+
32+
private byte[] jceHkdfExpand(Hash hash, FloeKey prk, byte[] info, int len) {
33+
try {
34+
Mac mac = Mac.getInstance(hash.getJceName());
35+
mac.init(prk.getKey());
36+
mac.update(info);
37+
mac.update((byte) 1);
38+
byte[] bytes = mac.doFinal();
39+
if (bytes.length != len) {
40+
return Arrays.copyOf(bytes, len);
41+
}
42+
return bytes;
43+
} catch (NoSuchAlgorithmException | InvalidKeyException e) {
44+
throw new RuntimeException(e);
45+
}
46+
}
47+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import javax.crypto.SecretKey;
4+
5+
class FloeKey {
6+
private final SecretKey key;
7+
8+
FloeKey(SecretKey key) {
9+
this.key = key;
10+
}
11+
12+
SecretKey getKey() {
13+
return key;
14+
}
15+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import java.nio.ByteBuffer;
4+
import java.nio.ByteOrder;
5+
6+
public class FloeParameterSpec {
7+
private final Aead aead;
8+
private final Hash hash;
9+
private final int encryptedSegmentLength;
10+
private final FloeIvLength floeIvLength;
11+
private final FloeRandom floeRandom;
12+
13+
public FloeParameterSpec(Aead aead, Hash hash, int encryptedSegmentLength, int floeIvLength) {
14+
this(
15+
aead, hash, encryptedSegmentLength, new FloeIvLength(floeIvLength), new SecureFloeRandom());
16+
}
17+
18+
FloeParameterSpec(
19+
Aead aead,
20+
Hash hash,
21+
int encryptedSegmentLength,
22+
FloeIvLength floeIvLength,
23+
FloeRandom floeRandom) {
24+
this.aead = aead;
25+
this.hash = hash;
26+
this.encryptedSegmentLength = encryptedSegmentLength;
27+
this.floeIvLength = floeIvLength;
28+
this.floeRandom = floeRandom;
29+
}
30+
31+
byte[] paramEncode() {
32+
ByteBuffer result = ByteBuffer.allocate(10).order(ByteOrder.BIG_ENDIAN);
33+
result.put(aead.getId());
34+
result.put(hash.getId());
35+
result.putInt(encryptedSegmentLength);
36+
result.putInt(floeIvLength.getLength());
37+
return result.array();
38+
}
39+
40+
public Hash getHash() {
41+
return hash;
42+
}
43+
44+
public FloeIvLength getFloeIvLength() {
45+
return floeIvLength;
46+
}
47+
48+
FloeRandom getFloeRandom() {
49+
return floeRandom;
50+
}
51+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import java.nio.charset.StandardCharsets;
4+
5+
public enum FloePurpose {
6+
HEADER_TAG("HEADER_TAG:".getBytes(StandardCharsets.UTF_8));
7+
8+
private final byte[] bytes;
9+
10+
FloePurpose(byte[] bytes) {
11+
this.bytes = bytes;
12+
}
13+
14+
public byte[] getBytes() {
15+
return bytes;
16+
}
17+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,5 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
interface FloeRandom {
4+
byte[] ofLength(int length);
5+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,21 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
public enum Hash {
4+
SHA384((byte) 0, "HmacSHA384");
5+
6+
private byte id;
7+
private final String jceName;
8+
9+
Hash(byte id, String jceName) {
10+
this.id = id;
11+
this.jceName = jceName;
12+
}
13+
14+
byte getId() {
15+
return id;
16+
}
17+
18+
public String getJceName() {
19+
return jceName;
20+
}
21+
}
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,15 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import java.security.SecureRandom;
4+
5+
class SecureFloeRandom implements FloeRandom {
6+
private static final ThreadLocal<SecureRandom> random =
7+
ThreadLocal.withInitial(SecureRandom::new);
8+
9+
@Override
10+
public byte[] ofLength(int length) {
11+
byte[] bytes = new byte[length];
12+
random.get().nextBytes(bytes);
13+
return bytes;
14+
}
15+
}

0 commit comments

Comments
 (0)