[PATCH v5 4/5] MdeMoudlePkg/CapsulePei: Optimize GetScatterGatherHeadEntries

["Gao, Zhichao" <zhichao.gao@intel.com>]

REF: https://bugzilla.tianocore.org/show_bug.cgi?id=1853

Rename the MACRO from MAX_SG_LIST_HEADS to DEFAULT_SG_LIST_HEADS.
GetScatterGatherHeadEntries: use allocated buffer instead of fixed
array to handle the condition which the SG list is larger then the
array size.

Cc: Jian J Wang <jian.j.wang@intel.com>
Cc: Hao A Wu <hao.a.wu@intel.com>
Cc: Ray Ni <ray.ni@intel.com>
Cc: Star Zeng <star.zeng@intel.com>
Cc: Liming Gao <liming.gao@intel.com>
Cc: Sean Brogan <sean.brogan@microsoft.com>
Cc: Michael Turner <Michael.Turner@microsoft.com>
Cc: Bret Barkelew <Bret.Barkelew@microsoft.com>
Cc: Leif Lindholm <leif.lindholm@linaro.org>
Signed-off-by: Zhichao gao <zhichao.gao@intel.com>
---
 .../Universal/CapsulePei/UefiCapsule.c        | 51 ++++++++++++++-----
 1 file changed, 37 insertions(+), 14 deletions(-)

diff --git a/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c b/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
index ce6d95a786..3ac95b7be6 100644
--- a/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
+++ b/MdeModulePkg/Universal/CapsulePei/UefiCapsule.c
@@ -10,6 +10,8 @@ SPDX-License-Identifier: BSD-2-Clause-Patent
 
 #include "Capsule.h"
 
+#define DEFAULT_SG_LIST_HEADS       (20)
+
 #ifdef MDE_CPU_IA32
 //
 // Global Descriptor Table (GDT)
@@ -841,8 +843,6 @@ AreCapsulesStaged (
   return FALSE;
 }
 
-#define MAX_SG_LIST_HEADS (20)
-
 /**
   Check all the variables for SG list heads and get the count and addresses.
 
@@ -861,17 +861,19 @@ GetScatterGatherHeadEntries (
   OUT EFI_PHYSICAL_ADDRESS **HeadList
   )
 {
-  EFI_STATUS                       Status;
-  UINTN                            Size;
-  UINTN                            Index;
-  UINTN                            TempIndex;
-  UINTN                            ValidIndex;
-  BOOLEAN                          Flag;
-  CHAR16                           CapsuleVarName[30];
-  CHAR16                           *TempVarName;
-  EFI_PHYSICAL_ADDRESS             CapsuleDataPtr64;
-  EFI_PEI_READ_ONLY_VARIABLE2_PPI  *PPIVariableServices;
-  EFI_PHYSICAL_ADDRESS             TempList[MAX_SG_LIST_HEADS];
+  EFI_STATUS                        Status;
+  UINTN                             Size;
+  UINTN                             Index;
+  UINTN                             TempIndex;
+  UINTN                             ValidIndex;
+  BOOLEAN                           Flag;
+  CHAR16                            CapsuleVarName[30];
+  CHAR16                            *TempVarName;
+  EFI_PHYSICAL_ADDRESS              CapsuleDataPtr64;
+  EFI_PEI_READ_ONLY_VARIABLE2_PPI   *PPIVariableServices;
+  EFI_PHYSICAL_ADDRESS              *TempList;
+  EFI_PHYSICAL_ADDRESS              *EnlargedTempList;
+  UINTN                             TempListLength;
 
   Index             = 0;
   TempVarName       = NULL;
@@ -901,12 +903,22 @@ GetScatterGatherHeadEntries (
     return Status;
   }
 
+  //
+  // Allocate memory for sg list head
+  //
+  TempListLength = DEFAULT_SG_LIST_HEADS * sizeof (EFI_PHYSICAL_ADDRESS);
+  TempList = AllocateZeroPool (TempListLength);
+  if (TempList == NULL) {
+    DEBUG((DEBUG_ERROR, "Failed to allocate memory\n"));
+    return EFI_OUT_OF_RESOURCES;
+  }
+
   //
   // setup var name buffer for update capsules
   //
   StrCpyS (CapsuleVarName, sizeof (CapsuleVarName) / sizeof (CHAR16), EFI_CAPSULE_VARIABLE_NAME);
   TempVarName = CapsuleVarName + StrLen (CapsuleVarName);
-  while (ValidIndex < MAX_SG_LIST_HEADS) {
+  while (TRUE) {
     if (Index != 0) {
       UnicodeValueToStringS (
         TempVarName,
@@ -948,6 +960,17 @@ GetScatterGatherHeadEntries (
       continue;
     }
 
+    //
+    // The TempList is full, enlarge it
+    //
+    if ((ValidIndex + 1) >= TempListLength) {
+      EnlargedTempList = AllocateZeroPool (TempListLength * 2);
+      CopyMem (EnlargedTempList, TempList, TempListLength);
+      FreePool (TempList);
+      TempList = EnlargedTempList;
+      TempListLength *= 2;
+    }
+
     //
     // add it to the cached list
     //
-- 
2.21.0.windows.1