use crate::{    json_log::LogRecord,    protocol::{        MESSAGE_LEN_SIZE, MessageId, message_len_from_buffer, read_message_with_len, write_message,    },    proxy::ProxyLaunchError,};use anyhow::{Context as _, Result, anyhow};use async_trait::async_trait;use collections::HashMap;use futures::{    AsyncReadExt as _, Future, FutureExt as _, StreamExt as _,    channel::{        mpsc::{self, Sender, UnboundedReceiver, UnboundedSender},        oneshot,    },    future::{BoxFuture, Shared},    select, select_biased,};use gpui::{    App, AppContext as _, AsyncApp, BackgroundExecutor, BorrowAppContext, Context, Entity,    EventEmitter, Global, SemanticVersion, Task, WeakEntity,};use itertools::Itertools;use parking_lot::Mutex; use release_channel::{AppCommitSha, AppVersion, ReleaseChannel};use rpc::{    AnyProtoClient, EntityMessageSubscriber, ErrorExt, ProtoClient, ProtoMessageHandlerSet,    RpcError,    proto::{self, Envelope, EnvelopedMessage, PeerId, RequestMessage, build_typed_envelope},};use schemars::JsonSchema;use serde::{Deserialize, Serialize};use smol::{    fs,    process::{self, Child, Stdio},};use std::{    any::TypeId,    collections::VecDeque,    fmt, iter,    ops::ControlFlow,    path::{Path, PathBuf},    sync::{        Arc, Weak,        atomic::{AtomicU32, AtomicU64, Ordering::SeqCst},    },    time::{Duration, Instant},};use tempfile::TempDir;use util::ResultExt; #[derive(    Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, serde::Serialize, serde::Deserialize,)]pub struct SshProjectId(pub u64); #[derive(Clone)]pub struct SshSocket {    connection_options: SshConnectionOptions,    socket_path: PathBuf,} #[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize, JsonSchema)]pub struct SshPortForwardOption {    #[serde(skip_serializing_if = "Option::is_none")]    pub local_host: Option<String>,    pub local_port: u16,    #[serde(skip_serializing_if = "Option::is_none")]    pub remote_host: Option<String>,    pub remote_port: u16,} #[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]pub struct SshConnectionOptions {    pub host: String,    pub username: Option<String>,    pub port: Option<u16>,    pub password: Option<String>,    pub args: Option<Vec<String>>,    pub port_forwards: Option<Vec<SshPortForwardOption>>,     pub nickname: Option<String>,    pub upload_binary_over_ssh: bool,} #[macro_export]macro_rules! shell_script {    ($fmt:expr, $($name:ident = $arg:expr),+ $(,)?) => {{        format!(            $fmt,            $(                $name = shlex::try_quote($arg).unwrap()            ),+        )    }};} fn parse_port_number(port_str: &str) -> Result<u16> {    port_str        .parse()        .with_context(|| format!("parsing port number: {port_str}"))} fn parse_port_forward_spec(spec: &str) -> Result<SshPortForwardOption> {    let parts: Vec<&str> = spec.split(':').collect();     match parts.len() {        4 => {            let local_port = parse_port_number(parts[1])?;            let remote_port = parse_port_number(parts[3])?;             Ok(SshPortForwardOption {                local_host: Some(parts[0].to_string()),                local_port,                remote_host: Some(parts[2].to_string()),                remote_port,            })        }        3 => {            let local_port = parse_port_number(parts[0])?;            let remote_port = parse_port_number(parts[2])?;             Ok(SshPortForwardOption {                local_host: None,                local_port,                remote_host: Some(parts[1].to_string()),                remote_port,            })        }        _ => anyhow::bail!("Invalid port forward format"),    }} impl SshConnectionOptions {    pub fn parse_command_line(input: &str) -> Result<Self> {        let input = input.trim_start_matches("ssh ");        let mut hostname: Option<String> = None;        let mut username: Option<String> = None;        let mut port: Option<u16> = None;        let mut args = Vec::new();        let mut port_forwards: Vec<SshPortForwardOption> = Vec::new();         // disallowed: -E, -e, -F, -f, -G, -g, -M, -N, -n, -O, -q, -S, -s, -T, -t, -V, -v, -W        const ALLOWED_OPTS: &[&str] = &[            "-4", "-6", "-A", "-a", "-C", "-K", "-k", "-X", "-x", "-Y", "-y",        ];        const ALLOWED_ARGS: &[&str] = &[            "-B", "-b", "-c", "-D", "-F", "-I", "-i", "-J", "-l", "-m", "-o", "-P", "-p", "-R",            "-w",        ];         let mut tokens = shlex::split(input).context("invalid input")?.into_iter();         'outer: while let Some(arg) = tokens.next() {            if ALLOWED_OPTS.contains(&(&arg as &str)) {                args.push(arg.to_string());                continue;            }            if arg == "-p" {                port = tokens.next().and_then(|arg| arg.parse().ok());                continue;            } else if let Some(p) = arg.strip_prefix("-p") {                port = p.parse().ok();                continue;            }            if arg == "-l" {                username = tokens.next();                continue;            } else if let Some(l) = arg.strip_prefix("-l") {                username = Some(l.to_string());                continue;            }            if arg == "-L" || arg.starts_with("-L") {                let forward_spec = if arg == "-L" {                    tokens.next()                } else {                    Some(arg.strip_prefix("-L").unwrap().to_string())                };                 if let Some(spec) = forward_spec {                    port_forwards.push(parse_port_forward_spec(&spec)?);                } else {                    anyhow::bail!("Missing port forward format");                }            }             for a in ALLOWED_ARGS {                if arg == *a {                    args.push(arg);                    if let Some(next) = tokens.next() {                        args.push(next);                    }                    continue 'outer;                } else if arg.starts_with(a) {                    args.push(arg);                    continue 'outer;                }            }            if arg.starts_with("-") || hostname.is_some() {                anyhow::bail!("unsupported argument: {:?}", arg);            }            let mut input = &arg as &str;            // Destination might be: username1@username2@ip2@ip1            if let Some((u, rest)) = input.rsplit_once('@') {                input = rest;                username = Some(u.to_string());            }            if let Some((rest, p)) = input.split_once(':') {                input = rest;                port = p.parse().ok()            }            hostname = Some(input.to_string())        }         let Some(hostname) = hostname else {            anyhow::bail!("missing hostname");        };         let port_forwards = match port_forwards.len() {            0 => None,            _ => Some(port_forwards),        };         Ok(Self {            host: hostname.to_string(),            username: username.clone(),            port,            port_forwards,            args: Some(args),            password: None,            nickname: None,            upload_binary_over_ssh: false,        })    }     pub fn ssh_url(&self) -> String {        let mut result = String::from("ssh://");        if let Some(username) = &self.username {            // Username might be: username1@username2@ip2            let username = urlencoding::encode(username);            result.push_str(&username);            result.push('@');        }        result.push_str(&self.host);        if let Some(port) = self.port {            result.push(':');            result.push_str(&port.to_string());        }        result    }     pub fn additional_args(&self) -> Vec<String> {        let mut args = self.args.iter().flatten().cloned().collect::<Vec<String>>();         if let Some(forwards) = &self.port_forwards {            args.extend(forwards.iter().map(|pf| {                let local_host = match &pf.local_host {                    Some(host) => host,                    None => "localhost",                };                let remote_host = match &pf.remote_host {                    Some(host) => host,                    None => "localhost",                };                 format!(                    "-L{}:{}:{}:{}",                    local_host, pf.local_port, remote_host, pf.remote_port                )            }));        }         args    }     fn scp_url(&self) -> String {        if let Some(username) = &self.username {            format!("{}@{}", username, self.host)        } else {            self.host.clone()        }    }     pub fn connection_string(&self) -> String {        let host = if let Some(username) = &self.username {            format!("{}@{}", username, self.host)        } else {            self.host.clone()        };        if let Some(port) = &self.port {            format!("{}:{}", host, port)        } else {            host        }    }} #[derive(Copy, Clone, Debug)]pub struct SshPlatform {    pub os: &'static str,    pub arch: &'static str,} impl SshPlatform {    pub fn triple(&self) -> Option<String> {        Some(format!(            "{}-{}",            self.arch,            match self.os {                "linux" => "unknown-linux-gnu",                "macos" => "apple-darwin",                _ => return None,            }        ))    }} pub trait SshClientDelegate: Send + Sync {    fn ask_password(&self, prompt: String, tx: oneshot::Sender<String>, cx: &mut AsyncApp);    fn get_download_params(        &self,        platform: SshPlatform,        release_channel: ReleaseChannel,        version: Option<SemanticVersion>,        cx: &mut AsyncApp,    ) -> Task<Result<Option<(String, String)>>>;     fn download_server_binary_locally(        &self,        platform: SshPlatform,        release_channel: ReleaseChannel,        version: Option<SemanticVersion>,        cx: &mut AsyncApp,    ) -> Task<Result<PathBuf>>;    fn set_status(&self, status: Option<&str>, cx: &mut AsyncApp);} impl SshSocket {    // :WARNING: ssh unquotes arguments when executing on the remote :WARNING:    // e.g. $ ssh host sh -c 'ls -l' is equivalent to $ ssh host sh -c ls -l    // and passes -l as an argument to sh, not to ls.    // Furthermore, some setups (e.g. Coder) will change directory when SSH'ing    // into a machine. You must use `cd` to get back to $HOME.    // You need to do it like this: $ ssh host "cd; sh -c 'ls -l /tmp'"    fn ssh_command(&self, program: &str, args: &[&str]) -> process::Command {        let mut command = util::command::new_smol_command("ssh");        let to_run = iter::once(&program)            .chain(args.iter())            .map(|token| {                // We're trying to work with: sh, bash, zsh, fish, tcsh, ...?                debug_assert!(                    !token.contains('\n'),                    "multiline arguments do not work in all shells"                );                shlex::try_quote(token).unwrap()            })            .join(" ");        let to_run = format!("cd; {to_run}");        log::debug!("ssh {} {:?}", self.connection_options.ssh_url(), to_run);        self.ssh_options(&mut command)            .arg(self.connection_options.ssh_url())            .arg(to_run);        command    }     async fn run_command(&self, program: &str, args: &[&str]) -> Result<String> {        let output = self.ssh_command(program, args).output().await?;        anyhow::ensure!(            output.status.success(),            "failed to run command: {}",            String::from_utf8_lossy(&output.stderr)        );        Ok(String::from_utf8_lossy(&output.stdout).to_string())    }     fn ssh_options<'a>(&self, command: &'a mut process::Command) -> &'a mut process::Command {        command            .stdin(Stdio::piped())            .stdout(Stdio::piped())            .stderr(Stdio::piped())            .args(["-o", "ControlMaster=no", "-o"])            .arg(format!("ControlPath={}", self.socket_path.display()))    }     fn ssh_args(&self) -> Vec<String> {        vec![            "-o".to_string(),            "ControlMaster=no".to_string(),            "-o".to_string(),            format!("ControlPath={}", self.socket_path.display()),            self.connection_options.ssh_url(),        ]    }} const MAX_MISSED_HEARTBEATS: usize = 5;const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5); const MAX_RECONNECT_ATTEMPTS: usize = 3; enum State {    Connecting,    Connected {        ssh_connection: Arc<dyn RemoteConnection>,        delegate: Arc<dyn SshClientDelegate>,         multiplex_task: Task<Result<()>>,        heartbeat_task: Task<Result<()>>,    },    HeartbeatMissed {        missed_heartbeats: usize,         ssh_connection: Arc<dyn RemoteConnection>,        delegate: Arc<dyn SshClientDelegate>,         multiplex_task: Task<Result<()>>,        heartbeat_task: Task<Result<()>>,    },    Reconnecting,    ReconnectFailed {        ssh_connection: Arc<dyn RemoteConnection>,        delegate: Arc<dyn SshClientDelegate>,         error: anyhow::Error,        attempts: usize,    },    ReconnectExhausted,    ServerNotRunning,} impl fmt::Display for State {    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {        match self {            Self::Connecting => write!(f, "connecting"),            Self::Connected { .. } => write!(f, "connected"),            Self::Reconnecting => write!(f, "reconnecting"),            Self::ReconnectFailed { .. } => write!(f, "reconnect failed"),            Self::ReconnectExhausted => write!(f, "reconnect exhausted"),            Self::HeartbeatMissed { .. } => write!(f, "heartbeat missed"),            Self::ServerNotRunning { .. } => write!(f, "server not running"),        }    }} impl State {    fn ssh_connection(&self) -> Option<&dyn RemoteConnection> {        match self {            Self::Connected { ssh_connection, .. } => Some(ssh_connection.as_ref()),            Self::HeartbeatMissed { ssh_connection, .. } => Some(ssh_connection.as_ref()),            Self::ReconnectFailed { ssh_connection, .. } => Some(ssh_connection.as_ref()),            _ => None,        }    }     fn can_reconnect(&self) -> bool {        match self {            Self::Connected { .. }            | Self::HeartbeatMissed { .. }            | Self::ReconnectFailed { .. } => true,            State::Connecting            | State::Reconnecting            | State::ReconnectExhausted            | State::ServerNotRunning => false,        }    }     fn is_reconnect_failed(&self) -> bool {        matches!(self, Self::ReconnectFailed { .. })    }     fn is_reconnect_exhausted(&self) -> bool {        matches!(self, Self::ReconnectExhausted { .. })    }     fn is_server_not_running(&self) -> bool {        matches!(self, Self::ServerNotRunning)    }     fn is_reconnecting(&self) -> bool {        matches!(self, Self::Reconnecting { .. })    }     fn heartbeat_recovered(self) -> Self {        match self {            Self::HeartbeatMissed {                ssh_connection,                delegate,                multiplex_task,                heartbeat_task,                ..            } => Self::Connected {                ssh_connection,                delegate,                multiplex_task,                heartbeat_task,            },            _ => self,

Lines 1–500 of 2585. Load more