From 9c4765ca95476acfc0d9522fa1e5f3cf042badf4 Mon Sep 17 00:00:00 2001
From: Ales Katona <ales@katona.me>
Date: Sat, 16 Jan 2021 22:09:54 -0800
Subject: [PATCH] better error handling

---
 src/common/loadable.rs    |  5 ++++
 src/error.rs              | 58 ++++++++++++++++++++++++++++++++++-----
 src/formats/elf/output.rs | 22 +++++++++++----
 src/linker.rs             |  1 +
 4 files changed, 74 insertions(+), 12 deletions(-)

diff --git a/src/common/loadable.rs b/src/common/loadable.rs
index bb366dd..9d4dcfc 100644
--- a/src/common/loadable.rs
+++ b/src/common/loadable.rs
@@ -38,6 +38,7 @@ pub struct Loadable {
     rodata: SegmentSections,
     data: SegmentSections,
     bss: SegmentSections,
+    pub start_offset: Option<u64>,
 }
 
 impl Loadable {
@@ -60,6 +61,10 @@ impl Loadable {
         Self::section_names().len()
     }
 
+    pub const fn segment_count() -> usize {
+        3
+    }
+
     pub fn segment_sections(&self) -> impl Iterator<Item = &SegmentSections> {
         let text = once(&self.text);
         let rodata = once(&self.rodata);
diff --git a/src/error.rs b/src/error.rs
index 4ba7b2e..f3c678f 100644
--- a/src/error.rs
+++ b/src/error.rs
@@ -10,16 +10,33 @@ pub enum Error {
     InvalidSectionName,
     InvalidSectionData,
     ParseError(Box<dyn std::error::Error>),
-    TryFromIntError, //
+    TryFromIntError,
     LinkingError(Trace),
 }
 
 // "backtrace" for error origin
-#[derive(Debug)]
+#[derive(Debug, Default)]
 pub struct Trace {
-    origin: PathBuf,
-    offset: u64, // 0 indicates unknown/invalid
-    source_info: Option<Box<dyn SourceInfo>>,
+    pub code: u32, // specific trace code, should be >= 700
+    pub message: &'static str,
+    pub origin: PathBuf,
+    pub offset: u64, // 0 indicates unknown/invalid
+    pub source_info: Option<Box<dyn SourceInfo>>,
+}
+
+impl Error {
+    // used for programmatic input/output handling, IDE's etc.
+    pub fn code(&self) -> u32 {
+        match self {
+            Error::IOError(_) => 100,
+            Error::InvalidObjectType(_) => 200,
+            Error::InvalidSectionName => 300,
+            Error::InvalidSectionData => 400,
+            Error::ParseError(_) => 500,
+            Error::TryFromIntError => 600,
+            Error::LinkingError(trace) => trace.code,
+        }
+    }
 }
 
 pub trait SourceInfo: Debug + Display {
@@ -42,12 +59,39 @@ impl From<std::num::TryFromIntError> for Error {
 
 impl Display for Error {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "TODO!")
+        write!(f, "[{}] ", self.code())?;
+
+        match self {
+            Error::IOError(err) => write!(f, "IO error: {}", err),
+            Error::InvalidObjectType(ot) => write!(f, "Invalid object type: {}", ot),
+            Error::InvalidSectionName => write!(f, "Invalid section name"),
+            Error::InvalidSectionData => write!(f, "Invalid section data"),
+            Error::ParseError(err) => write!(f, "Parse error: {}", err),
+            Error::TryFromIntError => write!(f, "Integer conversion error"),
+            Error::LinkingError(trace) => write!(f, "Linking error: {}", trace),
+        }
     }
 }
 
 impl Display for Trace {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
-        write!(f, "TODO!")
+        let mut new_line = ""; // new line in case of previous info only
+
+        if let Some(origin) = self.origin.to_str() {
+            new_line = "\n";
+            write!(f, "\n\t{}", origin)?;
+        }
+
+        if self.offset > 0 {
+            new_line = "\n";
+            write!(f, "\n@offset: {}", self.offset)?;
+        }
+
+        if let Some(si) = &self.source_info {
+            new_line = "\n";
+            write!(f, "\n{}", si)?;
+        }
+
+        write!(f, "{}{}", new_line, self.message)
     }
 }
diff --git a/src/formats/elf/output.rs b/src/formats/elf/output.rs
index a13d052..480b693 100644
--- a/src/formats/elf/output.rs
+++ b/src/formats/elf/output.rs
@@ -47,10 +47,10 @@ impl Output<ElfObject> for ElfOutput {
 
         let strtab = make_strtab();
         let mut ehdr = make_elf_header();
-        ehdr.e_shnum = Loadable::section_count() as u16 + 1;
-        ehdr.e_phnum = 3; // TODO: probably move to loadable?
-        ehdr.e_shstrndx = ehdr.e_shnum - 1;
-        ehdr.e_entry = 4096; // TODO: find .start symbol eventual location
+        ehdr.e_shnum = Loadable::section_count() as u16 + 1; // +1 for .shstrtab
+        ehdr.e_phnum = Loadable::segment_count() as u16;
+        ehdr.e_shstrndx = ehdr.e_shnum - 1; // .shstrab is always last
+        ehdr.e_entry = get_start_offset(loadable)?;
         ehdr.e_phentsize = PHS;
         ehdr.e_shentsize = SHS;
         ehdr.e_phoff = EHS;
@@ -77,7 +77,6 @@ impl Output<ElfObject> for ElfOutput {
         let mut name_idx = 0;
         // write section header table (text + rodata + data + bss)
         for sections in loadable.segment_sections() {
-            // TODO: find out data_offset for these headers
             let sh = make_section_header(
                 strtab.1[name_idx],
                 ShType::ProgBits,
@@ -216,6 +215,19 @@ fn make_strtab() -> (Vec<u8>, Vec<usize>) {
     (strtab_bytes, indexes)
 }
 
+fn get_start_offset(loadable: &Loadable) -> Result<u64, Error> {
+    use crate::error::Trace;
+
+    loadable.start_offset.ok_or_else(|| {
+        let trace = Trace {
+            message: "Program entrypoint not found",
+            ..Default::default()
+        };
+
+        Error::LinkingError(trace)
+    })
+}
+
 fn expand_path(path: &Path) -> Result<&str, Error> {
     use std::io::Error as IOError;
     use std::io::ErrorKind;
diff --git a/src/linker.rs b/src/linker.rs
index e88f4a9..288c9d8 100644
--- a/src/linker.rs
+++ b/src/linker.rs
@@ -48,6 +48,7 @@ where
             }
         }
 
+        self.loadable.start_offset = Some(4096); // TODO: get from .start symbol location
         self.output.finalize(&self.relocatables, &self.loadable)
     }