From 9c001af44fd15d6fc779fa611c1bf3278dfd46b8 Mon Sep 17 00:00:00 2001
From: Connor Peet <connor@peet.io>
Date: Wed, 28 Feb 2024 15:48:08 -0800
Subject: [PATCH 1/2] cli: use connection token file in serve-web

---
 cli/src/commands/args.rs      |  3 +++
 cli/src/commands/serve_web.rs | 49 +++++++++++++++++++++++++++--------
 cli/src/util/errors.rs        |  2 ++
 3 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/cli/src/commands/args.rs b/cli/src/commands/args.rs
index a69d021cab866..a64f196089d28 100644
--- a/cli/src/commands/args.rs
+++ b/cli/src/commands/args.rs
@@ -197,6 +197,9 @@ pub struct ServeWebArgs {
 	/// A secret that must be included with all requests.
 	#[clap(long)]
 	pub connection_token: Option<String>,
+	/// A file containing a secret that must be included with all requests.
+	#[clap(long)]
+	pub connection_token_file: Option<String>,
 	/// Run without a connection token. Only use this if the connection is secured by other means.
 	#[clap(long)]
 	pub without_connection_token: bool,
diff --git a/cli/src/commands/serve_web.rs b/cli/src/commands/serve_web.rs
index 959763a431d10..9b4a0c47b15d7 100644
--- a/cli/src/commands/serve_web.rs
+++ b/cli/src/commands/serve_web.rs
@@ -6,13 +6,14 @@
 use std::collections::HashMap;
 use std::convert::Infallible;
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 
 use const_format::concatcp;
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{Body, Request, Response, Server};
+use tempfile::{tempdir, TempDir};
 use tokio::io::{AsyncBufReadExt, BufReader};
 use tokio::pin;
 
@@ -76,16 +77,21 @@ pub async fn serve_web(ctx: CommandContext, mut args: ServeWebArgs) -> Result<i3
 	legal::require_consent(&ctx.paths, args.accept_server_license_terms)?;
 
 	let platform: crate::update_service::Platform = PreReqChecker::new().verify().await?;
-
-	if !args.without_connection_token {
+	let token_file = if !args.without_connection_token {
 		// Ensure there's a defined connection token, since if multiple server versions
 		// are excuted, they will need to have a single shared token.
-		args.connection_token = Some(
-			args.connection_token
-				.clone()
-				.unwrap_or_else(|| uuid::Uuid::new_v4().to_string()),
-		);
-	}
+		let connection_token = args
+			.connection_token
+			.clone()
+			.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
+		let tf = ConnectionTokenFile::new(&connection_token)
+			.map_err(CodeError::CouldNotCreateConnectionTokenFile)?;
+		args.connection_token = Some(connection_token);
+		args.connection_token_file = Some(tf.path().to_string_lossy().to_string());
+		Some(tf)
+	} else {
+		None
+	};
 
 	let cm = ConnectionManager::new(&ctx, platform, args.clone());
 	let key = get_server_key_half(&ctx.paths);
@@ -137,6 +143,7 @@ pub async fn serve_web(ctx: CommandContext, mut args: ServeWebArgs) -> Result<i3
 	};
 
 	r.map_err(CodeError::CouldNotListenOnInterface)?;
+	drop(token_file); // ensure it lives long enough
 
 	Ok(0)
 }
@@ -704,8 +711,10 @@ impl ConnectionManager {
 		if args.args.without_connection_token {
 			cmd.arg("--without-connection-token");
 		}
-		if let Some(ct) = &args.args.connection_token {
-			cmd.arg("--connection-token");
+		// Note: intentional that we don't pass --connection-token here, we always
+		// convert it into the file variant.
+		if let Some(ct) = &args.args.connection_token_file {
+			cmd.arg("--connection-token-file");
 			cmd.arg(ct);
 		}
 
@@ -779,3 +788,21 @@ struct StartArgs {
 	release: Release,
 	opener: BarrierOpener<Result<StartData, String>>,
 }
+
+struct ConnectionTokenFile {
+	path: PathBuf,
+	_dir: TempDir, // implements Drop to delete the dir
+}
+
+impl ConnectionTokenFile {
+	fn new(connection_token: &str) -> std::io::Result<Self> {
+		let d = tempdir()?;
+		let path = d.path().join("connection-token");
+		std::fs::write(&path, connection_token)?;
+		Ok(ConnectionTokenFile { path, _dir: d })
+	}
+
+	fn path(&self) -> &Path {
+		&self.path
+	}
+}
diff --git a/cli/src/util/errors.rs b/cli/src/util/errors.rs
index 38d9b36f54bf5..684dbcb88d2de 100644
--- a/cli/src/util/errors.rs
+++ b/cli/src/util/errors.rs
@@ -512,6 +512,8 @@ pub enum CodeError {
 	// todo: can be specialized when update service is moved to CodeErrors
 	#[error("Could not check for update: {0}")]
 	UpdateCheckFailed(String),
+	#[error("Could not write connection token file: {0}")]
+	CouldNotCreateConnectionTokenFile(std::io::Error)
 }
 
 makeAnyError!(

From 8e3246289ce1525befd47cf5f1590ce5bc5da067 Mon Sep 17 00:00:00 2001
From: Connor Peet <connor@peet.io>
Date: Wed, 6 Mar 2024 08:51:57 -0800
Subject: [PATCH 2/2] address comments

---
 cli/src/commands/serve_web.rs | 59 ++++++++++++++++++-----------------
 1 file changed, 31 insertions(+), 28 deletions(-)

diff --git a/cli/src/commands/serve_web.rs b/cli/src/commands/serve_web.rs
index 9b4a0c47b15d7..d9e525af6c693 100644
--- a/cli/src/commands/serve_web.rs
+++ b/cli/src/commands/serve_web.rs
@@ -5,6 +5,8 @@
 
 use std::collections::HashMap;
 use std::convert::Infallible;
+use std::fs;
+use std::io::{Read, Write};
 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex};
@@ -13,7 +15,6 @@ use std::time::{Duration, Instant};
 use const_format::concatcp;
 use hyper::service::{make_service_fn, service_fn};
 use hyper::{Body, Request, Response, Server};
-use tempfile::{tempdir, TempDir};
 use tokio::io::{AsyncBufReadExt, BufReader};
 use tokio::pin;
 
@@ -77,21 +78,15 @@ pub async fn serve_web(ctx: CommandContext, mut args: ServeWebArgs) -> Result<i3
 	legal::require_consent(&ctx.paths, args.accept_server_license_terms)?;
 
 	let platform: crate::update_service::Platform = PreReqChecker::new().verify().await?;
-	let token_file = if !args.without_connection_token {
+	if !args.without_connection_token {
 		// Ensure there's a defined connection token, since if multiple server versions
 		// are excuted, they will need to have a single shared token.
-		let connection_token = args
-			.connection_token
-			.clone()
-			.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
-		let tf = ConnectionTokenFile::new(&connection_token)
+		let token_path = ctx.paths.root().join("serve-web-token");
+		let token = mint_connection_token(&token_path, args.connection_token.clone())
 			.map_err(CodeError::CouldNotCreateConnectionTokenFile)?;
-		args.connection_token = Some(connection_token);
-		args.connection_token_file = Some(tf.path().to_string_lossy().to_string());
-		Some(tf)
-	} else {
-		None
-	};
+		args.connection_token = Some(token);
+		args.connection_token_file = Some(token_path.to_string_lossy().to_string());
+	}
 
 	let cm = ConnectionManager::new(&ctx, platform, args.clone());
 	let key = get_server_key_half(&ctx.paths);
@@ -143,7 +138,6 @@ pub async fn serve_web(ctx: CommandContext, mut args: ServeWebArgs) -> Result<i3
 	};
 
 	r.map_err(CodeError::CouldNotListenOnInterface)?;
-	drop(token_file); // ensure it lives long enough
 
 	Ok(0)
 }
@@ -789,20 +783,29 @@ struct StartArgs {
 	opener: BarrierOpener<Result<StartData, String>>,
 }
 
-struct ConnectionTokenFile {
-	path: PathBuf,
-	_dir: TempDir, // implements Drop to delete the dir
-}
-
-impl ConnectionTokenFile {
-	fn new(connection_token: &str) -> std::io::Result<Self> {
-		let d = tempdir()?;
-		let path = d.path().join("connection-token");
-		std::fs::write(&path, connection_token)?;
-		Ok(ConnectionTokenFile { path, _dir: d })
+fn mint_connection_token(path: &Path, prefer_token: Option<String>) -> std::io::Result<String> {
+	#[cfg(not(windows))]
+	use std::os::unix::fs::OpenOptionsExt;
+
+	let mut f = fs::OpenOptions::new();
+	f.create(true);
+	f.write(true);
+	f.read(true);
+	#[cfg(not(windows))]
+	f.mode(0o600);
+	let mut f = f.open(path)?;
+
+	if prefer_token.is_none() {
+		let mut t = String::new();
+		f.read_to_string(&mut t)?;
+		let t = t.trim();
+		if !t.is_empty() {
+			return Ok(t.to_string());
+		}
 	}
 
-	fn path(&self) -> &Path {
-		&self.path
-	}
+	f.set_len(0)?;
+	let prefer_token = prefer_token.unwrap_or_else(|| uuid::Uuid::new_v4().to_string());
+	f.write_all(prefer_token.as_bytes())?;
+	Ok(prefer_token)
 }