@@ -35,7 +35,7 @@ use std::{
35
35
env,
36
36
ffi:: { OsString , c_void} ,
37
37
iter,
38
- path:: PathBuf ,
38
+ path:: { Path , PathBuf } ,
39
39
ptr,
40
40
sync:: {
41
41
Arc , Mutex , OnceLock ,
@@ -48,7 +48,7 @@ use std::{
48
48
49
49
use windows:: {
50
50
Win32 :: {
51
- Foundation :: { HANDLE , LUID } ,
51
+ Foundation :: { HANDLE , LUID , STILL_ACTIVE } ,
52
52
Security :: {
53
53
AdjustTokenPrivileges , DuplicateTokenEx , LookupPrivilegeValueW , SE_PRIVILEGE_ENABLED ,
54
54
SE_PRIVILEGE_REMOVED , SE_TCB_NAME , SecurityIdentification , TOKEN_ACCESS_MASK ,
@@ -454,7 +454,7 @@ fn run_as<T>(session_id: u32, cb: impl FnOnce(HANDLE) -> Result<T>) -> Result<T>
454
454
}
455
455
456
456
/// Newtype over handle which closes the HANDLE on drop.
457
- #[ derive( Default ) ]
457
+ #[ derive( Clone , Default ) ]
458
458
struct OwnedHandle ( HANDLE ) ;
459
459
460
460
unsafe impl Send for OwnedHandle { }
@@ -474,46 +474,103 @@ impl From<HANDLE> for OwnedHandle {
474
474
475
475
impl Drop for OwnedHandle {
476
476
fn drop ( & mut self ) {
477
- unsafe {
478
- self . 0 . free ( ) ;
477
+ if self . is_valid ( ) {
478
+ unsafe {
479
+ self . 0 . free ( ) ;
480
+ }
479
481
}
480
482
}
481
483
}
482
484
483
- /// A child process. Tries to kill the process when dropped.
484
- struct Child ( OwnedHandle ) ;
485
+ /// A quick wrapper over a child process. Child is terminated when this is dropped
486
+ #[ derive( Clone ) ]
487
+ struct Child ( Arc < OwnedHandle > ) ;
485
488
486
489
impl Child {
487
- fn new ( ) -> Self {
488
- Self ( OwnedHandle :: default ( ) )
490
+ /// Spawn a child as the user provided by token
491
+ fn spawn_as ( token : HANDLE , exe : impl AsRef < Path > , args : impl AsRef < str > ) -> Result < Self > {
492
+ let mut exe = exe
493
+ . as_ref ( )
494
+ . to_string_lossy ( )
495
+ . to_string ( )
496
+ . encode_utf16 ( )
497
+ . chain ( iter:: once ( 0 ) )
498
+ . collect :: < Vec < _ > > ( ) ;
499
+
500
+ let exe = PWSTR ( exe. as_mut_ptr ( ) ) ;
501
+
502
+ let mut _args = args. as_ref ( ) ;
503
+ let mut args;
504
+
505
+ let args = if !_args. is_empty ( ) {
506
+ args = _args
507
+ . encode_utf16 ( )
508
+ . chain ( iter:: once ( 0 ) )
509
+ . collect :: < Vec < _ > > ( ) ;
510
+
511
+ Some ( PWSTR ( args. as_mut_ptr ( ) ) )
512
+ } else {
513
+ None
514
+ } ;
515
+
516
+ let env_block = EnvBlock :: new ( token) ?;
517
+
518
+ let mut process_info = PROCESS_INFORMATION :: default ( ) ;
519
+ unsafe {
520
+ CreateProcessAsUserW (
521
+ Some ( token) ,
522
+ exe,
523
+ args,
524
+ None ,
525
+ None ,
526
+ false ,
527
+ // CREATE_UNICODE_ENVIRONMENT is required if we pass env block.
528
+ // <https://learn.microsoft.com/en-us/windows/win32/api/userenv/nf-userenv-createenvironmentblock#remarks>
529
+ //
530
+ // CREATE_NO_WINDOW causes all child processes to not show a visible
531
+ // console window. <https://stackoverflow.com/a/71364777/9423933>
532
+ CREATE_UNICODE_ENVIRONMENT | CREATE_NO_WINDOW ,
533
+ Some ( env_block. 0 ) ,
534
+ None ,
535
+ & STARTUPINFOW :: default ( ) ,
536
+ & mut process_info,
537
+ ) ?;
538
+ }
539
+
540
+ // It is safe for EnvBlock to drop after calling CreateProcessAsUser.
541
+ // <https://learn.microsoft.com/en-us/windows/win32/api/userenv/nf-userenv-createenvironmentblock#remarks>
542
+
543
+ let this = Self ( Arc :: new ( OwnedHandle ( process_info. hProcess ) ) ) ;
544
+
545
+ Ok ( this)
489
546
}
490
547
491
- fn kill ( & mut self ) -> Result < ( ) > {
492
- if self . 0 . is_valid ( ) {
493
- unsafe {
494
- TerminateProcess ( self . 0 . 0 , 0 ) ?;
495
- }
548
+ fn is_alive ( & self ) -> bool {
549
+ self . exit_code ( ) as i32 == STILL_ACTIVE . 0
550
+ }
496
551
497
- self . 0 = OwnedHandle :: default ( ) ;
498
- }
552
+ fn exit_code ( & self ) -> u32 {
553
+ let mut code = 0u32 ;
554
+ _ = unsafe { GetExitCodeProcess ( self . 0 . 0 , & mut code) } ;
499
555
500
- Ok ( ( ) )
556
+ code
501
557
}
502
558
503
- fn reset ( & mut self ) {
504
- self . 0 = OwnedHandle :: default ( ) ;
559
+ fn kill ( & self ) {
560
+ _ = unsafe { TerminateProcess ( self . 0 . 0 , 0 ) } ;
505
561
}
506
- }
507
562
508
- impl Default for Child {
509
- fn default ( ) -> Self {
510
- Self :: new ( )
563
+ /// wait until child dies
564
+ fn wait ( & self ) {
565
+ unsafe {
566
+ WaitForSingleObject ( self . 0 . 0 , INFINITE ) ;
567
+ }
511
568
}
512
569
}
513
570
514
571
impl Drop for Child {
515
572
fn drop ( & mut self ) {
516
- _ = self . kill ( ) ;
573
+ self . kill ( ) ;
517
574
}
518
575
}
519
576
@@ -542,10 +599,8 @@ impl Drop for EnvBlock {
542
599
/// Manages the child daemon, by spawning / stopping it, or reporting abnormal exit, and allowing
543
600
/// wait().
544
601
struct Spawner {
545
- // Whether a child daemon is running.
546
- running : Arc < AtomicBool > ,
547
602
// Holds the actual process of the running child daemon.
548
- child : Arc < Mutex < Child > > ,
603
+ child : Arc < Mutex < Option < Child > > > ,
549
604
// Whether the process has exited without our request.
550
605
dirty : Arc < AtomicBool > ,
551
606
// Used to differentiate between requested stop() and if process is dirty (see above).
@@ -561,7 +616,6 @@ impl Spawner {
561
616
let ( wait_tx, wait_rx) = channel ( ) ;
562
617
563
618
Self {
564
- running : Arc :: default ( ) ,
565
619
child : Arc :: default ( ) ,
566
620
dirty : Arc :: default ( ) ,
567
621
request_stop : Arc :: default ( ) ,
@@ -572,7 +626,11 @@ impl Spawner {
572
626
573
627
/// Is the child daemon running?
574
628
fn running ( & self ) -> bool {
575
- self . running . load ( Ordering :: Relaxed )
629
+ self . child
630
+ . lock ( )
631
+ . unwrap ( )
632
+ . as_ref ( )
633
+ . is_some_and ( |c| c. is_alive ( ) )
576
634
}
577
635
578
636
/// Stop the spawned daemon.
@@ -581,25 +639,18 @@ impl Spawner {
581
639
/// make _sure_ you put this stop() _after_ you change the `while` condition to false
582
640
/// otherwise any condition change will not be observed.
583
641
fn stop ( & self ) {
584
- let mut child = self . child . lock ( ) . unwrap ( ) ;
585
-
586
642
// Request a normal stop. This is not an abnormal process exit.
587
643
self . request_stop . store ( true , Ordering :: Relaxed ) ;
588
- match child. kill ( ) {
589
- Ok ( _) => {
590
- debug ! ( "stop() kill" ) ;
591
- self . running . store ( false , Ordering :: Relaxed ) ;
592
- // Signal the wait() to exit so a `while` condition is checked at least once more.
593
- // As long as `while` conditions have been changed _before_ the call to stop(),
594
- // the changed condition will be observed.
595
- _ = self . wait_tx . send ( ( ) ) ;
596
- }
597
644
598
- Err ( e) => {
599
- self . running . store ( false , Ordering :: Relaxed ) ;
600
- error ! ( "failed to stop(): {e}" ) ;
601
- }
645
+ if let Some ( child) = self . child . lock ( ) . unwrap ( ) . as_ref ( ) {
646
+ child. kill ( ) ;
647
+ debug ! ( "stop() killed" ) ;
602
648
}
649
+
650
+ // Signal the wait() to exit so a `while` condition is checked at least once more.
651
+ // As long as `while` conditions have been changed _before_ the call to stop(),
652
+ // the changed condition will be observed.
653
+ _ = self . wait_tx . send ( ( ) ) ;
603
654
}
604
655
605
656
/// Wait for child process to exit.
@@ -614,97 +665,55 @@ impl Spawner {
614
665
615
666
/// Try to spawn a child daemon.
616
667
fn start ( & self , session : u32 ) -> Result < ( ) > {
617
- let running = self . running . clone ( ) ;
618
668
let child = self . child . clone ( ) ;
619
669
let waiter = self . wait_tx . clone ( ) ;
620
670
let dirty = self . dirty . clone ( ) ;
621
671
let request_stop = self . request_stop . clone ( ) ;
672
+
622
673
_ = thread:: spawn ( move || {
623
674
request_stop. store ( false , Ordering :: Relaxed ) ;
624
675
625
676
let res = run_as ( session, move |token| {
626
- let mut arguments = Vec :: new ( ) ;
627
-
628
677
let config = CONFIG
629
678
. get ( )
630
679
. ok_or_else ( || eyre ! ( "failed to get CONFIG" ) ) ?
631
680
. clone ( ) ;
632
681
682
+ // Try to get the path to the current binary
683
+ let current_exe = env:: current_exe ( ) ?;
684
+
685
+ let mut args = Vec :: new ( ) ;
686
+
687
+ // first argument MUST be the exe name
688
+ args. push ( format ! ( r#""{exe}""# , exe = current_exe. display( ) ) ) ;
689
+
633
690
if let Some ( config) = & config. config_path {
634
- arguments . push ( "--config" . to_string ( ) ) ;
635
- arguments . push ( format ! ( r#""{}""# , config. to_string_lossy( ) . into_owned( ) ) ) ;
691
+ args . push ( "--config" . to_string ( ) ) ;
692
+ args . push ( format ! ( r#""{}""# , config. to_string_lossy( ) . into_owned( ) ) ) ;
636
693
}
637
694
638
695
if let Some ( profile) = & config. profile {
639
- arguments . push ( "--profile" . to_string ( ) ) ;
640
- arguments . push ( format ! ( r#""{profile}""# ) ) ;
696
+ args . push ( "--profile" . to_string ( ) ) ;
697
+ args . push ( format ! ( r#""{profile}""# ) ) ;
641
698
}
642
699
643
- let arguments = arguments . join ( " " ) ;
700
+ let args = args . join ( " " ) ;
644
701
645
- // Try to get the path to the current binary
646
- let mut current_exe = env:: current_exe ( ) ?
647
- . to_string_lossy ( )
648
- . to_string ( )
649
- . encode_utf16 ( )
650
- . chain ( iter:: once ( 0 ) )
651
- . collect :: < Vec < _ > > ( ) ;
652
-
653
- let mut arguments = arguments
654
- . encode_utf16 ( )
655
- . chain ( iter:: once ( 0 ) )
656
- . collect :: < Vec < _ > > ( ) ;
657
-
658
- let env_block = EnvBlock :: new ( token) ?;
659
-
660
- let mut process_info = PROCESS_INFORMATION :: default ( ) ;
661
- unsafe {
662
- CreateProcessAsUserW (
663
- Some ( token) ,
664
- PWSTR ( current_exe. as_mut_ptr ( ) ) ,
665
- Some ( PWSTR ( arguments. as_mut_ptr ( ) ) ) ,
666
- None ,
667
- None ,
668
- false ,
669
- // CREATE_UNICODE_ENVIRONMENT is required if we pass env block.
670
- // <https://learn.microsoft.com/en-us/windows/win32/api/userenv/nf-userenv-createenvironmentblock#remarks>
671
- //
672
- // CREATE_NO_WINDOW causes all child processes to not show a visible
673
- // console window. <https://stackoverflow.com/a/71364777/9423933>
674
- CREATE_UNICODE_ENVIRONMENT | CREATE_NO_WINDOW ,
675
- Some ( env_block. 0 ) ,
676
- None ,
677
- & STARTUPINFOW :: default ( ) ,
678
- & mut process_info,
679
- ) ?;
680
- }
681
-
682
- // It is safe to drop this after calling CreateProcessAsUser.
683
- // <https://learn.microsoft.com/en-us/windows/win32/api/userenv/nf-userenv-createenvironmentblock#remarks>
684
- drop ( env_block) ;
702
+ let spawned_child = Child :: spawn_as ( token, current_exe, args) ?;
685
703
686
704
// Store the child process.
687
705
{
688
706
let mut lock = child. lock ( ) . unwrap ( ) ;
689
- * lock = Child ( process_info . hProcess . into ( ) ) ;
707
+ * lock = Some ( spawned_child . clone ( ) ) ;
690
708
}
691
709
692
- running. store ( true , Ordering :: Relaxed ) ;
693
-
694
710
// Wait until the process exits.
695
- unsafe {
696
- WaitForSingleObject ( process_info. hProcess , INFINITE ) ;
697
- }
698
-
699
- running. store ( false , Ordering :: Relaxed ) ;
711
+ spawned_child. wait ( ) ;
700
712
701
713
// Check if process exited on its own without our explicit request (`stop()` was not
702
714
// called).
703
715
if !request_stop. swap ( false , Ordering :: Relaxed ) {
704
- let mut code = 0u32 ;
705
- unsafe {
706
- GetExitCodeProcess ( process_info. hProcess , & mut code) ?;
707
- }
716
+ let code = spawned_child. exit_code ( ) ;
708
717
709
718
debug ! ( "spawner code {code}" ) ;
710
719
@@ -717,7 +726,7 @@ impl Spawner {
717
726
}
718
727
}
719
728
720
- child. lock ( ) . unwrap ( ) . reset ( ) ;
729
+ child. lock ( ) . unwrap ( ) . take ( ) ;
721
730
722
731
Ok ( ( ) )
723
732
} ) ;
0 commit comments