From a6a99c416f89d2350ea3e0b3f3b497ef38c561c9 Mon Sep 17 00:00:00 2001
From: Rahul Frias <rfrias@google.com>
Date: Mon, 21 Mar 2022 14:37:58 -0700
Subject: [PATCH 1/3] Free EC key in test session destructor

[ Merged from http://go/wvgerrit/148030 ]

Bug: 224375138
Test: GtsMediaDrmTests
Change-Id: Ic1d42756cb1a03fa52d90c5fb0bd496e87a43c65
---
 libwvdrmengine/oemcrypto/test/oec_session_util.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/libwvdrmengine/oemcrypto/test/oec_session_util.cpp b/libwvdrmengine/oemcrypto/test/oec_session_util.cpp
index a323e8a6..a01c1fd3 100644
--- a/libwvdrmengine/oemcrypto/test/oec_session_util.cpp
+++ b/libwvdrmengine/oemcrypto/test/oec_session_util.cpp
@@ -1407,6 +1407,7 @@ Session::Session()
 Session::~Session() {
   if (!forced_session_id_ && open_) close();
   if (public_rsa_) RSA_free(public_rsa_);
+  if (public_ec_) EC_KEY_free(public_ec_);
 }
 
 void Session::open() {

From b07523f4c6323d9c9a12e3e6b6dd70dcde22d028 Mon Sep 17 00:00:00 2001
From: Rahul Frias <rfrias@google.com>
Date: Mon, 21 Mar 2022 14:41:01 -0700
Subject: [PATCH 2/3] Constrain the generated key type in testing

[ Merge of http://go/wvgerrit/148157 ]

Fail the test if the returned type is other values.

Bug: 224375138
Test: GtsMediaDrmTests
Change-Id: I4abad9d69865cac99654d3dedd443463dd728a58
---
 libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp b/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp
index cb27efc5..0141bb64 100644
--- a/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp
+++ b/libwvdrmengine/oemcrypto/test/oemcrypto_test.cpp
@@ -1340,6 +1340,8 @@ TEST_F(OEMCryptoProv40Test, GenerateCertificateKeyPairSuccess) {
   } else if (key_type == OEMCrypto_PrivateKeyType::OEMCrypto_ECC_Private_Key) {
     ASSERT_NO_FATAL_FAILURE(
         s.SetEcPublicKey(public_key.data(), public_key_size));
+  } else {
+    FAIL() << "Unknown private key type: " << static_cast<int>(key_type);
   }
 }
 

From 587516dfa27fbed9dd27009414f2b0aff43509a0 Mon Sep 17 00:00:00 2001
From: Rahul Frias <rfrias@google.com>
Date: Mon, 21 Mar 2022 14:43:54 -0700
Subject: [PATCH 3/3] Return error first if the returned status is not ok

[ Merge of http://go/wvgerrit/148249 ]

Bug: 224375138
Test: GtsMediaDrmTests
Change-Id: I10f0c1bc28342d6cd5ae7d373ef554321bb6d3c8
---
 .../cdm/core/src/crypto_session.cpp           | 23 +++++++++++++------
 1 file changed, 16 insertions(+), 7 deletions(-)

diff --git a/libwvdrmengine/cdm/core/src/crypto_session.cpp b/libwvdrmengine/cdm/core/src/crypto_session.cpp
index 6cefeaa3..f8e11ad2 100644
--- a/libwvdrmengine/cdm/core/src/crypto_session.cpp
+++ b/libwvdrmengine/cdm/core/src/crypto_session.cpp
@@ -1447,8 +1447,13 @@ CdmResponseType CryptoSession::GetBootCertificateChain(
           &additional_signature_length);
     });
   }
-  return MapOEMCryptoResult(sts, GET_BOOT_CERTIFICATE_CHAIN_ERROR,
-                            "GetBootCertificateChain");
+  if (sts != OEMCrypto_SUCCESS) {
+    return MapOEMCryptoResult(sts, GET_BOOT_CERTIFICATE_CHAIN_ERROR,
+                              "GetBootCertificateChain");
+  }
+  bcc->resize(bcc_length);
+  additional_signature->resize(additional_signature_length);
+  return NO_ERROR;
 }
 
 CdmResponseType CryptoSession::GenerateCertificateKeyPair(
@@ -1494,6 +1499,11 @@ CdmResponseType CryptoSession::GenerateCertificateKeyPair(
             &wrapped_private_key_length, &oemcrypto_key_type),
         metrics_, oemcrypto_generate_certificate_key_pair_, status);
   });
+  if (status != OEMCrypto_SUCCESS) {
+    return MapOEMCryptoResult(status, GENERATE_CERTIFICATE_KEY_PAIR_ERROR,
+                              "GenerateCertificateKeyPair");
+  }
+
   public_key->resize(public_key_length);
   public_key_signature->resize(public_key_signature_length);
   wrapped_private_key->resize(wrapped_private_key_length);
@@ -1505,13 +1515,10 @@ CdmResponseType CryptoSession::GenerateCertificateKeyPair(
   } else {
     LOGE("Unexpected key type returned from GenerateCertificateKeyPair: %d",
          static_cast<int>(oemcrypto_key_type));
-    return MapOEMCryptoResult(status,
-                              GENERATE_CERTIFICATE_KEY_PAIR_UNKNOWN_TYPE_ERROR,
-                              "GenerateCertificateKeyPair");
+    return GENERATE_CERTIFICATE_KEY_PAIR_UNKNOWN_TYPE_ERROR;
   }
 
-  return MapOEMCryptoResult(status, GENERATE_CERTIFICATE_KEY_PAIR_ERROR,
-                            "GenerateCertificateKeyPair");
+  return NO_ERROR;
 }
 
 CdmResponseType CryptoSession::LoadOemCertificatePrivateKey(
@@ -1539,7 +1546,9 @@ CdmResponseType CryptoSession::LoadOemCertificatePrivateKey(
 // Private.
 CdmResponseType CryptoSession::SelectKey(const std::string& key_id,
                                          CdmCipherMode cipher_mode) {
+  RETURN_IF_NOT_OPEN(CRYPTO_SESSION_NOT_OPEN);
   const OEMCryptoResult sts = WithOecSessionLock("SelectKey", [&] {
+    RETURN_IF_NULL(key_session_, OEMCrypto_ERROR_INVALID_SESSION);
     return key_session_->SelectKey(key_id, cipher_mode);
   });