Skip to content

Commit c7d178e

Browse files
SalusaSecondusSalusaSecondus
committed
Reorder calls to Cipher object for better support of non-standard JCA providers
1 parent 056b3de commit c7d178e

File tree

2 files changed

+84
-25
lines changed

2 files changed

+84
-25
lines changed

‎src/main/java/com/amazonaws/services/dynamodbv2/datamodeling/encryption/DynamoDBEncryptor.java

Lines changed: 47 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,8 @@
3131
import java.util.HashMap;
3232
import java.util.Map;
3333
import java.util.Set;
34+
import java.util.concurrent.ConcurrentHashMap;
35+
import java.util.function.Function;
3436

3537
import javax.crypto.Cipher;
3638
import javax.crypto.SecretKey;
@@ -58,7 +60,16 @@ public class DynamoDBEncryptor {
5860
private static final String DEFAULT_DESCRIPTION_BASE = "amzn-ddb-map-"; // Same as the Mapper
5961
private static final Charset UTF8 = Charset.forName("UTF-8");
6062
private static final String SYMMETRIC_ENCRYPTION_MODE = "/CBC/PKCS5Padding";
61-
63+
private static final ConcurrentHashMap<String, Integer> BLOCK_SIZE_CACHE = new ConcurrentHashMap<>();
64+
private static final Function<String, Integer> BLOCK_SIZE_CALCULATOR = (transformation) -> {
65+
try {
66+
final Cipher c = Cipher.getInstance(transformation);
67+
return c.getBlockSize();
68+
} catch (final GeneralSecurityException ex) {
69+
throw new IllegalArgumentException("Algorithm does not exist", ex);
70+
}
71+
};
72+
6273
private static final int CURRENT_VERSION = 0;
6374

6475
private String signatureFieldName = DEFAULT_SIGNATURE_FIELD;
@@ -339,7 +350,7 @@ private void actualDecryption(Map<String, AttributeValue> itemAttributes,
339350
final String encryptionMode = encryptionKey != null ? encryptionKey.getAlgorithm() +
340351
materialDescription.get(symmetricEncryptionModeHeader) : null;
341352
Cipher cipher = null;
342-
int ivSize = -1;
353+
int blockSize = -1;
343354

344355
for (Map.Entry<String, AttributeValue> entry: itemAttributes.entrySet()) {
345356
Set<EncryptionFlags> flags = attributeFlags.get(entry.getKey());
@@ -354,15 +365,13 @@ private void actualDecryption(Map<String, AttributeValue> itemAttributes,
354365
plainText = ByteBuffer.wrap(((DelegatedKey)encryptionKey).decrypt(toByteArray(cipherText), null, encryptionMode));
355366
} else {
356367
if (cipher == null) {
357-
cipher = Cipher.getInstance(
358-
encryptionMode);
359-
ivSize = cipher.getBlockSize();
368+
blockSize = getBlockSize(encryptionMode);
369+
cipher = Cipher.getInstance(encryptionMode);
360370
}
361-
byte[] iv = new byte[ivSize];
371+
byte[] iv = new byte[blockSize];
362372
cipherText.get(iv);
363373
cipher.init(Cipher.DECRYPT_MODE, encryptionKey, new IvParameterSpec(iv), Utils.getRng());
364-
plainText = ByteBuffer.allocate(
365-
cipher.getOutputSize(cipherText.remaining()));
374+
plainText = ByteBuffer.allocate(cipher.getOutputSize(cipherText.remaining()));
366375
cipher.doFinal(cipherText, plainText);
367376
plainText.rewind();
368377
}
@@ -371,6 +380,10 @@ private void actualDecryption(Map<String, AttributeValue> itemAttributes,
371380
}
372381
}
373382

383+
protected int getBlockSize(final String encryptionMode) {
384+
return BLOCK_SIZE_CACHE.computeIfAbsent(encryptionMode, BLOCK_SIZE_CALCULATOR);
385+
}
386+
374387
/**
375388
* This method has the side effect of replacing the plaintext
376389
* attribute-values of "itemAttributes" with ciphertext attribute-values
@@ -388,7 +401,7 @@ private void actualEncryption(Map<String, AttributeValue> itemAttributes,
388401
encryptionMode = encryptionKey.getAlgorithm() + SYMMETRIC_ENCRYPTION_MODE;
389402
}
390403
Cipher cipher = null;
391-
int ivSize = -1;
404+
int blockSize = -1;
392405

393406
for (Map.Entry<String, AttributeValue> entry: itemAttributes.entrySet()) {
394407
Set<EncryptionFlags> flags = attributeFlags.get(entry.getKey());
@@ -405,16 +418,22 @@ private void actualEncryption(Map<String, AttributeValue> itemAttributes,
405418
dk.encrypt(toByteArray(plainText), null, encryptionMode));
406419
} else {
407420
if (cipher == null) {
421+
blockSize = getBlockSize(encryptionMode);
408422
cipher = Cipher.getInstance(encryptionMode);
409-
ivSize = cipher.getBlockSize();
410423
}
411424
// Encryption format: <iv><ciphertext>
412425
// Note a unique iv is generated per attribute
413-
byte[] iv = Utils.getRandom(ivSize);
414-
cipher.init(Cipher.ENCRYPT_MODE, encryptionKey, new IvParameterSpec(iv), Utils.getRng());
415-
cipherText = ByteBuffer.allocate(ivSize + cipher.getOutputSize(plainText.remaining()));
416-
cipherText.put(iv);
426+
cipher.init(Cipher.ENCRYPT_MODE, encryptionKey, Utils.getRng());
427+
cipherText = ByteBuffer.allocate(blockSize + cipher.getOutputSize(plainText.remaining()));
428+
cipherText.position(blockSize);
417429
cipher.doFinal(plainText, cipherText);
430+
cipherText.flip();
431+
final byte[] iv = cipher.getIV();
432+
if (iv.length != blockSize) {
433+
throw new IllegalStateException(String.format("Generated IV length (%d) not equal to block size (%d)",
434+
iv.length, blockSize));
435+
}
436+
cipherText.put(iv);
418437
cipherText.rewind();
419438
}
420439
// Replace the plaintext attribute value with the encrypted content
@@ -539,17 +558,22 @@ protected static Map<String, String> unmarshallDescription(AttributeValue attrib
539558
attributeValue.getB().reset();
540559
}
541560
}
542-
561+
543562
private static byte[] toByteArray(ByteBuffer buffer) {
544-
if (buffer.hasArray()) {
563+
buffer = buffer.duplicate();
564+
// We can only return the array directly if:
565+
// 1. The ByteBuffer exposes an array
566+
// 2. The ByteBuffer starts at the beginning of the array
567+
// 3. The ByteBuffer uses the entire array
568+
if (buffer.hasArray() && buffer.arrayOffset() == 0) {
545569
byte[] result = buffer.array();
546-
buffer.rewind();
547-
return result;
548-
} else {
549-
byte[] result = new byte[buffer.remaining()];
550-
buffer.get(result);
551-
buffer.rewind();
552-
return result;
570+
if (buffer.remaining() == result.length) {
571+
return result;
572+
}
553573
}
574+
575+
byte[] result = new byte[buffer.remaining()];
576+
buffer.get(result);
577+
return result;
554578
}
555579
}

‎src/test/java/com/amazonaws/services/dynamodbv2/datamodeling/encryption/DynamoDBEncryptorTest.java

Lines changed: 37 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,20 +14,21 @@
1414
*/
1515
package com.amazonaws.services.dynamodbv2.datamodeling.encryption;
1616

17+
import static org.junit.Assert.assertArrayEquals;
1718
import static org.junit.Assert.assertEquals;
1819
import static org.junit.Assert.assertNotNull;
1920
import static org.junit.Assert.assertNull;
2021
import static org.junit.Assert.assertThat;
2122
import static org.junit.Assert.assertTrue;
2223

24+
import java.lang.reflect.Method;
2325
import java.nio.ByteBuffer;
2426
import java.security.GeneralSecurityException;
2527
import java.security.InvalidAlgorithmParameterException;
2628
import java.security.KeyPair;
2729
import java.security.KeyPairGenerator;
2830
import java.security.NoSuchAlgorithmException;
2931
import java.security.NoSuchProviderException;
30-
import java.security.SecureRandom;
3132
import java.security.Security;
3233
import java.security.SignatureException;
3334
import java.util.Collection;
@@ -118,7 +119,7 @@ public void testSetMaterialDescriptionFieldName() {
118119

119120
@Test
120121
public void fullEncryption() throws GeneralSecurityException {
121-
Map<String, AttributeValue> encryptedAttributes =
122+
Map<String, AttributeValue> encryptedAttributes =
122123
encryptor.encryptAllFieldsExcept(Collections.unmodifiableMap(attribs), context, "hashKey", "rangeKey", "version");
123124
assertThat(encryptedAttributes, AttrMatcher.invert(attribs));
124125

@@ -298,6 +299,40 @@ public void EcdsaSignedOnlyBadSignature() throws GeneralSecurityException {
298299
encryptor.decryptAllFieldsExcept(encryptedAttributes, context, attribs.keySet().toArray(new String[0]));
299300
}
300301

302+
@Test
303+
public void toByteArray() throws ReflectiveOperationException {
304+
final byte[] expected = new byte[] {0, 1, 2, 3, 4, 5};
305+
assertToByteArray("Wrap", expected, ByteBuffer.wrap(expected));
306+
assertToByteArray("Wrap-RO", expected, ByteBuffer.wrap(expected).asReadOnlyBuffer());
307+
308+
assertToByteArray("Wrap-Truncated-Sliced", expected, ByteBuffer.wrap(new byte[] {0, 1, 2, 3, 4, 5, 6}, 0, 6).slice());
309+
assertToByteArray("Wrap-Offset-Sliced", expected, ByteBuffer.wrap(new byte[] {6, 0, 1, 2, 3, 4, 5, 6}, 1, 6).slice());
310+
assertToByteArray("Wrap-Truncated", expected, ByteBuffer.wrap(new byte[] {0, 1, 2, 3, 4, 5, 6}, 0, 6));
311+
assertToByteArray("Wrap-Offset", expected, ByteBuffer.wrap(new byte[] {6, 0, 1, 2, 3, 4, 5, 6}, 1, 6));
312+
313+
ByteBuffer buff = ByteBuffer.allocate(expected.length + 10);
314+
buff.put(expected);
315+
buff.flip();
316+
assertToByteArray("Normal", expected, buff);
317+
318+
buff = ByteBuffer.allocateDirect(expected.length + 10);
319+
buff.put(expected);
320+
buff.flip();
321+
assertToByteArray("Direct", expected, buff);
322+
}
323+
324+
private void assertToByteArray(final String msg, final byte[] expected, final ByteBuffer testValue) throws ReflectiveOperationException {
325+
Method m = DynamoDBEncryptor.class.getDeclaredMethod("toByteArray", ByteBuffer.class);
326+
m.setAccessible(true);
327+
328+
int oldPosition = testValue.position();
329+
int oldLimit = testValue.limit();
330+
331+
assertArrayEquals(msg + ":Array", expected, (byte[]) m.invoke(null, testValue));
332+
assertEquals(msg + ":Position", oldPosition, testValue.position());
333+
assertEquals(msg + ":Limit", oldLimit, testValue.limit());
334+
}
335+
301336
private void assertAttrEquals(AttributeValue o1, AttributeValue o2) {
302337
Assert.assertEquals(o1.getB(), o2.getB());
303338
assertSetsEqual(o1.getBS(), o2.getBS());

0 commit comments

Comments
 (0)