1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
//! Library for allocating RWX memory on Unix and Windows
//!
//! ```
//! use virtual_memory::*;
//!
//! let buf = &[
//!     //mov eax, 1337
//!     0xb8, 0x39, 0x05, 0x00, 0x00,
//!     //ret
//!     0xc3,
//! ];
//!
//! let mut memory = VirtualMemory::new(buf.len()).expect("failed to allocate rwx memory");
//! memory.copy_from_slice(buf);
//!
//! let f: extern "C" fn() -> u32 = unsafe { std::mem::transmute(memory.as_ptr()) };
//! assert_eq!(f(), 1337);
//! ```

#![no_std]

use core::{
    error,
    ffi::c_void,
    fmt,
    ops::{Deref, DerefMut},
    ptr, result, slice,
};

/// Custom `Result` type that is used by this crate
pub type Result<T, E = Error> = result::Result<T, E>;

/// Describes possible errors, currently it's only memory allocation error
#[derive(Debug)]
pub enum Error {
    Allocation,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Error::Allocation => write!(f, "memory allocation error"),
        }
    }
}

impl error::Error for Error {}

/// Wraps OS-specific functionality related to allocation RWX memory
pub struct VirtualMemory {
    ptr: *mut c_void,
    len: usize,
}

impl VirtualMemory {
    /// Trying to allocate RWX memory
    pub fn new(len: usize) -> Result<Self> {
        #[cfg(unix)]
        {
            let ptr = unsafe {
                libc::mmap(
                    ptr::null_mut(),
                    len,
                    libc::PROT_READ | libc::PROT_WRITE | libc::PROT_EXEC,
                    libc::MAP_PRIVATE | libc::MAP_ANON,
                    -1,
                    0,
                )
            };

            if ptr == libc::MAP_FAILED {
                return Err(Error::Allocation);
            }

            Ok(Self { ptr, len })
        }

        #[cfg(windows)]
        {
            use windows_sys::Win32::System::Memory::{
                VirtualAlloc, MEM_COMMIT, MEM_RESERVE, PAGE_EXECUTE_READWRITE,
            };

            let ptr = unsafe {
                VirtualAlloc(
                    ptr::null(),
                    len,
                    MEM_COMMIT | MEM_RESERVE,
                    PAGE_EXECUTE_READWRITE,
                )
            };

            if ptr.is_null() {
                return Err(Error::Allocation);
            }

            Ok(Self { ptr, len })
        }
    }

    /// Returns pointer to this memory
    #[inline]
    pub fn ptr(&self) -> *const u8 {
        self.ptr as *const u8
    }

    /// Returns mutable pointer to this memory
    #[inline]
    pub fn mut_ptr(&mut self) -> *mut u8 {
        self.ptr as *mut u8
    }

    /// Returns length of allocated memory
    #[inline]
    #[allow(clippy::len_without_is_empty)]
    pub fn len(&self) -> usize {
        self.len
    }
}

/// Allows to pass `VirtualMemory` as `&[u8]`
impl Deref for VirtualMemory {
    type Target = [u8];

    #[inline]
    fn deref(&self) -> &Self::Target {
        unsafe { slice::from_raw_parts(self.ptr as _, self.len) }
    }
}

/// Allows to pass `VirtualMemory` as `&mut [u8]`
impl DerefMut for VirtualMemory {
    #[inline]
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { slice::from_raw_parts_mut(self.ptr as _, self.len) }
    }
}

/// Acts as destructor (OS-specific)
impl Drop for VirtualMemory {
    fn drop(&mut self) {
        #[cfg(unix)]
        unsafe {
            libc::munmap(self.ptr, self.len);
        }

        #[cfg(windows)]
        unsafe {
            use windows_sys::Win32::System::Memory::{VirtualFree, MEM_RELEASE};

            VirtualFree(self.ptr, 0, MEM_RELEASE);
        }
    }
}