Catch more allocation failure cases

Also move ReseedRng() to use the stack rather than the heap.

Bug: 216213939
Bug: 215451239
Change-Id: I62ee8a5853acdaf49ff1908c625ac988d5b209c8

3 files changed, 23 insertions, 6 deletions

diff --git a/ipc/keymaster_ipc.cpp b/ipc/keymaster_ipc.cpp
index 0d4051b..789e1e8 100644
--- a/ipc/keymaster_ipc.cpp
+++ b/ipc/keymaster_ipc.cpp
@@ -678,6 +678,9 @@ static long handle_msg(keymaster_chan_ctx* ctx) {
     // allocate msg_buf, with one extra byte for null-terminator
     keymaster::UniquePtr<uint8_t[]> msg_buf(new (std::nothrow)
                                                     uint8_t[msg_inf.len + 1]);
+    if (msg_buf.get() == nullptr) {
+        return ERR_NO_MEMORY;
+    }
     msg_buf[msg_inf.len] = 0;
 
     /* read msg content */
diff --git a/secure_storage_manager.cpp b/secure_storage_manager.cpp
index 5403831..fa3d2a8 100644
--- a/secure_storage_manager.cpp
+++ b/secure_storage_manager.cpp
@@ -348,6 +348,9 @@ keymaster_error_t SecureStorageManager::WriteAtapKeyAndCertsToStorage(
     }
     UniquePtr<AttestationKey> attestation_key(
             new (std::nothrow) AttestationKey(AttestationKey_init_zero));
+    if (attestation_key.get() == nullptr) {
+        return KM_ERROR_MEMORY_ALLOCATION_FAILED;
+    }
     attestation_key->has_key = true;
     attestation_key->key.size = key_size;
     memcpy(attestation_key->key.bytes, key, key_size);
@@ -519,6 +522,9 @@ keymaster_error_t SecureStorageManager::SetAttestationIds(
         const SetAttestationIdsRequest& request) {
     AttestationIds* attestation_ids_p =
             new (std::nothrow) AttestationIds(AttestationIds_init_zero);
+    if (attestation_ids_p == nullptr) {
+        return KM_ERROR_MEMORY_ALLOCATION_FAILED;
+    }
     UniquePtr<AttestationIds> attestation_ids(attestation_ids_p);
     if (request.brand.buffer_size() > kAttestationIdLengthMax) {
         LOG_E("Error: Brand ID too large: %d", request.brand.buffer_size());
@@ -817,6 +823,13 @@ keymaster_error_t SecureStorageManager::TranslateLegacyFormat() {
         // New attribute file exists, nothing to do.
         return KM_ERROR_OK;
     }
+
+    UniquePtr<KeymasterAttributes> km_attributes(new (
+            std::nothrow) KeymasterAttributes(KeymasterAttributes_init_zero));
+    if (km_attributes.get() == nullptr) {
+        return KM_ERROR_MEMORY_ALLOCATION_FAILED;
+    }
+
     AttestationKeySlot key_slots[] = {
             AttestationKeySlot::kRsa,        AttestationKeySlot::kEcdsa,
             AttestationKeySlot::kEddsa,      AttestationKeySlot::kEpid,
@@ -832,6 +845,9 @@ keymaster_error_t SecureStorageManager::TranslateLegacyFormat() {
         AttestationKeySlot key_slot = key_slots[i];
         UniquePtr<AttestationKey> attestation_key(
                 new (std::nothrow) AttestationKey(AttestationKey_init_zero));
+        if (attestation_key.get() == nullptr) {
+            return KM_ERROR_MEMORY_ALLOCATION_FAILED;
+        }
         snprintf(key_file, kStorageIdLengthMax, "%s.%s", kLegacyAttestKeyPrefix,
                  GetKeySlotStr(key_slot));
         err = LegacySecureStorageRead(key_file, attestation_key->key.bytes,
@@ -881,8 +897,6 @@ keymaster_error_t SecureStorageManager::TranslateLegacyFormat() {
         }
     }
 
-    UniquePtr<KeymasterAttributes> km_attributes(new (
-            std::nothrow) KeymasterAttributes(KeymasterAttributes_init_zero));
     uint32_t product_id_size;
     err = LegacySecureStorageRead(kLegacyProductIdFileName,
                                   km_attributes->product_id.bytes,
diff --git a/trusty_keymaster_context.cpp b/trusty_keymaster_context.cpp
index c5ecd19..f0643c4 100644
--- a/trusty_keymaster_context.cpp
+++ b/trusty_keymaster_context.cpp
@@ -772,14 +772,14 @@ bool TrustyKeymasterContext::ShouldReseedRng() const {
 }
 
 bool TrustyKeymasterContext::ReseedRng() {
-    UniquePtr<uint8_t[]> rand_seed(new (std::nothrow) uint8_t[kRngReseedSize]);
-    memset(rand_seed.get(), 0, kRngReseedSize);
-    if (trusty_rng_hw_rand(rand_seed.get(), kRngReseedSize) != 0) {
+    uint8_t rand_seed[kRngReseedSize];
+    memset(rand_seed, 0, kRngReseedSize);
+    if (trusty_rng_hw_rand(rand_seed, kRngReseedSize) != 0) {
         LOG_E("Failed to get bytes from HW RNG", 0);
         return false;
     }
     LOG_I("Reseeding with %d bytes from HW RNG", kRngReseedSize);
-    trusty_rng_add_entropy(rand_seed.get(), kRngReseedSize);
+    trusty_rng_add_entropy(rand_seed, kRngReseedSize);
 
     rng_initialized_ = true;
     return true;