Skip to content

Commit e983f8f

Browse files
ajewellamzajewellamz
authored
feat: improve perfomance by careful use of maps (#197)
* feat: improve performance by careful use of maps
1 parent 7cde314 commit e983f8f

File tree

7 files changed

+65
-58
lines changed

7 files changed

+65
-58
lines changed

‎DynamoDbEncryption/dafny/DynamoDbEncryption/src/DynamoToStruct.dfy

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,7 @@ module DynamoToStruct {
7171
&& kv.1.content.Terminal.value == TopLevelAttributeToBytes(item[kv.0]).value
7272

7373
{
74-
var structuredMap := map kv <- item.Items | true :: kv.0 := AttrToStructured(kv.1);
74+
var structuredMap := map k <- item :: k := AttrToStructured(item[k]);
7575
MapKeysMatchItems(item);
7676
MapError(SimplifyMapValue(structuredMap))
7777
}
@@ -101,11 +101,11 @@ module DynamoToStruct {
101101
&& kv.1 == StructuredToAttr(s[kv.0]).value
102102
{
103103
if forall k <- s.Keys :: IsValid_AttributeName(k) then
104-
var structuredData := map kv <- s.Items | true :: kv.0 := StructuredToAttr(kv.1);
104+
var structuredData := map k <- s :: k := StructuredToAttr(s[k]);
105105
MapKeysMatchItems(s);
106106
MapError(SimplifyMapValue(structuredData))
107107
else
108-
var badNames := set k <- s.Keys | !IsValid_AttributeName(k) :: k;
108+
var badNames := set k <- s | !IsValid_AttributeName(k) :: k;
109109
OneBadKey(s, badNames, IsValid_AttributeName);
110110
var orderedAttrNames := SetToOrderedSequence(badNames, CharLess);
111111
var attrNameList := Join(orderedAttrNames, ",");
@@ -1081,11 +1081,11 @@ module DynamoToStruct {
10811081
}
10821082

10831083
function method FlattenValueMap<X,Y>(m : map<X, Result<Y,string>>): map<X,Y> {
1084-
map kv <- m.Items | kv.1.Success? :: kv.0 := kv.1.value
1084+
map k <- m | m[k].Success? :: k := m[k].value
10851085
}
10861086

10871087
function method FlattenErrors<X,Y>(m : map<X, Result<Y,string>>): set<string> {
1088-
set k <- m.Values | k.Failure? :: k.error
1088+
set k <- m | m[k].Failure? :: m[k].error
10891089
}
10901090

10911091
lemma OneBadResult<X,Y>(m : map<X, Result<Y,string>>)
@@ -1128,7 +1128,7 @@ module DynamoToStruct {
11281128
ensures ret.Success? ==> |ret.value.Keys| == |m.Keys|
11291129
ensures ret.Success? ==> |ret.value| == |m|
11301130
{
1131-
if forall v <- m.Values :: v.Success? then
1131+
if forall k <- m :: m[k].Success? then
11321132
var result := FlattenValueMap(m);
11331133
MapKeysMatchItems(m);
11341134
Success(result)

‎DynamoDbEncryption/dafny/DynamoDbItemEncryptor/src/AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations.dfy

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -438,7 +438,7 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
438438
ret.value.content.SchemaMap[k].content ==
439439
CSE.CryptoSchemaContent.Action(config.attributeActions[k]))
440440
{
441-
var schema := map kv <- item.Items | true :: kv.0 := GetCryptoSchemaAction(config, kv.0);
441+
var schema := map k <- item :: k := GetCryptoSchemaAction(config, k);
442442
DynamoToStruct.MapKeysMatchItems(item);
443443
DynamoToStruct.SimplifyMapValueSuccess(schema);
444444
var actionMapRes := DynamoToStruct.SimplifyMapValue(schema);
@@ -500,8 +500,8 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
500500
//# An item MUST be determined to be plaintext if it does not contain
501501
//# attributes with the names "aws_dbe_head" and "aws_dbe_foot".
502502
predicate method IsPlaintextItem(ddbItem: ComAmazonawsDynamodbTypes.AttributeMap) {
503-
&& StructuredEncryptionUtil.HeaderField !in ddbItem.Keys
504-
&& StructuredEncryptionUtil.FooterField !in ddbItem.Keys
503+
&& StructuredEncryptionUtil.HeaderField !in ddbItem
504+
&& StructuredEncryptionUtil.FooterField !in ddbItem
505505
}
506506

507507
function method ConvertCryptoSchemaToAttributeActions(config: ValidConfig, schema: CSE.CryptoSchema)
@@ -517,7 +517,7 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
517517
DynamoDbItemEncryptorException( message := "Recieved unexpected Crypto Schema: mismatch with signature scope"));
518518
:- Need(forall k <- schema.content.SchemaMap :: ComAmazonawsDynamodbTypes.IsValid_AttributeName(k),
519519
DynamoDbItemEncryptorException( message := "Recieved unexpected Crypto Schema: Invalid attribute names"));
520-
Success(map k <- schema.content.SchemaMap.Keys | true :: k := schema.content.SchemaMap[k].content.Action)
520+
Success(map k <- schema.content.SchemaMap :: k := schema.content.SchemaMap[k].content.Action)
521521
}
522522

523523
predicate EncryptItemEnsuresPublicly(input: EncryptItemInput, output: Result<EncryptItemOutput, Error>)
@@ -651,8 +651,6 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
651651
return Success(passthroughOutput);
652652
}
653653

654-
assert {:split_here} true;
655-
656654
var plaintextStructure :- DynamoToStruct.ItemToStructured(input.plaintextItem)
657655
.MapFailure(e => Error.AwsCryptographyDbEncryptionSdkDynamoDb(e));
658656
var context :- MakeEncryptionContext(config, plaintextStructure);

‎DynamoDbEncryption/dafny/DynamoDbItemEncryptor/src/Index.dfy

Lines changed: 23 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@ module
1616
import Operations = AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations
1717
import SE = StructuredEncryptionUtil
1818
import InternalLegacyConfig
19+
import SortedSets
20+
import DDB = ComAmazonawsDynamodbTypes
21+
1922

2023
// There is no sensible default, so construct something simple but invalid at runtime.
2124
function method DefaultDynamoDbItemEncryptorConfig(): DynamoDbItemEncryptorConfig
@@ -90,37 +93,40 @@ module
9093
message := "Sort key attribute action MUST be SIGN_ONLY"
9194
));
9295

93-
var attributeActions' := config.attributeActions;
94-
while attributeActions'.Keys != {}
95-
invariant forall attribute <- (config.attributeActions - attributeActions'.Keys)
96-
:: Operations.ForwardCompatibleAttributeAction(
97-
attribute,
98-
config.attributeActions[attribute],
99-
config.allowedUnauthenticatedAttributes,
100-
config.allowedUnauthenticatedAttributePrefix)
101-
invariant forall attribute <- (config.attributeActions - attributeActions'.Keys)
102-
:: UnreservedPrefix(attribute)
96+
var attributeNames : seq<DDB.AttributeName> := SortedSets.ComputeSetToOrderedSequence2(config.attributeActions.Keys, CharLess);
97+
for i := 0 to |attributeNames|
98+
invariant forall j | 0 <= j < i ::
99+
&& UnreservedPrefix(attributeNames[j])
100+
&& (Operations.ForwardCompatibleAttributeAction(
101+
attributeNames[j],
102+
config.attributeActions[attributeNames[j]],
103+
config.allowedUnauthenticatedAttributes,
104+
config.allowedUnauthenticatedAttributePrefix))
103105
{
104-
var attribute :| attribute in attributeActions';
105-
var action := config.attributeActions[attribute];
106+
var attributeName := attributeNames[i];
107+
var action := config.attributeActions[attributeName];
106108
if !(Operations.ForwardCompatibleAttributeAction(
107-
attribute,
109+
attributeName,
108110
action,
109111
config.allowedUnauthenticatedAttributes,
110112
config.allowedUnauthenticatedAttributePrefix
111113
))
112114
{
113115
return Failure(DynamoDbItemEncryptorException(
114-
message := Operations.ExplainNotForwardCompatible(attribute, action, config.allowedUnauthenticatedAttributes, config.allowedUnauthenticatedAttributePrefix)
116+
message := Operations.ExplainNotForwardCompatible(attributeName, action, config.allowedUnauthenticatedAttributes, config.allowedUnauthenticatedAttributePrefix)
115117
));
116118
}
117-
if !UnreservedPrefix(attribute) {
119+
if !UnreservedPrefix(attributeName) {
118120
return Failure(DynamoDbItemEncryptorException(
119-
message := "Attribute: " + attribute + " is reserved, and may not be configured."
121+
message := "Attribute: " + attributeName + " is reserved, and may not be configured."
120122
));
121123
}
122-
attributeActions' := attributeActions' - {attribute};
124+
assert UnreservedPrefix(attributeName);
125+
assert UnreservedPrefix(attributeNames[i]);
123126
}
127+
assert (forall attribute <- attributeNames :: UnreservedPrefix(attribute));
128+
assert (forall attribute <- config.attributeActions.Keys :: UnreservedPrefix(attribute));
129+
assert (forall attribute <- config.attributeActions.Keys :: !(ReservedPrefix <= attribute));
124130

125131
// Create the structured encryption client
126132
var structuredEncryptionRes := StructuredEncryption.StructuredEncryption();

‎DynamoDbEncryption/dafny/StructuredEncryption/src/AwsCryptographyDbEncryptionSdkStructuredEncryptionOperations.dfy

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -285,7 +285,7 @@ module AwsCryptographyDbEncryptionSdkStructuredEncryptionOperations refines Abst
285285
{
286286
reveal Maps.Injective();
287287
Paths.SimpleCanonUnique(tableName);
288-
map k <- data.Keys | schema[k].content.Action != DO_NOTHING :: Paths.SimpleCanon(tableName, k) := k
288+
map k <- data | schema[k].content.Action != DO_NOTHING :: Paths.SimpleCanon(tableName, k) := k
289289
}
290290

291291
// construct the EncryptCanon
@@ -323,7 +323,7 @@ module AwsCryptographyDbEncryptionSdkStructuredEncryptionOperations refines Abst
323323
var data_c := map k <- fieldMap :: k := data[fieldMap[k]];
324324
var signedFields_c := SortedSets.ComputeSetToOrderedSequence2(data_c.Keys, ByteLess);
325325
var encFields_c := FilterEncrypt(signedFields_c, fieldMap, schema);
326-
var trimmedSchema := map k <- fieldMap.Values :: k := schema[k];
326+
var trimmedSchema := map k <- fieldMap :: fieldMap[k] := schema[fieldMap[k]];
327327

328328
assert |data_c| == |trimmedSchema| by {
329329
assert data_c.Keys == fieldMap.Keys;
@@ -379,7 +379,7 @@ module AwsCryptographyDbEncryptionSdkStructuredEncryptionOperations refines Abst
379379

380380
reveal Maps.Injective();
381381
Paths.SimpleCanonUnique(input.tableName);
382-
var fieldMap := map k <- input.data.Keys | input.authSchema[k].content.Action == SIGN ::
382+
var fieldMap := map k <- input.data | input.authSchema[k].content.Action == SIGN ::
383383
Paths.SimpleCanon(input.tableName, k) := k;
384384
assert Maps.Injective(fieldMap);
385385

@@ -401,8 +401,8 @@ module AwsCryptographyDbEncryptionSdkStructuredEncryptionOperations refines Abst
401401
var encFields_c : seq<CanonicalPath> := FilterEncrypted(signedFields_c, input.legend);
402402
:- Need(|encFields_c| < (UINT32_LIMIT / 3), E("Too many encrypted fields."));
403403

404-
var actionMap := map v <- fieldMap.Values ::
405-
v := if Paths.SimpleCanon(input.tableName, v) in encFields_c then
404+
var actionMap := map k <- fieldMap ::
405+
fieldMap[k] := if Paths.SimpleCanon(input.tableName, fieldMap[k]) in encFields_c then
406406
CryptoSchema(
407407
content := CryptoSchemaContent.Action(ENCRYPT_AND_SIGN),
408408
attributes := None
@@ -543,7 +543,7 @@ module AwsCryptographyDbEncryptionSdkStructuredEncryptionOperations refines Abst
543543
:- Need(|canonData.encFields_c| < (UINT32_LIMIT / 3), E("Too many encrypted fields"));
544544
var encryptedItems :- Crypt.Encrypt(config.primitives, alg, key, head, canonData.encFields_c, canonData.data_c);
545545

546-
var result : map<string, StructuredData> := map k | k in plainRecord.Keys :: k :=
546+
var result : map<string, StructuredData> := map k <- plainRecord | true :: k :=
547547
var c := Paths.SimpleCanon(input.tableName, k);
548548
if c in encryptedItems then
549549
encryptedItems[c]
@@ -839,7 +839,7 @@ module AwsCryptographyDbEncryptionSdkStructuredEncryptionOperations refines Abst
839839
canonData.signedFields_c, canonData.encFields_c, map[], canonData.data_c, headerSerialized);
840840
var decryptedItems :- Crypt.Decrypt(config.primitives, postCMMAlg, key, head, canonData.encFields_c, canonData.data_c);
841841

842-
var result : map<string, StructuredData> := map k | k in encRecord.Keys :: k :=
842+
var result : map<string, StructuredData> := map k <- encRecord | true :: k :=
843843
var c := Paths.SimpleCanon(input.tableName, k);
844844
if c in decryptedItems then
845845
decryptedItems[c]

‎DynamoDbEncryption/dafny/StructuredEncryption/src/Crypt.dfy

Lines changed: 20 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ module StructuredEncryptionCrypt {
1313
import opened StandardLibrary.UInt
1414
import opened AwsCryptographyDbEncryptionSdkStructuredEncryptionTypes
1515
import opened StructuredEncryptionUtil
16+
import opened DafnyLibraries
1617

1718
import CMP = AwsCryptographyMaterialProvidersTypes
1819
import Prim = AwsCryptographyPrimitivesTypes
@@ -23,7 +24,6 @@ module StructuredEncryptionCrypt {
2324
import AesKdfCtr
2425
import Seq
2526

26-
2727
function method FieldKey(HKDFOutput : Bytes, offset : uint32)
2828
: (ret : Result<Bytes, Error>)
2929
requires |HKDFOutput| == KeySize
@@ -231,41 +231,44 @@ module StructuredEncryptionCrypt {
231231
//# The calculated Field Root MUST have length equal to the
232232
//# [algorithm suite's encryption key length](../../submodules/MaterialProviders/aws-encryption-sdk-specification/framework/algorithm-suites.md#algorithm-suites-encryption-settings).
233233
assert |fieldRootKey| == AlgorithmSuites.GetEncryptKeyLength(alg) as int;
234-
var result := CryptList(mode, client, alg, fieldRootKey, 0, fieldNames, data, map[]);
234+
var result := CryptList(mode, client, alg, fieldRootKey, fieldNames, data);
235235
return result;
236236
}
237237

238238
// Encrypt or Decrypt each entry in keys, putting results in output
239-
method {:tailrecursion} CryptList(
239+
method CryptList(
240240
mode : EncryptionSelector,
241241
client: Primitives.AtomicPrimitivesClient,
242242
alg : CMP.AlgorithmSuiteInfo,
243243
fieldRootKey : Key,
244-
offset : uint32,
245244
fieldNames : seq<CanonicalPath>,
246-
input : StructuredDataCanon,
247-
output : StructuredDataCanon
245+
input : StructuredDataCanon
248246
)
249247
returns (ret : Result<StructuredDataCanon, Error>)
250248
requires forall k <- fieldNames :: k in input
251-
requires (|fieldNames| + offset as nat) * 3 < UINT32_LIMIT
249+
requires (|fieldNames| as nat) * 3 < UINT32_LIMIT
252250
decreases |fieldNames|
253251

254252
modifies client.Modifies - {client.History} , client.History`AESEncrypt, client.History`AESDecrypt
255253
requires client.ValidState()
256254
ensures client.ValidState()
257255
{
258-
if |fieldNames| == 0 {
259-
return Success(output);
256+
// It is very inefficient to manually build Dafny maps in methods, so use
257+
// a MutableMap to build the key value pairs then convert back to a Dafny map.
258+
var mutMap : MutableMap<CanonicalPath, StructuredDataTerminalType> := new MutableMap();
259+
for i := 0 to |fieldNames| {
260+
var data;
261+
var fieldName := fieldNames[i];
262+
if mode == DoEncrypt {
263+
data :- EncryptTerminal(client, alg, fieldRootKey, i as uint32, fieldName, input[fieldName].content.Terminal);
264+
} else {
265+
data :- DecryptTerminal(client, alg, fieldRootKey, i as uint32, fieldName, input[fieldName].content.Terminal);
266+
}
267+
mutMap.Put(fieldName, data);
260268
}
261-
var data;
262-
if mode == DoEncrypt {
263-
data :- EncryptTerminal(client, alg, fieldRootKey, offset, fieldNames[0], input[fieldNames[0]].content.Terminal);
264-
} else {
265-
data :- DecryptTerminal(client, alg, fieldRootKey, offset, fieldNames[0], input[fieldNames[0]].content.Terminal);
266-
}
267-
var result := CryptList(mode, client, alg, fieldRootKey, offset+1, fieldNames[1..], input, output[fieldNames[0] := data]);
268-
return result;
269+
var mutMapItems := mutMap.content(); // Have to initialize this separately, otherwise the map comprehension will do something very inefficient
270+
var output : StructuredDataCanon := map k <- mutMapItems :: k := mutMap.Select(k);
271+
return Success(output);
269272
}
270273

271274
// Encrypt a single Terminal

‎DynamoDbEncryption/dafny/StructuredEncryption/src/Header.dfy

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -338,9 +338,9 @@ module StructuredEncryptionHeader {
338338
}
339339

340340
function method MyMap<X, Y, Z>(f: X -> Y, m: map<X, Z>): map<Y, Z>
341-
requires forall a, b | a in m.Keys && b in m.Keys :: a != b ==> f(a) != f(b)
341+
requires forall a, b | a in m && b in m :: a != b ==> f(a) != f(b)
342342
{
343-
map k | k in m.Keys :: f(k) := m[k]
343+
map k | k in m :: f(k) := m[k]
344344
}
345345

346346
// Create a Legend from the Schema

‎DynamoDbEncryption/dafny/StructuredEncryption/src/Util.dfy

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -126,21 +126,21 @@ module StructuredEncryptionUtil {
126126
function method CryptoSchemaMapIsFlat(data : CryptoSchemaMap) : (ret : bool)
127127
ensures ret ==> (forall v <- data.Values :: v.content.Action?)
128128
{
129-
forall v <- data.Values :: v.content.Action?
129+
forall k <- data :: data[k].content.Action?
130130
}
131131

132132
// Schema must contain only Actions
133133
function method AuthSchemaIsFlat(data : AuthenticateSchemaMap) : (ret : bool)
134134
ensures ret ==> (forall v <- data.Values :: v.content.Action?)
135135
{
136-
forall v <- data.Values :: v.content.Action?
136+
forall k <- data :: data[k].content.Action?
137137
}
138138

139139
// Map must contain only Terminals
140140
function method DataMapIsFlat(data : StructuredDataMap) : (ret : bool)
141141
ensures ret ==> (forall v <- data.Values :: v.content.Terminal?)
142142
{
143-
forall v <- data.Values :: v.content.Terminal?
143+
forall k <- data :: data[k].content.Terminal?
144144
}
145145

146146
// attribute is "authorized", a.k.a. included in the signature

0 commit comments

Comments
 (0)