feat: limit size of Items and Structures handled (#185)

* feat: limit size of Items and Structures handled

* tests for too many attributes

* test for deep structure

Author: ajewellamz

DynamoDbEncryption/dafny/DynamoDbEncryption/src/DynamoToStruct.dfy

@@ -11,6 +11,7 @@ module DynamoToStruct {
   import opened Wrappers
   import opened StandardLibrary
   import opened StandardLibrary.UInt
+  import opened DynamoDbEncryptionUtil
   import AwsCryptographyDbEncryptionSdkDynamoDbTypes
   import UTF8
   import SortedSets
@@ -264,19 +265,20 @@ module DynamoToStruct {
 
   // convert AttributeValue to byte sequence
   // if `prefix` is true, prefix sequence with TypeID and Length
-  function method {:opaque} AttrToBytes(a : AttributeValue, prefix : bool) : (ret : Result<seq<uint8>, string>)
+  function method {:opaque} AttrToBytes(a : AttributeValue, prefix : bool, depth : nat := 1) : (ret : Result<seq<uint8>, string>)
     decreases a
     ensures ret.Success? && prefix ==> 6 <= |ret.value|
+    ensures MAX_STRUCTURE_DEPTH < depth ==> ret.Failure?
 
     //= specification/dynamodb-encryption-client/ddb-attribute-serialization.md#boolean
     //= type=implication
     //# Boolean MUST be serialized as:
     //# - `0x00` if the value is `false`
     //# - `0x01` if the value is `true`
-    ensures a.BOOL? && !prefix ==>
+    ensures a.BOOL? && !prefix && depth <= MAX_STRUCTURE_DEPTH ==>
       && (a.BOOL  ==> ret.Success? && |ret.value| == BOOL_LEN && ret.value[0] == 1)
       && (!a.BOOL ==> ret.Success? && |ret.value| == BOOL_LEN && ret.value[0] == 0)
-    ensures a.BOOL? && prefix ==>
+    ensures a.BOOL? && prefix && depth <= MAX_STRUCTURE_DEPTH ==>
       && (a.BOOL  ==> (ret.Success? && |ret.value| == PREFIX_LEN+BOOL_LEN && ret.value[PREFIX_LEN] == 1
           && ret.value[0..TYPEID_LEN] == BOOLEAN && ret.value[TYPEID_LEN..PREFIX_LEN] == [0,0,0,1]))
       && (!a.BOOL ==> (ret.Success? && |ret.value| == PREFIX_LEN+BOOL_LEN && ret.value[PREFIX_LEN] == 0
@@ -286,8 +288,8 @@ module DynamoToStruct {
     //= type=implication
     //# Binary MUST be serialized with the identity function;
     //# or more plainly, Binary Attribute Values are used as is.
-    ensures a.B? && !prefix ==> ret.Success? && ret.value == a.B
-    ensures a.B? && prefix && ret.Success? ==>
+    ensures a.B? && !prefix && depth <= MAX_STRUCTURE_DEPTH ==> ret.Success? && ret.value == a.B
+    ensures a.B? && prefix && ret.Success? && depth <= MAX_STRUCTURE_DEPTH ==>
       && ret.value[PREFIX_LEN..] == a.B
       && ret.value[0..TYPEID_LEN] == BINARY
       && U32ToBigEndian(|a.B|).Success?
@@ -297,8 +299,8 @@ module DynamoToStruct {
     //= specification/dynamodb-encryption-client/ddb-attribute-serialization.md#null
     //= type=implication
     //# Null MUST be serialized as a zero-length byte string.
-    ensures a.NULL? && !prefix ==> ret.Success? && |ret.value| == 0
-    ensures a.NULL? &&  prefix ==> ret.Success? && |ret.value| == PREFIX_LEN && ret.value[0..TYPEID_LEN] == NULL && ret.value[TYPEID_LEN..PREFIX_LEN] == [0,0,0,0]
+    ensures a.NULL? && !prefix && depth <= MAX_STRUCTURE_DEPTH ==> ret.Success? && |ret.value| == 0
+    ensures a.NULL? &&  prefix && depth <= MAX_STRUCTURE_DEPTH ==> ret.Success? && |ret.value| == PREFIX_LEN && ret.value[0..TYPEID_LEN] == NULL && ret.value[TYPEID_LEN..PREFIX_LEN] == [0,0,0,0]
 
     //= specification/dynamodb-encryption-client/ddb-attribute-serialization.md#string
     //= type=implication
@@ -463,15 +465,16 @@ module DynamoToStruct {
       && (|a.M| == 0 ==> |ret.value| == PREFIX_LEN + LENGTH_LEN)
 
   {
+    :- Need(depth <= MAX_STRUCTURE_DEPTH, "Depth of attribute structure to serialize exceeds limit of " + MAX_STRUCTURE_DEPTH_STR);
     var baseBytes :- match a {
       case S(s) => UTF8.Encode(s)
       case N(n) => var nn :- Norm.NormalizeNumber(n); UTF8.Encode(nn)
       case B(b) => Success(b)
       case SS(ss) => StringSetAttrToBytes(ss)
       case NS(ns) => NumberSetAttrToBytes(ns)
       case BS(bs) => BinarySetAttrToBytes(bs)
-      case M(m) => MapAttrToBytes(a, m)
-      case L(l) => ListAttrToBytes(l)
+      case M(m) => MapAttrToBytes(a, m, depth)
+      case L(l) => ListAttrToBytes(l, depth)
       case NULL(n) => Success([])
       case BOOL(b) => Success([BoolToUint8(b)])
     };
@@ -516,7 +519,7 @@ module DynamoToStruct {
   // along with the corresponding precondition,
   // lets Dafny find the correct termination metric.
   // See "The Parent Trick" for details: <https://leino.science/papers/krml283.html>.
-  function method MapAttrToBytes(ghost parent: AttributeValue, m: MapAttributeValue): (ret: Result<seq<uint8>, string>)
+  function method MapAttrToBytes(ghost parent: AttributeValue, m: MapAttributeValue, depth : nat): (ret: Result<seq<uint8>, string>)
     requires forall kv <- m.Items :: kv.1 < parent
   {
     //= specification/dynamodb-encryption-client/ddb-attribute-serialization.md#value-type
@@ -532,17 +535,17 @@ module DynamoToStruct {
     //= specification/dynamodb-encryption-client/ddb-attribute-serialization.md#map-value
     //# A Map MAY hold any DynamoDB Attribute Value data type,
     //# and MAY hold values of different types.
-    var bytesResults := map kv <- m.Items :: kv.0 := AttrToBytes(kv.1, true);
+    var bytesResults := map kv <- m.Items :: kv.0 := AttrToBytes(kv.1, true, depth+1);
     var count :- U32ToBigEndian(|m|);
     var bytes :- SimplifyMapValue(bytesResults);
     var body :- CollectMap(bytes);
     Success(count + body)
   }
 
-  function method ListAttrToBytes(l: ListAttributeValue): (ret: Result<seq<uint8>, string>)
+  function method ListAttrToBytes(l: ListAttributeValue, depth : nat): (ret: Result<seq<uint8>, string>)
   {
     var count :- U32ToBigEndian(|l|);
-    var body :- CollectList(l);
+    var body :- CollectList(l, depth);
     Success(count + body)
   }
 
@@ -672,6 +675,7 @@ module DynamoToStruct {
   //# and MAY hold values of different types.
   function method {:opaque} CollectList(
     listToSerialize : ListAttributeValue,
+    depth : nat,
     serialized : seq<uint8> := []
   )
     : (ret : Result<seq<uint8>, string>)
@@ -681,8 +685,8 @@ module DynamoToStruct {
     if |listToSerialize| == 0 then
       Success(serialized)
     else
-      var val :- AttrToBytes(listToSerialize[0], true);
-      CollectList(listToSerialize[1..], serialized + val)
+      var val :- AttrToBytes(listToSerialize[0], true, depth+1);
+      CollectList(listToSerialize[1..], depth, serialized + val)
   }
 
   function method SerializeMapItem(key : string, value : seq<uint8>) : (ret : Result<seq<uint8>, string>)
@@ -881,7 +885,8 @@ module DynamoToStruct {
   function method {:vcs_split_on_every_assert} {:opaque} DeserializeList(
     serialized : seq<uint8>, 
     remainingCount : nat, 
-    origSerializedSize : nat, 
+    ghost origSerializedSize : nat,
+    depth : nat,
     resultList : AttrValueAndLength)
     : (ret : Result<AttrValueAndLength, string>)
     requires resultList.val.L?
@@ -902,17 +907,18 @@ module DynamoToStruct {
       if |serialized| < len then
         Failure("Out of bytes reading Content of List element")
       else
-        var nval :- BytesToAttr(serialized[..len], TerminalTypeId, false);
+        var nval :- BytesToAttr(serialized[..len], TerminalTypeId, false, depth+1);
         var nattr := AttributeValue.L(resultList.val.L + [nval.val]);
-        DeserializeList(serialized[len..], remainingCount-1, origSerializedSize, AttrValueAndLength(nattr, resultList.len + len + 6))
+        DeserializeList(serialized[len..], remainingCount-1, origSerializedSize, depth, AttrValueAndLength(nattr, resultList.len + len + 6))
   }
 
   // Bytes to Map
   // Can't be {:tailrecursion} because it calls BytesToAttr which might again call DeserializeMap
   function method {:vcs_split_on_every_assert} {:opaque} DeserializeMap(
     serialized : seq<uint8>,
     remainingCount : nat,
-    origSerializedSize : nat,
+    ghost origSerializedSize : nat,
+    depth : nat,
     resultMap : AttrValueAndLength)
     : (ret : Result<AttrValueAndLength, string>)
     requires resultMap.val.M?
@@ -948,7 +954,7 @@ module DynamoToStruct {
       var serialized := serialized[2..];
 
       // get value and construct result
-      var nval :- BytesToAttr(serialized, TerminalTypeId_value, true);
+      var nval :- BytesToAttr(serialized, TerminalTypeId_value, true, depth+1);
       var serialized := serialized[nval.len..];
 
       //= specification/dynamodb-encryption-client/ddb-attribute-serialization.md#key-value-pair-entries
@@ -960,15 +966,23 @@ module DynamoToStruct {
       var nattr := AttributeValue.M(resultMap.val.M[key := nval.val]);
       var newResultMap := AttrValueAndLength(nattr, resultMap.len + nval.len + 8 + len);
       assert |serialized| + newResultMap.len == origSerializedSize;
-      DeserializeMap(serialized, remainingCount - 1, origSerializedSize, newResultMap)
+      DeserializeMap(serialized, remainingCount - 1, origSerializedSize, depth, newResultMap)
   }
 
   // Bytes to AttributeValue
   // Can't be {:tailrecursion} because it calls DeserializeList and DeserializeMap which then call BytesToAttr
-  function method {:vcs_split_on_every_assert} {:opaque} BytesToAttr(value : seq<uint8>, typeId : TerminalTypeId, hasLen : bool) : (ret : Result<AttrValueAndLength, string>)
+  function method {:vcs_split_on_every_assert} {:opaque} BytesToAttr(
+    value : seq<uint8>,
+    typeId : TerminalTypeId,
+    hasLen : bool,
+    depth : nat := 1
+  )
+    : (ret : Result<AttrValueAndLength, string>)
     ensures ret.Success? ==> ret.value.len <= |value|
+    ensures MAX_STRUCTURE_DEPTH < depth ==> ret.Failure?
     decreases |value|
   {
+    :- Need(depth <= MAX_STRUCTURE_DEPTH, "Depth of attribute structure to deserialize exceeds limit of " + MAX_STRUCTURE_DEPTH_STR);
     var len :- if hasLen then
         if |value| < LENGTH_LEN then
           Failure("Out of bytes reading length")
@@ -1043,15 +1057,15 @@ module DynamoToStruct {
       else
         var len :- BigEndianToU32(value);
         var value := value[LENGTH_LEN..];
-        DeserializeMap(value, len, |value| + LENGTH_LEN + lengthBytes,  AttrValueAndLength(AttributeValue.M(map[]), LENGTH_LEN + lengthBytes))
+        DeserializeMap(value, len, |value| + LENGTH_LEN + lengthBytes, depth, AttrValueAndLength(AttributeValue.M(map[]), LENGTH_LEN + lengthBytes))
 
     else if typeId == LIST then
       if |value| < LENGTH_LEN then
         Failure("List Structured Data has less than 4 bytes")
       else
         var len :- BigEndianToU32(value);
         var value := value[LENGTH_LEN..];
-        DeserializeList(value, len, |value| + LENGTH_LEN + lengthBytes, AttrValueAndLength(AttributeValue.L([]), LENGTH_LEN + lengthBytes))
+        DeserializeList(value, len, |value| + LENGTH_LEN + lengthBytes, depth, AttrValueAndLength(AttributeValue.L([]), LENGTH_LEN + lengthBytes))
 
     else
       Failure("Unsupported TerminalTypeId")

DynamoDbEncryption/dafny/DynamoDbEncryption/src/Util.dfy

@@ -14,6 +14,9 @@ module DynamoDbEncryptionUtil {
   const BeaconPrefix := "aws_dbe_b_"
   const VersionPrefix := "aws_dbe_v_"
 
+  const MAX_STRUCTURE_DEPTH := 32
+  const MAX_STRUCTURE_DEPTH_STR := "32"
+
   type HmacKeyMap = map<string, Bytes>
 
   // For Multi-Tenant Queries, it's OK to have no KeyId in the query

DynamoDbEncryption/dafny/DynamoDbEncryption/test/DynamoToStruct.dfy

@@ -11,6 +11,7 @@ module DynamoToStructTest {
   import opened ComAmazonawsDynamodbTypes
   import opened AwsCryptographyDbEncryptionSdkStructuredEncryptionTypes
   import opened StandardLibrary.UInt
+  import opened DynamoDbEncryptionUtil
 
   method DoFail(data : seq<uint8>, typeId : TerminalTypeId)
   {
@@ -46,13 +47,13 @@ module DynamoToStructTest {
     DoFail([], LIST);
   }
 
-    const k := 'k' as uint8;
-    const e := 'e' as uint8;
-    const y := 'y' as uint8;
-    const A := 'A' as uint8;
-    const B := 'B' as uint8;
-    const C := 'C' as uint8;
-    const D := 'D' as uint8;
+  const k := 'k' as uint8;
+  const e := 'e' as uint8;
+  const y := 'y' as uint8;
+  const A := 'A' as uint8;
+  const B := 'B' as uint8;
+  const C := 'C' as uint8;
+  const D := 'D' as uint8;
 
   method {:test} TestBadType() {
     DoSucceed([0,0,0,1, 0,0, 0,0,0,0], LIST, 5);
@@ -262,13 +263,15 @@ module DynamoToStructTest {
     //= type=test
     //# A Map MAY hold any DynamoDB Attribute Value data type,
     //# and MAY hold values of different types.
-    var encodedMapData := StructuredDataTerminal(value :=
-      [0,0,0,4,
+    var encodedMapData := StructuredDataTerminal(
+      value := [
+        0,0,0,4,
         0,1, 0,0,0,4, k,e,y,A, 0xff,0xff, 0,0,0,5, 1,2,3,4,5,
         0,1, 0,0,0,4, k,e,y,B, 0,0, 0,0,0,0,
         0,1, 0,0,0,4, k,e,y,C, 0,4, 0,0,0,1, 0,
-        0,1, 0,0,0,4, k,e,y,D, 3,0, 0,0,0,28, 0,0,0,3, 0xff,0xff, 0,0,0,5, 1,2,3,4,5, 0,0, 0,0,0,0, 0,4, 0,0,0,1, 0],
-    typeId := [2,0]);
+        0,1, 0,0,0,4, k,e,y,D, 3,0, 0,0,0,28, 0,0,0,3, 0xff,0xff, 0,0,0,5, 1,2,3,4,5, 0,0, 0,0,0,0, 0,4, 0,0,0,1, 0
+      ],
+      typeId := [2,0]);
     var encodedMapValue := StructuredData(content := Terminal(encodedMapData), attributes := None);
     var mapStruct := AttrToStructured(mapValue);
     expect mapStruct.Success?;
@@ -277,7 +280,7 @@ module DynamoToStructTest {
     var newMapValue := StructuredToAttr(mapStruct.value);
     expect newMapValue.Success?;
     expect newMapValue.value == mapValue;
-}
+  }
 
   //= specification/dynamodb-encryption-client/ddb-item-conversion.md#overview
   //= type=test
@@ -307,4 +310,34 @@ module DynamoToStructTest {
     var nAttrMap :- expect StructuredToItem(struct);
     expect attrMap == nAttrMap;
   }
+
+  method {:test} TestMaxDepth() {
+    var value := AttributeValue.S("hello");
+    for i := 0 to (MAX_STRUCTURE_DEPTH-1) {
+      if i % 2 == 0 {
+        value := AttributeValue.M(map["key" := value]);
+      } else {
+        value := AttributeValue.L([value]);
+      }
+    }
+    var attrMap : AttributeMap := map["key1" := value];
+    var struct :- expect ItemToStructured(attrMap);
+    var nAttrMap :- expect StructuredToItem(struct);
+    expect attrMap == nAttrMap;
+  }
+
+  method {:test} TestTooDeep() {
+    var value := AttributeValue.S("hello");
+    for i := 0 to MAX_STRUCTURE_DEPTH {
+      if i % 2 == 0 {
+        value := AttributeValue.M(map["key" := value]);
+      } else {
+        value := AttributeValue.L([value]);
+      }
+    }
+    var attrMap : AttributeMap := map["key1" := value];
+    var struct := ItemToStructured(attrMap);
+    expect struct.Failure?;
+    expect struct.error == E("Depth of attribute structure to serialize exceeds limit of 32");
+  }
 }

DynamoDbEncryption/dafny/DynamoDbEncryptionTransforms/test/TestFixtures.dfy

@@ -83,14 +83,14 @@ module TestFixtures {
     map["bar" := CSE.SIGN_ONLY, "encrypt" := CSE.ENCRYPT_AND_SIGN, "sign" := CSE.SIGN_ONLY]
   }
 
-  method GetEncryptorConfig() returns (output : DynamoDbItemEncryptorConfig) {
+  method GetEncryptorConfigFromActions(actions : AttributeActions) returns (output : DynamoDbItemEncryptorConfig) {
     var keyring := GetKmsKeyring();
     var logicalTableName := GetTableName("foo");
     output := DynamoDbItemEncryptorConfig(
       logicalTableName := logicalTableName,
       partitionKeyName := "bar",
       sortKeyName := None(),
-      attributeActions := GetAttributeActions(),
+      attributeActions := actions,
       allowedUnauthenticatedAttributes := Some(["nothing"]),
       allowedUnauthenticatedAttributePrefix := None(),
       keyring := Some(keyring),
@@ -101,6 +101,10 @@ module TestFixtures {
     );
   }
 
+  method GetEncryptorConfig() returns (output : DynamoDbItemEncryptorConfig) {
+    output := GetEncryptorConfigFromActions(GetAttributeActions());
+  }
+
   method GetDynamoDbItemEncryptorFrom(config : DynamoDbItemEncryptorConfig)
     returns (encryptor: DynamoDbItemEncryptor.DynamoDbItemEncryptorClient)
     ensures encryptor.ValidState()

DynamoDbEncryption/dafny/DynamoDbItemEncryptor/src/AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations.dfy

@@ -28,6 +28,7 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
   import DDBE = AwsCryptographyDbEncryptionSdkDynamoDbTypes
   import DynamoDbEncryptionUtil
   import StructuredEncryptionUtil
+  import StandardLibrary.String
 
   datatype Config = Config(
     nameonly cmpClient : MaterialProviders.MaterialProvidersClient,
@@ -614,13 +615,19 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
       ==>
         && output.value.encryptedItem == input.plaintextItem
         && output.value.parsedHeader == None
+
+    ensures output.Success? ==> |input.plaintextItem| <= MAX_ATTRIBUTE_COUNT
   {
     :- Need(
       && config.partitionKeyName in input.plaintextItem
       && (config.sortKeyName.None? || config.sortKeyName.value in input.plaintextItem)
-    , DynamoDbItemEncryptorException( message := "Configuration missmatch partition or sort key does not exist in item."));
+    , DynamoDbItemEncryptorException( message := "Configuration mismatch partition or sort key does not exist in item."));
 
-    assert {:split_here} true;
+    if |input.plaintextItem| > MAX_ATTRIBUTE_COUNT {
+      var actCount := String.Base10Int2String(|input.plaintextItem|);
+      var maxCount := String.Base10Int2String(MAX_ATTRIBUTE_COUNT);
+      return Failure(E("Item to encrypt had " + actCount + " attributes, but maximum allowed is " + maxCount));
+    }
 
     //= specification/dynamodb-encryption-client/encrypt-item.md#behavior
     //# If a [Legacy Policy](./ddb-table-encryption-config.md#legacy-policy) of
@@ -824,10 +831,17 @@ module AwsCryptographyDbEncryptionSdkDynamoDbItemEncryptorOperations refines Abs
         && output.value.plaintextItem == input.encryptedItem
         && output.value.parsedHeader == None
   {
+    var realCount := |set k <- input.encryptedItem | !(ReservedPrefix <= k)|;
+    if realCount > MAX_ATTRIBUTE_COUNT {
+      var actCount := String.Base10Int2String(realCount);
+      var maxCount := String.Base10Int2String(MAX_ATTRIBUTE_COUNT);
+      return Failure(E("Item to decrypt had " + actCount + " attributes, but maximum allowed is " + maxCount));
+    }
+
     :- Need(
       && config.partitionKeyName in input.encryptedItem
       && (config.sortKeyName.None? || config.sortKeyName.value in input.encryptedItem)
-    , DynamoDbItemEncryptorException( message := "Configuration missmatch partition or sort key does not exist in item."));
+    , DynamoDbItemEncryptorException( message := "Configuration mismatch partition or sort key does not exist in item."));
 
     //= specification/dynamodb-encryption-client/decrypt-item.md#behavior
     //# If a [Legacy Policy](./ddb-table-encryption-config.md#legacy-policy) of

DynamoDbEncryption/dafny/DynamoDbItemEncryptor/src/Util.dfy

@@ -12,6 +12,7 @@ module DynamoDbItemEncryptorUtil {
   const ReservedPrefix := "aws_dbe_"
   const BeaconPrefix := ReservedPrefix + "b_"
   const VersionPrefix := ReservedPrefix + "v_"
+  const MAX_ATTRIBUTE_COUNT := 100
 
   function method E(msg : string) : Error
   {