Merge pull request microsoft#185385 from microsoft/cli-service-improvements

Cli service improvements

Author: connor4312

cli/src/commands/args.rs

@@ -6,7 +6,7 @@
 use std::{fmt, path::PathBuf};
 
 use crate::{constants, log, options, tunnels::code_server::CodeServerArgs};
-use clap::{ValueEnum, Args, Parser, Subcommand};
+use clap::{Args, Parser, Subcommand, ValueEnum};
 use const_format::concatcp;
 
 const CLI_NAME: &str = concatcp!(constants::PRODUCT_NAME_LONG, " CLI");

cli/src/commands/tunnels.rs

@@ -4,7 +4,8 @@
  *--------------------------------------------------------------------------------------------*/
 
 use async_trait::async_trait;
-use base64::{Engine as _, engine::general_purpose as b64};
+use base64::{engine::general_purpose as b64, Engine as _};
+use serde::Serialize;
 use sha2::{Digest, Sha256};
 use std::{str::FromStr, time::Duration};
 use sysinfo::Pid;
@@ -247,26 +248,39 @@ pub async fn kill(ctx: CommandContext) -> Result<i32, AnyError> {
 	.map_err(|e| e.into())
 }
 
+#[derive(Serialize)]
+pub struct StatusOutput {
+	pub tunnel: Option<protocol::singleton::TunnelState>,
+	pub service_installed: bool,
+}
+
 pub async fn status(ctx: CommandContext) -> Result<i32, AnyError> {
-	let status = do_single_rpc_call::<_, protocol::singleton::Status>(
+	let tunnel_status = do_single_rpc_call::<_, protocol::singleton::Status>(
 		&ctx.paths.tunnel_lockfile(),
 		ctx.log.clone(),
 		protocol::singleton::METHOD_STATUS,
 		protocol::EmptyObject {},
 	)
 	.await;
 
-	match status {
-		Err(CodeError::NoRunningTunnel) => {
-			ctx.log.result(CodeError::NoRunningTunnel.to_string());
-			Ok(1)
-		}
-		Err(e) => Err(e.into()),
-		Ok(s) => {
-			ctx.log.result(serde_json::to_string(&s).unwrap());
-			Ok(0)
-		}
-	}
+	let service_installed = create_service_manager(ctx.log.clone(), &ctx.paths)
+		.is_installed()
+		.await
+		.unwrap_or(false);
+
+	ctx.log.result(
+		serde_json::to_string(&StatusOutput {
+			service_installed,
+			tunnel: match tunnel_status {
+				Ok(s) => Some(s.tunnel),
+				Err(CodeError::NoRunningTunnel) => None,
+				Err(e) => return Err(e.into()),
+			},
+		})
+		.unwrap(),
+	);
+
+	Ok(0)
 }
 
 /// Removes unused servers.
@@ -331,6 +345,13 @@ 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
 
+	debug!(
+		log,
+		"Starting tunnel with `{} {}`",
+		APPLICATION_NAME,
+		std::env::args().collect::<Vec<_>>().join(" ")
+	);
+
 	// 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.
@@ -421,7 +442,10 @@ async fn serve_with_csa(
 
 				return Ok(exit.code().unwrap_or(1));
 			}
-			Next::Exit => return Ok(0),
+			Next::Exit => {
+				debug!(log, "Tunnel shut down");
+				return Ok(0);
+			}
 			Next::Restart => continue,
 		}
 	}

cli/src/lib.rs

@@ -18,8 +18,8 @@ pub mod tunnels;
 pub mod update_service;
 pub mod util;
 
-mod download_cache;
 mod async_pipe;
+mod download_cache;
 mod json_rpc;
 mod msgpack_rpc;
 mod rpc;

cli/src/log.rs

@@ -159,9 +159,21 @@ pub struct FileLogSink {
 	file: Arc<std::sync::Mutex<std::fs::File>>,
 }
 
+const FILE_LOG_SIZE_LIMIT: u64 = 1024 * 1024 * 10; // 10MB
+
 impl FileLogSink {
 	pub fn new(level: Level, path: &Path) -> std::io::Result<Self> {
-		let file = std::fs::File::create(path)?;
+		// Truncate the service log occasionally to avoid growing infinitely
+		if matches!(path.metadata(), Ok(m) if m.len() > FILE_LOG_SIZE_LIMIT) {
+			// ignore errors, can happen if another process is writing right now
+			let _ = std::fs::remove_file(path);
+		}
+
+		let file = std::fs::OpenOptions::new()
+			.append(true)
+			.create(true)
+			.open(path)?;
+
 		Ok(Self {
 			level,
 			file: Arc::new(std::sync::Mutex::new(file)),

cli/src/tunnels/service.rs

@@ -41,6 +41,9 @@ pub trait ServiceManager {
 	/// Show logs from the running service to standard out.
 	async fn show_logs(&self) -> Result<(), AnyError>;
 
+	/// Gets whether the tunnel service is installed.
+	async fn is_installed(&self) -> Result<bool, AnyError>;
+
 	/// Unregisters the current executable as a service.
 	async fn unregister(&self) -> Result<(), AnyError>;
 }

cli/src/tunnels/service_linux.rs

@@ -17,7 +17,7 @@ use crate::{
 	constants::{APPLICATION_NAME, PRODUCT_NAME_LONG},
 	log,
 	state::LauncherPaths,
-	util::errors::{wrap, AnyError},
+	util::errors::{wrap, AnyError, DbusConnectFailedError},
 };
 
 use super::ServiceManager;
@@ -40,7 +40,7 @@ impl SystemdService {
 	async fn connect() -> Result<Connection, AnyError> {
 		let connection = Connection::session()
 			.await
-			.map_err(|e| wrap(e, "error creating dbus session"))?;
+			.map_err(|e| DbusConnectFailedError(e.to_string()))?;
 		Ok(connection)
 	}
 
@@ -110,9 +110,27 @@ impl ServiceManager for SystemdService {
 
 		info!(self.log, "Tunnel service successfully started");
 
+		if std::env::var("SSH_CLIENT").is_ok() || std::env::var("SSH_TTY").is_ok() {
+			info!(self.log, "Tip: run `sudo loginctl enable-linger $USER` to ensure the service stays running after you disconnect.");
+		}
+
 		Ok(())
 	}
 
+	async fn is_installed(&self) -> Result<bool, AnyError> {
+		let connection = SystemdService::connect().await?;
+		let proxy = SystemdService::proxy(&connection).await?;
+		let state = proxy
+			.get_unit_file_state(SystemdService::service_name_string())
+			.await;
+
+		if let Ok(s) = state {
+			Ok(s == "enabled")
+		} else {
+			Ok(false)
+		}
+	}
+
 	async fn run(
 		self,
 		launcher_paths: crate::state::LauncherPaths,
@@ -219,6 +237,8 @@ trait SystemdManagerDbus {
 		force: bool,
 	) -> zbus::Result<(bool, Vec<(String, String, String)>)>;
 
+	fn get_unit_file_state(&self, file: String) -> zbus::Result<String>;
+
 	fn link_unit_files(
 		&self,
 		files: Vec<String>,

cli/src/tunnels/service_macos.rs

@@ -75,6 +75,11 @@ impl ServiceManager for LaunchdService {
 		handle.run_service(self.log, launcher_paths).await
 	}
 
+	async fn is_installed(&self) -> Result<bool, AnyError> {
+		let cmd = capture_command_and_check_status("launchctl", &["list"]).await?;
+		Ok(String::from_utf8_lossy(&cmd.stdout).contains(&get_service_label()))
+	}
+
 	async fn unregister(&self) -> Result<(), crate::util::errors::AnyError> {
 		let service_file = get_service_file_path()?;
 

cli/src/tunnels/service_windows.rs

@@ -114,6 +114,11 @@ impl CliServiceManager for WindowsService {
 		Ok(())
 	}
 
+	async fn is_installed(&self) -> Result<bool, AnyError> {
+		let key = WindowsService::open_key()?;
+		Ok(key.get_raw_value(TUNNEL_ACTIVITY_NAME).is_ok())
+	}
+
 	async fn unregister(&self) -> Result<(), AnyError> {
 		let key = WindowsService::open_key()?;
 		key.delete_value(TUNNEL_ACTIVITY_NAME)

cli/src/util/errors.rs

@@ -442,6 +442,26 @@ macro_rules! makeAnyError {
     };
 }
 
+#[derive(Debug)]
+pub struct DbusConnectFailedError(pub String);
+
+impl Display for DbusConnectFailedError {
+	fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+		let mut str = String::new();
+		str.push_str("Error creating dbus session. This command uses systemd for managing services, you should check that systemd is installed and under your user.");
+
+		if std::env::var("WSL_DISTRO_NAME").is_ok() {
+			str.push_str("\n\nTo enable systemd on WSL, check out: https://devblogs.microsoft.com/commandline/systemd-support-is-now-available-in-wsl/.\n\n");
+		}
+
+		str.push_str("If running `systemctl status` works, systemd is ok, but your session dbus may not be. You might need to:\n\n- Install the `dbus-user-session` package, and reboot if it was not installed\n- Start the user dbus session with `systemctl --user enable dbus --now`.\n\nThe error encountered was: ");
+		str.push_str(&self.0);
+		str.push('\n');
+
+		write!(f, "{}", str)
+	}
+}
+
 /// Internal errors in the VS Code CLI.
 /// Note: other error should be migrated to this type gradually
 #[derive(Error, Debug)]
@@ -522,7 +542,8 @@ makeAnyError!(
 	MissingHomeDirectory,
 	OAuthError,
 	InvalidRpcDataError,
-	CodeError
+	CodeError,
+	DbusConnectFailedError
 );
 
 impl From<reqwest::Error> for AnyError {

cli/src/util/tar.rs

@@ -6,7 +6,7 @@ use crate::util::errors::{wrap, WrappedError};
 
 use flate2::read::GzDecoder;
 use std::fs;
-use std::io::{Seek, SeekFrom};
+use std::io::Seek;
 use std::path::{Path, PathBuf};
 use tar::Archive;
 
@@ -65,7 +65,7 @@ where
 
 	// reset since skip logic read the tar already:
 	tar_gz
-		.seek(SeekFrom::Start(0))
+		.rewind()
 		.map_err(|e| wrap(e, "error resetting seek position"))?;
 
 	let tar = GzDecoder::new(tar_gz);