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

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

[Zhang, Shenglei]

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>
---
 ShellPkg/Application/Shell/FileHandleWrappers.c | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

diff --git a/ShellPkg/Application/Shell/FileHandleWrappers.c b/ShellPkg/Application/Shell/FileHandleWrappers.c
index 587556c42495..673031c9c73a 100644
--- a/ShellPkg/Application/Shell/FileHandleWrappers.c
+++ b/ShellPkg/Application/Shell/FileHandleWrappers.c
@@ -1644,7 +1644,10 @@ FileInterfaceMemWrite(
     //
     if ((UINTN)(MemFile->Position + (*BufferSize)) > (UINTN)(MemFile->BufferSize)) {
       MemFile->Buffer = ReallocatePool((UINTN)(MemFile->BufferSize), (UINTN)(MemFile->BufferSize) + (*BufferSize) + MEM_WRITE_REALLOC_OVERHEAD, MemFile->Buffer);
-      MemFile->BufferSize += (*BufferSize) + MEM_WRITE_REALLOC_OVERHEAD;
+      if (MemFile->Buffer == NULL){
+        return EFI_OUT_OF_RESOURCES;
+      }
+      MemFile->BufferSize += (*BufferSize) + MEM_WRITE_REALLOC_OVERHEAD;
     }
     CopyMem(((UINT8*)MemFile->Buffer) + MemFile->Position, Buffer, *BufferSize);
     MemFile->Position += (*BufferSize);
@@ -1661,6 +1664,9 @@ FileInterfaceMemWrite(
     AsciiSPrint(AsciiBuffer, *BufferSize, "%S", Buffer);
     if ((UINTN)(MemFile->Position + AsciiStrSize(AsciiBuffer)) > (UINTN)(MemFile->BufferSize)) {
       MemFile->Buffer = ReallocatePool((UINTN)(MemFile->BufferSize), (UINTN)(MemFile->BufferSize) + AsciiStrSize(AsciiBuffer) + MEM_WRITE_REALLOC_OVERHEAD, MemFile->Buffer);
+      if (MemFile->Buffer == NULL){
+        return EFI_OUT_OF_RESOURCES;
+      }
       MemFile->BufferSize += AsciiStrSize(AsciiBuffer) + MEM_WRITE_REALLOC_OVERHEAD;
     }
     CopyMem(((UINT8*)MemFile->Buffer) + MemFile->Position, AsciiBuffer, AsciiStrSize(AsciiBuffer));
-- 
2.18.0.windows.1

Re: [PATCH] ShellPkg/Shell/FileHandleWrappers.c: Add check for MemFile->Buffer

[Gao, Zhichao]

> -----Original Message-----
> From: Zhang, Shenglei
> Sent: Saturday, October 12, 2019 3:51 PM
> To: devel@edk2.groups.io
> Cc: Ni, Ray <ray.ni@intel.com>; Gao, Zhichao <zhichao.gao@intel.com>
> Subject: [PATCH] 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>
> ---
>  ShellPkg/Application/Shell/FileHandleWrappers.c | 8 +++++++-
>  1 file changed, 7 insertions(+), 1 deletion(-)
> 
> diff --git a/ShellPkg/Application/Shell/FileHandleWrappers.c
> b/ShellPkg/Application/Shell/FileHandleWrappers.c
> index 587556c42495..673031c9c73a 100644
> --- a/ShellPkg/Application/Shell/FileHandleWrappers.c
> +++ b/ShellPkg/Application/Shell/FileHandleWrappers.c
> @@ -1644,7 +1644,10 @@ FileInterfaceMemWrite(
>      //
>      if ((UINTN)(MemFile->Position + (*BufferSize)) > (UINTN)(MemFile-
> >BufferSize)) {
>        MemFile->Buffer = ReallocatePool((UINTN)(MemFile->BufferSize),
> (UINTN)(MemFile->BufferSize) + (*BufferSize) +
> MEM_WRITE_REALLOC_OVERHEAD, MemFile->Buffer);
> -      MemFile->BufferSize += (*BufferSize) +
> MEM_WRITE_REALLOC_OVERHEAD;
> +      if (MemFile->Buffer == NULL){
> +        return EFI_OUT_OF_RESOURCES;
> +      }
> +      MemFile->BufferSize += (*BufferSize) +
> + MEM_WRITE_REALLOC_OVERHEAD;
>      }
>      CopyMem(((UINT8*)MemFile->Buffer) + MemFile->Position, Buffer,
> *BufferSize);
>      MemFile->Position += (*BufferSize); @@ -1661,6 +1664,9 @@
> FileInterfaceMemWrite(
>      AsciiSPrint(AsciiBuffer, *BufferSize, "%S", Buffer);
>      if ((UINTN)(MemFile->Position + AsciiStrSize(AsciiBuffer)) >
> (UINTN)(MemFile->BufferSize)) {
>        MemFile->Buffer = ReallocatePool((UINTN)(MemFile->BufferSize),
> (UINTN)(MemFile->BufferSize) + AsciiStrSize(AsciiBuffer) +
> MEM_WRITE_REALLOC_OVERHEAD, MemFile->Buffer);
> +      if (MemFile->Buffer == NULL){
> +        return EFI_OUT_OF_RESOURCES;
> +      }

Before return the function, we should free all the allocated buffer, here missing the AsciiBuffer.

Thanks,
Zhichao

>        MemFile->BufferSize += AsciiStrSize(AsciiBuffer) +
> MEM_WRITE_REALLOC_OVERHEAD;
>      }
>      CopyMem(((UINT8*)MemFile->Buffer) + MemFile->Position, AsciiBuffer,
> AsciiStrSize(AsciiBuffer));
> --
> 2.18.0.windows.1