From 757cedd0e0143cb6c0a5ef7129505651a5523729 Mon Sep 17 00:00:00 2001
From: Erin <erin@erindesu.cz>
Date: Thu, 23 Nov 2023 17:35:11 +0100
Subject: [PATCH] Moved module

---
 hbxrt/src/linux.rs | 67 ++++++++++++++++++++++++++++++++++++++++++++++
 hbxrt/src/main.rs  | 56 +++-----------------------------------
 2 files changed, 71 insertions(+), 52 deletions(-)
 create mode 100644 hbxrt/src/linux.rs

diff --git a/hbxrt/src/linux.rs b/hbxrt/src/linux.rs
new file mode 100644
index 0000000..949d971
--- /dev/null
+++ b/hbxrt/src/linux.rs
@@ -0,0 +1,67 @@
+use {
+    nix::sys::mman::{mmap, MapFlags, ProtFlags},
+    std::{fs::File, num::NonZeroUsize, path::Path, process::exit},
+};
+
+/// Allocate stack for program
+pub unsafe fn alloc_stack(size: usize) -> nix::Result<*mut u8> {
+    unsafe {
+        Ok(mmap::<std::fs::File>(
+            None,
+            NonZeroUsize::new(size).expect("Stack size should be > 0"),
+            ProtFlags::PROT_GROWSDOWN | ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
+            MapFlags::MAP_GROWSDOWN
+                | MapFlags::MAP_STACK
+                | MapFlags::MAP_ANON
+                | MapFlags::MAP_PRIVATE,
+            None,
+            0,
+        )?
+        .cast())
+    }
+}
+
+/// Memory map bytecode
+pub unsafe fn mmap_bytecode(path: impl AsRef<Path>) -> Result<*mut u8, Box<dyn std::error::Error>> {
+    let file = File::open(&path)?;
+    Ok(unsafe {
+        mmap(
+            None,
+            NonZeroUsize::new(file.metadata()?.len() as usize).ok_or("File is empty")?,
+            ProtFlags::PROT_READ,
+            MapFlags::MAP_PRIVATE,
+            Some(&file),
+            0,
+        )?
+        .cast()
+    })
+}
+
+/// Set handler for page fault
+pub unsafe fn hook_pagefault() -> nix::Result<()> {
+    unsafe {
+        use nix::sys::signal;
+
+        extern "C" fn action(
+            _: std::ffi::c_int,
+            info: *mut nix::libc::siginfo_t,
+            _: *mut std::ffi::c_void,
+        ) {
+            unsafe {
+                eprintln!("[E] Memory access fault at {:p}", (*info).si_addr());
+                exit(2);
+            }
+        }
+
+        signal::sigaction(
+            signal::Signal::SIGSEGV,
+            &nix::sys::signal::SigAction::new(
+                signal::SigHandler::SigAction(action),
+                signal::SaFlags::SA_NODEFER,
+                nix::sys::signalfd::SigSet::empty(),
+            ),
+        )?;
+    }
+
+    Ok(())
+}
diff --git a/hbxrt/src/main.rs b/hbxrt/src/main.rs
index d2935cc..121af6a 100644
--- a/hbxrt/src/main.rs
+++ b/hbxrt/src/main.rs
@@ -2,6 +2,7 @@
 
 #![deny(unsafe_op_in_unsafe_fn)]
 
+mod linux;
 mod mem;
 
 use {
@@ -20,67 +21,18 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     };
 
     // Allocate stack
-    const STACK_SIZE: usize = 1024 * 1024 * 2;
-
-    let stack_ptr = unsafe {
-        mmap::<std::fs::File>(
-            None,
-            NonZeroUsize::new(STACK_SIZE).expect("Stack size should be > 0"),
-            ProtFlags::PROT_GROWSDOWN | ProtFlags::PROT_READ | ProtFlags::PROT_WRITE,
-            MapFlags::MAP_GROWSDOWN
-                | MapFlags::MAP_STACK
-                | MapFlags::MAP_ANON
-                | MapFlags::MAP_PRIVATE,
-            None,
-            0,
-        )
-    }?;
-
+    let stack_ptr = unsafe { linux::alloc_stack(1024 * 1024 * 2) }?;
     eprintln!("[I] Stack allocated at {stack_ptr:p}");
 
     // Load program
     eprintln!("[I] Loading image from \"{image_path}\"");
-    let file = File::open(image_path)?;
-    let ptr = unsafe {
-        mmap(
-            None,
-            NonZeroUsize::new(file.metadata()?.len() as usize).ok_or("File is empty")?,
-            ProtFlags::PROT_READ,
-            MapFlags::MAP_PRIVATE,
-            Some(&file),
-            0,
-        )?
-    };
-
+    let ptr = unsafe { linux::mmap_bytecode(image_path) }?;
     eprintln!("[I] Image loaded at {ptr:p}");
 
     let mut vm = unsafe { Vm::<_, 0>::new(mem::HostMemory, Address::new(ptr as u64)) };
     vm.write_reg(254, stack_ptr as u64);
 
-    // Memory access fault handling
-    unsafe {
-        use nix::sys::signal;
-
-        extern "C" fn action(
-            _: std::ffi::c_int,
-            info: *mut nix::libc::siginfo_t,
-            _: *mut std::ffi::c_void,
-        ) {
-            unsafe {
-                eprintln!("[E] Memory access fault at {:p}", (*info).si_addr());
-                exit(2);
-            }
-        }
-
-        signal::sigaction(
-            signal::Signal::SIGSEGV,
-            &nix::sys::signal::SigAction::new(
-                signal::SigHandler::SigAction(action),
-                signal::SaFlags::SA_NODEFER,
-                nix::sys::signalfd::SigSet::empty(),
-            ),
-        )?;
-    }
+    unsafe { linux::hook_pagefault() }?;
 
     // Execute program
     let stat = loop {