chore(dafny): Add bucket beacon support

DynamoDbEncryption/dafny/DynamoDbEncryption/Model/AwsCryptographyDbEncryptionSdkDynamoDbTypes.dfy

@@ -46,12 +46,86 @@ module {:extern "software.amazon.cryptography.dbencryptionsdk.dynamodb.internald
     nameonly compoundBeacons: Option<CompoundBeaconList> := Option.None ,
     nameonly virtualFields: Option<VirtualFieldList> := Option.None ,
     nameonly encryptedParts: Option<EncryptedPartsList> := Option.None ,
-    nameonly signedParts: Option<SignedPartsList> := Option.None
+    nameonly signedParts: Option<SignedPartsList> := Option.None ,
+    nameonly numberOfBuckets: Option<BucketCount> := Option.None
   )
   type BeaconVersionList = x: seq<BeaconVersion> | IsValid_BeaconVersionList(x) witness *
   predicate method IsValid_BeaconVersionList(x: seq<BeaconVersion>) {
     ( 1 <= |x| <= 1 )
   }
+  type BucketCount = x: int32 | IsValid_BucketCount(x) witness *
+  predicate method IsValid_BucketCount(x: int32) {
+    ( 1 <= x <= 255 )
+  }
+  type BucketNumber = x: int32 | IsValid_BucketNumber(x) witness *
+  predicate method IsValid_BucketNumber(x: int32) {
+    ( 0 <= x <= 254 )
+  }
+  class IBucketSelectorCallHistory {
+    ghost constructor() {
+      GetBucketNumber := [];
+    }
+    ghost var GetBucketNumber: seq<DafnyCallEvent<GetBucketNumberInput, Result<GetBucketNumberOutput, Error>>>
+  }
+  trait {:termination false} IBucketSelector
+  {
+    // Helper to define any additional modifies/reads clauses.
+    // If your operations need to mutate state,
+    // add it in your constructor function:
+    // Modifies := {your, fields, here, History};
+    // If you do not need to mutate anything:
+    // Modifies := {History};
+
+    ghost const Modifies: set<object>
+    // For an unassigned field defined in a trait,
+    // Dafny can only assign a value in the constructor.
+    // This means that for Dafny to reason about this value,
+    // it needs some way to know (an invariant),
+    // about the state of the object.
+    // This builds on the Valid/Repr paradigm
+    // To make this kind requires safe to add
+    // to methods called from unverified code,
+    // the predicate MUST NOT take any arguments.
+    // This means that the correctness of this requires
+    // MUST only be evaluated by the class itself.
+    // If you require any additional mutation,
+    // then you MUST ensure everything you need in ValidState.
+    // You MUST also ensure ValidState in your constructor.
+    predicate ValidState()
+      ensures ValidState() ==> History in Modifies
+    ghost const History: IBucketSelectorCallHistory
+    predicate GetBucketNumberEnsuresPublicly(input: GetBucketNumberInput , output: Result<GetBucketNumberOutput, Error>)
+    // The public method to be called by library consumers
+    method GetBucketNumber ( input: GetBucketNumberInput )
+      returns (output: Result<GetBucketNumberOutput, Error>)
+      requires
+        && ValidState()
+      modifies Modifies - {History} ,
+               History`GetBucketNumber
+      // Dafny will skip type parameters when generating a default decreases clause.
+      decreases Modifies - {History}
+      ensures
+        && ValidState()
+      ensures GetBucketNumberEnsuresPublicly(input, output)
+      ensures History.GetBucketNumber == old(History.GetBucketNumber) + [DafnyCallEvent(input, output)]
+    {
+      output := GetBucketNumber' (input);
+      History.GetBucketNumber := History.GetBucketNumber + [DafnyCallEvent(input, output)];
+    }
+    // The method to implement in the concrete class.
+    method GetBucketNumber' ( input: GetBucketNumberInput )
+      returns (output: Result<GetBucketNumberOutput, Error>)
+      requires
+        && ValidState()
+      modifies Modifies - {History}
+      // Dafny will skip type parameters when generating a default decreases clause.
+      decreases Modifies - {History}
+      ensures
+        && ValidState()
+      ensures GetBucketNumberEnsuresPublicly(input, output)
+      ensures unchanged(History)
+
+  }
   type Char = x: string | IsValid_Char(x) witness *
   predicate method IsValid_Char(x: string) {
     ( 1 <= |x| <= 1 )
@@ -245,7 +319,8 @@ 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 plaintextOverride: Option<PlaintextOverride> := Option.None ,
+    nameonly bucketSelector: Option<IBucketSelector> := Option.None
   )
   type DynamoDbTableEncryptionConfigList = map<ComAmazonawsDynamodbTypes.TableName, DynamoDbTableEncryptionConfig>
   datatype DynamoDbTablesEncryptionConfig = | DynamoDbTablesEncryptionConfig (
@@ -272,6 +347,13 @@ module {:extern "software.amazon.cryptography.dbencryptionsdk.dynamodb.internald
   datatype GetBranchKeyIdFromDdbKeyOutput = | GetBranchKeyIdFromDdbKeyOutput (
     nameonly branchKeyId: string
   )
+  datatype GetBucketNumberInput = | GetBucketNumberInput (
+    nameonly item: ComAmazonawsDynamodbTypes.AttributeMap ,
+    nameonly numberOfBuckets: BucketCount
+  )
+  datatype GetBucketNumberOutput = | GetBucketNumberOutput (
+    nameonly bucketNumber: BucketNumber
+  )
   datatype GetEncryptedDataKeyDescriptionInput = | GetEncryptedDataKeyDescriptionInput (
     nameonly input: GetEncryptedDataKeyDescriptionUnion
   )
@@ -397,7 +479,8 @@ module {:extern "software.amazon.cryptography.dbencryptionsdk.dynamodb.internald
     nameonly name: string ,
     nameonly length: BeaconBitLength ,
     nameonly loc: Option<TerminalLocation> := Option.None ,
-    nameonly style: Option<BeaconStyle> := Option.None
+    nameonly style: Option<BeaconStyle> := Option.None ,
+    nameonly numberOfBuckets: Option<BucketCount> := Option.None
   )
   type StandardBeaconList = x: seq<StandardBeacon> | IsValid_StandardBeaconList(x) witness *
   predicate method IsValid_StandardBeaconList(x: seq<StandardBeacon>) {

DynamoDbEncryption/dafny/DynamoDbEncryption/Model/DynamoDbEncryption.smithy

@@ -49,6 +49,30 @@ service DynamoDbEncryption {
     errors: [ DynamoDbEncryptionException ]
 }
 
+resource BucketSelector {
+  operations: [GetBucketNumber]
+}
+
+@reference(resource: BucketSelector)
+structure BucketSelectorReference {}
+
+operation GetBucketNumber {
+  input: GetBucketNumberInput,
+  output: GetBucketNumberOutput,
+}
+
+structure GetBucketNumberInput {
+  @required
+  item: AttributeMap,
+  @required
+  numberOfBuckets : BucketCount
+}
+
+structure GetBucketNumberOutput {
+  @required
+  bucketNumber: BucketNumber
+}
+
 @javadoc("Returns encrypted data key description.")
 operation GetEncryptedDataKeyDescription {
     input: GetEncryptedDataKeyDescriptionInput,
@@ -206,7 +230,10 @@ structure DynamoDbTableEncryptionConfig {
     @javadoc("A configuration that override encryption and/or decryption to instead perform legacy encryption and/or decryption. Used as part of migration from version 2.x to version 3.x.")
     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
+    plaintextOverride: PlaintextOverride,
+
+    @javadoc("How to choose the bucket for an item. Default behavior is a random between 0 and numberOfBuckets.")
+    bucketSelector: BucketSelectorReference,
 }
 
 map AttributeActions {
@@ -286,6 +313,12 @@ integer BeaconBitLength
 @range(min: 1)
 integer VersionNumber
 
+@range(min: 1, max: 255)
+integer BucketCount
+
+@range(min: 0, max: 254)
+integer BucketNumber
+
 @length(min: 1, max: 1)
 string Char
 
@@ -669,6 +702,8 @@ structure StandardBeacon {
   loc : TerminalLocation,
   @javadoc("Optional augmented behavior.")
   style : BeaconStyle,
+  @javadoc("The number of separate buckets across which this particular beacon should be divided. Ths must be no greater than the global numberOfBuckets, and can never be changed once an item containing this beacon has been written.")
+  numberOfBuckets : BucketCount
 }
 
 //= specification/searchable-encryption/beacons.md#compound-beacon-initialization
@@ -796,6 +831,8 @@ structure BeaconVersion {
   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
 }
 
 //= specification/searchable-encryption/search-config.md#initialization