Begin to enable virtual memory in wasmtime

There is a big caveat here, which is that wasmtime's mprotect calls are
ignored, so this sandbox is very unsound: a wasm module can take over
the entire guest easily

Signed-off-by: Lucy Menon <[email protected]>

Author: syntactically

src/hyperlight_wasm_aot/src/main.rs

@@ -150,9 +150,5 @@ fn main() {
 fn get_config() -> Config {
     let mut config = Config::new();
     config.target("x86_64-unknown-none").unwrap();
-    config.memory_reservation(0);
-    config.memory_reservation_for_growth(0);
-    config.memory_guard_size(0);
-    config.guard_before_linear_memory(false);
     config
 }

src/wasm_runtime/src/component.rs

@@ -107,11 +107,9 @@ fn load_wasm_module_phys(function_call: &FunctionCall) -> Result<Vec<u8>> {
 
 #[no_mangle]
 pub extern "C" fn hyperlight_main() {
+    crate::platform::register_page_fault_handler();
+
     let mut config = Config::new();
-    config.memory_reservation(0);
-    config.memory_guard_size(0);
-    config.memory_reservation_for_growth(0);
-    config.guard_before_linear_memory(false);
     config.with_custom_code_memory(Some(alloc::sync::Arc::new(
         crate::platform::WasmtimeCodeMemory {}
     )));

src/wasm_runtime/src/module.rs

@@ -90,10 +90,6 @@ pub fn guest_dispatch_function(function_call: &FunctionCall) -> Result<Vec<u8>>
 
 fn init_wasm_runtime() -> Result<Vec<u8>> {
     let mut config = Config::new();
-    config.memory_reservation(0);
-    config.memory_guard_size(0);
-    config.memory_reservation_for_growth(0);
-    config.guard_before_linear_memory(false);
     config.with_custom_code_memory(Some(alloc::sync::Arc::new(
         crate::platform::WasmtimeCodeMemory {},
     )));
@@ -171,6 +167,8 @@ fn load_wasm_module_phys(function_call: &FunctionCall) -> Result<Vec<u8>> {
 #[no_mangle]
 #[allow(clippy::fn_to_numeric_cast)] // GuestFunctionDefinition expects a function pointer as i64
 pub extern "C" fn hyperlight_main() {
+    crate::platform::register_page_fault_handler();
+
     register_function(GuestFunctionDefinition::new(
         "PrintOutput".to_string(),
         vec![ParameterType::String],

src/wasm_runtime/src/platform.rs

@@ -14,11 +14,48 @@ See the License for the specific language governing permissions and
 limitations under the License.
 */
 
-use alloc::alloc::{alloc, dealloc, Layout};
 use core::ffi::c_void;
 use core::sync::atomic::{AtomicPtr, AtomicU64, Ordering};
 
 use hyperlight_guest_bin::exceptions::handler;
+use hyperlight_guest_bin::paging;
+
+// Extremely stupid virtual address allocator
+// 0x1_0000_0000 is where the module is
+// we start at
+// 0x100_0000_0000 and go up from there
+static FIRST_VADDR: AtomicU64 = AtomicU64::new(0x100_0000_0000u64);
+fn page_fault_handler(
+    _exception_number: u64,
+    info: *mut handler::ExceptionInfo,
+    _ctx: *mut handler::Context,
+    page_fault_address: u64,
+) -> bool {
+    let error_code = unsafe { (&raw const (*info).error_code).read_volatile() };
+    // TODO: check if this is a guard-region trap (which can't happen
+    // right now since we don't actually set the permissions properly
+    // in mprotect)
+
+    // TODO: replace this with some generic virtual memory area data
+    // structure in hyperlight core
+    if (error_code & 0x1) == 0x0 && page_fault_address >= 0x100_0000_0000u64 {
+        unsafe {
+            let phys_page = paging::alloc_phys_pages(1);
+            let virt_base = (page_fault_address & !0xFFF) as *mut u8;
+            paging::map_region(
+                phys_page,
+                virt_base,
+                hyperlight_guest_bin::OS_PAGE_SIZE as u64,
+            );
+            virt_base.write_bytes(0u8, hyperlight_guest_bin::OS_PAGE_SIZE as usize);
+        }
+        return true; // Try again!
+    }
+    false
+}
+pub(crate) fn register_page_fault_handler() {
+    handler::handlers[14].store(page_fault_handler as usize as u64, Ordering::Release);
+}
 
 // Wasmtime Embedding Interface
 
@@ -27,24 +64,26 @@ use hyperlight_guest_bin::exceptions::handler;
  * appropriate interrupt handler yet. Consequently, we configure
  * wasmtime not to use any guard region, and precommit memory. */
 #[no_mangle]
-pub extern "C" fn wasmtime_mmap_new(size: usize, _prot_flags: u32, ret: &mut *mut u8) -> i32 {
-    *ret = unsafe { alloc(Layout::from_size_align(size, 0x1000).unwrap()) };
+pub extern "C" fn wasmtime_mmap_new(_size: usize, _prot_flags: u32, ret: &mut *mut u8) -> i32 {
+    *ret = FIRST_VADDR.fetch_add(0x100_0000_0000, Ordering::Relaxed) as *mut u8;
     0
 }
 
 /* Because of the precommitted memory strategy, we can't generally
  * support remap */
 #[no_mangle]
 pub extern "C" fn wasmtime_mmap_remap(addr: *mut u8, size: usize, prot_flags: u32) -> i32 {
-    panic!(
-        "wasmtime_mmap_remap {:x} {:x} {:x}",
-        addr as usize, size, prot_flags
-    );
+    if size > 0x100_0000_0000 {
+        panic!(
+            "wasmtime_mmap_remap {:x} {:x} {:x}",
+            addr as usize, size, prot_flags
+        );
+    }
+    0
 }
 
 #[no_mangle]
-pub extern "C" fn wasmtime_munmap(ptr: *mut u8, size: usize) -> i32 {
-    unsafe { dealloc(ptr, Layout::from_size_align(size, 0x1000).unwrap()) };
+pub extern "C" fn wasmtime_munmap(_ptr: *mut u8, _size: usize) -> i32 {
     0
 }