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, }) }