feat: Add decrypt verification step after encryption

src/main/java/com/amazonaws/encryptionsdk/AwsCrypto.java

@@ -12,6 +12,7 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.nio.charset.StandardCharsets;
+import java.util.Arrays;
 import java.util.Collections;
 import java.util.Map;
 
@@ -78,6 +79,7 @@ public class AwsCrypto {
   private static final CommitmentPolicy DEFAULT_COMMITMENT_POLICY =
       CommitmentPolicy.RequireEncryptRequireDecrypt;
   private final CommitmentPolicy commitmentPolicy_;
+  private boolean isDecryptionCheckOnEncrypt_;
 
   /**
    * The maximum number of encrypted data keys to unwrap (resp. wrap) on decrypt (resp. encrypt), if
@@ -109,13 +111,15 @@ private AwsCrypto(Builder builder) {
     encryptionAlgorithm_ = builder.encryptionAlgorithm_;
     encryptionFrameSize_ = builder.encryptionFrameSize_;
     maxEncryptedDataKeys_ = builder.maxEncryptedDataKeys_;
+    isDecryptionCheckOnEncrypt_ = builder.isDecryptionCheckOnEncrypt_;
   }
 
   public static class Builder {
     private CryptoAlgorithm encryptionAlgorithm_;
     private int encryptionFrameSize_ = getDefaultFrameSize();
     private CommitmentPolicy commitmentPolicy_;
     private int maxEncryptedDataKeys_ = CiphertextHeaders.NO_MAX_ENCRYPTED_DATA_KEYS;
+    private boolean isDecryptionCheckOnEncrypt_ = false;
 
     private Builder() {}
 
@@ -124,6 +128,18 @@ private Builder(final AwsCrypto client) {
       encryptionFrameSize_ = client.encryptionFrameSize_;
       commitmentPolicy_ = client.commitmentPolicy_;
       maxEncryptedDataKeys_ = client.maxEncryptedDataKeys_;
+      isDecryptionCheckOnEncrypt_ = client.isDecryptionCheckOnEncrypt_;
+    }
+
+    /**
+     * Sets the {@boolean check} to perform a decryption check or not right after encryption
+     * to make sure no corruption occurred.
+     * @param check
+     * @return The Builder, for method chaining
+     */
+    public Builder withDecryptionCheckOnEncryption(boolean check) {
+      this.isDecryptionCheckOnEncrypt_ = check;
+      return this;
     }
 
     /**
@@ -234,6 +250,14 @@ public CryptoAlgorithm getEncryptionAlgorithm() {
     return encryptionAlgorithm_;
   }
 
+  public void setDecryptionCheckOnEncrypt(final boolean check) {
+    isDecryptionCheckOnEncrypt_ = check;
+  }
+
+  public boolean getDecryptionCheckOnEncrypt() {
+    return isDecryptionCheckOnEncrypt_;
+  }
+
   /**
    * Sets the framing size to use when <em>encrypting</em> data. This has no impact on decryption.
    * If {@code frameSize} is 0, then framing is disabled and the entire plaintext will be encrypted
@@ -365,7 +389,13 @@ public <K extends MasterKey<K>> CryptoResult<byte[], K> encryptData(
     final byte[] outBytes = Utils.truncate(out, outLen);
 
     //noinspection unchecked
-    return new CryptoResult(outBytes, cryptoHandler.getMasterKeys(), cryptoHandler.getHeaders());
+    CryptoResult cryptoResult = new CryptoResult(outBytes, cryptoHandler.getMasterKeys(), cryptoHandler.getHeaders());
+    if (isDecryptionCheckOnEncrypt_) {
+      CryptoResult<byte[], ?> decryptData = decryptData(materialsManager, outBytes);
+      if (!Arrays.equals(decryptData.getResult(), plaintext)) {
+        throw new AwsCryptoException("Decrypt verification step failed, encrypted plaintext couldn't successfully be decrypted.");      }
+    }
+    return cryptoResult;
   }
 
   /**

src/test/java/com/amazonaws/encryptionsdk/AwsCryptoTest.java

@@ -1341,4 +1341,11 @@ public void encryptDecryptStreamWithNoMaxEdks() throws IOException {
     IOUtils.copy(new ByteArrayInputStream(ciphertext), decryptStream);
     decryptStream.close();
   }
+
+  @Test
+  public void setDecryptionCheck() {
+    assertFalse(encryptionClient_.getDecryptionCheckOnEncrypt());
+    encryptionClient_.setDecryptionCheckOnEncrypt(true);
+    assertTrue(encryptionClient_.getDecryptionCheckOnEncrypt());
+  }
 }