Skip to content

Commit 7c7c8a1

Browse files
ajewellamzajewellamz
authored
fix: return plaintext items in UnprocessedItems in BatchWriteIttem (#1642)
* fix: return plaintext items in UnprocessedItems in BatchWriteIttem
1 parent a285eac commit 7c7c8a1

File tree

4 files changed

+331
-30
lines changed

4 files changed

+331
-30
lines changed

‎DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/src/BatchWriteItemTransform.dfy

+89-2
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ module BatchWriteItemTransform {
1313
import Seq
1414
import SortedSets
1515
import Util = DynamoDbEncryptionUtil
16+
import Types = AwsCryptographyDbEncryptionSdkDynamoDbTypes
1617

1718
method {:vcs_split_on_every_assert} Input(config: Config, input: BatchWriteItemInputTransformInput)
1819
returns (output: Result<BatchWriteItemInputTransformOutput, Error>)
@@ -80,11 +81,97 @@ module BatchWriteItemTransform {
8081
return Success(BatchWriteItemInputTransformOutput(transformedInput := input.sdkInput.(RequestItems := result)));
8182
}
8283

84+
// Given the encrypted item, find the original plaintext item, based on the never encrypted PK and SK.
85+
// We don't have to worry about duplicates, because DynamoDB will report an error if there are duplicates
86+
method GetOrigItem(
87+
tableConfig : ValidTableConfig,
88+
srcRequests : DDB.WriteRequests,
89+
itemReq : DDB.WriteRequest
90+
) returns (ret : Result<DDB.WriteRequest, Error>)
91+
requires itemReq.PutRequest.Some?
92+
ensures ret.Success? ==> ret.value.PutRequest.Some?
93+
{
94+
var partKey := tableConfig.partitionKeyName;
95+
var sortKey := tableConfig.sortKeyName;
96+
var item := itemReq.PutRequest.value.Item;
97+
:- Need(partKey in item, E("Required partition key not in unprocessed item"));
98+
:- Need(sortKey.None? || sortKey.value in item, E("Required sort key not in unprocessed item"));
99+
for i := 0 to |srcRequests| {
100+
if srcRequests[i].PutRequest.Some? {
101+
var req := srcRequests[i].PutRequest.value.Item;
102+
if partKey in req && req[partKey] == item[partKey] {
103+
if sortKey.None? || (sortKey.value in req && req[sortKey.value] == item[sortKey.value]) {
104+
return Success(srcRequests[i]);
105+
}
106+
}
107+
}
108+
}
109+
return Failure(E("Item in UnprocessedItems not found in original request."));
110+
}
111+
83112
method Output(config: Config, input: BatchWriteItemOutputTransformInput)
84113
returns (output: Result<BatchWriteItemOutputTransformOutput, Error>)
85-
ensures output.Success? && output.value.transformedOutput == input.sdkOutput
114+
115+
//= specification/dynamodb-encryption-client/ddb-sdk-integration.md#decrypt-after-batchwriteitem
116+
//= type=implication
117+
//# If there are no UnprocessedItems, then the BatchWriteItemOutput MUST be returned unchanged.
118+
ensures input.sdkOutput.UnprocessedItems.None? ==>
119+
&& output.Success?
120+
&& output.value.transformedOutput == input.sdkOutput
86121
{
87-
return Success(BatchWriteItemOutputTransformOutput(transformedOutput := input.sdkOutput));
122+
if input.sdkOutput.UnprocessedItems.None? {
123+
return Success(BatchWriteItemOutputTransformOutput(transformedOutput := input.sdkOutput));
124+
}
125+
126+
//= specification/dynamodb-encryption-client/ddb-sdk-integration.md#decrypt-after-batchwriteitem
127+
//# Each item in UnprocessedItems MUST be replaced by its original plaintext value.
128+
var srcItems := input.originalInput.RequestItems;
129+
var unprocessed := input.sdkOutput.UnprocessedItems.value;
130+
var tableNames := unprocessed.Keys;
131+
var result : map<DDB.TableArn, DDB.WriteRequests> := map[];
132+
var tableNamesSeq := SortedSets.ComputeSetToSequence(tableNames);
133+
ghost var tableNamesSet' := tableNames;
134+
var i := 0;
135+
while i < |tableNamesSeq|
136+
invariant Seq.HasNoDuplicates(tableNamesSeq)
137+
invariant forall j | i <= j < |tableNamesSeq| :: tableNamesSeq[j] in tableNamesSet'
138+
invariant |tableNamesSet'| == |tableNamesSeq| - i
139+
invariant tableNamesSet' <= unprocessed.Keys
140+
{
141+
var tableName := tableNamesSeq[i];
142+
143+
var writeRequests : DDB.WriteRequests := unprocessed[tableName];
144+
if !IsPlainWrite(config, tableName) {
145+
if tableName !in srcItems {
146+
return Failure(E(tableName + " in UnprocessedItems for BatchWriteItem response, but not in original request."));
147+
}
148+
var srcRequests := srcItems[tableName];
149+
var tableConfig := config.tableEncryptionConfigs[tableName];
150+
var origItems : seq<DDB.WriteRequest> := [];
151+
for x := 0 to |writeRequests|
152+
invariant |origItems| == x
153+
{
154+
var req : DDB.WriteRequest := writeRequests[x];
155+
if req.PutRequest.None? {
156+
// We only transform PutRequests, so no PutRequest ==> no change
157+
origItems := origItems + [req];
158+
} else {
159+
var orig_item :- GetOrigItem(tableConfig, srcRequests, req);
160+
origItems := origItems + [orig_item];
161+
}
162+
}
163+
writeRequests := origItems;
164+
}
165+
tableNamesSet' := tableNamesSet' - {tableName};
166+
i := i + 1;
167+
assert forall j | i <= j < |tableNamesSeq| :: tableNamesSeq[j] in tableNamesSet' by {
168+
reveal Seq.HasNoDuplicates();
169+
}
170+
result := result[tableName := writeRequests];
171+
}
172+
:- Need(|result| == |unprocessed|, E("Internal Error")); // Dafny gets too confused
173+
var new_output := input.sdkOutput.(UnprocessedItems := Some(result));
174+
return Success(BatchWriteItemOutputTransformOutput(transformedOutput := new_output));
88175
}
89176

90177
}

‎DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/test/BatchWriteItemTransform.dfy

+203-3
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ module BatchWriteItemTransformTest {
88
import opened DynamoDbEncryptionTransforms
99
import opened TestFixtures
1010
import DDB = ComAmazonawsDynamodbTypes
11-
import AwsCryptographyDbEncryptionSdkDynamoDbTransformsTypes
11+
import TTypes = AwsCryptographyDbEncryptionSdkDynamoDbTransformsTypes
1212

1313
method {:test} TestBatchWriteItemInputTransform() {
1414
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms();
@@ -25,7 +25,7 @@ module BatchWriteItemTransformTest {
2525
ReturnItemCollectionMetrics := None
2626
);
2727
var transformed := middlewareUnderTest.BatchWriteItemInputTransform(
28-
AwsCryptographyDbEncryptionSdkDynamoDbTransformsTypes.BatchWriteItemInputTransformInput(
28+
TTypes.BatchWriteItemInputTransformInput(
2929
sdkInput := input
3030
)
3131
);
@@ -34,6 +34,35 @@ module BatchWriteItemTransformTest {
3434
expect_equal("BatchWriteItemInput", transformed.value.transformedInput, input);
3535
}
3636

37+
38+
const Item1 : DDB.AttributeMap := map[
39+
"bar" := DDB.AttributeValue.S("bar1"),
40+
"sign" := DDB.AttributeValue.S("sign1"),
41+
"encrypt" := DDB.AttributeValue.S("encrypt1"),
42+
"plain" := DDB.AttributeValue.S("plain1")
43+
]
44+
45+
const Item2 : DDB.AttributeMap := map[
46+
"bar" := DDB.AttributeValue.S("bar2"),
47+
"sign" := DDB.AttributeValue.S("sign2"),
48+
"encrypt" := DDB.AttributeValue.S("encrypt2"),
49+
"plain" := DDB.AttributeValue.S("plain2")
50+
]
51+
52+
const Item3 : DDB.AttributeMap := map[
53+
"bar" := DDB.AttributeValue.S("bar3"),
54+
"sign" := DDB.AttributeValue.S("sign3"),
55+
"encrypt" := DDB.AttributeValue.S("encrypt3"),
56+
"plain" := DDB.AttributeValue.S("plain3")
57+
]
58+
59+
const Item4 : DDB.AttributeMap := map[
60+
"bar" := DDB.AttributeValue.S("bar4"),
61+
"sign" := DDB.AttributeValue.S("sign4"),
62+
"encrypt" := DDB.AttributeValue.S("encrypt4"),
63+
"plain" := DDB.AttributeValue.S("plain4")
64+
]
65+
3766
method {:test} TestBatchWriteItemOutputTransform() {
3867
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms();
3968
var output := DDB.BatchWriteItemOutput(
@@ -55,7 +84,7 @@ module BatchWriteItemTransformTest {
5584
ReturnItemCollectionMetrics := None
5685
);
5786
var transformed := middlewareUnderTest.BatchWriteItemOutputTransform(
58-
AwsCryptographyDbEncryptionSdkDynamoDbTransformsTypes.BatchWriteItemOutputTransformInput(
87+
TTypes.BatchWriteItemOutputTransformInput(
5988
sdkOutput := output,
6089
originalInput := input
6190
)
@@ -65,4 +94,175 @@ module BatchWriteItemTransformTest {
6594
expect_equal("BatchWriteItemOutput", transformed.value.transformedOutput, output);
6695
}
6796

97+
function method MakePut(item : DDB.AttributeMap) : DDB.WriteRequest
98+
{
99+
DDB.WriteRequest ( DeleteRequest := None, PutRequest := Some(DDB.PutRequest(Item := item)))
100+
}
101+
102+
//= specification/dynamodb-encryption-client/ddb-sdk-integration.md#decrypt-after-batchwriteitem
103+
//= type=test
104+
//# Each item in UnprocessedItems MUST be replaced by its original plaintext value.
105+
method {:test} TestBatchWriteItemOutputTransformUnprocessed() {
106+
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms();
107+
var tableName := GetTableName("foo");
108+
109+
var theRequests := GetBatchWriteItemRequestMap(map[tableName := [MakePut(Item1), MakePut(Item2), MakePut(Item3), MakePut(Item4)]]);
110+
111+
var input := DDB.BatchWriteItemInput(
112+
RequestItems := theRequests
113+
);
114+
var transInput :- expect middlewareUnderTest.BatchWriteItemInputTransform(
115+
TTypes.BatchWriteItemInputTransformInput(
116+
sdkInput := input
117+
)
118+
);
119+
120+
var unProcessed : DDB.BatchWriteItemRequestMap := transInput.transformedInput.RequestItems;
121+
expect unProcessed != input.RequestItems;
122+
var output := DDB.BatchWriteItemOutput(
123+
UnprocessedItems := Some(unProcessed)
124+
);
125+
126+
var transOutput :- expect middlewareUnderTest.BatchWriteItemOutputTransform(
127+
TTypes.BatchWriteItemOutputTransformInput(
128+
sdkOutput := output,
129+
originalInput := input
130+
)
131+
);
132+
133+
expect_equal("BatchWriteOutput", transOutput.transformedOutput.UnprocessedItems, Some(theRequests));
134+
}
135+
136+
method {:test} TestBatchWriteItemOutputTransformUnprocessed2() {
137+
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms();
138+
var tableName1 := GetTableName("foo");
139+
var tableName2 := GetTableName("baz");
140+
141+
var theRequests := GetBatchWriteItemRequestMap(map[tableName1 := [MakePut(Item1), MakePut(Item2)], tableName2 := [MakePut(Item3), MakePut(Item4)]]);
142+
143+
var input := DDB.BatchWriteItemInput(
144+
RequestItems := theRequests
145+
);
146+
var transInput :- expect middlewareUnderTest.BatchWriteItemInputTransform(
147+
TTypes.BatchWriteItemInputTransformInput(
148+
sdkInput := input
149+
)
150+
);
151+
152+
var unProcessed : DDB.BatchWriteItemRequestMap := transInput.transformedInput.RequestItems;
153+
expect unProcessed != input.RequestItems;
154+
var output := DDB.BatchWriteItemOutput(
155+
UnprocessedItems := Some(unProcessed)
156+
);
157+
158+
var transOutput :- expect middlewareUnderTest.BatchWriteItemOutputTransform(
159+
TTypes.BatchWriteItemOutputTransformInput(
160+
sdkOutput := output,
161+
originalInput := input
162+
)
163+
);
164+
165+
expect_equal("BatchWriteOutput", transOutput.transformedOutput.UnprocessedItems, Some(theRequests));
166+
}
167+
168+
method {:test} {:vcs_split_on_every_assert} TestBatchWriteItemOutputTransformUnprocessed3() {
169+
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms();
170+
var tableName1 := GetTableName("foo");
171+
var tableName2 := GetTableName("baz");
172+
173+
var theRequests := GetBatchWriteItemRequestMap(map[tableName1 := [MakePut(Item1), MakePut(Item2)], tableName2 := [MakePut(Item3), MakePut(Item4)]]);
174+
175+
var input := DDB.BatchWriteItemInput(
176+
RequestItems := theRequests
177+
);
178+
var transInput :- expect middlewareUnderTest.BatchWriteItemInputTransform(
179+
TTypes.BatchWriteItemInputTransformInput(
180+
sdkInput := input
181+
)
182+
);
183+
184+
expect tableName1 in transInput.transformedInput.RequestItems;
185+
expect tableName2 in transInput.transformedInput.RequestItems;
186+
187+
var list := map[
188+
tableName1 := transInput.transformedInput.RequestItems[tableName1][0..1],
189+
tableName2 := transInput.transformedInput.RequestItems[tableName2][0..1]
190+
];
191+
expect DDB.IsValid_BatchWriteItemRequestMap(list);
192+
var unProcessed : DDB.BatchWriteItemRequestMap := list;
193+
194+
var orig_list := map [
195+
tableName1 := [MakePut(Item1)],
196+
tableName2 := [MakePut(Item3)]
197+
];
198+
expect DDB.IsValid_BatchWriteItemRequestMap(orig_list);
199+
var originalUnProcessed : DDB.BatchWriteItemRequestMap := orig_list;
200+
201+
expect unProcessed != input.RequestItems;
202+
var output := DDB.BatchWriteItemOutput(
203+
UnprocessedItems := Some(unProcessed)
204+
);
205+
206+
var transOutput :- expect middlewareUnderTest.BatchWriteItemOutputTransform(
207+
TTypes.BatchWriteItemOutputTransformInput(
208+
sdkOutput := output,
209+
originalInput := input
210+
)
211+
);
212+
213+
expect_equal("BatchWriteOutput", transOutput.transformedOutput.UnprocessedItems, Some(originalUnProcessed));
214+
}
215+
216+
const Item1a : DDB.AttributeMap := map[
217+
"bar" := DDB.AttributeValue.S("bar1"),
218+
"sign" := DDB.AttributeValue.S("sign2"),
219+
"encrypt" := DDB.AttributeValue.S("encrypt2"),
220+
"plain" := DDB.AttributeValue.S("plain2")
221+
]
222+
223+
const Item1b : DDB.AttributeMap := map[
224+
"bar" := DDB.AttributeValue.S("bar1"),
225+
"sign" := DDB.AttributeValue.S("sign3"),
226+
"encrypt" := DDB.AttributeValue.S("encrypt3"),
227+
"plain" := DDB.AttributeValue.S("plain3")
228+
]
229+
230+
const Item1c : DDB.AttributeMap := map[
231+
"bar" := DDB.AttributeValue.S("bar1"),
232+
"sign" := DDB.AttributeValue.S("sign4"),
233+
"encrypt" := DDB.AttributeValue.S("encrypt4"),
234+
"plain" := DDB.AttributeValue.S("plain4")
235+
]
236+
237+
method {:test} TestBatchWriteItemOutputTransformUnprocessed4() {
238+
var middlewareUnderTest := TestFixtures.GetDynamoDbEncryptionTransforms2();
239+
var tableName1 := GetTableName("foo");
240+
var tableName2 := GetTableName("baz");
241+
242+
var theRequests := GetBatchWriteItemRequestMap(map[tableName1 := [MakePut(Item1), MakePut(Item1a)], tableName2 := [MakePut(Item1b), MakePut(Item1c)]]);
243+
244+
var input := DDB.BatchWriteItemInput(
245+
RequestItems := theRequests
246+
);
247+
var transInput :- expect middlewareUnderTest.BatchWriteItemInputTransform(
248+
TTypes.BatchWriteItemInputTransformInput(
249+
sdkInput := input
250+
)
251+
);
252+
253+
var unProcessed : DDB.BatchWriteItemRequestMap := transInput.transformedInput.RequestItems;
254+
expect unProcessed != input.RequestItems;
255+
var output := DDB.BatchWriteItemOutput(
256+
UnprocessedItems := Some(unProcessed)
257+
);
258+
259+
var transOutput :- expect middlewareUnderTest.BatchWriteItemOutputTransform(
260+
TTypes.BatchWriteItemOutputTransformInput(
261+
sdkOutput := output,
262+
originalInput := input
263+
)
264+
);
265+
266+
expect_equal("BatchWriteOutput", transOutput.transformedOutput.UnprocessedItems, Some(theRequests));
267+
}
68268
}

0 commit comments

Comments
 (0)