feat(SharedCache): Shared Cache for Searchable Encryption (#1476)

Author: RitvikKapila

DynamoDbEncryption/dafny/DynamoDbEncryption/Model/AwsCryptographyDbEncryptionSdkDynamoDbTypes.dfy

@@ -354,7 +354,8 @@ module {:extern "software.amazon.cryptography.dbencryptionsdk.dynamodb.internald
   datatype MultiKeyStore = | MultiKeyStore (
     nameonly keyFieldName: string ,
     nameonly cacheTTL: int32 ,
-    nameonly cache: Option<AwsCryptographyMaterialProvidersTypes.CacheType> := Option.None
+    nameonly cache: Option<AwsCryptographyMaterialProvidersTypes.CacheType> := Option.None ,
+    nameonly partitionId: Option<string> := Option.None
   )
   datatype PartOnly = | PartOnly (
 
@@ -388,7 +389,9 @@ module {:extern "software.amazon.cryptography.dbencryptionsdk.dynamodb.internald
   }
   datatype SingleKeyStore = | SingleKeyStore (
     nameonly keyId: string ,
-    nameonly cacheTTL: int32
+    nameonly cacheTTL: int32 ,
+    nameonly cache: Option<AwsCryptographyMaterialProvidersTypes.CacheType> := Option.None ,
+    nameonly partitionId: Option<string> := Option.None
   )
   datatype StandardBeacon = | StandardBeacon (
     nameonly name: string ,

DynamoDbEncryption/dafny/DynamoDbEncryption/Model/DynamoDbEncryption.smithy

@@ -708,6 +708,8 @@ structure KeyStoreReference {}
 //# On initialization of a Single Key Store, the caller MUST provide:
 //#  - [Beacon Key Id](#beacon-key-id)
 //#  - [cacheTTL](#cachettl)
+//#  - [cache](#key-store-cache)
+//#  - [partition-id](#partition-id)
 
 @javadoc("The configuration for using a single Beacon Key.")
 structure SingleKeyStore {
@@ -717,14 +719,19 @@ structure SingleKeyStore {
   @required
   @javadoc("How long (in seconds) the beacon key material is cached locally before it is re-retrieved from DynamoDB and re-authed with AWS KMS.")
   cacheTTL: Integer,
+  @documentation("Which type of local cache to use. Please see the [spec](https://github.com/aws/aws-database-encryption-sdk-dynamodb/blob/main/specification/searchable-encryption/search-config.md#key-store-cache) on how to provide a cache for a SingleKeyStore.")
+  cache : CacheType,
+  @documentation("Partition ID to distinguish Beacon Key Sources writing to a Shared cache. If the Partition ID is the same for two Beacon Key Sources, they can share the same cache entries in the Shared cache.")
+  partitionId: String
 }
 
 //= specification/searchable-encryption/search-config.md#multi-key-store-initialization
 //= type=implication
 //# On initialization of a Multi Key Store, the caller MUST provide:
 //#  - [Beacon Key Field Name](#beacon-key-field-name)
 //#  - [cacheTTL](#cachettl)
-//#  - [max cache size](#max-cache-size)
+//#  - [cache](#key-store-cache)
+//#  - [partition-id](#partition-id)
 
 @javadoc("The configuration for using multiple Beacon Keys.")
 structure MultiKeyStore {
@@ -735,7 +742,9 @@ structure MultiKeyStore {
   @javadoc("How long (in seconds) the beacon key material is cached locally before it is re-retrieved from DynamoDB and re-authed with AWS KMS.")
   cacheTTL: Integer,
   @javadoc("Which type of local cache to use.")
-  cache : CacheType
+  cache : CacheType,
+  @documentation("Partition ID to distinguish Beacon Key Sources writing to a Shared cache. If the Partition ID is the same for two Beacon Key Sources, they can share the same cache entries in the Shared cache.")
+  partitionId: String
 }
 
 //= specification/searchable-encryption/search-config.md#beacon-key-source

DynamoDbEncryption/dafny/DynamoDbEncryption/src/ConfigToInfo.dfy

@@ -38,6 +38,9 @@ module SearchConfigToInfo {
   method Convert(outer : DynamoDbTableEncryptionConfig)
     returns (output : Result<Option<I.ValidSearchInfo>, Error>)
     requires ValidSearchConfig(outer.search)
+    requires outer.search.Some? ==> ValidSharedCache(outer.search.value.versions[0].keySource)
+    modifies if outer.search.Some? then outer.search.value.versions[0].keyStore.Modifies else {}
+    ensures outer.search.Some? ==> ValidSharedCache(outer.search.value.versions[0].keySource)
     ensures output.Success? && output.value.Some? ==>
               && output.value.value.ValidState()
               && fresh(output.value.value.versions[0].keySource.client)
@@ -56,7 +59,8 @@ module SearchConfigToInfo {
     } else {
       :- Need(outer.search.value.writeVersion == 1, E("writeVersion must be '1'."));
       :- Need(|outer.search.value.versions| == 1, E("search config must be have exactly one version."));
-      var version :- ConvertVersion(outer, outer.search.value.versions[0]);
+      var beaconVersionConfig := outer.search.value.versions[0];
+      var version :- ConvertVersion(outer, beaconVersionConfig);
       var info := I.MakeSearchInfo(version);
       return Success(Some(info));
     }
@@ -74,6 +78,19 @@ module SearchConfigToInfo {
       forall b <- config.value.versions :: ValidBeaconVersion(b)
   }
 
+  // Valid state of the provided shared cache, if it exists
+  predicate {:opaque} ValidSharedCache(config: BeaconKeySource)
+  {
+    && (&& config.single?
+        && config.single.cache.Some?
+        && config.single.cache.value.Shared?
+        ==> && config.single.cache.value.Shared.ValidState())
+    && (&& config.multi?
+        && config.multi.cache.Some?
+        && config.multi.cache.value.Shared?
+        ==> && config.multi.cache.value.Shared.ValidState())
+  }
+
   // return true if, `keyFieldName` should be deleted from an item before writing
   function method ShouldDeleteKeyField(outer : DynamoDbTableEncryptionConfig, keyFieldName : string)
     : (ret : Result<bool, Error>)
@@ -101,7 +118,10 @@ module SearchConfigToInfo {
   )
     returns (output : Result<I.KeySource, Error>)
     modifies client.Modifies
+    modifies keyStore.Modifies
     requires client.ValidState()
+    requires ValidSharedCache(config)
+    ensures ValidSharedCache(config)
     ensures client.ValidState()
     ensures output.Success? ==>
               && output.value.ValidState()
@@ -119,18 +139,6 @@ module SearchConfigToInfo {
       && config.multi.keyFieldName in outer.attributeActionsOnEncrypt
       && outer.attributeActionsOnEncrypt[config.multi.keyFieldName] == SE.ENCRYPT_AND_SIGN
       ==> output.Failure?
-    // Not in Spec, but for now, SE does not support the Shared Cache Type
-    ensures
-      && config.multi?
-      && config.multi.cache.Some?
-      && config.multi.cache.value.Shared?
-      ==>
-        && output.Failure?
-           // If the failure was NOT caused by booting up the MPL
-        && !output.error.AwsCryptographyMaterialProviders?
-        ==>
-          && output.error.DynamoDbEncryptionException?
-          && output.error.message == "Searchable Encryption does not support the Shared Cache type at this time."
   {
     // TODO-FutureCleanUp : https://github.com/aws/aws-database-encryption-sdk-dynamodb/issues/1510
     // It is not-good that the MPL is initialized here;
@@ -140,23 +148,35 @@ module SearchConfigToInfo {
     var mpl :- mplR.MapFailure(e => AwsCryptographyMaterialProviders(e));
 
     //= specification/searchable-encryption/search-config.md#key-store-cache
-    //# For a [Single Key Store](#single-key-store-initialization) the [Entry Capacity](../../submodules/MaterialProviders/aws-encryption-sdk-specification/framework/cryptographic-materials-cache.md#entry-capacity)
-    //# MUST be 1
-    //# For a [Multi Key Store](#multi-key-store-initialization) the [Entry Capacity](../../submodules/MaterialProviders/aws-encryption-sdk-specification/framework/cryptographic-materials-cache.md#entry-capacity)
-    //# MUST be key store's max cache size.
+    //# For a Beacon Key Source a [CMC](../../submodules/MaterialProviders/aws-encryption-sdk-specification/framework/cryptographic-materials-cache.md)
+    //# MUST be created.
+    //# For a [Single Key Store](#single-key-store-initialization), either the user provides a cache, or we create a cache that has [Entry Capacity](../../submodules/MaterialProviders/aws-encryption-sdk-specification/framework/cryptographic-materials-cache.md#entry-capacity)
+    //# equal to 1.
+    //# For a [Multi Key Store](#multi-key-store-initialization), either the user provides a cache, or we create a cache that has [Entry Capacity](../../submodules/MaterialProviders/aws-encryption-sdk-specification/framework/cryptographic-materials-cache.md#entry-capacity)
+    //# equal to 1000.
     var cacheType : MPT.CacheType :=
       if config.multi? then
         if config.multi.cache.Some? then
           config.multi.cache.value
         else
           MPT.Default(Default := MPT.DefaultCache(entryCapacity := 1000))
+      else
+      if config.single.cache.Some? then
+        // If the user provides a CacheType, and it is NOT Shared,
+        // we SHOULD only allow an entryCapacity = 1
+        // because the SingleKeyStore only ever caches one value.
+        // That is, we SHOULD add a check here for entryCapacity = 1.
+        // However, that requires us to write an if block for each CacheType.
+        // Also, it does NOT matter what the entryCapacity is, because the cache
+        // can only hold one element at a time.
+        config.single.cache.value
       else
         MPT.Default(Default := MPT.DefaultCache(entryCapacity := 1));
 
     var cache;
     if cacheType.Shared? {
-      return Failure(DynamoDbEncryptionException(message:="Searchable Encryption does not support the Shared Cache type at this time."));
-      // cache := cacheType.Shared;
+      cache := cacheType.Shared;
+      reveal ValidSharedCache(config);
     } else {
       //= specification/searchable-encryption/search-config.md#key-store-cache
       //# For a Beacon Key Source a [CMC](../../submodules/MaterialProviders/aws-encryption-sdk-specification/framework/cryptographic-materials-cache.md)
@@ -168,20 +188,59 @@ module SearchConfigToInfo {
       cache :- maybeCache.MapFailure(e => AwsCryptographyMaterialProviders(e));
     }
 
+    var partitionIdBytes : seq<uint8>;
+
+    if config.multi? && config.multi.partitionId.Some? {
+      partitionIdBytes :- UTF8.Encode(config.multi.partitionId.value)
+      .MapFailure(
+        e => Error.DynamoDbEncryptionException(
+            message := "Could not UTF-8 Encode Partition ID from MultiKeyStore: " + e
+          )
+      );
+    }
+    else if config.single? && config.single.partitionId.Some? {
+      partitionIdBytes :- UTF8.Encode(config.single.partitionId.value)
+      .MapFailure(
+        e => Error.DynamoDbEncryptionException(
+            message := "Could not UTF-8 Encode Partition ID from SingleKeyStore: " + e
+          )
+      );
+    }
+    else {
+      partitionIdBytes :- I.GenerateUuidBytes();
+    }
+    var getKeyStoreInfoOutput? := keyStore.GetKeyStoreInfo();
+    var getKeyStoreInfoOutput :- getKeyStoreInfoOutput?.MapFailure(e => Error.AwsCryptographyKeyStore(e));
+    var logicalKeyStoreName : string := getKeyStoreInfoOutput.logicalKeyStoreName;
+    var logicalKeyStoreNameBytes : seq<uint8> :- UTF8.Encode(logicalKeyStoreName)
+    .MapFailure(
+      e => Error.DynamoDbEncryptionException(
+          message := "Could not UTF-8 Encode Logical Key Store Name: " + e
+        )
+    );
+
     if config.multi? {
       :- Need(0 < config.multi.cacheTTL, E("Beacon Cache TTL must be at least 1."));
       var deleteKey :- ShouldDeleteKeyField(outer, config.multi.keyFieldName);
-      output := Success(I.KeySource(client, keyStore, I.MultiLoc(config.multi.keyFieldName, deleteKey), cache, config.multi.cacheTTL as uint32));
+      output := Success(I.KeySource(client, keyStore, I.MultiLoc(config.multi.keyFieldName, deleteKey), cache, config.multi.cacheTTL as uint32, partitionIdBytes, logicalKeyStoreNameBytes));
     } else {
       :- Need(0 < config.single.cacheTTL, E("Beacon Cache TTL must be at least 1."));
-      output := Success(I.KeySource(client, keyStore, I.SingleLoc(config.single.keyId), cache, config.single.cacheTTL as uint32));
+      output := Success(I.KeySource(client, keyStore, I.SingleLoc(config.single.keyId), cache, config.single.cacheTTL as uint32, partitionIdBytes, logicalKeyStoreNameBytes));
+    }
+    assert output.value.ValidState() by {
+      // This axiom is important because it is not easy to prove
+      // keyStore.Modifies !! cache.Modifies for a shared cache.
+      assume {:axiom} keyStore.Modifies !! cache.Modifies;
     }
   }
 
   // convert configured BeaconVersion to internal BeaconVersion
   method ConvertVersion(outer : DynamoDbTableEncryptionConfig, config : BeaconVersion)
     returns (output : Result<I.ValidBeaconVersion, Error>)
     requires ValidBeaconVersion(config)
+    requires ValidSharedCache(config.keySource)
+    modifies config.keyStore.Modifies
+    ensures ValidSharedCache(config.keySource)
     ensures output.Success? ==>
               && output.value.ValidState()
               && fresh(output.value.keySource.client)