Merge pull request #16 from devdiv-microsoft/connor4312/msrc/use-token-file

MSRC Fix for February 2024 Patch Tuesday

Author: aiday-mar

cli/src/commands/args.rs

@@ -201,6 +201,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,

cli/src/commands/serve_web.rs

@@ -5,8 +5,10 @@
 
 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::PathBuf;
+use std::path::{Path, PathBuf};
 use std::sync::{Arc, Mutex};
 use std::time::{Duration, Instant};
 
@@ -76,15 +78,14 @@ 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 {
 		// 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 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(token);
+		args.connection_token_file = Some(token_path.to_string_lossy().to_string());
 	}
 
 	let cm = ConnectionManager::new(&ctx, platform, args.clone());
@@ -704,8 +705,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 +782,30 @@ struct StartArgs {
 	release: Release,
 	opener: BarrierOpener<Result<StartData, String>>,
 }
+
+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());
+		}
+	}
+
+	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)
+}

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!(