From c9c429969e0b4ff7cb0dc6f5318e828d7a3dbb04 Mon Sep 17 00:00:00 2001
From: Sergey Isakov <isakov-sl@bk.ru>
Date: Fri, 14 Feb 2020 07:48:30 +0300
Subject: [PATCH] caller responsibility to allocate memory

Signed-off-by: Sergey Isakov <isakov-sl@bk.ru>
---
 Protocols/AppleImageCodec/AppleImageCodec.c | 25 +++++++++++++++------
 1 file changed, 18 insertions(+), 7 deletions(-)

diff --git a/Protocols/AppleImageCodec/AppleImageCodec.c b/Protocols/AppleImageCodec/AppleImageCodec.c
index 8b1cabfd7..acaa46315 100644
--- a/Protocols/AppleImageCodec/AppleImageCodec.c
+++ b/Protocols/AppleImageCodec/AppleImageCodec.c
@@ -5,9 +5,11 @@
 
 **/
 
+#include <Uefi.h>
+#include <Library/BaseLib.h>
 #include <Library/UefiLib.h>
 #include <Library/BaseMemoryLib.h>
-#include <Library/MemoryAllocationLib.h>
+//#include <Library/MemoryAllocationLib.h>
 #include <Library/DebugLib.h>
 #include <Library/UefiBootServicesTableLib.h>
 
@@ -145,9 +147,11 @@ DecodeImageData (//IN APPLE_IMAGE_CODEC_PROTOCOL* This,
   UINT32         *RawImageDataSize
   )
 {
-  EFI_STATUS    Status;
+//  EFI_STATUS    Status;
   EG_IMAGE      *Image;
   INTN          Index;
+  UINT32        ImageDataSize;
+
   //automatic choose format
   if (!RawImageData || !RawImageDataSize) {
     return EFI_INVALID_PARAMETER;
@@ -159,13 +163,20 @@ DecodeImageData (//IN APPLE_IMAGE_CODEC_PROTOCOL* This,
     DBG("EFI_UNSUPPORTED\n");
     return EFI_UNSUPPORTED;
   }
-  
-  *RawImageDataSize = (UINT32)(Image->Width * Image->Height * sizeof(EFI_UGA_PIXEL));
-  Status = gBS->AllocatePool(EfiBootServicesData, *RawImageDataSize, (VOID **)RawImageData);
-  if (!EFI_ERROR(Status)) {
-    CopyMem(*RawImageData, (VOID*)Image->PixelData, *RawImageDataSize);
+
+  ImageDataSize = (UINT32)(Image->Width * Image->Height * sizeof(EFI_UGA_PIXEL));
+  if (*RawImageDataSize < ImageDataSize) {
+    *RawImageDataSize = ImageDataSize;
+    egFreeImage(Image);
+    return EFI_BUFFER_TOO_SMALL;
   }
   
+//  *RawImageDataSize = (UINT32)(Image->Width * Image->Height * sizeof(EFI_UGA_PIXEL));
+//  Status = gBS->AllocatePool(EfiBootServicesData, *RawImageDataSize, (VOID **)RawImageData);
+//  if (!EFI_ERROR(Status)) {
+  CopyMem(*RawImageData, (VOID*)Image->PixelData, ImageDataSize);
+//  }
+  
   DBG("EFI_SUCCESS, RawImageDataSize=%d\n", *RawImageDataSize);
   DBG("ImageBuffer=%p, ImageSize=%d\n", ImageBuffer, ImageSize);
   DBG("Decoded: W=%d, H=%d\n", Image->Width, Image->Height);