[llvm] [llvm][ELF]Add Shdr check for getBuildID (PR #126537)

[Ruoyu Qiu via llvm-commits]

https://github.com/cabbaken updated https://github.com/llvm/llvm-project/pull/126537

>From a19e1a05cde04298aee6732811353d1814d21eb8 Mon Sep 17 00:00:00 2001
From: Ruoyu Qiu <cabbaken at outlook.com>
Date: Mon, 10 Feb 2025 15:56:31 +0000
Subject: [PATCH 1/3] [llvm-objdump][ELF]Add Shdr buildID check(#126418)

Add Section Header check for getBuildID, fix crash
with invalid Program Header.

Signed-off-by: Ruoyu Qiu <cabbaken at outlook.com>
---
 llvm/lib/Object/BuildID.cpp | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/llvm/lib/Object/BuildID.cpp b/llvm/lib/Object/BuildID.cpp
index 89d6bc3ab550d..d446057129827 100644
--- a/llvm/lib/Object/BuildID.cpp
+++ b/llvm/lib/Object/BuildID.cpp
@@ -24,6 +24,20 @@ using namespace llvm::object;
 namespace {
 
 template <typename ELFT> BuildIDRef getBuildID(const ELFFile<ELFT> &Obj) {
+  auto Sections = cantFail(Obj.sections());
+  if (!Sections.empty()) {
+    for (const auto &S : Sections) {
+      if (S.sh_type != ELF::SHT_NOTE)
+        continue;
+      Error Err = Error::success();
+      for (auto N : Obj.notes(S, Err))
+        if (N.getType() == ELF::NT_GNU_BUILD_ID &&
+            N.getName() == ELF::ELF_NOTE_GNU)
+          return N.getDesc(S.sh_addralign);
+      consumeError(std::move(Err));
+    }
+  }
+
   auto PhdrsOrErr = Obj.program_headers();
   if (!PhdrsOrErr) {
     consumeError(PhdrsOrErr.takeError());

>From 169f9db3abaf765d58a4f7fd3e8cbc44469fafbe Mon Sep 17 00:00:00 2001
From: Ruoyu Qiu <cabbaken at outlook.com>
Date: Tue, 11 Feb 2025 08:29:19 +0000
Subject: [PATCH 2/3] [LLVM][ELF]Add Align check for getDesc()

Add Aligh check to avoid crashing if Aligh is 0.

Signed-off-by: Ruoyu Qiu <cabbaken at outlook.com>
---
 llvm/include/llvm/Object/ELFTypes.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/include/llvm/Object/ELFTypes.h b/llvm/include/llvm/Object/ELFTypes.h
index 87e4dbe448091..8441009900316 100644
--- a/llvm/include/llvm/Object/ELFTypes.h
+++ b/llvm/include/llvm/Object/ELFTypes.h
@@ -660,7 +660,7 @@ class Elf_Note_Impl {
 
   /// Get the note's descriptor.
   ArrayRef<uint8_t> getDesc(size_t Align) const {
-    if (!Nhdr.n_descsz)
+    if (!Nhdr.n_descsz || !Align)
       return ArrayRef<uint8_t>();
     return ArrayRef<uint8_t>(
         reinterpret_cast<const uint8_t *>(&Nhdr) +

>From 0c6d355e70a8d563c42a53b33dc87854d91b4cba Mon Sep 17 00:00:00 2001
From: Ruoyu Qiu <cabbaken at outlook.com>
Date: Wed, 12 Feb 2025 10:49:06 +0000
Subject: [PATCH 3/3] Optimize getBuildID() code style

Signed-off-by: Ruoyu Qiu <cabbaken at outlook.com>
---
 llvm/lib/Object/BuildID.cpp | 30 +++++++++++++++++-------------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/Object/BuildID.cpp b/llvm/lib/Object/BuildID.cpp
index d446057129827..08f1502b362b0 100644
--- a/llvm/lib/Object/BuildID.cpp
+++ b/llvm/lib/Object/BuildID.cpp
@@ -24,20 +24,27 @@ using namespace llvm::object;
 namespace {
 
 template <typename ELFT> BuildIDRef getBuildID(const ELFFile<ELFT> &Obj) {
+  auto findBuildID = [&Obj](const auto &ShdrOrPhdr,
+                            uint64_t Alignment) -> std::optional<BuildIDRef> {
+    Error Err = Error::success();
+    for (auto N : Obj.notes(ShdrOrPhdr, Err))
+      if (N.getType() == ELF::NT_GNU_BUILD_ID &&
+          N.getName() == ELF::ELF_NOTE_GNU)
+        return N.getDesc(Alignment);
+    consumeError(std::move(Err));
+    return std::nullopt;
+  };
+
   auto Sections = cantFail(Obj.sections());
   if (!Sections.empty()) {
     for (const auto &S : Sections) {
       if (S.sh_type != ELF::SHT_NOTE)
         continue;
-      Error Err = Error::success();
-      for (auto N : Obj.notes(S, Err))
-        if (N.getType() == ELF::NT_GNU_BUILD_ID &&
-            N.getName() == ELF::ELF_NOTE_GNU)
-          return N.getDesc(S.sh_addralign);
-      consumeError(std::move(Err));
+      auto ShdrRes = findBuildID(S, S.sh_addralign);
+      if (ShdrRes)
+        return ShdrRes.value();
     }
   }
-
   auto PhdrsOrErr = Obj.program_headers();
   if (!PhdrsOrErr) {
     consumeError(PhdrsOrErr.takeError());
@@ -46,12 +53,9 @@ template <typename ELFT> BuildIDRef getBuildID(const ELFFile<ELFT> &Obj) {
   for (const auto &P : *PhdrsOrErr) {
     if (P.p_type != ELF::PT_NOTE)
       continue;
-    Error Err = Error::success();
-    for (auto N : Obj.notes(P, Err))
-      if (N.getType() == ELF::NT_GNU_BUILD_ID &&
-          N.getName() == ELF::ELF_NOTE_GNU)
-        return N.getDesc(P.p_align);
-    consumeError(std::move(Err));
+    auto PhdrRes = findBuildID(P, P.p_align);
+    if (PhdrRes)
+      return PhdrRes.value();
   }
   return {};
 }