refactor: Address PR Comments

- Remove ValidAAD precondition on DefaultCMM.GetEncryptionMaterials - Add runtime check of encryptionContext validity to DefaultCMM.GetEncryptionMaterials - Move MakeKeyring from DefaultCMMTests to TestUtils
aws · Mar 5, 2020 · 00d87fc · 00d87fc
1 parent 6208d65
commit 00d87fc
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 40 deletions.
diff --git a/src/SDK/CMM/DefaultCMM.dfy b/src/SDK/CMM/DefaultCMM.dfy
@@ -44,7 +44,6 @@ module {:extern "DefaultCMMDef"} DefaultCMMDef {
     method GetEncryptionMaterials(materialsRequest: Materials.EncryptionMaterialsRequest)
                                   returns (res: Result<Materials.ValidEncryptionMaterials>)
       requires Valid()
-      requires ValidAAD(materialsRequest.encryptionContext)
       ensures Valid()
       ensures Materials.EC_PUBLIC_KEY_FIELD in materialsRequest.encryptionContext ==> res.Failure?
       ensures res.Success? && (materialsRequest.algorithmSuiteID.None? || materialsRequest.algorithmSuiteID.get.SignatureType().Some?) ==>
@@ -86,20 +85,21 @@ module {:extern "DefaultCMMDef"} DefaultCMMDef {
           <  { assert |signatureKeys.verificationKey| <= 3000; }
             UINT16_LIMIT;
           }
-
           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();
-          }
       }
 
+      // Check validity of the encryption context at runtime.
+      var len := MessageHeader.ComputeKVPairsLength(enc_ctx);
+      var validAAD := MessageHeader.ComputeValidAAD(enc_ctx);
+      if UINT16_LIMIT <= |enc_ctx| {
+        return Failure("encryption context has too many entries");
+      } else if UINT16_LIMIT <= len {
+        return Failure("encryption context too big");
+      } else if !validAAD {
+        return Failure("encryption context has invalid key-value pairs");
+      }
+      assert MessageHeader.ValidAAD(enc_ctx);
+
       var materials := Materials.EncryptionMaterials.WithoutDataKeys(enc_ctx, id, enc_sk);
       assert materials.encryptionContext == enc_ctx;
       materials :- keyring.OnEncrypt(materials);

diff --git a/test/SDK/DefaultCMM.dfy b/test/SDK/DefaultCMM.dfy
@@ -1,31 +1,27 @@
-include "../../src/SDK/Keyring/RawRSAKeyring.dfy"
 include "../../src/SDK/Materials.dfy"
 include "../../src/SDK/AlgorithmSuite.dfy"
 include "../../src/StandardLibrary/StandardLibrary.dfy"
 include "../../src/StandardLibrary/UInt.dfy"
 include "../../src/SDK/CMM/Defs.dfy"
 include "../../src/SDK/CMM/DefaultCMM.dfy"
-include "../../src/Crypto/RSAEncryption.dfy"
 include "../../src/Util/UTF8.dfy"
 include "../../src/SDK/MessageHeader.dfy"
+include "../Util/TestUtils.dfy"
 
 module {:extern "DefaultCMMTests"} DefaultCMMTests {
   import opened StandardLibrary
   import opened UInt = StandardLibrary.UInt
   import AlgorithmSuite
   import DefaultCMMDef
-  import RSA = RSAEncryption
-  import RawRSAKeyringDef
   import Materials
   import MessageHeader
   import UTF8
+  import TestUtils
 
   method {:test} TestDefaultCMMNoAlg() returns (res: Result<()>) {
-    var keyring :- MakeKeyring();
+    var keyring :- TestUtils.MakeRSAKeyring();
     var cmm := new DefaultCMMDef.DefaultCMM.OfKeyring(keyring);
     var encCtx: Materials.EncryptionContext := map[];
-    reveal MessageHeader.ValidAAD();
-    assert MessageHeader.ValidAAD(encCtx);
     var getEncMatOutput :- cmm.GetEncryptionMaterials(Materials.EncryptionMaterialsRequest(encCtx, None, None));
 
     expect getEncMatOutput.algorithmSuiteID == 0x0378, "GetEncryptionMaterials returned unexpected algorithm id";
@@ -44,11 +40,9 @@ module {:extern "DefaultCMMTests"} DefaultCMMTests {
   }
 
   method {:test} TestDefaultCMMWithAlg() returns (res: Result<()>) {
-    var keyring :- MakeKeyring();
+    var keyring :- TestUtils.MakeRSAKeyring();
     var cmm := new DefaultCMMDef.DefaultCMM.OfKeyring(keyring);
     var encCtx: Materials.EncryptionContext := map[];
-    reveal MessageHeader.ValidAAD();
-    assert MessageHeader.ValidAAD(encCtx);
     var getEncMatOutput :- cmm.GetEncryptionMaterials(Materials.EncryptionMaterialsRequest(encCtx, Some(AlgorithmSuite.AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384), None));
 
     expect getEncMatOutput.algorithmSuiteID == AlgorithmSuite.AES_256_GCM_IV12_TAG16_HKDF_SHA384_ECDSA_P384,
@@ -68,12 +62,10 @@ module {:extern "DefaultCMMTests"} DefaultCMMTests {
   }
 
   method {:test} TestDefaultCMMWithAlgNoSig() returns (res: Result<()>) {
-    var keyring :- MakeKeyring();
+    var keyring :- TestUtils.MakeRSAKeyring();
     var cmm := new DefaultCMMDef.DefaultCMM.OfKeyring(keyring);
 
     var encCtx: Materials.EncryptionContext := map[];
-    reveal MessageHeader.ValidAAD();
-    assert MessageHeader.ValidAAD(encCtx);
     var getEncMatOutput :- cmm.GetEncryptionMaterials(Materials.EncryptionMaterialsRequest(encCtx, Some(AlgorithmSuite.AES_256_GCM_IV12_TAG16_HKDF_SHA256_NO_SIGNATURE_ALG), None));
 
     expect getEncMatOutput.algorithmSuiteID == AlgorithmSuite.AES_256_GCM_IV12_TAG16_HKDF_SHA256_NO_SIGNATURE_ALG,
@@ -93,26 +85,12 @@ module {:extern "DefaultCMMTests"} DefaultCMMTests {
   }
 
   method {:test} TestDefaultCMMRejectsBadEncCtx() returns (res: Result<()>) {
-    var keyring :- MakeKeyring();
+    var keyring :- TestUtils.MakeRSAKeyring();
     var cmm := new DefaultCMMDef.DefaultCMM.OfKeyring(keyring);
     var encCtx: Materials.EncryptionContext := map[];
     encCtx := encCtx[Materials.EC_PUBLIC_KEY_FIELD := [0x00]];
-    assert |encCtx| < UINT16_LIMIT;
-    assert MessageHeader.ValidKVPairs(encCtx);
-    var validAAD := MessageHeader.ComputeValidAAD(encCtx);
-    expect validAAD;
     var shouldBeFail := cmm.GetEncryptionMaterials(Materials.EncryptionMaterialsRequest(encCtx, None, None));
 
     expect shouldBeFail.Failure?, "GetEncryptionMaterials returned Success with bad input";
   }
-
-  method MakeKeyring() returns (res: Result<RawRSAKeyringDef.RawRSAKeyring>)
-    ensures res.Success? ==> res.value.Valid()
-  {
-    var namespace :- UTF8.Encode("namespace");
-    var name :- UTF8.Encode("MyKeyring");
-    var ek, dk := RSA.GenerateKeyPair(2048, RSA.PKCS1);
-    var keyring := new RawRSAKeyringDef.RawRSAKeyring(namespace, name, RSA.PaddingMode.PKCS1, Some(ek), Some(dk));
-    return Success(keyring);
-  }
 }
diff --git a/test/Util/TestUtils.dfy b/test/Util/TestUtils.dfy
@@ -1,3 +1,5 @@
+include "../../src/SDK/Keyring/RawRSAKeyring.dfy"
+include "../../src/Crypto/RSAEncryption.dfy"
 include "../../src/StandardLibrary/StandardLibrary.dfy"
 include "../../src/StandardLibrary/UInt.dfy"
 include "../../src/Util/UTF8.dfy"
@@ -10,6 +12,8 @@ module {:extern "TestUtils"} TestUtils {
   import UTF8
   import Materials
   import MessageHeader
+  import RSA = RSAEncryption
+  import RawRSAKeyringDef
 
   const SHARED_TEST_KEY_ARN := "arn:aws:kms:us-west-2:658956600833:key/b3537ef1-d8dc-4780-9f5a-55776cbb2f7f";
 
@@ -83,4 +87,14 @@ module {:extern "TestUtils"} TestUtils {
     ensures MessageHeader.KVPairEntriesLength(kvPairs, 0, n) <= n * (4 + kvBound)
   {
   }
+
+  method MakeRSAKeyring() returns (res: Result<RawRSAKeyringDef.RawRSAKeyring>)
+    ensures res.Success? ==> res.value.Valid()
+  {
+    var namespace :- UTF8.Encode("namespace");
+    var name :- UTF8.Encode("MyKeyring");
+    var ek, dk := RSA.GenerateKeyPair(2048, RSA.PKCS1);
+    var keyring := new RawRSAKeyringDef.RawRSAKeyring(namespace, name, RSA.PaddingMode.PKCS1, Some(ek), Some(dk));
+    return Success(keyring);
+  }
 }