1658 lines · 56.8 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 collections::HashMap;use futures::{    channel::{        mpsc::{self, Sender, UnboundedReceiver, UnboundedSender},        oneshot,    },    future::BoxFuture,    select_biased, AsyncReadExt as _, AsyncWriteExt as _, Future, FutureExt as _, SinkExt,    StreamExt as _,};use gpui::{    AppContext, AsyncAppContext, Context, EventEmitter, Model, ModelContext, SemanticVersion, Task,    WeakModel,};use parking_lot::Mutex;use rpc::{    proto::{self, build_typed_envelope, Envelope, EnvelopedMessage, PeerId, RequestMessage},    AnyProtoClient, EntityMessageSubscriber, ProtoClient, ProtoMessageHandlerSet, RpcError,};use smol::{    fs,    process::{self, Child, Stdio},};use std::{    any::TypeId,    ffi::OsStr,    fmt,    ops::ControlFlow,    path::{Path, PathBuf},    sync::{        atomic::{AtomicU32, Ordering::SeqCst},        Arc,    },    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)]pub struct SshConnectionOptions {    pub host: String,    pub username: Option<String>,    pub port: Option<u16>,    pub password: Option<String>,    pub args: Option<Vec<String>>,} 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,            password: None,            args: Some(args),        })    }     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        }    }     // Uniquely identifies dev server projects on a remote host. Needs to be    // stable for the same dev server project.    pub fn dev_server_identifier(&self) -> String {        let mut identifier = format!("dev-server-{:?}", self.host);        if let Some(username) = self.username.as_ref() {            identifier.push('-');            identifier.push_str(&username);        }        identifier    }} #[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 remote_server_binary_path(        &self,        platform: SshPlatform,        cx: &mut AsyncAppContext,    ) -> Result<PathBuf>;    fn get_server_binary(        &self,        platform: SshPlatform,        cx: &mut AsyncAppContext,    ) -> oneshot::Receiver<Result<(PathBuf, SemanticVersion)>>;    fn set_status(&self, status: Option<&str>, cx: &mut AsyncAppContext);    fn set_error(&self, error_message: String, cx: &mut AsyncAppContext);} impl SshSocket {    fn ssh_command<S: AsRef<OsStr>>(&self, program: S) -> process::Command {        let mut command = process::Command::new("ssh");        self.ssh_options(&mut command)            .arg(self.connection_options.ssh_url())            .arg(program);        command    }     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(),        ]    }} async fn run_cmd(command: &mut process::Command) -> Result<String> {    let output = command.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)        ))    }} struct ChannelForwarder {    quit_tx: UnboundedSender<()>,    forwarding_task: Task<(UnboundedSender<Envelope>, UnboundedReceiver<Envelope>)>,} impl ChannelForwarder {    fn new(        mut incoming_tx: UnboundedSender<Envelope>,        mut outgoing_rx: UnboundedReceiver<Envelope>,        cx: &AsyncAppContext,    ) -> (Self, UnboundedSender<Envelope>, UnboundedReceiver<Envelope>) {        let (quit_tx, mut quit_rx) = mpsc::unbounded::<()>();         let (proxy_incoming_tx, mut proxy_incoming_rx) = mpsc::unbounded::<Envelope>();        let (mut proxy_outgoing_tx, proxy_outgoing_rx) = mpsc::unbounded::<Envelope>();         let forwarding_task = cx.background_executor().spawn(async move {            loop {                select_biased! {                    _ = quit_rx.next().fuse() => {                        break;                    },                    incoming_envelope = proxy_incoming_rx.next().fuse() => {                        if let Some(envelope) = incoming_envelope {                            if incoming_tx.send(envelope).await.is_err() {                                break;                            }                        } else {                            break;                        }                    }                    outgoing_envelope = outgoing_rx.next().fuse() => {                        if let Some(envelope) = outgoing_envelope {                            if proxy_outgoing_tx.send(envelope).await.is_err() {                                break;                            }                        } else {                            break;                        }                    }                }            }             (incoming_tx, outgoing_rx)        });         (            Self {                forwarding_task,                quit_tx,            },            proxy_incoming_tx,            proxy_outgoing_rx,        )    }     async fn into_channels(mut self) -> (UnboundedSender<Envelope>, UnboundedReceiver<Envelope>) {        let _ = self.quit_tx.send(()).await;        self.forwarding_task.await    }} 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: SshRemoteConnection,        delegate: Arc<dyn SshClientDelegate>,        forwarder: ChannelForwarder,         multiplex_task: Task<Result<()>>,        heartbeat_task: Task<Result<()>>,    },    HeartbeatMissed {        missed_heartbeats: usize,         ssh_connection: SshRemoteConnection,        delegate: Arc<dyn SshClientDelegate>,        forwarder: ChannelForwarder,         multiplex_task: Task<Result<()>>,        heartbeat_task: Task<Result<()>>,    },    Reconnecting,    ReconnectFailed {        ssh_connection: SshRemoteConnection,        delegate: Arc<dyn SshClientDelegate>,        forwarder: ChannelForwarder,         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<&SshRemoteConnection> {        match self {            Self::Connected { ssh_connection, .. } => Some(ssh_connection),            Self::HeartbeatMissed { ssh_connection, .. } => Some(ssh_connection),            Self::ReconnectFailed { ssh_connection, .. } => Some(ssh_connection),            _ => 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_reconnecting(&self) -> bool {        matches!(self, Self::Reconnecting { .. })    }     fn heartbeat_recovered(self) -> Self {        match self {            Self::HeartbeatMissed {                ssh_connection,                delegate,                forwarder,                multiplex_task,                heartbeat_task,                ..            } => Self::Connected {                ssh_connection,                delegate,                forwarder,                multiplex_task,                heartbeat_task,            },            _ => self,        }    }     fn heartbeat_missed(self) -> Self {        match self {            Self::Connected {                ssh_connection,                delegate,                forwarder,                multiplex_task,                heartbeat_task,            } => Self::HeartbeatMissed {                missed_heartbeats: 1,                ssh_connection,                delegate,                forwarder,                multiplex_task,                heartbeat_task,            },            Self::HeartbeatMissed {                missed_heartbeats,                ssh_connection,                delegate,                forwarder,                multiplex_task,                heartbeat_task,            } => Self::HeartbeatMissed {                missed_heartbeats: missed_heartbeats + 1,                ssh_connection,                delegate,                forwarder,                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,

Lines 1–500 of 1658.