[llvm] [llvm-objcopy] Remove references for empty section groups (PR #98106)

[Pranav Kant via llvm-commits]

https://github.com/pranavk created https://github.com/llvm/llvm-project/pull/98106

Otherwise, llvm-objcopy fails with use-after-free when built under sanitizers. Simple repro can be running the test ELF/remove-section-in-group.test under asan. This is due to symbol table references to empty section groups.

>From d5427f30bd08a49170cacf27a7dd50c5386bdcc5 Mon Sep 17 00:00:00 2001
From: Pranav Kant <prka at google.com>
Date: Tue, 9 Jul 2024 02:41:06 +0000
Subject: [PATCH] [llvm-objcopy] Remove references for empty section groups

Otherwise, llvm-objcopy fails with use-after-free when built under
sanitizers. Simple repro can be running the test ELF/remove-section-in-group.test
under asan. This is due to symbol table references to empty section groups.
---
 llvm/lib/ObjCopy/ELF/ELFObject.cpp | 13 +++++--------
 llvm/lib/ObjCopy/ELF/ELFObject.h   |  7 ++++++-
 2 files changed, 11 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/ObjCopy/ELF/ELFObject.cpp b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
index f9c5d2579be69..d222aa4f2cf3d 100644
--- a/llvm/lib/ObjCopy/ELF/ELFObject.cpp
+++ b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
@@ -2203,6 +2203,11 @@ Error Object::removeSections(
           if (auto ToRelSec = RelSec->getSection())
             return !ToRemove(*ToRelSec);
         }
+        // Remove empty group sections.
+        if (Sec->Type == ELF::SHT_GROUP) {
+          auto GroupSec = cast<GroupSection>(Sec.get());
+          return !llvm::all_of(GroupSec->members(), ToRemove);
+        }
         return true;
       });
   if (SymbolTable != nullptr && ToRemove(*SymbolTable))
@@ -2242,14 +2247,6 @@ Error Object::removeSections(
   // Now get rid of them altogether.
   Sections.erase(Iter, std::end(Sections));
 
-  // Finally erase empty SHT_GROUP sections.
-  llvm::erase_if(Sections, [](const SecPtr &Sec) {
-    if (auto GroupSec = dyn_cast<GroupSection>(Sec.get()))
-      return GroupSec->getMembersCount() == 0;
-
-    return false;
-  });
-
   return Error::success();
 }
 
diff --git a/llvm/lib/ObjCopy/ELF/ELFObject.h b/llvm/lib/ObjCopy/ELF/ELFObject.h
index 2a9f337c3f323..fd119cabcc92b 100644
--- a/llvm/lib/ObjCopy/ELF/ELFObject.h
+++ b/llvm/lib/ObjCopy/ELF/ELFObject.h
@@ -941,6 +941,9 @@ class GroupSection : public SectionBase {
   SmallVector<SectionBase *, 3> GroupMembers;
 
 public:
+  template <class T>
+  using ConstRange = iterator_range<pointee_iterator<
+      typename llvm::SmallVector<T *, 3>::const_iterator>>;
   // TODO: Contents is present in several classes of the hierarchy.
   // This needs to be refactored to avoid duplication.
   ArrayRef<uint8_t> Contents;
@@ -964,7 +967,9 @@ class GroupSection : public SectionBase {
       const DenseMap<SectionBase *, SectionBase *> &FromTo) override;
   void onRemove() override;
 
-  size_t getMembersCount() const { return GroupMembers.size(); }
+  ConstRange<SectionBase> members() const {
+    return make_pointee_range(GroupMembers);
+  }
 
   static bool classof(const SectionBase *S) {
     return S->OriginalType == ELF::SHT_GROUP;