Skip to content

Commit dd4ff7d

Browse files
sfc-gh-pfussfc-gh-pfus
committed
Implement processing segments
1 parent 2634bdd commit dd4ff7d

21 files changed

+496
-133
lines changed

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

+47-20
Original file line numberDiff line numberDiff line change
@@ -3,38 +3,42 @@
33
*/
44
package net.snowflake.client.jdbc.cloud.storage;
55

6-
import static java.nio.file.StandardOpenOption.CREATE;
7-
import static java.nio.file.StandardOpenOption.READ;
6+
import net.snowflake.client.core.SnowflakeJdbcInternalApi;
7+
import net.snowflake.client.jdbc.MatDesc;
8+
import net.snowflake.client.jdbc.cloud.storage.floe.AeadProvider;
9+
import net.snowflake.common.core.RemoteStoreFileEncryptionMaterial;
810

11+
import javax.crypto.BadPaddingException;
12+
import javax.crypto.Cipher;
13+
import javax.crypto.CipherInputStream;
14+
import javax.crypto.IllegalBlockSizeException;
15+
import javax.crypto.NoSuchPaddingException;
16+
import javax.crypto.SecretKey;
17+
import javax.crypto.spec.GCMParameterSpec;
18+
import javax.crypto.spec.SecretKeySpec;
919
import java.io.File;
1020
import java.io.FileOutputStream;
1121
import java.io.IOException;
1222
import java.io.InputStream;
1323
import java.io.OutputStream;
1424
import java.nio.channels.FileChannel;
1525
import java.nio.file.Files;
26+
import java.security.GeneralSecurityException;
1627
import java.security.InvalidAlgorithmParameterException;
1728
import java.security.InvalidKeyException;
1829
import java.security.NoSuchAlgorithmException;
1930
import java.security.SecureRandom;
2031
import java.util.Base64;
21-
import javax.crypto.BadPaddingException;
22-
import javax.crypto.Cipher;
23-
import javax.crypto.CipherInputStream;
24-
import javax.crypto.IllegalBlockSizeException;
25-
import javax.crypto.NoSuchPaddingException;
26-
import javax.crypto.SecretKey;
27-
import javax.crypto.spec.GCMParameterSpec;
28-
import javax.crypto.spec.SecretKeySpec;
29-
import net.snowflake.client.jdbc.MatDesc;
30-
import net.snowflake.common.core.RemoteStoreFileEncryptionMaterial;
3132

32-
class GcmEncryptionProvider {
33+
import static java.nio.file.StandardOpenOption.CREATE;
34+
import static java.nio.file.StandardOpenOption.READ;
35+
36+
@SnowflakeJdbcInternalApi
37+
public class GcmEncryptionProvider implements AeadProvider {
3338
private static final int TAG_LENGTH_IN_BITS = 128;
3439
private static final int IV_LENGTH_IN_BYTES = 12;
3540
private static final String AES = "AES";
36-
private static final String FILE_CIPHER = "AES/GCM/NoPadding";
37-
private static final String KEY_CIPHER = "AES/GCM/NoPadding";
41+
private static final String JCE_CIPHER_NAME = "AES/GCM/NoPadding";
3842
private static final int BUFFER_SIZE = 8 * 1024 * 1024; // 2 MB
3943
private static final ThreadLocal<SecureRandom> random =
4044
ThreadLocal.withInitial(SecureRandom::new);
@@ -85,7 +89,7 @@ private static byte[] encryptKey(byte[] kekBytes, byte[] keyBytes, byte[] keyIvD
8589
BadPaddingException, NoSuchPaddingException, NoSuchAlgorithmException {
8690
SecretKey kek = new SecretKeySpec(kekBytes, 0, kekBytes.length, AES);
8791
GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_IN_BITS, keyIvData);
88-
Cipher keyCipher = Cipher.getInstance(KEY_CIPHER);
92+
Cipher keyCipher = Cipher.getInstance(JCE_CIPHER_NAME);
8993
keyCipher.init(Cipher.ENCRYPT_MODE, kek, gcmParameterSpec);
9094
if (aad != null) {
9195
keyCipher.updateAAD(aad);
@@ -99,7 +103,7 @@ private static CipherInputStream encryptContent(
99103
NoSuchAlgorithmException {
100104
SecretKey fileKey = new SecretKeySpec(keyBytes, 0, keyBytes.length, AES);
101105
GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_IN_BITS, dataIvBytes);
102-
Cipher fileCipher = Cipher.getInstance(FILE_CIPHER);
106+
Cipher fileCipher = Cipher.getInstance(JCE_CIPHER_NAME);
103107
fileCipher.init(Cipher.ENCRYPT_MODE, fileKey, gcmParameterSpec);
104108
if (aad != null) {
105109
fileCipher.updateAAD(aad);
@@ -172,7 +176,7 @@ private static CipherInputStream decryptContentFromStream(
172176
NoSuchAlgorithmException {
173177
GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_IN_BITS, ivBytes);
174178
SecretKey fileKey = new SecretKeySpec(fileKeyBytes, AES);
175-
Cipher fileCipher = Cipher.getInstance(FILE_CIPHER);
179+
Cipher fileCipher = Cipher.getInstance(JCE_CIPHER_NAME);
176180
fileCipher.init(Cipher.DECRYPT_MODE, fileKey, gcmParameterSpec);
177181
if (aad != null) {
178182
fileCipher.updateAAD(aad);
@@ -187,7 +191,7 @@ private static void decryptContentFromFile(
187191
SecretKey fileKey = new SecretKeySpec(fileKeyBytes, AES);
188192
GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_IN_BITS, cekIvBytes);
189193
byte[] buffer = new byte[BUFFER_SIZE];
190-
Cipher fileCipher = Cipher.getInstance(FILE_CIPHER);
194+
Cipher fileCipher = Cipher.getInstance(JCE_CIPHER_NAME);
191195
fileCipher.init(Cipher.DECRYPT_MODE, fileKey, gcmParameterSpec);
192196
if (aad != null) {
193197
fileCipher.updateAAD(aad);
@@ -215,11 +219,34 @@ private static byte[] decryptKey(byte[] kekBytes, byte[] ivBytes, byte[] keyByte
215219
BadPaddingException, NoSuchPaddingException, NoSuchAlgorithmException {
216220
SecretKey kek = new SecretKeySpec(kekBytes, 0, kekBytes.length, AES);
217221
GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_IN_BITS, ivBytes);
218-
Cipher keyCipher = Cipher.getInstance(KEY_CIPHER);
222+
Cipher keyCipher = Cipher.getInstance(JCE_CIPHER_NAME);
219223
keyCipher.init(Cipher.DECRYPT_MODE, kek, gcmParameterSpec);
220224
if (aad != null) {
221225
keyCipher.updateAAD(aad);
222226
}
223227
return keyCipher.doFinal(keyBytes);
224228
}
229+
230+
// TODO refactor to reuse cipher (consider thread safety vs performance)
231+
@Override
232+
public byte[] encrypt(SecretKey key, byte[] iv, byte[] aad, byte[] plaintext) throws GeneralSecurityException {
233+
GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_IN_BITS, iv);
234+
Cipher keyCipher = Cipher.getInstance(JCE_CIPHER_NAME);
235+
keyCipher.init(Cipher.ENCRYPT_MODE, key, gcmParameterSpec);
236+
if (aad != null) {
237+
keyCipher.updateAAD(aad);
238+
}
239+
return keyCipher.doFinal(plaintext);
240+
}
241+
242+
@Override
243+
public byte[] decrypt(SecretKey key, byte[] iv, byte[] aad, byte[] ciphertext) throws GeneralSecurityException {
244+
GCMParameterSpec gcmParameterSpec = new GCMParameterSpec(TAG_LENGTH_IN_BITS, iv);
245+
Cipher keyCipher = Cipher.getInstance(JCE_CIPHER_NAME);
246+
keyCipher.init(Cipher.DECRYPT_MODE, key, gcmParameterSpec);
247+
if (aad != null) {
248+
keyCipher.updateAAD(aad);
249+
}
250+
return keyCipher.doFinal(ciphertext);
251+
}
225252
}

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

+36-3
Original file line numberDiff line numberDiff line change
@@ -1,16 +1,49 @@
11
package net.snowflake.client.jdbc.cloud.storage.floe;
22

3+
import net.snowflake.client.jdbc.cloud.storage.GcmEncryptionProvider;
4+
35
public enum Aead {
4-
AES_GCM_128((byte) 0),
5-
AES_GCM_256((byte) 1);
6+
// TODO confirm id
7+
AES_GCM_256((byte) 0, "AES/GCM/NoPadding", 32, 12, 16, new GcmEncryptionProvider()),
8+
AES_GCM_128((byte) 1, "AES/GCM/NoPadding", 16, 12, 16, new GcmEncryptionProvider());
69

710
private byte id;
11+
private String jceName;
12+
private int keyLength;
13+
private int ivLength;
14+
private int authTagLength;
15+
private AeadProvider aeadProvider;
816

9-
Aead(byte id) {
17+
Aead(byte id, String jceName, int keyLength, int ivLength, int authTagLength, AeadProvider aeadProvider) {
18+
this.jceName = jceName;
19+
this.keyLength = keyLength;
1020
this.id = id;
21+
this.ivLength = ivLength;
22+
this.authTagLength = authTagLength;
23+
this.aeadProvider = aeadProvider;
1124
}
1225

1326
byte getId() {
1427
return id;
1528
}
29+
30+
String getJceName() {
31+
return jceName;
32+
}
33+
34+
int getKeyLength() {
35+
return keyLength;
36+
}
37+
38+
int getIvLength() {
39+
return ivLength;
40+
}
41+
42+
int getAuthTagLength() {
43+
return authTagLength;
44+
}
45+
46+
AeadProvider getAeadProvider() {
47+
return aeadProvider;
48+
}
1649
}

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

+22
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import java.nio.ByteBuffer;
4+
5+
class AeadAad {
6+
private final byte[] bytes;
7+
8+
private AeadAad(long segmentCounter, byte terminalityByte) {
9+
ByteBuffer buf = ByteBuffer.allocate(9);
10+
buf.putLong(segmentCounter);
11+
buf.put(terminalityByte);
12+
this.bytes = buf.array();
13+
}
14+
15+
static AeadAad nonTerminal(long segmentCounter) {
16+
return new AeadAad(segmentCounter, (byte) 0);
17+
}
18+
19+
byte[] getBytes() {
20+
return bytes;
21+
}
22+
}

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

+25
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import java.nio.ByteBuffer;
4+
5+
class AeadIv {
6+
private final byte[] bytes;
7+
8+
AeadIv(byte[] bytes) {
9+
this.bytes = bytes;
10+
}
11+
12+
public static AeadIv generateRandom(FloeRandom floeRandom, int ivLength) {
13+
return new AeadIv(floeRandom.ofLength(ivLength));
14+
}
15+
16+
public static AeadIv from(ByteBuffer buffer, int ivLength) {
17+
byte[] bytes = new byte[ivLength];
18+
buffer.get(bytes);
19+
return new AeadIv(bytes);
20+
}
21+
22+
byte[] getBytes() {
23+
return bytes;
24+
}
25+
}

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

+15
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 AeadKey {
6+
private final SecretKey key;
7+
8+
AeadKey(SecretKey key) {
9+
this.key = key;
10+
}
11+
12+
SecretKey getKey() {
13+
return key;
14+
}
15+
}

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

+9
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,9 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import javax.crypto.SecretKey;
4+
import java.security.GeneralSecurityException;
5+
6+
public interface AeadProvider {
7+
byte[] encrypt(SecretKey key, byte[] iv, byte[] aad, byte[] plaintext) throws GeneralSecurityException;
8+
byte[] decrypt(SecretKey key, byte[] iv, byte[] aad, byte[] ciphertext) throws GeneralSecurityException;
9+
}

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

+37
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
package net.snowflake.client.jdbc.cloud.storage.floe;
2+
3+
import javax.crypto.SecretKey;
4+
import javax.crypto.spec.SecretKeySpec;
5+
6+
abstract class BaseSegmentProcessor {
7+
protected static final int NON_TERMINAL_SEGMENT_SIZE_MARKER = -1;
8+
protected static final int headerTagLength = 32;
9+
10+
protected final FloeParameterSpec parameterSpec;
11+
protected final FloeKey floeKey;
12+
protected final FloeAad floeAad;
13+
14+
protected final KeyDerivator floeKdf;
15+
16+
private AeadKey currentAeadKey;
17+
18+
BaseSegmentProcessor(FloeParameterSpec parameterSpec, FloeKey floeKey, FloeAad floeAad) {
19+
this.parameterSpec = parameterSpec;
20+
this.floeKey = floeKey;
21+
this.floeAad = floeAad;
22+
this.floeKdf = new KeyDerivator(parameterSpec);
23+
}
24+
25+
protected AeadKey getKey(FloeKey floeKey, FloeIv floeIv, FloeAad floeAad, long segmentCounter) {
26+
if (currentAeadKey == null || segmentCounter % parameterSpec.getKeyRotationModulo() == 0) {
27+
currentAeadKey = deriveKey(floeKey, floeIv, floeAad, segmentCounter);
28+
}
29+
return currentAeadKey;
30+
}
31+
32+
private AeadKey deriveKey(FloeKey floeKey, FloeIv floeIv, FloeAad floeAad, long segmentCounter) {
33+
byte[] keyBytes = floeKdf.hkdfExpand(floeKey, floeIv, floeAad, new DekTagFloePurpose(segmentCounter), parameterSpec.getAead().getKeyLength());
34+
SecretKey key = new SecretKeySpec(keyBytes, "AES"); // for now it is safe as we use only AES as AEAD
35+
return new AeadKey(key);
36+
}
37+
}

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

-18
This file was deleted.

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

+3-1
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,5 @@
11
package net.snowflake.client.jdbc.cloud.storage.floe;
22

3-
public interface FloeDecryptor {}
3+
public interface FloeDecryptor extends SegmentProcessor {
4+
5+
}

0 commit comments

Comments
 (0)