From 2e9b29199af78edae57ddb04647870c4b46b05d0 Mon Sep 17 00:00:00 2001
From: Philip Craig <philipjcraig@gmail.com>
Date: Mon, 18 Mar 2024 13:13:31 +1000
Subject: [PATCH] build/elf: improve handling of sh_offset for SHT_NOBITS

---
 src/build/elf.rs     | 33 +++++++++++++++------------------
 tests/build/elf.rs   | 39 +++++++++++++++++++++++++++++++++++++++
 tests/build/mod.rs   |  3 +++
 tests/integration.rs |  1 +
 4 files changed, 58 insertions(+), 18 deletions(-)
 create mode 100644 tests/build/elf.rs
 create mode 100644 tests/build/mod.rs

diff --git a/src/build/elf.rs b/src/build/elf.rs
index 04d2862b..3272c668 100644
--- a/src/build/elf.rs
+++ b/src/build/elf.rs
@@ -1094,20 +1094,19 @@ impl<'data> Builder<'data> {
                 .collect();
             // The data for alloc sections may need to be written in a different order
             // from their section headers.
-            alloc_sections.sort_by_key(|index| {
-                let section = &self.sections.get(out_sections[*index].id);
-                // SHT_NOBITS sections need to come before other sections at the same offset.
-                let file_size = if section.sh_type == elf::SHT_NOBITS {
-                    0
-                } else {
-                    section.sh_size
-                };
-                (section.sh_offset, file_size)
-            });
+            alloc_sections
+                .sort_by_key(|index| self.sections.get(out_sections[*index].id).sh_offset);
             for index in &alloc_sections {
                 let out_section = &mut out_sections[*index];
                 let section = &self.sections.get(out_section.id);
 
+                if section.sh_type == elf::SHT_NOBITS {
+                    // sh_offset is meaningless for SHT_NOBITS, so preserve the input
+                    // value without checking it.
+                    out_section.offset = section.sh_offset as usize;
+                    continue;
+                }
+
                 if section.sh_offset < writer.reserved_len() as u64 {
                     return Err(Error(format!(
                         "Unsupported sh_offset value 0x{:x} for section '{}', expected at least 0x{:x}",
@@ -1123,10 +1122,6 @@ impl<'data> Builder<'data> {
                     SectionData::Data(data) => {
                         writer.reserve(data.len(), section.sh_addralign as usize)
                     }
-                    SectionData::UninitializedData(_) => {
-                        // Note: unaligned input sh_offset was observed in practice.
-                        writer.reserve(0, 1)
-                    }
                     SectionData::DynamicRelocation(relocations) => writer
                         .reserve_relocations(relocations.len(), section.sh_type == elf::SHT_RELA),
                     SectionData::Note(data) => {
@@ -1191,9 +1186,7 @@ impl<'data> Builder<'data> {
                 SectionData::Data(data) => {
                     writer.reserve(data.len(), section.sh_addralign as usize)
                 }
-                SectionData::UninitializedData(_) => {
-                    writer.reserve(0, section.sh_addralign as usize)
-                }
+                SectionData::UninitializedData(_) => writer.reserved_len(),
                 SectionData::Note(data) => {
                     writer.reserve(data.len(), section.sh_addralign as usize)
                 }
@@ -1268,12 +1261,16 @@ impl<'data> Builder<'data> {
             for index in &alloc_sections {
                 let out_section = &mut out_sections[*index];
                 let section = self.sections.get(out_section.id);
+
+                if section.sh_type == elf::SHT_NOBITS {
+                    continue;
+                }
+
                 writer.pad_until(out_section.offset);
                 match &section.data {
                     SectionData::Data(data) => {
                         writer.write(data);
                     }
-                    SectionData::UninitializedData(_) => {}
                     SectionData::DynamicRelocation(relocations) => {
                         for rel in relocations {
                             let r_sym = if let Some(symbol) = rel.symbol {
diff --git a/tests/build/elf.rs b/tests/build/elf.rs
new file mode 100644
index 00000000..be56e675
--- /dev/null
+++ b/tests/build/elf.rs
@@ -0,0 +1,39 @@
+use object::{build, elf};
+
+// Test that offset 0 is supported for SHT_NOBITS sections.
+#[test]
+fn test_nobits_offset() {
+    let mut builder = build::elf::Builder::new(object::Endianness::Little, true);
+    builder.header.e_type = elf::ET_EXEC;
+    builder.header.e_phoff = 0x40;
+
+    let section = builder.sections.add();
+    section.name = b".shstrtab"[..].into();
+    section.sh_type = elf::SHT_STRTAB;
+    section.data = build::elf::SectionData::SectionString;
+
+    let section = builder.sections.add();
+    section.name = b".bss"[..].into();
+    section.sh_type = elf::SHT_NOBITS;
+    section.sh_flags = (elf::SHF_ALLOC | elf::SHF_WRITE) as u64;
+    section.sh_addr = 0x1000;
+    section.sh_offset = 0;
+    section.sh_size = 0x1000;
+    section.sh_addralign = 16;
+    section.data = build::elf::SectionData::UninitializedData(0x1000);
+    let section_id = section.id();
+
+    let segment = builder.segments.add();
+    segment.p_type = elf::PT_LOAD;
+    segment.p_flags = elf::PF_R | elf::PF_W;
+    segment.p_offset = 0x1000;
+    segment.p_vaddr = 0x1000;
+    segment.p_paddr = 0x1000;
+    segment.p_filesz = 0;
+    segment.p_memsz = 0x1000;
+    segment.p_align = 16;
+    segment.sections.push(section_id);
+
+    let mut buf = Vec::new();
+    builder.write(&mut buf).unwrap();
+}
diff --git a/tests/build/mod.rs b/tests/build/mod.rs
new file mode 100644
index 00000000..6a777367
--- /dev/null
+++ b/tests/build/mod.rs
@@ -0,0 +1,3 @@
+#![cfg(feature = "build")]
+
+mod elf;
diff --git a/tests/integration.rs b/tests/integration.rs
index 6ebcb547..560ba626 100644
--- a/tests/integration.rs
+++ b/tests/integration.rs
@@ -1,2 +1,3 @@
+mod build;
 mod read;
 mod round_trip;