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