Skip to content

fix: return plaintext items in UnprocessedItems in BatchWriteIttem #1642

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Feb 11, 2025
Merged

fix: return plaintext items in UnprocessedItems in BatchWriteIttem #1642

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4446d78
fix: return plaintext items in UnprocessedItems in BatchWriteIttem
ajewellamz Feb 11, 2025
4bb289a
m
ajewellamz Feb 11, 2025
07946e7
m
ajewellamz Feb 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 89 additions & 2 deletions DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/src/BatchWriteItemTransform.dfy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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>)
Expand Down Expand Up @@ -80,11 +81,97 @@ module BatchWriteItemTransform {
return Success(BatchWriteItemInputTransformOutput(transformedInput := input.sdkInput.(RequestItems := result)));
}

// Given the encrypted item, find the original plaintext item, based on the never encrypted PK and SK.
// We don't have to worry about duplicates, because DynamoDB will report an error if there are duplicates
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

//= specification/dynamodb-encryption-client/ddb-sdk-integration.md#decrypt-after-batchwriteitem
//= type=implication
//# If there are no UnprocessedItems, then the BatchWriteItemOutput MUST be returned unchanged.
ensures input.sdkOutput.UnprocessedItems.None? ==>
&& output.Success?
&& output.value.transformedOutput == input.sdkOutput
{
return Success(BatchWriteItemOutputTransformOutput(transformedOutput := input.sdkOutput));
if input.sdkOutput.UnprocessedItems.None? {
return Success(BatchWriteItemOutputTransformOutput(transformedOutput := input.sdkOutput));
}

//= specification/dynamodb-encryption-client/ddb-sdk-integration.md#decrypt-after-batchwriteitem
//# Each item in UnprocessedItems MUST be replaced by its original plaintext value.
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));
}

}
206 changes: 203 additions & 3 deletions DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/test/BatchWriteItemTransform.dfy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ module BatchWriteItemTransformTest {
import opened DynamoDbEncryptionTransforms
import opened TestFixtures
import DDB = ComAmazonawsDynamodbTypes
import AwsCryptographyDbEncryptionSdkDynamoDbTransformsTypes
import TTypes = AwsCryptographyDbEncryptionSdkDynamoDbTransformsTypes

method {:test} TestBatchWriteItemInputTransform() {
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms();
Expand All @@ -25,7 +25,7 @@ module BatchWriteItemTransformTest {
ReturnItemCollectionMetrics := None
);
var transformed := middlewareUnderTest.BatchWriteItemInputTransform(
AwsCryptographyDbEncryptionSdkDynamoDbTransformsTypes.BatchWriteItemInputTransformInput(
TTypes.BatchWriteItemInputTransformInput(
sdkInput := input
)
);
Expand All @@ -34,6 +34,35 @@ module BatchWriteItemTransformTest {
expect_equal("BatchWriteItemInput", transformed.value.transformedInput, input);
}


const Item1 : DDB.AttributeMap := map[
"bar" := DDB.AttributeValue.S("bar1"),
"sign" := DDB.AttributeValue.S("sign1"),
"encrypt" := DDB.AttributeValue.S("encrypt1"),
"plain" := DDB.AttributeValue.S("plain1")
]

const Item2 : DDB.AttributeMap := map[
"bar" := DDB.AttributeValue.S("bar2"),
"sign" := DDB.AttributeValue.S("sign2"),
"encrypt" := DDB.AttributeValue.S("encrypt2"),
"plain" := DDB.AttributeValue.S("plain2")
]

const Item3 : DDB.AttributeMap := map[
"bar" := DDB.AttributeValue.S("bar3"),
"sign" := DDB.AttributeValue.S("sign3"),
"encrypt" := DDB.AttributeValue.S("encrypt3"),
"plain" := DDB.AttributeValue.S("plain3")
]

const Item4 : DDB.AttributeMap := map[
"bar" := DDB.AttributeValue.S("bar4"),
"sign" := DDB.AttributeValue.S("sign4"),
"encrypt" := DDB.AttributeValue.S("encrypt4"),
"plain" := DDB.AttributeValue.S("plain4")
]

method {:test} TestBatchWriteItemOutputTransform() {
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms();
var output := DDB.BatchWriteItemOutput(
Expand All @@ -55,7 +84,7 @@ module BatchWriteItemTransformTest {
ReturnItemCollectionMetrics := None
);
var transformed := middlewareUnderTest.BatchWriteItemOutputTransform(
AwsCryptographyDbEncryptionSdkDynamoDbTransformsTypes.BatchWriteItemOutputTransformInput(
TTypes.BatchWriteItemOutputTransformInput(
sdkOutput := output,
originalInput := input
)
Expand All @@ -65,4 +94,175 @@ module BatchWriteItemTransformTest {
expect_equal("BatchWriteItemOutput", transformed.value.transformedOutput, output);
}

function method MakePut(item : DDB.AttributeMap) : DDB.WriteRequest
{
DDB.WriteRequest ( DeleteRequest := None, PutRequest := Some(DDB.PutRequest(Item := item)))
}

//= specification/dynamodb-encryption-client/ddb-sdk-integration.md#decrypt-after-batchwriteitem
//= type=test
//# Each item in UnprocessedItems MUST be replaced by its original plaintext value.
method {:test} TestBatchWriteItemOutputTransformUnprocessed() {
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms();
var tableName := GetTableName("foo");

var theRequests := GetBatchWriteItemRequestMap(map[tableName := [MakePut(Item1), MakePut(Item2), MakePut(Item3), MakePut(Item4)]]);

var input := DDB.BatchWriteItemInput(
RequestItems := theRequests
);
var transInput :- expect middlewareUnderTest.BatchWriteItemInputTransform(
TTypes.BatchWriteItemInputTransformInput(
sdkInput := input
)
);

var unProcessed : DDB.BatchWriteItemRequestMap := transInput.transformedInput.RequestItems;
expect unProcessed != input.RequestItems;
var output := DDB.BatchWriteItemOutput(
UnprocessedItems := Some(unProcessed)
);

var transOutput :- expect middlewareUnderTest.BatchWriteItemOutputTransform(
TTypes.BatchWriteItemOutputTransformInput(
sdkOutput := output,
originalInput := input
)
);

expect_equal("BatchWriteOutput", transOutput.transformedOutput.UnprocessedItems, Some(theRequests));
}

method {:test} TestBatchWriteItemOutputTransformUnprocessed2() {
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms();
var tableName1 := GetTableName("foo");
var tableName2 := GetTableName("baz");

var theRequests := GetBatchWriteItemRequestMap(map[tableName1 := [MakePut(Item1), MakePut(Item2)], tableName2 := [MakePut(Item3), MakePut(Item4)]]);

var input := DDB.BatchWriteItemInput(
RequestItems := theRequests
);
var transInput :- expect middlewareUnderTest.BatchWriteItemInputTransform(
TTypes.BatchWriteItemInputTransformInput(
sdkInput := input
)
);

var unProcessed : DDB.BatchWriteItemRequestMap := transInput.transformedInput.RequestItems;
expect unProcessed != input.RequestItems;
var output := DDB.BatchWriteItemOutput(
UnprocessedItems := Some(unProcessed)
);

var transOutput :- expect middlewareUnderTest.BatchWriteItemOutputTransform(
TTypes.BatchWriteItemOutputTransformInput(
sdkOutput := output,
originalInput := input
)
);

expect_equal("BatchWriteOutput", transOutput.transformedOutput.UnprocessedItems, Some(theRequests));
}

method {:test} {:vcs_split_on_every_assert} TestBatchWriteItemOutputTransformUnprocessed3() {
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms();
var tableName1 := GetTableName("foo");
var tableName2 := GetTableName("baz");

var theRequests := GetBatchWriteItemRequestMap(map[tableName1 := [MakePut(Item1), MakePut(Item2)], tableName2 := [MakePut(Item3), MakePut(Item4)]]);

var input := DDB.BatchWriteItemInput(
RequestItems := theRequests
);
var transInput :- expect middlewareUnderTest.BatchWriteItemInputTransform(
TTypes.BatchWriteItemInputTransformInput(
sdkInput := input
)
);

expect tableName1 in transInput.transformedInput.RequestItems;
expect tableName2 in transInput.transformedInput.RequestItems;

var list := map[
tableName1 := transInput.transformedInput.RequestItems[tableName1][0..1],
tableName2 := transInput.transformedInput.RequestItems[tableName2][0..1]
];
expect DDB.IsValid_BatchWriteItemRequestMap(list);
var unProcessed : DDB.BatchWriteItemRequestMap := list;

var orig_list := map [
tableName1 := [MakePut(Item1)],
tableName2 := [MakePut(Item3)]
];
expect DDB.IsValid_BatchWriteItemRequestMap(orig_list);
var originalUnProcessed : DDB.BatchWriteItemRequestMap := orig_list;

expect unProcessed != input.RequestItems;
var output := DDB.BatchWriteItemOutput(
UnprocessedItems := Some(unProcessed)
);

var transOutput :- expect middlewareUnderTest.BatchWriteItemOutputTransform(
TTypes.BatchWriteItemOutputTransformInput(
sdkOutput := output,
originalInput := input
)
);

expect_equal("BatchWriteOutput", transOutput.transformedOutput.UnprocessedItems, Some(originalUnProcessed));
}

const Item1a : DDB.AttributeMap := map[
"bar" := DDB.AttributeValue.S("bar1"),
"sign" := DDB.AttributeValue.S("sign2"),
"encrypt" := DDB.AttributeValue.S("encrypt2"),
"plain" := DDB.AttributeValue.S("plain2")
]

const Item1b : DDB.AttributeMap := map[
"bar" := DDB.AttributeValue.S("bar1"),
"sign" := DDB.AttributeValue.S("sign3"),
"encrypt" := DDB.AttributeValue.S("encrypt3"),
"plain" := DDB.AttributeValue.S("plain3")
]

const Item1c : DDB.AttributeMap := map[
"bar" := DDB.AttributeValue.S("bar1"),
"sign" := DDB.AttributeValue.S("sign4"),
"encrypt" := DDB.AttributeValue.S("encrypt4"),
"plain" := DDB.AttributeValue.S("plain4")
]

method {:test} TestBatchWriteItemOutputTransformUnprocessed4() {
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms2();
var tableName1 := GetTableName("foo");
var tableName2 := GetTableName("baz");

var theRequests := GetBatchWriteItemRequestMap(map[tableName1 := [MakePut(Item1), MakePut(Item1a)], tableName2 := [MakePut(Item1b), MakePut(Item1c)]]);

var input := DDB.BatchWriteItemInput(
RequestItems := theRequests
);
var transInput :- expect middlewareUnderTest.BatchWriteItemInputTransform(
TTypes.BatchWriteItemInputTransformInput(
sdkInput := input
)
);

var unProcessed : DDB.BatchWriteItemRequestMap := transInput.transformedInput.RequestItems;
expect unProcessed != input.RequestItems;
var output := DDB.BatchWriteItemOutput(
UnprocessedItems := Some(unProcessed)
);

var transOutput :- expect middlewareUnderTest.BatchWriteItemOutputTransform(
TTypes.BatchWriteItemOutputTransformInput(
sdkOutput := output,
originalInput := input
)
);

expect_equal("BatchWriteOutput", transOutput.transformedOutput.UnprocessedItems, Some(theRequests));
}
}
Loading
Loading