[llvm] [BOLT] Fix memory leak in BinarySection (PR #82520)

[Maksim Panchenko via llvm-commits]

https://github.com/maksfb created https://github.com/llvm/llvm-project/pull/82520

The change in #80950 exposed a memory leak in BinarySection. Let BinarySection manage memory passed via updateContents() unless a valid SectionID is set indicating that the contents are managed by JITLink.

>From 33341e03e501eb19f52ce963d25be6bd95c611aa Mon Sep 17 00:00:00 2001
From: Maksim Panchenko <maks at fb.com>
Date: Wed, 21 Feb 2024 00:22:48 -0800
Subject: [PATCH] [BOLT] Fix memory leak in BinarySection

The change in #80950 exposed a memory leak in BinarySection. Let
BinarySection manage memory passed via updateContents() unless a valid
SectionID is set indicating that the contents are managed by JITLink.
---
 bolt/include/bolt/Core/BinarySection.h | 20 +++++++++++++-------
 bolt/lib/Core/BinarySection.cpp        | 13 +------------
 bolt/lib/Rewrite/RewriteInstance.cpp   |  9 +++------
 bolt/unittests/Core/BinaryContext.cpp  |  9 +++++----
 4 files changed, 22 insertions(+), 29 deletions(-)

diff --git a/bolt/include/bolt/Core/BinarySection.h b/bolt/include/bolt/Core/BinarySection.h
index a85dbf28950e31..0f179877bd3df3 100644
--- a/bolt/include/bolt/Core/BinarySection.h
+++ b/bolt/include/bolt/Core/BinarySection.h
@@ -139,10 +139,7 @@ class BinarySection {
     Alignment = NewAlignment;
     ELFType = NewELFType;
     ELFFlags = NewELFFlags;
-    OutputSize = NewSize;
-    OutputContents = StringRef(reinterpret_cast<const char *>(NewData),
-                               NewData ? NewSize : 0);
-    IsFinalized = true;
+    updateContents(NewData, NewSize);
   }
 
 public:
@@ -484,9 +481,18 @@ class BinarySection {
   void flushPendingRelocations(raw_pwrite_stream &OS,
                                SymbolResolverFuncTy Resolver);
 
-  /// Change contents of the section.
-  void updateContents(const uint8_t *Data, size_t NewSize) {
-    OutputContents = StringRef(reinterpret_cast<const char *>(Data), NewSize);
+  /// Change contents of the section. Unless the section has a valid SectionID,
+  /// the memory passed in \p NewData will be managed by the instance of
+  /// BinarySection.
+  void updateContents(const uint8_t *NewData, size_t NewSize) {
+    if (getOutputData() && !hasValidSectionID() &&
+        (!hasSectionRef() ||
+         OutputContents.data() != getContentsOrQuit(Section).data())) {
+      delete[] getOutputData();
+    }
+
+    OutputContents = StringRef(reinterpret_cast<const char *>(NewData),
+                               NewData ? NewSize : 0);
     OutputSize = NewSize;
     IsFinalized = true;
   }
diff --git a/bolt/lib/Core/BinarySection.cpp b/bolt/lib/Core/BinarySection.cpp
index 564c63e81914c3..9ad49ca1b3a038 100644
--- a/bolt/lib/Core/BinarySection.cpp
+++ b/bolt/lib/Core/BinarySection.cpp
@@ -190,18 +190,7 @@ void BinarySection::flushPendingRelocations(raw_pwrite_stream &OS,
   clearList(PendingRelocations);
 }
 
-BinarySection::~BinarySection() {
-  if (isReordered()) {
-    delete[] getData();
-    return;
-  }
-
-  if (!isAllocatable() && !hasValidSectionID() &&
-      (!hasSectionRef() ||
-       OutputContents.data() != getContentsOrQuit(Section).data())) {
-    delete[] getOutputData();
-  }
-}
+BinarySection::~BinarySection() { updateContents(nullptr, 0); }
 
 void BinarySection::clearRelocations() { clearList(Relocations); }
 
diff --git a/bolt/lib/Rewrite/RewriteInstance.cpp b/bolt/lib/Rewrite/RewriteInstance.cpp
index 954c0fc86fa175..cde195c1739074 100644
--- a/bolt/lib/Rewrite/RewriteInstance.cpp
+++ b/bolt/lib/Rewrite/RewriteInstance.cpp
@@ -4092,12 +4092,9 @@ void RewriteInstance::rewriteNoteSections() {
       return getNewValueForSymbol(S->getName());
     });
 
-    // Set/modify section info.
-    BinarySection &NewSection = BC->registerOrUpdateNoteSection(
-        SectionName, SectionData, Size, Section.sh_addralign,
-        !BSec->isWritable(), BSec->getELFType());
-    NewSection.setOutputAddress(0);
-    NewSection.setOutputFileOffset(NextAvailableOffset);
+    // Section contents are no longer needed, but we need to update the size so
+    // that it will be reflected in the section header table.
+    BSec->updateContents(nullptr, Size);
 
     NextAvailableOffset += Size;
   }
diff --git a/bolt/unittests/Core/BinaryContext.cpp b/bolt/unittests/Core/BinaryContext.cpp
index 1fbb07bca966a7..94ee65e63a1dcf 100644
--- a/bolt/unittests/Core/BinaryContext.cpp
+++ b/bolt/unittests/Core/BinaryContext.cpp
@@ -77,10 +77,11 @@ TEST_P(BinaryContextTester, FlushPendingRelocCALL26) {
   // 12: bl func2
   // 16: func2
 
-  char Data[20] = {};
+  constexpr size_t DataSize = 20;
+  uint8_t *Data = new uint8_t[DataSize];
   BinarySection &BS = BC->registerOrUpdateSection(
-      ".text", ELF::SHT_PROGBITS, ELF::SHF_EXECINSTR | ELF::SHF_ALLOC,
-      (uint8_t *)Data, sizeof(Data), 4);
+      ".text", ELF::SHT_PROGBITS, ELF::SHF_EXECINSTR | ELF::SHF_ALLOC, Data,
+      DataSize, 4);
   MCSymbol *RelSymbol1 = BC->getOrCreateGlobalSymbol(4, "Func1");
   ASSERT_TRUE(RelSymbol1);
   BS.addRelocation(8, RelSymbol1, ELF::R_AARCH64_CALL26, 0, 0, true);
@@ -89,7 +90,7 @@ TEST_P(BinaryContextTester, FlushPendingRelocCALL26) {
   BS.addRelocation(12, RelSymbol2, ELF::R_AARCH64_CALL26, 0, 0, true);
 
   std::error_code EC;
-  SmallVector<char> Vect(sizeof(Data));
+  SmallVector<char> Vect(DataSize);
   raw_svector_ostream OS(Vect);
 
   BS.flushPendingRelocations(OS, [&](const MCSymbol *S) {