From 66f30ed29f851db914d31d132df80a96c9182732 Mon Sep 17 00:00:00 2001
From: Aaron Tomb <aarotomb@amazon.com>
Date: Fri, 23 Aug 2024 13:19:25 -0700
Subject: [PATCH] chore: remove case-based if statement (#630)

Remove a case-based if statement, as it caused errors about non-determinism.

Reduce resource use to fix a verification failure along the way.
---
 .../src/Keyrings/RawECDHKeyring.dfy             | 17 ++++++++++-------
 AwsCryptographyPrimitives/test/TestECDH.dfy     |  2 +-
 2 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/RawECDHKeyring.dfy b/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/RawECDHKeyring.dfy
index fdaede43f..af11bcd01 100644
--- a/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/RawECDHKeyring.dfy
+++ b/AwsCryptographicMaterialProviders/dafny/AwsCryptographicMaterialProviders/src/Keyrings/RawECDHKeyring.dfy
@@ -603,21 +603,24 @@ module {:options "/functionSyntax:4" } RawECDHKeyring {
 
       var sharedSecretPublicKey: seq<uint8>;
       var sharedSecretPrivateKey: seq<uint8>;
-      if {
-        case this.keyAgreementScheme.PublicKeyDiscovery? =>
+      match this.keyAgreementScheme {
+        case PublicKeyDiscovery(publicKeyDiscovery) => {
           sharedSecretPublicKey := senderPublicKey;
-          sharedSecretPrivateKey := this.keyAgreementScheme.PublicKeyDiscovery.recipientStaticPrivateKey;
-        case this.keyAgreementScheme.RawPrivateKeyToStaticPublicKey? =>
-          sharedSecretPrivateKey := this.keyAgreementScheme.RawPrivateKeyToStaticPublicKey.senderStaticPrivateKey;
-          if this.keyAgreementScheme.RawPrivateKeyToStaticPublicKey.recipientPublicKey == recipientPublicKey {
+          sharedSecretPrivateKey := publicKeyDiscovery.recipientStaticPrivateKey;
+        }
+        case RawPrivateKeyToStaticPublicKey(rawPrivateKeyToStaticPublicKey) => {
+          sharedSecretPrivateKey := rawPrivateKeyToStaticPublicKey.senderStaticPrivateKey;
+          if rawPrivateKeyToStaticPublicKey.recipientPublicKey == recipientPublicKey {
             sharedSecretPublicKey := recipientPublicKey;
           } else {
             sharedSecretPublicKey := senderPublicKey;
           }
-        case this.keyAgreementScheme.EphemeralPrivateKeyToStaticPublicKey? =>
+        }
+        case EphemeralPrivateKeyToStaticPublicKey(_) => {
           //= aws-encryption-sdk-specification/framework/key-agreement-schemas.md#ephemeralprivatekeytostaticpublickey
           //# On decrypt, the keyring MUST fail.
           return Failure(E("EphemeralPrivateKeyToStaticPublicKey Not allowed on decrypt"));
+        }
       }
 
       var _ :- ValidatePublicKey(
diff --git a/AwsCryptographyPrimitives/test/TestECDH.dfy b/AwsCryptographyPrimitives/test/TestECDH.dfy
index f3935caf9..1f13b4549 100644
--- a/AwsCryptographyPrimitives/test/TestECDH.dfy
+++ b/AwsCryptographyPrimitives/test/TestECDH.dfy
@@ -466,7 +466,7 @@ module TestECDH {
     for i := 0 to |curves|
     {
       var curve := curves[i];
-      var originalPublicKey := derX509PublicKeys[i];
+      var originalPublicKey := expectLooseHexString(derX509PublicKeys[i]);
       var publicKeyBytes := HexStrings.FromHexString(originalPublicKey);
       var compressedKey := expectLooseHexString(compressedKeys[i]);
       var compressedKeyBytes := HexStrings.FromHexString(compressedKey);