stop tunnel when executable gets deleted (microsoft#181505)

stop when executable gets deleted

Author: aeschli

cli/src/commands/tunnels.rs

@@ -303,21 +303,23 @@ async fn serve_with_csa(
 	log = log.tee(log_broadcast.clone());
 	log::install_global_logger(log.clone()); // re-install so that library logs are captured
 
-	let shutdown = match gateway_args
-		.parent_process_id
-		.and_then(|p| Pid::from_str(&p).ok())
-	{
-		Some(pid) => ShutdownRequest::create_rx([
-			ShutdownRequest::CtrlC,
-			ShutdownRequest::ParentProcessKilled(pid),
-		]),
-		None => ShutdownRequest::create_rx([ShutdownRequest::CtrlC]),
-	};
-
 	// Intentionally read before starting the server. If the server updated and
 	// respawn is requested, the old binary will get renamed, and then
 	// current_exe will point to the wrong path.
 	let current_exe = std::env::current_exe().unwrap();
+
+	let mut vec = vec![
+		ShutdownRequest::CtrlC,
+		ShutdownRequest::ExeUninstalled(current_exe.to_owned()),
+	];
+	if let Some(p) = gateway_args
+		.parent_process_id
+		.and_then(|p| Pid::from_str(&p).ok())
+	{
+		vec.push(ShutdownRequest::ParentProcessKilled(p));
+	}
+	let shutdown = ShutdownRequest::create_rx(vec);
+
 	let server = loop {
 		if shutdown.is_open() {
 			return Ok(0);

cli/src/tunnels/shutdown_signal.rs

@@ -4,11 +4,11 @@
  *--------------------------------------------------------------------------------------------*/
 
 use futures::{stream::FuturesUnordered, StreamExt};
-use std::fmt;
+use std::{fmt, path::PathBuf};
 use sysinfo::Pid;
 
 use crate::util::{
-	machine::wait_until_process_exits,
+	machine::{wait_until_exe_deleted, wait_until_process_exits},
 	sync::{new_barrier, Barrier, Receivable},
 };
 
@@ -17,6 +17,7 @@ use crate::util::{
 pub enum ShutdownSignal {
 	CtrlC,
 	ParentProcessKilled(Pid),
+	ExeUninstalled,
 	ServiceStopped,
 	RpcShutdownRequested,
 	RpcRestartRequested,
@@ -29,6 +30,9 @@ impl fmt::Display for ShutdownSignal {
 			ShutdownSignal::ParentProcessKilled(p) => {
 				write!(f, "Parent process {} no longer exists", p)
 			}
+			ShutdownSignal::ExeUninstalled => {
+				write!(f, "Executable no longer exists")
+			}
 			ShutdownSignal::ServiceStopped => write!(f, "Service stopped"),
 			ShutdownSignal::RpcShutdownRequested => write!(f, "RPC client requested shutdown"),
 			ShutdownSignal::RpcRestartRequested => {
@@ -41,6 +45,7 @@ impl fmt::Display for ShutdownSignal {
 pub enum ShutdownRequest {
 	CtrlC,
 	ParentProcessKilled(Pid),
+	ExeUninstalled(PathBuf),
 	Derived(Box<dyn Receivable<ShutdownSignal> + Send>),
 }
 
@@ -56,6 +61,10 @@ impl ShutdownRequest {
 				wait_until_process_exits(pid, 2000).await;
 				Some(ShutdownSignal::ParentProcessKilled(pid))
 			}
+			ShutdownRequest::ExeUninstalled(exe_path) => {
+				wait_until_exe_deleted(&exe_path, 2000).await;
+				Some(ShutdownSignal::ExeUninstalled)
+			}
 			ShutdownRequest::Derived(mut rx) => rx.recv_msg().await,
 		}
 	}

cli/src/util/machine.rs

@@ -52,3 +52,10 @@ pub fn find_running_process(name: &Path) -> Option<u32> {
 	}
 	None
 }
+
+pub async fn wait_until_exe_deleted(current_exe: &Path, poll_ms: u64) {
+	let duration = Duration::from_millis(poll_ms);
+	while current_exe.exists() {
+		tokio::time::sleep(duration).await;
+	}
+}