fix: Allow reading old plaintext records after adding new encrypted field to schema (#152)

* fix : NPE if reading record without a signature fields

* Fix code formatting

* Add integration test for missing signature fields in document.

* Updated decrypt logic to be less nested.
Improve unit tests

* Add IT test for missing "MaterialDesc" fields.

* Fix code format.

* Remove unwanted if condition and add better unit tests.

* Rename Test class to better represent the intent.

* Fix Unit test

* Remove `if` condition which throw a new error.

* Rename IT class to be better

* Rename unit test method to be better

* Correct formatting

Co-authored-by: Bhuvan Gupta <[email protected]>

Author: bhuvangu

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

@@ -234,7 +234,7 @@ public Map<String, AttributeValue> decryptRecord(
       Map<String, Set<EncryptionFlags>> attributeFlags,
       EncryptionContext context)
       throws GeneralSecurityException {
-    if (attributeFlags.isEmpty()) {
+    if (!itemContainsFieldsToDecryptOrSign(itemAttributes.keySet(), attributeFlags)) {
       return itemAttributes;
     }
     // Copy to avoid changing anyone elses objects
@@ -291,6 +291,13 @@ public Map<String, AttributeValue> decryptRecord(
     return itemAttributes;
   }
 
+  private boolean itemContainsFieldsToDecryptOrSign(
+      Set<String> attributeNamesToCheck, Map<String, Set<EncryptionFlags>> attributeFlags) {
+    return attributeNamesToCheck.stream()
+        .filter(attributeFlags::containsKey)
+        .anyMatch(attributeName -> !attributeFlags.get(attributeName).isEmpty());
+  }
+
   /**
    * Returns the encrypted (and signed) record, which is a map of item attributes. There is no side
    * effect on the input parameters upon calling this method.

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

@@ -15,6 +15,7 @@
 package com.amazonaws.services.dynamodbv2.datamodeling.encryption;
 
 import static com.amazonaws.services.dynamodbv2.datamodeling.encryption.utils.EncryptionContextOperators.overrideEncryptionContextTableName;
+import static java.util.stream.Collectors.toMap;
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.not;
@@ -23,6 +24,7 @@
 import static org.testng.AssertJUnit.assertNotNull;
 import static org.testng.AssertJUnit.assertNull;
 import static org.testng.AssertJUnit.assertTrue;
+import static org.testng.collections.Sets.newHashSet;
 
 import com.amazonaws.services.dynamodbv2.datamodeling.encryption.materials.DecryptionMaterials;
 import com.amazonaws.services.dynamodbv2.datamodeling.encryption.materials.EncryptionMaterials;
@@ -54,6 +56,7 @@
 import org.bouncycastle.jce.ECNamedCurveTable;
 import org.bouncycastle.jce.provider.BouncyCastleProvider;
 import org.bouncycastle.jce.spec.ECParameterSpec;
+import org.mockito.internal.util.collections.Sets;
 import org.testng.Assert;
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.BeforeMethod;
@@ -498,6 +501,32 @@ public void toByteArray() throws ReflectiveOperationException {
     assertToByteArray("Direct", expected, buff);
   }
 
+  @Test
+  public void testDecryptWithPlaintextItem() throws GeneralSecurityException {
+    Map<String, Set<EncryptionFlags>> attributeWithEmptyEncryptionFlags =
+        attribs.keySet().stream().collect(toMap(k -> k, k -> newHashSet()));
+
+    Map<String, AttributeValue> decryptedAttributes =
+        encryptor.decryptRecord(attribs, attributeWithEmptyEncryptionFlags, context);
+    assertThat(decryptedAttributes, AttrMatcher.match(attribs));
+  }
+
+  /*
+  Test decrypt with a map that contains a new key (not included in attribs) with an encryption flag set that contains ENCRYPT and SIGN.
+   */
+  @Test
+  public void testDecryptWithPlainTextItemAndAdditionNewAttributeHavingEncryptionFlag()
+      throws GeneralSecurityException {
+    Map<String, Set<EncryptionFlags>> attributeWithEmptyEncryptionFlags =
+        attribs.keySet().stream().collect(toMap(k -> k, k -> newHashSet()));
+    attributeWithEmptyEncryptionFlags.put(
+        "newAttribute", Sets.newSet(EncryptionFlags.ENCRYPT, EncryptionFlags.SIGN));
+
+    Map<String, AttributeValue> decryptedAttributes =
+        encryptor.decryptRecord(attribs, attributeWithEmptyEncryptionFlags, context);
+    assertThat(decryptedAttributes, AttrMatcher.match(attribs));
+  }
+
   private void assertToByteArray(
       final String msg, final byte[] expected, final ByteBuffer testValue)
       throws ReflectiveOperationException {

sdk1/src/test/java/com/amazonaws/services/dynamodbv2/mapper/integration/PlaintextItemITCase.java

@@ -0,0 +1,115 @@
+/*
+ * Copyright 2015 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except
+ * in compliance with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+package com.amazonaws.services.dynamodbv2.mapper.integration;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNull;
+
+import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBAttribute;
+import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBHashKey;
+import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBMapper;
+import com.amazonaws.services.dynamodbv2.datamodeling.DynamoDBTable;
+import com.amazonaws.services.dynamodbv2.datamodeling.encryption.DoNotTouch;
+import com.amazonaws.services.dynamodbv2.mapper.encryption.TestDynamoDBMapperFactory;
+import com.amazonaws.services.dynamodbv2.model.AttributeValue;
+import com.amazonaws.services.dynamodbv2.model.PutItemRequest;
+import java.util.HashMap;
+import java.util.Map;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+public class PlaintextItemITCase extends DynamoDBMapperCryptoIntegrationTestBase {
+  private static final String STRING_ATTRIBUTE = "stringAttribute";
+  private static Map<String, AttributeValue> plaintextItem = new HashMap<>();
+  // Test data
+  static {
+    plaintextItem.put(KEY_NAME, new AttributeValue().withS("" + startKey++));
+    plaintextItem.put(STRING_ATTRIBUTE, new AttributeValue().withS("" + startKey++));
+  }
+
+  @BeforeClass
+  public static void setUp() throws Exception {
+    DynamoDBMapperCryptoIntegrationTestBase.setUp();
+    // Insert the data
+    dynamo.putItem(new PutItemRequest(TABLE_NAME, plaintextItem));
+  }
+
+  @Test
+  public void testLoadWithPlaintextItem() {
+    DynamoDBMapper util = TestDynamoDBMapperFactory.createDynamoDBMapper(dynamo);
+    UntouchedTable load = util.load(UntouchedTable.class, plaintextItem.get(KEY_NAME).getS());
+
+    assertEquals(load.getKey(), plaintextItem.get(KEY_NAME).getS());
+    assertEquals(load.getStringAttribute(), plaintextItem.get(STRING_ATTRIBUTE).getS());
+  }
+
+  @Test
+  public void testLoadWithPlaintextItemWithModelHavingNewEncryptedAttribute() {
+    DynamoDBMapper util = TestDynamoDBMapperFactory.createDynamoDBMapper(dynamo);
+    UntouchedWithNewEncryptedAttributeTable load =
+        util.load(
+            UntouchedWithNewEncryptedAttributeTable.class, plaintextItem.get(KEY_NAME).getS());
+
+    assertEquals(load.getKey(), plaintextItem.get(KEY_NAME).getS());
+    assertEquals(load.getStringAttribute(), plaintextItem.get(STRING_ATTRIBUTE).getS());
+    assertNull(load.getNewAttribute());
+  }
+
+  @DynamoDBTable(tableName = "aws-java-sdk-util-crypto")
+  public static class UntouchedTable {
+
+    private String key;
+
+    private String stringAttribute;
+
+    @DynamoDBHashKey
+    @DoNotTouch
+    public String getKey() {
+      return key;
+    }
+
+    public void setKey(String key) {
+      this.key = key;
+    }
+
+    @DynamoDBAttribute
+    @DoNotTouch
+    public String getStringAttribute() {
+      return stringAttribute;
+    }
+
+    public void setStringAttribute(String stringAttribute) {
+      this.stringAttribute = stringAttribute;
+    }
+
+    @Override
+    public boolean equals(Object o) {
+      if (this == o) return true;
+      if (o == null || getClass() != o.getClass()) return false;
+      UntouchedTable that = (UntouchedTable) o;
+      return key.equals(that.key) && stringAttribute.equals(that.stringAttribute);
+    }
+  }
+
+  public static final class UntouchedWithNewEncryptedAttributeTable extends UntouchedTable {
+    private String newAttribute;
+
+    public String getNewAttribute() {
+      return newAttribute;
+    }
+
+    public void setNewAttribute(String newAttribute) {
+      this.newAttribute = newAttribute;
+    }
+  }
+}