cli: use connection token for CLI connections (microsoft#167426)

This uses a hash of the tunnel ID to create the connection token, which
should be sufficient to resolve the issues.

We also now publish the protocol version in the tunnel tags, since the
connection token must be supplied in the resolver, which is before we
start connecting to the tunnel.

See microsoft/vscode-internalbacklog#3287

Author: connor4312

cli/Cargo.lock

@@ -47,6 +47,8 @@ url = "2.3"
 async-trait = "0.1"
 log = "0.4"
 const_format = "0.2"
+sha2 = "0.10"
+base64 = "0.13"
 
 [build-dependencies]
 serde = { version = "1.0" }

cli/Cargo.toml

@@ -4,6 +4,7 @@
  *--------------------------------------------------------------------------------------------*/
 
 use async_trait::async_trait;
+use sha2::{Digest, Sha256};
 use std::fmt;
 use std::str::FromStr;
 use sysinfo::{Pid, SystemExt};
@@ -18,6 +19,7 @@ use super::{
 	CommandContext,
 };
 
+use crate::tunnels::dev_tunnels::ActiveTunnel;
 use crate::{
 	auth::Auth,
 	log::{self, Logger},
@@ -234,11 +236,18 @@ pub async fn serve(ctx: CommandContext, gateway_args: TunnelServeArgs) -> Result
 	serve_with_csa(paths, log, gateway_args, csa, None).await
 }
 
+fn get_connection_token(tunnel: &ActiveTunnel) -> String {
+	let mut hash = Sha256::new();
+	hash.update(tunnel.id.as_bytes());
+	let result = hash.finalize();
+	base64::encode_config(result, base64::URL_SAFE_NO_PAD)
+}
+
 async fn serve_with_csa(
 	paths: LauncherPaths,
 	log: Logger,
 	gateway_args: TunnelServeArgs,
-	csa: CodeServerArgs,
+	mut csa: CodeServerArgs,
 	shutdown_rx: Option<mpsc::UnboundedReceiver<ShutdownSignal>>,
 ) -> Result<i32, AnyError> {
 	// Intentionally read before starting the server. If the server updated and
@@ -256,6 +265,8 @@ async fn serve_with_csa(
 			.await
 	}?;
 
+	csa.connection_token = Some(get_connection_token(&tunnel));
+
 	let shutdown_tx = if let Some(tx) = shutdown_rx {
 		tx
 	} else {

cli/src/commands/tunnels.rs

@@ -17,7 +17,13 @@ pub const CONTROL_PORT: u16 = 31545;
 ///  1 - Initial protocol version
 ///  2 - Addition of `serve.compressed` property to control whether servermsg's
 ///      are compressed bidirectionally.
-pub const PROTOCOL_VERSION: u32 = 2;
+///  3 - The server's connection token is set to a SHA256 hash of the tunnel ID
+pub const PROTOCOL_VERSION: u32 = 3;
+
+/// Prefix for the tunnel tag that includes the version.
+pub const PROTOCOL_VERSION_TAG_PREFIX: &str = "protocolv";
+/// Tag for the current protocol version, which is included in dev tunnels.
+pub const PROTOCOL_VERSION_TAG: &str = concatcp!("protocolv", PROTOCOL_VERSION);
 
 pub const VSCODE_CLI_VERSION: Option<&'static str> = option_env!("VSCODE_CLI_VERSION");
 pub const VSCODE_CLI_AI_KEY: Option<&'static str> = option_env!("VSCODE_CLI_AI_KEY");

cli/src/constants.rs

@@ -534,7 +534,6 @@ impl<'a, Http: SimpleHttp + Send + Sync + Clone + 'static> ServerBuilder<'a, Htt
 
 		let mut cmd = self.get_base_command();
 		cmd.arg("--start-server")
-			.arg("--without-connection-token")
 			.arg("--enable-remote-auto-shutdown")
 			.arg(format!("--socket-path={}", socket.display()));
 

cli/src/tunnels/code_server.rs

@@ -3,7 +3,9 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 use crate::auth;
-use crate::constants::{CONTROL_PORT, TUNNEL_SERVICE_USER_AGENT};
+use crate::constants::{
+	CONTROL_PORT, PROTOCOL_VERSION_TAG, PROTOCOL_VERSION_TAG_PREFIX, TUNNEL_SERVICE_USER_AGENT,
+};
 use crate::state::{LauncherPaths, PersistedState};
 use crate::util::errors::{
 	wrap, AnyError, DevTunnelError, InvalidTunnelName, TunnelCreationFailed, WrappedError,
@@ -138,6 +140,8 @@ pub struct DevTunnels {
 pub struct ActiveTunnel {
 	/// Name of the tunnel
 	pub name: String,
+	/// Underlying dev tunnels ID
+	pub id: String,
 	manager: ActiveTunnelManager,
 }
 
@@ -378,7 +382,7 @@ impl DevTunnels {
 		preferred_name: Option<String>,
 		use_random_name: bool,
 	) -> Result<ActiveTunnel, AnyError> {
-		let (tunnel, persisted) = match self.launcher_tunnel.load() {
+		let (mut tunnel, persisted) = match self.launcher_tunnel.load() {
 			Some(mut persisted) => {
 				if let Some(name) = preferred_name {
 					if persisted.name.ne(&name) {
@@ -404,6 +408,12 @@ impl DevTunnels {
 			}
 		};
 
+		if !tunnel.tags.iter().any(|t| t == PROTOCOL_VERSION_TAG) {
+			tunnel = self
+				.update_protocol_version_tag(tunnel, &HOST_TUNNEL_REQUEST_OPTIONS)
+				.await?;
+		}
+
 		let locator = TunnelLocator::try_from(&tunnel).unwrap();
 		let host_token = get_host_token_from_tunnel(&tunnel);
 
@@ -462,7 +472,11 @@ impl DevTunnels {
 		let mut tried_recycle = false;
 
 		let new_tunnel = Tunnel {
-			tags: vec![name.to_string(), VSCODE_CLI_TUNNEL_TAG.to_string()],
+			tags: vec![
+				name.to_string(),
+				PROTOCOL_VERSION_TAG.to_string(),
+				VSCODE_CLI_TUNNEL_TAG.to_string(),
+			],
 			..Default::default()
 		};
 
@@ -507,6 +521,40 @@ impl DevTunnels {
 		}
 	}
 
+	/// Ensures the tunnel contains a tag for the current PROTCOL_VERSION, and no
+	/// other version tags.
+	async fn update_protocol_version_tag(
+		&self,
+		tunnel: Tunnel,
+		options: &TunnelRequestOptions,
+	) -> Result<Tunnel, AnyError> {
+		debug!(
+			self.log,
+			"Updating tunnel protocol version tag to {}", PROTOCOL_VERSION_TAG
+		);
+		let mut new_tags: Vec<String> = tunnel
+			.tags
+			.into_iter()
+			.filter(|t| !t.starts_with(PROTOCOL_VERSION_TAG_PREFIX))
+			.collect();
+		new_tags.push(PROTOCOL_VERSION_TAG.to_string());
+
+		let tunnel_update = Tunnel {
+			tags: new_tags,
+			tunnel_id: tunnel.tunnel_id.clone(),
+			cluster_id: tunnel.cluster_id.clone(),
+			..Default::default()
+		};
+
+		let result = spanf!(
+			self.log,
+			self.log.span("dev-tunnel.protocol-tag-update"),
+			self.client.update_tunnel(&tunnel_update, options)
+		);
+
+		result.map_err(|e| wrap(e, "tunnel tag update failed").into())
+	}
+
 	/// Tries to delete an unused tunnel, and then creates a tunnel with the
 	/// given `new_name`.
 	async fn try_recycle_tunnel(&mut self) -> Result<bool, AnyError> {
@@ -690,6 +738,7 @@ impl DevTunnels {
 
 		Ok(ActiveTunnel {
 			name: tunnel_details.name.clone(),
+			id: tunnel_details.id.clone(),
 			manager,
 		})
 	}