SignedCapsulePkg SystemFirmwareUpdateDxe: Fix failure caused by d69d922

[mirror_edk2.git] / SignedCapsulePkg / Universal / SystemFirmwareUpdate / SystemFirmwareUpdateDxe.c

diff --git a/SignedCapsulePkg/Universal/SystemFirmwareUpdate/SystemFirmwareUpdateDxe.c b/SignedCapsulePkg/Universal/SystemFirmwareUpdate/SystemFirmwareUpdateDxe.c
index d0b1c9913ca8d16d6ef322447346d39adbabf2d6..fa0c5f03ffdda0c624346e3260cdd6ef1735c97b 100644 (file)
--- a/SignedCapsulePkg/Universal/SystemFirmwareUpdate/SystemFirmwareUpdateDxe.c
+++ b/SignedCapsulePkg/Universal/SystemFirmwareUpdate/SystemFirmwareUpdateDxe.c
@@ -602,6 +602,7 @@ FindMatchingFmpHandles (
   )\r
 {\r
   EFI_STATUS                     Status;\r
+  UINTN                          TempHandleCount;\r
   EFI_HANDLE                     *HandleBuffer;\r
   UINTN                          Index;\r
   UINTN                          Index2;\r
@@ -613,20 +614,20 @@ FindMatchingFmpHandles (
   BOOLEAN                        MatchFound;\r
 \r
   *HandleCount  = 0;\r
+  TempHandleCount = 0;\r
   HandleBuffer = NULL;\r
   Status = gBS->LocateHandleBuffer (\r
                    ByProtocol,\r
                    ProtocolGuid,\r
                    NULL,\r
-                   HandleCount,\r
+                   &TempHandleCount,\r
                    &HandleBuffer\r
                    );\r
   if (EFI_ERROR (Status)) {\r
-    *HandleCount  = 0;\r
     return NULL;\r
   }\r
 \r
-  for (Index = 0; Index < *HandleCount; Index++) {\r
+  for (Index = 0; Index < TempHandleCount; Index++) {\r
     OriginalFmpImageInfoBuf = GetFmpImageDescriptors (\r
                                 HandleBuffer[Index],\r
                                 ProtocolGuid,\r
@@ -657,12 +658,21 @@ FindMatchingFmpHandles (
       //\r
       FmpImageInfoBuf = (EFI_FIRMWARE_IMAGE_DESCRIPTOR *)(((UINT8 *)FmpImageInfoBuf) + DescriptorSize);\r
     }\r
-    if (!MatchFound) {\r
-      HandleBuffer[Index] = NULL;\r
+    if (MatchFound) {\r
+      HandleBuffer[*HandleCount] = HandleBuffer[Index];\r
+      (*HandleCount)++;\r
     }\r
 \r
     FreePool (OriginalFmpImageInfoBuf);\r
   }\r
+\r
+  if ((*HandleCount) == 0) {\r
+    //\r
+    // No any matching handle.\r
+    //\r
+    FreePool (HandleBuffer);\r
+    return NULL;\r
+  }\r
   return HandleBuffer;\r
 }\r
 \r
@@ -801,6 +811,7 @@ SystemFirmwareUpdateMainDxe (
     // Install System FMP protocol onto handle with matching FMP Protocol\r
     //\r
     DEBUG ((DEBUG_INFO, "SystemFirmwareUpdateDxe: Install System FMP onto matching FMP handle\n"));\r
+    mSystemFmpPrivate->Handle = HandleBuffer[0];\r
     Status = gBS->InstallMultipleProtocolInterfaces (\r
                     &HandleBuffer[0],\r
                     &gSystemFmpProtocolGuid,\r