ssh_session.rs

   1use crate::{
   2    json_log::LogRecord,
   3    protocol::{
   4        message_len_from_buffer, read_message_with_len, write_message, MessageId, MESSAGE_LEN_SIZE,
   5    },
   6    proxy::ProxyLaunchError,
   7};
   8use anyhow::{anyhow, Context as _, Result};
   9use async_trait::async_trait;
  10use collections::HashMap;
  11use futures::{
  12    channel::{
  13        mpsc::{self, Sender, UnboundedReceiver, UnboundedSender},
  14        oneshot,
  15    },
  16    future::{BoxFuture, Shared},
  17    select, select_biased, AsyncReadExt as _, Future, FutureExt as _, StreamExt as _,
  18};
  19use gpui::{
  20    AppContext, AsyncAppContext, BorrowAppContext, Context, EventEmitter, Global, Model,
  21    ModelContext, SemanticVersion, Task, WeakModel,
  22};
  23use itertools::Itertools;
  24use parking_lot::Mutex;
  25use paths;
  26use release_channel::{AppCommitSha, AppVersion, ReleaseChannel};
  27use rpc::{
  28    proto::{self, build_typed_envelope, Envelope, EnvelopedMessage, PeerId, RequestMessage},
  29    AnyProtoClient, EntityMessageSubscriber, ErrorExt, ProtoClient, ProtoMessageHandlerSet,
  30    RpcError,
  31};
  32use smol::{
  33    fs,
  34    process::{self, Child, Stdio},
  35};
  36use std::{
  37    any::TypeId,
  38    collections::VecDeque,
  39    fmt, iter,
  40    ops::ControlFlow,
  41    path::{Path, PathBuf},
  42    sync::{
  43        atomic::{AtomicU32, Ordering::SeqCst},
  44        Arc, Weak,
  45    },
  46    time::{Duration, Instant},
  47};
  48use tempfile::TempDir;
  49use util::ResultExt;
  50
  51#[derive(
  52    Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, serde::Serialize, serde::Deserialize,
  53)]
  54pub struct SshProjectId(pub u64);
  55
  56#[derive(Clone)]
  57pub struct SshSocket {
  58    connection_options: SshConnectionOptions,
  59    socket_path: PathBuf,
  60}
  61
  62#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
  63pub struct SshConnectionOptions {
  64    pub host: String,
  65    pub username: Option<String>,
  66    pub port: Option<u16>,
  67    pub password: Option<String>,
  68    pub args: Option<Vec<String>>,
  69
  70    pub nickname: Option<String>,
  71    pub upload_binary_over_ssh: bool,
  72}
  73
  74#[macro_export]
  75macro_rules! shell_script {
  76    ($fmt:expr, $($name:ident = $arg:expr),+ $(,)?) => {{
  77        format!(
  78            $fmt,
  79            $(
  80                $name = shlex::try_quote($arg).unwrap()
  81            ),+
  82        )
  83    }};
  84}
  85
  86impl SshConnectionOptions {
  87    pub fn parse_command_line(input: &str) -> Result<Self> {
  88        let input = input.trim_start_matches("ssh ");
  89        let mut hostname: Option<String> = None;
  90        let mut username: Option<String> = None;
  91        let mut port: Option<u16> = None;
  92        let mut args = Vec::new();
  93
  94        // disallowed: -E, -e, -F, -f, -G, -g, -M, -N, -n, -O, -q, -S, -s, -T, -t, -V, -v, -W
  95        const ALLOWED_OPTS: &[&str] = &[
  96            "-4", "-6", "-A", "-a", "-C", "-K", "-k", "-X", "-x", "-Y", "-y",
  97        ];
  98        const ALLOWED_ARGS: &[&str] = &[
  99            "-B", "-b", "-c", "-D", "-I", "-i", "-J", "-L", "-l", "-m", "-o", "-P", "-p", "-R",
 100            "-w",
 101        ];
 102
 103        let mut tokens = shlex::split(input)
 104            .ok_or_else(|| anyhow!("invalid input"))?
 105            .into_iter();
 106
 107        'outer: while let Some(arg) = tokens.next() {
 108            if ALLOWED_OPTS.contains(&(&arg as &str)) {
 109                args.push(arg.to_string());
 110                continue;
 111            }
 112            if arg == "-p" {
 113                port = tokens.next().and_then(|arg| arg.parse().ok());
 114                continue;
 115            } else if let Some(p) = arg.strip_prefix("-p") {
 116                port = p.parse().ok();
 117                continue;
 118            }
 119            if arg == "-l" {
 120                username = tokens.next();
 121                continue;
 122            } else if let Some(l) = arg.strip_prefix("-l") {
 123                username = Some(l.to_string());
 124                continue;
 125            }
 126            for a in ALLOWED_ARGS {
 127                if arg == *a {
 128                    args.push(arg);
 129                    if let Some(next) = tokens.next() {
 130                        args.push(next);
 131                    }
 132                    continue 'outer;
 133                } else if arg.starts_with(a) {
 134                    args.push(arg);
 135                    continue 'outer;
 136                }
 137            }
 138            if arg.starts_with("-") || hostname.is_some() {
 139                anyhow::bail!("unsupported argument: {:?}", arg);
 140            }
 141            let mut input = &arg as &str;
 142            if let Some((u, rest)) = input.split_once('@') {
 143                input = rest;
 144                username = Some(u.to_string());
 145            }
 146            if let Some((rest, p)) = input.split_once(':') {
 147                input = rest;
 148                port = p.parse().ok()
 149            }
 150            hostname = Some(input.to_string())
 151        }
 152
 153        let Some(hostname) = hostname else {
 154            anyhow::bail!("missing hostname");
 155        };
 156
 157        Ok(Self {
 158            host: hostname.to_string(),
 159            username: username.clone(),
 160            port,
 161            args: Some(args),
 162            password: None,
 163            nickname: None,
 164            upload_binary_over_ssh: false,
 165        })
 166    }
 167
 168    pub fn ssh_url(&self) -> String {
 169        let mut result = String::from("ssh://");
 170        if let Some(username) = &self.username {
 171            result.push_str(username);
 172            result.push('@');
 173        }
 174        result.push_str(&self.host);
 175        if let Some(port) = self.port {
 176            result.push(':');
 177            result.push_str(&port.to_string());
 178        }
 179        result
 180    }
 181
 182    pub fn additional_args(&self) -> Option<&Vec<String>> {
 183        self.args.as_ref()
 184    }
 185
 186    fn scp_url(&self) -> String {
 187        if let Some(username) = &self.username {
 188            format!("{}@{}", username, self.host)
 189        } else {
 190            self.host.clone()
 191        }
 192    }
 193
 194    pub fn connection_string(&self) -> String {
 195        let host = if let Some(username) = &self.username {
 196            format!("{}@{}", username, self.host)
 197        } else {
 198            self.host.clone()
 199        };
 200        if let Some(port) = &self.port {
 201            format!("{}:{}", host, port)
 202        } else {
 203            host
 204        }
 205    }
 206}
 207
 208#[derive(Copy, Clone, Debug)]
 209pub struct SshPlatform {
 210    pub os: &'static str,
 211    pub arch: &'static str,
 212}
 213
 214impl SshPlatform {
 215    pub fn triple(&self) -> Option<String> {
 216        Some(format!(
 217            "{}-{}",
 218            self.arch,
 219            match self.os {
 220                "linux" => "unknown-linux-gnu",
 221                "macos" => "apple-darwin",
 222                _ => return None,
 223            }
 224        ))
 225    }
 226}
 227
 228pub trait SshClientDelegate: Send + Sync {
 229    fn ask_password(
 230        &self,
 231        prompt: String,
 232        cx: &mut AsyncAppContext,
 233    ) -> oneshot::Receiver<Result<String>>;
 234    fn get_download_params(
 235        &self,
 236        platform: SshPlatform,
 237        release_channel: ReleaseChannel,
 238        version: Option<SemanticVersion>,
 239        cx: &mut AsyncAppContext,
 240    ) -> Task<Result<Option<(String, String)>>>;
 241
 242    fn download_server_binary_locally(
 243        &self,
 244        platform: SshPlatform,
 245        release_channel: ReleaseChannel,
 246        version: Option<SemanticVersion>,
 247        cx: &mut AsyncAppContext,
 248    ) -> Task<Result<PathBuf>>;
 249    fn set_status(&self, status: Option<&str>, cx: &mut AsyncAppContext);
 250}
 251
 252impl SshSocket {
 253    // :WARNING: ssh unquotes arguments when executing on the remote :WARNING:
 254    // e.g. $ ssh host sh -c 'ls -l' is equivalent to $ ssh host sh -c ls -l
 255    // and passes -l as an argument to sh, not to ls.
 256    // You need to do it like this: $ ssh host "sh -c 'ls -l /tmp'"
 257    fn ssh_command(&self, program: &str, args: &[&str]) -> process::Command {
 258        let mut command = process::Command::new("ssh");
 259        let to_run = iter::once(&program)
 260            .chain(args.iter())
 261            .map(|token| {
 262                // We're trying to work with: sh, bash, zsh, fish, tcsh, ...?
 263                debug_assert!(
 264                    !token.contains('\n'),
 265                    "multiline arguments do not work in all shells"
 266                );
 267                shlex::try_quote(token).unwrap()
 268            })
 269            .join(" ");
 270        log::debug!("ssh {} {:?}", self.connection_options.ssh_url(), to_run);
 271        self.ssh_options(&mut command)
 272            .arg(self.connection_options.ssh_url())
 273            .arg(to_run);
 274        command
 275    }
 276
 277    async fn run_command(&self, program: &str, args: &[&str]) -> Result<String> {
 278        let output = self.ssh_command(program, args).output().await?;
 279        if output.status.success() {
 280            Ok(String::from_utf8_lossy(&output.stdout).to_string())
 281        } else {
 282            Err(anyhow!(
 283                "failed to run command: {}",
 284                String::from_utf8_lossy(&output.stderr)
 285            ))
 286        }
 287    }
 288
 289    fn ssh_options<'a>(&self, command: &'a mut process::Command) -> &'a mut process::Command {
 290        command
 291            .stdin(Stdio::piped())
 292            .stdout(Stdio::piped())
 293            .stderr(Stdio::piped())
 294            .args(["-o", "ControlMaster=no", "-o"])
 295            .arg(format!("ControlPath={}", self.socket_path.display()))
 296    }
 297
 298    fn ssh_args(&self) -> Vec<String> {
 299        vec![
 300            "-o".to_string(),
 301            "ControlMaster=no".to_string(),
 302            "-o".to_string(),
 303            format!("ControlPath={}", self.socket_path.display()),
 304            self.connection_options.ssh_url(),
 305        ]
 306    }
 307}
 308
 309const MAX_MISSED_HEARTBEATS: usize = 5;
 310const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
 311const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5);
 312
 313const MAX_RECONNECT_ATTEMPTS: usize = 3;
 314
 315enum State {
 316    Connecting,
 317    Connected {
 318        ssh_connection: Arc<dyn RemoteConnection>,
 319        delegate: Arc<dyn SshClientDelegate>,
 320
 321        multiplex_task: Task<Result<()>>,
 322        heartbeat_task: Task<Result<()>>,
 323    },
 324    HeartbeatMissed {
 325        missed_heartbeats: usize,
 326
 327        ssh_connection: Arc<dyn RemoteConnection>,
 328        delegate: Arc<dyn SshClientDelegate>,
 329
 330        multiplex_task: Task<Result<()>>,
 331        heartbeat_task: Task<Result<()>>,
 332    },
 333    Reconnecting,
 334    ReconnectFailed {
 335        ssh_connection: Arc<dyn RemoteConnection>,
 336        delegate: Arc<dyn SshClientDelegate>,
 337
 338        error: anyhow::Error,
 339        attempts: usize,
 340    },
 341    ReconnectExhausted,
 342    ServerNotRunning,
 343}
 344
 345impl fmt::Display for State {
 346    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 347        match self {
 348            Self::Connecting => write!(f, "connecting"),
 349            Self::Connected { .. } => write!(f, "connected"),
 350            Self::Reconnecting => write!(f, "reconnecting"),
 351            Self::ReconnectFailed { .. } => write!(f, "reconnect failed"),
 352            Self::ReconnectExhausted => write!(f, "reconnect exhausted"),
 353            Self::HeartbeatMissed { .. } => write!(f, "heartbeat missed"),
 354            Self::ServerNotRunning { .. } => write!(f, "server not running"),
 355        }
 356    }
 357}
 358
 359impl State {
 360    fn ssh_connection(&self) -> Option<&dyn RemoteConnection> {
 361        match self {
 362            Self::Connected { ssh_connection, .. } => Some(ssh_connection.as_ref()),
 363            Self::HeartbeatMissed { ssh_connection, .. } => Some(ssh_connection.as_ref()),
 364            Self::ReconnectFailed { ssh_connection, .. } => Some(ssh_connection.as_ref()),
 365            _ => None,
 366        }
 367    }
 368
 369    fn can_reconnect(&self) -> bool {
 370        match self {
 371            Self::Connected { .. }
 372            | Self::HeartbeatMissed { .. }
 373            | Self::ReconnectFailed { .. } => true,
 374            State::Connecting
 375            | State::Reconnecting
 376            | State::ReconnectExhausted
 377            | State::ServerNotRunning => false,
 378        }
 379    }
 380
 381    fn is_reconnect_failed(&self) -> bool {
 382        matches!(self, Self::ReconnectFailed { .. })
 383    }
 384
 385    fn is_reconnect_exhausted(&self) -> bool {
 386        matches!(self, Self::ReconnectExhausted { .. })
 387    }
 388
 389    fn is_server_not_running(&self) -> bool {
 390        matches!(self, Self::ServerNotRunning)
 391    }
 392
 393    fn is_reconnecting(&self) -> bool {
 394        matches!(self, Self::Reconnecting { .. })
 395    }
 396
 397    fn heartbeat_recovered(self) -> Self {
 398        match self {
 399            Self::HeartbeatMissed {
 400                ssh_connection,
 401                delegate,
 402                multiplex_task,
 403                heartbeat_task,
 404                ..
 405            } => Self::Connected {
 406                ssh_connection,
 407                delegate,
 408                multiplex_task,
 409                heartbeat_task,
 410            },
 411            _ => self,
 412        }
 413    }
 414
 415    fn heartbeat_missed(self) -> Self {
 416        match self {
 417            Self::Connected {
 418                ssh_connection,
 419                delegate,
 420                multiplex_task,
 421                heartbeat_task,
 422            } => Self::HeartbeatMissed {
 423                missed_heartbeats: 1,
 424                ssh_connection,
 425                delegate,
 426                multiplex_task,
 427                heartbeat_task,
 428            },
 429            Self::HeartbeatMissed {
 430                missed_heartbeats,
 431                ssh_connection,
 432                delegate,
 433                multiplex_task,
 434                heartbeat_task,
 435            } => Self::HeartbeatMissed {
 436                missed_heartbeats: missed_heartbeats + 1,
 437                ssh_connection,
 438                delegate,
 439                multiplex_task,
 440                heartbeat_task,
 441            },
 442            _ => self,
 443        }
 444    }
 445}
 446
 447/// The state of the ssh connection.
 448#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 449pub enum ConnectionState {
 450    Connecting,
 451    Connected,
 452    HeartbeatMissed,
 453    Reconnecting,
 454    Disconnected,
 455}
 456
 457impl From<&State> for ConnectionState {
 458    fn from(value: &State) -> Self {
 459        match value {
 460            State::Connecting => Self::Connecting,
 461            State::Connected { .. } => Self::Connected,
 462            State::Reconnecting | State::ReconnectFailed { .. } => Self::Reconnecting,
 463            State::HeartbeatMissed { .. } => Self::HeartbeatMissed,
 464            State::ReconnectExhausted => Self::Disconnected,
 465            State::ServerNotRunning => Self::Disconnected,
 466        }
 467    }
 468}
 469
 470pub struct SshRemoteClient {
 471    client: Arc<ChannelClient>,
 472    unique_identifier: String,
 473    connection_options: SshConnectionOptions,
 474    state: Arc<Mutex<Option<State>>>,
 475}
 476
 477#[derive(Debug)]
 478pub enum SshRemoteEvent {
 479    Disconnected,
 480}
 481
 482impl EventEmitter<SshRemoteEvent> for SshRemoteClient {}
 483
 484// Identifies the socket on the remote server so that reconnects
 485// can re-join the same project.
 486pub enum ConnectionIdentifier {
 487    Setup,
 488    Workspace(i64),
 489}
 490
 491impl ConnectionIdentifier {
 492    // This string gets used in a socket name, and so must be relatively short.
 493    // The total length of:
 494    //   /home/{username}/.local/share/zed/server_state/{name}/stdout.sock
 495    // Must be less than about 100 characters
 496    //   https://unix.stackexchange.com/questions/367008/why-is-socket-path-length-limited-to-a-hundred-chars
 497    // So our strings should be at most 20 characters or so.
 498    fn to_string(&self, cx: &AppContext) -> String {
 499        let identifier_prefix = match ReleaseChannel::global(cx) {
 500            ReleaseChannel::Stable => "".to_string(),
 501            release_channel => format!("{}-", release_channel.dev_name()),
 502        };
 503        match self {
 504            Self::Setup => format!("{identifier_prefix}setup"),
 505            Self::Workspace(workspace_id) => {
 506                format!("{identifier_prefix}workspace-{workspace_id}",)
 507            }
 508        }
 509    }
 510}
 511
 512impl SshRemoteClient {
 513    pub fn new(
 514        unique_identifier: ConnectionIdentifier,
 515        connection_options: SshConnectionOptions,
 516        cancellation: oneshot::Receiver<()>,
 517        delegate: Arc<dyn SshClientDelegate>,
 518        cx: &mut AppContext,
 519    ) -> Task<Result<Option<Model<Self>>>> {
 520        let unique_identifier = unique_identifier.to_string(cx);
 521        cx.spawn(|mut cx| async move {
 522            let success = Box::pin(async move {
 523                let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
 524                let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
 525                let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
 526
 527                let client =
 528                    cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "client"))?;
 529                let this = cx.new_model(|_| Self {
 530                    client: client.clone(),
 531                    unique_identifier: unique_identifier.clone(),
 532                    connection_options: connection_options.clone(),
 533                    state: Arc::new(Mutex::new(Some(State::Connecting))),
 534                })?;
 535
 536                let ssh_connection = cx
 537                    .update(|cx| {
 538                        cx.update_default_global(|pool: &mut ConnectionPool, cx| {
 539                            pool.connect(connection_options, &delegate, cx)
 540                        })
 541                    })?
 542                    .await
 543                    .map_err(|e| e.cloned())?;
 544
 545                let io_task = ssh_connection.start_proxy(
 546                    unique_identifier,
 547                    false,
 548                    incoming_tx,
 549                    outgoing_rx,
 550                    connection_activity_tx,
 551                    delegate.clone(),
 552                    &mut cx,
 553                );
 554
 555                let multiplex_task = Self::monitor(this.downgrade(), io_task, &cx);
 556
 557                if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
 558                    log::error!("failed to establish connection: {}", error);
 559                    return Err(error);
 560                }
 561
 562                let heartbeat_task =
 563                    Self::heartbeat(this.downgrade(), connection_activity_rx, &mut cx);
 564
 565                this.update(&mut cx, |this, _| {
 566                    *this.state.lock() = Some(State::Connected {
 567                        ssh_connection,
 568                        delegate,
 569                        multiplex_task,
 570                        heartbeat_task,
 571                    });
 572                })?;
 573
 574                Ok(Some(this))
 575            });
 576
 577            select! {
 578                _ = cancellation.fuse() => {
 579                    Ok(None)
 580                }
 581                result = success.fuse() =>  result
 582            }
 583        })
 584    }
 585
 586    pub fn shutdown_processes<T: RequestMessage>(
 587        &self,
 588        shutdown_request: Option<T>,
 589    ) -> Option<impl Future<Output = ()>> {
 590        let state = self.state.lock().take()?;
 591        log::info!("shutting down ssh processes");
 592
 593        let State::Connected {
 594            multiplex_task,
 595            heartbeat_task,
 596            ssh_connection,
 597            delegate,
 598        } = state
 599        else {
 600            return None;
 601        };
 602
 603        let client = self.client.clone();
 604
 605        Some(async move {
 606            if let Some(shutdown_request) = shutdown_request {
 607                client.send(shutdown_request).log_err();
 608                // We wait 50ms instead of waiting for a response, because
 609                // waiting for a response would require us to wait on the main thread
 610                // which we want to avoid in an `on_app_quit` callback.
 611                smol::Timer::after(Duration::from_millis(50)).await;
 612            }
 613
 614            // Drop `multiplex_task` because it owns our ssh_proxy_process, which is a
 615            // child of master_process.
 616            drop(multiplex_task);
 617            // Now drop the rest of state, which kills master process.
 618            drop(heartbeat_task);
 619            drop(ssh_connection);
 620            drop(delegate);
 621        })
 622    }
 623
 624    fn reconnect(&mut self, cx: &mut ModelContext<Self>) -> Result<()> {
 625        let mut lock = self.state.lock();
 626
 627        let can_reconnect = lock
 628            .as_ref()
 629            .map(|state| state.can_reconnect())
 630            .unwrap_or(false);
 631        if !can_reconnect {
 632            let error = if let Some(state) = lock.as_ref() {
 633                format!("invalid state, cannot reconnect while in state {state}")
 634            } else {
 635                "no state set".to_string()
 636            };
 637            log::info!("aborting reconnect, because not in state that allows reconnecting");
 638            return Err(anyhow!(error));
 639        }
 640
 641        let state = lock.take().unwrap();
 642        let (attempts, ssh_connection, delegate) = match state {
 643            State::Connected {
 644                ssh_connection,
 645                delegate,
 646                multiplex_task,
 647                heartbeat_task,
 648            }
 649            | State::HeartbeatMissed {
 650                ssh_connection,
 651                delegate,
 652                multiplex_task,
 653                heartbeat_task,
 654                ..
 655            } => {
 656                drop(multiplex_task);
 657                drop(heartbeat_task);
 658                (0, ssh_connection, delegate)
 659            }
 660            State::ReconnectFailed {
 661                attempts,
 662                ssh_connection,
 663                delegate,
 664                ..
 665            } => (attempts, ssh_connection, delegate),
 666            State::Connecting
 667            | State::Reconnecting
 668            | State::ReconnectExhausted
 669            | State::ServerNotRunning => unreachable!(),
 670        };
 671
 672        let attempts = attempts + 1;
 673        if attempts > MAX_RECONNECT_ATTEMPTS {
 674            log::error!(
 675                "Failed to reconnect to after {} attempts, giving up",
 676                MAX_RECONNECT_ATTEMPTS
 677            );
 678            drop(lock);
 679            self.set_state(State::ReconnectExhausted, cx);
 680            return Ok(());
 681        }
 682        drop(lock);
 683
 684        self.set_state(State::Reconnecting, cx);
 685
 686        log::info!("Trying to reconnect to ssh server... Attempt {}", attempts);
 687
 688        let unique_identifier = self.unique_identifier.clone();
 689        let client = self.client.clone();
 690        let reconnect_task = cx.spawn(|this, mut cx| async move {
 691            macro_rules! failed {
 692                ($error:expr, $attempts:expr, $ssh_connection:expr, $delegate:expr) => {
 693                    return State::ReconnectFailed {
 694                        error: anyhow!($error),
 695                        attempts: $attempts,
 696                        ssh_connection: $ssh_connection,
 697                        delegate: $delegate,
 698                    };
 699                };
 700            }
 701
 702            if let Err(error) = ssh_connection
 703                .kill()
 704                .await
 705                .context("Failed to kill ssh process")
 706            {
 707                failed!(error, attempts, ssh_connection, delegate);
 708            };
 709
 710            let connection_options = ssh_connection.connection_options();
 711
 712            let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
 713            let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
 714            let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
 715
 716            let (ssh_connection, io_task) = match async {
 717                let ssh_connection = cx
 718                    .update_global(|pool: &mut ConnectionPool, cx| {
 719                        pool.connect(connection_options, &delegate, cx)
 720                    })?
 721                    .await
 722                    .map_err(|error| error.cloned())?;
 723
 724                let io_task = ssh_connection.start_proxy(
 725                    unique_identifier,
 726                    true,
 727                    incoming_tx,
 728                    outgoing_rx,
 729                    connection_activity_tx,
 730                    delegate.clone(),
 731                    &mut cx,
 732                );
 733                anyhow::Ok((ssh_connection, io_task))
 734            }
 735            .await
 736            {
 737                Ok((ssh_connection, io_task)) => (ssh_connection, io_task),
 738                Err(error) => {
 739                    failed!(error, attempts, ssh_connection, delegate);
 740                }
 741            };
 742
 743            let multiplex_task = Self::monitor(this.clone(), io_task, &cx);
 744            client.reconnect(incoming_rx, outgoing_tx, &cx);
 745
 746            if let Err(error) = client.resync(HEARTBEAT_TIMEOUT).await {
 747                failed!(error, attempts, ssh_connection, delegate);
 748            };
 749
 750            State::Connected {
 751                ssh_connection,
 752                delegate,
 753                multiplex_task,
 754                heartbeat_task: Self::heartbeat(this.clone(), connection_activity_rx, &mut cx),
 755            }
 756        });
 757
 758        cx.spawn(|this, mut cx| async move {
 759            let new_state = reconnect_task.await;
 760            this.update(&mut cx, |this, cx| {
 761                this.try_set_state(cx, |old_state| {
 762                    if old_state.is_reconnecting() {
 763                        match &new_state {
 764                            State::Connecting
 765                            | State::Reconnecting { .. }
 766                            | State::HeartbeatMissed { .. }
 767                            | State::ServerNotRunning => {}
 768                            State::Connected { .. } => {
 769                                log::info!("Successfully reconnected");
 770                            }
 771                            State::ReconnectFailed {
 772                                error, attempts, ..
 773                            } => {
 774                                log::error!(
 775                                    "Reconnect attempt {} failed: {:?}. Starting new attempt...",
 776                                    attempts,
 777                                    error
 778                                );
 779                            }
 780                            State::ReconnectExhausted => {
 781                                log::error!("Reconnect attempt failed and all attempts exhausted");
 782                            }
 783                        }
 784                        Some(new_state)
 785                    } else {
 786                        None
 787                    }
 788                });
 789
 790                if this.state_is(State::is_reconnect_failed) {
 791                    this.reconnect(cx)
 792                } else if this.state_is(State::is_reconnect_exhausted) {
 793                    Ok(())
 794                } else {
 795                    log::debug!("State has transition from Reconnecting into new state while attempting reconnect.");
 796                    Ok(())
 797                }
 798            })
 799        })
 800        .detach_and_log_err(cx);
 801
 802        Ok(())
 803    }
 804
 805    fn heartbeat(
 806        this: WeakModel<Self>,
 807        mut connection_activity_rx: mpsc::Receiver<()>,
 808        cx: &mut AsyncAppContext,
 809    ) -> Task<Result<()>> {
 810        let Ok(client) = this.update(cx, |this, _| this.client.clone()) else {
 811            return Task::ready(Err(anyhow!("SshRemoteClient lost")));
 812        };
 813
 814        cx.spawn(|mut cx| {
 815            let this = this.clone();
 816            async move {
 817                let mut missed_heartbeats = 0;
 818
 819                let keepalive_timer = cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse();
 820                futures::pin_mut!(keepalive_timer);
 821
 822                loop {
 823                    select_biased! {
 824                        result = connection_activity_rx.next().fuse() => {
 825                            if result.is_none() {
 826                                log::warn!("ssh heartbeat: connection activity channel has been dropped. stopping.");
 827                                return Ok(());
 828                            }
 829
 830                            if missed_heartbeats != 0 {
 831                                missed_heartbeats = 0;
 832                                this.update(&mut cx, |this, mut cx| {
 833                                    this.handle_heartbeat_result(missed_heartbeats, &mut cx)
 834                                })?;
 835                            }
 836                        }
 837                        _ = keepalive_timer => {
 838                            log::debug!("Sending heartbeat to server...");
 839
 840                            let result = select_biased! {
 841                                _ = connection_activity_rx.next().fuse() => {
 842                                    Ok(())
 843                                }
 844                                ping_result = client.ping(HEARTBEAT_TIMEOUT).fuse() => {
 845                                    ping_result
 846                                }
 847                            };
 848
 849                            if result.is_err() {
 850                                missed_heartbeats += 1;
 851                                log::warn!(
 852                                    "No heartbeat from server after {:?}. Missed heartbeat {} out of {}.",
 853                                    HEARTBEAT_TIMEOUT,
 854                                    missed_heartbeats,
 855                                    MAX_MISSED_HEARTBEATS
 856                                );
 857                            } else if missed_heartbeats != 0 {
 858                                missed_heartbeats = 0;
 859                            } else {
 860                                continue;
 861                            }
 862
 863                            let result = this.update(&mut cx, |this, mut cx| {
 864                                this.handle_heartbeat_result(missed_heartbeats, &mut cx)
 865                            })?;
 866                            if result.is_break() {
 867                                return Ok(());
 868                            }
 869                        }
 870                    }
 871
 872                    keepalive_timer.set(cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse());
 873                }
 874            }
 875        })
 876    }
 877
 878    fn handle_heartbeat_result(
 879        &mut self,
 880        missed_heartbeats: usize,
 881        cx: &mut ModelContext<Self>,
 882    ) -> ControlFlow<()> {
 883        let state = self.state.lock().take().unwrap();
 884        let next_state = if missed_heartbeats > 0 {
 885            state.heartbeat_missed()
 886        } else {
 887            state.heartbeat_recovered()
 888        };
 889
 890        self.set_state(next_state, cx);
 891
 892        if missed_heartbeats >= MAX_MISSED_HEARTBEATS {
 893            log::error!(
 894                "Missed last {} heartbeats. Reconnecting...",
 895                missed_heartbeats
 896            );
 897
 898            self.reconnect(cx)
 899                .context("failed to start reconnect process after missing heartbeats")
 900                .log_err();
 901            ControlFlow::Break(())
 902        } else {
 903            ControlFlow::Continue(())
 904        }
 905    }
 906
 907    fn monitor(
 908        this: WeakModel<Self>,
 909        io_task: Task<Result<i32>>,
 910        cx: &AsyncAppContext,
 911    ) -> Task<Result<()>> {
 912        cx.spawn(|mut cx| async move {
 913            let result = io_task.await;
 914
 915            match result {
 916                Ok(exit_code) => {
 917                    if let Some(error) = ProxyLaunchError::from_exit_code(exit_code) {
 918                        match error {
 919                            ProxyLaunchError::ServerNotRunning => {
 920                                log::error!("failed to reconnect because server is not running");
 921                                this.update(&mut cx, |this, cx| {
 922                                    this.set_state(State::ServerNotRunning, cx);
 923                                })?;
 924                            }
 925                        }
 926                    } else if exit_code > 0 {
 927                        log::error!("proxy process terminated unexpectedly");
 928                        this.update(&mut cx, |this, cx| {
 929                            this.reconnect(cx).ok();
 930                        })?;
 931                    }
 932                }
 933                Err(error) => {
 934                    log::warn!("ssh io task died with error: {:?}. reconnecting...", error);
 935                    this.update(&mut cx, |this, cx| {
 936                        this.reconnect(cx).ok();
 937                    })?;
 938                }
 939            }
 940
 941            Ok(())
 942        })
 943    }
 944
 945    fn state_is(&self, check: impl FnOnce(&State) -> bool) -> bool {
 946        self.state.lock().as_ref().map_or(false, check)
 947    }
 948
 949    fn try_set_state(
 950        &self,
 951        cx: &mut ModelContext<Self>,
 952        map: impl FnOnce(&State) -> Option<State>,
 953    ) {
 954        let mut lock = self.state.lock();
 955        let new_state = lock.as_ref().and_then(map);
 956
 957        if let Some(new_state) = new_state {
 958            lock.replace(new_state);
 959            cx.notify();
 960        }
 961    }
 962
 963    fn set_state(&self, state: State, cx: &mut ModelContext<Self>) {
 964        log::info!("setting state to '{}'", &state);
 965
 966        let is_reconnect_exhausted = state.is_reconnect_exhausted();
 967        let is_server_not_running = state.is_server_not_running();
 968        self.state.lock().replace(state);
 969
 970        if is_reconnect_exhausted || is_server_not_running {
 971            cx.emit(SshRemoteEvent::Disconnected);
 972        }
 973        cx.notify();
 974    }
 975
 976    pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Model<E>) {
 977        self.client.subscribe_to_entity(remote_id, entity);
 978    }
 979
 980    pub fn ssh_args(&self) -> Option<Vec<String>> {
 981        self.state
 982            .lock()
 983            .as_ref()
 984            .and_then(|state| state.ssh_connection())
 985            .map(|ssh_connection| ssh_connection.ssh_args())
 986    }
 987
 988    pub fn proto_client(&self) -> AnyProtoClient {
 989        self.client.clone().into()
 990    }
 991
 992    pub fn connection_string(&self) -> String {
 993        self.connection_options.connection_string()
 994    }
 995
 996    pub fn connection_options(&self) -> SshConnectionOptions {
 997        self.connection_options.clone()
 998    }
 999
1000    pub fn connection_state(&self) -> ConnectionState {
1001        self.state
1002            .lock()
1003            .as_ref()
1004            .map(ConnectionState::from)
1005            .unwrap_or(ConnectionState::Disconnected)
1006    }
1007
1008    pub fn is_disconnected(&self) -> bool {
1009        self.connection_state() == ConnectionState::Disconnected
1010    }
1011
1012    #[cfg(any(test, feature = "test-support"))]
1013    pub fn simulate_disconnect(&self, client_cx: &mut AppContext) -> Task<()> {
1014        let opts = self.connection_options();
1015        client_cx.spawn(|cx| async move {
1016            let connection = cx
1017                .update_global(|c: &mut ConnectionPool, _| {
1018                    if let Some(ConnectionPoolEntry::Connecting(c)) = c.connections.get(&opts) {
1019                        c.clone()
1020                    } else {
1021                        panic!("missing test connection")
1022                    }
1023                })
1024                .unwrap()
1025                .await
1026                .unwrap();
1027
1028            connection.simulate_disconnect(&cx);
1029        })
1030    }
1031
1032    #[cfg(any(test, feature = "test-support"))]
1033    pub fn fake_server(
1034        client_cx: &mut gpui::TestAppContext,
1035        server_cx: &mut gpui::TestAppContext,
1036    ) -> (SshConnectionOptions, Arc<ChannelClient>) {
1037        let port = client_cx
1038            .update(|cx| cx.default_global::<ConnectionPool>().connections.len() as u16 + 1);
1039        let opts = SshConnectionOptions {
1040            host: "<fake>".to_string(),
1041            port: Some(port),
1042            ..Default::default()
1043        };
1044        let (outgoing_tx, _) = mpsc::unbounded::<Envelope>();
1045        let (_, incoming_rx) = mpsc::unbounded::<Envelope>();
1046        let server_client =
1047            server_cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "fake-server"));
1048        let connection: Arc<dyn RemoteConnection> = Arc::new(fake::FakeRemoteConnection {
1049            connection_options: opts.clone(),
1050            server_cx: fake::SendableCx::new(server_cx),
1051            server_channel: server_client.clone(),
1052        });
1053
1054        client_cx.update(|cx| {
1055            cx.update_default_global(|c: &mut ConnectionPool, cx| {
1056                c.connections.insert(
1057                    opts.clone(),
1058                    ConnectionPoolEntry::Connecting(
1059                        cx.foreground_executor()
1060                            .spawn({
1061                                let connection = connection.clone();
1062                                async move { Ok(connection.clone()) }
1063                            })
1064                            .shared(),
1065                    ),
1066                );
1067            })
1068        });
1069
1070        (opts, server_client)
1071    }
1072
1073    #[cfg(any(test, feature = "test-support"))]
1074    pub async fn fake_client(
1075        opts: SshConnectionOptions,
1076        client_cx: &mut gpui::TestAppContext,
1077    ) -> Model<Self> {
1078        let (_tx, rx) = oneshot::channel();
1079        client_cx
1080            .update(|cx| {
1081                Self::new(
1082                    ConnectionIdentifier::Setup,
1083                    opts,
1084                    rx,
1085                    Arc::new(fake::Delegate),
1086                    cx,
1087                )
1088            })
1089            .await
1090            .unwrap()
1091            .unwrap()
1092    }
1093}
1094
1095enum ConnectionPoolEntry {
1096    Connecting(Shared<Task<Result<Arc<dyn RemoteConnection>, Arc<anyhow::Error>>>>),
1097    Connected(Weak<dyn RemoteConnection>),
1098}
1099
1100#[derive(Default)]
1101struct ConnectionPool {
1102    connections: HashMap<SshConnectionOptions, ConnectionPoolEntry>,
1103}
1104
1105impl Global for ConnectionPool {}
1106
1107impl ConnectionPool {
1108    pub fn connect(
1109        &mut self,
1110        opts: SshConnectionOptions,
1111        delegate: &Arc<dyn SshClientDelegate>,
1112        cx: &mut AppContext,
1113    ) -> Shared<Task<Result<Arc<dyn RemoteConnection>, Arc<anyhow::Error>>>> {
1114        let connection = self.connections.get(&opts);
1115        match connection {
1116            Some(ConnectionPoolEntry::Connecting(task)) => {
1117                let delegate = delegate.clone();
1118                cx.spawn(|mut cx| async move {
1119                    delegate.set_status(Some("Waiting for existing connection attempt"), &mut cx);
1120                })
1121                .detach();
1122                return task.clone();
1123            }
1124            Some(ConnectionPoolEntry::Connected(ssh)) => {
1125                if let Some(ssh) = ssh.upgrade() {
1126                    if !ssh.has_been_killed() {
1127                        return Task::ready(Ok(ssh)).shared();
1128                    }
1129                }
1130                self.connections.remove(&opts);
1131            }
1132            None => {}
1133        }
1134
1135        let task = cx
1136            .spawn({
1137                let opts = opts.clone();
1138                let delegate = delegate.clone();
1139                |mut cx| async move {
1140                    let connection = SshRemoteConnection::new(opts.clone(), delegate, &mut cx)
1141                        .await
1142                        .map(|connection| Arc::new(connection) as Arc<dyn RemoteConnection>);
1143
1144                    cx.update_global(|pool: &mut Self, _| {
1145                        debug_assert!(matches!(
1146                            pool.connections.get(&opts),
1147                            Some(ConnectionPoolEntry::Connecting(_))
1148                        ));
1149                        match connection {
1150                            Ok(connection) => {
1151                                pool.connections.insert(
1152                                    opts.clone(),
1153                                    ConnectionPoolEntry::Connected(Arc::downgrade(&connection)),
1154                                );
1155                                Ok(connection)
1156                            }
1157                            Err(error) => {
1158                                pool.connections.remove(&opts);
1159                                Err(Arc::new(error))
1160                            }
1161                        }
1162                    })?
1163                }
1164            })
1165            .shared();
1166
1167        self.connections
1168            .insert(opts.clone(), ConnectionPoolEntry::Connecting(task.clone()));
1169        task
1170    }
1171}
1172
1173impl From<SshRemoteClient> for AnyProtoClient {
1174    fn from(client: SshRemoteClient) -> Self {
1175        AnyProtoClient::new(client.client.clone())
1176    }
1177}
1178
1179#[async_trait(?Send)]
1180trait RemoteConnection: Send + Sync {
1181    #[allow(clippy::too_many_arguments)]
1182    fn start_proxy(
1183        &self,
1184        unique_identifier: String,
1185        reconnect: bool,
1186        incoming_tx: UnboundedSender<Envelope>,
1187        outgoing_rx: UnboundedReceiver<Envelope>,
1188        connection_activity_tx: Sender<()>,
1189        delegate: Arc<dyn SshClientDelegate>,
1190        cx: &mut AsyncAppContext,
1191    ) -> Task<Result<i32>>;
1192    async fn kill(&self) -> Result<()>;
1193    fn has_been_killed(&self) -> bool;
1194    fn ssh_args(&self) -> Vec<String>;
1195    fn connection_options(&self) -> SshConnectionOptions;
1196
1197    #[cfg(any(test, feature = "test-support"))]
1198    fn simulate_disconnect(&self, _: &AsyncAppContext) {}
1199}
1200
1201struct SshRemoteConnection {
1202    socket: SshSocket,
1203    master_process: Mutex<Option<process::Child>>,
1204    remote_binary_path: Option<PathBuf>,
1205    _temp_dir: TempDir,
1206}
1207
1208#[async_trait(?Send)]
1209impl RemoteConnection for SshRemoteConnection {
1210    async fn kill(&self) -> Result<()> {
1211        let Some(mut process) = self.master_process.lock().take() else {
1212            return Ok(());
1213        };
1214        process.kill().ok();
1215        process.status().await?;
1216        Ok(())
1217    }
1218
1219    fn has_been_killed(&self) -> bool {
1220        self.master_process.lock().is_none()
1221    }
1222
1223    fn ssh_args(&self) -> Vec<String> {
1224        self.socket.ssh_args()
1225    }
1226
1227    fn connection_options(&self) -> SshConnectionOptions {
1228        self.socket.connection_options.clone()
1229    }
1230    fn start_proxy(
1231        &self,
1232        unique_identifier: String,
1233        reconnect: bool,
1234        incoming_tx: UnboundedSender<Envelope>,
1235        outgoing_rx: UnboundedReceiver<Envelope>,
1236        connection_activity_tx: Sender<()>,
1237        delegate: Arc<dyn SshClientDelegate>,
1238        cx: &mut AsyncAppContext,
1239    ) -> Task<Result<i32>> {
1240        delegate.set_status(Some("Starting proxy"), cx);
1241
1242        let Some(remote_binary_path) = self.remote_binary_path.clone() else {
1243            return Task::ready(Err(anyhow!("Remote binary path not set")));
1244        };
1245
1246        let mut start_proxy_command = shell_script!(
1247            "exec {binary_path} proxy --identifier {identifier}",
1248            binary_path = &remote_binary_path.to_string_lossy(),
1249            identifier = &unique_identifier,
1250        );
1251
1252        if let Some(rust_log) = std::env::var("RUST_LOG").ok() {
1253            start_proxy_command = format!(
1254                "RUST_LOG={} {}",
1255                shlex::try_quote(&rust_log).unwrap(),
1256                start_proxy_command
1257            )
1258        }
1259        if let Some(rust_backtrace) = std::env::var("RUST_BACKTRACE").ok() {
1260            start_proxy_command = format!(
1261                "RUST_BACKTRACE={} {}",
1262                shlex::try_quote(&rust_backtrace).unwrap(),
1263                start_proxy_command
1264            )
1265        }
1266        if reconnect {
1267            start_proxy_command.push_str(" --reconnect");
1268        }
1269
1270        let ssh_proxy_process = match self
1271            .socket
1272            .ssh_command("sh", &["-c", &start_proxy_command])
1273            // IMPORTANT: we kill this process when we drop the task that uses it.
1274            .kill_on_drop(true)
1275            .spawn()
1276        {
1277            Ok(process) => process,
1278            Err(error) => {
1279                return Task::ready(Err(anyhow!("failed to spawn remote server: {}", error)))
1280            }
1281        };
1282
1283        Self::multiplex(
1284            ssh_proxy_process,
1285            incoming_tx,
1286            outgoing_rx,
1287            connection_activity_tx,
1288            &cx,
1289        )
1290    }
1291}
1292
1293impl SshRemoteConnection {
1294    #[cfg(not(unix))]
1295    async fn new(
1296        _connection_options: SshConnectionOptions,
1297        _delegate: Arc<dyn SshClientDelegate>,
1298        _cx: &mut AsyncAppContext,
1299    ) -> Result<Self> {
1300        Err(anyhow!("ssh is not supported on this platform"))
1301    }
1302
1303    #[cfg(unix)]
1304    async fn new(
1305        connection_options: SshConnectionOptions,
1306        delegate: Arc<dyn SshClientDelegate>,
1307        cx: &mut AsyncAppContext,
1308    ) -> Result<Self> {
1309        use futures::AsyncWriteExt as _;
1310        use futures::{io::BufReader, AsyncBufReadExt as _};
1311        use smol::net::unix::UnixStream;
1312        use smol::{fs::unix::PermissionsExt as _, net::unix::UnixListener};
1313        use util::ResultExt as _;
1314
1315        delegate.set_status(Some("Connecting"), cx);
1316
1317        let url = connection_options.ssh_url();
1318        let temp_dir = tempfile::Builder::new()
1319            .prefix("zed-ssh-session")
1320            .tempdir()?;
1321
1322        // Create a domain socket listener to handle requests from the askpass program.
1323        let askpass_socket = temp_dir.path().join("askpass.sock");
1324        let (askpass_opened_tx, askpass_opened_rx) = oneshot::channel::<()>();
1325        let listener =
1326            UnixListener::bind(&askpass_socket).context("failed to create askpass socket")?;
1327
1328        let (askpass_kill_master_tx, askpass_kill_master_rx) = oneshot::channel::<UnixStream>();
1329        let mut kill_tx = Some(askpass_kill_master_tx);
1330
1331        let askpass_task = cx.spawn({
1332            let delegate = delegate.clone();
1333            |mut cx| async move {
1334                let mut askpass_opened_tx = Some(askpass_opened_tx);
1335
1336                while let Ok((mut stream, _)) = listener.accept().await {
1337                    if let Some(askpass_opened_tx) = askpass_opened_tx.take() {
1338                        askpass_opened_tx.send(()).ok();
1339                    }
1340                    let mut buffer = Vec::new();
1341                    let mut reader = BufReader::new(&mut stream);
1342                    if reader.read_until(b'\0', &mut buffer).await.is_err() {
1343                        buffer.clear();
1344                    }
1345                    let password_prompt = String::from_utf8_lossy(&buffer);
1346                    if let Some(password) = delegate
1347                        .ask_password(password_prompt.to_string(), &mut cx)
1348                        .await
1349                        .context("failed to get ssh password")
1350                        .and_then(|p| p)
1351                        .log_err()
1352                    {
1353                        stream.write_all(password.as_bytes()).await.log_err();
1354                    } else {
1355                        if let Some(kill_tx) = kill_tx.take() {
1356                            kill_tx.send(stream).log_err();
1357                            break;
1358                        }
1359                    }
1360                }
1361            }
1362        });
1363
1364        // Create an askpass script that communicates back to this process.
1365        let askpass_script = format!(
1366            "{shebang}\n{print_args} | nc -U {askpass_socket} 2> /dev/null \n",
1367            askpass_socket = askpass_socket.display(),
1368            print_args = "printf '%s\\0' \"$@\"",
1369            shebang = "#!/bin/sh",
1370        );
1371        let askpass_script_path = temp_dir.path().join("askpass.sh");
1372        fs::write(&askpass_script_path, askpass_script).await?;
1373        fs::set_permissions(&askpass_script_path, std::fs::Permissions::from_mode(0o755)).await?;
1374
1375        // Start the master SSH process, which does not do anything except for establish
1376        // the connection and keep it open, allowing other ssh commands to reuse it
1377        // via a control socket.
1378        let socket_path = temp_dir.path().join("ssh.sock");
1379
1380        let mut master_process = process::Command::new("ssh")
1381            .stdin(Stdio::null())
1382            .stdout(Stdio::piped())
1383            .stderr(Stdio::piped())
1384            .env("SSH_ASKPASS_REQUIRE", "force")
1385            .env("SSH_ASKPASS", &askpass_script_path)
1386            .args(connection_options.additional_args().unwrap_or(&Vec::new()))
1387            .args([
1388                "-N",
1389                "-o",
1390                "ControlPersist=no",
1391                "-o",
1392                "ControlMaster=yes",
1393                "-o",
1394            ])
1395            .arg(format!("ControlPath={}", socket_path.display()))
1396            .arg(&url)
1397            .kill_on_drop(true)
1398            .spawn()?;
1399
1400        // Wait for this ssh process to close its stdout, indicating that authentication
1401        // has completed.
1402        let mut stdout = master_process.stdout.take().unwrap();
1403        let mut output = Vec::new();
1404        let connection_timeout = Duration::from_secs(10);
1405
1406        let result = select_biased! {
1407            _ = askpass_opened_rx.fuse() => {
1408                select_biased! {
1409                    stream = askpass_kill_master_rx.fuse() => {
1410                        master_process.kill().ok();
1411                        drop(stream);
1412                        Err(anyhow!("SSH connection canceled"))
1413                    }
1414                    // If the askpass script has opened, that means the user is typing
1415                    // their password, in which case we don't want to timeout anymore,
1416                    // since we know a connection has been established.
1417                    result = stdout.read_to_end(&mut output).fuse() => {
1418                        result?;
1419                        Ok(())
1420                    }
1421                }
1422            }
1423            _ = stdout.read_to_end(&mut output).fuse() => {
1424                Ok(())
1425            }
1426            _ = futures::FutureExt::fuse(smol::Timer::after(connection_timeout)) => {
1427                Err(anyhow!("Exceeded {:?} timeout trying to connect to host", connection_timeout))
1428            }
1429        };
1430
1431        if let Err(e) = result {
1432            return Err(e.context("Failed to connect to host"));
1433        }
1434
1435        drop(askpass_task);
1436
1437        if master_process.try_status()?.is_some() {
1438            output.clear();
1439            let mut stderr = master_process.stderr.take().unwrap();
1440            stderr.read_to_end(&mut output).await?;
1441
1442            let error_message = format!(
1443                "failed to connect: {}",
1444                String::from_utf8_lossy(&output).trim()
1445            );
1446            Err(anyhow!(error_message))?;
1447        }
1448
1449        let socket = SshSocket {
1450            connection_options,
1451            socket_path,
1452        };
1453
1454        let mut this = Self {
1455            socket,
1456            master_process: Mutex::new(Some(master_process)),
1457            _temp_dir: temp_dir,
1458            remote_binary_path: None,
1459        };
1460
1461        let (release_channel, version, commit) = cx.update(|cx| {
1462            (
1463                ReleaseChannel::global(cx),
1464                AppVersion::global(cx),
1465                AppCommitSha::try_global(cx),
1466            )
1467        })?;
1468        this.remote_binary_path = Some(
1469            this.ensure_server_binary(&delegate, release_channel, version, commit, cx)
1470                .await?,
1471        );
1472
1473        Ok(this)
1474    }
1475
1476    async fn platform(&self) -> Result<SshPlatform> {
1477        let uname = self.socket.run_command("uname", &["-sm"]).await?;
1478        let Some((os, arch)) = uname.split_once(" ") else {
1479            Err(anyhow!("unknown uname: {uname:?}"))?
1480        };
1481
1482        let os = match os.trim() {
1483            "Darwin" => "macos",
1484            "Linux" => "linux",
1485            _ => Err(anyhow!("unknown uname os {os:?}"))?,
1486        };
1487        let arch = if arch.starts_with("arm") || arch.starts_with("aarch64") {
1488            "aarch64"
1489        } else if arch.starts_with("x86") || arch.starts_with("i686") {
1490            "x86_64"
1491        } else {
1492            Err(anyhow!("unknown uname architecture {arch:?}"))?
1493        };
1494
1495        Ok(SshPlatform { os, arch })
1496    }
1497
1498    fn multiplex(
1499        mut ssh_proxy_process: Child,
1500        incoming_tx: UnboundedSender<Envelope>,
1501        mut outgoing_rx: UnboundedReceiver<Envelope>,
1502        mut connection_activity_tx: Sender<()>,
1503        cx: &AsyncAppContext,
1504    ) -> Task<Result<i32>> {
1505        let mut child_stderr = ssh_proxy_process.stderr.take().unwrap();
1506        let mut child_stdout = ssh_proxy_process.stdout.take().unwrap();
1507        let mut child_stdin = ssh_proxy_process.stdin.take().unwrap();
1508
1509        let mut stdin_buffer = Vec::new();
1510        let mut stdout_buffer = Vec::new();
1511        let mut stderr_buffer = Vec::new();
1512        let mut stderr_offset = 0;
1513
1514        let stdin_task = cx.background_executor().spawn(async move {
1515            while let Some(outgoing) = outgoing_rx.next().await {
1516                write_message(&mut child_stdin, &mut stdin_buffer, outgoing).await?;
1517            }
1518            anyhow::Ok(())
1519        });
1520
1521        let stdout_task = cx.background_executor().spawn({
1522            let mut connection_activity_tx = connection_activity_tx.clone();
1523            async move {
1524                loop {
1525                    stdout_buffer.resize(MESSAGE_LEN_SIZE, 0);
1526                    let len = child_stdout.read(&mut stdout_buffer).await?;
1527
1528                    if len == 0 {
1529                        return anyhow::Ok(());
1530                    }
1531
1532                    if len < MESSAGE_LEN_SIZE {
1533                        child_stdout.read_exact(&mut stdout_buffer[len..]).await?;
1534                    }
1535
1536                    let message_len = message_len_from_buffer(&stdout_buffer);
1537                    let envelope =
1538                        read_message_with_len(&mut child_stdout, &mut stdout_buffer, message_len)
1539                            .await?;
1540                    connection_activity_tx.try_send(()).ok();
1541                    incoming_tx.unbounded_send(envelope).ok();
1542                }
1543            }
1544        });
1545
1546        let stderr_task: Task<anyhow::Result<()>> = cx.background_executor().spawn(async move {
1547            loop {
1548                stderr_buffer.resize(stderr_offset + 1024, 0);
1549
1550                let len = child_stderr
1551                    .read(&mut stderr_buffer[stderr_offset..])
1552                    .await?;
1553                if len == 0 {
1554                    return anyhow::Ok(());
1555                }
1556
1557                stderr_offset += len;
1558                let mut start_ix = 0;
1559                while let Some(ix) = stderr_buffer[start_ix..stderr_offset]
1560                    .iter()
1561                    .position(|b| b == &b'\n')
1562                {
1563                    let line_ix = start_ix + ix;
1564                    let content = &stderr_buffer[start_ix..line_ix];
1565                    start_ix = line_ix + 1;
1566                    if let Ok(record) = serde_json::from_slice::<LogRecord>(content) {
1567                        record.log(log::logger())
1568                    } else {
1569                        eprintln!("(remote) {}", String::from_utf8_lossy(content));
1570                    }
1571                }
1572                stderr_buffer.drain(0..start_ix);
1573                stderr_offset -= start_ix;
1574
1575                connection_activity_tx.try_send(()).ok();
1576            }
1577        });
1578
1579        cx.spawn(|_| async move {
1580            let result = futures::select! {
1581                result = stdin_task.fuse() => {
1582                    result.context("stdin")
1583                }
1584                result = stdout_task.fuse() => {
1585                    result.context("stdout")
1586                }
1587                result = stderr_task.fuse() => {
1588                    result.context("stderr")
1589                }
1590            };
1591
1592            let status = ssh_proxy_process.status().await?.code().unwrap_or(1);
1593            match result {
1594                Ok(_) => Ok(status),
1595                Err(error) => Err(error),
1596            }
1597        })
1598    }
1599
1600    #[allow(unused)]
1601    async fn ensure_server_binary(
1602        &self,
1603        delegate: &Arc<dyn SshClientDelegate>,
1604        release_channel: ReleaseChannel,
1605        version: SemanticVersion,
1606        commit: Option<AppCommitSha>,
1607        cx: &mut AsyncAppContext,
1608    ) -> Result<PathBuf> {
1609        let version_str = match release_channel {
1610            ReleaseChannel::Nightly => {
1611                let commit = commit.map(|s| s.0.to_string()).unwrap_or_default();
1612
1613                format!("{}-{}", version, commit)
1614            }
1615            ReleaseChannel::Dev => "build".to_string(),
1616            _ => version.to_string(),
1617        };
1618        let binary_name = format!(
1619            "zed-remote-server-{}-{}",
1620            release_channel.dev_name(),
1621            version_str
1622        );
1623        let dst_path = paths::remote_server_dir_relative().join(binary_name);
1624        let tmp_path_gz = PathBuf::from(format!(
1625            "{}-download-{}.gz",
1626            dst_path.to_string_lossy(),
1627            std::process::id()
1628        ));
1629
1630        #[cfg(debug_assertions)]
1631        if std::env::var("ZED_BUILD_REMOTE_SERVER").is_ok() {
1632            let src_path = self
1633                .build_local(self.platform().await?, delegate, cx)
1634                .await?;
1635            self.upload_local_server_binary(&src_path, &tmp_path_gz, delegate, cx)
1636                .await?;
1637            self.extract_server_binary(&dst_path, &tmp_path_gz, delegate, cx)
1638                .await?;
1639            return Ok(dst_path);
1640        }
1641
1642        if self
1643            .socket
1644            .run_command(&dst_path.to_string_lossy(), &["version"])
1645            .await
1646            .is_ok()
1647        {
1648            return Ok(dst_path);
1649        }
1650
1651        let wanted_version = cx.update(|cx| match release_channel {
1652            ReleaseChannel::Nightly => Ok(None),
1653            ReleaseChannel::Dev => {
1654                anyhow::bail!(
1655                    "ZED_BUILD_REMOTE_SERVER is not set and no remote server exists at ({:?})",
1656                    dst_path
1657                )
1658            }
1659            _ => Ok(Some(AppVersion::global(cx))),
1660        })??;
1661
1662        let platform = self.platform().await?;
1663
1664        if !self.socket.connection_options.upload_binary_over_ssh {
1665            if let Some((url, body)) = delegate
1666                .get_download_params(platform, release_channel, wanted_version, cx)
1667                .await?
1668            {
1669                match self
1670                    .download_binary_on_server(&url, &body, &tmp_path_gz, delegate, cx)
1671                    .await
1672                {
1673                    Ok(_) => {
1674                        self.extract_server_binary(&dst_path, &tmp_path_gz, delegate, cx)
1675                            .await?;
1676                        return Ok(dst_path);
1677                    }
1678                    Err(e) => {
1679                        log::error!(
1680                            "Failed to download binary on server, attempting to upload server: {}",
1681                            e
1682                        )
1683                    }
1684                }
1685            }
1686        }
1687
1688        let src_path = delegate
1689            .download_server_binary_locally(platform, release_channel, wanted_version, cx)
1690            .await?;
1691        self.upload_local_server_binary(&src_path, &tmp_path_gz, delegate, cx)
1692            .await?;
1693        self.extract_server_binary(&dst_path, &tmp_path_gz, delegate, cx)
1694            .await?;
1695        return Ok(dst_path);
1696    }
1697
1698    async fn download_binary_on_server(
1699        &self,
1700        url: &str,
1701        body: &str,
1702        tmp_path_gz: &Path,
1703        delegate: &Arc<dyn SshClientDelegate>,
1704        cx: &mut AsyncAppContext,
1705    ) -> Result<()> {
1706        if let Some(parent) = tmp_path_gz.parent() {
1707            self.socket
1708                .run_command("mkdir", &["-p", &parent.to_string_lossy()])
1709                .await?;
1710        }
1711
1712        delegate.set_status(Some("Downloading remote development server on host"), cx);
1713
1714        match self
1715            .socket
1716            .run_command(
1717                "curl",
1718                &[
1719                    "-f",
1720                    "-L",
1721                    "-X",
1722                    "GET",
1723                    "-H",
1724                    "Content-Type: application/json",
1725                    "-d",
1726                    &body,
1727                    &url,
1728                    "-o",
1729                    &tmp_path_gz.to_string_lossy(),
1730                ],
1731            )
1732            .await
1733        {
1734            Ok(_) => {}
1735            Err(e) => {
1736                if self.socket.run_command("which", &["curl"]).await.is_ok() {
1737                    return Err(e);
1738                }
1739
1740                match self
1741                    .socket
1742                    .run_command(
1743                        "wget",
1744                        &[
1745                            "--max-redirect=5",
1746                            "--method=GET",
1747                            "--header=Content-Type: application/json",
1748                            "--body-data",
1749                            &body,
1750                            &url,
1751                            "-O",
1752                            &tmp_path_gz.to_string_lossy(),
1753                        ],
1754                    )
1755                    .await
1756                {
1757                    Ok(_) => {}
1758                    Err(e) => {
1759                        if self.socket.run_command("which", &["wget"]).await.is_ok() {
1760                            return Err(e);
1761                        } else {
1762                            anyhow::bail!("Neither curl nor wget is available");
1763                        }
1764                    }
1765                }
1766            }
1767        }
1768
1769        Ok(())
1770    }
1771
1772    async fn upload_local_server_binary(
1773        &self,
1774        src_path: &Path,
1775        tmp_path_gz: &Path,
1776        delegate: &Arc<dyn SshClientDelegate>,
1777        cx: &mut AsyncAppContext,
1778    ) -> Result<()> {
1779        if let Some(parent) = tmp_path_gz.parent() {
1780            self.socket
1781                .run_command("mkdir", &["-p", &parent.to_string_lossy()])
1782                .await?;
1783        }
1784
1785        let src_stat = fs::metadata(&src_path).await?;
1786        let size = src_stat.len();
1787
1788        let t0 = Instant::now();
1789        delegate.set_status(Some("Uploading remote development server"), cx);
1790        log::info!(
1791            "uploading remote development server to {:?} ({}kb)",
1792            tmp_path_gz,
1793            size / 1024
1794        );
1795        self.upload_file(&src_path, &tmp_path_gz)
1796            .await
1797            .context("failed to upload server binary")?;
1798        log::info!("uploaded remote development server in {:?}", t0.elapsed());
1799        Ok(())
1800    }
1801
1802    async fn extract_server_binary(
1803        &self,
1804        dst_path: &Path,
1805        tmp_path_gz: &Path,
1806        delegate: &Arc<dyn SshClientDelegate>,
1807        cx: &mut AsyncAppContext,
1808    ) -> Result<()> {
1809        delegate.set_status(Some("Extracting remote development server"), cx);
1810        let server_mode = 0o755;
1811
1812        let script = shell_script!(
1813            "gunzip -f {tmp_path_gz} && chmod {server_mode} {tmp_path} && mv {tmp_path} {dst_path}",
1814            tmp_path_gz = &tmp_path_gz.to_string_lossy(),
1815            tmp_path = &tmp_path_gz.to_string_lossy().strip_suffix(".gz").unwrap(),
1816            server_mode = &format!("{:o}", server_mode),
1817            dst_path = &dst_path.to_string_lossy()
1818        );
1819        self.socket.run_command("sh", &["-c", &script]).await?;
1820        Ok(())
1821    }
1822
1823    async fn upload_file(&self, src_path: &Path, dest_path: &Path) -> Result<()> {
1824        log::debug!("uploading file {:?} to {:?}", src_path, dest_path);
1825        let mut command = process::Command::new("scp");
1826        let output = self
1827            .socket
1828            .ssh_options(&mut command)
1829            .args(
1830                self.socket
1831                    .connection_options
1832                    .port
1833                    .map(|port| vec!["-P".to_string(), port.to_string()])
1834                    .unwrap_or_default(),
1835            )
1836            .arg(src_path)
1837            .arg(format!(
1838                "{}:{}",
1839                self.socket.connection_options.scp_url(),
1840                dest_path.display()
1841            ))
1842            .output()
1843            .await?;
1844
1845        if output.status.success() {
1846            Ok(())
1847        } else {
1848            Err(anyhow!(
1849                "failed to upload file {} -> {}: {}",
1850                src_path.display(),
1851                dest_path.display(),
1852                String::from_utf8_lossy(&output.stderr)
1853            ))
1854        }
1855    }
1856
1857    #[cfg(debug_assertions)]
1858    async fn build_local(
1859        &self,
1860        platform: SshPlatform,
1861        delegate: &Arc<dyn SshClientDelegate>,
1862        cx: &mut AsyncAppContext,
1863    ) -> Result<PathBuf> {
1864        use smol::process::{Command, Stdio};
1865
1866        async fn run_cmd(command: &mut Command) -> Result<()> {
1867            let output = command
1868                .kill_on_drop(true)
1869                .stderr(Stdio::inherit())
1870                .output()
1871                .await?;
1872            if !output.status.success() {
1873                Err(anyhow!("Failed to run command: {:?}", command))?;
1874            }
1875            Ok(())
1876        }
1877
1878        if platform.arch == std::env::consts::ARCH && platform.os == std::env::consts::OS {
1879            delegate.set_status(Some("Building remote server binary from source"), cx);
1880            log::info!("building remote server binary from source");
1881            run_cmd(Command::new("cargo").args([
1882                "build",
1883                "--package",
1884                "remote_server",
1885                "--features",
1886                "debug-embed",
1887                "--target-dir",
1888                "target/remote_server",
1889            ]))
1890            .await?;
1891
1892            delegate.set_status(Some("Compressing binary"), cx);
1893
1894            run_cmd(Command::new("gzip").args([
1895                "-9",
1896                "-f",
1897                "target/remote_server/debug/remote_server",
1898            ]))
1899            .await?;
1900
1901            let path = std::env::current_dir()?.join("target/remote_server/debug/remote_server.gz");
1902            return Ok(path);
1903        }
1904        let Some(triple) = platform.triple() else {
1905            anyhow::bail!("can't cross compile for: {:?}", platform);
1906        };
1907        smol::fs::create_dir_all("target/remote_server").await?;
1908
1909        delegate.set_status(Some("Installing cross.rs for cross-compilation"), cx);
1910        log::info!("installing cross");
1911        run_cmd(Command::new("cargo").args([
1912            "install",
1913            "cross",
1914            "--git",
1915            "https://github.com/cross-rs/cross",
1916        ]))
1917        .await?;
1918
1919        delegate.set_status(
1920            Some(&format!(
1921                "Building remote server binary from source for {} with Docker",
1922                &triple
1923            )),
1924            cx,
1925        );
1926        log::info!("building remote server binary from source for {}", &triple);
1927        run_cmd(
1928            Command::new("cross")
1929                .args([
1930                    "build",
1931                    "--package",
1932                    "remote_server",
1933                    "--features",
1934                    "debug-embed",
1935                    "--target-dir",
1936                    "target/remote_server",
1937                    "--target",
1938                    &triple,
1939                ])
1940                .env(
1941                    "CROSS_CONTAINER_OPTS",
1942                    "--mount type=bind,src=./target,dst=/app/target",
1943                ),
1944        )
1945        .await?;
1946
1947        delegate.set_status(Some("Compressing binary"), cx);
1948
1949        run_cmd(Command::new("gzip").args([
1950            "-9",
1951            "-f",
1952            &format!("target/remote_server/{}/debug/remote_server", triple),
1953        ]))
1954        .await?;
1955
1956        let path = std::env::current_dir()?.join(format!(
1957            "target/remote_server/{}/debug/remote_server.gz",
1958            triple
1959        ));
1960
1961        return Ok(path);
1962    }
1963}
1964
1965type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>;
1966
1967pub struct ChannelClient {
1968    next_message_id: AtomicU32,
1969    outgoing_tx: Mutex<mpsc::UnboundedSender<Envelope>>,
1970    buffer: Mutex<VecDeque<Envelope>>,
1971    response_channels: ResponseChannels,
1972    message_handlers: Mutex<ProtoMessageHandlerSet>,
1973    max_received: AtomicU32,
1974    name: &'static str,
1975    task: Mutex<Task<Result<()>>>,
1976}
1977
1978impl ChannelClient {
1979    pub fn new(
1980        incoming_rx: mpsc::UnboundedReceiver<Envelope>,
1981        outgoing_tx: mpsc::UnboundedSender<Envelope>,
1982        cx: &AppContext,
1983        name: &'static str,
1984    ) -> Arc<Self> {
1985        Arc::new_cyclic(|this| Self {
1986            outgoing_tx: Mutex::new(outgoing_tx),
1987            next_message_id: AtomicU32::new(0),
1988            max_received: AtomicU32::new(0),
1989            response_channels: ResponseChannels::default(),
1990            message_handlers: Default::default(),
1991            buffer: Mutex::new(VecDeque::new()),
1992            name,
1993            task: Mutex::new(Self::start_handling_messages(
1994                this.clone(),
1995                incoming_rx,
1996                &cx.to_async(),
1997            )),
1998        })
1999    }
2000
2001    fn start_handling_messages(
2002        this: Weak<Self>,
2003        mut incoming_rx: mpsc::UnboundedReceiver<Envelope>,
2004        cx: &AsyncAppContext,
2005    ) -> Task<Result<()>> {
2006        cx.spawn(|cx| async move {
2007            let peer_id = PeerId { owner_id: 0, id: 0 };
2008            while let Some(incoming) = incoming_rx.next().await {
2009                let Some(this) = this.upgrade() else {
2010                    return anyhow::Ok(());
2011                };
2012                if let Some(ack_id) = incoming.ack_id {
2013                    let mut buffer = this.buffer.lock();
2014                    while buffer.front().is_some_and(|msg| msg.id <= ack_id) {
2015                        buffer.pop_front();
2016                    }
2017                }
2018                if let Some(proto::envelope::Payload::FlushBufferedMessages(_)) = &incoming.payload
2019                {
2020                    log::debug!(
2021                        "{}:ssh message received. name:FlushBufferedMessages",
2022                        this.name
2023                    );
2024                    {
2025                        let buffer = this.buffer.lock();
2026                        for envelope in buffer.iter() {
2027                            this.outgoing_tx
2028                                .lock()
2029                                .unbounded_send(envelope.clone())
2030                                .ok();
2031                        }
2032                    }
2033                    let mut envelope = proto::Ack {}.into_envelope(0, Some(incoming.id), None);
2034                    envelope.id = this.next_message_id.fetch_add(1, SeqCst);
2035                    this.outgoing_tx.lock().unbounded_send(envelope).ok();
2036                    continue;
2037                }
2038
2039                this.max_received.store(incoming.id, SeqCst);
2040
2041                if let Some(request_id) = incoming.responding_to {
2042                    let request_id = MessageId(request_id);
2043                    let sender = this.response_channels.lock().remove(&request_id);
2044                    if let Some(sender) = sender {
2045                        let (tx, rx) = oneshot::channel();
2046                        if incoming.payload.is_some() {
2047                            sender.send((incoming, tx)).ok();
2048                        }
2049                        rx.await.ok();
2050                    }
2051                } else if let Some(envelope) =
2052                    build_typed_envelope(peer_id, Instant::now(), incoming)
2053                {
2054                    let type_name = envelope.payload_type_name();
2055                    if let Some(future) = ProtoMessageHandlerSet::handle_message(
2056                        &this.message_handlers,
2057                        envelope,
2058                        this.clone().into(),
2059                        cx.clone(),
2060                    ) {
2061                        log::debug!("{}:ssh message received. name:{type_name}", this.name);
2062                        cx.foreground_executor()
2063                            .spawn(async move {
2064                                match future.await {
2065                                    Ok(_) => {
2066                                        log::debug!(
2067                                            "{}:ssh message handled. name:{type_name}",
2068                                            this.name
2069                                        );
2070                                    }
2071                                    Err(error) => {
2072                                        log::error!(
2073                                            "{}:error handling message. type:{}, error:{}",
2074                                            this.name,
2075                                            type_name,
2076                                            format!("{error:#}").lines().fold(
2077                                                String::new(),
2078                                                |mut message, line| {
2079                                                    if !message.is_empty() {
2080                                                        message.push(' ');
2081                                                    }
2082                                                    message.push_str(line);
2083                                                    message
2084                                                }
2085                                            )
2086                                        );
2087                                    }
2088                                }
2089                            })
2090                            .detach()
2091                    } else {
2092                        log::error!("{}:unhandled ssh message name:{type_name}", this.name);
2093                    }
2094                }
2095            }
2096            anyhow::Ok(())
2097        })
2098    }
2099
2100    pub fn reconnect(
2101        self: &Arc<Self>,
2102        incoming_rx: UnboundedReceiver<Envelope>,
2103        outgoing_tx: UnboundedSender<Envelope>,
2104        cx: &AsyncAppContext,
2105    ) {
2106        *self.outgoing_tx.lock() = outgoing_tx;
2107        *self.task.lock() = Self::start_handling_messages(Arc::downgrade(self), incoming_rx, cx);
2108    }
2109
2110    pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Model<E>) {
2111        let id = (TypeId::of::<E>(), remote_id);
2112
2113        let mut message_handlers = self.message_handlers.lock();
2114        if message_handlers
2115            .entities_by_type_and_remote_id
2116            .contains_key(&id)
2117        {
2118            panic!("already subscribed to entity");
2119        }
2120
2121        message_handlers.entities_by_type_and_remote_id.insert(
2122            id,
2123            EntityMessageSubscriber::Entity {
2124                handle: entity.downgrade().into(),
2125            },
2126        );
2127    }
2128
2129    pub fn request<T: RequestMessage>(
2130        &self,
2131        payload: T,
2132    ) -> impl 'static + Future<Output = Result<T::Response>> {
2133        self.request_internal(payload, true)
2134    }
2135
2136    fn request_internal<T: RequestMessage>(
2137        &self,
2138        payload: T,
2139        use_buffer: bool,
2140    ) -> impl 'static + Future<Output = Result<T::Response>> {
2141        log::debug!("ssh request start. name:{}", T::NAME);
2142        let response =
2143            self.request_dynamic(payload.into_envelope(0, None, None), T::NAME, use_buffer);
2144        async move {
2145            let response = response.await?;
2146            log::debug!("ssh request finish. name:{}", T::NAME);
2147            T::Response::from_envelope(response)
2148                .ok_or_else(|| anyhow!("received a response of the wrong type"))
2149        }
2150    }
2151
2152    pub async fn resync(&self, timeout: Duration) -> Result<()> {
2153        smol::future::or(
2154            async {
2155                self.request_internal(proto::FlushBufferedMessages {}, false)
2156                    .await?;
2157
2158                for envelope in self.buffer.lock().iter() {
2159                    self.outgoing_tx
2160                        .lock()
2161                        .unbounded_send(envelope.clone())
2162                        .ok();
2163                }
2164                Ok(())
2165            },
2166            async {
2167                smol::Timer::after(timeout).await;
2168                Err(anyhow!("Timeout detected"))
2169            },
2170        )
2171        .await
2172    }
2173
2174    pub async fn ping(&self, timeout: Duration) -> Result<()> {
2175        smol::future::or(
2176            async {
2177                self.request(proto::Ping {}).await?;
2178                Ok(())
2179            },
2180            async {
2181                smol::Timer::after(timeout).await;
2182                Err(anyhow!("Timeout detected"))
2183            },
2184        )
2185        .await
2186    }
2187
2188    pub fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {
2189        log::debug!("ssh send name:{}", T::NAME);
2190        self.send_dynamic(payload.into_envelope(0, None, None))
2191    }
2192
2193    fn request_dynamic(
2194        &self,
2195        mut envelope: proto::Envelope,
2196        type_name: &'static str,
2197        use_buffer: bool,
2198    ) -> impl 'static + Future<Output = Result<proto::Envelope>> {
2199        envelope.id = self.next_message_id.fetch_add(1, SeqCst);
2200        let (tx, rx) = oneshot::channel();
2201        let mut response_channels_lock = self.response_channels.lock();
2202        response_channels_lock.insert(MessageId(envelope.id), tx);
2203        drop(response_channels_lock);
2204
2205        let result = if use_buffer {
2206            self.send_buffered(envelope)
2207        } else {
2208            self.send_unbuffered(envelope)
2209        };
2210        async move {
2211            if let Err(error) = &result {
2212                log::error!("failed to send message: {}", error);
2213                return Err(anyhow!("failed to send message: {}", error));
2214            }
2215
2216            let response = rx.await.context("connection lost")?.0;
2217            if let Some(proto::envelope::Payload::Error(error)) = &response.payload {
2218                return Err(RpcError::from_proto(error, type_name));
2219            }
2220            Ok(response)
2221        }
2222    }
2223
2224    pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {
2225        envelope.id = self.next_message_id.fetch_add(1, SeqCst);
2226        self.send_buffered(envelope)
2227    }
2228
2229    fn send_buffered(&self, mut envelope: proto::Envelope) -> Result<()> {
2230        envelope.ack_id = Some(self.max_received.load(SeqCst));
2231        self.buffer.lock().push_back(envelope.clone());
2232        // ignore errors on send (happen while we're reconnecting)
2233        // assume that the global "disconnected" overlay is sufficient.
2234        self.outgoing_tx.lock().unbounded_send(envelope).ok();
2235        Ok(())
2236    }
2237
2238    fn send_unbuffered(&self, mut envelope: proto::Envelope) -> Result<()> {
2239        envelope.ack_id = Some(self.max_received.load(SeqCst));
2240        self.outgoing_tx.lock().unbounded_send(envelope).ok();
2241        Ok(())
2242    }
2243}
2244
2245impl ProtoClient for ChannelClient {
2246    fn request(
2247        &self,
2248        envelope: proto::Envelope,
2249        request_type: &'static str,
2250    ) -> BoxFuture<'static, Result<proto::Envelope>> {
2251        self.request_dynamic(envelope, request_type, true).boxed()
2252    }
2253
2254    fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
2255        self.send_dynamic(envelope)
2256    }
2257
2258    fn send_response(&self, envelope: Envelope, _message_type: &'static str) -> anyhow::Result<()> {
2259        self.send_dynamic(envelope)
2260    }
2261
2262    fn message_handler_set(&self) -> &Mutex<ProtoMessageHandlerSet> {
2263        &self.message_handlers
2264    }
2265
2266    fn is_via_collab(&self) -> bool {
2267        false
2268    }
2269}
2270
2271#[cfg(any(test, feature = "test-support"))]
2272mod fake {
2273    use std::{path::PathBuf, sync::Arc};
2274
2275    use anyhow::Result;
2276    use async_trait::async_trait;
2277    use futures::{
2278        channel::{
2279            mpsc::{self, Sender},
2280            oneshot,
2281        },
2282        select_biased, FutureExt, SinkExt, StreamExt,
2283    };
2284    use gpui::{AsyncAppContext, SemanticVersion, Task, TestAppContext};
2285    use release_channel::ReleaseChannel;
2286    use rpc::proto::Envelope;
2287
2288    use super::{
2289        ChannelClient, RemoteConnection, SshClientDelegate, SshConnectionOptions, SshPlatform,
2290    };
2291
2292    pub(super) struct FakeRemoteConnection {
2293        pub(super) connection_options: SshConnectionOptions,
2294        pub(super) server_channel: Arc<ChannelClient>,
2295        pub(super) server_cx: SendableCx,
2296    }
2297
2298    pub(super) struct SendableCx(AsyncAppContext);
2299    impl SendableCx {
2300        // SAFETY: When run in test mode, GPUI is always single threaded.
2301        pub(super) fn new(cx: &TestAppContext) -> Self {
2302            Self(cx.to_async())
2303        }
2304
2305        // SAFETY: Enforce that we're on the main thread by requiring a valid AsyncAppContext
2306        fn get(&self, _: &AsyncAppContext) -> AsyncAppContext {
2307            self.0.clone()
2308        }
2309    }
2310
2311    // SAFETY: There is no way to access a SendableCx from a different thread, see [`SendableCx::new`] and [`SendableCx::get`]
2312    unsafe impl Send for SendableCx {}
2313    unsafe impl Sync for SendableCx {}
2314
2315    #[async_trait(?Send)]
2316    impl RemoteConnection for FakeRemoteConnection {
2317        async fn kill(&self) -> Result<()> {
2318            Ok(())
2319        }
2320
2321        fn has_been_killed(&self) -> bool {
2322            false
2323        }
2324
2325        fn ssh_args(&self) -> Vec<String> {
2326            Vec::new()
2327        }
2328
2329        fn connection_options(&self) -> SshConnectionOptions {
2330            self.connection_options.clone()
2331        }
2332
2333        fn simulate_disconnect(&self, cx: &AsyncAppContext) {
2334            let (outgoing_tx, _) = mpsc::unbounded::<Envelope>();
2335            let (_, incoming_rx) = mpsc::unbounded::<Envelope>();
2336            self.server_channel
2337                .reconnect(incoming_rx, outgoing_tx, &self.server_cx.get(&cx));
2338        }
2339
2340        fn start_proxy(
2341            &self,
2342
2343            _unique_identifier: String,
2344            _reconnect: bool,
2345            mut client_incoming_tx: mpsc::UnboundedSender<Envelope>,
2346            mut client_outgoing_rx: mpsc::UnboundedReceiver<Envelope>,
2347            mut connection_activity_tx: Sender<()>,
2348            _delegate: Arc<dyn SshClientDelegate>,
2349            cx: &mut AsyncAppContext,
2350        ) -> Task<Result<i32>> {
2351            let (mut server_incoming_tx, server_incoming_rx) = mpsc::unbounded::<Envelope>();
2352            let (server_outgoing_tx, mut server_outgoing_rx) = mpsc::unbounded::<Envelope>();
2353
2354            self.server_channel.reconnect(
2355                server_incoming_rx,
2356                server_outgoing_tx,
2357                &self.server_cx.get(cx),
2358            );
2359
2360            cx.background_executor().spawn(async move {
2361                loop {
2362                    select_biased! {
2363                        server_to_client = server_outgoing_rx.next().fuse() => {
2364                            let Some(server_to_client) = server_to_client else {
2365                                return Ok(1)
2366                            };
2367                            connection_activity_tx.try_send(()).ok();
2368                            client_incoming_tx.send(server_to_client).await.ok();
2369                        }
2370                        client_to_server = client_outgoing_rx.next().fuse() => {
2371                            let Some(client_to_server) = client_to_server else {
2372                                return Ok(1)
2373                            };
2374                            server_incoming_tx.send(client_to_server).await.ok();
2375                        }
2376                    }
2377                }
2378            })
2379        }
2380    }
2381
2382    pub(super) struct Delegate;
2383
2384    impl SshClientDelegate for Delegate {
2385        fn ask_password(
2386            &self,
2387            _: String,
2388            _: &mut AsyncAppContext,
2389        ) -> oneshot::Receiver<Result<String>> {
2390            unreachable!()
2391        }
2392
2393        fn download_server_binary_locally(
2394            &self,
2395            _: SshPlatform,
2396            _: ReleaseChannel,
2397            _: Option<SemanticVersion>,
2398            _: &mut AsyncAppContext,
2399        ) -> Task<Result<PathBuf>> {
2400            unreachable!()
2401        }
2402
2403        fn get_download_params(
2404            &self,
2405            _platform: SshPlatform,
2406            _release_channel: ReleaseChannel,
2407            _version: Option<SemanticVersion>,
2408            _cx: &mut AsyncAppContext,
2409        ) -> Task<Result<Option<(String, String)>>> {
2410            unreachable!()
2411        }
2412
2413        fn set_status(&self, _: Option<&str>, _: &mut AsyncAppContext) {}
2414    }
2415}