From c36414b131dfd0a1ca51f10f87a18955bc110ff2 Mon Sep 17 00:00:00 2001
From: Nhi Pham <nhi@os.amperecomputing.com>
Date: Wed, 29 May 2024 09:46:25 +0700
Subject: [PATCH] MdeModulePkg/DxeCapsuleLibFmp: Fix crash if no ESRT is found

The ESRT table is not required in UEFI firmware. In such cases, the
table may not be present in the UEFI Configuration Table. The mEsrtTable
is to check if the IsNestedFmpCapsule() function is invoked at runtime
to determine whether to use gEsrtManagementProtocolGuid or the ESRT
table from the Configuration Table. Unfortunately, the check does not
cover situations where the ESRT is not present, potentially resulting in
a kernel crash. This patch is intended to fix this issue.

Signed-off-by: Nhi Pham <nhi@os.amperecomputing.com>
---
 .../Library/DxeCapsuleLibFmp/DxeCapsuleLib.c        | 13 +++++++++++++
 .../Library/DxeCapsuleLibFmp/DxeCapsuleRuntime.c    |  4 ++++
 2 files changed, 17 insertions(+)

diff --git a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
index 2433c76a8c..36efd64c4c 100644
--- a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
+++ b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
@@ -11,6 +11,7 @@
   performs basic validation.
 
   Copyright (c) 2016 - 2019, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2024, Ampere Computing LLC. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -48,6 +49,8 @@ EFI_EVENT  mDxeCapsuleLibEndOfDxeEvent = NULL;
 
 EDKII_FIRMWARE_MANAGEMENT_PROGRESS_PROTOCOL  *mFmpProgress = NULL;
 
+extern BOOLEAN  mDxeCapsuleLibReadyToBootEvent;
+
 /**
   Initialize capsule related variables.
 **/
@@ -1402,6 +1405,16 @@ IsNestedFmpCapsule (
       }
     }
   } else {
+    if (mDxeCapsuleLibReadyToBootEvent) {
+      //
+      // The ESRT table (mEsrtTable) in the Configuration Table would be located
+      // at the ReadyToBoot event if it exists. Hence, it should return here to
+      // avoid a crash due to calling gBS->LocateProtocol () at runtime in case
+      // there is no ERST table installed.
+      //
+      return FALSE;
+    }
+
     //
     // Check ESRT protocol
     //
diff --git a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleRuntime.c b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleRuntime.c
index 44f30c16c2..efff714d01 100644
--- a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleRuntime.c
+++ b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleRuntime.c
@@ -2,6 +2,7 @@
   Capsule library runtime support.
 
   Copyright (c) 2016 - 2017, Intel Corporation. All rights reserved.<BR>
+  Copyright (c) 2024, Ampere Computing LLC. All rights reserved.<BR>
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
 **/
@@ -23,6 +24,7 @@
 extern EFI_SYSTEM_RESOURCE_TABLE  *mEsrtTable;
 EFI_EVENT                         mDxeRuntimeCapsuleLibVirtualAddressChangeEvent = NULL;
 EFI_EVENT                         mDxeRuntimeCapsuleLibReadyToBootEvent          = NULL;
+BOOLEAN                           mDxeCapsuleLibReadyToBootEvent                 = FALSE;
 
 /**
   Convert EsrtTable physical address to virtual address.
@@ -93,6 +95,8 @@ DxeCapsuleLibReadyToBootEventNotify (
     //
     mEsrtTable->FwResourceCountMax = mEsrtTable->FwResourceCount;
   }
+
+  mDxeCapsuleLibReadyToBootEvent = TRUE;
 }
 
 /**