feat: allow indirect attribute names with MultiKeyStore (#1208)

* feat: allow indirect attribute names with MultiKeyStore

Author: ajewellamz

DynamoDbEncryption/dafny/DynamoDbEncryption/src/FilterExpr.dfy

@@ -1533,7 +1533,7 @@ module DynamoDBFilterExpr {
     else
       var name := if names.Some? && attr.value.s in names.value then names.value[attr.value.s] else attr.value.s;
       var keyIdField := bv.keySource.keyLoc.keyName;
-      if keyIdField == attr.value.s then
+      if keyIdField == name then
         Some(value)
       else
         KeyIdFromPart(bv, keyIdField, attr.value.s, value)

DynamoDbEncryption/dafny/DynamoDbEncryption/test/BeaconTestFixtures.dfy

@@ -130,9 +130,9 @@ module BeaconTestFixtures {
     var keyStoreConfig := KTypes.KeyStoreConfig(
       id := None,
       kmsConfiguration := kmsConfig,
-      logicalKeyStoreName := "foo",
+      logicalKeyStoreName := "KeyStoreDdbTable",
       grantTokens := None,
-      ddbTableName := "foo",
+      ddbTableName := "KeyStoreDdbTable",
       ddbClient := Some(ddbClient),
       kmsClient := Some(kmsClient)
     );
@@ -177,6 +177,24 @@ module BeaconTestFixtures {
       );
   }
 
+  method GetLotsaBeaconsMulti() returns (output : BeaconVersion)
+    ensures output.keyStore.ValidState()
+    ensures fresh(output.keyStore.Modifies)
+    ensures output.version == 1
+  {
+    var store := GetKeyStore();
+    return BeaconVersion (
+        version := 1,
+        keyStore := store,
+        keySource := multi(MultiKeyStore(keyFieldName := "TheKeyField", cacheTTL := 42)),
+        standardBeacons := [std2, std4, std6, NameTitleBeacon, NameB, TitleB],
+        compoundBeacons := Some([NameTitle, YearName, Mixed, JustSigned]),
+        virtualFields := Some([NameTitleField]),
+        encryptedParts := None,
+        signedParts := None
+      );
+  }
+
   const EmptyTableConfig := DynamoDbTableEncryptionConfig (
                               logicalTableName := "Foo",
                               partitionKeyName := "foo",
@@ -200,7 +218,8 @@ module BeaconTestFixtures {
                              "Title" := SE.ENCRYPT_AND_SIGN,
                              "TooBad" := SE.ENCRYPT_AND_SIGN,
                              "Year" := SE.SIGN_ONLY,
-                             "Date" := SE.SIGN_ONLY
+                             "Date" := SE.SIGN_ONLY,
+                             "TheKeyField" := SE.SIGN_ONLY
                            ]
                            )
 
@@ -223,6 +242,21 @@ module BeaconTestFixtures {
     return SI.KeySource(client, version.keyStore, SI.LiteralLoc(keys), cache, 0);
   }
 
+  method GetMultiSource(keyName : string, version : BeaconVersion) returns (output : SI.KeySource)
+    requires version.keyStore.ValidState()
+    ensures output.ValidState()
+    ensures version.keyStore == output.store
+    ensures fresh(output.client.Modifies)
+  {
+    var client :- expect Primitives.AtomicPrimitives();
+    var mpl :- expect MaterialProviders.MaterialProviders();
+    var input := MPT.CreateCryptographicMaterialsCacheInput(
+      cache := MPT.Default(Default := MPT.DefaultCache(entryCapacity := 3))
+    );
+    var cache :- expect mpl.CreateCryptographicMaterialsCache(input);
+    return SI.KeySource(client, version.keyStore, SI.MultiLoc(keyName, false), cache, 0);
+  }
+
   const SimpleItem : DDB.AttributeMap := map[
                                            "std2" := Std2String,
                                            "std4" := Std4String,

DynamoDbEncryption/dafny/DynamoDbEncryption/test/DDBSupport.dfy

@@ -32,4 +32,60 @@ module TestDDBSupport {
     expect newItem == SimpleItem + expectedNew;
   }
 
+  // DynamoDB String :: cast string to DDB.AttributeValue.S
+  function method DS(x : string) : DDB.AttributeValue
+  {
+    DDB.AttributeValue.S(x)
+  }
+
+  method {:test} TestMulti() {
+    var version := GetLotsaBeaconsMulti();
+    var src := GetMultiSource("TheKeyField", version);
+    var bv :- expect ConvertVersionWithSource(FullTableConfig, version, src);
+    var search := SI.SearchInfo([bv], 0);
+    var expressionAttributeValues : map<string, AttributeValue> := map[
+      ":value" := DS("0ad21413-51aa-42e1-9c3d-6a4b1edf7e10")
+    ];
+    var queryInput := DDB.QueryInput (
+      TableName := "SomeTable",
+      ExpressionAttributeValues := Some(expressionAttributeValues),
+      KeyConditionExpression := Some("TheKeyField = :value")
+    );
+    var result :- expect QueryInputForBeacons(Some(search), FullTableConfig.attributeActionsOnEncrypt, queryInput);
+
+    // Verify Success with branch key id plus beacon
+    expressionAttributeValues := map[
+      ":value" := DS("0ad21413-51aa-42e1-9c3d-6a4b1edf7e10"),
+      ":other" := DS("junk")
+    ];
+    queryInput := DDB.QueryInput (
+      TableName := "foo",
+      ExpressionAttributeValues := Some(expressionAttributeValues),
+      KeyConditionExpression := Some("TheKeyField = :value AND std2 = :other")
+    );
+    result :- expect QueryInputForBeacons(Some(search), FullTableConfig.attributeActionsOnEncrypt, queryInput);
+
+    // Verify Failure with beacon but no branch key id
+    queryInput := DDB.QueryInput (
+      TableName := "foo",
+      ExpressionAttributeValues := Some(expressionAttributeValues),
+      KeyConditionExpression := Some("std2 = :other")
+    );
+    var result2 := QueryInputForBeacons(Some(search), FullTableConfig.attributeActionsOnEncrypt, queryInput);
+    expect result2 == Failure(AwsCryptographyDbEncryptionSdkDynamoDbTypes.Error.DynamoDbEncryptionException(
+                                message := "Need KeyId because of beacon std2 but no KeyId found in query"));
+
+    // Verify Success, even when field names are indirect via ExpressionAttributeNames
+    var expressionAttributeNames := map[
+      "#beacon" := "std2",
+      "#keyfield" := "TheKeyField"
+    ];
+    queryInput := DDB.QueryInput (
+      TableName := "foo",
+      ExpressionAttributeNames := Some(expressionAttributeNames),
+      ExpressionAttributeValues := Some(expressionAttributeValues),
+      KeyConditionExpression := Some("#keyfield = :value AND #beacon = :other")
+    );
+    result :- expect QueryInputForBeacons(Some(search), FullTableConfig.attributeActionsOnEncrypt, queryInput);
+  }
 }