From 03faa0f551ba85cd6c7abc66ef25c24f6784eda7 Mon Sep 17 00:00:00 2001
From: Haojian Zhuang <haojian.zhuang@linaro.org>
Date: Sun, 11 Feb 2018 18:42:12 +0800
Subject: [PATCH 1/2] EmbeddedPkg/AndroidFastbootApp: fix overflow on fill buf

Fix overflow on fill buffer.

Signed-off-by: Haojian Zhuang <haojian.zhuang@linaro.org>
---
 .../AndroidFastboot/AndroidFastbootApp.c      | 19 +++++++++++++------
 1 file changed, 13 insertions(+), 6 deletions(-)

diff --git a/EmbeddedPkg/Application/AndroidFastboot/AndroidFastbootApp.c b/EmbeddedPkg/Application/AndroidFastboot/AndroidFastbootApp.c
index 8495536307c1..56496568563e 100644
--- a/EmbeddedPkg/Application/AndroidFastboot/AndroidFastbootApp.c
+++ b/EmbeddedPkg/Application/AndroidFastboot/AndroidFastbootApp.c
@@ -180,7 +180,7 @@ FlashSparseImage (
   )
 {
   EFI_STATUS        Status = EFI_SUCCESS;
-  UINTN             Chunk, Offset = 0, Index;
+  UINTN             Chunk, Offset = 0, Left, Count;
   VOID             *Image;
   CHUNK_HEADER     *ChunkHeader;
   UINT32            FillBuf[FILL_BUF_SIZE];
@@ -209,20 +209,27 @@ FlashSparseImage (
       Offset += ChunkHeader->ChunkSize * SparseHeader->BlockSize;
       break;
     case CHUNK_TYPE_FILL:
-      SetMem32 (FillBuf, FILL_BUF_SIZE * sizeof (UINT32), *(UINT32 *)Image);
-      Image += sizeof (UINT32);
-      for (Index = 0; Index < ChunkHeader->ChunkSize; Index++) {
+      Left = ChunkHeader->ChunkSize * SparseHeader->BlockSize;
+      while (Left > 0) {
+        if (Left > FILL_BUF_SIZE * sizeof (UINT32)) {
+          Count = FILL_BUF_SIZE * sizeof (UINT32);
+        } else {
+          Count = Left;
+        }
+        SetMem32 (FillBuf, Count, *(UINT32 *)Image);
         Status = mPlatform->FlashPartitionEx (
                               PartitionName,
                               Offset,
-                              SparseHeader->BlockSize,
+                              Count,
                               FillBuf
                               );
         if (EFI_ERROR (Status)) {
           return Status;
         }
-        Offset += SparseHeader->BlockSize;
+        Offset += Count;
+        Left = Left - Count;
       }
+      Image += sizeof (UINT32);
       break;
     case CHUNK_TYPE_DONT_CARE:
       Offset += ChunkHeader->ChunkSize * SparseHeader->BlockSize;

From c17a622f7823f6b7663e2763661d585c12431b82 Mon Sep 17 00:00:00 2001
From: Haojian Zhuang <haojian.zhuang@linaro.org>
Date: Sun, 11 Feb 2018 19:10:45 +0800
Subject: [PATCH 2/2] EmbeddedPkg/AndroidFastbootApp: increase fill buf

Increase the fill buffer that could increase the performance.

Signed-off-by: Haojian Zhuang <haojian.zhuang@linaro.org>
---
 .../AndroidFastboot/AndroidFastbootApp.c      | 33 ++++++++++++++++---
 1 file changed, 28 insertions(+), 5 deletions(-)

diff --git a/EmbeddedPkg/Application/AndroidFastboot/AndroidFastbootApp.c b/EmbeddedPkg/Application/AndroidFastboot/AndroidFastbootApp.c
index 56496568563e..9074c6c7794d 100644
--- a/EmbeddedPkg/Application/AndroidFastboot/AndroidFastbootApp.c
+++ b/EmbeddedPkg/Application/AndroidFastboot/AndroidFastbootApp.c
@@ -33,7 +33,8 @@
 #define CHUNK_TYPE_DONT_CARE        0xCAC3
 #define CHUNK_TYPE_CRC32            0xCAC4
 
-#define FILL_BUF_SIZE               1024
+#define FILL_BUF_SIZE               (16 * 1024 * 1024)
+#define SPARSE_BLOCK_SIZE           4096
 
 #define IS_DEVICE_PATH_NODE(node,type,subtype) (((node)->Type == (type)) && ((node)->SubType == (subtype)))
 
@@ -180,14 +181,35 @@ FlashSparseImage (
   )
 {
   EFI_STATUS        Status = EFI_SUCCESS;
-  UINTN             Chunk, Offset = 0, Left, Count;
+  UINTN             Chunk, Offset = 0, Left, Count, FillBufSize;
   VOID             *Image;
   CHUNK_HEADER     *ChunkHeader;
-  UINT32            FillBuf[FILL_BUF_SIZE];
+  VOID             *FillBuf;
   CHAR16            OutputString[FASTBOOT_STRING_MAX_LENGTH];
 
   Image = (VOID *)SparseHeader;
   Image += SparseHeader->FileHeaderSize;
+
+  // allocate the fill buf with dynamic size
+  FillBufSize = FILL_BUF_SIZE;
+  while (FillBufSize >= SPARSE_BLOCK_SIZE) {
+    FillBuf = AllocatePool (FillBufSize);
+    if (FillBuf == NULL) {
+      FillBufSize = FillBufSize >> 1;
+    } else {
+      break;
+    }
+  };
+  if (FillBufSize < SPARSE_BLOCK_SIZE) {
+    UnicodeSPrint (
+      OutputString,
+      sizeof (OutputString),
+      L"Fail to allocate the fill buffer\n"
+      );
+    mTextOut->OutputString (mTextOut, OutputString);
+    return EFI_BUFFER_TOO_SMALL;
+  }
+
   for (Chunk = 0; Chunk < SparseHeader->TotalChunks; Chunk++) {
     ChunkHeader = (CHUNK_HEADER *)Image;
     DEBUG ((DEBUG_INFO, "Chunk #%d - Type: 0x%x Size: %d TotalSize: %d Offset %d\n",
@@ -211,8 +233,8 @@ FlashSparseImage (
     case CHUNK_TYPE_FILL:
       Left = ChunkHeader->ChunkSize * SparseHeader->BlockSize;
       while (Left > 0) {
-        if (Left > FILL_BUF_SIZE * sizeof (UINT32)) {
-          Count = FILL_BUF_SIZE * sizeof (UINT32);
+        if (Left > FILL_BUF_SIZE) {
+          Count = FILL_BUF_SIZE;
         } else {
           Count = Left;
         }
@@ -245,6 +267,7 @@ FlashSparseImage (
       break;
     }
   }
+  FreePool ((VOID *)FillBuf);
   return Status;
 }