2512 lines · 84.2 KB · Rust
RawDownload
use crate::{    json_log::LogRecord,    protocol::{        message_len_from_buffer, read_message_with_len, write_message, MessageId, MESSAGE_LEN_SIZE,    },    proxy::ProxyLaunchError,};use anyhow::{anyhow, Context as _, Result};use async_trait::async_trait;use collections::HashMap;use futures::{    channel::{        mpsc::{self, Sender, UnboundedReceiver, UnboundedSender},        oneshot,    },    future::{BoxFuture, Shared},    select, select_biased, AsyncReadExt as _, Future, FutureExt as _, StreamExt as _,};use gpui::{    AppContext, AsyncAppContext, BorrowAppContext, Context, EventEmitter, Global, Model,    ModelContext, SemanticVersion, Task, WeakModel,};use itertools::Itertools;use parking_lot::Mutex;use paths;use release_channel::{AppCommitSha, AppVersion, ReleaseChannel};use rpc::{    proto::{self, build_typed_envelope, Envelope, EnvelopedMessage, PeerId, RequestMessage},    AnyProtoClient, EntityMessageSubscriber, ErrorExt, ProtoClient, ProtoMessageHandlerSet,    RpcError,};use smol::{    fs,    process::{self, Child, Stdio},};use std::{    any::TypeId,    collections::VecDeque,    fmt, iter,    ops::ControlFlow,    path::{Path, PathBuf},    sync::{        atomic::{AtomicU32, AtomicU64, Ordering::SeqCst},        Arc, Weak,    },    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, 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 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()            ),+        )    }};} 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();         // 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", "-I", "-i", "-J", "-L", "-l", "-m", "-o", "-P", "-p", "-R",            "-w",        ];         let mut tokens = shlex::split(input)            .ok_or_else(|| anyhow!("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;            }            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;            if let Some((u, rest)) = input.split_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");        };         Ok(Self {            host: hostname.to_string(),            username: username.clone(),            port,            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 {            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) -> Option<&Vec<String>> {        self.args.as_ref()    }     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,        cx: &mut AsyncAppContext,    ) -> oneshot::Receiver<Result<String>>;    fn get_download_params(        &self,        platform: SshPlatform,        release_channel: ReleaseChannel,        version: Option<SemanticVersion>,        cx: &mut AsyncAppContext,    ) -> Task<Result<Option<(String, String)>>>;     fn download_server_binary_locally(        &self,        platform: SshPlatform,        release_channel: ReleaseChannel,        version: Option<SemanticVersion>,        cx: &mut AsyncAppContext,    ) -> Task<Result<PathBuf>>;    fn set_status(&self, status: Option<&str>, cx: &mut AsyncAppContext);} 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.    // You need to do it like this: $ ssh host "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(" ");        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?;        if output.status.success() {            Ok(String::from_utf8_lossy(&output.stdout).to_string())        } else {            Err(anyhow!(                "failed to run command: {}",                String::from_utf8_lossy(&output.stderr)            ))        }    }     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,        }    }     fn heartbeat_missed(self) -> Self {        match self {            Self::Connected {                ssh_connection,                delegate,                multiplex_task,                heartbeat_task,            } => Self::HeartbeatMissed {                missed_heartbeats: 1,                ssh_connection,                delegate,                multiplex_task,                heartbeat_task,            },            Self::HeartbeatMissed {                missed_heartbeats,                ssh_connection,                delegate,                multiplex_task,                heartbeat_task,            } => Self::HeartbeatMissed {                missed_heartbeats: missed_heartbeats + 1,                ssh_connection,                delegate,                multiplex_task,                heartbeat_task,            },            _ => self,        }    }} /// The state of the ssh connection.#[derive(Clone, Copy, Debug, PartialEq, Eq)]pub enum ConnectionState {    Connecting,    Connected,    HeartbeatMissed,    Reconnecting,    Disconnected,} impl From<&State> for ConnectionState {    fn from(value: &State) -> Self {        match value {            State::Connecting => Self::Connecting,            State::Connected { .. } => Self::Connected,            State::Reconnecting | State::ReconnectFailed { .. } => Self::Reconnecting,            State::HeartbeatMissed { .. } => Self::HeartbeatMissed,            State::ReconnectExhausted => Self::Disconnected,            State::ServerNotRunning => Self::Disconnected,        }    }} pub struct SshRemoteClient {    client: Arc<ChannelClient>,    unique_identifier: String,    connection_options: SshConnectionOptions,    state: Arc<Mutex<Option<State>>>,} #[derive(Debug)]pub enum SshRemoteEvent {    Disconnected,} impl EventEmitter<SshRemoteEvent> for SshRemoteClient {} // Identifies the socket on the remote server so that reconnects// can re-join the same project.pub enum ConnectionIdentifier {    Setup(u64),    Workspace(i64),} static NEXT_ID: AtomicU64 = AtomicU64::new(1); impl ConnectionIdentifier {    pub fn setup() -> Self {        Self::Setup(NEXT_ID.fetch_add(1, SeqCst))    }    // This string gets used in a socket name, and so must be relatively short.    // The total length of:    //   /home/{username}/.local/share/zed/server_state/{name}/stdout.sock    // Must be less than about 100 characters

Lines 1–500 of 2512.