<commit_before>/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/
use crate::commands::tunnels::ShutdownSignal;
use crate::constants::{CONTROL_PORT, PROTOCOL_VERSION, VSCODE_CLI_VERSION};
use crate::log;
use crate::self_update::SelfUpdate;
use crate::state::LauncherPaths;
use crate::update_service::{Platform, UpdateService};
use crate::util::errors::{
	wrap, AnyError, MismatchedLaunchModeError, NoAttachedServerError, ServerWriteError,
};
use crate::util::io::SilentCopyProgress;
use crate::util::sync::{new_barrier, Barrier};
use opentelemetry::trace::SpanKind;
use opentelemetry::KeyValue;
use serde::Serialize;
use std::convert::Infallible;
use std::env;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Instant;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
use tokio::pin;
use tokio::sync::{mpsc, Mutex};

use super::code_server::{
	AnyCodeServer, CodeServerArgs, ServerBuilder, ServerParamsRaw, SocketCodeServer,
};
use super::dev_tunnels::ActiveTunnel;
use super::paths::prune_stopped_servers;
use super::port_forwarder::{PortForwarding, PortForwardingProcessor};
use super::protocol::{
	CallServerHttpParams, CallServerHttpResult, ClientRequestMethod, EmptyResult, ErrorResponse,
	ForwardParams, ForwardResult, GetHostnameResponse, RefServerMessageParams, ResponseError,
	ServeParams, ServerLog, ServerMessageParams, ServerRequestMethod, SuccessResponse,
	ToClientRequest, ToServerRequest, UnforwardParams, UpdateParams, UpdateResult, VersionParams,
};
use super::server_bridge::{get_socket_rw_stream, FromServerMessage, ServerBridge};

type ServerBridgeList = Option<Vec<(u16, ServerBridge)>>;
type ServerBridgeListLock = Arc<Mutex<ServerBridgeList>>;

struct HandlerContext {
	/// Exit barrier for the socket.
	closer: Barrier<()>,
	/// Log handle for the server
	log: log::Logger,
	/// A loopback channel to talk to the TCP server task.
	server_tx: mpsc::Sender<ServerSignal>,
	/// A loopback channel to talk to the socket server task.
	socket_tx: mpsc::Sender<SocketSignal>,
	/// Configured launcher paths.
	launcher_paths: LauncherPaths,
	/// Connected VS Code Server
	code_server: Option<SocketCodeServer>,
	/// Potentially many "websocket" connections to client
	server_bridges: ServerBridgeListLock,
	// the cli arguments used to start the code server
	code_server_args: CodeServerArgs,
	/// counter for the number of bytes received from the socket
	rx_counter: Arc<AtomicUsize>,
	/// port forwarding functionality
	port_forwarding: PortForwarding,
	/// install platform for the VS Code server
	platform: Platform,
}

impl HandlerContext {
	async fn dispose(self) {
		let bridges: ServerBridgeList = {
			let mut lock = self.server_bridges.lock().await;
			let bridges = lock.take();
			*lock = None;
			bridges
		};

		if let Some(b) = bridges {
			for (_, bridge) in b {
				if let Err(e) = bridge.close().await {
					warning!(
						self.log,
						"Could not properly dispose of connection context: {}",
						e
					)
				} else {
					debug!(self.log, "Closed server bridge.");
				}
			}
		}

		info!(self.log, "Disposed of connection to running server.");
	}
}

enum ServerSignal {
	/// Signalled when the server has been updated and we want to respawn.
	/// We'd generally need to stop and then restart the launcher, but the
	/// program might be managed by a supervisor like systemd. Instead, we
	/// will stop the TCP listener and spawn the launcher again as a subprocess
	/// with the same arguments we used.
	Respawn,
}

struct CloseReason(String);

enum SocketSignal {
	/// Signals bytes to send to the socket.
	Send(Vec<u8>),
	/// Closes the socket (e.g. as a result of an error)
	CloseWith(CloseReason),
	/// Disposes ServerBridge corresponding to an ID
	CloseServerBridge(u16),
}

impl SocketSignal {
	fn from_message<T>(msg: &T) -> Self
	where
		T: Serialize + ?Sized,
	{
		SocketSignal::Send(rmp_serde::to_vec_named(msg).unwrap())
	}
}

impl FromServerMessage for SocketSignal {
	fn from_server_message(i: u16, body: &[u8]) -> Self {
		SocketSignal::from_message(&ToClientRequest {
			id: None,
			params: ClientRequestMethod::servermsg(RefServerMessageParams { i, body }),
		})
	}

	fn from_closed_server_bridge(i: u16) -> Self {
		SocketSignal::CloseServerBridge(i)
	}
}

pub struct ServerTermination {
	/// Whether the server should be respawned in a new binary (see ServerSignal.Respawn).
	pub respawn: bool,
	pub tunnel: ActiveTunnel,
}

fn print_listening(log: &log::Logger, tunnel_name: &str) {
	debug!(log, "VS Code Server is listening for incoming connections");

	let extension_name = "+ms-vscode.remote-server";
	let home_dir = dirs::home_dir().unwrap_or_else(|| PathBuf::from(""));
	let current_dir = env::current_dir().unwrap_or_else(|_| PathBuf::from(""));

	let dir = if home_dir == current_dir {
		PathBuf::from("")
	} else {
		current_dir
	};

	let mut addr = url::Url::parse("https://insiders.vscode.dev").unwrap();
	{
		let mut ps = addr.path_segments_mut().unwrap();
		ps.push(extension_name);
		ps.push(tunnel_name);
		for segment in &dir {
			let as_str = segment.to_string_lossy();
			if !(as_str.len() == 1 && as_str.starts_with(std::path::MAIN_SEPARATOR)) {<commit_msg>tunnel: Generate vanity URL (#164710)<commit_after>/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *  Licensed under the MIT License. See License.txt in the project root for license information.
 *--------------------------------------------------------------------------------------------*/
use crate::commands::tunnels::ShutdownSignal;
use crate::constants::{CONTROL_PORT, PROTOCOL_VERSION, VSCODE_CLI_VERSION};
use crate::log;
use crate::self_update::SelfUpdate;
use crate::state::LauncherPaths;
use crate::update_service::{Platform, UpdateService};
use crate::util::errors::{
	wrap, AnyError, MismatchedLaunchModeError, NoAttachedServerError, ServerWriteError,
};
use crate::util::io::SilentCopyProgress;
use crate::util::sync::{new_barrier, Barrier};
use opentelemetry::trace::SpanKind;
use opentelemetry::KeyValue;
use serde::Serialize;
use std::convert::Infallible;
use std::env;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::Instant;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt, BufReader};
use tokio::pin;
use tokio::sync::{mpsc, Mutex};

use super::code_server::{
	AnyCodeServer, CodeServerArgs, ServerBuilder, ServerParamsRaw, SocketCodeServer,
};
use super::dev_tunnels::ActiveTunnel;
use super::paths::prune_stopped_servers;
use super::port_forwarder::{PortForwarding, PortForwardingProcessor};
use super::protocol::{
	CallServerHttpParams, CallServerHttpResult, ClientRequestMethod, EmptyResult, ErrorResponse,
	ForwardParams, ForwardResult, GetHostnameResponse, RefServerMessageParams, ResponseError,
	ServeParams, ServerLog, ServerMessageParams, ServerRequestMethod, SuccessResponse,
	ToClientRequest, ToServerRequest, UnforwardParams, UpdateParams, UpdateResult, VersionParams,
};
use super::server_bridge::{get_socket_rw_stream, FromServerMessage, ServerBridge};

type ServerBridgeList = Option<Vec<(u16, ServerBridge)>>;
type ServerBridgeListLock = Arc<Mutex<ServerBridgeList>>;

struct HandlerContext {
	/// Exit barrier for the socket.
	closer: Barrier<()>,
	/// Log handle for the server
	log: log::Logger,
	/// A loopback channel to talk to the TCP server task.
	server_tx: mpsc::Sender<ServerSignal>,
	/// A loopback channel to talk to the socket server task.
	socket_tx: mpsc::Sender<SocketSignal>,
	/// Configured launcher paths.
	launcher_paths: LauncherPaths,
	/// Connected VS Code Server
	code_server: Option<SocketCodeServer>,
	/// Potentially many "websocket" connections to client
	server_bridges: ServerBridgeListLock,
	// the cli arguments used to start the code server
	code_server_args: CodeServerArgs,
	/// counter for the number of bytes received from the socket
	rx_counter: Arc<AtomicUsize>,
	/// port forwarding functionality
	port_forwarding: PortForwarding,
	/// install platform for the VS Code server
	platform: Platform,
}

impl HandlerContext {
	async fn dispose(self) {
		let bridges: ServerBridgeList = {
			let mut lock = self.server_bridges.lock().await;
			let bridges = lock.take();
			*lock = None;
			bridges
		};

		if let Some(b) = bridges {
			for (_, bridge) in b {
				if let Err(e) = bridge.close().await {
					warning!(
						self.log,
						"Could not properly dispose of connection context: {}",
						e
					)
				} else {
					debug!(self.log, "Closed server bridge.");
				}
			}
		}

		info!(self.log, "Disposed of connection to running server.");
	}
}

enum ServerSignal {
	/// Signalled when the server has been updated and we want to respawn.
	/// We'd generally need to stop and then restart the launcher, but the
	/// program might be managed by a supervisor like systemd. Instead, we
	/// will stop the TCP listener and spawn the launcher again as a subprocess
	/// with the same arguments we used.
	Respawn,
}

struct CloseReason(String);

enum SocketSignal {
	/// Signals bytes to send to the socket.
	Send(Vec<u8>),
	/// Closes the socket (e.g. as a result of an error)
	CloseWith(CloseReason),
	/// Disposes ServerBridge corresponding to an ID
	CloseServerBridge(u16),
}

impl SocketSignal {
	fn from_message<T>(msg: &T) -> Self
	where
		T: Serialize + ?Sized,
	{
		SocketSignal::Send(rmp_serde::to_vec_named(msg).unwrap())
	}
}

impl FromServerMessage for SocketSignal {
	fn from_server_message(i: u16, body: &[u8]) -> Self {
		SocketSignal::from_message(&ToClientRequest {
			id: None,
			params: ClientRequestMethod::servermsg(RefServerMessageParams { i, body }),
		})
	}

	fn from_closed_server_bridge(i: u16) -> Self {
		SocketSignal::CloseServerBridge(i)
	}
}

pub struct ServerTermination {
	/// Whether the server should be respawned in a new binary (see ServerSignal.Respawn).
	pub respawn: bool,
	pub tunnel: ActiveTunnel,
}

fn print_listening(log: &log::Logger, tunnel_name: &str) {
	debug!(log, "VS Code Server is listening for incoming connections");

	let home_dir = dirs::home_dir().unwrap_or_else(|| PathBuf::from(""));
	let current_dir = env::current_dir().unwrap_or_else(|_| PathBuf::from(""));

	let dir = if home_dir == current_dir {
		PathBuf::from("")
	} else {
		current_dir
	};

	let mut addr = url::Url::parse("https://insiders.vscode.dev").unwrap();
	{
		let mut ps = addr.path_segments_mut().unwrap();
		ps.push("tunnel");
		ps.push(tunnel_name);
		for segment in &dir {
			let as_str = segment.to_string_lossy();
			if !(as_str.len() == 1 && as_str.starts_with(std::path::MAIN_SEPARATOR)) {
{'fraction_non_alphanumeric': 0.152994555353902, 'fraction_numerical': 0.002540834845735027, 'mean_word_length': 4.255603242727706, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 31, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6794146', 'n_tokens_mistral': 3555, 'n_tokens_neox': 3384, 'n_words': 941}