m

Author: ajewellamz

DynamoDbEncryption/dafny/DynamoDbEncryption/Model/AwsCryptographyDbEncryptionSdkDynamoDbTypes.dfy

@@ -47,7 +47,9 @@ module {:extern "software.amazon.cryptography.dbencryptionsdk.dynamodb.internald
     nameonly virtualFields: Option<VirtualFieldList> := Option.None ,
     nameonly encryptedParts: Option<EncryptedPartsList> := Option.None ,
     nameonly signedParts: Option<SignedPartsList> := Option.None ,
-    nameonly numberOfBuckets: Option<BucketCount> := Option.None
+    nameonly maximumNumberOfBuckets: Option<BucketCount> := Option.None ,
+    nameonly defaultNumberOfBuckets: Option<BucketCount> := Option.None ,
+    nameonly bucketSelector: Option<IBucketSelector> := Option.None
   )
   type BeaconVersionList = x: seq<BeaconVersion> | IsValid_BeaconVersionList(x) witness *
   predicate method IsValid_BeaconVersionList(x: seq<BeaconVersion>) {
@@ -166,9 +168,11 @@ module {:extern "software.amazon.cryptography.dbencryptionsdk.dynamodb.internald
     ghost constructor() {
       CreateDynamoDbEncryptionBranchKeyIdSupplier := [];
       GetEncryptedDataKeyDescription := [];
+      GetNumberOfQueries := [];
     }
     ghost var CreateDynamoDbEncryptionBranchKeyIdSupplier: seq<DafnyCallEvent<CreateDynamoDbEncryptionBranchKeyIdSupplierInput, Result<CreateDynamoDbEncryptionBranchKeyIdSupplierOutput, Error>>>
     ghost var GetEncryptedDataKeyDescription: seq<DafnyCallEvent<GetEncryptedDataKeyDescriptionInput, Result<GetEncryptedDataKeyDescriptionOutput, Error>>>
+    ghost var GetNumberOfQueries: seq<DafnyCallEvent<GetNumberOfQueriesInput, Result<GetNumberOfQueriesOutput, Error>>>
   }
   trait {:termination false} IDynamoDbEncryptionClient
   {
@@ -238,6 +242,21 @@ module {:extern "software.amazon.cryptography.dbencryptionsdk.dynamodb.internald
       ensures GetEncryptedDataKeyDescriptionEnsuresPublicly(input, output)
       ensures History.GetEncryptedDataKeyDescription == old(History.GetEncryptedDataKeyDescription) + [DafnyCallEvent(input, output)]
 
+    predicate GetNumberOfQueriesEnsuresPublicly(input: GetNumberOfQueriesInput , output: Result<GetNumberOfQueriesOutput, Error>)
+    // The public method to be called by library consumers
+    method GetNumberOfQueries ( input: GetNumberOfQueriesInput )
+      returns (output: Result<GetNumberOfQueriesOutput, Error>)
+      requires
+        && ValidState()
+      modifies Modifies - {History} ,
+               History`GetNumberOfQueries
+      // Dafny will skip type parameters when generating a default decreases clause.
+      decreases Modifies - {History}
+      ensures
+        && ValidState()
+      ensures GetNumberOfQueriesEnsuresPublicly(input, output)
+      ensures History.GetNumberOfQueries == old(History.GetNumberOfQueries) + [DafnyCallEvent(input, output)]
+
   }
   datatype DynamoDbEncryptionConfig = | DynamoDbEncryptionConfig (
 
@@ -319,8 +338,7 @@ module {:extern "software.amazon.cryptography.dbencryptionsdk.dynamodb.internald
     nameonly keyring: Option<AwsCryptographyMaterialProvidersTypes.IKeyring> := Option.None ,
     nameonly cmm: Option<AwsCryptographyMaterialProvidersTypes.ICryptographicMaterialsManager> := Option.None ,
     nameonly legacyOverride: Option<LegacyOverride> := Option.None ,
-    nameonly plaintextOverride: Option<PlaintextOverride> := Option.None ,
-    nameonly bucketSelector: Option<IBucketSelector> := Option.None
+    nameonly plaintextOverride: Option<PlaintextOverride> := Option.None
   )
   type DynamoDbTableEncryptionConfigList = map<ComAmazonawsDynamodbTypes.TableName, DynamoDbTableEncryptionConfig>
   datatype DynamoDbTablesEncryptionConfig = | DynamoDbTablesEncryptionConfig (
@@ -363,6 +381,12 @@ module {:extern "software.amazon.cryptography.dbencryptionsdk.dynamodb.internald
   datatype GetEncryptedDataKeyDescriptionUnion =
     | header(header: seq<uint8>)
     | item(item: ComAmazonawsDynamodbTypes.AttributeMap)
+  datatype GetNumberOfQueriesInput = | GetNumberOfQueriesInput (
+    nameonly input: ComAmazonawsDynamodbTypes.QueryInput
+  )
+  datatype GetNumberOfQueriesOutput = | GetNumberOfQueriesOutput (
+    nameonly numberOfQueries: BucketCount
+  )
   datatype GetPrefix = | GetPrefix (
     nameonly length: int32
   )
@@ -662,6 +686,26 @@ abstract module AbstractAwsCryptographyDbEncryptionSdkDynamoDbService
       History.GetEncryptedDataKeyDescription := History.GetEncryptedDataKeyDescription + [DafnyCallEvent(input, output)];
     }
 
+    predicate GetNumberOfQueriesEnsuresPublicly(input: GetNumberOfQueriesInput , output: Result<GetNumberOfQueriesOutput, Error>)
+    {Operations.GetNumberOfQueriesEnsuresPublicly(input, output)}
+    // The public method to be called by library consumers
+    method GetNumberOfQueries ( input: GetNumberOfQueriesInput )
+      returns (output: Result<GetNumberOfQueriesOutput, Error>)
+      requires
+        && ValidState()
+      modifies Modifies - {History} ,
+               History`GetNumberOfQueries
+      // Dafny will skip type parameters when generating a default decreases clause.
+      decreases Modifies - {History}
+      ensures
+        && ValidState()
+      ensures GetNumberOfQueriesEnsuresPublicly(input, output)
+      ensures History.GetNumberOfQueries == old(History.GetNumberOfQueries) + [DafnyCallEvent(input, output)]
+    {
+      output := Operations.GetNumberOfQueries(config, input);
+      History.GetNumberOfQueries := History.GetNumberOfQueries + [DafnyCallEvent(input, output)];
+    }
+
   }
 }
 abstract module AbstractAwsCryptographyDbEncryptionSdkDynamoDbOperations {
@@ -711,4 +755,20 @@ abstract module AbstractAwsCryptographyDbEncryptionSdkDynamoDbOperations {
     ensures
       && ValidInternalConfig?(config)
     ensures GetEncryptedDataKeyDescriptionEnsuresPublicly(input, output)
+
+
+  predicate GetNumberOfQueriesEnsuresPublicly(input: GetNumberOfQueriesInput , output: Result<GetNumberOfQueriesOutput, Error>)
+  // The private method to be refined by the library developer
+
+
+  method GetNumberOfQueries ( config: InternalConfig , input: GetNumberOfQueriesInput )
+    returns (output: Result<GetNumberOfQueriesOutput, Error>)
+    requires
+      && ValidInternalConfig?(config)
+    modifies ModifiesInternalConfig(config)
+    // Dafny will skip type parameters when generating a default decreases clause.
+    decreases ModifiesInternalConfig(config)
+    ensures
+      && ValidInternalConfig?(config)
+    ensures GetNumberOfQueriesEnsuresPublicly(input, output)
 }

DynamoDbEncryption/dafny/DynamoDbEncryption/Model/DynamoDbEncryption.smithy

@@ -21,6 +21,7 @@ use aws.cryptography.dbEncryptionSdk.structuredEncryption#CryptoAction
 
 use com.amazonaws.dynamodb#DynamoDB_20120810
 use com.amazonaws.dynamodb#AttributeMap
+use com.amazonaws.dynamodb#QueryInput
 use com.amazonaws.dynamodb#TableName
 use com.amazonaws.dynamodb#AttributeName
 use com.amazonaws.dynamodb#Key
@@ -45,7 +46,7 @@ use aws.cryptography.materialProviders#AwsCryptographicMaterialProviders
 )
 service DynamoDbEncryption {
     version: "2024-04-02",
-    operations: [ CreateDynamoDbEncryptionBranchKeyIdSupplier, GetEncryptedDataKeyDescription],
+    operations: [ CreateDynamoDbEncryptionBranchKeyIdSupplier, GetEncryptedDataKeyDescription, GetNumberOfQueries],
     errors: [ DynamoDbEncryptionException ]
 }
 
@@ -73,6 +74,20 @@ structure GetBucketNumberOutput {
   bucketNumber: BucketNumber
 }
 
+@javadoc("Return the necessary number of query operations for this query, based on bucket usage.")
+operation GetNumberOfQueries {
+    input: GetNumberOfQueriesInput,
+    output: GetNumberOfQueriesOutput,
+}
+structure GetNumberOfQueriesInput {
+    @required
+    input: QueryInput
+}
+structure GetNumberOfQueriesOutput {
+    @required
+    numberOfQueries: BucketCount
+}
+
 @javadoc("Returns encrypted data key description.")
 operation GetEncryptedDataKeyDescription {
     input: GetEncryptedDataKeyDescriptionInput,
@@ -231,9 +246,6 @@ structure DynamoDbTableEncryptionConfig {
     legacyOverride: LegacyOverride,
     @javadoc("A configuration that override encryption and/or decryption to instead passthrough and write and/or read plaintext. Used to update plaintext tables to fully use client-side encryption.")
     plaintextOverride: PlaintextOverride,
-
-    @javadoc("How to choose the bucket for an item. Default behavior is a random between 0 and numberOfBuckets.")
-    bucketSelector: BucketSelectorReference,
 }
 
 map AttributeActions {
@@ -274,7 +286,7 @@ structure LegacyDynamoDbEncryptorReference {}
 @javadoc("A configuration for overriding encryption and/or decryption to instead perform legacy encryption and decryption.")
 structure LegacyOverride {
     @required
-    @javadoc("A policy which configurates whether legacy behavior overrides encryption and/or decryption.")
+    @javadoc("A policy which configures whether legacy behavior overrides encryption and/or decryption.")
     policy: LegacyPolicy,
     @required
     @javadoc("A configuration for the legacy DynamoDB Encryption Client's Encryptor.")
@@ -509,7 +521,7 @@ structure GetSegment {
   @javadoc("The characters to split on.")
   split : Char,
   @required
-  @javadoc("The index of the split string result to return. 0 represents the segment before the first split character. -1 respresents the segment after the last split character.")
+  @javadoc("The index of the split string result to return. 0 represents the segment before the first split character. -1 represents the segment after the last split character.")
   index : Integer
 }
 
@@ -669,7 +681,7 @@ structure Constructor {
 //# - A name -- a string
 //# - A required flag -- a boolean
 
-@javadoc("A part of a Compound Becaon Construction.")
+@javadoc("A part of a Compound Beacon Construction.")
 structure ConstructorPart {
   @required
   @javadoc("The name of the Encrypted Part or Signed Part for which this constructor part gets a value.")
@@ -824,15 +836,22 @@ structure BeaconVersion {
 
   @javadoc("The Compound Beacons to be written with items.")
   compoundBeacons : CompoundBeaconList,
-  @javadoc("The Virtual Fields to be calculated, supporting other searchable enryption configurations.")
+  @javadoc("The Virtual Fields to be calculated, supporting other searchable encryption configurations.")
   virtualFields : VirtualFieldList,
 
   @javadoc("The list of Encrypted Parts that may be included in any compound beacon.")
   encryptedParts : EncryptedPartsList,
   @javadoc("The list of Signed Parts that may be included in any compound beacon.")
   signedParts : SignedPartsList,
+
   @javadoc("The number of separate buckets across which beacons should be divided.")
-  numberOfBuckets : BucketCount
+  maximumNumberOfBuckets : BucketCount,
+
+  @javadoc("The number of buckets for any beacon that doesn't specify a numberOfBuckets")
+  defaultNumberOfBuckets : BucketCount,
+
+  @javadoc("How to choose the bucket for an item. Default behavior is a random between 0 and maximumNumberOfBuckets.")
+  bucketSelector: BucketSelectorReference,
 }
 
 //= specification/searchable-encryption/search-config.md#initialization

DynamoDbEncryption/dafny/DynamoDbEncryption/src/AwsCryptographyDbEncryptionSdkDynamoDbOperations.dfy

@@ -43,6 +43,17 @@ module AwsCryptographyDbEncryptionSdkDynamoDbOperations refines AbstractAwsCrypt
         )
       );
   }
+  predicate GetNumberOfQueriesEnsuresPublicly(input: GetNumberOfQueriesInput , output: Result<GetNumberOfQueriesOutput, Error>)
+  {true}
+
+  method GetNumberOfQueries(config: InternalConfig, input: GetNumberOfQueriesInput)
+    returns (output: Result<GetNumberOfQueriesOutput, Error>)
+    ensures GetNumberOfQueriesEnsuresPublicly(input, output)
+  {
+    return Success(GetNumberOfQueriesOutput(numberOfQueries := 1));
+  }
+
+
   predicate GetEncryptedDataKeyDescriptionEnsuresPublicly(input: GetEncryptedDataKeyDescriptionInput , output: Result<GetEncryptedDataKeyDescriptionOutput, Error>)
   {true}
 

DynamoDbEncryption/dafny/DynamoDbEncryption/src/ConfigToInfo.dfy

@@ -26,6 +26,7 @@ module SearchConfigToInfo {
   import opened StandardLibrary.MemoryMath
   import MaterialProviders
   import SortedSets
+  import Random
 
   import I = SearchableEncryptionInfo
   import V = DdbVirtualFields
@@ -266,6 +267,47 @@ module SearchConfigToInfo {
     output := ConvertVersionWithSource(outer, config, source);
   }
 
+  class DefaultBucketSelector extends IBucketSelector
+  {
+    predicate ValidState()
+      ensures ValidState() ==> History in Modifies
+    { History in Modifies }
+
+    constructor ()
+      ensures ValidState() && fresh(History) && fresh(Modifies)
+    {
+      History := new IBucketSelectorCallHistory();
+      Modifies := { History };
+    }
+
+    predicate GetBucketNumberEnsuresPublicly (
+      input: GetBucketNumberInput ,
+      output: Result<GetBucketNumberOutput, Error> )
+      : (outcome: bool)
+    {
+      true
+    }
+
+    method GetBucketNumber'(input: GetBucketNumberInput)
+      returns (output: Result<GetBucketNumberOutput, Error> )
+      requires ValidState()
+      modifies Modifies - {History}
+      decreases Modifies - {History}
+      ensures ValidState()
+      ensures GetBucketNumberEnsuresPublicly(input, output)
+      ensures unchanged(History)
+    {
+      if input.numberOfBuckets == 1 {
+        return Success(GetBucketNumberOutput(bucketNumber := 0));
+      } else {
+        var randR := Random.GenerateBytes(1);
+        var rand : seq<uint8> :- randR.MapFailure(e => DynamoDbEncryptionException(message := "Failed to get random byte"));
+        var bucket := (rand[0] % (input.numberOfBuckets as uint8)) as BucketNumber;
+        return Success(GetBucketNumberOutput(bucketNumber := bucket));
+      }
+    }
+  }
+
   // convert configured BeaconVersion to internal BeaconVersion
   method ConvertVersionWithSource(
     outer : DynamoDbTableEncryptionConfig,
@@ -280,8 +322,19 @@ module SearchConfigToInfo {
               && output.value.ValidState()
               && output.value.keySource == source
   {
+    var maxBuckets : BucketCount := config.maximumNumberOfBuckets.UnwrapOr(1);
+    var defaultBuckets : BucketCount := config.defaultNumberOfBuckets.UnwrapOr(maxBuckets);
+    :- Need(0 <= maxBuckets as nat < MAX_BUCKET_COUNT, E("Invalid number of buckets specified, " + Base10Int2String(maxBuckets as int) + ", must be 0 < maximumNumberOfBuckets <= 255."));
+    // Zero is invalid, but in Java we can't distinguish None from Some(0)
+    if maxBuckets == 0 {
+      maxBuckets := 1;
+    }
+    if defaultBuckets == 0 {
+      defaultBuckets := maxBuckets;
+    }
+
     var virtualFields :- ConvertVirtualFields(outer, config.virtualFields);
-    var std :- AddStandardBeacons(config.standardBeacons, outer, source.client, virtualFields);
+    var std :- AddStandardBeacons(config.standardBeacons, outer, source.client, virtualFields, maxBuckets, defaultBuckets);
 
     var signed := if config.signedParts.Some? then config.signedParts.value else [];
 
@@ -315,20 +368,22 @@ module SearchConfigToInfo {
         return Failure(E("A beacon key field name of " + name + " was configured, but there's also a virtual field of that name."));
       }
     }
+      var bucketSelector;
+      if outer.search.Some? && outer.search.value.versions[0].bucketSelector.Some? {
+        bucketSelector := outer.search.value.versions[0].bucketSelector.value;
+        assume {:axiom} bucketSelector.ValidState();
+      } else {
+        bucketSelector := new DefaultBucketSelector();
+      }
 
-    var numBuckets : int := config.numberOfBuckets.UnwrapOr(1) as int;
-    :- Need(0 <= numBuckets < MAX_BUCKET_COUNT, E("Invalid number of buckets specified, " + Base10Int2String(numBuckets) + ", must be 0 < numberOfBuckets <= 255."));
-    // Zero is invalid, but in Java we can't distinguish None from Some(0)
-    if numBuckets == 0 {
-      numBuckets := 1;
-    }
     return I.MakeBeaconVersion(
         config.version as I.VersionNumber,
         source,
         beacons,
         virtualFields,
         outer.attributeActionsOnEncrypt,
-        numBuckets as BucketCount
+        bucketSelector,
+        maxBuckets
       );
   }
 
@@ -526,31 +581,28 @@ module SearchConfigToInfo {
       Failure(E("Beacon " + name + " is shared to " + share + " which is not defined."))
   }
 
-  function method GetBucketCount(outer : DynamoDbTableEncryptionConfig, inner : Option<BucketCount>, name : string) :
+  function method GetBucketCount(outer : DynamoDbTableEncryptionConfig, inner : Option<BucketCount>, name : string, maxBuckets : BucketCount, defaultBuckets : BucketCount) :
     Result<BucketCount,Error>
   {
     if outer.search.None? || |outer.search.value.versions| == 0 then
       Success(1)
     else
-      var num := outer.search.value.versions[0].numberOfBuckets;
-      if BucketCountNone(num) then
-        if !BucketCountNone(inner) then
-          Failure(E("Constrained numberOfBuckets for  " + name + " is " + Base10Int2String(inner.value as int) + " but there is not global numberOfBuckets set."))
-        else
-          Success(1)
-      else if BucketCountNone(inner) then
-        Success(num.value)
-      else if inner.value < num.value then
+      if BucketCountNone(inner) then
+        Success(defaultBuckets)
+      else if inner.value < maxBuckets then
         Success(inner.value)
       else
-        Failure(E("Constrained numberOfBuckets for  " + name + " is " + Base10Int2String(inner.value as int) + " but it must be less than the global numberOfBuckets " + Base10Int2String(num.value as int)))
+        Failure(E("Constrained numberOfBuckets for  " + name + " is " + Base10Int2String(inner.value as int) + " but it must be less than the maximumNumberOfBuckets " + Base10Int2String(maxBuckets as int)))
   }
+
   // convert configured StandardBeacons to internal Beacons
   method {:tailrecursion} AddStandardBeacons(
     beacons : seq<StandardBeacon>,
     outer : DynamoDbTableEncryptionConfig,
     client: Primitives.AtomicPrimitivesClient,
     virtualFields : V.VirtualFieldMap,
+    maxBuckets : BucketCount,
+    defaultBuckets : BucketCount,
     converted : I.BeaconMap := map[])
     returns (output : Result<I.BeaconMap, Error>)
     modifies client.Modifies
@@ -608,7 +660,7 @@ module SearchConfigToInfo {
         case sharedSet(t) => share := Some(t.other); isAsSet := true;
       }
     }
-    var bucketCount :- GetBucketCount(outer, beacons[0].numberOfBuckets, beacons[0].name);
+    var bucketCount :- GetBucketCount(outer, beacons[0].numberOfBuckets, beacons[0].name, maxBuckets, defaultBuckets);
     var newBeacon :- B.MakeStandardBeacon(client, beacons[0].name, beacons[0].length as B.BeaconLength, locString,
                                           isPartOnly, isAsSet, share, bucketCount);
 
@@ -628,7 +680,7 @@ module SearchConfigToInfo {
                        + ", but virtual field " + badField.value + " is already defined on that single location."));
     }
 
-    output := AddStandardBeacons(beacons[1..], outer, client, virtualFields, converted[beacons[0].name := I.Standard(newBeacon)]);
+    output := AddStandardBeacons(beacons[1..], outer, client, virtualFields, maxBuckets, defaultBuckets, converted[beacons[0].name := I.Standard(newBeacon)]);
   }
 
   // optional location, defaults to name