Feature - Memory management on Windows

Cargo.toml

@@ -22,20 +22,21 @@ include = [
 arbitrary = { version = "1.0", optional = true, features = ["derive"] }
 byteorder = "1.2"
 combine = "3.8.1"
+gdbstub = { version = "0.6.2", optional = true }
 goblin = "0.5.1"
 hash32 = "0.2.0"
 libc = { version = "0.2", optional = true }
 log = "0.4.2"
 rand = { version = "0.8.5", features = ["small_rng"]}
+rustc-demangle = "0.1"
 scroll = "0.11"
 thiserror = "1.0.26"
-rustc-demangle = "0.1"
-gdbstub = { version = "0.6.2", optional = true }
+winapi = { version = "0.3", features = ["memoryapi", "sysinfoapi", "winnt", "errhandlingapi"], optional = true }
 
 [features]
 default = ["jit"]
 fuzzer-not-safe-for-production = ["arbitrary"]
-jit = ["libc"]
+jit = ["libc", "winapi"]
 debugger = ["gdbstub"]
 
 [dev-dependencies]

src/jit.rs

@@ -10,15 +10,16 @@
 // the MIT license <http://opensource.org/licenses/MIT>, at your option. This file may not be
 // copied, modified, or distributed except according to those terms.
 
-extern crate libc;
-
 use rand::{rngs::SmallRng, Rng, SeedableRng};
 use std::{fmt::Debug, marker::PhantomData, mem, ptr};
 
 use crate::{
     ebpf::{self, FIRST_SCRATCH_REG, FRAME_PTR_REG, INSN_SIZE, SCRATCH_REGS, STACK_PTR_REG},
     elf::Executable,
     error::EbpfError,
+    memory_management::{
+        allocate_pages, free_pages, get_system_page_size, protect_pages, round_to_page_size,
+    },
     memory_region::{AccessType, MemoryMapping},
     vm::{Config, ContextObject, ProgramResult, RuntimeEnvironment, SyscallFunction},
     x86::*,
@@ -39,59 +40,18 @@ pub struct JitProgram<C: ContextObject> {
     _marker: PhantomData<C>,
 }
 
-#[cfg(not(target_os = "windows"))]
-macro_rules! libc_error_guard {
-    (succeeded?, mmap, $addr:expr, $($arg:expr),*) => {{
-        *$addr = libc::mmap(*$addr, $($arg),*);
-        *$addr != libc::MAP_FAILED
-    }};
-    (succeeded?, $function:ident, $($arg:expr),*) => {
-        libc::$function($($arg),*) == 0
-    };
-    ($function:ident, $($arg:expr),*) => {{
-        const RETRY_COUNT: usize = 3;
-        for i in 0..RETRY_COUNT {
-            if libc_error_guard!(succeeded?, $function, $($arg),*) {
-                break;
-            } else if i + 1 == RETRY_COUNT {
-                let args = vec![$(format!("{:?}", $arg)),*];
-                #[cfg(any(target_os = "freebsd", target_os = "ios", target_os = "macos"))]
-                let errno = *libc::__error();
-                #[cfg(any(target_os = "android", target_os = "netbsd", target_os = "openbsd"))]
-                let errno = *libc::__errno();
-                #[cfg(target_os = "linux")]
-                let errno = *libc::__errno_location();
-                return Err(EbpfError::LibcInvocationFailed(stringify!($function), args, errno));
-            }
-        }
-    }};
-}
-
-fn round_to_page_size(value: usize, page_size: usize) -> usize {
-    (value + page_size - 1) / page_size * page_size
-}
-
 impl<C: ContextObject> JitProgram<C> {
     fn new(pc: usize, code_size: usize) -> Result<Self, EbpfError> {
+        let page_size = get_system_page_size();
+        let pc_loc_table_size = round_to_page_size(pc * 8, page_size);
+        let over_allocated_code_size = round_to_page_size(code_size, page_size);
         unsafe {
-            let page_size = libc::sysconf(libc::_SC_PAGESIZE) as usize;
-            let pc_loc_table_size = round_to_page_size(pc * 8, page_size);
-            let over_allocated_code_size = round_to_page_size(code_size, page_size);
-            let mut raw: *mut libc::c_void = std::ptr::null_mut();
-            libc_error_guard!(
-                mmap,
-                &mut raw,
-                pc_loc_table_size + over_allocated_code_size,
-                libc::PROT_READ | libc::PROT_WRITE,
-                libc::MAP_ANONYMOUS | libc::MAP_PRIVATE,
-                0,
-                0
-            );
+            let raw = allocate_pages(pc_loc_table_size + over_allocated_code_size)?;
             Ok(Self {
                 page_size,
                 pc_section: std::slice::from_raw_parts_mut(raw as *mut usize, pc),
                 text_section: std::slice::from_raw_parts_mut(
-                    raw.add(pc_loc_table_size) as *mut u8,
+                    (raw as *mut u8).add(pc_loc_table_size),
                     over_allocated_code_size,
                 ),
                 _marker: PhantomData::default(),
@@ -103,38 +63,31 @@ impl<C: ContextObject> JitProgram<C> {
         if self.page_size == 0 {
             return Ok(());
         }
+        let raw = self.pc_section.as_ptr() as *mut u8;
+        let pc_loc_table_size = round_to_page_size(self.pc_section.len() * 8, self.page_size);
+        let over_allocated_code_size = round_to_page_size(self.text_section.len(), self.page_size);
+        let code_size = round_to_page_size(text_section_usage, self.page_size);
         unsafe {
-            let raw = self.pc_section.as_ptr() as *mut u8;
-            let pc_loc_table_size = round_to_page_size(self.pc_section.len() * 8, self.page_size);
-            let over_allocated_code_size =
-                round_to_page_size(self.text_section.len(), self.page_size);
-            let code_size = round_to_page_size(text_section_usage, self.page_size);
-            if over_allocated_code_size > code_size {
-                libc_error_guard!(
-                    munmap,
-                    raw.add(pc_loc_table_size).add(code_size) as *mut _,
-                    over_allocated_code_size - code_size
-                );
-            }
+            // Fill with debugger traps
             std::ptr::write_bytes(
                 raw.add(pc_loc_table_size).add(text_section_usage),
                 0xcc,
                 code_size - text_section_usage,
-            ); // Fill with debugger traps
+            );
+            if over_allocated_code_size > code_size {
+                free_pages(
+                    raw.add(pc_loc_table_size).add(code_size),
+                    over_allocated_code_size - code_size,
+                )?;
+            }
             self.text_section =
                 std::slice::from_raw_parts_mut(raw.add(pc_loc_table_size), text_section_usage);
-            libc_error_guard!(
-                mprotect,
-                self.pc_section.as_mut_ptr() as *mut _,
+            protect_pages(
+                self.pc_section.as_mut_ptr() as *mut u8,
                 pc_loc_table_size,
-                libc::PROT_READ
-            );
-            libc_error_guard!(
-                mprotect,
-                self.text_section.as_mut_ptr() as *mut _,
-                code_size,
-                libc::PROT_EXEC | libc::PROT_READ
-            );
+                false,
+            )?;
+            protect_pages(self.text_section.as_mut_ptr(), code_size, true)?;
         }
         Ok(())
     }
@@ -200,10 +153,9 @@ impl<C: ContextObject> Drop for JitProgram<C> {
         let pc_loc_table_size = round_to_page_size(self.pc_section.len() * 8, self.page_size);
         let code_size = round_to_page_size(self.text_section.len(), self.page_size);
         if pc_loc_table_size + code_size > 0 {
-            #[cfg(not(target_os = "windows"))]
             unsafe {
-                libc::munmap(
-                    self.pc_section.as_ptr() as *mut _,
+                let _ = free_pages(
+                    self.pc_section.as_ptr() as *mut u8,
                     pc_loc_table_size + code_size,
                 );
             }

src/lib.rs

@@ -40,6 +40,8 @@ pub mod insn_builder;
 pub mod interpreter;
 #[cfg(all(feature = "jit", not(target_os = "windows"), target_arch = "x86_64"))]
 mod jit;
+#[cfg(feature = "jit")]
+mod memory_management;
 pub mod memory_region;
 pub mod static_analysis;
 pub mod syscalls;

src/memory_management.rs

@@ -0,0 +1,166 @@
+// Copyright 2022 Solana Maintainers <maintainers@solana.com>
+//
+// Licensed under the Apache License, Version 2.0 <http://www.apache.org/licenses/LICENSE-2.0> or
+// the MIT license <http://opensource.org/licenses/MIT>, at your option. This file may not be
+// copied, modified, or distributed except according to those terms.
+
+#![cfg_attr(target_os = "windows", allow(dead_code))]
+
+use crate::error::EbpfError;
+
+#[cfg(not(target_os = "windows"))]
+extern crate libc;
+#[cfg(not(target_os = "windows"))]
+use libc::c_void;
+
+#[cfg(target_os = "windows")]
+use winapi::{
+    ctypes::c_void,
+    shared::minwindef,
+    um::{
+        errhandlingapi::GetLastError,
+        memoryapi::{VirtualAlloc, VirtualFree, VirtualProtect},
+        sysinfoapi::{GetSystemInfo, SYSTEM_INFO},
+        winnt,
+    },
+};
+
+#[cfg(not(target_os = "windows"))]
+macro_rules! libc_error_guard {
+    (succeeded?, mmap, $addr:expr, $($arg:expr),*) => {{
+        *$addr = libc::mmap(*$addr, $($arg),*);
+        *$addr != libc::MAP_FAILED
+    }};
+    (succeeded?, $function:ident, $($arg:expr),*) => {
+        libc::$function($($arg),*) == 0
+    };
+    ($function:ident, $($arg:expr),* $(,)?) => {{
+        const RETRY_COUNT: usize = 3;
+        for i in 0..RETRY_COUNT {
+            if libc_error_guard!(succeeded?, $function, $($arg),*) {
+                break;
+            } else if i.saturating_add(1) == RETRY_COUNT {
+                let args = vec![$(format!("{:?}", $arg)),*];
+                #[cfg(any(target_os = "freebsd", target_os = "ios", target_os = "macos"))]
+                let errno = *libc::__error();
+                #[cfg(any(target_os = "android", target_os = "netbsd", target_os = "openbsd"))]
+                let errno = *libc::__errno();
+                #[cfg(target_os = "linux")]
+                let errno = *libc::__errno_location();
+                return Err(EbpfError::LibcInvocationFailed(stringify!($function), args, errno));
+            }
+        }
+    }};
+}
+
+#[cfg(target_os = "windows")]
+macro_rules! winapi_error_guard {
+    (succeeded?, VirtualAlloc, $addr:expr, $($arg:expr),*) => {{
+        *$addr = VirtualAlloc(*$addr, $($arg),*);
+        !(*$addr).is_null()
+    }};
+    (succeeded?, $function:ident, $($arg:expr),*) => {
+        $function($($arg),*) != 0
+    };
+    ($function:ident, $($arg:expr),* $(,)?) => {{
+        if !winapi_error_guard!(succeeded?, $function, $($arg),*) {
+            let args = vec![$(format!("{:?}", $arg)),*];
+            let errno = GetLastError();
+            return Err(EbpfError::LibcInvocationFailed(stringify!($function), args, errno as i32));
+        }
+    }};
+}
+
+pub fn get_system_page_size() -> usize {
+    #[cfg(not(target_os = "windows"))]
+    unsafe {
+        libc::sysconf(libc::_SC_PAGESIZE) as usize
+    }
+    #[cfg(target_os = "windows")]
+    unsafe {
+        let mut system_info: SYSTEM_INFO = std::mem::zeroed();
+        GetSystemInfo(&mut system_info);
+        system_info.dwPageSize as usize
+    }
+}
+
+pub fn round_to_page_size(value: usize, page_size: usize) -> usize {
+    value
+        .saturating_add(page_size)
+        .saturating_sub(1)
+        .saturating_div(page_size)
+        .saturating_mul(page_size)
+}
+
+pub unsafe fn allocate_pages(size_in_bytes: usize) -> Result<*mut u8, EbpfError> {
+    let mut raw: *mut c_void = std::ptr::null_mut();
+    #[cfg(not(target_os = "windows"))]
+    libc_error_guard!(
+        mmap,
+        &mut raw,
+        size_in_bytes,
+        libc::PROT_READ | libc::PROT_WRITE,
+        libc::MAP_ANONYMOUS | libc::MAP_PRIVATE,
+        0,
+        0,
+    );
+    #[cfg(target_os = "windows")]
+    winapi_error_guard!(
+        VirtualAlloc,
+        &mut raw,
+        size_in_bytes,
+        winnt::MEM_RESERVE | winnt::MEM_COMMIT,
+        winnt::PAGE_READWRITE,
+    );
+    Ok(raw as *mut u8)
+}
+
+pub unsafe fn free_pages(raw: *mut u8, size_in_bytes: usize) -> Result<(), EbpfError> {
+    #[cfg(not(target_os = "windows"))]
+    libc_error_guard!(munmap, raw as *mut _, size_in_bytes);
+    #[cfg(target_os = "windows")]
+    winapi_error_guard!(
+        VirtualFree,
+        raw as *mut _,
+        size_in_bytes,
+        winnt::MEM_RELEASE, // winnt::MEM_DECOMMIT
+    );
+    Ok(())
+}
+
+pub unsafe fn protect_pages(
+    raw: *mut u8,
+    size_in_bytes: usize,
+    executable_flag: bool,
+) -> Result<(), EbpfError> {
+    #[cfg(not(target_os = "windows"))]
+    {
+        libc_error_guard!(
+            mprotect,
+            raw as *mut _,
+            size_in_bytes,
+            if executable_flag {
+                libc::PROT_EXEC | libc::PROT_READ
+            } else {
+                libc::PROT_READ
+            },
+        );
+    }
+    #[cfg(target_os = "windows")]
+    {
+        let mut old: minwindef::DWORD = 0;
+        let ptr_old: *mut minwindef::DWORD = &mut old;
+        winapi_error_guard!(
+            VirtualProtect,
+            raw as *mut _,
+            size_in_bytes,
+            if executable_flag {
+                winnt::PAGE_EXECUTE_READ
+            } else {
+                winnt::PAGE_READONLY
+            },
+            ptr_old,
+        );
+    }
+    Ok(())
+}