K8:

[mirror_edk2.git] / MdeModulePkg / Library / UefiHiiLib / HiiLib.c

diff --git a/MdeModulePkg/Library/UefiHiiLib/HiiLib.c b/MdeModulePkg/Library/UefiHiiLib/HiiLib.c
index f5544d9fce21982e6fd4684bb1fc148551d2a507..83e33d2a6ee67b7501df037fe0a83ab93dba2a27 100644 (file)
--- a/MdeModulePkg/Library/UefiHiiLib/HiiLib.c
+++ b/MdeModulePkg/Library/UefiHiiLib/HiiLib.c
@@ -281,13 +281,13 @@ HiiLibGetHiiHandles (
   OUT    EFI_HII_HANDLE            **HiiHandleBuffer\r
   )\r
 {\r
-  UINTN       BufferLength;\r
   EFI_STATUS  Status;\r
 \r
   ASSERT (HandleBufferLength != NULL);\r
   ASSERT (HiiHandleBuffer != NULL);\r
 \r
-  BufferLength = 0;\r
+  *HandleBufferLength = 0;\r
+  *HiiHandleBuffer = NULL;\r
 \r
   //\r
   // Try to find the actual buffer size for HiiHandle Buffer.\r
@@ -296,27 +296,26 @@ HiiLibGetHiiHandles (
                                  mHiiDatabaseProt,\r
                                  EFI_HII_PACKAGE_TYPE_ALL,\r
                                  NULL,\r
-                                 &BufferLength,\r
+                                 HandleBufferLength,\r
                                  *HiiHandleBuffer\r
                                  );\r
-\r
+  \r
   if (Status == EFI_BUFFER_TOO_SMALL) {\r
-      *HiiHandleBuffer = AllocateZeroPool (BufferLength);\r
+      *HiiHandleBuffer = AllocateZeroPool (*HandleBufferLength);\r
       ASSERT (*HiiHandleBuffer != NULL);\r
       Status = mHiiDatabaseProt->ListPackageLists (\r
                                      mHiiDatabaseProt,\r
                                      EFI_HII_PACKAGE_TYPE_ALL,\r
                                      NULL,\r
-                                     &BufferLength,\r
+                                     HandleBufferLength,\r
                                      *HiiHandleBuffer\r
                                      );\r
-      //\r
-      // we should not fail here.\r
-      //\r
-      ASSERT_EFI_ERROR (Status);\r
-  }\r
 \r
-  *HandleBufferLength = BufferLength;\r
+      if (EFI_ERROR (Status)) {\r
+        FreePool (*HiiHandleBuffer);\r
+        *HiiHandleBuffer = NULL;\r
+      }\r
+  }\r
 \r
   return Status;\r
 }\r