pelite\mmap/
windows.rs

1
2use std::{io, mem, ptr};
3use std::ffi::OsStr;
4use std::path::Path;
5use std::os::windows::ffi::OsStrExt;
6use std::os::windows::io::{AsRawHandle, RawHandle};
7
8use winapi::um::fileapi::{CreateFileW, OPEN_EXISTING};
9use winapi::um::memoryapi::{CreateFileMappingW, MapViewOfFile, UnmapViewOfFile, VirtualQuery, FILE_MAP_READ, FILE_MAP_COPY};
10use winapi::um::handleapi::{CloseHandle, INVALID_HANDLE_VALUE};
11use winapi::shared::ntdef::{NULL, HANDLE};
12use winapi::shared::minwindef::{LPVOID};
13use winapi::um::winnt::{PAGE_READONLY, SEC_IMAGE, GENERIC_READ, FILE_SHARE_READ, FILE_ATTRIBUTE_NORMAL};
14
15//----------------------------------------------------------------
16
17/// Memory mapped image.
18pub struct ImageMap {
19	handle: HANDLE,
20	bytes: *mut [u8],
21}
22impl ImageMap {
23	/// Maps the executable image into memory with correctly aligned sections.
24	pub fn open<P: AsRef<Path> + ?Sized>(path: &P) -> io::Result<ImageMap> {
25		unsafe { Self::_open(path.as_ref()) }
26	}
27	unsafe fn _open(path: &Path) -> io::Result<ImageMap> {
28		// Get its file handle
29		let file = {
30			// Get the path as a nul terminated wide string
31			let path: &OsStr = path.as_ref();
32			let mut wpath: Vec<u16> = path.encode_wide().collect();
33			wpath.push(0);
34			CreateFileW(wpath.as_ptr(), GENERIC_READ, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL)
35		};
36		if file != INVALID_HANDLE_VALUE {
37			// Create the image file mapping, `SEC_IMAGE` does its magic thing
38			let map = CreateFileMappingW(file, ptr::null_mut(), PAGE_READONLY | SEC_IMAGE, 0, 0, ptr::null());
39			CloseHandle(file);
40			if map != NULL {
41				// Map view of the file
42				let view = MapViewOfFile(map, FILE_MAP_COPY, 0, 0, 0);
43				if view != ptr::null_mut() {
44					// Trust the OS with correctly mapping the image.
45					// Trust me to have read and understood the documentation.
46					// There is no validation and 64bit headers are used because the offsets are the same for PE32.
47					use crate::image::{IMAGE_DOS_HEADER, IMAGE_NT_HEADERS64};
48					let dos_header = view as *const IMAGE_DOS_HEADER;
49					let nt_header = (view as usize + (*dos_header).e_lfanew as usize) as *const IMAGE_NT_HEADERS64;
50					let size_of = (*nt_header).OptionalHeader.SizeOfImage;
51					let bytes = ptr::slice_from_raw_parts_mut(view as *mut u8, size_of as usize);
52					return Ok(ImageMap { handle: map, bytes });
53				}
54				let err = io::Error::last_os_error();
55				CloseHandle(map);
56				return Err(err);
57			}
58		}
59		Err(io::Error::last_os_error())
60	}
61}
62impl AsRawHandle for ImageMap {
63	fn as_raw_handle(&self) -> RawHandle {
64		self.handle as RawHandle
65	}
66}
67impl AsRef<[u8]> for ImageMap {
68	fn as_ref(&self) -> &[u8] {
69		unsafe { &*self.bytes }
70	}
71}
72impl Drop for ImageMap {
73	fn drop(&mut self) {
74		unsafe {
75			UnmapViewOfFile((*self.bytes).as_ptr() as LPVOID);
76			CloseHandle(self.handle);
77		}
78	}
79}
80
81//----------------------------------------------------------------
82
83/// Memory mapped file.
84pub struct FileMap {
85	handle: HANDLE,
86	bytes: *mut [u8],
87}
88impl FileMap {
89	/// Maps the whole file into memory.
90	pub fn open<P: AsRef<Path> + ?Sized>(path: &P) -> io::Result<FileMap> {
91		unsafe { Self::_open(path.as_ref()) }
92	}
93	unsafe fn _open(path: &Path) -> io::Result<FileMap> {
94		// Get its file handle
95		let file = {
96			// Get the path as a nul terminated wide string
97			let path: &OsStr = path.as_ref();
98			let mut wpath: Vec<u16> = path.encode_wide().collect();
99			wpath.push(0);
100			CreateFileW(wpath.as_ptr(), GENERIC_READ, FILE_SHARE_READ, ptr::null_mut(), OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL)
101		};
102		if file == INVALID_HANDLE_VALUE {
103			return Err(io::Error::last_os_error());
104		}
105		// Create the memory file mapping
106		let map = CreateFileMappingW(file, ptr::null_mut(), PAGE_READONLY, 0, 0, ptr::null());
107		CloseHandle(file);
108		if map == NULL {
109			return Err(io::Error::last_os_error());
110		}
111		// Map view of the file
112		let view = MapViewOfFile(map, FILE_MAP_READ, 0, 0, 0);
113		if view == ptr::null_mut() {
114			let err = io::Error::last_os_error();
115			CloseHandle(map);
116			return Err(err);
117		}
118		// Get the size of the file mapping, should never fail...
119		let mut mem_basic_info = mem::zeroed();
120		let vq_result = VirtualQuery(view, &mut mem_basic_info, mem::size_of_val(&mem_basic_info));
121		debug_assert_eq!(vq_result, mem::size_of_val(&mem_basic_info));
122		// Now have enough information to construct the FileMap
123		let bytes = ptr::slice_from_raw_parts_mut(view as *mut u8, mem_basic_info.RegionSize as usize);
124		Ok(FileMap { handle: map, bytes })
125	}
126}
127impl AsRawHandle for FileMap {
128	fn as_raw_handle(&self) -> RawHandle {
129		self.handle as RawHandle
130	}
131}
132impl AsRef<[u8]> for FileMap {
133	fn as_ref(&self) -> &[u8] {
134		unsafe { &*self.bytes }
135	}
136}
137impl Drop for FileMap {
138	fn drop(&mut self) {
139		unsafe {
140			UnmapViewOfFile((*self.bytes).as_ptr() as LPVOID);
141			CloseHandle(self.handle);
142		}
143	}
144}