fix: return plaintext items in UnprocessedItems in BatchWriteIttem

Author: ajewellamz

DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/src/BatchWriteItemTransform.dfy

@@ -13,6 +13,7 @@ module BatchWriteItemTransform {
   import Seq
   import SortedSets
   import Util = DynamoDbEncryptionUtil
+  import Types = AwsCryptographyDbEncryptionSdkDynamoDbTypes
 
   method {:vcs_split_on_every_assert} Input(config: Config, input: BatchWriteItemInputTransformInput)
     returns (output: Result<BatchWriteItemInputTransformOutput, Error>)
@@ -80,11 +81,86 @@ module BatchWriteItemTransform {
     return Success(BatchWriteItemInputTransformOutput(transformedInput := input.sdkInput.(RequestItems := result)));
   }
 
+  method GetOrigItem(
+    tableConfig : ValidTableConfig,
+    srcRequests : DDB.WriteRequests,
+    itemReq : DDB.WriteRequest
+    ) returns (ret : Result<DDB.WriteRequest, Error>)
+    requires itemReq.PutRequest.Some?
+    ensures ret.Success? ==> ret.value.PutRequest.Some?
+  {
+    var partKey := tableConfig.partitionKeyName;
+    var sortKey := tableConfig.sortKeyName;
+    var item := itemReq.PutRequest.value.Item;
+    :- Need(partKey in item, E("Required partition key not in unprocessed item"));
+    :- Need(sortKey.None? || sortKey.value in item, E("Required sort key not in unprocessed item"));
+    for i := 0 to |srcRequests| {
+      if srcRequests[i].PutRequest.Some? {
+        var req := srcRequests[i].PutRequest.value.Item;
+        if partKey in req  && req[partKey] == item[partKey] {
+          if sortKey.None? || (sortKey.value in req  && req[sortKey.value] == item[sortKey.value]) {
+            return Success(srcRequests[i]);
+          }
+        }
+      }
+    }
+    return Failure(E("Item in UnprocessedItems not found in original request."));
+  }
+
   method Output(config: Config, input: BatchWriteItemOutputTransformInput)
     returns (output: Result<BatchWriteItemOutputTransformOutput, Error>)
-    ensures output.Success? && output.value.transformedOutput == input.sdkOutput
   {
-    return Success(BatchWriteItemOutputTransformOutput(transformedOutput := input.sdkOutput));
+    if input.sdkOutput.UnprocessedItems.None? {
+      return Success(BatchWriteItemOutputTransformOutput(transformedOutput := input.sdkOutput));
+    }
+
+    var srcItems := input.originalInput.RequestItems;
+    var unprocessed := input.sdkOutput.UnprocessedItems.value;
+    var tableNames := unprocessed.Keys;
+    var result : map<DDB.TableArn, DDB.WriteRequests> := map[];
+    var tableNamesSeq := SortedSets.ComputeSetToSequence(tableNames);
+    ghost var tableNamesSet' := tableNames;
+    var i := 0;
+    while i < |tableNamesSeq|
+      invariant Seq.HasNoDuplicates(tableNamesSeq)
+      invariant forall j | i <= j < |tableNamesSeq| :: tableNamesSeq[j] in tableNamesSet'
+      invariant |tableNamesSet'| == |tableNamesSeq| - i
+      invariant tableNamesSet' <= unprocessed.Keys
+    {
+      var tableName := tableNamesSeq[i];
+
+      var writeRequests : DDB.WriteRequests := unprocessed[tableName];
+      if !IsPlainWrite(config, tableName) {
+        if tableName !in srcItems {
+          return Failure(E(tableName + " in UnprocessedItems for BatchWriteItem response, but not in original request."));
+        }
+        var srcRequests := srcItems[tableName];
+        var tableConfig := config.tableEncryptionConfigs[tableName];
+        var origItems : seq<DDB.WriteRequest> := [];
+        for x := 0 to |writeRequests|
+          invariant |origItems| == x
+        {
+          var req : DDB.WriteRequest := writeRequests[x];
+          if req.PutRequest.None? {
+            // We only transform PutRequests, so no PutRequest ==> no change
+            origItems := origItems + [req];
+          } else {
+            var orig_item :- GetOrigItem(tableConfig, srcRequests, req);
+            origItems := origItems + [orig_item];
+          }
+        }
+        writeRequests := origItems;
+      }
+      tableNamesSet' := tableNamesSet' - {tableName};
+      i := i + 1;
+      assert forall j | i <= j < |tableNamesSeq| :: tableNamesSeq[j] in tableNamesSet' by {
+        reveal Seq.HasNoDuplicates();
+      }
+      result := result[tableName := writeRequests];
+    }
+    :- Need(|result| ==  |unprocessed|, E("Internal Error")); // Dafny gets too confused
+    var new_output := input.sdkOutput.(UnprocessedItems := Some(result));
+    return Success(BatchWriteItemOutputTransformOutput(transformedOutput := new_output));
   }
 
 }