@@ -18,14 +18,17 @@ use core::ffi::c_void;
1818use  std:: fmt; 
1919use  std:: fmt:: { Debug ,  Formatter } ; 
2020use  std:: string:: String ; 
21+ use  std:: sync:: atomic:: { AtomicBool ,  Ordering } ; 
22+ use  std:: sync:: Arc ; 
2123
2224use  hyperlight_common:: mem:: PAGE_SIZE_USIZE ; 
2325use  log:: LevelFilter ; 
2426use  tracing:: { instrument,  Span } ; 
2527use  windows:: Win32 :: System :: Hypervisor :: { 
26-     WHvX64RegisterCr0 ,  WHvX64RegisterCr3 ,  WHvX64RegisterCr4 ,  WHvX64RegisterCs ,  WHvX64RegisterEfer , 
27-     WHV_MEMORY_ACCESS_TYPE ,  WHV_PARTITION_HANDLE ,  WHV_REGISTER_VALUE ,  WHV_RUN_VP_EXIT_CONTEXT , 
28-     WHV_RUN_VP_EXIT_REASON ,  WHV_X64_SEGMENT_REGISTER ,  WHV_X64_SEGMENT_REGISTER_0 , 
28+     WHvCancelRunVirtualProcessor ,  WHvX64RegisterCr0 ,  WHvX64RegisterCr3 ,  WHvX64RegisterCr4 , 
29+     WHvX64RegisterCs ,  WHvX64RegisterEfer ,  WHV_MEMORY_ACCESS_TYPE ,  WHV_PARTITION_HANDLE , 
30+     WHV_REGISTER_VALUE ,  WHV_RUN_VP_EXIT_CONTEXT ,  WHV_RUN_VP_EXIT_REASON ,  WHV_X64_SEGMENT_REGISTER , 
31+     WHV_X64_SEGMENT_REGISTER_0 , 
2932} ; 
3033
3134use  super :: fpu:: { FP_TAG_WORD_DEFAULT ,  MXCSR_DEFAULT } ; 
@@ -37,11 +40,11 @@ use super::surrogate_process_manager::*;
3740use  super :: windows_hypervisor_platform:: { VMPartition ,  VMProcessor } ; 
3841use  super :: wrappers:: { HandleWrapper ,  WHvFPURegisters } ; 
3942use  super :: { 
40-     HyperlightExit ,  Hypervisor ,  VirtualCPU ,  CR0_AM ,  CR0_ET ,  CR0_MP ,  CR0_NE ,  CR0_PE ,  CR0_PG ,  CR0_WP , 
41-     CR4_OSFXSR ,  CR4_OSXMMEXCPT ,  CR4_PAE ,  EFER_LMA ,  EFER_LME ,  EFER_NX ,  EFER_SCE , 
43+     HyperlightExit ,  Hypervisor ,  InterruptHandle ,  VirtualCPU ,  CR0_AM ,  CR0_ET ,  CR0_MP ,  CR0_NE , 
44+     CR0_PE ,  CR0_PG ,  CR0_WP ,  CR4_OSFXSR ,  CR4_OSXMMEXCPT ,  CR4_PAE ,  EFER_LMA ,  EFER_LME ,  EFER_NX , 
45+     EFER_SCE , 
4246} ; 
4347use  crate :: hypervisor:: fpu:: FP_CONTROL_WORD_DEFAULT ; 
44- use  crate :: hypervisor:: hypervisor_handler:: HypervisorHandler ; 
4548use  crate :: hypervisor:: wrappers:: WHvGeneralRegisters ; 
4649use  crate :: mem:: memory_region:: { MemoryRegion ,  MemoryRegionFlags } ; 
4750use  crate :: mem:: ptr:: { GuestPtr ,  RawPtr } ; 
@@ -56,6 +59,7 @@ pub(crate) struct HypervWindowsDriver {
5659    entrypoint :  u64 , 
5760    orig_rsp :  GuestPtr , 
5861    mem_regions :  Vec < MemoryRegion > , 
62+     interrupt_handle :  Arc < WindowsInterruptHandle > , 
5963} 
6064/* This does not automatically impl Send/Sync because the host 
6165 * address of the shared memory region is a raw pointer, which are 
@@ -90,6 +94,7 @@ impl HypervWindowsDriver {
9094
9195        let  mut  proc = VMProcessor :: new ( partition) ?; 
9296        Self :: setup_initial_sregs ( & mut  proc,  pml4_address) ?; 
97+         let  partition_handle = proc. get_partition_hdl ( ) ; 
9398
9499        // subtract 2 pages for the guard pages, since when we copy memory to and from surrogate process, 
95100        // we don't want to copy the guard pages themselves (that would cause access violation) 
@@ -102,6 +107,11 @@ impl HypervWindowsDriver {
102107            entrypoint, 
103108            orig_rsp :  GuestPtr :: try_from ( RawPtr :: from ( rsp) ) ?, 
104109            mem_regions, 
110+             interrupt_handle :  Arc :: new ( WindowsInterruptHandle  { 
111+                 running :  AtomicBool :: new ( false ) , 
112+                 partition_handle, 
113+                 dropped :  AtomicBool :: new ( false ) , 
114+             } ) , 
105115        } ) 
106116    } 
107117
@@ -151,11 +161,6 @@ impl HypervWindowsDriver {
151161        error. push_str ( & format ! ( "Registers: \n {:#?}" ,  self . processor. get_regs( ) ?) ) ; 
152162        Ok ( error) 
153163    } 
154- 
155-     #[ instrument( skip_all,  parent = Span :: current( ) ,  level = "Trace" ) ]  
156-     pub ( crate )  fn  get_partition_hdl ( & self )  -> WHV_PARTITION_HANDLE  { 
157-         self . processor . get_partition_hdl ( ) 
158-     } 
159164} 
160165
161166impl  Debug  for  HypervWindowsDriver  { 
@@ -307,7 +312,6 @@ impl Hypervisor for HypervWindowsDriver {
307312        page_size :  u32 , 
308313        outb_hdl :  OutBHandlerWrapper , 
309314        mem_access_hdl :  MemAccessHandlerWrapper , 
310-         hv_handler :  Option < HypervisorHandler > , 
311315        max_guest_log_level :  Option < LevelFilter > , 
312316        #[ cfg( gdb) ]   dbg_mem_access_hdl :  DbgMemAccessHandlerWrapper , 
313317    )  -> Result < ( ) >  { 
@@ -333,7 +337,6 @@ impl Hypervisor for HypervWindowsDriver {
333337
334338        VirtualCPU :: run ( 
335339            self . as_mut_hypervisor ( ) , 
336-             hv_handler, 
337340            outb_hdl, 
338341            mem_access_hdl, 
339342            #[ cfg( gdb) ]  
@@ -349,7 +352,6 @@ impl Hypervisor for HypervWindowsDriver {
349352        dispatch_func_addr :  RawPtr , 
350353        outb_hdl :  OutBHandlerWrapper , 
351354        mem_access_hdl :  MemAccessHandlerWrapper , 
352-         hv_handler :  Option < HypervisorHandler > , 
353355        #[ cfg( gdb) ]   dbg_mem_access_hdl :  DbgMemAccessHandlerWrapper , 
354356    )  -> Result < ( ) >  { 
355357        // Reset general purpose registers, then set RIP and RSP 
@@ -371,7 +373,6 @@ impl Hypervisor for HypervWindowsDriver {
371373
372374        VirtualCPU :: run ( 
373375            self . as_mut_hypervisor ( ) , 
374-             hv_handler, 
375376            outb_hdl, 
376377            mem_access_hdl, 
377378            #[ cfg( gdb) ]  
@@ -407,7 +408,11 @@ impl Hypervisor for HypervWindowsDriver {
407408
408409    #[ instrument( err( Debug ) ,  skip_all,  parent = Span :: current( ) ,  level = "Trace" ) ]  
409410    fn  run ( & mut  self )  -> Result < super :: HyperlightExit >  { 
411+         self . interrupt_handle . running . store ( true ,  Ordering :: Relaxed ) ; 
410412        let  exit_context:  WHV_RUN_VP_EXIT_CONTEXT  = self . processor . run ( ) ?; 
413+         self . interrupt_handle 
414+             . running 
415+             . store ( false ,  Ordering :: Relaxed ) ; 
411416
412417        let  result = match  exit_context. ExitReason  { 
413418            // WHvRunVpExitReasonX64IoPortAccess 
@@ -481,8 +486,8 @@ impl Hypervisor for HypervWindowsDriver {
481486        Ok ( result) 
482487    } 
483488
484-     fn  get_partition_handle ( & self )  -> WHV_PARTITION_HANDLE  { 
485-         self . processor . get_partition_hdl ( ) 
489+     fn  interrupt_handle ( & self )  -> Arc < dyn   InterruptHandle >  { 
490+         self . interrupt_handle . clone ( ) 
486491    } 
487492
488493    #[ instrument( skip_all,  parent = Span :: current( ) ,  level = "Trace" ) ]  
@@ -496,28 +501,26 @@ impl Hypervisor for HypervWindowsDriver {
496501    } 
497502} 
498503
499- #[ cfg( test) ]  
500- pub  mod  tests { 
501-     use  std:: sync:: { Arc ,  Mutex } ; 
504+ impl  Drop  for  HypervWindowsDriver  { 
505+     fn  drop ( & mut  self )  { 
506+         self . interrupt_handle . dropped . store ( true ,  Ordering :: Relaxed ) ; 
507+     } 
508+ } 
502509
503-     use  serial_test:: serial; 
510+ pub  struct  WindowsInterruptHandle  { 
511+     // `WHvCancelRunVirtualProcessor()` will return Ok even if the vcpu is not running, which is the reason we need this flag. 
512+     running :  AtomicBool , 
513+     partition_handle :  WHV_PARTITION_HANDLE , 
514+     dropped :  AtomicBool , 
515+ } 
504516
505-     use  crate :: hypervisor:: handlers:: { MemAccessHandler ,  OutBHandler } ; 
506-     use  crate :: hypervisor:: tests:: test_initialise; 
507-     use  crate :: Result ; 
517+ impl  InterruptHandle  for  WindowsInterruptHandle  { 
518+     fn  kill ( & self )  -> bool  { 
519+         self . running . load ( Ordering :: Relaxed ) 
520+             && unsafe  {  WHvCancelRunVirtualProcessor ( self . partition_handle ,  0 ,  0 ) . is_ok ( )  } 
521+     } 
508522
509-     #[ test]  
510-     #[ serial]  
511-     fn  test_init ( )  { 
512-         let  outb_handler = { 
513-             let  func:  Box < dyn  FnMut ( u16 ,  u32 )  -> Result < ( ) >  + Send >  =
514-                 Box :: new ( |_,  _| -> Result < ( ) >  {  Ok ( ( ) )  } ) ; 
515-             Arc :: new ( Mutex :: new ( OutBHandler :: from ( func) ) ) 
516-         } ; 
517-         let  mem_access_handler = { 
518-             let  func:  Box < dyn  FnMut ( )  -> Result < ( ) >  + Send >  = Box :: new ( || -> Result < ( ) >  {  Ok ( ( ) )  } ) ; 
519-             Arc :: new ( Mutex :: new ( MemAccessHandler :: from ( func) ) ) 
520-         } ; 
521-         test_initialise ( outb_handler,  mem_access_handler) . unwrap ( ) ; 
523+     fn  dropped ( & self )  -> bool  { 
524+         self . dropped . load ( Ordering :: Relaxed ) 
522525    } 
523526} 
0 commit comments