m

Author: ajewellamz

DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/src/AttributeResolver.dfy

@@ -32,7 +32,7 @@ module AttributeResolver {
         );
     } else {
       var tableConfig := config.tableEncryptionConfigs[input.TableName];
-      var bucket := GetRandomBucket(tableConfig); // TODO - should this be input?
+      var bucket :- GetRandomBucket(tableConfig);
       var vf :- GetVirtualFields(tableConfig.search.value, input.Item, input.Version);
       var cb :- GetCompoundBeacons(tableConfig.search.value, input.Item, input.Version, bucket);
       return Success(

DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/src/BatchWriteItemTransform.dfy

@@ -54,7 +54,7 @@ module BatchWriteItemTransform {
             //# The Item MUST be [writable](ddb-support.md#writable).
             var _ :- IsWriteable(tableConfig, req.PutRequest.value.Item);
 
-            var bucket := GetRandomBucket(tableConfig);
+            var bucket :- GetRandomBucket(tableConfig);
             var item :- AddSignedBeacons(tableConfig, req.PutRequest.value.Item, bucket);
 
             var encryptRes := tableConfig.itemEncryptor.EncryptItem(EncTypes.EncryptItemInput(plaintextItem:=item));

DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/src/DdbMiddlewareConfig.dfy

@@ -12,6 +12,7 @@ module DdbMiddlewareConfig {
   import SearchableEncryptionInfo
   import DDB = ComAmazonawsDynamodbTypes
   import HexStrings
+  import Random
 
   datatype TableConfig = TableConfig(
     physicalTableName: ComAmazonawsDynamodbTypes.TableName,
@@ -30,16 +31,27 @@ module DdbMiddlewareConfig {
     || config.tableEncryptionConfigs[tableName].plaintextOverride == AwsCryptographyDbEncryptionSdkDynamoDbTypes.PlaintextOverride.FORCE_PLAINTEXT_WRITE_ALLOW_PLAINTEXT_READ
   }
 
-  function method GetRandomBucket(config : TableConfig) : seq<uint8>
+  method GetRandomBucket(config : TableConfig) returns (output : Result<seq<uint8>, Error>)
   {
-    if config.search.Some? then
-      var numBuckets : uint32 := config.search.value.versions[0].numBuckets;
-      if numBuckets <= 1 then
-        []
-      else
-        [1] // TODO
-    else
-      []
+    if config.search.None? {
+      return Success([]);
+    }
+    var numBuckets : uint32 := config.search.value.versions[0].numBuckets;
+    if numBuckets <= 1 {
+      return Success([]);
+    }
+    if (numBuckets > 255) {
+      return Failure(E("Number of buckets exceeds 255"));
+    }
+
+    var randR := Random.GenerateBytes(1);
+    var rand : seq<uint8> :- randR.MapFailure(e => E("Failed to get random byte"));
+    var bucket := rand[0] % (numBuckets as uint8);
+    if bucket == 0 {
+      return Success([]);
+    } else {
+      return Success([bucket as uint8]);
+    }
   }
 
   predicate ValidTableConfig?(config: TableConfig) {

DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/src/PutItemTransform.dfy

@@ -58,7 +58,7 @@ module PutItemTransform {
                                      input.sdkInput.ConditionExpression,
                                      input.sdkInput.ExpressionAttributeNames,
                                      input.sdkInput.ExpressionAttributeValues);
-    var bucket := GetRandomBucket(tableConfig);
+    var bucket :- GetRandomBucket(tableConfig);
     var item :- AddSignedBeacons(tableConfig, input.sdkInput.Item, bucket);
     var encryptRes := tableConfig.itemEncryptor.EncryptItem(
       EncTypes.EncryptItemInput(plaintextItem:=item)

DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/src/ScanTransform.dfy

@@ -43,7 +43,7 @@ module ScanTransform {
       :- Need(input.sdkInput.ConditionalOperator.None?, E("Legacy parameter 'ConditionalOperator' not supported in UpdateItem with Encryption"));
       var tableConfig := config.tableEncryptionConfigs[input.sdkInput.TableName];
 
-    var bucket := []; // TODO
+      var bucket := []; // TODO
       var finalResult :- ScanInputForBeacons(tableConfig, input.sdkInput, bucket);
       return Success(ScanInputTransformOutput(transformedInput := finalResult));
     }

DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/src/TransactWriteItemsTransform.dfy

@@ -89,7 +89,7 @@ module TransactWriteItemsTransform {
                                          item.Put.value.ExpressionAttributeNames,
                                          item.Put.value.ExpressionAttributeValues);
 
-        var bucket := GetRandomBucket(tableConfig);
+        var bucket :- GetRandomBucket(tableConfig);
         var beaconItem :- AddSignedBeacons(tableConfig, item.Put.value.Item, bucket);
 
         var encryptRes := tableConfig.itemEncryptor.EncryptItem(