[PATCH] MdeModulePkg/CapsuleLib: Fix runtime issue

[Jiewen Yao <jiewen.yao@intel.com>]

The previous patch has problem on handling capsule
request at runtime.

This patch fixed the issue.

Cc: Chao Zhang <chao.b.zhang@intel.com>
Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Jiewen Yao <jiewen.yao@intel.com>
---
 MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c | 53 ++++++++++++--------
 1 file changed, 32 insertions(+), 21 deletions(-)

diff --git a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
index 6344214..9ed0be3 100644
--- a/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
+++ b/MdeModulePkg/Library/DxeCapsuleLibFmp/DxeCapsuleLib.c
@@ -1453,33 +1453,44 @@ IsNestedFmpCapsule (
   EFI_SYSTEM_RESOURCE_ENTRY  Entry;
 
   EsrtGuidFound = FALSE;
-
-  //
-  // Check ESRT protocol
-  //
-  Status = gBS->LocateProtocol(&gEsrtManagementProtocolGuid, NULL, (VOID **)&EsrtProtocol);
-  if (!EFI_ERROR(Status)) {
-    Status = EsrtProtocol->GetEsrtEntry(&CapsuleHeader->CapsuleGuid, &Entry);
-    if (!EFI_ERROR(Status)) {
-      EsrtGuidFound = TRUE;
-    }
-  }
-
-  //
-  // Check ESRT configuration table
-  //
-  if (!EsrtGuidFound) {
-    Status = EfiGetSystemConfigurationTable(&gEfiSystemResourceTableGuid, (VOID **)&Esrt);
-    if (!EFI_ERROR(Status)) {
-      ASSERT (Esrt != NULL);
-      EsrtEntry = (VOID *)(Esrt + 1);
-      for (Index = 0; Index < Esrt->FwResourceCount; Index++, EsrtEntry++) {
+  if (mIsVirtualAddrConverted) {
+    if(mEsrtTable != NULL) {
+      EsrtEntry = (EFI_SYSTEM_RESOURCE_ENTRY *)(mEsrtTable + 1);
+      for (Index = 0; Index < mEsrtTable->FwResourceCount ; Index++, EsrtEntry++) {
         if (CompareGuid(&EsrtEntry->FwClass, &CapsuleHeader->CapsuleGuid)) {
           EsrtGuidFound = TRUE;
           break;
         }
       }
     }
+  } else {
+    //
+    // Check ESRT protocol
+    //
+    Status = gBS->LocateProtocol(&gEsrtManagementProtocolGuid, NULL, (VOID **)&EsrtProtocol);
+    if (!EFI_ERROR(Status)) {
+      Status = EsrtProtocol->GetEsrtEntry(&CapsuleHeader->CapsuleGuid, &Entry);
+      if (!EFI_ERROR(Status)) {
+        EsrtGuidFound = TRUE;
+      }
+    }
+
+    //
+    // Check ESRT configuration table
+    //
+    if (!EsrtGuidFound) {
+      Status = EfiGetSystemConfigurationTable(&gEfiSystemResourceTableGuid, (VOID **)&Esrt);
+      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;
+          }
+        }
+      }
+    }
   }
   if (!EsrtGuidFound) {
     return FALSE;
-- 
2.7.4.windows.1