chore: further performance improvements

DynamoDbEncryption/dafny/DynamoDbEncryption/src/DynamoToStruct.dfy

@@ -77,7 +77,7 @@ module DynamoToStruct {
                 && StructuredToAttr(s[kv.0]).Success?
                 && kv.1 == StructuredToAttr(s[kv.0]).value
   {
-    if forall k <- s.Keys :: IsValid_AttributeName(k) then
+    if forall k <- s :: IsValid_AttributeName(k) then
       var structuredData := map k <- s :: k := StructuredToAttr(s[k]);
       MapKeysMatchItems(s);
       MapError(SimplifyMapValue(structuredData))
@@ -477,7 +477,7 @@ module DynamoToStruct {
   // lets Dafny find the correct termination metric.
   // See "The Parent Trick" for details: <https://leino.science/papers/krml283.html>.
   function method MapAttrToBytes(ghost parent: AttributeValue, m: MapAttributeValue, depth : nat): (ret: Result<seq<uint8>, string>)
-    requires forall kv <- m.Items :: kv.1 < parent
+    requires forall k <- m :: m[k] < parent
   {
     //= specification/dynamodb-encryption-client/ddb-attribute-serialization.md#value-type
     //# Value Type MUST be the [Type ID](#type-id) of the type of [Map Value](#map-value).
@@ -492,7 +492,8 @@ module DynamoToStruct {
     //= specification/dynamodb-encryption-client/ddb-attribute-serialization.md#map-value
     //# A Map MAY hold any DynamoDB Attribute Value data type,
     //# and MAY hold values of different types.
-    var bytesResults := map kv <- m.Items :: kv.0 := AttrToBytes(kv.1, true, depth+1);
+    var bytesResults := map k <- m :: k := AttrToBytes(m[k], true, depth+1);
+
     var count :- U32ToBigEndian(|m|);
     var bytes :- SimplifyMapValue(bytesResults);
     var body :- CollectMap(bytes);

DynamoDbEncryption/dafny/StructuredEncryption/src/Header.dfy

@@ -57,12 +57,23 @@ module StructuredEncryptionHeader {
   type Legend = x : seq<LegendByte> | |x| < UINT16_LIMIT
   type CMPUtf8Bytes = x : CMP.Utf8Bytes | |x| < UINT16_LIMIT
 
-  predicate method IsVersion2Schema(data : CanonCryptoList)
-  {
-    exists x <- data :: x.action == SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT
+  predicate method {:tailrecursion} IsVersion2Schema(data : CanonCryptoList, pos : uint32 := 0) : (ret : bool)
+    requires |data| < UINT32_LIMIT
+    requires pos <= |data| as uint32
+    requires forall i | 0 <= i < pos :: data[i].action != SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT
+    decreases |data| as uint32 - pos
+    ensures ret <==> (exists x <- data :: x.action == SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT)
+  {
+    if pos == |data| as uint32 then
+      false
+    else if data[pos].action == SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT then
+      true
+    else
+      IsVersion2Schema(data, pos+1)
   }
 
   function method VersionFromSchema(data : CanonCryptoList) : (ret : Version)
+    requires |data| < UINT32_LIMIT
     ensures (exists x <- data :: x.action == SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT) <==> (ret == 2)
     ensures !(exists x <- data :: x.action == SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT) <==> (ret == 1)
   {
@@ -229,8 +240,9 @@ module StructuredEncryptionHeader {
     //# If any [Crypto Action](./structures.md#crypto-action) is configured as
     //# [SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT Crypto Action](./structures.md#sign_and_include_in_encryption_context)
     //# the Version MUST be 0x02; otherwise, Version MUST be 0x01.
-    ensures ret.Success? ==> ret.value.version == VersionFromSchema(schema)
+    ensures ret.Success? ==> |schema| < UINT32_LIMIT && ret.value.version == VersionFromSchema(schema)
   {
+    :- Need(|schema| < UINT32_LIMIT, E("Unexpected large schema"));
     :- Need(ValidEncryptionContext(mat.encryptionContext), E("Invalid Encryption Context"));
     :- Need(0 < |mat.encryptedDataKeys|, E("There must be at least one data key"));
     :- Need(|mat.encryptedDataKeys| < UINT8_LIMIT, E("Too many data keys."));

DynamoDbEncryption/dafny/StructuredEncryption/src/Util.dfy

@@ -220,7 +220,7 @@ module StructuredEncryptionUtil {
   // attribute is "authorized", a.k.a. included in the signature
   predicate method IsAuthAttr(x : CryptoAction)
   {
-    x.ENCRYPT_AND_SIGN? || x.SIGN_AND_INCLUDE_IN_ENCRYPTION_CONTEXT? || x.SIGN_ONLY?
+    !x.DO_NOTHING?
   }
 
   // wrap a value in a StructuredData

TestVectors/dafny/DDBEncryption/src/Index.dfy

@@ -61,6 +61,7 @@ module WrappedDDBEncryptionMain {
     config :- expect AddJson(config, "data.json", keyVectors);
     config :- expect AddJson(config, "iotest.json", keyVectors);
     config :- expect AddJson(config, "PermTest.json", keyVectors);
+    config :- expect AddJson(config, "large_records.json", keyVectors);
     config.RunAllTests(keyVectors);
   }
 }

TestVectors/dafny/DDBEncryption/src/JsonConfig.dfy

@@ -46,9 +46,15 @@ module {:options "-functionSyntax:4"} JsonConfig {
     item : DDB.AttributeMap
   )
 
+  datatype LargeRecord = LargeRecord (
+    name : string,
+    count : nat,
+    item : DDB.AttributeMap
+  )
+
   datatype TableConfig = TableConfig (
     name : ConfigName,
-    config : Types.DynamoDbTableEncryptionConfig,
+    config : AwsCryptographyDbEncryptionSdkDynamoDbTypes.DynamoDbTableEncryptionConfig,
     vanilla : bool
   )
 
@@ -1410,7 +1416,6 @@ module {:options "-functionSyntax:4"} JsonConfig {
     return Success(results);
   }
 
-
   method GetRecord(data : JSON) returns (output : Result<Record, string>)
   {
     var item :- JsonToDdbItem(data);
@@ -1424,4 +1429,43 @@ module {:options "-functionSyntax:4"} JsonConfig {
     var num :- StrToNat(hash.N);
     return Success(Record(num, item));
   }
+
+  method GetLarge(name : string, data : JSON) returns (output : Result<LargeRecord, string>)
+  {
+    :- Need(data.Object?, "LargeRecord must be a JSON object.");
+    var count := 0;
+    var item : DDB.AttributeMap := map[];
+    for i := 0 to |data.obj| {
+      var obj := data.obj[i];
+      match obj.0 {
+        case "Count" =>
+          :- Need(obj.1.Number?, "GetPrefix length must be a number");
+          count :- DecimalToInt(obj.1.num);
+        case "Item" => var src :- JsonToDdbItem(obj.1); item := src;
+        case _ => return Failure("Unexpected part of a LargeRecord : '" + obj.0 + "'");
+      }
+    }
+    if (count <= 0) {
+      return Failure("Missing or invalid Count in LargeRecord : '" + name + "'");
+    }
+    if (|item| == 0) {
+      return Failure("Missing or Empty LargeRecord : '" + name + "'");
+    }
+    var record := LargeRecord(name, count, item);
+    return Success(record);
+  }
+
+  method GetLarges(data : JSON) returns (output : Result<seq<LargeRecord>, string>)
+  {
+    :- Need(data.Object?, "Larges must be a JSON object.");
+    var results : seq<LargeRecord> := [];
+    for i := 0 to |data.obj| {
+      var obj := data.obj[i];
+      var record :- GetLarge(obj.0, obj.1);
+      results := results + [record];
+    }
+    return Success(results);
+  }
+
+
 }

TestVectors/dafny/DDBEncryption/src/TestVectors.dfy

@@ -41,12 +41,18 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
   import MPT = AwsCryptographyMaterialProvidersTypes
   import Primitives = AtomicPrimitives
   import ParseJsonManifests
+  import Time
+  import Trans = AwsCryptographyDbEncryptionSdkDynamoDbTransformsTypes
+  import TransOp = AwsCryptographyDbEncryptionSdkDynamoDbTransformsOperations
+  import DdbMiddlewareConfig
+  import DynamoDbEncryptionTransforms
 
 
   datatype TestVectorConfig = TestVectorConfig (
     schemaOnEncrypt : DDB.CreateTableInput,
     globalRecords : seq<Record>,
     tableEncryptionConfigs : map<ConfigName, TableConfig>,
+    largeEncryptionConfigs : map<ConfigName, TableConfig>,
     queries : seq<SimpleQuery>,
     names : DDB.ExpressionAttributeNameMap,
     values : DDB.ExpressionAttributeValueMap,
@@ -58,7 +64,8 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
     writeTests : seq<WriteTest>,
     roundTripTests : seq<RoundTripTest>,
     decryptTests : seq<DecryptTest>,
-    strings : seq<string>
+    strings : seq<string>,
+    large : seq<LargeRecord>
   ) {
 
     method RunAllTests(keyVectors: KeyVectors.KeyVectorsClient)
@@ -69,6 +76,7 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
       print "DBE Test Vectors\n";
       print |globalRecords|, " records.\n";
       print |tableEncryptionConfigs|, " tableEncryptionConfigs.\n";
+      print |largeEncryptionConfigs|, " largeEncryptionConfigs.\n";
       print |queries|, " queries.\n";
       print |names|, " names.\n";
       print |values|, " values.\n";
@@ -78,6 +86,7 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
       print |configsForIoTest|, " configsForIoTest.\n";
       print |configsForModTest|, " configsForModTest.\n";
       print |strings|, " strings.\n";
+      print |large|, " large.\n";
       if |roundTripTests| != 0 {
         print |roundTripTests[0].configs|, " configs and ", |roundTripTests[0].records|, " records for round trip.\n";
       }
@@ -107,6 +116,7 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
         return;
       }
       StringOrdering();
+      LargeTests();
       BasicIoTest();
       RunIoTests();
       BasicQueryTest();
@@ -484,6 +494,87 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
       }
     }
 
+    const TestConfigs : set<string> := {"all"}
+    const TestRecords : set<string> := {"all"}
+
+    predicate DoTestConfig(name : string)
+    {
+      "all" in TestConfigs || name in TestConfigs
+    }
+
+    predicate DoTestRecord(name : string)
+    {
+      "all" in TestRecords || name in TestRecords
+    }
+
+    method LargeTests()
+    {
+      print "LargeTests\n";
+      DoLargeTest("do_nothing");
+      DoLargeTest("do_nothing_nosign");
+      DoLargeTest("full_encrypt");
+      DoLargeTest("full_encrypt_nosign");
+      DoLargeTest("full_sign");
+      DoLargeTest("full_sign_nosign");
+    }
+
+    method DoLargeTest(config : string)
+    {
+      if !DoTestConfig(config) {
+        return;
+      }
+      expect config in largeEncryptionConfigs;
+      var tconfig := largeEncryptionConfigs[config];
+      var configs := Types.DynamoDbTablesEncryptionConfig (
+        tableEncryptionConfigs := map[TableName := tconfig.config]
+      );
+      // because there are lots of pre-conditions on configs
+      assume {:axiom} false;
+      var client :- expect DynamoDbEncryptionTransforms.DynamoDbEncryptionTransforms(configs);
+      LargeTestsClient(client, config);
+    }
+
+    method LargeTestsClient(client : Trans.IDynamoDbEncryptionTransformsClient, config : string)
+      requires client.ValidState()
+      ensures client.ValidState()
+      modifies client.Modifies
+    {
+      for i := 0 to |large| {
+        RunLargeTest(large[i], client, config);
+      }
+    }
+
+    method RunLargeTest(record : LargeRecord, client : Trans.IDynamoDbEncryptionTransformsClient, config : string)
+      requires client.ValidState()
+      ensures client.ValidState()
+      modifies client.Modifies
+    {
+      if !DoTestRecord(record.name) {
+        return;
+      }
+
+      var time := Time.GetAbsoluteTime();
+      for i := 0 to record.count {
+        var put_input_input := Trans.PutItemInputTransformInput ( sdkInput := DDB.PutItemInput (TableName := TableName, Item := record.item));
+        var put_input_output :- expect client.PutItemInputTransform(put_input_input);
+      }
+      var elapsed := Time.TimeSince(time);
+      Time.PrintTimeLong(elapsed, "Large Encrypt " + record.name + "(" + Base10Int2String(record.count) + ") " + config);
+
+      var put_input_input := Trans.PutItemInputTransformInput ( sdkInput := DDB.PutItemInput (TableName := TableName, Item := record.item));
+      var put_input_output :- expect client.PutItemInputTransform(put_input_input);
+      time := Time.GetAbsoluteTime();
+      for i := 0 to record.count {
+        var orig_get_input := DDB.GetItemInput(TableName := TableName, Key := map[]);
+        var get_output := DDB.GetItemOutput(Item := Some(put_input_output.transformedInput.Item));
+        var trans_get_input := Trans.GetItemOutputTransformInput(sdkOutput := get_output, originalInput := orig_get_input);
+        var put_output :- expect client.GetItemOutputTransform(trans_get_input);
+
+      }
+      elapsed := Time.TimeSince(time);
+      Time.PrintTimeLong(elapsed, "Large Decrypt " + record.name + "(" + Base10Int2String(record.count) + ") " + config);
+    }
+
     method RoundTripTests()
     {
       print "RoundTripTests\n";
@@ -999,7 +1090,7 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
 
   function MakeEmptyTestVector() : TestVectorConfig
   {
-    TestVectorConfig(MakeCreateTableInput(), [], map[], [], map[], map[], [], [], [], [], [], [], [], [], [])
+    TestVectorConfig(MakeCreateTableInput(), [], map[], map[], [], map[], map[], [], [], [], [], [], [], [], [], [], [])
   }
 
   method ParseTestVector(data : JSON, prev : TestVectorConfig, keyVectors: KeyVectors.KeyVectorsClient)
@@ -1020,10 +1111,12 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
     var ioTests : seq<IoTest> := [];
     var gsi : seq<DDB.GlobalSecondaryIndex> := [];
     var tableEncryptionConfigs : map<string, TableConfig> := map[];
+    var largeEncryptionConfigs : map<string, TableConfig> := map[];
     var writeTests : seq<WriteTest> := [];
     var roundTripTests : seq<RoundTripTest> := [];
     var decryptTests : seq<DecryptTest> := [];
     var strings : seq<string> := [];
+    var large : seq<LargeRecord> := [];
 
     for i := 0 to |data.obj| {
       match data.obj[i].0 {
@@ -1038,10 +1131,12 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
         case "IoTests" => ioTests :- GetIoTests(data.obj[i].1, keyVectors);
         case "GSI" => gsi :- GetGSIs(data.obj[i].1);
         case "tableEncryptionConfigs" => tableEncryptionConfigs :- GetTableConfigs(data.obj[i].1, keyVectors);
+        case "largeEncryptionConfigs" => largeEncryptionConfigs :- GetTableConfigs(data.obj[i].1, keyVectors);
         case "WriteTests" => writeTests :- GetWriteTests(data.obj[i].1, keyVectors);
         case "RoundTripTest" => roundTripTests :- GetRoundTripTests(data.obj[i].1, keyVectors);
         case "DecryptTests" => decryptTests :- GetDecryptTests(data.obj[i].1, keyVectors);
         case "Strings" => strings :- GetStrings(data.obj[i].1);
+        case "Large" => large :- GetLarges(data.obj[i].1);
         case _ => return Failure("Unexpected top level tag " + data.obj[i].0);
       }
     }
@@ -1052,6 +1147,7 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
         schemaOnEncrypt := newSchema,
         globalRecords := prev.globalRecords + records,
         tableEncryptionConfigs := prev.tableEncryptionConfigs + tableEncryptionConfigs,
+        largeEncryptionConfigs := prev.largeEncryptionConfigs + largeEncryptionConfigs,
         queries := prev.queries + queries,
         failingQueries := prev.failingQueries + failingQueries,
         names := prev.names + names,
@@ -1063,7 +1159,8 @@ module {:options "-functionSyntax:4"} DdbEncryptionTestVectors {
         writeTests := prev.writeTests + writeTests,
         roundTripTests := prev.roundTripTests + roundTripTests,
         decryptTests := prev.decryptTests + decryptTests,
-        strings := prev.strings + strings
+        strings := prev.strings + strings,
+        large := prev.large + large
       )
     );
   }