From e099e4658c5e429978b4fb36f74ca38a3055274a Mon Sep 17 00:00:00 2001
From: Sergej Schumilo <sergej@schumilo.de>
Date: Tue, 12 Apr 2022 10:33:21 +0200
Subject: [PATCH] add set_input function

---
 libnyx/src/ffi.rs |  5 ++---
 libnyx/src/lib.rs | 28 ++++++++++++++++++++++++++++
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/libnyx/src/ffi.rs b/libnyx/src/ffi.rs
index 0c7d83b..6fb97b9 100644
--- a/libnyx/src/ffi.rs
+++ b/libnyx/src/ffi.rs
@@ -181,9 +181,8 @@ pub extern "C" fn nyx_set_afl_input(nyx_process: * mut NyxProcess, buffer: *mut
         assert!((nyx_process as usize) % std::mem::align_of::<NyxProcess>() == 0);
         assert!((buffer as usize) % std::mem::align_of::<u8>() == 0);
 
-        std::ptr::copy(&size, (*nyx_process).process.payload.as_mut_ptr() as *mut u32, 1 as usize);
-        std::ptr::copy(buffer, (*nyx_process).process.payload[std::mem::size_of::<u32>()..].as_mut_ptr(), std::cmp::min(size as usize, (*nyx_process).input_buffer_size()));
-    }
+        (*nyx_process).set_input_ptr(buffer, size);
+   }
 }
 
 
diff --git a/libnyx/src/lib.rs b/libnyx/src/lib.rs
index 6d47d9c..f6daf4d 100644
--- a/libnyx/src/lib.rs
+++ b/libnyx/src/lib.rs
@@ -24,6 +24,23 @@ pub enum NyxReturnValue {
     Abort,      // Abort hypercall called
 }
 
+impl fmt::Display for NyxReturnValue {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+
+        let nyx_return_value_str = match self {
+            NyxReturnValue::Normal                => "Normal",
+            NyxReturnValue::Crash                 => "Crash",
+            NyxReturnValue::Timeout               => "Timeout",
+            NyxReturnValue::InvalidWriteToPayload => "InvalidWriteToPayload",
+            NyxReturnValue::Abort                 => "Abort",
+            NyxReturnValue::Error                 => "Error",
+            _                                     => "Unknown",
+        };
+
+        write!(f, "{}", nyx_return_value_str)
+    }
+}
+
 pub struct NyxProcess {
     process: QemuProcess,
 }
@@ -275,4 +292,15 @@ impl NyxProcess {
             }
         }
     }
+
+    pub fn set_input_ptr(&mut self, buffer: *const u8, size: u32) {
+        unsafe{
+            std::ptr::copy(&size, self.process.payload.as_mut_ptr() as *mut u32, 1 as usize);
+            std::ptr::copy(buffer, self.process.payload[std::mem::size_of::<u32>()..].as_mut_ptr(), std::cmp::min(size as usize, self.input_buffer_size()));
+        }
+    }
+    
+    pub fn set_input(&mut self, buffer: &[u8], size: u32) {
+        self.set_input_ptr(buffer.as_ptr(), size);
+    }
 }