cli: add --install-extension command for use with tunnel server (#207741)

* cli: add --install-extension command for use with tunnel server

* fix clippy

cli/src/bin/code/legacy_args.rs

@@ -42,6 +42,9 @@ pub fn try_parse_legacy(
 				}
 			}
 		} else if let Ok(value) = arg.to_value() {
+			if value == "tunnel" {
+				return None;
+			}
 			if let Some(last_arg) = &last_arg {
 				args.get_mut(last_arg)
 					.expect("expected to have last arg")

cli/src/commands/args.rs

@@ -655,6 +655,17 @@ pub struct TunnelServeArgs {
 	/// If set, the user accepts the server license terms and the server will be started without a user prompt.
 	#[clap(long)]
 	pub accept_server_license_terms: bool,
+
+	/// Requests that extensions be preloaded and installed on connecting servers.
+	#[clap(long)]
+	pub install_extension: Vec<String>,
+}
+
+impl TunnelServeArgs {
+	pub fn apply_to_server_args(&self, csa: &mut CodeServerArgs) {
+		csa.install_extensions
+			.extend_from_slice(&self.install_extension);
+	}
 }
 
 #[derive(Args, Debug, Clone)]

cli/src/commands/tunnels.rs

@@ -50,7 +50,11 @@ use crate::{
 		AuthRequired, Next, ServeStreamParams, ServiceContainer, ServiceManager,
 	},
 	util::{
-		app_lock::AppMutex, command::new_std_command, errors::{wrap, AnyError, CodeError}, machine::canonical_exe, prereqs::PreReqChecker
+		app_lock::AppMutex,
+		command::new_std_command,
+		errors::{wrap, AnyError, CodeError},
+		machine::canonical_exe,
+		prereqs::PreReqChecker,
 	},
 };
 use crate::{
@@ -227,8 +231,7 @@ pub async fn service(
 			// likewise for license consent
 			legal::require_consent(&ctx.paths, args.accept_server_license_terms)?;
 
-			let current_exe =
-				canonical_exe().map_err(|e| wrap(e, "could not get current exe"))?;
+			let current_exe = canonical_exe().map_err(|e| wrap(e, "could not get current exe"))?;
 
 			manager
 				.register(
@@ -404,7 +407,8 @@ pub async fn serve(ctx: CommandContext, gateway_args: TunnelServeArgs) -> Result
 
 	legal::require_consent(&paths, gateway_args.accept_server_license_terms)?;
 
-	let csa = (&args).into();
+	let mut csa = (&args).into();
+	gateway_args.apply_to_server_args(&mut csa);
 	let result = serve_with_csa(paths, log, gateway_args, csa, TUNNEL_CLI_LOCK_NAME).await;
 	drop(no_sleep);
 

cli/src/tunnels/code_server.rs

@@ -15,7 +15,8 @@ use crate::update_service::{
 	unzip_downloaded_release, Platform, Release, TargetKind, UpdateService,
 };
 use crate::util::command::{
-	capture_command, capture_command_and_check_status, kill_tree, new_script_command,
+	capture_command, capture_command_and_check_status, check_output_status, kill_tree,
+	new_script_command,
 };
 use crate::util::errors::{wrap, AnyError, CodeError, ExtensionInstallFailed, WrappedError};
 use crate::util::http::{self, BoxedHttp};
@@ -488,6 +489,28 @@ impl<'a> ServerBuilder<'a> {
 		})
 	}
 
+	/// Runs the command that just installs extensions and exits.
+	pub async fn install_extensions(&self) -> Result<(), AnyError> {
+		// cmd already has --install-extensions from base
+		let mut cmd = self.get_base_command();
+		let cmd_str = || {
+			self.server_params
+				.code_server_args
+				.command_arguments()
+				.join(" ")
+		};
+
+		let r = cmd.output().await.map_err(|e| CodeError::CommandFailed {
+			command: cmd_str(),
+			code: -1,
+			output: e.to_string(),
+		})?;
+
+		check_output_status(r, cmd_str)?;
+
+		Ok(())
+	}
+
 	pub async fn listen_on_default_socket(&self) -> Result<SocketCodeServer, AnyError> {
 		let requested_file = get_socket_name();
 		self.listen_on_socket(&requested_file).await

cli/src/tunnels/control_server.rs

@@ -6,6 +6,7 @@ use crate::async_pipe::get_socket_rw_stream;
 use crate::constants::{CONTROL_PORT, PRODUCT_NAME_LONG};
 use crate::log;
 use crate::msgpack_rpc::{new_msgpack_rpc, start_msgpack_rpc, MsgPackCodec, MsgPackSerializer};
+use crate::options::Quality;
 use crate::rpc::{MaybeSync, RpcBuilder, RpcCaller, RpcDispatcher};
 use crate::self_update::SelfUpdate;
 use crate::state::LauncherPaths;
@@ -144,6 +145,31 @@ pub struct ServerTermination {
 	pub tunnel: ActiveTunnel,
 }
 
+async fn preload_extensions(
+	log: &log::Logger,
+	platform: Platform,
+	mut args: CodeServerArgs,
+	launcher_paths: LauncherPaths,
+) -> Result<(), AnyError> {
+	args.start_server = false;
+
+	let params_raw = ServerParamsRaw {
+		commit_id: None,
+		quality: Quality::Stable,
+		code_server_args: args.clone(),
+		headless: true,
+		platform,
+	};
+
+	// cannot use delegated HTTP here since there's no remote connection yet
+	let http = Arc::new(ReqwestSimpleHttp::new());
+	let resolved = params_raw.resolve(log, http.clone()).await?;
+	let sb = ServerBuilder::new(log, &resolved, &launcher_paths, http.clone());
+
+	sb.setup().await?;
+	sb.install_extensions().await
+}
+
 // Runs the launcher server. Exits on a ctrl+c or when requested by a user.
 // Note that client connections may not be closed when this returns; use
 // `close_all_clients()` on the ServerTermination to make this happen.
@@ -160,6 +186,26 @@ pub async fn serve(
 	let (tx, mut rx) = mpsc::channel::<ServerSignal>(4);
 	let (exit_barrier, signal_exit) = new_barrier();
 
+	if !code_server_args.install_extensions.is_empty() {
+		info!(
+			log,
+			"Preloading extensions using stable server: {:?}", code_server_args.install_extensions
+		);
+		let log = log.clone();
+		let code_server_args = code_server_args.clone();
+		let launcher_paths = launcher_paths.clone();
+		// This is run async to the primary tunnel setup to be speedy.
+		tokio::spawn(async move {
+			if let Err(e) =
+				preload_extensions(&log, platform, code_server_args, launcher_paths).await
+			{
+				warning!(log, "Failed to preload extensions: {:?}", e);
+			} else {
+				info!(log, "Extension install complete");
+			}
+		});
+	}
+
 	loop {
 		tokio::select! {
 			Ok(reason) = shutdown_rx.wait() => {