Catch more allocation failure cases

Bug: 216213939
Bug: 215451239
Test: VtsAidlKeyMintTargetTest
Ignore-AOSP-First: Cherry pick from AOSP
Change-Id: I604c980dba7644e4012a358b6f2dc204cb215c3c
(cherry picked from commit 01d2e13c563a82e5cb0dc1adacd0ca5b41322497)

3 files changed, 15 insertions, 0 deletions

diff --git a/km_openssl/ec_key_factory.cpp b/km_openssl/ec_key_factory.cpp
index 2e0e2f1..4ec3175 100644
--- a/km_openssl/ec_key_factory.cpp
+++ b/km_openssl/ec_key_factory.cpp
@@ -234,12 +234,18 @@ keymaster_error_t EcKeyFactory::ImportKey(const AuthorizationSet& key_descriptio
     switch (EVP_PKEY_type(pkey->type)) {
     case EVP_PKEY_ED25519:
         key.reset(new (std::nothrow) Ed25519Key(*hw_enforced, *sw_enforced, this));
+        if (key.get() == nullptr) {
+            return KM_ERROR_MEMORY_ALLOCATION_FAILED;
+        }
         if (!key->EvpToInternal(pkey.get())) {
             return KM_ERROR_UNSUPPORTED_KEY_FORMAT;
         }
         break;
     case EVP_PKEY_X25519:
         key.reset(new (std::nothrow) X25519Key(*hw_enforced, *sw_enforced, this));
+        if (key.get() == nullptr) {
+            return KM_ERROR_MEMORY_ALLOCATION_FAILED;
+        }
         if (!key->EvpToInternal(pkey.get())) {
             return KM_ERROR_UNSUPPORTED_KEY_FORMAT;
         }
@@ -249,6 +255,9 @@ keymaster_error_t EcKeyFactory::ImportKey(const AuthorizationSet& key_descriptio
         if (!ec_key.get()) return KM_ERROR_INVALID_ARGUMENT;
 
         key.reset(new (std::nothrow) EcKey(*hw_enforced, *sw_enforced, this, move(ec_key)));
+        if (key.get() == nullptr) {
+            return KM_ERROR_MEMORY_ALLOCATION_FAILED;
+        }
         break;
     }
     default:
diff --git a/ng/AndroidRemotelyProvisionedComponentDevice.cpp b/ng/AndroidRemotelyProvisionedComponentDevice.cpp
index 0c2d842..54ea70c 100644
--- a/ng/AndroidRemotelyProvisionedComponentDevice.cpp
+++ b/ng/AndroidRemotelyProvisionedComponentDevice.cpp
@@ -113,6 +113,9 @@ ScopedAStatus AndroidRemotelyProvisionedComponentDevice::generateCertificateRequ
     request.test_mode = testMode;
     request.num_keys = keysToSign.size();
     request.keys_to_sign_array = new (std::nothrow) KeymasterBlob[keysToSign.size()];
+    if (request.keys_to_sign_array == nullptr) {
+        return km_utils::kmError2ScopedAStatus(KM_ERROR_MEMORY_ALLOCATION_FAILED);
+    }
     for (size_t i = 0; i < keysToSign.size(); i++) {
         request.SetKeyToSign(i, keysToSign[i].macedKey.data(), keysToSign[i].macedKey.size());
     }
diff --git a/ng/AndroidSharedSecret.cpp b/ng/AndroidSharedSecret.cpp
index b2cb85d..c800666 100644
--- a/ng/AndroidSharedSecret.cpp
+++ b/ng/AndroidSharedSecret.cpp
@@ -46,6 +46,9 @@ ScopedAStatus AndroidSharedSecret::computeSharedSecret(const vector<SharedSecret
     ComputeSharedHmacRequest request(impl_->message_version());
     request.params_array.params_array =
         new (std::nothrow) keymaster::HmacSharingParameters[params.size()];
+    if (request.params_array.params_array == nullptr) {
+        return kmError2ScopedAStatus(KM_ERROR_MEMORY_ALLOCATION_FAILED);
+    }
     request.params_array.num_params = params.size();
     for (size_t i = 0; i < params.size(); ++i) {
         request.params_array.params_array[i].seed = {params[i].seed.data(), params[i].seed.size()};