[llvm-objcopy] [COFF] Add support for removing sections

[mstorsjo]

Summary

This is designed similarly to removing symbols; each section gets a unique id, which is used to map symbols to sections through removals/additions, which is then mapped back to proper sequential indices when writing out.

Compared to removal of symbols, section definition symbols add a bit of extra code, as they contain section numbers in the section aux data as well. For associative comdat symbols, that needs to be updated just like the normal section number.

I've tried to keep the removal behaviour sensible, that is: When removing a section, all symbols pointing to that section is removed. (This matches GNU objcopy.) If there were relocations against those symbols in a section that is kept, writing errors out. (GNU objcopy on the other hand will just leave the relocation pointing at any random symbol that is left, but I prefer erroring out in this case.)

If the target section of a associative comdat is removed, the associative comdat can also be removed automatically (as nothing would include it). Nothing that really is used in the wild I guess, but makes the case of manually removing a specific section more convenient.

[rupprecht]

test/tools/llvm-objcopy/COFF/remove-section.yaml

29 ↗

(On Diff #181654)

Can you keep the RUN sections together?

tools/llvm-objcopy/COFF/COFFObjcopy.cpp

32 ↗

(On Diff #181654)

Use explicit capture of &Config

tools/llvm-objcopy/COFF/Object.cpp

99 ↗

(On Diff #181654)

Use explicit capture of &ToRemove, &RemovedSections

117 ↗

(On Diff #181654)

I think it might be clearer to use an iterative approach, and possibly be able to remove the Internal var, e.g. logically something like this:

do {
  Sections.erase(...);
  Symbols.erase(...);
  ToRemove = [](...){ return AssociatedSections.count(Sec.UniqueId) == 1; }
} while (!AssociatedSections.empty());
updateSections();
updateSymbols();

I dunno though. It might be a little awkward, but recursive stuff just seems a little harder to debug, so it might be worth it.

tools/llvm-objcopy/COFF/Object.h

81 ↗

(On Diff #181654)

Internal is not very clear -- i.e. it's clear that it's for internal use, but not really clear as to how it affects the call. (Ditto for removeSections).

87 ↗

(On Diff #181654)

nit: extra blank line

111 ↗

(On Diff #181654)

nit: extra blank line

tools/llvm-objcopy/COFF/Reader.cpp

116 ↗

(On Diff #181654)

Consider having the error flow as the final condition instead of in the middle

[mstorsjo]

Applied @rupprecht's suggestions.

test/tools/llvm-objcopy/COFF/remove-section.yaml

29 ↗

(On Diff #181654)

Sure - I tried to group them with the corresponding checks, but I see why it can be good to have all the RUN commands at the top as well.

tools/llvm-objcopy/COFF/Object.cpp

117 ↗

(On Diff #181654)

Ok, that works - by updating the ToRemove predicate it becomes pretty palatable.

A second non-recursive alternative would to just remove sections and symbols for the associated sections once. I'm not entirely sure if it's valid (or supported by other tools) to have chains of associated sections, but handling that came for free in the recursive form (and by updating ToRemove).

[mstorsjo]

Tweaked the scope of lambdas in removeSections to avoid use-after-free of stack variables in the implementation of the lambda captures, as pointed out by asan.

[rupprecht]

It'd be nice if someone else could take a pass too, especially for any COFF-specific aspects, but LGTM

[mstorsjo]

In D56683#1360543, @rupprecht wrote:

It'd be nice if someone else could take a pass too, especially for any COFF-specific aspects, but LGTM

Thanks! @rnk - can you give this a peek? The main COFF specific details are updating section definition aux symbols if present, tracking the target of associative comdat sections. And if removing the target of an associative comdat, the associated section can be removed as well, to avoid leaving it dangling.

[mstorsjo]

Adding another COFF-knowledgeable reviewer

[alexander-shaposhnikov]

tools/llvm-objcopy/COFF/Writer.cpp

43 ↗

(On Diff #182137)

general comment (plus maybe others (Jake, James, Jordan) add their opinion): I know that smth like this ("finalize*" methods) has been going on in the llvm-objcopy for ELF for a while (in particular, the abstract class Writer has had the method "finalize" from the early beginning), however, this brings up the following question: what is the contract between Object and Writer ? in particular, to what extent an instance of Object should be "ready" to be consumable by a writer.

[mstorsjo]

tools/llvm-objcopy/COFF/Writer.cpp

43 ↗

(On Diff #182137)

Well, it's not very strictly typed out, but in general, while the Object class tries to maintain internal consistency as far as possible, some of the fields in the file format level structs have broken-out forms in the Symbol/Section/Relocation classes. In these cases, the broken-out field is the correct one that is maintained, and the corresponding field in the file format level struct is only set once the final layout of the file is known.

[mstorsjo]

@rnk or @smeenai - This is awaiting an ack from a COFF perspective.

[alexander-shaposhnikov]

tools/llvm-objcopy/COFF/Writer.cpp

43 ↗

(On Diff #182137)

so this is still a concern for me (and the existing ELF code too). So basically several people have mentioned interest in building a library from the code in llvm-objcopy, I'm thinking more and more about the current interfaces & design. Let me try to explain what makes me worried: so basically the idea of having these separate abstractions (Reader, Object, Writer) provides two important "extension points": we can read multiple input formats and we can write multiple output formats (at least potentially). For example, if one day smb decides to implement YAMLWriter i think it'll be unfortunate if he has to reimplement all this finalize* logic. Another thought - this introduces more subtle coupling between these classes (the assumptions which Writer makes about Object are non-obvious, at least for me). I think the same applies to ELF. @jakehehrlich , @jhenderson, @rupprecht - maybe you have some other thoughts / I'm missing something ?

[rnk]

In D56683#1363573, @mstorsjo wrote:

@rnk or @smeenai - This is awaiting an ack from a COFF perspective.

Going forward, I'd say that @mstosjo has enough COFF expertise that we don't need an additional reviewer for that. Please keep getting reviews from @rupprecht and other objcopy owners so the COFF port fits well into the objcopy design, but otherwise, forge onwards without waiting for me. Feel free to ping me if you have specific COFF questions.

[mstorsjo]

tools/llvm-objcopy/COFF/Writer.cpp

43 ↗

(On Diff #182137)

Very fair points, and it might be good to move some amount of finalization back from Writer into Object.

Are you ok with me moving forward with this patch though, and we can look at that as a separate refactoring concern at a later point?

[alexander-shaposhnikov]

tools/llvm-objcopy/COFF/Writer.cpp

43 ↗

(On Diff #182137)

yeah, I'm okay with that (don't want to block this diff), I just wanted to take advantage of this discussion to bring this question up / get some feedback from other people.

[rupprecht]

tools/llvm-objcopy/COFF/Writer.cpp

43 ↗

(On Diff #182137)

Sure, there's definitely opportunity for cleanup...

• The split between write() and finalize() doesn't seem that important -- finalize() is always followed by write(), I don't see why we can't just combine them.
• Having some sort of finalize/write logic in writers seems necessary; different ways of outputting the object is going to have different requirements (e.g. final bookkeeping)
• But, perhaps Object itself should have a finalize() method for generic things (e.g. assigning sections indices). The balance of what goes in an Object vs Writer finalize method is going to be more of an art than a science, I think.
• I think intra-coupling is actually the worse culprit, e.g. look at BinaryELFBuilder (to use my own code as a bad example), which requires its own private methods to be called in a specific error.
• Also I'm not sure how much (if at all) we can generalize this -- creating a generic Object super interface that is unique across ELF/COFF/MachO will likely be too restrictive. Probably something we should just enforce in code review.

All that said, I don't think anything blocks this patch, but it's good to bring it up.

[jhenderson]

@jakehehrlich , @jhenderson, @rupprecht - maybe you have some other thoughts / I'm missing something ?

Just a quick drive-by comment to add to what @rupprecht said. Conceptually to me, the Object should be "final" when all of its details are consistent to represent a valid Object. Section indices are a good example, since there are sections within the Object, but until finalize is called, they don't have valid index values. As a result, a finalize method on the Object probably makes sense to achieve this. Writer meanwhile wants to handle the extra information that is specific to the file format, and not really generic in any meaningful sense (e.g. program headers). I agree that it's going to be hard to have a hard-and-fast rule though.

llvm/trunk/test/tools/llvm-objcopy/COFF/remove-section.test

+# RUN: yaml2obj %s > %t.in.o
+#
+# RUN: llvm-objdump -section-headers %t.in.o | FileCheck %s --check-prefixes=SECTIONS-PRE
+# RUN: llvm-objdump -t %t.in.o | FileCheck %s --check-prefixes=SYMBOLS-PRE
+#
+# RUN: llvm-objcopy -R .bss %t.in.o %t.remove-bss.o
+# RUN: llvm-objdump -section-headers %t.remove-bss.o | FileCheck %s --check-prefix=SECTIONS-REMOVE-BSS
+# RUN: llvm-objdump -t %t.remove-bss.o | FileCheck %s --check-prefix=SYMBOLS-REMOVE-BSS
+#
+# RUN: llvm-objcopy --remove-section .bss %t.in.o %t.cmp.o
+# RUN: cmp %t.remove-bss.o %t.cmp.o
+#
+# RUN: llvm-objcopy -R .text %t.in.o %t.remove-text.o
+# RUN: llvm-objdump -section-headers %t.remove-text.o | FileCheck %s --check-prefix=SECTIONS-REMOVE-TEXT
+# RUN: llvm-objdump -t %t.remove-text.o | FileCheck %s --check-prefix=SYMBOLS-REMOVE-TEXT
+#
+# RUN: not llvm-objcopy -R .comdat %t.in.o %t.remove-comdat.o 2>&1 | FileCheck %s --check-prefix=ERROR-RELOC
+#
+# RUN: llvm-objcopy -R .text -R .comdat %t.in.o %t.remove-text-comdat.o
+# RUN: llvm-objdump -section-headers %t.remove-text-comdat.o | FileCheck %s --check-prefix=SECTIONS-REMOVE-TEXT-COMDAT
+# RUN: llvm-objdump -t %t.remove-text-comdat.o | FileCheck %s --check-prefix=SYMBOLS-REMOVE-TEXT-COMDAT
+#
+#
+# SECTIONS-PRE: Sections:
+# SECTIONS-PRE-NEXT: Idx Name
+# SECTIONS-PRE-NEXT: 0 .text
+# SECTIONS-PRE-NEXT: 1 .bss
+# SECTIONS-PRE-NEXT: 2 .comdat
+# SECTIONS-PRE-NEXT: 3 .associative
+# SECTIONS-PRE-EMPTY:
+#
+# SYMBOLS-PRE: SYMBOL TABLE:
+# SYMBOLS-PRE-NEXT: {{.*}}(sec -1){{.*}} @feat.00
+# SYMBOLS-PRE-NEXT: {{.*}}(sec 1){{.*}} .text
+# SYMBOLS-PRE-NEXT: AUX scnlen {{.*}} assoc 1 comdat 0
+# SYMBOLS-PRE-NEXT: {{.*}}(sec 2){{.*}} .bss
+# SYMBOLS-PRE-NEXT: AUX scnlen {{.*}} assoc 2 comdat 0
+# SYMBOLS-PRE-NEXT: {{.*}}(sec 4){{.*}} .associative
+# SYMBOLS-PRE-NEXT: AUX scnlen {{.*}} assoc 3 comdat 5
+# SYMBOLS-PRE-NEXT: {{.*}}(sec 3){{.*}} .comdat
+# SYMBOLS-PRE-NEXT: AUX scnlen {{.*}} assoc 3 comdat 2
+# SYMBOLS-PRE-NEXT: {{.*}}(sec 3){{.*}} foo
+# SYMBOLS-PRE-NEXT: {{.*}}(sec 1){{.*}} main
+# SYMBOLS-PRE-EMPTY:
+#
+#
+# Removing the .bss section removes one symbol and its aux symbol,
+# and updates the section indices in symbols pointing to later
+# symbols, including the aux section defintitions.
+#
+# Testing that the absolute symbol @feat.00 survives the section number
+# mangling.
+#
+# SECTIONS-REMOVE-BSS: Sections:
+# SECTIONS-REMOVE-BSS-NEXT: Idx Name
+# SECTIONS-REMOVE-BSS-NEXT: 0 .text
+# SECTIONS-REMOVE-BSS-NEXT: 1 .comdat
+# SECTIONS-REMOVE-BSS-NEXT: 2 .associative
+# SECTIONS-REMOVE-BSS-EMPTY:
+#
+# SYMBOLS-REMOVE-BSS: SYMBOL TABLE:
+# SYMBOLS-REMOVE-BSS-NEXT: {{.*}}(sec -1){{.*}} @feat.00
+# SYMBOLS-REMOVE-BSS-NEXT: {{.*}}(sec 1){{.*}} .text
+# SYMBOLS-REMOVE-BSS-NEXT: AUX scnlen {{.*}} assoc 1 comdat 0
+# SYMBOLS-REMOVE-BSS-NEXT: {{.*}}(sec 3){{.*}} .associative
+# SYMBOLS-REMOVE-BSS-NEXT: AUX scnlen {{.*}} assoc 2 comdat 5
+# SYMBOLS-REMOVE-BSS-NEXT: {{.*}}(sec 2){{.*}} .comdat
+# SYMBOLS-REMOVE-BSS-NEXT: AUX scnlen {{.*}} assoc 2 comdat 2
+# SYMBOLS-REMOVE-BSS-NEXT: {{.*}}(sec 2){{.*}} foo
+# SYMBOLS-REMOVE-BSS-NEXT: {{.*}}(sec 1){{.*}} main
+# SYMBOLS-REMOVE-BSS-EMPTY:
+#
+#
+# Removing the .text section is ok and just removes the external symbol
+# referring to it.
+#
+# SECTIONS-REMOVE-TEXT: Sections:
+# SECTIONS-REMOVE-TEXT-NEXT: Idx Name
+# SECTIONS-REMOVE-TEXT-NEXT: 0 .bss
+# SECTIONS-REMOVE-TEXT-NEXT: 1 .comdat
+# SECTIONS-REMOVE-TEXT-NEXT: 2 .associative
+# SECTIONS-REMOVE-TEXT-EMPTY:
+#
+# SYMBOLS-REMOVE-TEXT: SYMBOL TABLE:
+# SYMBOLS-REMOVE-TEXT-NEXT: {{.*}}(sec -1){{.*}} @feat.00
+# SYMBOLS-REMOVE-TEXT-NEXT: {{.*}}(sec 1){{.*}} .bss
+# SYMBOLS-REMOVE-TEXT-NEXT: AUX scnlen {{.*}} assoc 1 comdat 0
+# SYMBOLS-REMOVE-TEXT-NEXT: {{.*}}(sec 3){{.*}} .associative
+# SYMBOLS-REMOVE-TEXT-NEXT: AUX scnlen {{.*}} assoc 2 comdat 5
+# SYMBOLS-REMOVE-TEXT-NEXT: {{.*}}(sec 2){{.*}} .comdat
+# SYMBOLS-REMOVE-TEXT-NEXT: AUX scnlen {{.*}} assoc 2 comdat 2
+# SYMBOLS-REMOVE-TEXT-NEXT: {{.*}}(sec 2){{.*}} foo
+# SYMBOLS-REMOVE-TEXT-EMPTY:
+#
+#
+# Removing the .comdat section fails, since the .text section has relocations
+# against it.
+#
+# ERROR-RELOC: Relocation target foo ({{.*}}) not found
+#
+#
+# Removing the .comdat section and .text (with a relocation against .comdat)
+# works, as it also removes the .associative section transitively.
+#
+# SECTIONS-REMOVE-TEXT-COMDAT: Sections:
+# SECTIONS-REMOVE-TEXT-COMDAT-NEXT: Idx Name
+# SECTIONS-REMOVE-TEXT-COMDAT-NEXT: 0 .bss
+# SECTIONS-REMOVE-TEXT-COMDAT-EMPTY:
+#
+# SYMBOLS-REMOVE-TEXT-COMDAT: SYMBOL TABLE:
+# SYMBOLS-REMOVE-TEXT-COMDAT-NEXT: {{.*}}(sec -1){{.*}} @feat.00
+# SYMBOLS-REMOVE-TEXT-COMDAT-NEXT: {{.*}}(sec 1){{.*}} .bss
+# SYMBOLS-REMOVE-TEXT-COMDAT-NEXT: AUX scnlen {{.*}} assoc 1 comdat 0
+# SYMBOLS-REMOVE-TEXT-COMDAT-EMPTY:
+
+--- !COFF
+header:          
+  Machine:         IMAGE_FILE_MACHINE_AMD64
+  Characteristics: [  ]
+sections:        
+  - Name:            .text
+    Characteristics: [  ]
+    Alignment:       4
+    SectionData:     488B0500000000C3
+    Relocations:     
+      - VirtualAddress:  3
+        SymbolName:      foo
+        Type:            IMAGE_REL_AMD64_REL32
+  - Name:            .bss
+    Characteristics: [  ]
+    Alignment:       4
+    SectionData:     ''
+  - Name:            .comdat
+    Characteristics: [ IMAGE_SCN_LNK_COMDAT ]
+    Alignment:       1
+    SectionData:     '2A000000'
+  - Name:            .associative
+    Characteristics: [ IMAGE_SCN_LNK_COMDAT ]
+    Alignment:       1
+    SectionData:     '0000000000000000'
+symbols:         
+  - Name:            '@feat.00'
+    Value:           0
+    SectionNumber:   -1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+  - Name:            .text
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition: 
+      Length:          8
+      NumberOfRelocations: 1
+      NumberOfLinenumbers: 0
+      CheckSum:        583624169
+      Number:          1
+  - Name:            .bss
+    Value:           0
+    SectionNumber:   2
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition: 
+      Length:          0
+      NumberOfRelocations: 0
+      NumberOfLinenumbers: 0
+      CheckSum:        0
+      Number:          2
+  - Name:            .associative
+    Value:           0
+    SectionNumber:   4
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition: 
+      Length:          8
+      NumberOfRelocations: 0
+      NumberOfLinenumbers: 0
+      CheckSum:        0
+      Number:          3
+      Selection:       IMAGE_COMDAT_SELECT_ASSOCIATIVE
+  - Name:            .comdat
+    Value:           0
+    SectionNumber:   3
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_STATIC
+    SectionDefinition: 
+      Length:          4
+      NumberOfRelocations: 0
+      NumberOfLinenumbers: 0
+      CheckSum:        3482275674
+      Number:          3
+      Selection:       IMAGE_COMDAT_SELECT_ANY
+  - Name:            foo
+    Value:           0
+    SectionNumber:   3
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+  - Name:            main
+    Value:           0
+    SectionNumber:   1
+    SimpleType:      IMAGE_SYM_TYPE_NULL
+    ComplexType:     IMAGE_SYM_DTYPE_NULL
+    StorageClass:    IMAGE_SYM_CLASS_EXTERNAL
+...

llvm/trunk/tools/llvm-objcopy/COFF/COFFObjcopy.cpp

 namespace llvm {
 namespace objcopy {
 namespace coff {
 
 using namespace object;
 using namespace COFF;
 
 static Error handleArgs(const CopyConfig &Config, Object &Obj) {
+  // Perform the actual section removals.
+  Obj.removeSections([&Config](const Section &Sec) {
+    if (is_contained(Config.ToRemove, Sec.Name))
+      return true;
+
+    return false;
+  });
+
   // StripAll removes all symbols and thus also removes all relocations.
   if (Config.StripAll || Config.StripAllGNU)
-    for (Section &Sec : Obj.Sections)
+    for (Section &Sec : Obj.getMutableSections())
       Sec.Relocs.clear();
 
   // If we need to do per-symbol removals, initialize the Referenced field.
   if (Config.StripUnneeded || Config.DiscardAll ||
       !Config.SymbolsToRemove.empty())
     if (Error E = Obj.markSymbols())
       return E;
 

llvm/trunk/tools/llvm-objcopy/COFF/Object.h

struct Relocation {
   StringRef TargetName; // Used for diagnostics only
 };
 
 struct Section {
   object::coff_section Header;
   ArrayRef<uint8_t> Contents;
   std::vector<Relocation> Relocs;
   StringRef Name;
+  ssize_t UniqueId;
+  size_t Index;
 };
 
 struct Symbol {
   object::coff_symbol32 Sym;
   StringRef Name;
-  ArrayRef<uint8_t> AuxData;
+  std::vector<uint8_t> AuxData;
+  ssize_t TargetSectionId;
+  ssize_t AssociativeComdatTargetSectionId = 0;
   size_t UniqueId;
   size_t RawIndex;
   bool Referenced;
 };
 
 struct Object {
   bool IsPE = false;
 
   object::dos_header DosHeader;
   ArrayRef<uint8_t> DosStub;
 
   object::coff_file_header CoffFileHeader;
 
   bool Is64 = false;
   object::pe32plus_header PeHeader;
   uint32_t BaseOfData = 0; // pe32plus_header lacks this field.
 
   std::vector<object::data_directory> DataDirectories;
-  std::vector<Section> Sections;
 
   ArrayRef<Symbol> getSymbols() const { return Symbols; }
   // This allows mutating individual Symbols, but not mutating the list
   // of symbols itself.
   iterator_range<std::vector<Symbol>::iterator> getMutableSymbols() {
     return make_range(Symbols.begin(), Symbols.end());
   }
 
   const Symbol *findSymbol(size_t UniqueId) const;
 
   void addSymbols(ArrayRef<Symbol> NewSymbols);
   void removeSymbols(function_ref<bool(const Symbol &)> ToRemove);
 
   // Set the Referenced field on all Symbols, based on relocations in
   // all sections.
   Error markSymbols();
 
+  ArrayRef<Section> getSections() const { return Sections; }
+  // This allows mutating individual Sections, but not mutating the list
+  // of symbols itself.
+  iterator_range<std::vector<Section>::iterator> getMutableSections() {
+    return make_range(Sections.begin(), Sections.end());
+  }
+
+  const Section *findSection(ssize_t UniqueId) const;
+
+  void addSections(ArrayRef<Section> NewSections);
+  void removeSections(function_ref<bool(const Section &)> ToRemove);
+
 private:
   std::vector<Symbol> Symbols;
   DenseMap<size_t, Symbol *> SymbolMap;
 
   size_t NextSymbolUniqueId = 0;
 
+  std::vector<Section> Sections;
+  DenseMap<ssize_t, Section *> SectionMap;
+
+  ssize_t NextSectionUniqueId = 1; // Allow a UniqueId 0 to mean undefined.
+
   // Update SymbolMap and RawIndex in each Symbol.
   void updateSymbols();
+
+  // Update SectionMap and Index in each Section.
+  void updateSections();
 };
 
 // Copy between coff_symbol16 and coff_symbol32.
 // The source and destination files can use either coff_symbol16 or
 // coff_symbol32, while we always store them as coff_symbol32 in the
 // intermediate data structure.
 template <class Symbol1Ty, class Symbol2Ty>
 void copySymbol(Symbol1Ty &Dest, const Symbol2Ty &Src) {

llvm/trunk/tools/llvm-objcopy/COFF/Object.cpp

 //===- Object.cpp ---------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Object.h"
+#include "llvm/ADT/DenseSet.h"
 #include <algorithm>
 
 namespace llvm {
 namespace objcopy {
 namespace coff {
 
 using namespace object;
 
for (const Relocation &R : Sec.Relocs) {
                                            " not found",
                                        object_error::invalid_symbol_index);
       It->second->Referenced = true;
     }
   }
   return Error::success();
 }
 
+void Object::addSections(ArrayRef<Section> NewSections) {
+  for (Section S : NewSections) {
+    S.UniqueId = NextSectionUniqueId++;
+    Sections.emplace_back(S);
+  }
+  updateSections();
+}
+
+void Object::updateSections() {
+  SectionMap = DenseMap<ssize_t, Section *>(Sections.size());
+  size_t Index = 1;
+  for (Section &S : Sections) {
+    SectionMap[S.UniqueId] = &S;
+    S.Index = Index++;
+  }
+}
+
+const Section *Object::findSection(ssize_t UniqueId) const {
+  auto It = SectionMap.find(UniqueId);
+  if (It == SectionMap.end())
+    return nullptr;
+  return It->second;
+}
+
+void Object::removeSections(function_ref<bool(const Section &)> ToRemove) {
+  DenseSet<ssize_t> AssociatedSections;
+  auto RemoveAssociated = [&AssociatedSections](const Section &Sec) {
+    return AssociatedSections.count(Sec.UniqueId) == 1;
+  };
+  do {
+    DenseSet<ssize_t> RemovedSections;
+    Sections.erase(
+        std::remove_if(std::begin(Sections), std::end(Sections),
+                       [ToRemove, &RemovedSections](const Section &Sec) {
+                         bool Remove = ToRemove(Sec);
+                         if (Remove)
+                           RemovedSections.insert(Sec.UniqueId);
+                         return Remove;
+                       }),
+        std::end(Sections));
+    // Remove all symbols referring to the removed sections.
+    AssociatedSections.clear();
+    Symbols.erase(
+        std::remove_if(
+            std::begin(Symbols), std::end(Symbols),
+            [&RemovedSections, &AssociatedSections](const Symbol &Sym) {
+              // If there are sections that are associative to a removed
+              // section,
+              // remove those as well as nothing will include them (and we can't
+              // leave them dangling).
+              if (RemovedSections.count(Sym.AssociativeComdatTargetSectionId) ==
+                  1)
+                AssociatedSections.insert(Sym.TargetSectionId);
+              return RemovedSections.count(Sym.TargetSectionId) == 1;
+            }),
+        std::end(Symbols));
+    ToRemove = RemoveAssociated;
+  } while (!AssociatedSections.empty());
+  updateSections();
+  updateSymbols();
+}
+
 } // end namespace coff
 } // end namespace objcopy
 } // end namespace llvm

llvm/trunk/tools/llvm-objcopy/COFF/Reader.cpp

 //===- Reader.cpp ---------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "Reader.h"
 #include "Object.h"
 #include "llvm-objcopy.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/BinaryFormat/COFF.h"
 #include "llvm/Object/COFF.h"
 #include "llvm/Support/ErrorHandling.h"
 #include <cstddef>
 #include <cstdint>
 
 namespace llvm {
 namespace objcopy {
 namespace coff {
 
 using namespace object;
+using namespace COFF;
 
 Error COFFReader::readExecutableHeaders(Object &Obj) const {
   const dos_header *DH = COFFObj.getDOSHeader();
   Obj.Is64 = COFFObj.is64();
   if (!DH)
     return Error::success();
 
   Obj.IsPE = true;
for (size_t I = 0; I < Obj.PeHeader.NumberOfRvaAndSize; I++) {
     if (auto EC = COFFObj.getDataDirectory(I, Dir))
       return errorCodeToError(EC);
     Obj.DataDirectories.emplace_back(*Dir);
   }
   return Error::success();
 }
 
 Error COFFReader::readSections(Object &Obj) const {
+  std::vector<Section> Sections;
   // Section indexing starts from 1.
   for (size_t I = 1, E = COFFObj.getNumberOfSections(); I <= E; I++) {
     const coff_section *Sec;
     if (auto EC = COFFObj.getSection(I, Sec))
       return errorCodeToError(EC);
-    Obj.Sections.push_back(Section());
-    Section &S = Obj.Sections.back();
+    Sections.push_back(Section());
+    Section &S = Sections.back();
     S.Header = *Sec;
     if (auto EC = COFFObj.getSectionContents(Sec, S.Contents))
       return errorCodeToError(EC);
     ArrayRef<coff_relocation> Relocs = COFFObj.getRelocations(Sec);
     for (const coff_relocation &R : Relocs)
       S.Relocs.push_back(R);
     if (auto EC = COFFObj.getSectionName(Sec, S.Name))
       return errorCodeToError(EC);
     if (Sec->hasExtendedRelocations())
       return make_error<StringError>("Extended relocations not supported yet",
                                      object_error::parse_failed);
   }
+  Obj.addSections(Sections);
   return Error::success();
 }
 
 Error COFFReader::readSymbols(Object &Obj, bool IsBigObj) const {
   std::vector<Symbol> Symbols;
   Symbols.reserve(COFFObj.getRawNumberOfSymbols());
+  ArrayRef<Section> Sections = Obj.getSections();
   for (uint32_t I = 0, E = COFFObj.getRawNumberOfSymbols(); I < E;) {
     Expected<COFFSymbolRef> SymOrErr = COFFObj.getSymbol(I);
     if (!SymOrErr)
       return SymOrErr.takeError();
     COFFSymbolRef SymRef = *SymOrErr;
 
     Symbols.push_back(Symbol());
     Symbol &Sym = Symbols.back();
     // Copy symbols from the original form into an intermediate coff_symbol32.
     if (IsBigObj)
       copySymbol(Sym.Sym,
                  *reinterpret_cast<const coff_symbol32 *>(SymRef.getRawPtr()));
     else
       copySymbol(Sym.Sym,
                  *reinterpret_cast<const coff_symbol16 *>(SymRef.getRawPtr()));
     if (auto EC = COFFObj.getSymbolName(SymRef, Sym.Name))
       return errorCodeToError(EC);
     Sym.AuxData = COFFObj.getSymbolAuxData(SymRef);
     assert((Sym.AuxData.size() %
             (IsBigObj ? sizeof(coff_symbol32) : sizeof(coff_symbol16))) == 0);
+    // Find the unique id of the section
+    if (SymRef.getSectionNumber() <=
+        0) // Special symbol (undefined/absolute/debug)
+      Sym.TargetSectionId = SymRef.getSectionNumber();
+    else if (static_cast<uint32_t>(SymRef.getSectionNumber() - 1) <
+             Sections.size())
+      Sym.TargetSectionId = Sections[SymRef.getSectionNumber() - 1].UniqueId;
+    else
+      return make_error<StringError>("Section number out of range",
+                                     object_error::parse_failed);
+    // For section definitions, check if it is comdat associative, and if
+    // it is, find the target section unique id.
+    const coff_aux_section_definition *SD = SymRef.getSectionDefinition();
+    if (SD && SD->Selection == IMAGE_COMDAT_SELECT_ASSOCIATIVE) {
+      int32_t Index = SD->getNumber(IsBigObj);
+      if (Index <= 0 || static_cast<uint32_t>(Index - 1) >= Sections.size())
+        return make_error<StringError>("Unexpected associative section index",
+                                       object_error::parse_failed);
+      Sym.AssociativeComdatTargetSectionId = Sections[Index - 1].UniqueId;
+    }
     I += 1 + SymRef.getNumberOfAuxSymbols();
   }
   Obj.addSymbols(Symbols);
   return Error::success();
 }
 
 Error COFFReader::setRelocTargets(Object &Obj) const {
   std::vector<const Symbol *> RawSymbolTable;
   for (const Symbol &Sym : Obj.getSymbols()) {
     RawSymbolTable.push_back(&Sym);
     for (size_t I = 0; I < Sym.Sym.NumberOfAuxSymbols; I++)
       RawSymbolTable.push_back(nullptr);
   }
-  for (Section &Sec : Obj.Sections) {
+  for (Section &Sec : Obj.getMutableSections()) {
     for (Relocation &R : Sec.Relocs) {
       if (R.Reloc.SymbolTableIndex >= RawSymbolTable.size())
         return make_error<StringError>("SymbolTableIndex out of range",
                                        object_error::parse_failed);
       const Symbol *Sym = RawSymbolTable[R.Reloc.SymbolTableIndex];
       if (Sym == nullptr)
         return make_error<StringError>("Invalid SymbolTableIndex",
                                        object_error::parse_failed);

llvm/trunk/tools/llvm-objcopy/COFF/Writer.h

class COFFWriter {
   Buffer &Buf;
 
   size_t FileSize;
   size_t FileAlignment;
   size_t SizeOfInitializedData;
   StringTableBuilder StrTabBuilder;
 
   Error finalizeRelocTargets();
+  Error finalizeSectionNumbers();
   void layoutSections();
   size_t finalizeStringTable();
   template <class SymbolTy> std::pair<size_t, size_t> finalizeSymbolTable();
 
   Error finalize(bool IsBigObj);
 
   void writeHeaders(bool IsBigObj);
   void writeSections();

llvm/trunk/tools/llvm-objcopy/COFF/Writer.cpp

 namespace llvm {
 namespace objcopy {
 namespace coff {
 
 using namespace object;
 using namespace COFF;
 
 Error COFFWriter::finalizeRelocTargets() {
-  for (Section &Sec : Obj.Sections) {
+  for (Section &Sec : Obj.getMutableSections()) {
     for (Relocation &R : Sec.Relocs) {
       const Symbol *Sym = Obj.findSymbol(R.Target);
       if (Sym == nullptr)
         return make_error<StringError>("Relocation target " + R.TargetName +
                                            " (" + Twine(R.Target) +
                                            ") not found",
                                        object_error::invalid_symbol_index);
       R.Reloc.SymbolTableIndex = Sym->RawIndex;
     }
   }
   return Error::success();
 }
 
+Error COFFWriter::finalizeSectionNumbers() {
+  for (Symbol &Sym : Obj.getMutableSymbols()) {
+    if (Sym.TargetSectionId <= 0) {
+      // Undefined, or a special kind of symbol. These negative values
+      // are stored in the SectionNumber field which is unsigned.
+      Sym.Sym.SectionNumber = static_cast<uint32_t>(Sym.TargetSectionId);
+    } else {
+      const Section *Sec = Obj.findSection(Sym.TargetSectionId);
+      if (Sec == nullptr)
+        return make_error<StringError>("Symbol " + Sym.Name +
+                                           " points to a removed section",
+                                       object_error::invalid_symbol_index);
+      Sym.Sym.SectionNumber = Sec->Index;
+
+      if (Sym.Sym.NumberOfAuxSymbols == 1 &&
+          Sym.Sym.StorageClass == IMAGE_SYM_CLASS_STATIC) {
+        coff_aux_section_definition *SD =
+            reinterpret_cast<coff_aux_section_definition *>(Sym.AuxData.data());
+        uint32_t SDSectionNumber;
+        if (Sym.AssociativeComdatTargetSectionId == 0) {
+          // Not a comdat associative section; just set the Number field to
+          // the number of the section itself.
+          SDSectionNumber = Sec->Index;
+        } else {
+          Sec = Obj.findSection(Sym.AssociativeComdatTargetSectionId);
+          if (Sec == nullptr)
+            return make_error<StringError>(
+                "Symbol " + Sym.Name + " is associative to a removed section",
+                object_error::invalid_symbol_index);
+          SDSectionNumber = Sec->Index;
+        }
+        // Update the section definition with the new section number.
+        SD->NumberLowPart = static_cast<uint16_t>(SDSectionNumber);
+        SD->NumberHighPart = static_cast<uint16_t>(SDSectionNumber >> 16);
+      }
+    }
+  }
+  return Error::success();
+}
+
 void COFFWriter::layoutSections() {
-  for (auto &S : Obj.Sections) {
+  for (auto &S : Obj.getMutableSections()) {
     if (S.Header.SizeOfRawData > 0)
       S.Header.PointerToRawData = FileSize;
     FileSize += S.Header.SizeOfRawData; // For executables, this is already
                                         // aligned to FileAlignment.
     S.Header.NumberOfRelocations = S.Relocs.size();
     S.Header.PointerToRelocations =
         S.Header.NumberOfRelocations > 0 ? FileSize : 0;
     FileSize += S.Relocs.size() * sizeof(coff_relocation);
     FileSize = alignTo(FileSize, FileAlignment);
 
     if (S.Header.Characteristics & IMAGE_SCN_CNT_INITIALIZED_DATA)
       SizeOfInitializedData += S.Header.SizeOfRawData;
   }
 }
 
 size_t COFFWriter::finalizeStringTable() {
-  for (auto &S : Obj.Sections)
+  for (const auto &S : Obj.getSections())
     if (S.Name.size() > COFF::NameSize)
       StrTabBuilder.add(S.Name);
 
   for (const auto &S : Obj.getSymbols())
     if (S.Name.size() > COFF::NameSize)
       StrTabBuilder.add(S.Name);
 
   StrTabBuilder.finalize();
 
-  for (auto &S : Obj.Sections) {
+  for (auto &S : Obj.getMutableSections()) {
     if (S.Name.size() > COFF::NameSize) {
       snprintf(S.Header.Name, sizeof(S.Header.Name), "/%d",
                (int)StrTabBuilder.getOffset(S.Name));
     } else {
       strncpy(S.Header.Name, S.Name.data(), COFF::NameSize);
     }
   }
   for (auto &S : Obj.getMutableSymbols()) {
std::pair<size_t, size_t> COFFWriter::finalizeSymbolTable() {
   for (const auto &S : Obj.getSymbols())
     SymTabSize += S.AuxData.size();
   return std::make_pair(SymTabSize, sizeof(SymbolTy));
 }
 
 Error COFFWriter::finalize(bool IsBigObj) {
   if (Error E = finalizeRelocTargets())
     return E;
+  if (Error E = finalizeSectionNumbers())
+    return E;
 
   size_t SizeOfHeaders = 0;
   FileAlignment = 1;
   size_t PeHeaderSize = 0;
   if (Obj.IsPE) {
     Obj.DosHeader.AddressOfNewExeHeader =
         sizeof(Obj.DosHeader) + Obj.DosStub.size();
     SizeOfHeaders += Obj.DosHeader.AddressOfNewExeHeader + sizeof(PEMagic);
 
     FileAlignment = Obj.PeHeader.FileAlignment;
     Obj.PeHeader.NumberOfRvaAndSize = Obj.DataDirectories.size();
 
     PeHeaderSize = Obj.Is64 ? sizeof(pe32plus_header) : sizeof(pe32_header);
     SizeOfHeaders +=
         PeHeaderSize + sizeof(data_directory) * Obj.DataDirectories.size();
   }
-  Obj.CoffFileHeader.NumberOfSections = Obj.Sections.size();
+  Obj.CoffFileHeader.NumberOfSections = Obj.getSections().size();
   SizeOfHeaders +=
       IsBigObj ? sizeof(coff_bigobj_file_header) : sizeof(coff_file_header);
-  SizeOfHeaders += sizeof(coff_section) * Obj.Sections.size();
+  SizeOfHeaders += sizeof(coff_section) * Obj.getSections().size();
   SizeOfHeaders = alignTo(SizeOfHeaders, FileAlignment);
 
   Obj.CoffFileHeader.SizeOfOptionalHeader =
       PeHeaderSize + sizeof(data_directory) * Obj.DataDirectories.size();
 
   FileSize = SizeOfHeaders;
   SizeOfInitializedData = 0;
 
   layoutSections();
 
   if (Obj.IsPE) {
     Obj.PeHeader.SizeOfHeaders = SizeOfHeaders;
     Obj.PeHeader.SizeOfInitializedData = SizeOfInitializedData;
 
-    if (!Obj.Sections.empty()) {
-      const Section &S = Obj.Sections.back();
+    if (!Obj.getSections().empty()) {
+      const Section &S = Obj.getSections().back();
       Obj.PeHeader.SizeOfImage =
           alignTo(S.Header.VirtualAddress + S.Header.VirtualSize,
                   Obj.PeHeader.SectionAlignment);
     }
 
     // If the PE header had a checksum, clear it, since it isn't valid
     // any longer. (We don't calculate a new one.)
     Obj.PeHeader.CheckSum = 0;
if (!IsBigObj) {
     BigObjHeader.TimeDateStamp = Obj.CoffFileHeader.TimeDateStamp;
     memcpy(BigObjHeader.UUID, BigObjMagic, sizeof(BigObjMagic));
     BigObjHeader.unused1 = 0;
     BigObjHeader.unused2 = 0;
     BigObjHeader.unused3 = 0;
     BigObjHeader.unused4 = 0;
     // The value in Obj.CoffFileHeader.NumberOfSections is truncated, thus
     // get the original one instead.
-    BigObjHeader.NumberOfSections = Obj.Sections.size();
+    BigObjHeader.NumberOfSections = Obj.getSections().size();
     BigObjHeader.PointerToSymbolTable = Obj.CoffFileHeader.PointerToSymbolTable;
     BigObjHeader.NumberOfSymbols = Obj.CoffFileHeader.NumberOfSymbols;
 
     memcpy(Ptr, &BigObjHeader, sizeof(BigObjHeader));
     Ptr += sizeof(BigObjHeader);
   }
   if (Obj.IsPE) {
     if (Obj.Is64) {
       memcpy(Ptr, &Obj.PeHeader, sizeof(Obj.PeHeader));
       Ptr += sizeof(Obj.PeHeader);
     } else {
       pe32_header PeHeader;
       copyPeHeader(PeHeader, Obj.PeHeader);
       // The pe32plus_header (stored in Object) lacks the BaseOfData field.
       PeHeader.BaseOfData = Obj.BaseOfData;
 
       memcpy(Ptr, &PeHeader, sizeof(PeHeader));
       Ptr += sizeof(PeHeader);
     }
     for (const auto &DD : Obj.DataDirectories) {
       memcpy(Ptr, &DD, sizeof(DD));
       Ptr += sizeof(DD);
     }
   }
-  for (const auto &S : Obj.Sections) {
+  for (const auto &S : Obj.getSections()) {
     memcpy(Ptr, &S.Header, sizeof(S.Header));
     Ptr += sizeof(S.Header);
   }
 }
 
 void COFFWriter::writeSections() {
-  for (const auto &S : Obj.Sections) {
+  for (const auto &S : Obj.getSections()) {
     uint8_t *Ptr = Buf.getBufferStart() + S.Header.PointerToRawData;
     std::copy(S.Contents.begin(), S.Contents.end(), Ptr);
 
     // For executable sections, pad the remainder of the raw data size with
     // 0xcc, which is int3 on x86.
     if ((S.Header.Characteristics & IMAGE_SCN_CNT_CODE) &&
         S.Header.SizeOfRawData > S.Contents.size())
       memset(Ptr + S.Contents.size(), 0xcc,

 // the debug_directory structs in there, and set the PointerToRawData field
 // in all of them, according to their new physical location in the file.
 Error COFFWriter::patchDebugDirectory() {
   if (Obj.DataDirectories.size() < DEBUG_DIRECTORY)
     return Error::success();
   const data_directory *Dir = &Obj.DataDirectories[DEBUG_DIRECTORY];
   if (Dir->Size <= 0)
     return Error::success();
-  for (const auto &S : Obj.Sections) {
+  for (const auto &S : Obj.getSections()) {
     if (Dir->RelativeVirtualAddress >= S.Header.VirtualAddress &&
         Dir->RelativeVirtualAddress <
             S.Header.VirtualAddress + S.Header.SizeOfRawData) {
       if (Dir->RelativeVirtualAddress + Dir->Size >
           S.Header.VirtualAddress + S.Header.SizeOfRawData)
         return make_error<StringError>(
             "Debug directory extends past end of section",
             object_error::parse_failed);
if (Dir->RelativeVirtualAddress >= S.Header.VirtualAddress &&
       return Error::success();
     }
   }
   return make_error<StringError>("Debug directory not found",
                                  object_error::parse_failed);
 }
 
 Error COFFWriter::write() {
-  bool IsBigObj = Obj.Sections.size() > MaxNumberOfSections16;
+  bool IsBigObj = Obj.getSections().size() > MaxNumberOfSections16;
   if (IsBigObj && Obj.IsPE)
     return make_error<StringError>("Too many sections for executable",
                                    object_error::parse_failed);
   return write(IsBigObj);
 }
 
 } // end namespace coff
 } // end namespace objcopy
 } // end namespace llvm