[edk2] [PATCH 11/27] ShellPkg/EfiCompress: Handle memory allocation failure

[Ruiyu Ni]

Contributed-under: TianoCore Contribution Agreement 1.0
Signed-off-by: Ruiyu Ni <ruiyu...@intel.com>
Cc: Jaben Carsey <jaben.car...@intel.com>
---
 .../UefiShellDebug1CommandsLib/EfiCompress.c       | 28 +++++++++++++---------
 1 file changed, 17 insertions(+), 11 deletions(-)

diff --git a/ShellPkg/Library/UefiShellDebug1CommandsLib/EfiCompress.c 
b/ShellPkg/Library/UefiShellDebug1CommandsLib/EfiCompress.c
index 85ae732..1ffbd92 100644
--- a/ShellPkg/Library/UefiShellDebug1CommandsLib/EfiCompress.c
+++ b/ShellPkg/Library/UefiShellDebug1CommandsLib/EfiCompress.c
@@ -116,20 +116,26 @@ ShellCommandRunEfiCompress (
           Status = gEfiShellProtocol->GetFileSize(InShellFileHandle, &InSize);
           ASSERT_EFI_ERROR(Status);
           InBuffer = AllocateZeroPool((UINTN)InSize);
-          ASSERT(InBuffer != NULL);
-          InSize2 = (UINTN)InSize;
-          Status = gEfiShellProtocol->ReadFile(InShellFileHandle, &InSize2, 
InBuffer);
-          InSize = InSize2;
-          ASSERT_EFI_ERROR(Status);
-          Status = Compress(InBuffer, InSize, OutBuffer, &OutSize);
-          if (Status == EFI_BUFFER_TOO_SMALL) {
-            OutBuffer = AllocateZeroPool((UINTN)OutSize);
-            ASSERT(OutBuffer != NULL);
-            Status = Compress(InBuffer, InSize, OutBuffer, &OutSize);
+          if (InBuffer == NULL) {
+            Status = EFI_OUT_OF_RESOURCES;
+          } else {
+            InSize2 = (UINTN) InSize;
+            Status = gEfiShellProtocol->ReadFile (InShellFileHandle, &InSize2, 
InBuffer);
+            InSize = InSize2;
+            ASSERT_EFI_ERROR (Status);
+            Status = Compress (InBuffer, InSize, OutBuffer, &OutSize);
+            if (Status == EFI_BUFFER_TOO_SMALL) {
+              OutBuffer = AllocateZeroPool ((UINTN) OutSize);
+              if (OutBuffer == NULL) {
+                Status = EFI_OUT_OF_RESOURCES;
+              } else {
+                Status = Compress (InBuffer, InSize, OutBuffer, &OutSize);
+              }
+            }
           }
           if (EFI_ERROR(Status)) {
             ShellPrintHiiEx(-1, -1, NULL, STRING_TOKEN 
(STR_EFI_COMPRESS_FAIL), gShellDebug1HiiHandle, Status);
-            ShellStatus = SHELL_DEVICE_ERROR;
+            ShellStatus = ((Status == EFI_OUT_OF_RESOURCES) ? 
SHELL_OUT_OF_RESOURCES : SHELL_DEVICE_ERROR);
           } else {
             OutSize2 = (UINTN)OutSize;
             Status = gEfiShellProtocol->WriteFile(OutShellFileHandle, 
&OutSize2, OutBuffer);
-- 
2.9.0.windows.1

_______________________________________________
edk2-devel mailing list
edk2-devel@lists.01.org
https://lists.01.org/mailman/listinfo/edk2-devel