fix: Trade AssumeValidAAD for run-time checks and proofs (#203)

* fix: Trade AssumeValidAAD for run-time checks and proofs

DefaultCMM.GetEncryptionMaterials is given a valid encryption context, to which it may
append the mapping `[reservedField := enc_vk]`. This would invalidate the encryption
context if the new entry causes problems with the size limits.

* `reservedField` is not UTF8 or is too long (but we know it's fine, because it's a constant)

* `enc_vk` is not UTF8 or is too long after it has been Base64 encoded (but we know it's fine, because of the new and
  tested postcondition of `Signature.KeyGen`, and because of some new specifications and lemmas regarding `Base64.Encode`)

* the updated encryption context exceeds the maximum number of allowed mappings (this is now run-time checked)

* the updated encryption context contains too many bytes altogether (this is now run-time checked, which involves
   walking over all the entries in the encryption context--possibly expensive)

Fixes #79

* Remove AssumeValidAAD altogether

The previous uses of AssumeValidAAD in test/SDK/Serialize.dfy have now been replaced
by proofs, found in test/Util/TestUtils.dfy.

* Respond to PR comment

* Replace extern postcondition with runtime check

src/Crypto/Signature.dfy

@@ -13,9 +13,28 @@ module {:extern "Signature"} Signature {
         case ECDSA_P256 => 71
         case ECDSA_P384 => 103
       }
+
+      function method FieldSize(): nat {
+        match this
+        case ECDSA_P256 => assert 1 + (256 + 7) / 8 == 33; 33
+        case ECDSA_P384 => assert 1 + (384 + 7) / 8 == 49; 49
+      }
+    }
+
+    method KeyGen(s: ECDSAParams) returns (res: Result<SignatureKeyPair>)
+      ensures match res
+        case Success(sigKeyPair) => |sigKeyPair.verificationKey| == s.FieldSize()
+        case Failure(_) => true
+    {
+      var sigKeyPair :- ExternKeyGen(s);
+      if |sigKeyPair.verificationKey| == s.FieldSize() {
+        return Success(sigKeyPair);
+      } else {
+        return Failure("Incorrect verification-key length from ExternKeyGen.");
+      }
     }
 
-    method {:extern "Signature.ECDSA", "KeyGen"} KeyGen(s : ECDSAParams) returns (res: Result<SignatureKeyPair>)
+    method {:extern "Signature.ECDSA", "ExternKeyGen"} ExternKeyGen(s: ECDSAParams) returns (res: Result<SignatureKeyPair>)
 
     method {:extern "Signature.ECDSA", "Sign"} Sign(s: ECDSAParams, key: seq<uint8>, digest: seq<uint8>) returns (sig: Result<seq<uint8>>)
 

src/SDK/CMM/DefaultCMM.dfy

@@ -69,14 +69,32 @@ module {:extern "DefaultCMMDef"} DefaultCMMDef {
           var signatureKeys :- Signature.KeyGen(param);
           enc_sk := Some(signatureKeys.signingKey);
           var enc_vk :- UTF8.Encode(Base64.Encode(signatureKeys.verificationKey));
+          calc {
+            |enc_vk|;
+          ==  { assert UTF8.IsASCIIString(Base64.Encode(signatureKeys.verificationKey)); }
+            |Base64.Encode(signatureKeys.verificationKey)|;
+          <=  { Base64.EncodeLengthBound(signatureKeys.verificationKey); }
+            |signatureKeys.verificationKey| / 3 * 4 + 4;
+          <  { assert |signatureKeys.verificationKey| <= 3000; }
+            UINT16_LIMIT;
+          }
+
           var reservedField := Materials.EC_PUBLIC_KEY_FIELD;
           assert reservedField in Materials.ReservedKeyValues;
           assert forall i :: i in materialsRequest.encryptionContext.Keys ==> i != reservedField;
           enc_ctx := enc_ctx[reservedField := enc_vk];
+          var len := MessageHeader.ComputeKVPairsLength(enc_ctx);
+          if UINT16_LIMIT <= |enc_ctx| {
+            return Failure("encryption context has too many entries");
+          }
+          if UINT16_LIMIT <= len {
+            return Failure("encryption context too big");
+          }
+          assert ValidAAD(enc_ctx) by {
+            reveal MessageHeader.ValidAAD();
+          }
       }
 
-      MessageHeader.AssumeValidAAD(enc_ctx);  // TODO: we should check this (https://github.com/awslabs/aws-encryption-sdk-dafny/issues/79)
-
       var materials := Materials.EncryptionMaterials.WithoutDataKeys(enc_ctx, id, enc_sk);
       assert materials.encryptionContext == enc_ctx;
       materials :- kr.OnEncrypt(materials);

src/SDK/MessageHeader.dfy

@@ -290,9 +290,6 @@ module {:extern "MessageHeader"} MessageHeader {
     && (forall key :: key in encryptionContext.Keys ==> ValidKVPair((key, encryptionContext[key])))
   }
 
-  lemma {:axiom} AssumeValidAAD(encryptionContext: Materials.EncryptionContext)  // TODO: this should be removed and replaced by something usable
-    ensures ValidAAD(encryptionContext)
-
   predicate ValidFrameLength(frameLength: uint32, contentType: ContentType) {
     match contentType
     case NonFramed => frameLength == 0

src/StandardLibrary/Base64.dfy

@@ -367,12 +367,12 @@ module Base64 {
     if IsBase64String(s) then Success(DecodeValid(s)) else Failure("The encoding is malformed")
   }
 
-  predicate StringIs8Bit(s: string) {
-    forall i :: 0 <= i < |s| ==> s[i] < 256 as char
+  predicate StringIs7Bit(s: string) {
+    forall i :: 0 <= i < |s| ==> s[i] < 128 as char
   }
 
   function method Encode(b: seq<uint8>): (s: seq<char>)
-    ensures StringIs8Bit(s)
+    ensures StringIs7Bit(s)
     ensures |s| % 4 == 0
     ensures IsBase64String(s)
     // Rather than ensure Decode(s) == Success(b) directly, lemmas are used to verify this property
@@ -384,4 +384,53 @@ module Base64 {
     else if |b| % 3 == 1 then EncodeUnpadded(b[..(|b| - 1)]) + Encode2Padding(b[(|b| - 1)..])
     else EncodeUnpadded(b[..(|b| - 2)]) + Encode1Padding(b[(|b| - 2)..])
   }
+
+  lemma EncodeLengthExact(b: seq<uint8>)
+    ensures var s := Encode(b);
+      && (|b| % 3 == 0 ==> |s| == |b| / 3 * 4)
+      && (|b| % 3 != 0 ==> |s| == |b| / 3 * 4 + 4)
+  {
+    var s := Encode(b);
+    if |b| % 3 == 0 {
+      assert s == EncodeUnpadded(b);
+      assert |s| == |b| / 3 * 4;
+    } else if |b| % 3 == 1 {
+      assert s == EncodeUnpadded(b[..(|b| - 1)]) + Encode2Padding(b[(|b| - 1)..]);
+      calc {
+        |s|;
+      ==
+        |EncodeUnpadded(b[..(|b| - 1)])| + |Encode2Padding(b[(|b| - 1)..])|;
+      ==  { assert |Encode2Padding(b[(|b| - 1)..])| == 4; }
+        |EncodeUnpadded(b[..(|b| - 1)])| + 4;
+      ==  { assert |EncodeUnpadded(b[..(|b| - 1)])| == |b[..(|b| - 1)]| / 3 * 4; }
+        |b[..(|b| - 1)]| / 3 * 4 + 4;
+      ==  { assert |b[..(|b| - 1)]| == |b| - 1; }
+        (|b| - 1) / 3 * 4 + 4;
+      ==  { assert (|b| - 1) / 3 == |b| / 3; }
+        |b| / 3 * 4 + 4;
+      }
+    } else {
+      assert s == EncodeUnpadded(b[..(|b| - 2)]) + Encode1Padding(b[(|b| - 2)..]);
+      calc {
+        |s|;
+      ==
+        |EncodeUnpadded(b[..(|b| - 2)])| + |Encode1Padding(b[(|b| - 2)..])|;
+      ==  { assert |Encode1Padding(b[(|b| - 2)..])| == 4; }
+        |EncodeUnpadded(b[..(|b| - 2)])| + 4;
+      ==  { assert |EncodeUnpadded(b[..(|b| - 2)])| == |b[..(|b| - 2)]| / 3 * 4; }
+        |b[..(|b| - 2)]| / 3 * 4 + 4;
+      ==  { assert |b[..(|b| - 2)]| == |b| - 2; }
+        (|b| - 2) / 3 * 4 + 4;
+      ==  { assert (|b| - 2) / 3 == |b| / 3; }
+        |b| / 3 * 4 + 4;
+      }
+    }
+  }
+
+  lemma EncodeLengthBound(b: seq<uint8>)
+    ensures var s := Encode(b);
+      |s| <= |b| / 3 * 4 + 4
+  {
+    EncodeLengthExact(b);
+  }
 }

src/extern/dotnet/Signature.cs

@@ -28,7 +28,7 @@ public ECDSAUnsupportedParametersException(ECDSAParams parameters)
     }
 
     public partial class ECDSA {
-        public static STL.Result<SignatureKeyPair> KeyGen(ECDSAParams x) {
+        public static STL.Result<SignatureKeyPair> ExternKeyGen(ECDSAParams x) {
             try {
                 ECKeyPairGenerator generator = new ECKeyPairGenerator();
                 SecureRandom rng = new SecureRandom();

test/Crypto/Signature.dfy

@@ -9,8 +9,14 @@ module TestSignature {
   import Signature
   import UTF8
 
+  method RequireGoodKeyLengths(s: Signature.ECDSAParams, sigKeyPair: Signature.SignatureKeyPair) {
+    // The following is a declared postcondition of the KeyGen method:
+    expect |sigKeyPair.verificationKey| == s.FieldSize();
+  }
+
   method YCompression(s: Signature.ECDSAParams, fieldSize: nat) {
     var res :- expect Signature.KeyGen(s);
+    RequireGoodKeyLengths(s, res);
     var public, secret := res.verificationKey, res.signingKey;
     // This is the declared postcondition of the natively implemented KenGen method, plus a condition
     // about zero-padding:
@@ -30,6 +36,7 @@ module TestSignature {
   method VerifyMessage(params: Signature.ECDSAParams) {
     var message :- expect UTF8.Encode("Hello, World!");
     var keys :- expect Signature.KeyGen(params);
+    RequireGoodKeyLengths(params, keys);
 
     var digest :- expect Signature.Digest(params, message);
     var signature :- expect Signature.Sign(params, keys.signingKey, digest);

test/SDK/Serialize.dfy

@@ -18,7 +18,7 @@ module TestSerialize {
     var keyB :- expect UTF8.Encode("keyB");
     var valB :- expect UTF8.Encode("valB");
     var encryptionContext := map[keyB := valB, keyA := valA];
-    MessageHeader.AssumeValidAAD(encryptionContext);
+    TestUtils.ValidSmallEncryptionContext(encryptionContext);
 
     var expectedSerializedAAD := [0, 26, 0, 2, 0, 4, 107, 101, 121, 65, 0, 4, 118, 97, 108, 65, 0, 4, 107, 101, 121, 66, 0, 4, 118, 97, 108, 66];
 
@@ -40,7 +40,6 @@ module TestSerialize {
   method {:test} TestSerializeLargeValidEC() {   
     var wr := new Streams.ByteWriter();
     var encCtx := TestUtils.GenerateLargeValidEncryptionContext();
-    MessageHeader.AssumeValidAAD(encCtx);
 
     var len :- expect SerializeAAD(wr, encCtx);
     expect len == 4 + |encCtx| as int * 7;
@@ -53,7 +52,7 @@ module TestSerialize {
     var keyB :- expect UTF8.Encode("keyB");
     var valB :- expect UTF8.Encode("valB");
     var encryptionContext := map[keyB := valB, keyA := valA];
-    MessageHeader.AssumeValidAAD(encryptionContext);
+    TestUtils.ValidSmallEncryptionContext(encryptionContext);
 
     var expectedSerializedAAD := [0, 2, 0, 4, 107, 101, 121, 65, 0, 4, 118, 97, 108, 65, 0, 4, 107, 101, 121, 66, 0, 4, 118, 97, 108, 66];
 

test/Util/TestUtils.dfy

@@ -2,12 +2,14 @@ include "../../src/StandardLibrary/StandardLibrary.dfy"
 include "../../src/StandardLibrary/UInt.dfy"
 include "../../src/Util/UTF8.dfy"
 include "../../src/SDK/Materials.dfy"
+include "../../src/SDK/MessageHeader.dfy"
 
 module {:extern "TestUtils"} TestUtils {
   import opened StandardLibrary
   import opened UInt = StandardLibrary.UInt
   import UTF8
   import Materials
+  import MessageHeader
 
   const SHARED_TEST_KEY_ARN := "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f";
 
@@ -21,20 +23,23 @@ module {:extern "TestUtils"} TestUtils {
   // Generates a large encryption context that approaches the upper bounds of
   // what is able to be serialized in the message format.
   // Building a map item by item is slow in dafny, so this method should be used sparingly.
-  method GenerateLargeValidEncryptionContext() returns (r: Materials.EncryptionContext) {
+  method GenerateLargeValidEncryptionContext() returns (r: Materials.EncryptionContext)
+    ensures MessageHeader.ValidAAD(r)
+  {
     // KVPairsMaxSize - KVPairsLenLen / KVPairLen ==> 
     // (2^16 - 1 - 2) / (2 + 2 + 2 + 1) ==> (2^16 - 3) / 7 ==> 9361
     // which is close to the max number of pairs you can stuff into a valid AAD.
     // We look for 9361 valid 2 byte UTF8 sequences (sticking to 2 bytes for simplicity).
     assert (0x1_0000 - 1 - 2) / (2 + 2 + 2 + 1) == (0x1_0000 - 3) / 7 == 9361;
     var numMaxPairs := 9361;
     var val :- expect UTF8.Encode("a");
-    var encCtx : map<UTF8.ValidUTF8Bytes, UTF8.ValidUTF8Bytes> := map[];
+    var encCtx: map<UTF8.ValidUTF8Bytes, UTF8.ValidUTF8Bytes> := map[];
 
     // Instead of proving termination for looking for Valid UTF8 sequences,
     // provide an upper bound to while loop
-    var i : int := 0;
+    var i := 0;
     while |encCtx| < numMaxPairs && i < 0x1_0000
+      invariant forall k :: k in encCtx ==> |k| + |encCtx[k]| == 3
       decreases 0x1_0000 - i
     {
       var key := UInt16ToSeq(i as uint16);
@@ -46,6 +51,36 @@ module {:extern "TestUtils"} TestUtils {
     // Check that we actually built a encCtx of the correct size
     expect |encCtx| == numMaxPairs;
 
+    assert MessageHeader.ValidAAD(encCtx) by {
+      reveal MessageHeader.ValidAAD();
+      assert MessageHeader.KVPairsLength(encCtx) < UINT16_LIMIT by {
+        var keys: seq<UTF8.ValidUTF8Bytes> := SetToOrderedSequence(encCtx.Keys, UInt.UInt8Less);
+        var kvPairs := seq(|keys|, i requires 0 <= i < |keys| => (keys[i], encCtx[keys[i]]));
+        KVPairsLengthBound(kvPairs, |kvPairs|, 3);
+        assert MessageHeader.KVPairEntriesLength(kvPairs, 0, |kvPairs|) <= 2 + numMaxPairs * 7;
+      }
+    }
     return encCtx;
   }
+
+  lemma ValidSmallEncryptionContext(encryptionContext: Materials.EncryptionContext)
+    requires |encryptionContext| <= 5
+    requires forall k :: k in encryptionContext.Keys ==> |k| < 100 && |encryptionContext[k]| < 100
+    ensures MessageHeader.ValidAAD(encryptionContext)
+  {
+    reveal MessageHeader.ValidAAD();
+    assert MessageHeader.KVPairsLength(encryptionContext) < UINT16_LIMIT by {
+      var keys: seq<UTF8.ValidUTF8Bytes> := SetToOrderedSequence(encryptionContext.Keys, UInt.UInt8Less);
+      var kvPairs := seq(|keys|, i requires 0 <= i < |keys| => (keys[i], encryptionContext[keys[i]]));
+      KVPairsLengthBound(kvPairs, |kvPairs|, 200);
+      assert MessageHeader.KVPairEntriesLength(kvPairs, 0, |kvPairs|) <= 5 * 204;
+    }
+  }
+
+  lemma KVPairsLengthBound(kvPairs: MessageHeader.EncryptionContextSeq, n: nat, kvBound: int)
+    requires n <= |kvPairs|
+    requires forall i :: 0 <= i < n ==> |kvPairs[i].0| + |kvPairs[i].1| <= kvBound
+    ensures MessageHeader.KVPairEntriesLength(kvPairs, 0, n) <= n * (4 + kvBound)
+  {
+  }
 }