From adcffbdce98c3e764e447f83ac0f5629e31e0555 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Mon, 28 Nov 2022 12:03:15 -0500 Subject: [PATCH] cli: use connection token for CLI connections (#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 https://github.com/microsoft/vscode-internalbacklog/issues/3287 --- cli/Cargo.lock | 2 ++ cli/Cargo.toml | 2 ++ cli/src/commands/tunnels.rs | 13 +++++++- cli/src/constants.rs | 8 ++++- cli/src/tunnels/code_server.rs | 1 - cli/src/tunnels/dev_tunnels.rs | 55 ++++++++++++++++++++++++++++++++-- 6 files changed, 75 insertions(+), 6 deletions(-) diff --git a/cli/Cargo.lock b/cli/Cargo.lock index 4b7ab3e3c23..2026626c5b3 100644 --- a/cli/Cargo.lock +++ b/cli/Cargo.lock @@ -229,6 +229,7 @@ version = "0.1.0" dependencies = [ "async-trait", "atty", + "base64", "chrono", "clap", "clap_lex", @@ -254,6 +255,7 @@ dependencies = [ "serde", "serde_bytes", "serde_json", + "sha2", "sysinfo", "tar", "tempfile", diff --git a/cli/Cargo.toml b/cli/Cargo.toml index c288e847514..a632e10e432 100644 --- a/cli/Cargo.toml +++ b/cli/Cargo.toml @@ -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" } diff --git a/cli/src/commands/tunnels.rs b/cli/src/commands/tunnels.rs index 01e4de805c7..5b0bc0d38fc 100644 --- a/cli/src/commands/tunnels.rs +++ b/cli/src/commands/tunnels.rs @@ -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>, ) -> Result { // 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 { diff --git a/cli/src/constants.rs b/cli/src/constants.rs index 1707d9c93d0..3596e152cda 100644 --- a/cli/src/constants.rs +++ b/cli/src/constants.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"); diff --git a/cli/src/tunnels/code_server.rs b/cli/src/tunnels/code_server.rs index d720edb73f3..74a8b00efbe 100644 --- a/cli/src/tunnels/code_server.rs +++ b/cli/src/tunnels/code_server.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())); diff --git a/cli/src/tunnels/dev_tunnels.rs b/cli/src/tunnels/dev_tunnels.rs index 911bf547e2f..aeee42cb911 100644 --- a/cli/src/tunnels/dev_tunnels.rs +++ b/cli/src/tunnels/dev_tunnels.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, use_random_name: bool, ) -> Result { - 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 { + debug!( + self.log, + "Updating tunnel protocol version tag to {}", PROTOCOL_VERSION_TAG + ); + let mut new_tags: Vec = 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 { @@ -690,6 +738,7 @@ impl DevTunnels { Ok(ActiveTunnel { name: tunnel_details.name.clone(), + id: tunnel_details.id.clone(), manager, }) }