add partition ID and cache

DynamoDbEncryption/dafny/DynamoDbEncryption/Model/AwsCryptographyDbEncryptionSdkDynamoDbTypes.dfy

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

DynamoDbEncryption/dafny/DynamoDbEncryption/Model/DynamoDbEncryption.smithy

@@ -713,11 +713,12 @@ structure SingleKeyStore {
   @required
   @javadoc("The Beacon Key ID.")
   keyId : String,
-  @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. Provide only one of cacheTTL or cache.")
+  @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. Provide only one of cacheTTL or cache.")
+  @documentation("Provide the Shared Cache for Searchable Encryption.")
   cache : CacheType,
-  @documentation("Partition ID to share DynamoDB Interceptors. TODO: Update description")
+  @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
 }
 
@@ -738,7 +739,7 @@ structure MultiKeyStore {
   cacheTTL: Integer,
   @javadoc("Which type of local cache to use.")
   cache : CacheType,
-  @documentation("Partition ID to share DynamoDB Interceptors. TODO: Update description")
+  @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
 }
 

DynamoDbEncryption/dafny/DynamoDbEncryption/src/ConfigToInfo.dfy

@@ -136,9 +136,11 @@ module SearchConfigToInfo {
         else
           MPT.Default(Default := MPT.DefaultCache(entryCapacity := 1000))
       else
-        MPT.Default(Default := MPT.DefaultCache(entryCapacity := 1));
+        if config.single.cache.Some? then
+          config.single.cache.value
+        else
+          MPT.Default(Default := MPT.DefaultCache(entryCapacity := 1));
 
-    // TODO : Add check that customers only provide either cacheTTL or cache in case of SingleKeyStore
     var cache;
     if cacheType.Shared? {
       cache := cacheType.Shared;
@@ -155,15 +157,21 @@ module SearchConfigToInfo {
 
     var partitionIdBytes : seq<uint8>;
 
-    if outer.keyring.Some? {
-      if outer.keyring.value.partitionId.Some? {
-        partitionIdBytes :- UTF8.Encode(outer.keyring.value.partitionId.value)
-        .MapFailure(
-          e => Error.DynamoDbEncryptionException(
-              message := "Could not UTF-8 Encode Partition ID: " + e
-            )
-        );
-      }
+    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();

DynamoDbEncryption/dafny/DynamoDbEncryption/src/SearchInfo.dfy

@@ -155,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, partitionIdBytes);
+        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"));
@@ -165,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, partitionIdBytes); return Success(Keys(theMap));
+          case KeyId(id) => var theMap :- getKeysCache(stdNames, id, cacheTTL as MP.PositiveLong, partitionIdBytes); return Success(Keys(theMap));
         }
       }
     }
@@ -182,9 +182,22 @@ 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,
+      cacheTTL : MP.PositiveLong,
       partitionIdBytes : seq<uint8>
     )
       returns (output : Result<HmacKeyMap, Error>)
@@ -269,8 +282,21 @@ module SearchableEncryptionInfo {
         return Failure(AwsCryptographyMaterialProviders(AwsCryptographyMaterialProviders:=getCacheOutput.error));
       }
 
-      // TODO: Add cacheEntryWithinLimits
-      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(

DynamoDbEncryption/runtimes/java/src/main/smithy-generated/software/amazon/cryptography/dbencryptionsdk/dynamodb/ToDafny.java

@@ -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(

DynamoDbEncryption/runtimes/java/src/main/smithy-generated/software/amazon/cryptography/dbencryptionsdk/dynamodb/ToNative.java

@@ -574,6 +574,13 @@ public static MultiKeyStore MultiKeyStore(
         )
       );
     }
+    if (dafnyValue.dtor_partitionId().is_Some()) {
+      nativeBuilder.partitionId(
+        software.amazon.smithy.dafny.conversion.ToNative.Simple.String(
+          dafnyValue.dtor_partitionId().dtor_value()
+        )
+      );
+    }
     return nativeBuilder.build();
   }
 
@@ -653,6 +660,20 @@ public static SingleKeyStore SingleKeyStore(
       )
     );
     nativeBuilder.cacheTTL((dafnyValue.dtor_cacheTTL()));
+    if (dafnyValue.dtor_cache().is_Some()) {
+      nativeBuilder.cache(
+        software.amazon.cryptography.materialproviders.ToNative.CacheType(
+          dafnyValue.dtor_cache().dtor_value()
+        )
+      );
+    }
+    if (dafnyValue.dtor_partitionId().is_Some()) {
+      nativeBuilder.partitionId(
+        software.amazon.smithy.dafny.conversion.ToNative.Simple.String(
+          dafnyValue.dtor_partitionId().dtor_value()
+        )
+      );
+    }
     return nativeBuilder.build();
   }
 

DynamoDbEncryption/runtimes/java/src/main/smithy-generated/software/amazon/cryptography/dbencryptionsdk/dynamodb/model/MultiKeyStore.java

@@ -26,10 +26,16 @@ public class MultiKeyStore {
    */
   private final CacheType cache;
 
+  /**
+   * 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.
+   */
+  private final String partitionId;
+
   protected MultiKeyStore(BuilderImpl builder) {
     this.keyFieldName = builder.keyFieldName();
     this.cacheTTL = builder.cacheTTL();
     this.cache = builder.cache();
+    this.partitionId = builder.partitionId();
   }
 
   /**
@@ -53,6 +59,13 @@ public CacheType cache() {
     return this.cache;
   }
 
+  /**
+   * @return 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.
+   */
+  public String partitionId() {
+    return this.partitionId;
+  }
+
   public Builder toBuilder() {
     return new BuilderImpl(this);
   }
@@ -92,6 +105,16 @@ public interface Builder {
      */
     CacheType cache();
 
+    /**
+     * @param partitionId 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.
+     */
+    Builder partitionId(String partitionId);
+
+    /**
+     * @return 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.
+     */
+    String partitionId();
+
     MultiKeyStore build();
   }
 
@@ -103,12 +126,15 @@ static class BuilderImpl implements Builder {
 
     protected CacheType cache;
 
+    protected String partitionId;
+
     protected BuilderImpl() {}
 
     protected BuilderImpl(MultiKeyStore model) {
       this.keyFieldName = model.keyFieldName();
       this.cacheTTL = model.cacheTTL();
       this.cache = model.cache();
+      this.partitionId = model.partitionId();
     }
 
     public Builder keyFieldName(String keyFieldName) {
@@ -138,6 +164,15 @@ public CacheType cache() {
       return this.cache;
     }
 
+    public Builder partitionId(String partitionId) {
+      this.partitionId = partitionId;
+      return this;
+    }
+
+    public String partitionId() {
+      return this.partitionId;
+    }
+
     public MultiKeyStore build() {
       if (Objects.isNull(this.keyFieldName())) {
         throw new IllegalArgumentException(