feat: allow indirect attribute names with MultiKeyStore

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

@@ -141,6 +141,29 @@ module BeaconTestFixtures {
     return store;
   }
 
+  method GetMultiKeyStore() returns (output : SI.ValidStore)
+    ensures fresh(output.Modifies)
+  {
+    var kmsClient :- expect KMS.KMSClient();
+    var ddbClient :- expect DDBC.DynamoDBClient();
+    var kmsConfig := KTypes.KMSConfiguration.kmsKeyArn(
+      "arn:aws:kms:us-west-2:370957321024:key/9d989aa2-2f9c-438c-a745-cc57d3ad0126"
+    );
+    var keyStoreConfig := KTypes.KeyStoreConfig(
+      id := None,
+      kmsConfiguration := kmsConfig,
+      logicalKeyStoreName := "KeyStoreDdbTable",
+      grantTokens := None,
+      ddbTableName := "KeyStoreDdbTable",
+      ddbClient := Some(ddbClient),
+      kmsClient := Some(kmsClient)
+    );
+
+    var store :- expect KeyStore.KeyStore(keyStoreConfig);
+    return store;
+  }
+
+
   method GetEmptyBeacons() returns (output : BeaconVersion)
     ensures output.keyStore.ValidState()
     ensures fresh(output.keyStore.Modifies)
@@ -177,6 +200,24 @@ module BeaconTestFixtures {
       );
   }
 
+  method GetLotsaBeaconsMulti() returns (output : BeaconVersion)
+    ensures output.keyStore.ValidState()
+    ensures fresh(output.keyStore.Modifies)
+    ensures output.version == 1
+  {
+    var store := GetMultiKeyStore();
+    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 +241,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 +265,26 @@ module BeaconTestFixtures {
     return SI.KeySource(client, version.keyStore, SI.LiteralLoc(keys), cache, 0);
   }
 
+  datatype KeyLocation =
+    | LiteralLoc (keys: HmacKeyMap)
+    | SingleLoc (keyId: string)
+    | MultiLoc (keyName : string, deleteKey : bool)
+
+  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,69 @@ module TestDDBSupport {
     expect newItem == SimpleItem + expectedNew;
   }
 
+  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);
+
+    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);
+
+    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"));
+
+    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);
+  }
 }
+// Map<String, String> expressionAttributesNames = new HashMap<>();
+// expressionAttributesNames.put("#compound", "firstNameCompound");
+
+// Map<String, AttributeValue> expressionAttributeValues = new HashMap<>();
+// expressionAttributeValues.put(":value", AttributeValue.builder().s("f1-l.f2-li.f3-lil.f4-lily.s-Placeholder - 3afba703-6345-4a25-b28a-ec22b1b79a35").build());
+
+// QueryRequest queryRequest = QueryRequest.builder()
+//                 .tableName(tableName)
+//                 .indexName("aws_dbe_b_firstNameCompound-index")
+//                 .keyConditionExpression("#compound = :value")
+//                 .expressionAttributeNames(expressionAttributesNames)
+//                 .expressionAttributeValues(expressionAttributeValues)
+//                 .build();