From 1bc2fb3eea1e082fd547cbbf5f9ae870dabbfa97 Mon Sep 17 00:00:00 2001
From: Goldfish64 <goldfishnet@aol.com>
Date: Sun, 26 Jul 2020 23:04:59 -0500
Subject: [PATCH] OcXmlLib: Fix previous commit

---
 Library/OcXmlLib/OcXmlLib.c | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/Library/OcXmlLib/OcXmlLib.c b/Library/OcXmlLib/OcXmlLib.c
index 9b0c530c..d84a346d 100755
--- a/Library/OcXmlLib/OcXmlLib.c
+++ b/Library/OcXmlLib/OcXmlLib.c
@@ -1153,30 +1153,33 @@ XmlDocumentExport (
   )
 {
   CHAR8   *Buffer;
-  CHAR8   *BufferXmlContent;
+  CHAR8   *NewBuffer;
   UINT32  AllocSize;
   UINT32  CurrentSize;
+  UINT32  NewSize;
 
   AllocSize = Document->Buffer.Length + 1;
-  if (PrependPlistInfo && OcOverflowAddU32 (AllocSize, L_STR_SIZE_NT (XML_PLIST_HEADER), &AllocSize)) {
-      return NULL;
-  }
   Buffer = AllocatePool (AllocSize);
   if (Buffer == NULL) {
     XML_USAGE_ERROR ("XmlDocumentExport::failed to allocate");
     return NULL;
   }
-  BufferXmlContent = PrependPlistInfo ? &Buffer[L_STR_LEN (XML_PLIST_HEADER)] : Buffer;
 
   CurrentSize = 0;
-  XmlNodeExportRecursive (Document->Root, &BufferXmlContent, &AllocSize, &CurrentSize, Skip);
+  XmlNodeExportRecursive (Document->Root, &Buffer, &AllocSize, &CurrentSize, Skip);
 
   if (PrependPlistInfo) {
-    if (OcOverflowAddU32 (CurrentSize, L_STR_SIZE_NT (XML_PLIST_HEADER), &CurrentSize)) {
+    if (OcOverflowTriAddU32 (CurrentSize, L_STR_SIZE_NT (XML_PLIST_HEADER), 1, &NewSize)) {
       FreePool (Buffer);
       return NULL;
     }
-    CopyMem (Buffer, XML_PLIST_HEADER, L_STR_SIZE_NT (XML_PLIST_HEADER));
+    NewBuffer = AllocatePool (NewSize);
+    CopyMem (&NewBuffer[L_STR_LEN (XML_PLIST_HEADER)], Buffer, CurrentSize);
+    CopyMem (NewBuffer, XML_PLIST_HEADER, L_STR_SIZE_NT (XML_PLIST_HEADER));
+    FreePool (Buffer);
+
+    CurrentSize = NewSize - 1;
+    Buffer      = NewBuffer;
   }
 
   if (Length != NULL) {