3 files changed, 19 insertions, 12 deletions
diff --git a/EdkUnixPkg/Dxe/PlatformBds/Generic/BootMaint/BmLib.c b/EdkUnixPkg/Dxe/PlatformBds/Generic/BootMaint/BmLib.c
index 415321efd7..e9245d2fb6 100644
--- a/EdkUnixPkg/Dxe/PlatformBds/Generic/BootMaint/BmLib.c
+++ b/EdkUnixPkg/Dxe/PlatformBds/Generic/BootMaint/BmLib.c
@@ -464,8 +464,12 @@ Returns:
   EFI_GUID          HiiGuid;
   EFI_HII_PROTOCOL  *Hii;
 
-  HandleBufferLength  = 0x1000;
+  //
+  // Initialize params.
+  //
+  HandleBufferLength  = 0;
   HiiHandleBuffer     = NULL;
+  
   Status = gBS->LocateProtocol (
                   &gEfiHiiProtocolGuid,
                   NULL,
@@ -478,12 +482,9 @@ Returns:
   //
   // Get all the Hii handles
   //
-  HiiHandleBuffer = AllocateZeroPool (HandleBufferLength);
-  ASSERT (HiiHandleBuffer != NULL);
-
-  Status = Hii->FindHandles (Hii, &HandleBufferLength, HiiHandleBuffer);
+  Status = BdsLibGetHiiHandles (Hii, &HandleBufferLength, &HiiHandleBuffer);
   ASSERT_EFI_ERROR (Status);
-
+  
   //
   // Get the Hii Handle that matches the StructureNode->ProducerName
   //
diff --git a/EdkUnixPkg/Dxe/PlatformBds/Generic/DeviceMngr/DeviceManager.c b/EdkUnixPkg/Dxe/PlatformBds/Generic/DeviceMngr/DeviceManager.c
index 3f648b7de6..02be71e789 100644
--- a/EdkUnixPkg/Dxe/PlatformBds/Generic/DeviceMngr/DeviceManager.c
+++ b/EdkUnixPkg/Dxe/PlatformBds/Generic/DeviceMngr/DeviceManager.c
@@ -201,6 +201,7 @@ Returns:
 
   IfrOptionList       = NULL;
   VideoOption         = NULL;
+  HiiHandles          = NULL;
   HandleBufferLength  = 0;
 
   //
@@ -268,8 +269,11 @@ Returns:
   CreateSubTitleOpCode (STR_EMPTY_STRING, &UpdateData->Data);
   Hii->UpdateForm (Hii, FPCallbackInfo.DevMgrHiiHandle, (EFI_FORM_LABEL) Count, TRUE, UpdateData);
 
-  HiiHandles = AllocateZeroPool (HandleBufferLength);
-  Hii->FindHandles (Hii, &HandleBufferLength, HiiHandles);
+  //
+  // Get all the Hii handles
+  //
+  Status = BdsLibGetHiiHandles (Hii, &HandleBufferLength, &HiiHandles);
+  ASSERT_EFI_ERROR (Status);
 
   for (Index = 1, BufferSize = 0; Index < HandleBufferLength; Index++) {
     //
@@ -487,6 +491,7 @@ Returns:
   }
 
   gBS->FreePool (UpdateData);
+  gBS->FreePool (HiiHandles);
 
   return Status;
 }
diff --git a/EdkUnixPkg/Dxe/PlatformBds/Generic/FrontPage.c b/EdkUnixPkg/Dxe/PlatformBds/Generic/FrontPage.c
index 6a3a34e2c9..56292f9e47 100644
--- a/EdkUnixPkg/Dxe/PlatformBds/Generic/FrontPage.c
+++ b/EdkUnixPkg/Dxe/PlatformBds/Generic/FrontPage.c
@@ -485,15 +485,16 @@ Returns:
   UINT16          Length;
   EFI_GUID        HiiGuid;
 
-  HandleBufferLength  = 0x1000;
+  //
+  // Initialize params.
+  //
+  HandleBufferLength  = 0;
   HiiHandleBuffer     = NULL;
 
   //
   // Get all the Hii handles
   //
-  HiiHandleBuffer = AllocateZeroPool (HandleBufferLength);
-
-  Status          = Hii->FindHandles (Hii, &HandleBufferLength, HiiHandleBuffer);
+  Status = BdsLibGetHiiHandles (Hii, &HandleBufferLength, &HiiHandleBuffer);
   ASSERT_EFI_ERROR (Status);
 
   //