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