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

[Pranav Kant via llvm-commits]

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

>From 91c329d6daf5c7b247111bd1911fe9d63211e695 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            |  5 ++
 llvm/lib/ObjCopy/ELF/ELFObject.h              |  7 +++
 .../ELF/remove-section-in-group.test          | 50 ++++++++++++++++++-
 3 files changed, 61 insertions(+), 1 deletion(-)

diff --git a/llvm/lib/ObjCopy/ELF/ELFObject.cpp b/llvm/lib/ObjCopy/ELF/ELFObject.cpp
index 5e6d19b9bfa54..fd55f974115a2 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))
diff --git a/llvm/lib/ObjCopy/ELF/ELFObject.h b/llvm/lib/ObjCopy/ELF/ELFObject.h
index e3c0e7abda16b..6ccf85387131e 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,6 +967,10 @@ class GroupSection : public SectionBase {
       const DenseMap<SectionBase *, SectionBase *> &FromTo) override;
   void onRemove() override;
 
+  ConstRange<SectionBase> members() const {
+    return make_pointee_range(GroupMembers);
+  }
+
   static bool classof(const SectionBase *S) {
     return S->OriginalType == ELF::SHT_GROUP;
   }
diff --git a/llvm/test/tools/llvm-objcopy/ELF/remove-section-in-group.test b/llvm/test/tools/llvm-objcopy/ELF/remove-section-in-group.test
index 9e683b9f68c93..20a836cae8bfe 100644
--- a/llvm/test/tools/llvm-objcopy/ELF/remove-section-in-group.test
+++ b/llvm/test/tools/llvm-objcopy/ELF/remove-section-in-group.test
@@ -1,6 +1,6 @@
 ## This checks that the group section is shrunk when its member is removed.
 
-# RUN: yaml2obj %s -o - \
+# RUN: yaml2obj --docnum=1 %s -o - \
 # RUN:   | llvm-objcopy -R .foo - - \
 # RUN:   | obj2yaml - \
 # RUN:   | FileCheck %s
@@ -35,3 +35,51 @@ Symbols:
   - Name:     foo_bar_grp
     Section:  .group
     Binding:  STB_GLOBAL
+
+# RUN: yaml2obj --docnum=2 %s -o %t
+# RUN: llvm-objcopy --remove-section=.debug_macro %t
+# RUN: llvm-readelf --section-groups %t | FileCheck %s --check-prefix=GROUP-REMOVED
+
+--- !ELF
+FileHeader:
+  Class:      ELFCLASS64
+  Data:       ELFDATA2LSB
+  Type:       ET_REL
+  Machine:    EM_X86_64
+Sections:
+  - Name:     .group
+    Type:     SHT_GROUP
+    Info:     foo_grp
+    Members:
+      - SectionOrType:  GRP_COMDAT
+      - SectionOrType:  .debug_macro
+  - Name:     .debug_macro
+    Type:     SHT_PROGBITS
+    Flags:    [ SHF_GROUP ]
+Symbols:
+  - Name:     foo_grp
+    Section:  .group
+
+# GROUP-REMOVED: There are no section groups in this file.
+
+# RUN: yaml2obj --docnum=3 %s -o %t
+# RUN: llvm-objcopy --remove-section=.group %t
+# RUN: llvm-readelf --section-groups %t | FileCheck %s --check-prefix=EMPTY-GROUP-REMOVED
+
+--- !ELF
+FileHeader:
+  Class:      ELFCLASS64
+  Data:       ELFDATA2LSB
+  Type:       ET_REL
+  Machine:    EM_X86_64
+Sections:
+  - Name:     .group
+    Type:     SHT_GROUP
+    Info:     foo_grp
+    Members:
+      - SectionOrType:  GRP_COMDAT
+Symbols:
+  - Name:     foo_grp
+    Section:  .group
+
+# EMPTY-GROUP-REMOVED: There are no section groups in this file.
\ No newline at end of file