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 collections::HashMap;
  10use futures::{
  11    channel::{
  12        mpsc::{self, Sender, UnboundedReceiver, UnboundedSender},
  13        oneshot,
  14    },
  15    future::BoxFuture,
  16    select_biased, AsyncReadExt as _, AsyncWriteExt as _, Future, FutureExt as _, SinkExt,
  17    StreamExt as _,
  18};
  19use gpui::{
  20    AppContext, AsyncAppContext, Context, EventEmitter, Model, ModelContext, SemanticVersion, Task,
  21    WeakModel,
  22};
  23use parking_lot::Mutex;
  24use rpc::{
  25    proto::{self, build_typed_envelope, Envelope, EnvelopedMessage, PeerId, RequestMessage},
  26    AnyProtoClient, EntityMessageSubscriber, ProtoClient, ProtoMessageHandlerSet, RpcError,
  27};
  28use smol::{
  29    fs,
  30    process::{self, Child, Stdio},
  31};
  32use std::{
  33    any::TypeId,
  34    ffi::OsStr,
  35    fmt,
  36    ops::ControlFlow,
  37    path::{Path, PathBuf},
  38    sync::{
  39        atomic::{AtomicU32, Ordering::SeqCst},
  40        Arc,
  41    },
  42    time::{Duration, Instant},
  43};
  44use tempfile::TempDir;
  45use util::ResultExt;
  46
  47#[derive(
  48    Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, serde::Serialize, serde::Deserialize,
  49)]
  50pub struct SshProjectId(pub u64);
  51
  52#[derive(Clone)]
  53pub struct SshSocket {
  54    connection_options: SshConnectionOptions,
  55    socket_path: PathBuf,
  56}
  57
  58#[derive(Debug, Default, Clone, PartialEq, Eq)]
  59pub struct SshConnectionOptions {
  60    pub host: String,
  61    pub username: Option<String>,
  62    pub port: Option<u16>,
  63    pub password: Option<String>,
  64}
  65
  66impl SshConnectionOptions {
  67    pub fn ssh_url(&self) -> String {
  68        let mut result = String::from("ssh://");
  69        if let Some(username) = &self.username {
  70            result.push_str(username);
  71            result.push('@');
  72        }
  73        result.push_str(&self.host);
  74        if let Some(port) = self.port {
  75            result.push(':');
  76            result.push_str(&port.to_string());
  77        }
  78        result
  79    }
  80
  81    fn scp_url(&self) -> String {
  82        if let Some(username) = &self.username {
  83            format!("{}@{}", username, self.host)
  84        } else {
  85            self.host.clone()
  86        }
  87    }
  88
  89    pub fn connection_string(&self) -> String {
  90        let host = if let Some(username) = &self.username {
  91            format!("{}@{}", username, self.host)
  92        } else {
  93            self.host.clone()
  94        };
  95        if let Some(port) = &self.port {
  96            format!("{}:{}", host, port)
  97        } else {
  98            host
  99        }
 100    }
 101
 102    // Uniquely identifies dev server projects on a remote host. Needs to be
 103    // stable for the same dev server project.
 104    pub fn dev_server_identifier(&self) -> String {
 105        let mut identifier = format!("dev-server-{:?}", self.host);
 106        if let Some(username) = self.username.as_ref() {
 107            identifier.push('-');
 108            identifier.push_str(&username);
 109        }
 110        identifier
 111    }
 112}
 113
 114#[derive(Copy, Clone, Debug)]
 115pub struct SshPlatform {
 116    pub os: &'static str,
 117    pub arch: &'static str,
 118}
 119
 120impl SshPlatform {
 121    pub fn triple(&self) -> Option<String> {
 122        Some(format!(
 123            "{}-{}",
 124            self.arch,
 125            match self.os {
 126                "linux" => "unknown-linux-gnu",
 127                "macos" => "apple-darwin",
 128                _ => return None,
 129            }
 130        ))
 131    }
 132}
 133
 134pub trait SshClientDelegate: Send + Sync {
 135    fn ask_password(
 136        &self,
 137        prompt: String,
 138        cx: &mut AsyncAppContext,
 139    ) -> oneshot::Receiver<Result<String>>;
 140    fn remote_server_binary_path(
 141        &self,
 142        platform: SshPlatform,
 143        cx: &mut AsyncAppContext,
 144    ) -> Result<PathBuf>;
 145    fn get_server_binary(
 146        &self,
 147        platform: SshPlatform,
 148        cx: &mut AsyncAppContext,
 149    ) -> oneshot::Receiver<Result<(PathBuf, SemanticVersion)>>;
 150    fn set_status(&self, status: Option<&str>, cx: &mut AsyncAppContext);
 151    fn set_error(&self, error_message: String, cx: &mut AsyncAppContext);
 152}
 153
 154impl SshSocket {
 155    fn ssh_command<S: AsRef<OsStr>>(&self, program: S) -> process::Command {
 156        let mut command = process::Command::new("ssh");
 157        self.ssh_options(&mut command)
 158            .arg(self.connection_options.ssh_url())
 159            .arg(program);
 160        command
 161    }
 162
 163    fn ssh_options<'a>(&self, command: &'a mut process::Command) -> &'a mut process::Command {
 164        command
 165            .stdin(Stdio::piped())
 166            .stdout(Stdio::piped())
 167            .stderr(Stdio::piped())
 168            .args(["-o", "ControlMaster=no", "-o"])
 169            .arg(format!("ControlPath={}", self.socket_path.display()))
 170    }
 171
 172    fn ssh_args(&self) -> Vec<String> {
 173        vec![
 174            "-o".to_string(),
 175            "ControlMaster=no".to_string(),
 176            "-o".to_string(),
 177            format!("ControlPath={}", self.socket_path.display()),
 178            self.connection_options.ssh_url(),
 179        ]
 180    }
 181}
 182
 183async fn run_cmd(command: &mut process::Command) -> Result<String> {
 184    let output = command.output().await?;
 185    if output.status.success() {
 186        Ok(String::from_utf8_lossy(&output.stdout).to_string())
 187    } else {
 188        Err(anyhow!(
 189            "failed to run command: {}",
 190            String::from_utf8_lossy(&output.stderr)
 191        ))
 192    }
 193}
 194
 195struct ChannelForwarder {
 196    quit_tx: UnboundedSender<()>,
 197    forwarding_task: Task<(UnboundedSender<Envelope>, UnboundedReceiver<Envelope>)>,
 198}
 199
 200impl ChannelForwarder {
 201    fn new(
 202        mut incoming_tx: UnboundedSender<Envelope>,
 203        mut outgoing_rx: UnboundedReceiver<Envelope>,
 204        cx: &AsyncAppContext,
 205    ) -> (Self, UnboundedSender<Envelope>, UnboundedReceiver<Envelope>) {
 206        let (quit_tx, mut quit_rx) = mpsc::unbounded::<()>();
 207
 208        let (proxy_incoming_tx, mut proxy_incoming_rx) = mpsc::unbounded::<Envelope>();
 209        let (mut proxy_outgoing_tx, proxy_outgoing_rx) = mpsc::unbounded::<Envelope>();
 210
 211        let forwarding_task = cx.background_executor().spawn(async move {
 212            loop {
 213                select_biased! {
 214                    _ = quit_rx.next().fuse() => {
 215                        break;
 216                    },
 217                    incoming_envelope = proxy_incoming_rx.next().fuse() => {
 218                        if let Some(envelope) = incoming_envelope {
 219                            if incoming_tx.send(envelope).await.is_err() {
 220                                break;
 221                            }
 222                        } else {
 223                            break;
 224                        }
 225                    }
 226                    outgoing_envelope = outgoing_rx.next().fuse() => {
 227                        if let Some(envelope) = outgoing_envelope {
 228                            if proxy_outgoing_tx.send(envelope).await.is_err() {
 229                                break;
 230                            }
 231                        } else {
 232                            break;
 233                        }
 234                    }
 235                }
 236            }
 237
 238            (incoming_tx, outgoing_rx)
 239        });
 240
 241        (
 242            Self {
 243                forwarding_task,
 244                quit_tx,
 245            },
 246            proxy_incoming_tx,
 247            proxy_outgoing_rx,
 248        )
 249    }
 250
 251    async fn into_channels(mut self) -> (UnboundedSender<Envelope>, UnboundedReceiver<Envelope>) {
 252        let _ = self.quit_tx.send(()).await;
 253        self.forwarding_task.await
 254    }
 255}
 256
 257const MAX_MISSED_HEARTBEATS: usize = 5;
 258const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
 259const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5);
 260
 261const MAX_RECONNECT_ATTEMPTS: usize = 3;
 262
 263enum State {
 264    Connecting,
 265    Connected {
 266        ssh_connection: SshRemoteConnection,
 267        delegate: Arc<dyn SshClientDelegate>,
 268        forwarder: ChannelForwarder,
 269
 270        multiplex_task: Task<Result<()>>,
 271        heartbeat_task: Task<Result<()>>,
 272    },
 273    HeartbeatMissed {
 274        missed_heartbeats: usize,
 275
 276        ssh_connection: SshRemoteConnection,
 277        delegate: Arc<dyn SshClientDelegate>,
 278        forwarder: ChannelForwarder,
 279
 280        multiplex_task: Task<Result<()>>,
 281        heartbeat_task: Task<Result<()>>,
 282    },
 283    Reconnecting,
 284    ReconnectFailed {
 285        ssh_connection: SshRemoteConnection,
 286        delegate: Arc<dyn SshClientDelegate>,
 287        forwarder: ChannelForwarder,
 288
 289        error: anyhow::Error,
 290        attempts: usize,
 291    },
 292    ReconnectExhausted,
 293    ServerNotRunning,
 294}
 295
 296impl fmt::Display for State {
 297    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
 298        match self {
 299            Self::Connecting => write!(f, "connecting"),
 300            Self::Connected { .. } => write!(f, "connected"),
 301            Self::Reconnecting => write!(f, "reconnecting"),
 302            Self::ReconnectFailed { .. } => write!(f, "reconnect failed"),
 303            Self::ReconnectExhausted => write!(f, "reconnect exhausted"),
 304            Self::HeartbeatMissed { .. } => write!(f, "heartbeat missed"),
 305            Self::ServerNotRunning { .. } => write!(f, "server not running"),
 306        }
 307    }
 308}
 309
 310impl State {
 311    fn ssh_connection(&self) -> Option<&SshRemoteConnection> {
 312        match self {
 313            Self::Connected { ssh_connection, .. } => Some(ssh_connection),
 314            Self::HeartbeatMissed { ssh_connection, .. } => Some(ssh_connection),
 315            Self::ReconnectFailed { ssh_connection, .. } => Some(ssh_connection),
 316            _ => None,
 317        }
 318    }
 319
 320    fn can_reconnect(&self) -> bool {
 321        match self {
 322            Self::Connected { .. }
 323            | Self::HeartbeatMissed { .. }
 324            | Self::ReconnectFailed { .. } => true,
 325            State::Connecting
 326            | State::Reconnecting
 327            | State::ReconnectExhausted
 328            | State::ServerNotRunning => false,
 329        }
 330    }
 331
 332    fn is_reconnect_failed(&self) -> bool {
 333        matches!(self, Self::ReconnectFailed { .. })
 334    }
 335
 336    fn is_reconnect_exhausted(&self) -> bool {
 337        matches!(self, Self::ReconnectExhausted { .. })
 338    }
 339
 340    fn is_reconnecting(&self) -> bool {
 341        matches!(self, Self::Reconnecting { .. })
 342    }
 343
 344    fn heartbeat_recovered(self) -> Self {
 345        match self {
 346            Self::HeartbeatMissed {
 347                ssh_connection,
 348                delegate,
 349                forwarder,
 350                multiplex_task,
 351                heartbeat_task,
 352                ..
 353            } => Self::Connected {
 354                ssh_connection,
 355                delegate,
 356                forwarder,
 357                multiplex_task,
 358                heartbeat_task,
 359            },
 360            _ => self,
 361        }
 362    }
 363
 364    fn heartbeat_missed(self) -> Self {
 365        match self {
 366            Self::Connected {
 367                ssh_connection,
 368                delegate,
 369                forwarder,
 370                multiplex_task,
 371                heartbeat_task,
 372            } => Self::HeartbeatMissed {
 373                missed_heartbeats: 1,
 374                ssh_connection,
 375                delegate,
 376                forwarder,
 377                multiplex_task,
 378                heartbeat_task,
 379            },
 380            Self::HeartbeatMissed {
 381                missed_heartbeats,
 382                ssh_connection,
 383                delegate,
 384                forwarder,
 385                multiplex_task,
 386                heartbeat_task,
 387            } => Self::HeartbeatMissed {
 388                missed_heartbeats: missed_heartbeats + 1,
 389                ssh_connection,
 390                delegate,
 391                forwarder,
 392                multiplex_task,
 393                heartbeat_task,
 394            },
 395            _ => self,
 396        }
 397    }
 398}
 399
 400/// The state of the ssh connection.
 401#[derive(Clone, Copy, Debug, PartialEq, Eq)]
 402pub enum ConnectionState {
 403    Connecting,
 404    Connected,
 405    HeartbeatMissed,
 406    Reconnecting,
 407    Disconnected,
 408}
 409
 410impl From<&State> for ConnectionState {
 411    fn from(value: &State) -> Self {
 412        match value {
 413            State::Connecting => Self::Connecting,
 414            State::Connected { .. } => Self::Connected,
 415            State::Reconnecting | State::ReconnectFailed { .. } => Self::Reconnecting,
 416            State::HeartbeatMissed { .. } => Self::HeartbeatMissed,
 417            State::ReconnectExhausted => Self::Disconnected,
 418            State::ServerNotRunning => Self::Disconnected,
 419        }
 420    }
 421}
 422
 423pub struct SshRemoteClient {
 424    client: Arc<ChannelClient>,
 425    unique_identifier: String,
 426    connection_options: SshConnectionOptions,
 427    state: Arc<Mutex<Option<State>>>,
 428}
 429
 430#[derive(Debug)]
 431pub enum SshRemoteEvent {
 432    Disconnected,
 433}
 434
 435impl EventEmitter<SshRemoteEvent> for SshRemoteClient {}
 436
 437impl SshRemoteClient {
 438    pub fn new(
 439        unique_identifier: String,
 440        connection_options: SshConnectionOptions,
 441        delegate: Arc<dyn SshClientDelegate>,
 442        cx: &AppContext,
 443    ) -> Task<Result<Model<Self>>> {
 444        cx.spawn(|mut cx| async move {
 445            let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
 446            let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
 447            let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
 448
 449            let client = cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx))?;
 450            let this = cx.new_model(|_| Self {
 451                client: client.clone(),
 452                unique_identifier: unique_identifier.clone(),
 453                connection_options: connection_options.clone(),
 454                state: Arc::new(Mutex::new(Some(State::Connecting))),
 455            })?;
 456
 457            let (proxy, proxy_incoming_tx, proxy_outgoing_rx) =
 458                ChannelForwarder::new(incoming_tx, outgoing_rx, &mut cx);
 459
 460            let (ssh_connection, ssh_proxy_process) = Self::establish_connection(
 461                unique_identifier,
 462                false,
 463                connection_options,
 464                delegate.clone(),
 465                &mut cx,
 466            )
 467            .await?;
 468
 469            let multiplex_task = Self::multiplex(
 470                this.downgrade(),
 471                ssh_proxy_process,
 472                proxy_incoming_tx,
 473                proxy_outgoing_rx,
 474                connection_activity_tx,
 475                &mut cx,
 476            );
 477
 478            if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
 479                log::error!("failed to establish connection: {}", error);
 480                delegate.set_error(error.to_string(), &mut cx);
 481                return Err(error);
 482            }
 483
 484            let heartbeat_task = Self::heartbeat(this.downgrade(), connection_activity_rx, &mut cx);
 485
 486            this.update(&mut cx, |this, _| {
 487                *this.state.lock() = Some(State::Connected {
 488                    ssh_connection,
 489                    delegate,
 490                    forwarder: proxy,
 491                    multiplex_task,
 492                    heartbeat_task,
 493                });
 494            })?;
 495
 496            Ok(this)
 497        })
 498    }
 499
 500    pub fn shutdown_processes<T: RequestMessage>(
 501        &self,
 502        shutdown_request: Option<T>,
 503    ) -> Option<impl Future<Output = ()>> {
 504        let state = self.state.lock().take()?;
 505        log::info!("shutting down ssh processes");
 506
 507        let State::Connected {
 508            multiplex_task,
 509            heartbeat_task,
 510            ssh_connection,
 511            delegate,
 512            forwarder,
 513        } = state
 514        else {
 515            return None;
 516        };
 517
 518        let client = self.client.clone();
 519
 520        Some(async move {
 521            if let Some(shutdown_request) = shutdown_request {
 522                client.send(shutdown_request).log_err();
 523                // We wait 50ms instead of waiting for a response, because
 524                // waiting for a response would require us to wait on the main thread
 525                // which we want to avoid in an `on_app_quit` callback.
 526                smol::Timer::after(Duration::from_millis(50)).await;
 527            }
 528
 529            // Drop `multiplex_task` because it owns our ssh_proxy_process, which is a
 530            // child of master_process.
 531            drop(multiplex_task);
 532            // Now drop the rest of state, which kills master process.
 533            drop(heartbeat_task);
 534            drop(ssh_connection);
 535            drop(delegate);
 536            drop(forwarder);
 537        })
 538    }
 539
 540    fn reconnect(&mut self, cx: &mut ModelContext<Self>) -> Result<()> {
 541        let mut lock = self.state.lock();
 542
 543        let can_reconnect = lock
 544            .as_ref()
 545            .map(|state| state.can_reconnect())
 546            .unwrap_or(false);
 547        if !can_reconnect {
 548            let error = if let Some(state) = lock.as_ref() {
 549                format!("invalid state, cannot reconnect while in state {state}")
 550            } else {
 551                "no state set".to_string()
 552            };
 553            log::info!("aborting reconnect, because not in state that allows reconnecting");
 554            return Err(anyhow!(error));
 555        }
 556
 557        let state = lock.take().unwrap();
 558        let (attempts, mut ssh_connection, delegate, forwarder) = match state {
 559            State::Connected {
 560                ssh_connection,
 561                delegate,
 562                forwarder,
 563                multiplex_task,
 564                heartbeat_task,
 565            }
 566            | State::HeartbeatMissed {
 567                ssh_connection,
 568                delegate,
 569                forwarder,
 570                multiplex_task,
 571                heartbeat_task,
 572                ..
 573            } => {
 574                drop(multiplex_task);
 575                drop(heartbeat_task);
 576                (0, ssh_connection, delegate, forwarder)
 577            }
 578            State::ReconnectFailed {
 579                attempts,
 580                ssh_connection,
 581                delegate,
 582                forwarder,
 583                ..
 584            } => (attempts, ssh_connection, delegate, forwarder),
 585            State::Connecting
 586            | State::Reconnecting
 587            | State::ReconnectExhausted
 588            | State::ServerNotRunning => unreachable!(),
 589        };
 590
 591        let attempts = attempts + 1;
 592        if attempts > MAX_RECONNECT_ATTEMPTS {
 593            log::error!(
 594                "Failed to reconnect to after {} attempts, giving up",
 595                MAX_RECONNECT_ATTEMPTS
 596            );
 597            drop(lock);
 598            self.set_state(State::ReconnectExhausted, cx);
 599            return Ok(());
 600        }
 601        drop(lock);
 602
 603        self.set_state(State::Reconnecting, cx);
 604
 605        log::info!("Trying to reconnect to ssh server... Attempt {}", attempts);
 606
 607        let identifier = self.unique_identifier.clone();
 608        let client = self.client.clone();
 609        let reconnect_task = cx.spawn(|this, mut cx| async move {
 610            macro_rules! failed {
 611                ($error:expr, $attempts:expr, $ssh_connection:expr, $delegate:expr, $forwarder:expr) => {
 612                    return State::ReconnectFailed {
 613                        error: anyhow!($error),
 614                        attempts: $attempts,
 615                        ssh_connection: $ssh_connection,
 616                        delegate: $delegate,
 617                        forwarder: $forwarder,
 618                    };
 619                };
 620            }
 621
 622            if let Err(error) = ssh_connection.master_process.kill() {
 623                failed!(error, attempts, ssh_connection, delegate, forwarder);
 624            };
 625
 626            if let Err(error) = ssh_connection
 627                .master_process
 628                .status()
 629                .await
 630                .context("Failed to kill ssh process")
 631            {
 632                failed!(error, attempts, ssh_connection, delegate, forwarder);
 633            }
 634
 635            let connection_options = ssh_connection.socket.connection_options.clone();
 636
 637            let (incoming_tx, outgoing_rx) = forwarder.into_channels().await;
 638            let (forwarder, proxy_incoming_tx, proxy_outgoing_rx) =
 639                ChannelForwarder::new(incoming_tx, outgoing_rx, &mut cx);
 640            let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
 641
 642            let (ssh_connection, ssh_process) = match Self::establish_connection(
 643                identifier,
 644                true,
 645                connection_options,
 646                delegate.clone(),
 647                &mut cx,
 648            )
 649            .await
 650            {
 651                Ok((ssh_connection, ssh_process)) => (ssh_connection, ssh_process),
 652                Err(error) => {
 653                    failed!(error, attempts, ssh_connection, delegate, forwarder);
 654                }
 655            };
 656
 657            let multiplex_task = Self::multiplex(
 658                this.clone(),
 659                ssh_process,
 660                proxy_incoming_tx,
 661                proxy_outgoing_rx,
 662                connection_activity_tx,
 663                &mut cx,
 664            );
 665
 666            if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
 667                failed!(error, attempts, ssh_connection, delegate, forwarder);
 668            };
 669
 670            State::Connected {
 671                ssh_connection,
 672                delegate,
 673                forwarder,
 674                multiplex_task,
 675                heartbeat_task: Self::heartbeat(this.clone(), connection_activity_rx, &mut cx),
 676            }
 677        });
 678
 679        cx.spawn(|this, mut cx| async move {
 680            let new_state = reconnect_task.await;
 681            this.update(&mut cx, |this, cx| {
 682                this.try_set_state(cx, |old_state| {
 683                    if old_state.is_reconnecting() {
 684                        match &new_state {
 685                            State::Connecting
 686                            | State::Reconnecting { .. }
 687                            | State::HeartbeatMissed { .. }
 688                            | State::ServerNotRunning => {}
 689                            State::Connected { .. } => {
 690                                log::info!("Successfully reconnected");
 691                            }
 692                            State::ReconnectFailed {
 693                                error, attempts, ..
 694                            } => {
 695                                log::error!(
 696                                    "Reconnect attempt {} failed: {:?}. Starting new attempt...",
 697                                    attempts,
 698                                    error
 699                                );
 700                            }
 701                            State::ReconnectExhausted => {
 702                                log::error!("Reconnect attempt failed and all attempts exhausted");
 703                            }
 704                        }
 705                        Some(new_state)
 706                    } else {
 707                        None
 708                    }
 709                });
 710
 711                if this.state_is(State::is_reconnect_failed) {
 712                    this.reconnect(cx)
 713                } else if this.state_is(State::is_reconnect_exhausted) {
 714                    cx.emit(SshRemoteEvent::Disconnected);
 715                    Ok(())
 716                } else {
 717                    log::debug!("State has transition from Reconnecting into new state while attempting reconnect. Ignoring new state.");
 718                    Ok(())
 719                }
 720            })
 721        })
 722        .detach_and_log_err(cx);
 723
 724        Ok(())
 725    }
 726
 727    fn heartbeat(
 728        this: WeakModel<Self>,
 729        mut connection_activity_rx: mpsc::Receiver<()>,
 730        cx: &mut AsyncAppContext,
 731    ) -> Task<Result<()>> {
 732        let Ok(client) = this.update(cx, |this, _| this.client.clone()) else {
 733            return Task::ready(Err(anyhow!("SshRemoteClient lost")));
 734        };
 735
 736        cx.spawn(|mut cx| {
 737            let this = this.clone();
 738            async move {
 739                let mut missed_heartbeats = 0;
 740
 741                let keepalive_timer = cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse();
 742                futures::pin_mut!(keepalive_timer);
 743
 744                loop {
 745                    select_biased! {
 746                        result = connection_activity_rx.next().fuse() => {
 747                            if result.is_none() {
 748                                log::warn!("ssh heartbeat: connection activity channel has been dropped. stopping.");
 749                                return Ok(());
 750                            }
 751                            keepalive_timer.set(cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse());
 752                        }
 753                        _ = keepalive_timer => {
 754                            log::debug!("Sending heartbeat to server...");
 755
 756                            let result = select_biased! {
 757                                _ = connection_activity_rx.next().fuse() => {
 758                                    Ok(())
 759                                }
 760                                ping_result = client.ping(HEARTBEAT_TIMEOUT).fuse() => {
 761                                    ping_result
 762                                }
 763                            };
 764                            if result.is_err() {
 765                                missed_heartbeats += 1;
 766                                log::warn!(
 767                                    "No heartbeat from server after {:?}. Missed heartbeat {} out of {}.",
 768                                    HEARTBEAT_TIMEOUT,
 769                                    missed_heartbeats,
 770                                    MAX_MISSED_HEARTBEATS
 771                                );
 772                            } else if missed_heartbeats != 0 {
 773                                missed_heartbeats = 0;
 774                            } else {
 775                                continue;
 776                            }
 777
 778                            let result = this.update(&mut cx, |this, mut cx| {
 779                                this.handle_heartbeat_result(missed_heartbeats, &mut cx)
 780                            })?;
 781                            if result.is_break() {
 782                                return Ok(());
 783                            }
 784                        }
 785                    }
 786                }
 787            }
 788        })
 789    }
 790
 791    fn handle_heartbeat_result(
 792        &mut self,
 793        missed_heartbeats: usize,
 794        cx: &mut ModelContext<Self>,
 795    ) -> ControlFlow<()> {
 796        let state = self.state.lock().take().unwrap();
 797        let next_state = if missed_heartbeats > 0 {
 798            state.heartbeat_missed()
 799        } else {
 800            state.heartbeat_recovered()
 801        };
 802
 803        self.set_state(next_state, cx);
 804
 805        if missed_heartbeats >= MAX_MISSED_HEARTBEATS {
 806            log::error!(
 807                "Missed last {} heartbeats. Reconnecting...",
 808                missed_heartbeats
 809            );
 810
 811            self.reconnect(cx)
 812                .context("failed to start reconnect process after missing heartbeats")
 813                .log_err();
 814            ControlFlow::Break(())
 815        } else {
 816            ControlFlow::Continue(())
 817        }
 818    }
 819
 820    fn multiplex(
 821        this: WeakModel<Self>,
 822        mut ssh_proxy_process: Child,
 823        incoming_tx: UnboundedSender<Envelope>,
 824        mut outgoing_rx: UnboundedReceiver<Envelope>,
 825        mut connection_activity_tx: Sender<()>,
 826        cx: &AsyncAppContext,
 827    ) -> Task<Result<()>> {
 828        let mut child_stderr = ssh_proxy_process.stderr.take().unwrap();
 829        let mut child_stdout = ssh_proxy_process.stdout.take().unwrap();
 830        let mut child_stdin = ssh_proxy_process.stdin.take().unwrap();
 831
 832        let io_task = cx.background_executor().spawn(async move {
 833            let mut stdin_buffer = Vec::new();
 834            let mut stdout_buffer = Vec::new();
 835            let mut stderr_buffer = Vec::new();
 836            let mut stderr_offset = 0;
 837
 838            loop {
 839                stdout_buffer.resize(MESSAGE_LEN_SIZE, 0);
 840                stderr_buffer.resize(stderr_offset + 1024, 0);
 841
 842                select_biased! {
 843                    outgoing = outgoing_rx.next().fuse() => {
 844                        let Some(outgoing) = outgoing else {
 845                            return anyhow::Ok(None);
 846                        };
 847
 848                        write_message(&mut child_stdin, &mut stdin_buffer, outgoing).await?;
 849                    }
 850
 851                    result = child_stdout.read(&mut stdout_buffer).fuse() => {
 852                        match result {
 853                            Ok(0) => {
 854                                child_stdin.close().await?;
 855                                outgoing_rx.close();
 856                                let status = ssh_proxy_process.status().await?;
 857                                // If we don't have a code, we assume process
 858                                // has been killed and treat it as non-zero exit
 859                                // code
 860                                return Ok(status.code().or_else(|| Some(1)));
 861                            }
 862                            Ok(len) => {
 863                                if len < stdout_buffer.len() {
 864                                    child_stdout.read_exact(&mut stdout_buffer[len..]).await?;
 865                                }
 866
 867                                let message_len = message_len_from_buffer(&stdout_buffer);
 868                                match read_message_with_len(&mut child_stdout, &mut stdout_buffer, message_len).await {
 869                                    Ok(envelope) => {
 870                                        connection_activity_tx.try_send(()).ok();
 871                                        incoming_tx.unbounded_send(envelope).ok();
 872                                    }
 873                                    Err(error) => {
 874                                        log::error!("error decoding message {error:?}");
 875                                    }
 876                                }
 877                            }
 878                            Err(error) => {
 879                                Err(anyhow!("error reading stdout: {error:?}"))?;
 880                            }
 881                        }
 882                    }
 883
 884                    result = child_stderr.read(&mut stderr_buffer[stderr_offset..]).fuse() => {
 885                        match result {
 886                            Ok(len) => {
 887                                stderr_offset += len;
 888                                let mut start_ix = 0;
 889                                while let Some(ix) = stderr_buffer[start_ix..stderr_offset].iter().position(|b| b == &b'\n') {
 890                                    let line_ix = start_ix + ix;
 891                                    let content = &stderr_buffer[start_ix..line_ix];
 892                                    start_ix = line_ix + 1;
 893                                    if let Ok(record) = serde_json::from_slice::<LogRecord>(content) {
 894                                        record.log(log::logger())
 895                                    } else {
 896                                        eprintln!("(remote) {}", String::from_utf8_lossy(content));
 897                                    }
 898                                }
 899                                stderr_buffer.drain(0..start_ix);
 900                                stderr_offset -= start_ix;
 901
 902                                connection_activity_tx.try_send(()).ok();
 903                            }
 904                            Err(error) => {
 905                                Err(anyhow!("error reading stderr: {error:?}"))?;
 906                            }
 907                        }
 908                    }
 909                }
 910            }
 911        });
 912
 913        cx.spawn(|mut cx| async move {
 914            let result = io_task.await;
 915
 916            match result {
 917                Ok(Some(exit_code)) => {
 918                    if let Some(error) = ProxyLaunchError::from_exit_code(exit_code) {
 919                        match error {
 920                            ProxyLaunchError::ServerNotRunning => {
 921                                log::error!("failed to reconnect because server is not running");
 922                                this.update(&mut cx, |this, cx| {
 923                                    this.set_state(State::ServerNotRunning, cx);
 924                                    cx.emit(SshRemoteEvent::Disconnected);
 925                                })?;
 926                            }
 927                        }
 928                    } else if exit_code > 0 {
 929                        log::error!("proxy process terminated unexpectedly");
 930                        this.update(&mut cx, |this, cx| {
 931                            this.reconnect(cx).ok();
 932                        })?;
 933                    }
 934                }
 935                Ok(None) => {}
 936                Err(error) => {
 937                    log::warn!("ssh io task died with error: {:?}. reconnecting...", error);
 938                    this.update(&mut cx, |this, cx| {
 939                        this.reconnect(cx).ok();
 940                    })?;
 941                }
 942            }
 943            Ok(())
 944        })
 945    }
 946
 947    fn state_is(&self, check: impl FnOnce(&State) -> bool) -> bool {
 948        self.state.lock().as_ref().map_or(false, check)
 949    }
 950
 951    fn try_set_state(
 952        &self,
 953        cx: &mut ModelContext<Self>,
 954        map: impl FnOnce(&State) -> Option<State>,
 955    ) {
 956        let mut lock = self.state.lock();
 957        let new_state = lock.as_ref().and_then(map);
 958
 959        if let Some(new_state) = new_state {
 960            lock.replace(new_state);
 961            cx.notify();
 962        }
 963    }
 964
 965    fn set_state(&self, state: State, cx: &mut ModelContext<Self>) {
 966        log::info!("setting state to '{}'", &state);
 967        self.state.lock().replace(state);
 968        cx.notify();
 969    }
 970
 971    async fn establish_connection(
 972        unique_identifier: String,
 973        reconnect: bool,
 974        connection_options: SshConnectionOptions,
 975        delegate: Arc<dyn SshClientDelegate>,
 976        cx: &mut AsyncAppContext,
 977    ) -> Result<(SshRemoteConnection, Child)> {
 978        let ssh_connection =
 979            SshRemoteConnection::new(connection_options, delegate.clone(), cx).await?;
 980
 981        let platform = ssh_connection.query_platform().await?;
 982        let (local_binary_path, version) = delegate.get_server_binary(platform, cx).await??;
 983        let remote_binary_path = delegate.remote_server_binary_path(platform, cx)?;
 984        ssh_connection
 985            .ensure_server_binary(
 986                &delegate,
 987                &local_binary_path,
 988                &remote_binary_path,
 989                version,
 990                cx,
 991            )
 992            .await?;
 993
 994        let socket = ssh_connection.socket.clone();
 995        run_cmd(socket.ssh_command(&remote_binary_path).arg("version")).await?;
 996
 997        delegate.set_status(Some("Starting proxy"), cx);
 998
 999        let mut start_proxy_command = format!(
1000            "RUST_LOG={} RUST_BACKTRACE={} {:?} proxy --identifier {}",
1001            std::env::var("RUST_LOG").unwrap_or_default(),
1002            std::env::var("RUST_BACKTRACE").unwrap_or_default(),
1003            remote_binary_path,
1004            unique_identifier,
1005        );
1006        if reconnect {
1007            start_proxy_command.push_str(" --reconnect");
1008        }
1009
1010        let ssh_proxy_process = socket
1011            .ssh_command(start_proxy_command)
1012            // IMPORTANT: we kill this process when we drop the task that uses it.
1013            .kill_on_drop(true)
1014            .spawn()
1015            .context("failed to spawn remote server")?;
1016
1017        Ok((ssh_connection, ssh_proxy_process))
1018    }
1019
1020    pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Model<E>) {
1021        self.client.subscribe_to_entity(remote_id, entity);
1022    }
1023
1024    pub fn ssh_args(&self) -> Option<Vec<String>> {
1025        self.state
1026            .lock()
1027            .as_ref()
1028            .and_then(|state| state.ssh_connection())
1029            .map(|ssh_connection| ssh_connection.socket.ssh_args())
1030    }
1031
1032    pub fn proto_client(&self) -> AnyProtoClient {
1033        self.client.clone().into()
1034    }
1035
1036    pub fn connection_string(&self) -> String {
1037        self.connection_options.connection_string()
1038    }
1039
1040    pub fn connection_options(&self) -> SshConnectionOptions {
1041        self.connection_options.clone()
1042    }
1043
1044    #[cfg(not(any(test, feature = "test-support")))]
1045    pub fn connection_state(&self) -> ConnectionState {
1046        self.state
1047            .lock()
1048            .as_ref()
1049            .map(ConnectionState::from)
1050            .unwrap_or(ConnectionState::Disconnected)
1051    }
1052
1053    #[cfg(any(test, feature = "test-support"))]
1054    pub fn connection_state(&self) -> ConnectionState {
1055        ConnectionState::Connected
1056    }
1057
1058    pub fn is_disconnected(&self) -> bool {
1059        self.connection_state() == ConnectionState::Disconnected
1060    }
1061
1062    #[cfg(any(test, feature = "test-support"))]
1063    pub fn fake(
1064        client_cx: &mut gpui::TestAppContext,
1065        server_cx: &mut gpui::TestAppContext,
1066    ) -> (Model<Self>, Arc<ChannelClient>) {
1067        use gpui::Context;
1068
1069        let (server_to_client_tx, server_to_client_rx) = mpsc::unbounded();
1070        let (client_to_server_tx, client_to_server_rx) = mpsc::unbounded();
1071
1072        (
1073            client_cx.update(|cx| {
1074                let client = ChannelClient::new(server_to_client_rx, client_to_server_tx, cx);
1075                cx.new_model(|_| Self {
1076                    client,
1077                    unique_identifier: "fake".to_string(),
1078                    connection_options: SshConnectionOptions::default(),
1079                    state: Arc::new(Mutex::new(None)),
1080                })
1081            }),
1082            server_cx.update(|cx| ChannelClient::new(client_to_server_rx, server_to_client_tx, cx)),
1083        )
1084    }
1085}
1086
1087impl From<SshRemoteClient> for AnyProtoClient {
1088    fn from(client: SshRemoteClient) -> Self {
1089        AnyProtoClient::new(client.client.clone())
1090    }
1091}
1092
1093struct SshRemoteConnection {
1094    socket: SshSocket,
1095    master_process: process::Child,
1096    _temp_dir: TempDir,
1097}
1098
1099impl Drop for SshRemoteConnection {
1100    fn drop(&mut self) {
1101        if let Err(error) = self.master_process.kill() {
1102            log::error!("failed to kill SSH master process: {}", error);
1103        }
1104    }
1105}
1106
1107impl SshRemoteConnection {
1108    #[cfg(not(unix))]
1109    async fn new(
1110        _connection_options: SshConnectionOptions,
1111        _delegate: Arc<dyn SshClientDelegate>,
1112        _cx: &mut AsyncAppContext,
1113    ) -> Result<Self> {
1114        Err(anyhow!("ssh is not supported on this platform"))
1115    }
1116
1117    #[cfg(unix)]
1118    async fn new(
1119        connection_options: SshConnectionOptions,
1120        delegate: Arc<dyn SshClientDelegate>,
1121        cx: &mut AsyncAppContext,
1122    ) -> Result<Self> {
1123        use futures::{io::BufReader, AsyncBufReadExt as _};
1124        use smol::{fs::unix::PermissionsExt as _, net::unix::UnixListener};
1125        use util::ResultExt as _;
1126
1127        delegate.set_status(Some("connecting"), cx);
1128
1129        let url = connection_options.ssh_url();
1130        let temp_dir = tempfile::Builder::new()
1131            .prefix("zed-ssh-session")
1132            .tempdir()?;
1133
1134        // Create a domain socket listener to handle requests from the askpass program.
1135        let askpass_socket = temp_dir.path().join("askpass.sock");
1136        let (askpass_opened_tx, askpass_opened_rx) = oneshot::channel::<()>();
1137        let listener =
1138            UnixListener::bind(&askpass_socket).context("failed to create askpass socket")?;
1139
1140        let askpass_task = cx.spawn({
1141            let delegate = delegate.clone();
1142            |mut cx| async move {
1143                let mut askpass_opened_tx = Some(askpass_opened_tx);
1144
1145                while let Ok((mut stream, _)) = listener.accept().await {
1146                    if let Some(askpass_opened_tx) = askpass_opened_tx.take() {
1147                        askpass_opened_tx.send(()).ok();
1148                    }
1149                    let mut buffer = Vec::new();
1150                    let mut reader = BufReader::new(&mut stream);
1151                    if reader.read_until(b'\0', &mut buffer).await.is_err() {
1152                        buffer.clear();
1153                    }
1154                    let password_prompt = String::from_utf8_lossy(&buffer);
1155                    if let Some(password) = delegate
1156                        .ask_password(password_prompt.to_string(), &mut cx)
1157                        .await
1158                        .context("failed to get ssh password")
1159                        .and_then(|p| p)
1160                        .log_err()
1161                    {
1162                        stream.write_all(password.as_bytes()).await.log_err();
1163                    }
1164                }
1165            }
1166        });
1167
1168        // Create an askpass script that communicates back to this process.
1169        let askpass_script = format!(
1170            "{shebang}\n{print_args} | nc -U {askpass_socket} 2> /dev/null \n",
1171            askpass_socket = askpass_socket.display(),
1172            print_args = "printf '%s\\0' \"$@\"",
1173            shebang = "#!/bin/sh",
1174        );
1175        let askpass_script_path = temp_dir.path().join("askpass.sh");
1176        fs::write(&askpass_script_path, askpass_script).await?;
1177        fs::set_permissions(&askpass_script_path, std::fs::Permissions::from_mode(0o755)).await?;
1178
1179        // Start the master SSH process, which does not do anything except for establish
1180        // the connection and keep it open, allowing other ssh commands to reuse it
1181        // via a control socket.
1182        let socket_path = temp_dir.path().join("ssh.sock");
1183        let mut master_process = process::Command::new("ssh")
1184            .stdin(Stdio::null())
1185            .stdout(Stdio::piped())
1186            .stderr(Stdio::piped())
1187            .env("SSH_ASKPASS_REQUIRE", "force")
1188            .env("SSH_ASKPASS", &askpass_script_path)
1189            .args([
1190                "-N",
1191                "-o",
1192                "ControlPersist=no",
1193                "-o",
1194                "ControlMaster=yes",
1195                "-o",
1196            ])
1197            .arg(format!("ControlPath={}", socket_path.display()))
1198            .arg(&url)
1199            .spawn()?;
1200
1201        // Wait for this ssh process to close its stdout, indicating that authentication
1202        // has completed.
1203        let stdout = master_process.stdout.as_mut().unwrap();
1204        let mut output = Vec::new();
1205        let connection_timeout = Duration::from_secs(10);
1206
1207        let result = select_biased! {
1208            _ = askpass_opened_rx.fuse() => {
1209                // If the askpass script has opened, that means the user is typing
1210                // their password, in which case we don't want to timeout anymore,
1211                // since we know a connection has been established.
1212                stdout.read_to_end(&mut output).await?;
1213                Ok(())
1214            }
1215            result = stdout.read_to_end(&mut output).fuse() => {
1216                result?;
1217                Ok(())
1218            }
1219            _ = futures::FutureExt::fuse(smol::Timer::after(connection_timeout)) => {
1220                Err(anyhow!("Exceeded {:?} timeout trying to connect to host", connection_timeout))
1221            }
1222        };
1223
1224        if let Err(e) = result {
1225            let error_message = format!("Failed to connect to host: {}.", e);
1226            delegate.set_error(error_message, cx);
1227            return Err(e);
1228        }
1229
1230        drop(askpass_task);
1231
1232        if master_process.try_status()?.is_some() {
1233            output.clear();
1234            let mut stderr = master_process.stderr.take().unwrap();
1235            stderr.read_to_end(&mut output).await?;
1236
1237            let error_message = format!("failed to connect: {}", String::from_utf8_lossy(&output));
1238            delegate.set_error(error_message.clone(), cx);
1239            Err(anyhow!(error_message))?;
1240        }
1241
1242        Ok(Self {
1243            socket: SshSocket {
1244                connection_options,
1245                socket_path,
1246            },
1247            master_process,
1248            _temp_dir: temp_dir,
1249        })
1250    }
1251
1252    async fn ensure_server_binary(
1253        &self,
1254        delegate: &Arc<dyn SshClientDelegate>,
1255        src_path: &Path,
1256        dst_path: &Path,
1257        version: SemanticVersion,
1258        cx: &mut AsyncAppContext,
1259    ) -> Result<()> {
1260        let mut dst_path_gz = dst_path.to_path_buf();
1261        dst_path_gz.set_extension("gz");
1262
1263        if let Some(parent) = dst_path.parent() {
1264            run_cmd(self.socket.ssh_command("mkdir").arg("-p").arg(parent)).await?;
1265        }
1266
1267        let mut server_binary_exists = false;
1268        if cfg!(not(debug_assertions)) {
1269            if let Ok(installed_version) =
1270                run_cmd(self.socket.ssh_command(dst_path).arg("version")).await
1271            {
1272                if installed_version.trim() == version.to_string() {
1273                    server_binary_exists = true;
1274                }
1275            }
1276        }
1277
1278        if server_binary_exists {
1279            log::info!("remote development server already present",);
1280            return Ok(());
1281        }
1282
1283        let src_stat = fs::metadata(src_path).await?;
1284        let size = src_stat.len();
1285        let server_mode = 0o755;
1286
1287        let t0 = Instant::now();
1288        delegate.set_status(Some("uploading remote development server"), cx);
1289        log::info!("uploading remote development server ({}kb)", size / 1024);
1290        self.upload_file(src_path, &dst_path_gz)
1291            .await
1292            .context("failed to upload server binary")?;
1293        log::info!("uploaded remote development server in {:?}", t0.elapsed());
1294
1295        delegate.set_status(Some("extracting remote development server"), cx);
1296        run_cmd(
1297            self.socket
1298                .ssh_command("gunzip")
1299                .arg("--force")
1300                .arg(&dst_path_gz),
1301        )
1302        .await?;
1303
1304        delegate.set_status(Some("unzipping remote development server"), cx);
1305        run_cmd(
1306            self.socket
1307                .ssh_command("chmod")
1308                .arg(format!("{:o}", server_mode))
1309                .arg(dst_path),
1310        )
1311        .await?;
1312
1313        Ok(())
1314    }
1315
1316    async fn query_platform(&self) -> Result<SshPlatform> {
1317        let os = run_cmd(self.socket.ssh_command("uname").arg("-s")).await?;
1318        let arch = run_cmd(self.socket.ssh_command("uname").arg("-m")).await?;
1319
1320        let os = match os.trim() {
1321            "Darwin" => "macos",
1322            "Linux" => "linux",
1323            _ => Err(anyhow!("unknown uname os {os:?}"))?,
1324        };
1325        let arch = if arch.starts_with("arm") || arch.starts_with("aarch64") {
1326            "aarch64"
1327        } else if arch.starts_with("x86") || arch.starts_with("i686") {
1328            "x86_64"
1329        } else {
1330            Err(anyhow!("unknown uname architecture {arch:?}"))?
1331        };
1332
1333        Ok(SshPlatform { os, arch })
1334    }
1335
1336    async fn upload_file(&self, src_path: &Path, dest_path: &Path) -> Result<()> {
1337        let mut command = process::Command::new("scp");
1338        let output = self
1339            .socket
1340            .ssh_options(&mut command)
1341            .args(
1342                self.socket
1343                    .connection_options
1344                    .port
1345                    .map(|port| vec!["-P".to_string(), port.to_string()])
1346                    .unwrap_or_default(),
1347            )
1348            .arg(src_path)
1349            .arg(format!(
1350                "{}:{}",
1351                self.socket.connection_options.scp_url(),
1352                dest_path.display()
1353            ))
1354            .output()
1355            .await?;
1356
1357        if output.status.success() {
1358            Ok(())
1359        } else {
1360            Err(anyhow!(
1361                "failed to upload file {} -> {}: {}",
1362                src_path.display(),
1363                dest_path.display(),
1364                String::from_utf8_lossy(&output.stderr)
1365            ))
1366        }
1367    }
1368}
1369
1370type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>;
1371
1372pub struct ChannelClient {
1373    next_message_id: AtomicU32,
1374    outgoing_tx: mpsc::UnboundedSender<Envelope>,
1375    response_channels: ResponseChannels,             // Lock
1376    message_handlers: Mutex<ProtoMessageHandlerSet>, // Lock
1377}
1378
1379impl ChannelClient {
1380    pub fn new(
1381        incoming_rx: mpsc::UnboundedReceiver<Envelope>,
1382        outgoing_tx: mpsc::UnboundedSender<Envelope>,
1383        cx: &AppContext,
1384    ) -> Arc<Self> {
1385        let this = Arc::new(Self {
1386            outgoing_tx,
1387            next_message_id: AtomicU32::new(0),
1388            response_channels: ResponseChannels::default(),
1389            message_handlers: Default::default(),
1390        });
1391
1392        Self::start_handling_messages(this.clone(), incoming_rx, cx);
1393
1394        this
1395    }
1396
1397    fn start_handling_messages(
1398        this: Arc<Self>,
1399        mut incoming_rx: mpsc::UnboundedReceiver<Envelope>,
1400        cx: &AppContext,
1401    ) {
1402        cx.spawn(|cx| {
1403            let this = Arc::downgrade(&this);
1404            async move {
1405                let peer_id = PeerId { owner_id: 0, id: 0 };
1406                while let Some(incoming) = incoming_rx.next().await {
1407                    let Some(this) = this.upgrade() else {
1408                        return anyhow::Ok(());
1409                    };
1410
1411                    if let Some(request_id) = incoming.responding_to {
1412                        let request_id = MessageId(request_id);
1413                        let sender = this.response_channels.lock().remove(&request_id);
1414                        if let Some(sender) = sender {
1415                            let (tx, rx) = oneshot::channel();
1416                            if incoming.payload.is_some() {
1417                                sender.send((incoming, tx)).ok();
1418                            }
1419                            rx.await.ok();
1420                        }
1421                    } else if let Some(envelope) =
1422                        build_typed_envelope(peer_id, Instant::now(), incoming)
1423                    {
1424                        let type_name = envelope.payload_type_name();
1425                        if let Some(future) = ProtoMessageHandlerSet::handle_message(
1426                            &this.message_handlers,
1427                            envelope,
1428                            this.clone().into(),
1429                            cx.clone(),
1430                        ) {
1431                            log::debug!("ssh message received. name:{type_name}");
1432                            cx.foreground_executor().spawn(async move {
1433                                match future.await {
1434                                    Ok(_) => {
1435                                        log::debug!("ssh message handled. name:{type_name}");
1436                                    }
1437                                    Err(error) => {
1438                                        log::error!(
1439                                            "error handling message. type:{type_name}, error:{error}",
1440                                        );
1441                                    }
1442                                }
1443                            }).detach();
1444
1445                        } else {
1446                            log::error!("unhandled ssh message name:{type_name}");
1447                        }
1448                    }
1449                }
1450                anyhow::Ok(())
1451            }
1452        })
1453        .detach();
1454    }
1455
1456    pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Model<E>) {
1457        let id = (TypeId::of::<E>(), remote_id);
1458
1459        let mut message_handlers = self.message_handlers.lock();
1460        if message_handlers
1461            .entities_by_type_and_remote_id
1462            .contains_key(&id)
1463        {
1464            panic!("already subscribed to entity");
1465        }
1466
1467        message_handlers.entities_by_type_and_remote_id.insert(
1468            id,
1469            EntityMessageSubscriber::Entity {
1470                handle: entity.downgrade().into(),
1471            },
1472        );
1473    }
1474
1475    pub fn request<T: RequestMessage>(
1476        &self,
1477        payload: T,
1478    ) -> impl 'static + Future<Output = Result<T::Response>> {
1479        log::debug!("ssh request start. name:{}", T::NAME);
1480        let response = self.request_dynamic(payload.into_envelope(0, None, None), T::NAME);
1481        async move {
1482            let response = response.await?;
1483            log::debug!("ssh request finish. name:{}", T::NAME);
1484            T::Response::from_envelope(response)
1485                .ok_or_else(|| anyhow!("received a response of the wrong type"))
1486        }
1487    }
1488
1489    pub async fn ping(&self, timeout: Duration) -> Result<()> {
1490        smol::future::or(
1491            async {
1492                self.request(proto::Ping {}).await?;
1493                Ok(())
1494            },
1495            async {
1496                smol::Timer::after(timeout).await;
1497                Err(anyhow!("Timeout detected"))
1498            },
1499        )
1500        .await
1501    }
1502
1503    pub fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {
1504        log::debug!("ssh send name:{}", T::NAME);
1505        self.send_dynamic(payload.into_envelope(0, None, None))
1506    }
1507
1508    pub fn request_dynamic(
1509        &self,
1510        mut envelope: proto::Envelope,
1511        type_name: &'static str,
1512    ) -> impl 'static + Future<Output = Result<proto::Envelope>> {
1513        envelope.id = self.next_message_id.fetch_add(1, SeqCst);
1514        let (tx, rx) = oneshot::channel();
1515        let mut response_channels_lock = self.response_channels.lock();
1516        response_channels_lock.insert(MessageId(envelope.id), tx);
1517        drop(response_channels_lock);
1518        let result = self.outgoing_tx.unbounded_send(envelope);
1519        async move {
1520            if let Err(error) = &result {
1521                log::error!("failed to send message: {}", error);
1522                return Err(anyhow!("failed to send message: {}", error));
1523            }
1524
1525            let response = rx.await.context("connection lost")?.0;
1526            if let Some(proto::envelope::Payload::Error(error)) = &response.payload {
1527                return Err(RpcError::from_proto(error, type_name));
1528            }
1529            Ok(response)
1530        }
1531    }
1532
1533    pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {
1534        envelope.id = self.next_message_id.fetch_add(1, SeqCst);
1535        self.outgoing_tx.unbounded_send(envelope)?;
1536        Ok(())
1537    }
1538}
1539
1540impl ProtoClient for ChannelClient {
1541    fn request(
1542        &self,
1543        envelope: proto::Envelope,
1544        request_type: &'static str,
1545    ) -> BoxFuture<'static, Result<proto::Envelope>> {
1546        self.request_dynamic(envelope, request_type).boxed()
1547    }
1548
1549    fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
1550        self.send_dynamic(envelope)
1551    }
1552
1553    fn send_response(&self, envelope: Envelope, _message_type: &'static str) -> anyhow::Result<()> {
1554        self.send_dynamic(envelope)
1555    }
1556
1557    fn message_handler_set(&self) -> &Mutex<ProtoMessageHandlerSet> {
1558        &self.message_handlers
1559    }
1560
1561    fn is_via_collab(&self) -> bool {
1562        false
1563    }
1564}