From 258f7de25ea5934ad456c3e8eb355309758d288b Mon Sep 17 00:00:00 2001
From: "Kinney, Michael D" <michael.d.kinney@intel.com>
Date: Tue, 6 Feb 2018 17:43:32 -0800
Subject: [PATCH] MdeModulePkg/DxeCapsuleLibFmp: Verify nested capsule with FMP

https://bugzilla.tianocore.org/show_bug.cgi?id=873

Update IsNestedFmpCapsule() to verify the CapsuleGuid in
the CapsuleHeader against the installed Firmware Management
Protocol instances.  The current logic that uses the ESRT
Table does not work because capsules are processed before
the ESRT Table is published at the Ready To Boot event.

Cc: Jiewen Yao <jiewen.yao@intel.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: Eric Dong <eric.dong@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.1
Signed-off-by: Michael D Kinney <michael.d.kinney@intel.com>
Reviewed-by: Jiewen Yao <jiewen.yao@intel.com>
Reviewed-by: Star Zeng <star.zeng@intel.com>
Reviewed-by: Bret Barkelew <Bret.Barkelew@microsoft.com>
(cherry picked from commit 7f0301e39a11852c8c26b6a700f0a3a0c381e95c)
---
 .../Library/DxeCapsuleLibFmp/DxeCapsuleLib.c  | 25 +++++++++++--------
 .../DxeCapsuleLibFmp/DxeCapsuleLib.inf        |  1 -
 2 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
index 2f397789b5..8f4e6b8951 100644
--- a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
+++ b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
@@ -1446,7 +1446,6 @@ IsNestedFmpCapsule (
   )
 {
   EFI_STATUS                 Status;
-  EFI_SYSTEM_RESOURCE_TABLE  *Esrt;
   EFI_SYSTEM_RESOURCE_ENTRY  *EsrtEntry;
   UINTN                      Index;
   BOOLEAN                    EsrtGuidFound;
@@ -1454,6 +1453,8 @@ IsNestedFmpCapsule (
   UINTN                      NestedCapsuleSize;
   ESRT_MANAGEMENT_PROTOCOL   *EsrtProtocol;
   EFI_SYSTEM_RESOURCE_ENTRY  Entry;
+  EFI_HANDLE                 *HandleBuffer;
+  UINTN                      NumberOfHandles;
 
   EsrtGuidFound = FALSE;
   if (mIsVirtualAddrConverted) {
@@ -1479,19 +1480,21 @@ IsNestedFmpCapsule (
     }
 
     //
-    // Check ESRT configuration table
+    // Check Firmware Management Protocols
     //
     if (!EsrtGuidFound) {
-      Status = EfiGetSystemConfigurationTable(&gEfiSystemResourceTableGuid, (VOID **)&Esrt);
+      HandleBuffer = NULL;
+      Status = GetFmpHandleBufferByType (
+                 &CapsuleHeader->CapsuleGuid,
+                 0,
+                 &NumberOfHandles,
+                 &HandleBuffer
+                 );
       if (!EFI_ERROR(Status)) {
-        ASSERT (Esrt != NULL);
-        EsrtEntry = (VOID *)(Esrt + 1);
-        for (Index = 0; Index < Esrt->FwResourceCount; Index++, EsrtEntry++) {
-          if (CompareGuid(&EsrtEntry->FwClass, &CapsuleHeader->CapsuleGuid)) {
-            EsrtGuidFound = TRUE;
-            break;
-          }
-        }
+        EsrtGuidFound = TRUE;
+      }
+      if (HandleBuffer != NULL) {
+        FreePool (HandleBuffer);
       }
     }
   }
diff --git a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.inf b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.inf
index a6cf54cb6b..f24d21af36 100644
--- a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.inf
+++ b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.inf
@@ -72,7 +72,6 @@
 [Guids]
   gEfiFmpCapsuleGuid                      ## SOMETIMES_CONSUMES ## GUID
   gWindowsUxCapsuleGuid                   ## SOMETIMES_CONSUMES ## GUID
-  gEfiSystemResourceTableGuid             ## SOMETIMES_CONSUMES ## GUID
   gEfiCapsuleReportGuid                   ## CONSUMES ## Variable
   gEfiCapsuleVendorGuid                   ## CONSUMES ## Variable
   gEfiEndOfDxeEventGroupGuid              ## CONSUMES ## Event