changes from feedback.

AFLplusplus · Aug 22, 2023 · 56896dd · 56896dd
1 parent 4963bb2
commit 56896dd
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 4 deletions.
diff --git a/libafl_qemu/Cargo.toml b/libafl_qemu/Cargo.toml
@@ -59,6 +59,7 @@ capstone = "0.11.0"
 pyo3 = { version = "0.18.3", optional = true }
 rangemap = "1.3"
 log = "0.4.20"
+once_cell = "1.18.0"
 
 [build-dependencies]
 pyo3-build-config = { version = "0.15", optional = true }

diff --git a/libafl_qemu/libafl_qemu_build/src/bindings.rs b/libafl_qemu/libafl_qemu_build/src/bindings.rs
@@ -120,6 +120,7 @@ pub fn generate(
         .allowlist_function("syx_snapshot_root_restore")
         .allowlist_function("syx_snapshot_dirty_list_add")
         .allowlist_function("device_list_all")
+        .allowlist_function("qemu_target_page_size")
         .blocklist_function("main_loop_wait") // bindgen issue #1313
         .parse_callbacks(Box::new(bindgen::CargoCallbacks));
 

diff --git a/libafl_qemu/libafl_qemu_sys/src/x86_64_stub_bindings.rs b/libafl_qemu/libafl_qemu_sys/src/x86_64_stub_bindings.rs
@@ -11100,6 +11100,12 @@ impl ::std::ops::BitAndAssign for qemu_plugin_mem_rw {
         self.0 &= rhs.0;
     }
 }
+
+extern "C" {
+    #[doc = " qemu_target_page_size - return the target's page size"]
+    pub fn qemu_target_page_size() -> usize;
+}
+
 #[repr(transparent)]
 #[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
 pub struct qemu_plugin_mem_rw(pub ::std::os::raw::c_uint);

diff --git a/libafl_qemu/src/emu.rs b/libafl_qemu/src/emu.rs
@@ -1,4 +1,4 @@
-//! Expose QEMU user `LibAFL` C api to Rust
+// Expose QEMU user `LibAFL` C api to Rust
 
 use core::{
     convert::Into,
@@ -14,6 +14,10 @@ use std::{
 };
 use std::{slice::from_raw_parts, str::from_utf8_unchecked};
 
+use once_cell::sync::OnceCell;
+
+static SNAPSHOT_PAGE_SIZE: OnceCell<usize> = OnceCell::new();
+
 #[cfg(emulation_mode = "usermode")]
 use libc::c_int;
 use num_enum::{IntoPrimitive, TryFromPrimitive};
@@ -743,14 +747,18 @@ impl CPU {
         self.ptr
     }
 
-    pub fn get_page_size(&self) -> usize {
+    pub fn page_size(&self) -> usize {
         #[cfg(emulation_mode = "usermode")]
         {
-            4096usize
+            *SNAPSHOT_PAGE_SIZE.get_or_init(|| {
+                unsafe { libc::sysconf(libc::_SC_PAGE_SIZE) }
+                    .try_into()
+                    .expect("Invalid page size")
+            })
         }
         #[cfg(emulation_mode = "systemmode")]
         {
-            libafl_qemu_sys::qemu_target_page_size()
+            *SNAPSHOT_PAGE_SIZE.get_or_init(|| unsafe { libafl_qemu_sys::qemu_target_page_size() })
         }
     }
 }