From 0d81a006fc63ee822028bdabc565a985c2f7390c Mon Sep 17 00:00:00 2001
From: wenxuanjun <wenxuanjun@yahoo.com>
Date: Sun, 3 Nov 2024 08:08:21 +0800
Subject: [PATCH] No alloc needed & Return entry point

---
 .cargo/config.toml |  2 +-
 Cargo.lock         | 21 ++++--------
 Cargo.toml         |  5 ++-
 src/lib.rs         | 85 +++++++++++++++++-----------------------------
 4 files changed, 41 insertions(+), 72 deletions(-)

diff --git a/.cargo/config.toml b/.cargo/config.toml
index 4718b08..57e1098 100644
--- a/.cargo/config.toml
+++ b/.cargo/config.toml
@@ -1,6 +1,6 @@
 [unstable]
 build-std-features = ["compiler-builtins-mem"]
-build-std = ["core", "compiler_builtins", "alloc"]
+build-std = ["core", "compiler_builtins"]
 
 [build]
 target = ["i686-unknown-none.json", "x86_64-unknown-none"]
diff --git a/Cargo.lock b/Cargo.lock
index 4abb117..febda39 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3,23 +3,14 @@
 version = 4
 
 [[package]]
-name = "libelf_parse"
-version = "0.0.0"
-dependencies = [
- "object",
-]
-
-[[package]]
-name = "memchr"
-version = "2.7.4"
+name = "elf"
+version = "0.7.4"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "78ca9ab1a0babb1e7d5695e3530886289c18cf2f87ec19a575a0abdce112e3a3"
+checksum = "4445909572dbd556c457c849c4ca58623d84b27c8fff1e74b0b4227d8b90d17b"
 
 [[package]]
-name = "object"
-version = "0.36.5"
-source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "aedf0a2d09c573ed1d8d85b30c119153926a2b36dce0ab28322c09a117a4683e"
+name = "libelf_parse"
+version = "0.0.0"
 dependencies = [
- "memchr",
+ "elf",
 ]
diff --git a/Cargo.toml b/Cargo.toml
index 25be964..693c640 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -16,7 +16,6 @@ opt-level = 3
 strip = true
 codegen-units = 1
 
-[dependencies.object]
-version = "0.36.5"
-features = ["read_core", "elf", "unaligned"]
+[dependencies.elf]
+version = "0.7.4"
 default-features = false
diff --git a/src/lib.rs b/src/lib.rs
index 03f46b2..52fbff4 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,43 +1,18 @@
 #![no_std]
 #![no_main]
-#![feature(alloc_error_handler)]
 
-extern crate alloc;
-
-use core::alloc::{GlobalAlloc, Layout};
-use core::ffi::c_void;
 use core::panic::PanicInfo;
 use core::slice::from_raw_parts;
-use object::{File, Object, ObjectSegment};
+
+use elf::abi::PT_LOAD;
+use elf::endian::NativeEndian;
+use elf::ElfBytes;
 
 #[panic_handler]
 unsafe fn panic(_info: &PanicInfo) -> ! {
     loop {}
 }
 
-#[global_allocator]
-static ALLOCATOR: Allocator = Allocator;
-
-#[alloc_error_handler]
-fn alloc_error_handler(layout: Layout) -> ! {
-    panic!("Allocation error: {:?}", layout);
-}
-
-struct Allocator;
-
-static mut MALLOC: Option<extern "C" fn(usize) -> *mut c_void> = None;
-static mut FREE: Option<extern "C" fn(*mut c_void)> = None;
-
-unsafe impl GlobalAlloc for Allocator {
-    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        MALLOC.unwrap()(layout.size()) as *mut u8
-    }
-
-    unsafe fn dealloc(&self, ptr: *mut u8, _layout: Layout) {
-        FREE.unwrap()(ptr as *mut c_void);
-    }
-}
-
 #[repr(C)]
 pub struct Segment {
     address: u64,
@@ -47,14 +22,19 @@ pub struct Segment {
 
 impl Segment {
     pub fn new(address: u64, size: u64, data: *const u8) -> Self {
-        Self { address, size, data }
+        Self {
+            address,
+            size,
+            data,
+        }
     }
 }
 
 #[repr(C)]
-pub enum ParseElfError {
-    None = 0,
+pub enum ElfParseResult {
+    EntryPoint(usize),
     InvalidElfData,
+    ElfContainsNoSegments,
     FailedToGetSegmentData,
     AllocFunctionNotProvided,
 }
@@ -63,32 +43,31 @@ pub enum ParseElfError {
 pub unsafe extern "C" fn parse_elf(
     elf_data: *const u8,
     elf_size: usize,
-    callback: extern "C" fn(segment: Segment),
-    malloc: Option<extern "C" fn(usize) -> *mut c_void>,
-    free: Option<extern "C" fn(*mut c_void)>,
-) -> ParseElfError {
-    if malloc.is_none() || free.is_none() {
-        return ParseElfError::AllocFunctionNotProvided;
-    }
-
-    MALLOC = malloc;
-    FREE = free;
-
+    mapping_callback: extern "C" fn(segment: Segment),
+) -> ElfParseResult {
     let buffer = from_raw_parts(elf_data, elf_size);
 
-    let binary = match File::parse(buffer) {
-        Ok(file) => file,
-        Err(_) => return ParseElfError::InvalidElfData,
+    let elf_file = match ElfBytes::<NativeEndian>::minimal_parse(buffer) {
+        Ok(elf) => elf,
+        Err(_) => return ElfParseResult::InvalidElfData,
+    };
+
+    let parsing_table = match elf_file.segments() {
+        Some(s) => s,
+        None => return ElfParseResult::ElfContainsNoSegments,
     };
 
-    for segment in binary.segments() {
-        let data = match segment.data() {
-            Ok(d) => d,
-            Err(_) => return ParseElfError::FailedToGetSegmentData,
-        };
+    for header in parsing_table {
+        if header.p_type == PT_LOAD {
+            let data = match elf_file.segment_data(&header) {
+                Ok(d) => d,
+                Err(_) => return ElfParseResult::FailedToGetSegmentData,
+            };
 
-        callback(Segment::new(segment.address(), segment.size(), data.as_ptr()));
+            let segment = Segment::new(header.p_vaddr, header.p_memsz, data.as_ptr());
+            mapping_callback(segment);
+        }
     }
 
-    ParseElfError::None
+    ElfParseResult::EntryPoint(elf_file.ehdr.e_entry as usize)
 }