use crate::{    SshConnectionOptions,    protocol::MessageId,    proxy::ProxyLaunchError,    transport::{        ssh::SshRemoteConnection,        wsl::{WslConnectionOptions, WslRemoteConnection},    },};use anyhow::{Context as _, Result, anyhow};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},    select, select_biased,};use gpui::{    App, AppContext as _, AsyncApp, BackgroundExecutor, BorrowAppContext, Context, Entity,    EventEmitter, Global, SemanticVersion, 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 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)]pub struct RemotePlatform {    pub os: &'static str,    pub arch: &'static str,} #[derive(Clone, Debug)]pub struct CommandTemplate {    pub program: String,    pub args: Vec<String>,    pub env: HashMap<String, String>,} pub trait RemoteClientDelegate: Send + Sync {    fn ask_password(&self, prompt: String, tx: oneshot::Sender<String>, cx: &mut AsyncApp);    fn get_download_params(        &self,        platform: RemotePlatform,        release_channel: ReleaseChannel,        version: Option<SemanticVersion>,        cx: &mut AsyncApp,    ) -> Task<Result<Option<(String, String)>>>;    fn download_server_binary_locally(        &self,        platform: RemotePlatform,        release_channel: ReleaseChannel,        version: Option<SemanticVersion>,        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 MAX_RECONNECT_ATTEMPTS: usize = 3; enum State {    Connecting,    Connected {        ssh_connection: Arc<dyn RemoteConnection>,        delegate: Arc<dyn RemoteClientDelegate>,         multiplex_task: Task<Result<()>>,        heartbeat_task: Task<Result<()>>,    },    HeartbeatMissed {        missed_heartbeats: usize,         ssh_connection: Arc<dyn RemoteConnection>,        delegate: Arc<dyn RemoteClientDelegate>,         multiplex_task: Task<Result<()>>,        heartbeat_task: Task<Result<()>>,    },    Reconnecting,    ReconnectFailed {        ssh_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 { ssh_connection, .. } => Some(ssh_connection.clone()),            Self::HeartbeatMissed { ssh_connection, .. } => Some(ssh_connection.clone()),            Self::ReconnectFailed { ssh_connection, .. } => Some(ssh_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 {                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 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}",)            }        }    }} impl RemoteClient {    pub fn ssh(        unique_identifier: ConnectionIdentifier,        connection_options: SshConnectionOptions,        cancellation: oneshot::Receiver<()>,        delegate: Arc<dyn RemoteClientDelegate>,        cx: &mut App,    ) -> Task<Result<Option<Entity<Self>>>> {        Self::new(            unique_identifier,            RemoteConnectionOptions::Ssh(connection_options),            cancellation,            delegate,            cx,        )    }     pub fn new(        unique_identifier: ConnectionIdentifier,        connection_options: RemoteConnectionOptions,        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"))?;                 let ssh_connection = cx                    .update(|cx| {                        cx.update_default_global(|pool: &mut ConnectionPool, cx| {                            pool.connect(connection_options.clone(), &delegate, cx)                        })                    })?                    .await                    .map_err(|e| e.cloned())?;                 let path_style = ssh_connection.path_style();                let this = cx.new(|_| Self {                    client: client.clone(),                    unique_identifier: unique_identifier.clone(),                    connection_options,                    path_style,                    state: Some(State::Connecting),                })?;                 let io_task = ssh_connection.start_proxy(                    unique_identifier,                    false,                    incoming_tx,                    outgoing_rx,                    connection_activity_tx,                    delegate.clone(),                    cx,                );                 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 {                        ssh_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,    ) -> AnyProtoClient {        ChannelClient::new(incoming_rx, outgoing_tx, cx, name).into()    }     pub fn shutdown_processes<T: RequestMessage>(        &mut self,        shutdown_request: Option<T>,        executor: BackgroundExecutor,    ) -> Option<impl Future<Output = ()> + use<T>> {        let state = self.state.take()?;        log::info!("shutting down ssh processes");         let State::Connected {            multiplex_task,            heartbeat_task,            ssh_connection,            delegate,        } = state        else {            return None;        };         let client = self.client.clone();         Some(async move {            if let Some(shutdown_request) = shutdown_request {                client.send(shutdown_request).log_err();                // We wait 50ms instead of waiting for a response, because                // waiting for a response would require us to wait on the main thread                // which we want to avoid in an `on_app_quit` callback.                executor.timer(Duration::from_millis(50)).await;            }             // Drop `multiplex_task` because it owns our ssh_proxy_process, which is a            // child of master_process.            drop(multiplex_task);            // Now drop the rest of state, which kills master process.            drop(heartbeat_task);            drop(ssh_connection);            drop(delegate);        })    }     fn reconnect(&mut self, cx: &mut Context<Self>) -> Result<()> {        let can_reconnect = self            .state            .as_ref()            .map(|state| state.can_reconnect())            .unwrap_or(false);        if !can_reconnect {            log::info!("aborting reconnect, because not in state that allows reconnecting");            let error = if let Some(state) = self.state.as_ref() {                format!("invalid state, cannot reconnect while in state {state}")            } else {                "no state set".to_string()            };            anyhow::bail!(error);        }         let state = self.state.take().unwrap();        let (attempts, remote_connection, delegate) = match state {            State::Connected {                ssh_connection,                delegate,                multiplex_task,                heartbeat_task,            }            | State::HeartbeatMissed {                ssh_connection,                delegate,                multiplex_task,                heartbeat_task,                ..            } => {                drop(multiplex_task);                drop(heartbeat_task);                (0, ssh_connection, delegate)            }            State::ReconnectFailed {                attempts,                ssh_connection,                delegate,                ..            } => (attempts, ssh_connection, delegate),            State::Connecting            | State::Reconnecting            | State::ReconnectExhausted            | State::ServerNotRunning => unreachable!(),        };         let attempts = attempts + 1;        if attempts > MAX_RECONNECT_ATTEMPTS {            log::error!(                "Failed to reconnect to after {} attempts, giving up",                MAX_RECONNECT_ATTEMPTS            );            self.set_state(State::ReconnectExhausted, cx);            return Ok(());        }         self.set_state(State::Reconnecting, cx);         log::info!("Trying to reconnect to ssh server... Attempt {}", attempts);         let unique_identifier = self.unique_identifier.clone();        let client = self.client.clone();        let reconnect_task = cx.spawn(async move |this, cx| {            macro_rules! failed {                ($error:expr, $attempts:expr, $ssh_connection:expr, $delegate:expr) => {                    return State::ReconnectFailed {                        error: anyhow!($error),                        attempts: $attempts,                        ssh_connection: $ssh_connection,

Lines 1–500 of 1541. Load more