1 files changed, 72 insertions, 0 deletions
diff --git a/llvm/unittests/Object/ELFTest.cpp b/llvm/unittests/Object/ELFTest.cpp
index faf855c..7c68ab5 100644
--- a/llvm/unittests/Object/ELFTest.cpp
+++ b/llvm/unittests/Object/ELFTest.cpp
@@ -7,6 +7,10 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Object/ELF.h"
+#include "llvm/Object/ELFObjectFile.h"
+#include "llvm/ObjectYAML/yaml2obj.h"
+#include "llvm/Support/Error.h"
+#include "llvm/Support/YAMLTraits.h"
 #include "llvm/Testing/Support/Error.h"
 #include "gtest/gtest.h"
 
@@ -310,3 +314,71 @@ TEST(ELFTest, Hash) {
   // presuming 32-bit long. Thus make sure that extra bit doesn't appear. 
   EXPECT_EQ(hashSysV("ZZZZZW9p"), 0U);
 }
+
+template <class ELFT>
+static Expected<ELFObjectFile<ELFT>> toBinary(SmallVectorImpl<char> &Storage,
+                                              StringRef Yaml) {
+  raw_svector_ostream OS(Storage);
+  yaml::Input YIn(Yaml);
+  if (!yaml::convertYAML(YIn, OS, [](const Twine &Msg) {}))
+    return createStringError(std::errc::invalid_argument,
+                             "unable to convert YAML");
+  return ELFObjectFile<ELFT>::create(MemoryBufferRef(OS.str(), "dummyELF"));
+}
+
+TEST(ELFObjectFileTest, ELFNoteIteratorOverflow) {
+  using Elf_Shdr_Range = ELFFile<ELF64LE>::Elf_Shdr_Range;
+  using Elf_Phdr_Range = ELFFile<ELF64LE>::Elf_Phdr_Range;
+
+  SmallString<0> Storage;
+  Expected<ELFObjectFile<ELF64LE>> ElfOrErr = toBinary<ELF64LE>(Storage, R"(
+--- !ELF
+FileHeader:
+  Class:          ELFCLASS64
+  Data:           ELFDATA2LSB
+  Type:           ET_EXEC
+  Machine:        EM_X86_64
+ProgramHeaders:
+  - Type:         PT_NOTE
+    FileSize:     0xffffffffffffff88
+    FirstSec:     .note.gnu.build-id
+    LastSec:      .note.gnu.build-id
+Sections:
+  - Name:         .note.gnu.build-id
+    Type:         SHT_NOTE
+    AddressAlign: 0x04
+    ShOffset:     0xffffffffffffff88
+    Notes:
+      - Name:     "GNU"
+        Desc:     "abb50d82b6bdc861"
+        Type:     3
+)");
+  ASSERT_THAT_EXPECTED(ElfOrErr, Succeeded());
+  ELFFile<ELF64LE> Obj = ElfOrErr.get().getELFFile();
+
+  auto CheckOverflow = [&](auto &&PhdrOrShdr, uint64_t Offset, uint64_t Size) {
+    Error Err = Error::success();
+    Obj.notes(PhdrOrShdr, Err);
+
+    std::string ErrMessage;
+    handleAllErrors(std::move(Err), [&](const ErrorInfoBase &EI) {
+      ErrMessage = EI.message();
+    });
+
+    EXPECT_EQ(ErrMessage, ("invalid offset (0x" + Twine::utohexstr(Offset) +
+                           ") or size (0x" + Twine::utohexstr(Size) + ")")
+                              .str());
+  };
+
+  Expected<Elf_Phdr_Range> PhdrsOrErr = Obj.program_headers();
+  EXPECT_FALSE(!PhdrsOrErr);
+  for (Elf_Phdr_Impl<ELF64LE> P : *PhdrsOrErr)
+    if (P.p_type == ELF::PT_NOTE)
+      CheckOverflow(P, P.p_offset, P.p_filesz);
+
+  Expected<Elf_Shdr_Range> ShdrsOrErr = Obj.sections();
+  EXPECT_FALSE(!ShdrsOrErr);
+  for (Elf_Shdr_Impl<ELF64LE> S : *ShdrsOrErr)
+    if (S.sh_type == ELF::SHT_NOTE)
+      CheckOverflow(S, S.sh_offset, S.sh_size);
+}