From cbe129795ea40e623178aa52384c7fd48a863ea5 Mon Sep 17 00:00:00 2001
From: Liam Murphy <liamjm@google.com>
Date: Thu, 26 Nov 2020 00:14:00 +0000
Subject: [PATCH] migrate to MemBuffer class from New() to remove leak

---
 src/pefile.cpp | 39 +++++++++++++++++++++++----------------
 src/pefile.h   |  3 +++
 2 files changed, 26 insertions(+), 16 deletions(-)

diff --git a/src/pefile.cpp b/src/pefile.cpp
index 87ef7b3f..7ce6dd83 100644
--- a/src/pefile.cpp
+++ b/src/pefile.cpp
@@ -1737,16 +1737,21 @@ void PeFile::Resource::build(const upx_rnode *node, unsigned &bpos,
     }
 }
 
-upx_byte *PeFile::Resource::build()
+upx_byte* PeFile::Resource::build()
 {
-    newstart = New(upx_byte, dirsize());
-    unsigned bpos = 0,spos = dsize;
-    build(root,bpos,spos,0);
+    mem.dealloc();
+    newstart = nullptr;
+    if (dirsize()) {
+      mem.alloc(dirsize());
+      newstart = static_cast<upx_byte *>(mem.getVoidPtr());
+      unsigned bpos = 0,spos = dsize;
+      build(root,bpos,spos,0);
 
-    // dirsize() is 4 bytes aligned, so we may need to zero
-    // up to 2 bytes to make valgrind happy
-    while (spos < dirsize())
-        newstart[spos++] = 0;
+      // dirsize() is 4 bytes aligned, so we may need to zero
+      // up to 2 bytes to make valgrind happy
+      while (spos < dirsize())
+          newstart[spos++] = 0;
+    }
 
     return newstart;
 }
@@ -1833,9 +1838,10 @@ void PeFile::processResources(Resource *res,unsigned newaddr)
     while (res->next())
         if (res->newoffs())
             res->newoffs() += newaddr;
-    upx_byte *p = res->build();
-    memcpy(oresources,p,res->dirsize());
-    delete [] p;
+    if (res->dirsize()) {
+      upx_byte *p = res->build();
+      memcpy(oresources,p,res->dirsize());
+    }
 }
 
 static bool match(unsigned itype, const unsigned char *ntype,
@@ -2767,12 +2773,13 @@ void PeFile::rebuildResources(upx_byte *& extrainfo, unsigned lastvaddr)
                 icondir_count = 0;
             }
         }
-    upx_byte *p = res.build();
-    OCHECK(obuf + ODADDR(PEDIR_RESOURCE) - rvamin, 16);
-    // write back when the original is zeroed
-    if (get_le32(obuf + ODADDR(PEDIR_RESOURCE) - rvamin + 12) == 0)
+    if (res.dirsize()) {
+      upx_byte *p = res.build();
+      OCHECK(obuf + ODADDR(PEDIR_RESOURCE) - rvamin, 16);
+      // write back when the original is zeroed
+      if (get_le32(obuf + ODADDR(PEDIR_RESOURCE) - rvamin + 12) == 0)
         omemcpy(obuf + ODADDR(PEDIR_RESOURCE) - rvamin, p, res.dirsize());
-    delete [] p;
+    }
 }
 
 template <typename LEXX, typename ord_mask_t>
diff --git a/src/pefile.h b/src/pefile.h
index 60551193..e6245cde 100644
--- a/src/pefile.h
+++ b/src/pefile.h
@@ -29,6 +29,8 @@
 #ifndef __UPX_PEFILE_H
 #define __UPX_PEFILE_H 1
 
+#include "mem.h"
+
 
 /*************************************************************************
 // general/pe handling
@@ -348,6 +350,7 @@ protected:
         struct upx_rbranch;
         struct upx_rleaf;
 
+        MemBuffer mem;
         const upx_byte *start;
         upx_byte   *newstart;
         upx_rnode  *root;