[ELF]Add overflow check to ELF note iterator

llvm/include/llvm/Object/ELF.h

@@ -407,7 +407,8 @@ class ELFFile {
   Elf_Note_Iterator notes_begin(const Elf_Phdr &Phdr, Error &Err) const {
     assert(Phdr.p_type == ELF::PT_NOTE && "Phdr is not of type PT_NOTE");
     ErrorAsOutParameter ErrAsOutParam(Err);
-    if (Phdr.p_offset + Phdr.p_filesz > getBufSize()) {
+    if (Phdr.p_offset + Phdr.p_filesz > getBufSize() ||
+        Phdr.p_offset + Phdr.p_filesz < Phdr.p_offset) {
       Err =
           createError("invalid offset (0x" + Twine::utohexstr(Phdr.p_offset) +
                       ") or size (0x" + Twine::utohexstr(Phdr.p_filesz) + ")");
@@ -435,7 +436,8 @@ class ELFFile {
   Elf_Note_Iterator notes_begin(const Elf_Shdr &Shdr, Error &Err) const {
     assert(Shdr.sh_type == ELF::SHT_NOTE && "Shdr is not of type SHT_NOTE");
     ErrorAsOutParameter ErrAsOutParam(Err);
-    if (Shdr.sh_offset + Shdr.sh_size > getBufSize()) {
+    if (Shdr.sh_offset + Shdr.sh_size > getBufSize() ||
+        Shdr.sh_offset + Shdr.sh_size < Shdr.sh_offset) {
       Err =
           createError("invalid offset (0x" + Twine::utohexstr(Shdr.sh_offset) +
                       ") or size (0x" + Twine::utohexstr(Shdr.sh_size) + ")");

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,72 @@ 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);
+}