ShellPkg/Shell/FileHandleWrappers.c: Add check for MemFile->Buffer

Add check for MemFile->Buffer.
Return EFI_OUT_OF_RESOURCES if MemFile->Buffer is NULL.

Cc: Ray Ni <ray.ni@intel.com>
Cc: Zhichao Gao <zhichao.gao@intel.com>
Signed-off-by: Shenglei Zhang <shenglei.zhang@intel.com>
Reviewed-by: Zhichao Gao <zhichao.gao@intel.com>

diff --git a/ShellPkg/Application/Shell/FileHandleWrappers.c b/ShellPkg/Application/Shell/FileHandleWrappers.c
index 587556c42495e945d79cac8492d0eb5ad6f62808..2d7bd7bec67e2cbcbd4e94a86757ddca7d700b7a 100644 (file)
--- a/ShellPkg/Application/Shell/FileHandleWrappers.c
+++ b/ShellPkg/Application/Shell/FileHandleWrappers.c
@@ -1644,6 +1644,9 @@ FileInterfaceMemWrite(
     //\r
     if ((UINTN)(MemFile->Position + (*BufferSize)) > (UINTN)(MemFile->BufferSize)) {\r
       MemFile->Buffer = ReallocatePool((UINTN)(MemFile->BufferSize), (UINTN)(MemFile->BufferSize) + (*BufferSize) + MEM_WRITE_REALLOC_OVERHEAD, MemFile->Buffer);\r
+      if (MemFile->Buffer == NULL){\r
+        return EFI_OUT_OF_RESOURCES;\r
+      }\r
       MemFile->BufferSize += (*BufferSize) + MEM_WRITE_REALLOC_OVERHEAD;\r
     }\r
     CopyMem(((UINT8*)MemFile->Buffer) + MemFile->Position, Buffer, *BufferSize);\r
@@ -1661,6 +1664,10 @@ FileInterfaceMemWrite(
     AsciiSPrint(AsciiBuffer, *BufferSize, "%S", Buffer);\r
     if ((UINTN)(MemFile->Position + AsciiStrSize(AsciiBuffer)) > (UINTN)(MemFile->BufferSize)) {\r
       MemFile->Buffer = ReallocatePool((UINTN)(MemFile->BufferSize), (UINTN)(MemFile->BufferSize) + AsciiStrSize(AsciiBuffer) + MEM_WRITE_REALLOC_OVERHEAD, MemFile->Buffer);\r
+      if (MemFile->Buffer == NULL){\r
+        FreePool(AsciiBuffer);\r
+        return EFI_OUT_OF_RESOURCES;\r
+      }\r
       MemFile->BufferSize += AsciiStrSize(AsciiBuffer) + MEM_WRITE_REALLOC_OVERHEAD;\r
     }\r
     CopyMem(((UINT8*)MemFile->Buffer) + MemFile->Position, AsciiBuffer, AsciiStrSize(AsciiBuffer));\r