aws · RitvikKapila · Jan 24, 2025 · Nov 20, 2024 · Nov 20, 2024 · Nov 28, 2024
@@ -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 ,

@@ -716,6 +716,10 @@ 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("Provide the Shared Cache for Searchable Encryption.")
+  cache : CacheType,
+  @documentation("Partition ID to distinguish Beacon Key Sources writing to a cache. If the Partition ID is the same for two Beacon Key Sources, they can share the same cache entries in the cache.")
+  partitionId: String
 }
 
 //= specification/searchable-encryption/search-config.md#multi-key-store-initialization
@@ -734,7 +738,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 cache. If the Partition ID is the same for two Beacon Key Sources, they can share the same cache entries in the cache.")
+  partitionId: String
 }
 
 //= specification/searchable-encryption/search-config.md#beacon-key-source

@@ -33,6 +33,7 @@ module SearchConfigToInfo {
   import SE = AwsCryptographyDbEncryptionSdkStructuredEncryptionTypes
   import MPT = AwsCryptographyMaterialProvidersTypes
   import Primitives = AtomicPrimitives
+  import UUID
 
   // convert configured SearchConfig to internal SearchInfo
   method Convert(outer : DynamoDbTableEncryptionConfig)
@@ -134,6 +135,9 @@ module SearchConfigToInfo {
           config.multi.cache.value
         else
           MPT.Default(Default := MPT.DefaultCache(entryCapacity := 1000))
+      else
+      if config.single.cache.Some? then
+        config.single.cache.value
       else
         MPT.Default(Default := MPT.DefaultCache(entryCapacity := 1));
 
@@ -151,13 +155,41 @@ 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
+          )
+      );
+    }
+    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 {
+      var uuid? := UUID.GenerateUUID();
+
+      var uuid :- uuid?
+      .MapFailure(e => Error.DynamoDbEncryptionException(message := e));
+
+      partitionIdBytes :- UUID.ToByteArray(uuid)
+      .MapFailure(e => Error.DynamoDbEncryptionException(message := 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));
     } 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));
     }
   }
 

@@ -26,6 +26,7 @@ module SearchableEncryptionInfo {
   import MP = AwsCryptographyMaterialProvidersTypes
   import KeyStoreTypes = AwsCryptographyKeyStoreTypes
   import SE = AwsCryptographyDbEncryptionSdkStructuredEncryptionTypes
+  import opened CacheConstants
 
   //= specification/searchable-encryption/search-config.md#version-number
   //= type=implication
@@ -137,7 +138,8 @@ module SearchableEncryptionInfo {
     store : ValidStore,
     keyLoc : KeyLocation,
     cache : MP.ICryptographicMaterialsCache,
-    cacheTTL : uint32
+    cacheTTL : uint32,
+    partitionIdBytes : seq<uint8>
   ) {
     function Modifies() : set<object> {
       client.Modifies + store.Modifies
@@ -153,7 +155,7 @@ module SearchableEncryptionInfo {
     {
       if keyLoc.SingleLoc? {
         :- Need(keyId.DontUseKeyId?, E("KeyID should not be supplied with a SingleKeyStore"));
-        var theMap :- getKeysCache(stdNames, keyLoc.keyId);
+        var theMap :- getKeysCache(stdNames, keyLoc.keyId, cacheTTL as MP.PositiveLong, partitionIdBytes);
         return Success(Keys(theMap));
       } else if keyLoc.LiteralLoc? {
         :- Need(keyId.DontUseKeyId?, E("KeyID should not be supplied with a LiteralKeyStore"));
@@ -163,7 +165,7 @@ module SearchableEncryptionInfo {
         match keyId {
           case DontUseKeyId => return Failure(E("KeyID must not be supplied with a MultiKeyStore"));
           case ShouldHaveKeyId => return Success(ShouldHaveKeys);
-          case KeyId(id) => var theMap :- getKeysCache(stdNames, id); return Success(Keys(theMap));
+          case KeyId(id) => var theMap :- getKeysCache(stdNames, id, cacheTTL as MP.PositiveLong, partitionIdBytes); return Success(Keys(theMap));
         }
       }
     }
@@ -180,9 +182,23 @@ module SearchableEncryptionInfo {
       return Success(keyLoc.keys);
     }
 
+    // Checks if (time_now - cache creation time of the extracted cache entry) is less than the allowed
+    // TTL of the current Beacon Key Source calling the getEntry method from the cache.
+    // Mitigates risk if another Beacon Key Source wrote the entry with a longer TTL.
+    predicate method cacheEntryWithinLimits(
+      creationTime: MP.PositiveLong,
+      now: MP.PositiveLong,
+      ttlSeconds: MP.PositiveLong
+    ): (output: bool)
+    {
+      now - creationTime <= ttlSeconds as MP.PositiveLong
+    }
+
     method getKeysCache(
       stdNames : seq<string>,
-      keyId : string
+      keyId : string,
+      cacheTTL : MP.PositiveLong,
+      partitionIdBytes : seq<uint8>
     )
       returns (output : Result<HmacKeyMap, Error>)
       requires Seq.HasNoDuplicates(stdNames)
@@ -241,8 +257,21 @@ module SearchableEncryptionInfo {
 
                    )
     {
+
+      // Resource ID: Searchable Encryption [0x02]
+      var resourceId : seq<uint8> := RESOURCE_ID_HIERARCHICAL_KEYRING;
+
+      // Scope ID: Searchable Encryption [0x03]
+      var scopeId : seq<uint8> := SCOPE_ID_SEARCHABLE_ENCRYPTION;
+
+      // Create the Suffix
       var keyIdBytesR := UTF8.Encode(keyId);
       var keyIdBytes :- keyIdBytesR.MapFailure(e => E(e));
+      var suffix : seq<uint8> := keyIdBytes;
+
+      // Append Resource Id, Scope Id, Partition Id, and Suffix to create the cache identifier
+      var identifier := resourceId + NULL_BYTE + scopeId + NULL_BYTE + partitionIdBytes + NULL_BYTE + suffix;
+
       var getCacheInput := MP.GetCacheEntryInput(identifier := keyIdBytes, bytesUsed := None);
       verifyValidStateCache(cache);
       assume {:axiom} cache.Modifies == {};
@@ -253,7 +282,21 @@ module SearchableEncryptionInfo {
         return Failure(AwsCryptographyMaterialProviders(AwsCryptographyMaterialProviders:=getCacheOutput.error));
       }
 
-      if getCacheOutput.Failure? {
+      var now := Time.GetCurrent();
+
+      // //= specification/searchable-encryption/search-config.md#<heading>
+      //# If using a `Shared` cache across multiple Beacon Key Sources,
+      //# different Key Sources having the same `beaconKey` can have different TTLs.
+      //# In such a case, the expiry time in the cache is set according to the Beacon Key Source that populated the cache.
+      //# There MUST be a check (cacheEntryWithinLimits) to make sure that for the cache entry found, who's TTL has NOT expired,
+      //# `time.now() - cacheEntryCreationTime <= ttlSeconds` is true and
+      //# valid for TTL of the Beacon Key Source getting the cache entry.
+      //# If this is NOT true, then we MUST treat the cache entry as expired.
+      if getCacheOutput.Failure? || !cacheEntryWithinLimits(
+           creationTime := getCacheOutput.value.creationTime,
+           now := now,
+           ttlSeconds := cacheTTL
+         ) {
         //= specification/searchable-encryption/search-config.md#beacon-keys
         //# Beacon keys MUST be obtained from the configured [Beacon Key Source](#beacon-key-source).
         var maybeRawBeaconKeyMaterials := store.GetBeaconKey(

@@ -26,7 +26,7 @@ module DynamoDbEncryptionUtil {
   //
   // Counterintuitively, DontUseKeys and DontUseKeyId are very different things.
   // DontUseKeyId is the usual thing for single tenant, meaning to use the pre-configured
-  // KeyId, rather than fnd a new one from somewhere.
+  // KeyId, rather than find a new one from somewhere.
   // DontUseKeys means to not hash the values into beacons,
   // but to leave them plaintext, which is necessary for the post-query filtering.
   datatype MaybeKeyMap =

@@ -242,7 +242,9 @@ module BeaconTestFixtures {
       cache := MPT.Default(Default := MPT.DefaultCache(entryCapacity := 3))
     );
     var cache :- expect mpl.CreateCryptographicMaterialsCache(input);
-    return SI.KeySource(client, version.keyStore, SI.LiteralLoc(keys), cache, 0);
+    // Create a test partitionIdBytes
+    var partitionIdBytes : seq<uint8> := [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
+    return SI.KeySource(client, version.keyStore, SI.LiteralLoc(keys), cache, 0, partitionIdBytes);
   }
 
   method GetMultiSource(keyName : string, version : BeaconVersion) returns (output : SI.KeySource)
@@ -257,7 +259,9 @@ module BeaconTestFixtures {
       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);
+    // Create a test partitionIdBytes
+    var partitionIdBytes : seq<uint8> := [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
+    return SI.KeySource(client, version.keyStore, SI.MultiLoc(keyName, false), cache, 0, partitionIdBytes);
   }
 
   const SimpleItem : DDB.AttributeMap := map[

@@ -764,10 +764,10 @@ abstract module AbstractAwsCryptographyDbEncryptionSdkDynamoDbTransformsService
                               tmp5.search.Some? ==>
                                 var tmps6 := set t6 | t6 in tmp5.search.value.versions;
                                 forall tmp6 :: tmp6 in tmps6 ==>
-                                                 tmp6.keySource.multi? ==>
-                                                   tmp6.keySource.multi.cache.Some? ==>
-                                                     tmp6.keySource.multi.cache.value.Shared? ==>
-                                                       tmp6.keySource.multi.cache.value.Shared.ValidState()
+                                                 tmp6.keySource.single? ==>
+                                                   tmp6.keySource.single.cache.Some? ==>
+                                                     tmp6.keySource.single.cache.value.Shared? ==>
+                                                       tmp6.keySource.single.cache.value.Shared.ValidState()
     modifies set tmps7 <- set t7 <- config.tableEncryptionConfigs.Values | true
                                                                            && t7.keyring.Some?
                             :: t7.keyring.value,
@@ -788,10 +788,10 @@ abstract module AbstractAwsCryptographyDbEncryptionSdkDynamoDbTransformsService
     modifies set tmps12 <- set t12 <- config.tableEncryptionConfigs.Values | true
                                                                              && t12.search.Some?
                              , t13 <- t12.search.value.versions | true
-                                                                  && t13.keySource.multi?
-                                                                  && t13.keySource.multi.cache.Some?
-                                                                  && t13.keySource.multi.cache.value.Shared?
-                             :: t13.keySource.multi.cache.value.Shared,
+                                                                  && t13.keySource.single?
+                                                                  && t13.keySource.single.cache.Some?
+                                                                  && t13.keySource.single.cache.value.Shared?
+                             :: t13.keySource.single.cache.value.Shared,
                obj <- tmps12.Modifies | obj in tmps12.Modifies :: obj
     ensures res.Success? ==>
               && fresh(res.value)
@@ -816,10 +816,10 @@ abstract module AbstractAwsCryptographyDbEncryptionSdkDynamoDbTransformsService
                        ) - ( set tmps19 <- set t19 <- config.tableEncryptionConfigs.Values | true
                                                                                              && t19.search.Some?
                                              , t20 <- t19.search.value.versions | true
-                                                                                  && t20.keySource.multi?
-                                                                                  && t20.keySource.multi.cache.Some?
-                                                                                  && t20.keySource.multi.cache.value.Shared?
-                                             :: t20.keySource.multi.cache.value.Shared,
+                                                                                  && t20.keySource.single?
+                                                                                  && t20.keySource.single.cache.Some?
+                                                                                  && t20.keySource.single.cache.value.Shared?
+                                             :: t20.keySource.single.cache.value.Shared,
                                obj <- tmps19.Modifies | obj in tmps19.Modifies :: obj
                        ) )
               && fresh(res.value.History)
@@ -847,10 +847,10 @@ abstract module AbstractAwsCryptographyDbEncryptionSdkDynamoDbTransformsService
                               tmp26.search.Some? ==>
                                 var tmps27 := set t27 | t27 in tmp26.search.value.versions;
                                 forall tmp27 :: tmp27 in tmps27 ==>
-                                                  tmp27.keySource.multi? ==>
-                                                    tmp27.keySource.multi.cache.Some? ==>
-                                                      tmp27.keySource.multi.cache.value.Shared? ==>
-                                                        tmp27.keySource.multi.cache.value.Shared.ValidState()
+                                                  tmp27.keySource.single? ==>
+                                                    tmp27.keySource.single.cache.Some? ==>
+                                                      tmp27.keySource.single.cache.value.Shared? ==>
+                                                        tmp27.keySource.single.cache.value.Shared.ValidState()
 
   // Helper functions for the benefit of native code to create a Success(client) without referring to Dafny internals
   function method CreateSuccessOfClient(client: IDynamoDbEncryptionTransformsClient): Result<IDynamoDbEncryptionTransformsClient, Error> {

@@ -150,7 +150,7 @@ module
       var tableName: string := tableNamesSeq[i];
 
       var inputConfig := config.tableEncryptionConfigs[tableName];
-      :- Need(inputConfig.logicalTableName !in allLogicalTableNames,  E("Duplicate logical table maped to multipule physical tables: " + inputConfig.logicalTableName));
+      :- Need(inputConfig.logicalTableName !in allLogicalTableNames,  E("Duplicate logical table mapped to multiple physical tables: " + inputConfig.logicalTableName));
 
       assert SearchConfigToInfo.ValidSearchConfig(inputConfig.search);
       SearchInModifies(config, tableName);

@@ -668,7 +668,19 @@ public static MultiKeyStore MultiKeyStore(
           )
         )
         : Option.create_None(CacheType._typeDescriptor());
-    return new MultiKeyStore(keyFieldName, cacheTTL, cache);
+    Option<DafnySequence<? extends Character>> partitionId;
+    partitionId =
+      Objects.nonNull(nativeValue.partitionId())
+        ? Option.create_Some(
+          DafnySequence._typeDescriptor(TypeDescriptor.CHAR),
+          software.amazon.smithy.dafny.conversion.ToDafny.Simple.CharacterSequence(
+            nativeValue.partitionId()
+          )
+        )
+        : Option.create_None(
+          DafnySequence._typeDescriptor(TypeDescriptor.CHAR)
+        );
+    return new MultiKeyStore(keyFieldName, cacheTTL, cache, partitionId);
   }
 
   public static PartOnly PartOnly(
@@ -747,7 +759,29 @@ public static SingleKeyStore SingleKeyStore(
       );
     Integer cacheTTL;
     cacheTTL = (nativeValue.cacheTTL());
-    return new SingleKeyStore(keyId, cacheTTL);
+    Option<CacheType> cache;
+    cache =
+      Objects.nonNull(nativeValue.cache())
+        ? Option.create_Some(
+          CacheType._typeDescriptor(),
+          software.amazon.cryptography.materialproviders.ToDafny.CacheType(
+            nativeValue.cache()
+          )
+        )
+        : Option.create_None(CacheType._typeDescriptor());
+    Option<DafnySequence<? extends Character>> partitionId;
+    partitionId =
+      Objects.nonNull(nativeValue.partitionId())
+        ? Option.create_Some(
+          DafnySequence._typeDescriptor(TypeDescriptor.CHAR),
+          software.amazon.smithy.dafny.conversion.ToDafny.Simple.CharacterSequence(
+            nativeValue.partitionId()
+          )
+        )
+        : Option.create_None(
+          DafnySequence._typeDescriptor(TypeDescriptor.CHAR)
+        );
+    return new SingleKeyStore(keyId, cacheTTL, cache, partitionId);
   }
 
   public static StandardBeacon StandardBeacon(