sweisdb commented on code in PR #40969:
URL: https://github.com/apache/spark/pull/40969#discussion_r1183127444
##########
sql/catalyst/src/main/java/org/apache/spark/sql/catalyst/expressions/ExpressionImplUtils.java:
##########
@@ -85,85 +117,68 @@ public static byte[] aesDecrypt(byte[] input, byte[] key,
UTF8String mode, UTF8S
return aesInternal(input, key, mode.toString(), padding.toString(),
Cipher.DECRYPT_MODE);
}
+ private static SecretKeySpec getSecretKeySpec(byte[] key) {
+ switch (key.length) {
+ case 16: case 24: case 32:
+ return new SecretKeySpec(key, 0, key.length, "AES");
+ default:
+ throw QueryExecutionErrors.invalidAesKeyLengthError(key.length);
+ }
+ }
+
+ private static byte[] generateIv(CipherMode mode) {
+ byte[] iv = new byte[mode.ivLength];
+ threadLocalSecureRandom.get().nextBytes(iv);
+ return iv;
+ }
+
+ private static AlgorithmParameterSpec getParamSpec(CipherMode mode, byte[]
input, int offset) {
+ switch (mode) {
+ case CBC:
+ return new IvParameterSpec(input, offset, mode.ivLength);
+ case GCM:
+ return new GCMParameterSpec(mode.tagLength, input, offset,
mode.ivLength);
+ default:
+ return null;
+ }
+ }
+
private static byte[] aesInternal(
byte[] input,
byte[] key,
String mode,
String padding,
int opmode) {
- SecretKeySpec secretKey;
-
- switch (key.length) {
- case 16:
- case 24:
- case 32:
- secretKey = new SecretKeySpec(key, 0, key.length, "AES");
- break;
- default:
- throw QueryExecutionErrors.invalidAesKeyLengthError(key.length);
- }
-
try {
- if (mode.equalsIgnoreCase("ECB") &&
- (padding.equalsIgnoreCase("PKCS") ||
padding.equalsIgnoreCase("DEFAULT"))) {
- Cipher cipher = Cipher.getInstance("AES/ECB/PKCS5Padding");
- cipher.init(opmode, secretKey);
- return cipher.doFinal(input, 0, input.length);
- } else if (mode.equalsIgnoreCase("GCM") &&
- (padding.equalsIgnoreCase("NONE") ||
padding.equalsIgnoreCase("DEFAULT"))) {
- Cipher cipher = Cipher.getInstance("AES/GCM/NoPadding");
- if (opmode == Cipher.ENCRYPT_MODE) {
- byte[] iv = new byte[GCM_IV_LEN];
- secureRandom.nextBytes(iv);
- GCMParameterSpec parameterSpec = new GCMParameterSpec(GCM_TAG_LEN,
iv);
- cipher.init(Cipher.ENCRYPT_MODE, secretKey, parameterSpec);
- byte[] encrypted = cipher.doFinal(input, 0, input.length);
+ SecretKeySpec secretKey = getSecretKeySpec(key);
+ CipherMode cipherMode = CipherMode.fromString(mode, padding);
+ Cipher cipher = Cipher.getInstance(cipherMode.transformation);
+ if (opmode == Cipher.ENCRYPT_MODE) {
+ // This IV will be 0-length for ECB
+ byte[] iv = generateIv(cipherMode);
+ if (cipherMode.usesSpec) {
+ AlgorithmParameterSpec algSpec = getParamSpec(cipherMode, iv, 0);
+ cipher.init(opmode, secretKey, algSpec);
+ } else {
+ cipher.init(opmode, secretKey);
+ }
+ byte[] encrypted = cipher.doFinal(input, 0, input.length);
+ if (iv.length > 0) {
ByteBuffer byteBuffer = ByteBuffer.allocate(iv.length +
encrypted.length);
byteBuffer.put(iv);
byteBuffer.put(encrypted);
return byteBuffer.array();
} else {
- assert(opmode == Cipher.DECRYPT_MODE);
- GCMParameterSpec parameterSpec = new GCMParameterSpec(GCM_TAG_LEN,
input, 0, GCM_IV_LEN);
- cipher.init(Cipher.DECRYPT_MODE, secretKey, parameterSpec);
- return cipher.doFinal(input, GCM_IV_LEN, input.length - GCM_IV_LEN);
+ return encrypted;
}
- } else if (mode.equalsIgnoreCase("CBC") &&
- (padding.equalsIgnoreCase("PKCS") ||
padding.equalsIgnoreCase("DEFAULT"))) {
- Cipher cipher = Cipher.getInstance("AES/CBC/PKCS5Padding");
- if (opmode == Cipher.ENCRYPT_MODE) {
- byte[] salt = new byte[CBC_SALT_LEN];
- secureRandom.nextBytes(salt);
- final byte[] keyAndIv = getKeyAndIv(key, salt);
- final byte[] keyValue = Arrays.copyOfRange(keyAndIv, 0, key.length);
- final byte[] iv = Arrays.copyOfRange(keyAndIv, key.length,
key.length + CBC_IV_LEN);
- cipher.init(
- Cipher.ENCRYPT_MODE,
- new SecretKeySpec(keyValue, "AES"),
- new IvParameterSpec(iv));
- byte[] encrypted = cipher.doFinal(input, 0, input.length);
- ByteBuffer byteBuffer = ByteBuffer.allocate(
- SALTED_MAGIC.length + CBC_SALT_LEN + encrypted.length);
- byteBuffer.put(SALTED_MAGIC);
- byteBuffer.put(salt);
- byteBuffer.put(encrypted);
- return byteBuffer.array();
+ } else if (opmode == Cipher.DECRYPT_MODE) {
Review Comment:
Makes sense. Will do.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
