1 files changed, 168 insertions, 30 deletions

diff --git a/lib/Object/COFFObjectFile.cpp b/lib/Object/COFFObjectFile.cpp
index 854664e679df..2c0f6dc2b1e9 100644
--- a/lib/Object/COFFObjectFile.cpp
+++ b/lib/Object/COFFObjectFile.cpp
@@ -937,29 +937,6 @@ iterator_range<base_reloc_iterator> COFFObjectFile::base_relocs() const {
 }
 
 std::error_code
-COFFObjectFile::getCOFFHeader(const coff_file_header *&Res) const {
-  Res = COFFHeader;
-  return std::error_code();
-}
-
-std::error_code
-COFFObjectFile::getCOFFBigObjHeader(const coff_bigobj_file_header *&Res) const {
-  Res = COFFBigObjHeader;
-  return std::error_code();
-}
-
-std::error_code COFFObjectFile::getPE32Header(const pe32_header *&Res) const {
-  Res = PE32Header;
-  return std::error_code();
-}
-
-std::error_code
-COFFObjectFile::getPE32PlusHeader(const pe32plus_header *&Res) const {
-  Res = PE32PlusHeader;
-  return std::error_code();
-}
-
-std::error_code
 COFFObjectFile::getDataDirectory(uint32_t Index,
                                  const data_directory *&Res) const {
   // Error if there's no data directory or the index is out of range.
@@ -994,11 +971,12 @@ std::error_code COFFObjectFile::getSection(int32_t Index,
 std::error_code COFFObjectFile::getSection(StringRef SectionName,
                                            const coff_section *&Result) const {
   Result = nullptr;
-  StringRef SecName;
   for (const SectionRef &Section : sections()) {
-    if (std::error_code E = Section.getName(SecName))
-      return E;
-    if (SecName == SectionName) {
+    auto NameOrErr = Section.getName();
+    if (!NameOrErr)
+      return errorToErrorCode(NameOrErr.takeError());
+
+    if (*NameOrErr == SectionName) {
       Result = getCOFFSection(Section);
       return std::error_code();
     }
@@ -1684,9 +1662,12 @@ std::error_code BaseRelocRef::getRVA(uint32_t &Result) const {
   return std::error_code();
 }
 
-#define RETURN_IF_ERROR(E)                                                     \
-  if (E)                                                                       \
-    return E;
+#define RETURN_IF_ERROR(Expr)                                                  \
+  do {                                                                         \
+    Error E = (Expr);                                                          \
+    if (E)                                                                     \
+      return std::move(E);                                                     \
+  } while (0)
 
 Expected<ArrayRef<UTF16>>
 ResourceSectionRef::getDirStringAtOffset(uint32_t Offset) {
@@ -1715,11 +1696,168 @@ ResourceSectionRef::getTableAtOffset(uint32_t Offset) {
   return *Table;
 }
 
+Expected<const coff_resource_dir_entry &>
+ResourceSectionRef::getTableEntryAtOffset(uint32_t Offset) {
+  const coff_resource_dir_entry *Entry = nullptr;
+
+  BinaryStreamReader Reader(BBS);
+  Reader.setOffset(Offset);
+  RETURN_IF_ERROR(Reader.readObject(Entry));
+  assert(Entry != nullptr);
+  return *Entry;
+}
+
+Expected<const coff_resource_data_entry &>
+ResourceSectionRef::getDataEntryAtOffset(uint32_t Offset) {
+  const coff_resource_data_entry *Entry = nullptr;
+
+  BinaryStreamReader Reader(BBS);
+  Reader.setOffset(Offset);
+  RETURN_IF_ERROR(Reader.readObject(Entry));
+  assert(Entry != nullptr);
+  return *Entry;
+}
+
 Expected<const coff_resource_dir_table &>
 ResourceSectionRef::getEntrySubDir(const coff_resource_dir_entry &Entry) {
+  assert(Entry.Offset.isSubDir());
   return getTableAtOffset(Entry.Offset.value());
 }
 
+Expected<const coff_resource_data_entry &>
+ResourceSectionRef::getEntryData(const coff_resource_dir_entry &Entry) {
+  assert(!Entry.Offset.isSubDir());
+  return getDataEntryAtOffset(Entry.Offset.value());
+}
+
 Expected<const coff_resource_dir_table &> ResourceSectionRef::getBaseTable() {
   return getTableAtOffset(0);
 }
+
+Expected<const coff_resource_dir_entry &>
+ResourceSectionRef::getTableEntry(const coff_resource_dir_table &Table,
+                                  uint32_t Index) {
+  if (Index >= (uint32_t)(Table.NumberOfNameEntries + Table.NumberOfIDEntries))
+    return createStringError(object_error::parse_failed, "index out of range");
+  const uint8_t *TablePtr = reinterpret_cast<const uint8_t *>(&Table);
+  ptrdiff_t TableOffset = TablePtr - BBS.data().data();
+  return getTableEntryAtOffset(TableOffset + sizeof(Table) +
+                               Index * sizeof(coff_resource_dir_entry));
+}
+
+Error ResourceSectionRef::load(const COFFObjectFile *O) {
+  for (const SectionRef &S : O->sections()) {
+    Expected<StringRef> Name = S.getName();
+    if (!Name)
+      return Name.takeError();
+
+    if (*Name == ".rsrc" || *Name == ".rsrc$01")
+      return load(O, S);
+  }
+  return createStringError(object_error::parse_failed,
+                           "no resource section found");
+}
+
+Error ResourceSectionRef::load(const COFFObjectFile *O, const SectionRef &S) {
+  Obj = O;
+  Section = S;
+  Expected<StringRef> Contents = Section.getContents();
+  if (!Contents)
+    return Contents.takeError();
+  BBS = BinaryByteStream(*Contents, support::little);
+  const coff_section *COFFSect = Obj->getCOFFSection(Section);
+  ArrayRef<coff_relocation> OrigRelocs = Obj->getRelocations(COFFSect);
+  Relocs.reserve(OrigRelocs.size());
+  for (const coff_relocation &R : OrigRelocs)
+    Relocs.push_back(&R);
+  std::sort(Relocs.begin(), Relocs.end(),
+            [](const coff_relocation *A, const coff_relocation *B) {
+              return A->VirtualAddress < B->VirtualAddress;
+            });
+  return Error::success();
+}
+
+Expected<StringRef>
+ResourceSectionRef::getContents(const coff_resource_data_entry &Entry) {
+  if (!Obj)
+    return createStringError(object_error::parse_failed, "no object provided");
+
+  // Find a potential relocation at the DataRVA field (first member of
+  // the coff_resource_data_entry struct).
+  const uint8_t *EntryPtr = reinterpret_cast<const uint8_t *>(&Entry);
+  ptrdiff_t EntryOffset = EntryPtr - BBS.data().data();
+  coff_relocation RelocTarget{ulittle32_t(EntryOffset), ulittle32_t(0),
+                              ulittle16_t(0)};
+  auto RelocsForOffset =
+      std::equal_range(Relocs.begin(), Relocs.end(), &RelocTarget,
+                       [](const coff_relocation *A, const coff_relocation *B) {
+                         return A->VirtualAddress < B->VirtualAddress;
+                       });
+
+  if (RelocsForOffset.first != RelocsForOffset.second) {
+    // We found a relocation with the right offset. Check that it does have
+    // the expected type.
+    const coff_relocation &R = **RelocsForOffset.first;
+    uint16_t RVAReloc;
+    switch (Obj->getMachine()) {
+    case COFF::IMAGE_FILE_MACHINE_I386:
+      RVAReloc = COFF::IMAGE_REL_I386_DIR32NB;
+      break;
+    case COFF::IMAGE_FILE_MACHINE_AMD64:
+      RVAReloc = COFF::IMAGE_REL_AMD64_ADDR32NB;
+      break;
+    case COFF::IMAGE_FILE_MACHINE_ARMNT:
+      RVAReloc = COFF::IMAGE_REL_ARM_ADDR32NB;
+      break;
+    case COFF::IMAGE_FILE_MACHINE_ARM64:
+      RVAReloc = COFF::IMAGE_REL_ARM64_ADDR32NB;
+      break;
+    default:
+      return createStringError(object_error::parse_failed,
+                               "unsupported architecture");
+    }
+    if (R.Type != RVAReloc)
+      return createStringError(object_error::parse_failed,
+                               "unexpected relocation type");
+    // Get the relocation's symbol
+    Expected<COFFSymbolRef> Sym = Obj->getSymbol(R.SymbolTableIndex);
+    if (!Sym)
+      return Sym.takeError();
+    const coff_section *Section = nullptr;
+    // And the symbol's section
+    if (std::error_code EC = Obj->getSection(Sym->getSectionNumber(), Section))
+      return errorCodeToError(EC);
+    // Add the initial value of DataRVA to the symbol's offset to find the
+    // data it points at.
+    uint64_t Offset = Entry.DataRVA + Sym->getValue();
+    ArrayRef<uint8_t> Contents;
+    if (Error E = Obj->getSectionContents(Section, Contents))
+      return std::move(E);
+    if (Offset + Entry.DataSize > Contents.size())
+      return createStringError(object_error::parse_failed,
+                               "data outside of section");
+    // Return a reference to the data inside the section.
+    return StringRef(reinterpret_cast<const char *>(Contents.data()) + Offset,
+                     Entry.DataSize);
+  } else {
+    // Relocatable objects need a relocation for the DataRVA field.
+    if (Obj->isRelocatableObject())
+      return createStringError(object_error::parse_failed,
+                               "no relocation found for DataRVA");
+
+    // Locate the section that contains the address that DataRVA points at.
+    uint64_t VA = Entry.DataRVA + Obj->getImageBase();
+    for (const SectionRef &S : Obj->sections()) {
+      if (VA >= S.getAddress() &&
+          VA + Entry.DataSize <= S.getAddress() + S.getSize()) {
+        uint64_t Offset = VA - S.getAddress();
+        Expected<StringRef> Contents = S.getContents();
+        if (!Contents)
+          return Contents.takeError();
+        return Contents->slice(Offset, Offset + Entry.DataSize);
+      }
+    }
+    return createStringError(object_error::parse_failed,
+                             "address not found in image");
+  }
+}