From 8978086615efe0a59fcaa3782fa46c6f11bc65b8 Mon Sep 17 00:00:00 2001
From: Nishanth1311 <41198847+Nishanth1311@users.noreply.github.com>
Date: Thu, 7 Dec 2023 04:44:01 +0530
Subject: [PATCH] Fix invalid access of NewRecoveryChallenge when memory
 allocation fails (#168)

## Description

If memory allocation for the NewChallenge variable should not be used.
It will result in Invalid access at the following line.
  NewChallenge->SerialNumber = 0;

With this change if memory allocation for the NewChallenge variable
fails, EFI_OUT_OF_RESOURCES is returned

- [ ] Impacts functionality? No
- [ ] Impacts security? No
- [ ] Breaking change? No
- [ ] Includes tests? No
- [ ] Includes documentation? No

## How This Was Tested
Build passes with this change

## Integration Instructions
N/A
---
 DfciPkg/Library/DfciRecoveryLib/DfciRecoveryLib.c | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/DfciPkg/Library/DfciRecoveryLib/DfciRecoveryLib.c b/DfciPkg/Library/DfciRecoveryLib/DfciRecoveryLib.c
index 65080927..7dcd0890 100644
--- a/DfciPkg/Library/DfciRecoveryLib/DfciRecoveryLib.c
+++ b/DfciPkg/Library/DfciRecoveryLib/DfciRecoveryLib.c
@@ -66,7 +66,10 @@ GetRecoveryChallenge (
   //
   // Locate the RNG Protocol. This will be needed for the nonce.
   Status = gBS->LocateProtocol (&gEfiRngProtocolGuid, NULL, (VOID **)&RngProtocol);
-  DEBUG ((DEBUG_VERBOSE, "%a: LocateProtocol(RNG) = %r\n", __FUNCTION__, Status));
+  if (EFI_ERROR (Status)) {
+    DEBUG ((DEBUG_ERROR, "%a: LocateProtocol(RNG) = %r\n", __FUNCTION__, Status));
+    return EFI_NOT_FOUND;
+  }
 
   //
   // From now on, don't proceed on errors.
@@ -74,11 +77,11 @@ GetRecoveryChallenge (
 
   //
   // Allocate the buffer...
-  if (!EFI_ERROR (Status)) {
-    NewChallenge = AllocatePool (sizeof (DFCI_RECOVERY_CHALLENGE) + DFCI_MULTI_STRING_MAX_SIZE);
-    if (NewChallenge == NULL) {
-      Status = EFI_OUT_OF_RESOURCES;
-    }
+  NewChallenge = AllocatePool (sizeof (DFCI_RECOVERY_CHALLENGE) + DFCI_MULTI_STRING_MAX_SIZE);
+
+  // Exit if we ran out of resources
+  if (NewChallenge == NULL) {
+    return EFI_OUT_OF_RESOURCES;
   }
 
   //