From eee5e7643a2481ee1dac9a7e75f922ccc4e40f40 Mon Sep 17 00:00:00 2001
From: Connor Peet <connor@peet.io>
Date: Fri, 20 Dec 2024 10:29:40 -0800
Subject: [PATCH] cli: propagate server-data-dir and extensions-dir when
 installing service (#236734)

Fixes #236195
---
 cli/src/bin/code/main.rs    |  4 +--
 cli/src/commands/tunnels.rs | 55 ++++++++++++++++++++++++-------------
 2 files changed, 38 insertions(+), 21 deletions(-)

diff --git a/cli/src/bin/code/main.rs b/cli/src/bin/code/main.rs
index 0e4809b78c7..b73d0aa885b 100644
--- a/cli/src/bin/code/main.rs
+++ b/cli/src/bin/code/main.rs
@@ -103,7 +103,7 @@ async fn main() -> Result<(), std::convert::Infallible> {
 				serve_web::serve_web(context!(), sw_args).await
 			}
 
-			Some(args::Commands::Tunnel(tunnel_args)) => match tunnel_args.subcommand {
+			Some(args::Commands::Tunnel(mut tunnel_args)) => match tunnel_args.subcommand.take() {
 				Some(args::TunnelSubcommand::Prune) => tunnels::prune(context!()).await,
 				Some(args::TunnelSubcommand::Unregister) => tunnels::unregister(context!()).await,
 				Some(args::TunnelSubcommand::Kill) => tunnels::kill(context!()).await,
@@ -116,7 +116,7 @@ async fn main() -> Result<(), std::convert::Infallible> {
 					tunnels::user(context!(), user_command).await
 				}
 				Some(args::TunnelSubcommand::Service(service_args)) => {
-					tunnels::service(context_no_logger(), service_args).await
+					tunnels::service(context_no_logger(), tunnel_args, service_args).await
 				}
 				Some(args::TunnelSubcommand::ForwardInternal(forward_args)) => {
 					tunnels::forward(context_no_logger(), forward_args).await
diff --git a/cli/src/commands/tunnels.rs b/cli/src/commands/tunnels.rs
index 2a0b4c7ddce..f52fa714793 100644
--- a/cli/src/commands/tunnels.rs
+++ b/cli/src/commands/tunnels.rs
@@ -21,7 +21,7 @@ use tokio::{
 
 use super::{
 	args::{
-		AuthProvider, CliCore, CommandShellArgs, ExistingTunnelArgs, TunnelForwardArgs,
+		AuthProvider, CliCore, CommandShellArgs, ExistingTunnelArgs, TunnelArgs, TunnelForwardArgs,
 		TunnelRenameArgs, TunnelServeArgs, TunnelServiceSubCommands, TunnelUserSubCommands,
 	},
 	CommandContext,
@@ -104,12 +104,16 @@ fn fulfill_existing_tunnel_args(
 }
 
 struct TunnelServiceContainer {
-	args: CliCore,
+	core_args: CliCore,
+	tunnel_args: TunnelArgs,
 }
 
 impl TunnelServiceContainer {
-	fn new(args: CliCore) -> Self {
-		Self { args }
+	fn new(core_args: CliCore, tunnel_args: TunnelArgs) -> Self {
+		Self {
+			core_args,
+			tunnel_args,
+		}
 	}
 }
 
@@ -120,7 +124,8 @@ impl ServiceContainer for TunnelServiceContainer {
 		log: log::Logger,
 		launcher_paths: LauncherPaths,
 	) -> Result<(), AnyError> {
-		let csa = (&self.args).into();
+		let mut csa = (&self.core_args).into();
+		self.tunnel_args.serve_args.server_args.apply_to(&mut csa);
 		serve_with_csa(
 			launcher_paths,
 			log,
@@ -242,8 +247,27 @@ async fn is_port_available(host: IpAddr, port: u16) -> bool {
 		.is_ok()
 }
 
+fn make_service_args<'a: 'c, 'b: 'c, 'c>(
+	root_path: &'a str,
+	tunnel_args: &'b TunnelArgs,
+) -> Vec<&'c str> {
+	let mut args = ["--verbose", "--cli-data-dir", root_path, "tunnel"].to_vec();
+
+	if let Some(d) = tunnel_args.serve_args.server_args.extensions_dir.as_ref() {
+		args.extend_from_slice(&["--extensions-dir", d]);
+	}
+	if let Some(d) = tunnel_args.serve_args.server_args.server_data_dir.as_ref() {
+		args.extend_from_slice(&["--server-data-dir", d]);
+	}
+
+	args.extend_from_slice(&["service", "internal-run"]);
+
+	args
+}
+
 pub async fn service(
 	ctx: CommandContext,
+	tunnel_args: TunnelArgs,
 	service_args: TunnelServiceSubCommands,
 ) -> Result<i32, AnyError> {
 	let manager = create_service_manager(ctx.log.clone(), &ctx.paths);
@@ -265,20 +289,10 @@ pub async fn service(
 			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 root_path = ctx.paths.root().as_os_str().to_string_lossy();
+			let args = make_service_args(&root_path, &tunnel_args);
 
-			manager
-				.register(
-					current_exe,
-					&[
-						"--verbose",
-						"--cli-data-dir",
-						ctx.paths.root().as_os_str().to_string_lossy().as_ref(),
-						"tunnel",
-						"service",
-						"internal-run",
-					],
-				)
-				.await?;
+			manager.register(current_exe, &args).await?;
 			ctx.log.result(format!("Service successfully installed! You can use `{APPLICATION_NAME} tunnel service log` to monitor it, and `{APPLICATION_NAME} tunnel service uninstall` to remove it."));
 		}
 		TunnelServiceSubCommands::Uninstall => {
@@ -289,7 +303,10 @@ pub async fn service(
 		}
 		TunnelServiceSubCommands::InternalRun => {
 			manager
-				.run(ctx.paths.clone(), TunnelServiceContainer::new(ctx.args))
+				.run(
+					ctx.paths.clone(),
+					TunnelServiceContainer::new(ctx.args, tunnel_args),
+				)
 				.await?;
 		}
 	}