[yaml2obj][ELF] Add support for program headers

This change adds basic support for program headers.

I need to do some testing which requires generating program headers but
I can't use ld.lld or clang to produce programs that have headers. I'd
also like to test some strange things that those programs may never
produce.

Patch by Jake Ehrlich

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@308520 91177308-0d34-0410-b5e6-96231b3b80d8

1 files changed, 93 insertions, 5 deletions

diff --git a/tools/yaml2obj/yaml2elf.cpp b/tools/yaml2obj/yaml2elf.cpp
index c89f768ed6f..1d28a1abffc 100644
--- a/tools/yaml2obj/yaml2elf.cpp
+++ b/tools/yaml2obj/yaml2elf.cpp
@@ -99,6 +99,7 @@ namespace {
 template <class ELFT>
 class ELFState {
   typedef typename object::ELFFile<ELFT>::Elf_Ehdr Elf_Ehdr;
+  typedef typename object::ELFFile<ELFT>::Elf_Phdr Elf_Phdr;
   typedef typename object::ELFFile<ELFT>::Elf_Shdr Elf_Shdr;
   typedef typename object::ELFFile<ELFT>::Elf_Sym Elf_Sym;
   typedef typename object::ELFFile<ELFT>::Elf_Rel Elf_Rel;
@@ -118,6 +119,7 @@ class ELFState {
   bool buildSymbolIndex(std::size_t &StartIndex,
                         const std::vector<ELFYAML::Symbol> &Symbols);
   void initELFHeader(Elf_Ehdr &Header);
+  void initProgramHeaders(std::vector<Elf_Phdr> &PHeaders);
   bool initSectionHeaders(std::vector<Elf_Shdr> &SHeaders,
                           ContiguousBlobAccumulator &CBA);
   void initSymtabSectionHeader(Elf_Shdr &SHeader,
@@ -125,6 +127,8 @@ class ELFState {
   void initStrtabSectionHeader(Elf_Shdr &SHeader, StringRef Name,
                                StringTableBuilder &STB,
                                ContiguousBlobAccumulator &CBA);
+  void setProgramHeaderLayout(std::vector<Elf_Phdr> &PHeaders,
+                              std::vector<Elf_Shdr> &SHeaders);
   void addSymbols(const std::vector<ELFYAML::Symbol> &Symbols,
                   std::vector<Elf_Sym> &Syms, unsigned SymbolBinding);
   void writeSectionContent(Elf_Shdr &SHeader,
@@ -173,16 +177,32 @@ void ELFState<ELFT>::initELFHeader(Elf_Ehdr &Header) {
   Header.e_machine = Doc.Header.Machine;
   Header.e_version = EV_CURRENT;
   Header.e_entry = Doc.Header.Entry;
+  Header.e_phoff = sizeof(Header);
   Header.e_flags = Doc.Header.Flags;
   Header.e_ehsize = sizeof(Elf_Ehdr);
+  Header.e_phentsize = sizeof(Elf_Phdr);
+  Header.e_phnum = Doc.ProgramHeaders.size();
   Header.e_shentsize = sizeof(Elf_Shdr);
-  // Immediately following the ELF header.
-  Header.e_shoff = sizeof(Header);
+  // Immediately following the ELF header and program headers.
+  Header.e_shoff =
+      sizeof(Header) + sizeof(Elf_Phdr) * Doc.ProgramHeaders.size();
   Header.e_shnum = getSectionCount();
   Header.e_shstrndx = getDotShStrTabSecNo();
 }
 
 template <class ELFT>
+void ELFState<ELFT>::initProgramHeaders(std::vector<Elf_Phdr> &PHeaders) {
+  for (const auto &YamlPhdr : Doc.ProgramHeaders) {
+    Elf_Phdr Phdr;
+    Phdr.p_type = YamlPhdr.Type;
+    Phdr.p_flags = YamlPhdr.Flags;
+    Phdr.p_vaddr = YamlPhdr.VAddr;
+    Phdr.p_paddr = YamlPhdr.PAddr;
+    PHeaders.push_back(Phdr);
+  }
+}
+
+template <class ELFT>
 bool ELFState<ELFT>::initSectionHeaders(std::vector<Elf_Shdr> &SHeaders,
                                         ContiguousBlobAccumulator &CBA) {
   // Ensure SHN_UNDEF entry is present. An all-zero section header is a
@@ -311,6 +331,67 @@ void ELFState<ELFT>::initStrtabSectionHeader(Elf_Shdr &SHeader, StringRef Name,
 }
 
 template <class ELFT>
+void ELFState<ELFT>::setProgramHeaderLayout(std::vector<Elf_Phdr> &PHeaders,
+                                            std::vector<Elf_Shdr> &SHeaders) {
+  uint32_t PhdrIdx = 0;
+  for (auto &YamlPhdr : Doc.ProgramHeaders) {
+    auto &PHeader = PHeaders[PhdrIdx++];
+
+    if (YamlPhdr.Sections.size())
+      PHeader.p_offset = UINT32_MAX;
+    else
+      PHeader.p_offset = 0;
+
+    // Find the minimum offset for the program header.
+    for (auto SecName : YamlPhdr.Sections) {
+      uint32_t Index = 0;
+      SN2I.lookup(SecName.Section, Index);
+      const auto &SHeader = SHeaders[Index];
+      PHeader.p_offset = std::min(PHeader.p_offset, SHeader.sh_offset);
+    }
+
+    // Find the maximum offset of the end of a section in order to set p_filesz.
+    PHeader.p_filesz = 0;
+    for (auto SecName : YamlPhdr.Sections) {
+      uint32_t Index = 0;
+      SN2I.lookup(SecName.Section, Index);
+      const auto &SHeader = SHeaders[Index];
+      uint64_t EndOfSection;
+      if (SHeader.sh_type == llvm::ELF::SHT_NOBITS)
+        EndOfSection = SHeader.sh_offset;
+      else
+        EndOfSection = SHeader.sh_offset + SHeader.sh_size;
+      uint64_t EndOfSegment = PHeader.p_offset + PHeader.p_filesz;
+      EndOfSegment = std::max(EndOfSegment, EndOfSection);
+      PHeader.p_filesz = EndOfSegment - PHeader.p_offset;
+    }
+
+    // Find the memory size by adding the size of sections at the end of the
+    // segment. These should be empty (size of zero) and NOBITS sections.
+    PHeader.p_memsz = PHeader.p_filesz;
+    for (auto SecName : YamlPhdr.Sections) {
+      uint32_t Index = 0;
+      SN2I.lookup(SecName.Section, Index);
+      const auto &SHeader = SHeaders[Index];
+      if (SHeader.sh_offset == PHeader.p_offset + PHeader.p_filesz)
+        PHeader.p_memsz += SHeader.sh_size;
+    }
+
+    // Set the alignment of the segment to be the same as the maximum alignment
+    // of the the sections with the same offset so that by default the segment
+    // has a valid and sensible alignment.
+    PHeader.p_align = 1;
+    for (auto SecName : YamlPhdr.Sections) {
+      uint32_t Index = 0;
+      SN2I.lookup(SecName.Section, Index);
+      const auto &SHeader = SHeaders[Index];
+      if (SHeader.sh_offset == PHeader.p_offset)
+        PHeader.p_align = std::max(PHeader.p_align, SHeader.sh_addralign);
+    }
+  }
+}
+
+template <class ELFT>
 void ELFState<ELFT>::addSymbols(const std::vector<ELFYAML::Symbol> &Symbols,
                                 std::vector<Elf_Sym> &Syms,
                                 unsigned SymbolBinding) {
@@ -508,12 +589,15 @@ int ELFState<ELFT>::writeELF(raw_ostream &OS, const ELFYAML::Object &Doc) {
   State.initELFHeader(Header);
 
   // TODO: Flesh out section header support.
-  // TODO: Program headers.
+
+  std::vector<Elf_Phdr> PHeaders;
+  State.initProgramHeaders(PHeaders);
 
   // XXX: This offset is tightly coupled with the order that we write
   // things to `OS`.
-  const size_t SectionContentBeginOffset =
-      Header.e_ehsize + Header.e_shentsize * Header.e_shnum;
+  const size_t SectionContentBeginOffset = Header.e_ehsize +
+                                           Header.e_phentsize * Header.e_phnum +
+                                           Header.e_shentsize * Header.e_shnum;
   ContiguousBlobAccumulator CBA(SectionContentBeginOffset);
 
   // Doc might not contain .symtab, .strtab and .shstrtab sections,
@@ -543,7 +627,11 @@ int ELFState<ELFT>::writeELF(raw_ostream &OS, const ELFYAML::Object &Doc) {
                                 CBA);
   SHeaders.push_back(ShStrTabSHeader);
 
+  // Now we can decide segment offsets
+  State.setProgramHeaderLayout(PHeaders, SHeaders);
+
   OS.write((const char *)&Header, sizeof(Header));
+  writeArrayData(OS, makeArrayRef(PHeaders));
   writeArrayData(OS, makeArrayRef(SHeaders));
   CBA.writeBlobToStream(OS);
   return 0;