@@ -15,60 +15,162 @@ limitations under the License.
1515*/
1616
1717use core:: ffi:: c_void;
18+ use std:: collections:: HashMap ;
19+ use std:: collections:: hash_map:: Entry ;
1820
21+ use hyperlight_common:: mem:: PAGE_SIZE_USIZE ;
1922use tracing:: { Span , instrument} ;
2023use windows:: Win32 :: Foundation :: HANDLE ;
2124use windows:: Win32 :: System :: Memory :: {
22- MEMORY_MAPPED_VIEW_ADDRESS , UNMAP_VIEW_OF_FILE_FLAGS , UnmapViewOfFile2 ,
25+ MEMORY_MAPPED_VIEW_ADDRESS , MapViewOfFileNuma2 , PAGE_NOACCESS , PAGE_PROTECTION_FLAGS ,
26+ PAGE_READWRITE , UNMAP_VIEW_OF_FILE_FLAGS , UnmapViewOfFile2 , VirtualProtectEx ,
2327} ;
28+ use windows:: Win32 :: System :: SystemServices :: NUMA_NO_PREFERRED_NODE ;
2429
2530use super :: surrogate_process_manager:: get_surrogate_process_manager;
2631use super :: wrappers:: HandleWrapper ;
32+ use crate :: HyperlightError :: WindowsAPIError ;
33+ use crate :: { Result , log_then_return} ;
34+
35+ #[ derive( Debug ) ]
36+ pub ( crate ) struct HandleMapping {
37+ pub ( crate ) use_count : u64 ,
38+ pub ( crate ) surrogate_base : * mut c_void ,
39+ }
2740
2841/// Contains details of a surrogate process to be used by a Sandbox for providing memory to a HyperV VM on Windows.
2942/// See surrogate_process_manager for details on why this is needed.
3043#[ derive( Debug ) ]
3144pub ( super ) struct SurrogateProcess {
32- /// The address of memory allocated in the surrogate process to be mapped to the VM.
33- /// This includes the first guard page
34- pub ( crate ) allocated_address : * mut c_void ,
45+ /// The various mappings between handles in the host and surrogate process
46+ pub ( crate ) mappings : HashMap < usize , HandleMapping > ,
3547 /// The handle to the surrogate process.
3648pub ( crate ) process_handle : HandleWrapper ,
3749}
3850
3951impl SurrogateProcess {
4052 #[ instrument( skip_all, parent = Span :: current( ) , level= "Trace" ) ]
41- pub ( super ) fn new ( allocated_address : * mut c_void , process_handle : HANDLE ) -> Self {
53+ pub ( super ) fn new ( process_handle : HANDLE ) -> Self {
4254 Self {
43- allocated_address ,
55+ mappings : HashMap :: new ( ) ,
4456 process_handle : HandleWrapper :: from ( process_handle) ,
4557 }
4658 }
59+
60+ pub ( super ) fn map (
61+ & mut self ,
62+ handle : HandleWrapper ,
63+ host_base : usize ,
64+ host_size : usize ,
65+ ) -> Result < * mut c_void > {
66+ match self . mappings . entry ( host_base) {
67+ Entry :: Occupied ( mut oe) => {
68+ oe. get_mut ( ) . use_count += 1 ;
69+ Ok ( oe. get ( ) . surrogate_base )
70+ }
71+ Entry :: Vacant ( ve) => {
72+ // Use MapViewOfFile2 to map memory into the surrogate process, the MapViewOfFile2 API is implemented in as an inline function in a windows header file
73+ // (see https://learn.microsoft.com/en-us/windows/win32/api/memoryapi/nf-memoryapi-mapviewoffile2#remarks) so we use the same API it uses in the header file here instead of
74+ // MapViewOfFile2 which does not exist in the rust crate (see https://github.com/microsoft/windows-rs/issues/2595)
75+ let surrogate_base = unsafe {
76+ MapViewOfFileNuma2 (
77+ handle. into ( ) ,
78+ self . process_handle . into ( ) ,
79+ 0 ,
80+ None ,
81+ host_size,
82+ 0 ,
83+ PAGE_READWRITE . 0 ,
84+ NUMA_NO_PREFERRED_NODE ,
85+ )
86+ } ;
87+ let mut unused_out_old_prot_flags = PAGE_PROTECTION_FLAGS ( 0 ) ;
88+
89+ // the first page of the raw_size is the guard page
90+ let first_guard_page_start = surrogate_base. Value ;
91+ if let Err ( e) = unsafe {
92+ VirtualProtectEx (
93+ self . process_handle . into ( ) ,
94+ first_guard_page_start,
95+ PAGE_SIZE_USIZE ,
96+ PAGE_NOACCESS ,
97+ & mut unused_out_old_prot_flags,
98+ )
99+ } {
100+ log_then_return ! ( WindowsAPIError ( e. clone( ) ) ) ;
101+ }
102+
103+ // the last page of the raw_size is the guard page
104+ let last_guard_page_start =
105+ unsafe { first_guard_page_start. add ( host_size - PAGE_SIZE_USIZE ) } ;
106+ if let Err ( e) = unsafe {
107+ VirtualProtectEx (
108+ self . process_handle . into ( ) ,
109+ last_guard_page_start,
110+ PAGE_SIZE_USIZE ,
111+ PAGE_NOACCESS ,
112+ & mut unused_out_old_prot_flags,
113+ )
114+ } {
115+ log_then_return ! ( WindowsAPIError ( e. clone( ) ) ) ;
116+ }
117+ ve. insert ( HandleMapping {
118+ use_count : 1 ,
119+ surrogate_base : surrogate_base. Value ,
120+ } ) ;
121+ Ok ( surrogate_base. Value )
122+ }
123+ }
124+ }
125+
126+ pub ( super ) fn unmap ( & mut self , host_base : usize ) {
127+ match self . mappings . entry ( host_base) {
128+ Entry :: Occupied ( mut oe) => {
129+ oe. get_mut ( ) . use_count -= 1 ;
130+ if oe. get ( ) . use_count == 0 {
131+ let entry = oe. remove ( ) ;
132+ self . unmap_helper ( entry. surrogate_base ) ;
133+ }
134+ }
135+ Entry :: Vacant ( _) => {
136+ #[ cfg( debug_assertions) ]
137+ panic ! ( "Attempted to unmap from surrogate a region that was never mapped" )
138+ }
139+ }
140+ }
141+
142+ fn unmap_helper ( & self , surrogate_base : * mut c_void ) {
143+ let memory_mapped_view_address = MEMORY_MAPPED_VIEW_ADDRESS {
144+ Value : surrogate_base,
145+ } ;
146+ let flags = UNMAP_VIEW_OF_FILE_FLAGS ( 0 ) ;
147+ if let Err ( e) = unsafe {
148+ UnmapViewOfFile2 (
149+ self . process_handle . into ( ) ,
150+ memory_mapped_view_address,
151+ flags,
152+ )
153+ } {
154+ tracing:: error!(
155+ "Failed to free surrogate process resources (UnmapViewOfFile2 failed): {:?}" ,
156+ e
157+ ) ;
158+ }
159+ }
47160}
48161
49162impl Default for SurrogateProcess {
50163 #[ instrument( skip_all, parent = Span :: current( ) , level= "Trace" ) ]
51164 fn default ( ) -> Self {
52- let allocated_address = std:: ptr:: null_mut ( ) ;
53- Self :: new ( allocated_address, Default :: default ( ) )
165+ Self :: new ( Default :: default ( ) )
54166 }
55167}
56168
57169impl Drop for SurrogateProcess {
58170 #[ instrument( skip_all, parent = Span :: current( ) , level= "Trace" ) ]
59171 fn drop ( & mut self ) {
60- let process_handle: HANDLE = self . process_handle . into ( ) ;
61- let memory_mapped_view_address = MEMORY_MAPPED_VIEW_ADDRESS {
62- Value : self . allocated_address ,
63- } ;
64- let flags = UNMAP_VIEW_OF_FILE_FLAGS ( 0 ) ;
65- if let Err ( e) =
66- unsafe { UnmapViewOfFile2 ( process_handle, memory_mapped_view_address, flags) }
67- {
68- tracing:: error!(
69- "Failed to free surrogate process resources (UnmapViewOfFile2 failed): {:?}" ,
70- e
71- ) ;
172+ for mapping in self . mappings . values ( ) {
173+ self . unmap_helper ( mapping. surrogate_base ) ;
72174 }
73175
74176 // we need to do this take so we can take ownership
0 commit comments