From 6b6d261d4e93ed7e590865f8dd9c9e162f2c455c Mon Sep 17 00:00:00 2001
From: Ales Katona <ales@katona.me>
Date: Wed, 14 Apr 2021 10:27:08 -0700
Subject: [PATCH] fix section matching offsets + sizes

---
 src/common.rs             | 21 +++++++++++----------
 src/common/loadable.rs    |  4 ++++
 src/formats/elf/output.rs | 32 +++++++++++++++++++++++++-------
 3 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/src/common.rs b/src/common.rs
index 5d10f41..a30dc72 100644
--- a/src/common.rs
+++ b/src/common.rs
@@ -4,12 +4,7 @@ mod relocatable;
 mod section;
 mod symbol;
 
-use std::{
-    convert::TryInto,
-    fs::File,
-    io::{BufWriter, Seek, SeekFrom},
-    path::Path,
-};
+use std::path::Path;
 
 pub use loadable::*;
 pub use output::*;
@@ -30,13 +25,19 @@ pub fn expand_path(path: &Path) -> Result<&str, Error> {
     })
 }
 
-pub fn pad_to_next_page(writer: &mut BufWriter<File>, offset: usize) -> Result<usize, Error> {
+pub fn page_padding(offset: usize) -> usize {
     let page_size = page_size::get();
 
     let padding = page_size - (offset % page_size);
 
-    eprintln!("Padding from: {} with: {}", offset, padding);
-    writer.seek(SeekFrom::Current(padding.try_into()?))?;
+    padding
+}
+
+pub fn next_page(offset: usize) -> usize {
+    let page_size = page_size::get();
 
-    Ok(padding)
+    let page_count = (offset / page_size) + 1;
+    
+    page_count * page_size
 }
+
diff --git a/src/common/loadable.rs b/src/common/loadable.rs
index caa2323..fd1f909 100644
--- a/src/common/loadable.rs
+++ b/src/common/loadable.rs
@@ -133,6 +133,10 @@ impl Loadable {
         }
     }
 
+    pub const fn sections_count() -> usize {
+        4
+    }
+
     pub fn segment_sections(&self) -> impl Iterator<Item = &SegmentSections> {
         let text = once(&self.text);
         let rodata = once(&self.rodata);
diff --git a/src/formats/elf/output.rs b/src/formats/elf/output.rs
index f43cf72..f06bca2 100644
--- a/src/formats/elf/output.rs
+++ b/src/formats/elf/output.rs
@@ -14,7 +14,7 @@ use elf_utilities::{
 
 use crate::{
     common::{
-        expand_path, pad_to_next_page, Loadable, Output, Relocatable, SegmentType, SymbolValue,
+        expand_path, next_page, page_padding, Loadable, Output, Relocatable, SegmentType, SymbolValue,
     },
     error::Error,
 };
@@ -50,7 +50,7 @@ impl ElfOutput {
         const SHS: u16 = size_of::<Shdr64>() as u16;
 
         // write .shstrtab header
-        let data_offset = (offset + usize::from(SHS)) as u64;
+        let data_offset = offset + usize::from(SHS);
         let strtab_header = make_section_header(
             shstrtab.1[name_idx],
             ShType::StrTab,
@@ -77,7 +77,7 @@ impl Output<ElfObject> for ElfOutput {
 
         let shstrtab = make_shstrtab();
         let mut ehdr = make_elf_header();
-        ehdr.e_shnum = SECTION_NAMES.len() as u16;
+        ehdr.e_shnum = (SECTION_NAMES.len() + 1) as u16;
         ehdr.e_phnum = 3u16; // .text, .rodata + .data, .bss
         ehdr.e_shstrndx = ehdr.e_shnum - 1; // .shstrab is always last
         ehdr.e_entry = loadable.start_offset()? as u64;
@@ -104,17 +104,23 @@ impl Output<ElfObject> for ElfOutput {
         offset += self.writer.write(&ph.to_le_bytes())?;
 
         eprintln!("SH start: {}", offset);
-        let mut name_idx = 0;
+        let mut name_idx = 0; // TODO: fix to make this empty!
+        let mut sec_offset = next_page(offset + (SHS as usize) * Loadable::sections_count());
+        eprintln!("Sec_offset: {}", sec_offset);
+        // write 0-section
+        let empty_sh = make_section_header(0, ShType::Null, 0, 0);
+        offset += self.writer.write(&empty_sh.to_le_bytes())?;
         // write section header table (text + rodata + data + bss)
         for sections in loadable.segment_sections() {
             let sh = make_section_header(
                 shstrtab.1[name_idx],
                 ShType::ProgBits,
-                0,
+                sec_offset,
                 sections.data_size(),
             );
             name_idx += 1;
             offset += self.writer.write(&sh.to_le_bytes())?;
+            sec_offset += sections.data_size() as usize;
         }
 
         // .shstrtab as last section header written (+ the name data including)
@@ -130,6 +136,7 @@ impl Output<ElfObject> for ElfOutput {
                 .get(si.object_index)
                 .ok_or(Error::InvalidObjectIndex)?;
             let bytes = object.bytes(si.offset, si.file_size)?;
+            // eprintln!("Writing program section from {} offset {} size {}\n\t{:x?}", si.object_index, si.offset, bytes.len(), bytes);
             offset += self.writer.write(bytes)?;
         }
 
@@ -210,7 +217,7 @@ fn make_elf_header() -> Ehdr64 {
     ehdr
 }
 
-fn make_section_header(name_idx: usize, stype: ShType, offset: u64, size: u64) -> Shdr64 {
+fn make_section_header(name_idx: usize, stype: ShType, offset: usize, size: u64) -> Shdr64 {
     let mut h = Shdr64::default();
     h.set_type(stype);
 
@@ -222,7 +229,7 @@ fn make_section_header(name_idx: usize, stype: ShType, offset: u64, size: u64) -
     h.sh_link = 0;
     h.sh_size = size;
     // filled at finalize()
-    h.sh_offset = offset;
+    h.sh_offset = offset as u64;
 
     h
 }
@@ -270,3 +277,14 @@ fn make_shstrtab() -> (Vec<u8>, Vec<usize>) {
 
     (strtab_bytes, indexes)
 }
+
+fn pad_to_next_page(writer: &mut BufWriter<File>, offset: usize) -> Result<usize, Error> {
+    use std::convert::TryInto;
+
+    let padding = page_padding(offset);
+
+    eprintln!("Padding from: {} with: {}", offset, padding);
+    writer.seek(SeekFrom::Current(padding.try_into()?))?;
+
+    Ok(padding)
+}
\ No newline at end of file