[yaml2obj][obj2yaml] Support custom ELF section header string table name

This patch adds support for a new field in the FileHeader, which states
the name to use for the section header string table. This also allows
combining the string table with another string table in the object, e.g.
the symbol name string table. The field is optional. By default,
.shstrtab will continue to be used.

This partially fixes https://bugs.llvm.org/show_bug.cgi?id=50506.

Reviewed by: Higuoxing

Differential Revision: https://reviews.llvm.org/D104035

Author: jh7370

llvm/include/llvm/ObjectYAML/ELFYAML.h

@@ -118,6 +118,7 @@ struct FileHeader {
   Optional<ELF_EM> Machine;
   ELF_EF Flags;
   llvm::yaml::Hex64 Entry;
+  Optional<StringRef> SectionHeaderStringTable;
 
   Optional<llvm::yaml::Hex64> EPhOff;
   Optional<llvm::yaml::Hex16> EPhEntSize;

llvm/lib/ObjectYAML/ELFEmitter.cpp

@@ -177,15 +177,24 @@ template <class ELFT> class ELFState {
 
   enum class SymtabType { Static, Dynamic };
 
-  /// The future ".strtab" section.
+  /// The future symbol table string section.
   StringTableBuilder DotStrtab{StringTableBuilder::ELF};
 
-  /// The future ".shstrtab" section.
+  /// The future section header string table section, if a unique string table
+  /// is needed. Don't reference this variable direectly: use the
+  /// ShStrtabStrings member instead.
   StringTableBuilder DotShStrtab{StringTableBuilder::ELF};
 
-  /// The future ".dynstr" section.
+  /// The future dynamic symbol string section.
   StringTableBuilder DotDynstr{StringTableBuilder::ELF};
 
+  /// The name of the section header string table section. If it is .strtab or
+  /// .dynstr, the section header strings will be written to the same string
+  /// table as the static/dynamic symbols respectively. Otherwise a dedicated
+  /// section will be created with that name.
+  StringRef SectionHeaderStringTableName = ".shstrtab";
+  StringTableBuilder *ShStrtabStrings = &DotShStrtab;
+
   NameToIdxMap SN2I;
   NameToIdxMap SymN2I;
   NameToIdxMap DynSymN2I;
@@ -327,6 +336,18 @@ template <class T> static void zero(T &Obj) { memset(&Obj, 0, sizeof(Obj)); }
 template <class ELFT>
 ELFState<ELFT>::ELFState(ELFYAML::Object &D, yaml::ErrorHandler EH)
     : Doc(D), ErrHandler(EH) {
+  // The input may explicitly request to store the section header table strings
+  // in the same string table as dynamic or static symbol names. Set the
+  // ShStrtabStrings member accordingly.
+  if (Doc.Header.SectionHeaderStringTable) {
+    SectionHeaderStringTableName = *Doc.Header.SectionHeaderStringTable;
+    if (*Doc.Header.SectionHeaderStringTable == ".strtab")
+      ShStrtabStrings = &DotStrtab;
+    else if (*Doc.Header.SectionHeaderStringTable == ".dynstr")
+      ShStrtabStrings = &DotDynstr;
+    // Otherwise, the unique table will be used.
+  }
+
   std::vector<ELFYAML::Section *> Sections = Doc.getSections();
   // Insert SHT_NULL section implicitly when it is not defined in YAML.
   if (Sections.empty() || Sections.front()->Type != ELF::SHT_NULL)
@@ -363,19 +384,35 @@ ELFState<ELFT>::ELFState(ELFYAML::Object &D, yaml::ErrorHandler EH)
                   "' at YAML section/fill number " + Twine(I));
   }
 
-  std::vector<StringRef> ImplicitSections;
-  if (Doc.DynamicSymbols)
-    ImplicitSections.insert(ImplicitSections.end(), {".dynsym", ".dynstr"});
-  if (Doc.Symbols)
-    ImplicitSections.push_back(".symtab");
+  SmallSetVector<StringRef, 8> ImplicitSections;
+  if (Doc.DynamicSymbols) {
+    if (SectionHeaderStringTableName == ".dynsym")
+      reportError("cannot use '.dynsym' as the section header name table when "
+                  "there are dynamic symbols");
+    ImplicitSections.insert(".dynsym");
+    ImplicitSections.insert(".dynstr");
+  }
+  if (Doc.Symbols) {
+    if (SectionHeaderStringTableName == ".symtab")
+      reportError("cannot use '.symtab' as the section header name table when "
+                  "there are symbols");
+    ImplicitSections.insert(".symtab");
+  }
   if (Doc.DWARF)
     for (StringRef DebugSecName : Doc.DWARF->getNonEmptySectionNames()) {
       std::string SecName = ("." + DebugSecName).str();
-      ImplicitSections.push_back(StringRef(SecName).copy(StringAlloc));
+      // TODO: For .debug_str it should be possible to share the string table,
+      // in the same manner as the symbol string tables.
+      if (SectionHeaderStringTableName == SecName)
+        reportError("cannot use '" + SecName +
+                    "' as the section header name table when it is needed for "
+                    "DWARF output");
+      ImplicitSections.insert(StringRef(SecName).copy(StringAlloc));
     }
-  ImplicitSections.insert(ImplicitSections.end(), {".strtab"});
+  // TODO: Only create the .strtab here if any symbols have been requested.
+  ImplicitSections.insert(".strtab");
   if (!SecHdrTable || !SecHdrTable->NoHeaders.getValueOr(false))
-    ImplicitSections.insert(ImplicitSections.end(), {".shstrtab"});
+    ImplicitSections.insert(SectionHeaderStringTableName);
 
   // Insert placeholders for implicit sections that are not
   // defined explicitly in YAML.
@@ -387,7 +424,9 @@ ELFState<ELFT>::ELFState(ELFYAML::Object &D, yaml::ErrorHandler EH)
         ELFYAML::Chunk::ChunkKind::RawContent, true /*IsImplicit*/);
     Sec->Name = SecName;
 
-    if (SecName == ".dynsym")
+    if (SecName == SectionHeaderStringTableName)
+      Sec->Type = ELF::SHT_STRTAB;
+    else if (SecName == ".dynsym")
       Sec->Type = ELF::SHT_DYNSYM;
     else if (SecName == ".symtab")
       Sec->Type = ELF::SHT_SYMTAB;
@@ -480,8 +519,9 @@ void ELFState<ELFT>::writeELFHeader(raw_ostream &OS) {
 
   if (Doc.Header.EShStrNdx)
     Header.e_shstrndx = *Doc.Header.EShStrNdx;
-  else if (SectionHeaders.Offset && !ExcludedSectionHeaders.count(".shstrtab"))
-    Header.e_shstrndx = SN2I.get(".shstrtab");
+  else if (SectionHeaders.Offset &&
+           !ExcludedSectionHeaders.count(SectionHeaderStringTableName))
+    Header.e_shstrndx = SN2I.get(SectionHeaderStringTableName);
   else
     Header.e_shstrndx = 0;
 
@@ -615,16 +655,16 @@ bool ELFState<ELFT>::initImplicitHeader(ContiguousBlobAccumulator &CBA,
   if (Header.sh_offset)
     return false;
 
-  if (SecName == ".symtab")
-    initSymtabSectionHeader(Header, SymtabType::Static, CBA, YAMLSec);
-  else if (SecName == ".strtab")
+  if (SecName == ".strtab")
     initStrtabSectionHeader(Header, SecName, DotStrtab, CBA, YAMLSec);
-  else if (SecName == ".shstrtab")
-    initStrtabSectionHeader(Header, SecName, DotShStrtab, CBA, YAMLSec);
-  else if (SecName == ".dynsym")
-    initSymtabSectionHeader(Header, SymtabType::Dynamic, CBA, YAMLSec);
   else if (SecName == ".dynstr")
     initStrtabSectionHeader(Header, SecName, DotDynstr, CBA, YAMLSec);
+  else if (SecName == SectionHeaderStringTableName)
+    initStrtabSectionHeader(Header, SecName, *ShStrtabStrings, CBA, YAMLSec);
+  else if (SecName == ".symtab")
+    initSymtabSectionHeader(Header, SymtabType::Static, CBA, YAMLSec);
+  else if (SecName == ".dynsym")
+    initSymtabSectionHeader(Header, SymtabType::Dynamic, CBA, YAMLSec);
   else if (SecName.startswith(".debug_")) {
     // If a ".debug_*" section's type is a preserved one, e.g., SHT_DYNAMIC, we
     // will not treat it as a debug section.
@@ -671,7 +711,7 @@ uint64_t ELFState<ELFT>::getSectionNameOffset(StringRef Name) {
   // the string table.
   if (ExcludedSectionHeaders.count(Name))
     return 0;
-  return DotShStrtab.getOffset(Name);
+  return ShStrtabStrings->getOffset(Name);
 }
 
 static uint64_t writeContent(ContiguousBlobAccumulator &CBA,
@@ -995,7 +1035,7 @@ void ELFState<ELFT>::initStrtabSectionHeader(Elf_Shdr &SHeader, StringRef Name,
                                              StringTableBuilder &STB,
                                              ContiguousBlobAccumulator &CBA,
                                              ELFYAML::Section *YAMLSec) {
-  SHeader.sh_name = getSectionNameOffset(Name);
+  SHeader.sh_name = getSectionNameOffset(ELFYAML::dropUniqueSuffix(Name));
   SHeader.sh_type = YAMLSec ? YAMLSec->Type : ELF::SHT_STRTAB;
   SHeader.sh_addralign = YAMLSec ? (uint64_t)YAMLSec->AddressAlign : 1;
 
@@ -1817,10 +1857,8 @@ template <class ELFT> void ELFState<ELFT>::buildSectionIndex() {
       llvm_unreachable("buildSectionIndex() failed");
 
     if (!ExcludedSectionHeaders.count(S->Name))
-      DotShStrtab.add(ELFYAML::dropUniqueSuffix(S->Name));
+      ShStrtabStrings->add(ELFYAML::dropUniqueSuffix(S->Name));
   }
-
-  DotShStrtab.finalize();
 }
 
 template <class ELFT> void ELFState<ELFT>::buildSymbolIndexes() {
@@ -1870,6 +1908,11 @@ template <class ELFT> void ELFState<ELFT>::finalizeStrings() {
   }
 
   DotDynstr.finalize();
+
+  // Don't finalize the section header string table a second time if it has
+  // already been finalized due to being one of the symbol string tables.
+  if (ShStrtabStrings != &DotStrtab && ShStrtabStrings != &DotDynstr)
+    ShStrtabStrings->finalize();
 }
 
 template <class ELFT>
@@ -1879,14 +1922,16 @@ bool ELFState<ELFT>::writeELF(raw_ostream &OS, ELFYAML::Object &Doc,
   if (State.HasError)
     return false;
 
-  // Finalize .strtab and .dynstr sections. We do that early because want to
-  // finalize the string table builders before writing the content of the
-  // sections that might want to use them.
-  State.finalizeStrings();
-
+  // Build the section index, which adds sections to the section header string
+  // table first, so that we can finalise the section header string table.
   State.buildSectionIndex();
   State.buildSymbolIndexes();
 
+  // Finalize section header string table and the .strtab and .dynstr sections.
+  // We do this early because we want to finalize the string table builders
+  // before writing the content of the sections that might want to use them.
+  State.finalizeStrings();
+
   if (State.HasError)
     return false;
 

llvm/lib/ObjectYAML/ELFYAML.cpp

@@ -1007,6 +1007,7 @@ void MappingTraits<ELFYAML::FileHeader>::mapping(IO &IO,
   IO.mapOptional("Machine", FileHdr.Machine);
   IO.mapOptional("Flags", FileHdr.Flags, ELFYAML::ELF_EF(0));
   IO.mapOptional("Entry", FileHdr.Entry, Hex64(0));
+  IO.mapOptional("SectionHeaderStringTable", FileHdr.SectionHeaderStringTable);
 
   // obj2yaml does not dump these fields.
   assert(!IO.outputting() ||

llvm/test/tools/obj2yaml/ELF/invalid-section-name.yaml

@@ -11,6 +11,7 @@
 # CHECK-NEXT:   Class: ELFCLASS64
 # CHECK-NEXT:   Data:  ELFDATA2LSB
 # CHECK-NEXT:   Type:  ET_REL
+# CHECK-NEXT:   SectionHeaderStringTable: {{.*}}
 # CHECK-NEXT: Sections:
 # CHECK-NEXT:   - Name: "{{.*}}"
 # CHECK-NEXT:     Type: SHT_PROGBITS