1692 lines · 58.0 KB · Rust
RawDownload
#[cfg(any(test, feature = "test-support"))]use crate::transport::mock::ConnectGuard;use crate::{    SshConnectionOptions,    protocol::MessageId,    proxy::ProxyLaunchError,    transport::{        docker::{DockerConnectionOptions, DockerExecConnection},        ssh::SshRemoteConnection,        wsl::{WslConnectionOptions, WslRemoteConnection},    },};use anyhow::{Context as _, Result, anyhow};use askpass::EncryptedPassword;use async_trait::async_trait;use collections::HashMap;use futures::{    Future, FutureExt as _, StreamExt as _,    channel::{        mpsc::{self, Sender, UnboundedReceiver, UnboundedSender},        oneshot,    },    future::{BoxFuture, Shared, WeakShared},    select, select_biased,};use gpui::{    App, AppContext as _, AsyncApp, BackgroundExecutor, BorrowAppContext, Context, Entity,    EventEmitter, FutureExt, Global, Task, WeakEntity,};use parking_lot::Mutex; use release_channel::ReleaseChannel;use rpc::{    AnyProtoClient, ErrorExt, ProtoClient, ProtoMessageHandlerSet, RpcError,    proto::{self, Envelope, EnvelopedMessage, PeerId, RequestMessage, build_typed_envelope},};use semver::Version;use std::{    collections::VecDeque,    fmt,    ops::ControlFlow,    path::PathBuf,    sync::{        Arc, Weak,        atomic::{AtomicU32, AtomicU64, Ordering::SeqCst},    },    time::{Duration, Instant},};use util::{    ResultExt,    paths::{PathStyle, RemotePathBuf},}; #[derive(Copy, Clone, Debug, PartialEq, Eq)]pub enum RemoteOs {    Linux,    MacOs,    Windows,} impl RemoteOs {    pub fn as_str(&self) -> &'static str {        match self {            RemoteOs::Linux => "linux",            RemoteOs::MacOs => "macos",            RemoteOs::Windows => "windows",        }    }     pub fn is_windows(&self) -> bool {        matches!(self, RemoteOs::Windows)    }} impl std::fmt::Display for RemoteOs {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        f.write_str(self.as_str())    }} #[derive(Copy, Clone, Debug, PartialEq, Eq)]pub enum RemoteArch {    X86_64,    Aarch64,} impl RemoteArch {    pub fn as_str(&self) -> &'static str {        match self {            RemoteArch::X86_64 => "x86_64",            RemoteArch::Aarch64 => "aarch64",        }    }} impl std::fmt::Display for RemoteArch {    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {        f.write_str(self.as_str())    }} #[derive(Copy, Clone, Debug)]pub struct RemotePlatform {    pub os: RemoteOs,    pub arch: RemoteArch,} #[derive(Clone, Debug)]pub struct CommandTemplate {    pub program: String,    pub args: Vec<String>,    pub env: HashMap<String, String>,} /// Whether a command should be run with TTY allocation for interactive use.#[derive(Clone, Copy, Debug, PartialEq, Eq)]pub enum Interactive {    /// Allocate a pseudo-TTY for interactive terminal use.    Yes,    /// Do not allocate a TTY - for commands that communicate via piped stdio.    No,} pub trait RemoteClientDelegate: Send + Sync {    fn ask_password(        &self,        prompt: String,        tx: oneshot::Sender<EncryptedPassword>,        cx: &mut AsyncApp,    );    fn get_download_url(        &self,        platform: RemotePlatform,        release_channel: ReleaseChannel,        version: Option<Version>,        cx: &mut AsyncApp,    ) -> Task<Result<Option<String>>>;    fn download_server_binary_locally(        &self,        platform: RemotePlatform,        release_channel: ReleaseChannel,        version: Option<Version>,        cx: &mut AsyncApp,    ) -> Task<Result<PathBuf>>;    fn set_status(&self, status: Option<&str>, cx: &mut AsyncApp);} const MAX_MISSED_HEARTBEATS: usize = 5;const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5);const INITIAL_CONNECTION_TIMEOUT: Duration =    Duration::from_secs(if cfg!(debug_assertions) { 5 } else { 60 }); pub const MAX_RECONNECT_ATTEMPTS: usize = 3; enum State {    Connecting,    Connected {        remote_connection: Arc<dyn RemoteConnection>,        delegate: Arc<dyn RemoteClientDelegate>,         multiplex_task: Task<Result<()>>,        heartbeat_task: Task<Result<()>>,    },    HeartbeatMissed {        missed_heartbeats: usize,         remote_connection: Arc<dyn RemoteConnection>,        delegate: Arc<dyn RemoteClientDelegate>,         multiplex_task: Task<Result<()>>,        heartbeat_task: Task<Result<()>>,    },    Reconnecting,    ReconnectFailed {        remote_connection: Arc<dyn RemoteConnection>,        delegate: Arc<dyn RemoteClientDelegate>,         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 remote_connection(&self) -> Option<Arc<dyn RemoteConnection>> {        match self {            Self::Connected {                remote_connection, ..            } => Some(remote_connection.clone()),            Self::HeartbeatMissed {                remote_connection, ..            } => Some(remote_connection.clone()),            Self::ReconnectFailed {                remote_connection, ..            } => Some(remote_connection.clone()),            _ => 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 {                remote_connection,                delegate,                multiplex_task,                heartbeat_task,                ..            } => Self::Connected {                remote_connection,                delegate,                multiplex_task,                heartbeat_task,            },            _ => self,        }    }     fn heartbeat_missed(self) -> Self {        match self {            Self::Connected {                remote_connection,                delegate,                multiplex_task,                heartbeat_task,            } => Self::HeartbeatMissed {                missed_heartbeats: 1,                remote_connection,                delegate,                multiplex_task,                heartbeat_task,            },            Self::HeartbeatMissed {                missed_heartbeats,                remote_connection,                delegate,                multiplex_task,                heartbeat_task,            } => Self::HeartbeatMissed {                missed_heartbeats: missed_heartbeats + 1,                remote_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 RemoteClient {    client: Arc<ChannelClient>,    unique_identifier: String,    connection_options: RemoteConnectionOptions,    path_style: PathStyle,    state: Option<State>,} #[derive(Debug)]pub enum RemoteClientEvent {    Disconnected,} impl EventEmitter<RemoteClientEvent> for RemoteClient {} /// 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    //   https://unix.stackexchange.com/questions/367008/why-is-socket-path-length-limited-to-a-hundred-chars    // So our strings should be at most 20 characters or so.    fn to_string(&self, cx: &App) -> String {        let identifier_prefix = match ReleaseChannel::global(cx) {            ReleaseChannel::Stable => "".to_string(),            release_channel => format!("{}-", release_channel.dev_name()),        };        match self {            Self::Setup(setup_id) => format!("{identifier_prefix}setup-{setup_id}"),            Self::Workspace(workspace_id) => {                format!("{identifier_prefix}workspace-{workspace_id}",)            }        }    }} pub async fn connect(    connection_options: RemoteConnectionOptions,    delegate: Arc<dyn RemoteClientDelegate>,    cx: &mut AsyncApp,) -> Result<Arc<dyn RemoteConnection>> {    cx.update(|cx| {        cx.update_default_global(|pool: &mut ConnectionPool, cx| {            pool.connect(connection_options.clone(), delegate.clone(), cx)        })    })    .await    .map_err(|e| e.cloned())} impl RemoteClient {    pub fn new(        unique_identifier: ConnectionIdentifier,        remote_connection: Arc<dyn RemoteConnection>,        cancellation: oneshot::Receiver<()>,        delegate: Arc<dyn RemoteClientDelegate>,        cx: &mut App,    ) -> Task<Result<Option<Entity<Self>>>> {        let unique_identifier = unique_identifier.to_string(cx);        cx.spawn(async move |cx| {            let success = Box::pin(async move {                let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();                let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();                let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);                 let client = cx.update(|cx| {                    ChannelClient::new(                        incoming_rx,                        outgoing_tx,                        cx,                        "client",                        remote_connection.has_wsl_interop(),                    )                });                 let path_style = remote_connection.path_style();                let this = cx.new(|_| Self {                    client: client.clone(),                    unique_identifier: unique_identifier.clone(),                    connection_options: remote_connection.connection_options(),                    path_style,                    state: Some(State::Connecting),                });                 let io_task = remote_connection.start_proxy(                    unique_identifier,                    false,                    incoming_tx,                    outgoing_rx,                    connection_activity_tx,                    delegate.clone(),                    cx,                );                 let ready = client                    .wait_for_remote_started()                    .with_timeout(INITIAL_CONNECTION_TIMEOUT, cx.background_executor())                    .await;                match ready {                    Ok(Some(_)) => {}                    Ok(None) => {                        let mut error = "remote client exited before becoming ready".to_owned();                        if let Some(status) = io_task.now_or_never() {                            match status {                                Ok(exit_code) => {                                    error.push_str(&format!(", exit_code={exit_code:?}"))                                }                                Err(e) => error.push_str(&format!(", error={e:?}")),                            }                        }                        let error = anyhow::anyhow!("{error}");                        log::error!("failed to establish connection: {}", error);                        return Err(error);                    }                    Err(_) => {                        let mut error =                            "remote client did not become ready within the timeout".to_owned();                        if let Some(status) = io_task.now_or_never() {                            match status {                                Ok(exit_code) => {                                    error.push_str(&format!(", exit_code={exit_code:?}"))                                }                                Err(e) => error.push_str(&format!(", error={e:?}")),                            }                        }                        let error = anyhow::anyhow!("{error}");                        log::error!("failed to establish connection: {}", error);                        return Err(error);                    }                }                let multiplex_task = Self::monitor(this.downgrade(), io_task, cx);                if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {                    log::error!("failed to establish connection: {}", error);                    return Err(error);                }                 let heartbeat_task = Self::heartbeat(this.downgrade(), connection_activity_rx, cx);                 this.update(cx, |this, _| {                    this.state = Some(State::Connected {                        remote_connection,                        delegate,                        multiplex_task,                        heartbeat_task,                    });                });                 Ok(Some(this))            });             select! {                _ = cancellation.fuse() => {                    Ok(None)                }                result = success.fuse() =>  result            }        })    }     pub fn proto_client_from_channels(        incoming_rx: mpsc::UnboundedReceiver<Envelope>,        outgoing_tx: mpsc::UnboundedSender<Envelope>,        cx: &App,        name: &'static str,        has_wsl_interop: bool,    ) -> AnyProtoClient {        ChannelClient::new(incoming_rx, outgoing_tx, cx, name, has_wsl_interop).into()    }     pub fn shutdown_processes<T: RequestMessage>(        &mut self,

Lines 1–500 of 1692.