1use crate::{
   2    RemoteClientDelegate, RemotePlatform,
   3    remote_client::{CommandTemplate, RemoteConnection, RemoteConnectionOptions},
   4};
   5use anyhow::{Context as _, Result, anyhow};
   6use async_trait::async_trait;
   7use collections::HashMap;
   8use futures::{
   9    AsyncReadExt as _, FutureExt as _,
  10    channel::mpsc::{Sender, UnboundedReceiver, UnboundedSender},
  11    select_biased,
  12};
  13use gpui::{App, AppContext as _, AsyncApp, SemanticVersion, Task};
  14use parking_lot::Mutex;
  15use paths::remote_server_dir_relative;
  16use release_channel::{AppCommitSha, AppVersion, ReleaseChannel};
  17use rpc::proto::Envelope;
  18pub use settings::SshPortForwardOption;
  19use smol::{
  20    fs,
  21    process::{self, Child, Stdio},
  22};
  23use std::{
  24    path::{Path, PathBuf},
  25    sync::Arc,
  26    time::Instant,
  27};
  28use tempfile::TempDir;
  29use util::{
  30    paths::{PathStyle, RemotePathBuf},
  31    rel_path::RelPath,
  32};
  33
  34pub(crate) struct SshRemoteConnection {
  35    socket: SshSocket,
  36    master_process: Mutex<Option<Child>>,
  37    remote_binary_path: Option<Arc<RelPath>>,
  38    ssh_platform: RemotePlatform,
  39    ssh_path_style: PathStyle,
  40    ssh_shell: String,
  41    ssh_default_system_shell: String,
  42    _temp_dir: TempDir,
  43}
  44
  45#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
  46pub struct SshConnectionOptions {
  47    pub host: String,
  48    pub username: Option<String>,
  49    pub port: Option<u16>,
  50    pub password: Option<String>,
  51    pub args: Option<Vec<String>>,
  52    pub port_forwards: Option<Vec<SshPortForwardOption>>,
  53
  54    pub nickname: Option<String>,
  55    pub upload_binary_over_ssh: bool,
  56}
  57
  58impl From<settings::SshConnection> for SshConnectionOptions {
  59    fn from(val: settings::SshConnection) -> Self {
  60        SshConnectionOptions {
  61            host: val.host.into(),
  62            username: val.username,
  63            port: val.port,
  64            password: None,
  65            args: Some(val.args),
  66            nickname: val.nickname,
  67            upload_binary_over_ssh: val.upload_binary_over_ssh.unwrap_or_default(),
  68            port_forwards: val.port_forwards,
  69        }
  70    }
  71}
  72
  73struct SshSocket {
  74    connection_options: SshConnectionOptions,
  75    #[cfg(not(target_os = "windows"))]
  76    socket_path: std::path::PathBuf,
  77    envs: HashMap<String, String>,
  78    #[cfg(target_os = "windows")]
  79    _proxy: askpass::PasswordProxy,
  80}
  81
  82macro_rules! shell_script {
  83    ($fmt:expr, $($name:ident = $arg:expr),+ $(,)?) => {{
  84        format!(
  85            $fmt,
  86            $(
  87                $name = shlex::try_quote($arg).unwrap()
  88            ),+
  89        )
  90    }};
  91}
  92
  93#[async_trait(?Send)]
  94impl RemoteConnection for SshRemoteConnection {
  95    async fn kill(&self) -> Result<()> {
  96        let Some(mut process) = self.master_process.lock().take() else {
  97            return Ok(());
  98        };
  99        process.kill().ok();
 100        process.status().await?;
 101        Ok(())
 102    }
 103
 104    fn has_been_killed(&self) -> bool {
 105        self.master_process.lock().is_none()
 106    }
 107
 108    fn connection_options(&self) -> RemoteConnectionOptions {
 109        RemoteConnectionOptions::Ssh(self.socket.connection_options.clone())
 110    }
 111
 112    fn shell(&self) -> String {
 113        self.ssh_shell.clone()
 114    }
 115
 116    fn default_system_shell(&self) -> String {
 117        self.ssh_default_system_shell.clone()
 118    }
 119
 120    fn build_command(
 121        &self,
 122        input_program: Option<String>,
 123        input_args: &[String],
 124        input_env: &HashMap<String, String>,
 125        working_dir: Option<String>,
 126        port_forward: Option<(u16, String, u16)>,
 127    ) -> Result<CommandTemplate> {
 128        let Self {
 129            ssh_path_style,
 130            socket,
 131            ssh_shell,
 132            ..
 133        } = self;
 134        let env = socket.envs.clone();
 135        build_command(
 136            input_program,
 137            input_args,
 138            input_env,
 139            working_dir,
 140            port_forward,
 141            env,
 142            *ssh_path_style,
 143            ssh_shell,
 144            socket.ssh_args(),
 145        )
 146    }
 147
 148    fn build_forward_port_command(
 149        &self,
 150        local_port: u16,
 151        host: String,
 152        remote_port: u16,
 153    ) -> Result<CommandTemplate> {
 154        Ok(CommandTemplate {
 155            program: "ssh".into(),
 156            args: vec![
 157                "-N".into(),
 158                "-L".into(),
 159                format!("{local_port}:{host}:{remote_port}"),
 160            ],
 161            env: Default::default(),
 162        })
 163    }
 164
 165    fn upload_directory(
 166        &self,
 167        src_path: PathBuf,
 168        dest_path: RemotePathBuf,
 169        cx: &App,
 170    ) -> Task<Result<()>> {
 171        let mut command = util::command::new_smol_command("scp");
 172        let output = self
 173            .socket
 174            .ssh_options(&mut command)
 175            .args(
 176                self.socket
 177                    .connection_options
 178                    .port
 179                    .map(|port| vec!["-P".to_string(), port.to_string()])
 180                    .unwrap_or_default(),
 181            )
 182            .arg("-C")
 183            .arg("-r")
 184            .arg(&src_path)
 185            .arg(format!(
 186                "{}:{}",
 187                self.socket.connection_options.scp_url(),
 188                dest_path
 189            ))
 190            .output();
 191
 192        cx.background_spawn(async move {
 193            let output = output.await?;
 194
 195            anyhow::ensure!(
 196                output.status.success(),
 197                "failed to upload directory {} -> {}: {}",
 198                src_path.display(),
 199                dest_path.to_string(),
 200                String::from_utf8_lossy(&output.stderr)
 201            );
 202
 203            Ok(())
 204        })
 205    }
 206
 207    fn start_proxy(
 208        &self,
 209        unique_identifier: String,
 210        reconnect: bool,
 211        incoming_tx: UnboundedSender<Envelope>,
 212        outgoing_rx: UnboundedReceiver<Envelope>,
 213        connection_activity_tx: Sender<()>,
 214        delegate: Arc<dyn RemoteClientDelegate>,
 215        cx: &mut AsyncApp,
 216    ) -> Task<Result<i32>> {
 217        delegate.set_status(Some("Starting proxy"), cx);
 218
 219        let Some(remote_binary_path) = self.remote_binary_path.clone() else {
 220            return Task::ready(Err(anyhow!("Remote binary path not set")));
 221        };
 222
 223        let mut proxy_args = vec![];
 224        for env_var in ["RUST_LOG", "RUST_BACKTRACE", "ZED_GENERATE_MINIDUMPS"] {
 225            if let Some(value) = std::env::var(env_var).ok() {
 226                proxy_args.push(format!("{}='{}'", env_var, value));
 227            }
 228        }
 229        proxy_args.push(remote_binary_path.display(self.path_style()).into_owned());
 230        proxy_args.push("proxy".to_owned());
 231        proxy_args.push("--identifier".to_owned());
 232        proxy_args.push(unique_identifier);
 233
 234        if reconnect {
 235            proxy_args.push("--reconnect".to_owned());
 236        }
 237
 238        let ssh_proxy_process = match self
 239            .socket
 240            .ssh_command("env", &proxy_args)
 241            // IMPORTANT: we kill this process when we drop the task that uses it.
 242            .kill_on_drop(true)
 243            .spawn()
 244        {
 245            Ok(process) => process,
 246            Err(error) => {
 247                return Task::ready(Err(anyhow!("failed to spawn remote server: {}", error)));
 248            }
 249        };
 250
 251        super::handle_rpc_messages_over_child_process_stdio(
 252            ssh_proxy_process,
 253            incoming_tx,
 254            outgoing_rx,
 255            connection_activity_tx,
 256            cx,
 257        )
 258    }
 259
 260    fn path_style(&self) -> PathStyle {
 261        self.ssh_path_style
 262    }
 263}
 264
 265impl SshRemoteConnection {
 266    pub(crate) async fn new(
 267        connection_options: SshConnectionOptions,
 268        delegate: Arc<dyn RemoteClientDelegate>,
 269        cx: &mut AsyncApp,
 270    ) -> Result<Self> {
 271        use askpass::AskPassResult;
 272
 273        let url = connection_options.ssh_url();
 274
 275        let temp_dir = tempfile::Builder::new()
 276            .prefix("zed-ssh-session")
 277            .tempdir()?;
 278        let askpass_delegate = askpass::AskPassDelegate::new(cx, {
 279            let delegate = delegate.clone();
 280            move |prompt, tx, cx| delegate.ask_password(prompt, tx, cx)
 281        });
 282
 283        let mut askpass =
 284            askpass::AskPassSession::new(cx.background_executor(), askpass_delegate).await?;
 285
 286        delegate.set_status(Some("Connecting"), cx);
 287
 288        // Start the master SSH process, which does not do anything except for establish
 289        // the connection and keep it open, allowing other ssh commands to reuse it
 290        // via a control socket.
 291        #[cfg(not(target_os = "windows"))]
 292        let socket_path = temp_dir.path().join("ssh.sock");
 293
 294        let mut master_process = {
 295            #[cfg(not(target_os = "windows"))]
 296            let args = [
 297                "-N",
 298                "-o",
 299                "ControlPersist=no",
 300                "-o",
 301                "ControlMaster=yes",
 302                "-o",
 303            ];
 304            // On Windows, `ControlMaster` and `ControlPath` are not supported:
 305            // https://github.com/PowerShell/Win32-OpenSSH/issues/405
 306            // https://github.com/PowerShell/Win32-OpenSSH/wiki/Project-Scope
 307            #[cfg(target_os = "windows")]
 308            let args = ["-N"];
 309            let mut master_process = util::command::new_smol_command("ssh");
 310            master_process
 311                .kill_on_drop(true)
 312                .stdin(Stdio::null())
 313                .stdout(Stdio::piped())
 314                .stderr(Stdio::piped())
 315                .env("SSH_ASKPASS_REQUIRE", "force")
 316                .env("SSH_ASKPASS", askpass.script_path())
 317                .args(connection_options.additional_args())
 318                .args(args);
 319            #[cfg(not(target_os = "windows"))]
 320            master_process.arg(format!("ControlPath={}", socket_path.display()));
 321            master_process.arg(&url).spawn()?
 322        };
 323        // Wait for this ssh process to close its stdout, indicating that authentication
 324        // has completed.
 325        let mut stdout = master_process.stdout.take().unwrap();
 326        let mut output = Vec::new();
 327
 328        let result = select_biased! {
 329            result = askpass.run().fuse() => {
 330                match result {
 331                    AskPassResult::CancelledByUser => {
 332                        master_process.kill().ok();
 333                        anyhow::bail!("SSH connection canceled")
 334                    }
 335                    AskPassResult::Timedout => {
 336                        anyhow::bail!("connecting to host timed out")
 337                    }
 338                }
 339            }
 340            _ = stdout.read_to_end(&mut output).fuse() => {
 341                anyhow::Ok(())
 342            }
 343        };
 344
 345        if let Err(e) = result {
 346            return Err(e.context("Failed to connect to host"));
 347        }
 348
 349        if master_process.try_status()?.is_some() {
 350            output.clear();
 351            let mut stderr = master_process.stderr.take().unwrap();
 352            stderr.read_to_end(&mut output).await?;
 353
 354            let error_message = format!(
 355                "failed to connect: {}",
 356                String::from_utf8_lossy(&output).trim()
 357            );
 358            anyhow::bail!(error_message);
 359        }
 360
 361        #[cfg(not(target_os = "windows"))]
 362        let socket = SshSocket::new(connection_options, socket_path).await?;
 363        #[cfg(target_os = "windows")]
 364        let socket = SshSocket::new(
 365            connection_options,
 366            askpass
 367                .get_password()
 368                .or_else(|| askpass::EncryptedPassword::try_from("").ok())
 369                .context("Failed to fetch askpass password")?,
 370            cx.background_executor().clone(),
 371        )
 372        .await?;
 373        drop(askpass);
 374
 375        let ssh_platform = socket.platform().await?;
 376        let ssh_path_style = match ssh_platform.os {
 377            "windows" => PathStyle::Windows,
 378            _ => PathStyle::Posix,
 379        };
 380        let ssh_shell = socket.shell().await;
 381        let ssh_default_system_shell = String::from("/bin/sh");
 382
 383        let mut this = Self {
 384            socket,
 385            master_process: Mutex::new(Some(master_process)),
 386            _temp_dir: temp_dir,
 387            remote_binary_path: None,
 388            ssh_path_style,
 389            ssh_platform,
 390            ssh_shell,
 391            ssh_default_system_shell,
 392        };
 393
 394        let (release_channel, version, commit) = cx.update(|cx| {
 395            (
 396                ReleaseChannel::global(cx),
 397                AppVersion::global(cx),
 398                AppCommitSha::try_global(cx),
 399            )
 400        })?;
 401        this.remote_binary_path = Some(
 402            this.ensure_server_binary(&delegate, release_channel, version, commit, cx)
 403                .await?,
 404        );
 405
 406        Ok(this)
 407    }
 408
 409    async fn ensure_server_binary(
 410        &self,
 411        delegate: &Arc<dyn RemoteClientDelegate>,
 412        release_channel: ReleaseChannel,
 413        version: SemanticVersion,
 414        commit: Option<AppCommitSha>,
 415        cx: &mut AsyncApp,
 416    ) -> Result<Arc<RelPath>> {
 417        let version_str = match release_channel {
 418            ReleaseChannel::Nightly => {
 419                let commit = commit.map(|s| s.full()).unwrap_or_default();
 420                format!("{}-{}", version, commit)
 421            }
 422            ReleaseChannel::Dev => "build".to_string(),
 423            _ => version.to_string(),
 424        };
 425        let binary_name = format!(
 426            "zed-remote-server-{}-{}",
 427            release_channel.dev_name(),
 428            version_str
 429        );
 430        let dst_path =
 431            paths::remote_server_dir_relative().join(RelPath::unix(&binary_name).unwrap());
 432
 433        #[cfg(debug_assertions)]
 434        if let Some(remote_server_path) =
 435            super::build_remote_server_from_source(&self.ssh_platform, delegate.as_ref(), cx)
 436                .await?
 437        {
 438            let tmp_path = paths::remote_server_dir_relative().join(
 439                RelPath::unix(&format!(
 440                    "download-{}-{}",
 441                    std::process::id(),
 442                    remote_server_path.file_name().unwrap().to_string_lossy()
 443                ))
 444                .unwrap(),
 445            );
 446            self.upload_local_server_binary(&remote_server_path, &tmp_path, delegate, cx)
 447                .await?;
 448            self.extract_server_binary(&dst_path, &tmp_path, delegate, cx)
 449                .await?;
 450            return Ok(dst_path);
 451        }
 452
 453        if self
 454            .socket
 455            .run_command(&dst_path.display(self.path_style()), &["version"])
 456            .await
 457            .is_ok()
 458        {
 459            return Ok(dst_path);
 460        }
 461
 462        let wanted_version = cx.update(|cx| match release_channel {
 463            ReleaseChannel::Nightly => Ok(None),
 464            ReleaseChannel::Dev => {
 465                anyhow::bail!(
 466                    "ZED_BUILD_REMOTE_SERVER is not set and no remote server exists at ({:?})",
 467                    dst_path
 468                )
 469            }
 470            _ => Ok(Some(AppVersion::global(cx))),
 471        })??;
 472
 473        let tmp_path_gz = remote_server_dir_relative().join(
 474            RelPath::unix(&format!(
 475                "{}-download-{}.gz",
 476                binary_name,
 477                std::process::id()
 478            ))
 479            .unwrap(),
 480        );
 481        if !self.socket.connection_options.upload_binary_over_ssh
 482            && let Some((url, body)) = delegate
 483                .get_download_params(self.ssh_platform, release_channel, wanted_version, cx)
 484                .await?
 485        {
 486            match self
 487                .download_binary_on_server(&url, &body, &tmp_path_gz, delegate, cx)
 488                .await
 489            {
 490                Ok(_) => {
 491                    self.extract_server_binary(&dst_path, &tmp_path_gz, delegate, cx)
 492                        .await?;
 493                    return Ok(dst_path);
 494                }
 495                Err(e) => {
 496                    log::error!(
 497                        "Failed to download binary on server, attempting to upload server: {}",
 498                        e
 499                    )
 500                }
 501            }
 502        }
 503
 504        let src_path = delegate
 505            .download_server_binary_locally(self.ssh_platform, release_channel, wanted_version, cx)
 506            .await?;
 507        self.upload_local_server_binary(&src_path, &tmp_path_gz, delegate, cx)
 508            .await?;
 509        self.extract_server_binary(&dst_path, &tmp_path_gz, delegate, cx)
 510            .await?;
 511        Ok(dst_path)
 512    }
 513
 514    async fn download_binary_on_server(
 515        &self,
 516        url: &str,
 517        body: &str,
 518        tmp_path_gz: &RelPath,
 519        delegate: &Arc<dyn RemoteClientDelegate>,
 520        cx: &mut AsyncApp,
 521    ) -> Result<()> {
 522        if let Some(parent) = tmp_path_gz.parent() {
 523            self.socket
 524                .run_command("mkdir", &["-p", parent.display(self.path_style()).as_ref()])
 525                .await?;
 526        }
 527
 528        delegate.set_status(Some("Downloading remote development server on host"), cx);
 529
 530        match self
 531            .socket
 532            .run_command(
 533                "curl",
 534                &[
 535                    "-f",
 536                    "-L",
 537                    "-X",
 538                    "GET",
 539                    "-H",
 540                    "Content-Type: application/json",
 541                    "-d",
 542                    body,
 543                    url,
 544                    "-o",
 545                    &tmp_path_gz.display(self.path_style()),
 546                ],
 547            )
 548            .await
 549        {
 550            Ok(_) => {}
 551            Err(e) => {
 552                if self.socket.run_command("which", &["curl"]).await.is_ok() {
 553                    return Err(e);
 554                }
 555
 556                match self
 557                    .socket
 558                    .run_command(
 559                        "wget",
 560                        &[
 561                            "--header=Content-Type: application/json",
 562                            "--body-data",
 563                            body,
 564                            url,
 565                            "-O",
 566                            &tmp_path_gz.display(self.path_style()),
 567                        ],
 568                    )
 569                    .await
 570                {
 571                    Ok(_) => {}
 572                    Err(e) => {
 573                        if self.socket.run_command("which", &["wget"]).await.is_ok() {
 574                            return Err(e);
 575                        } else {
 576                            anyhow::bail!("Neither curl nor wget is available");
 577                        }
 578                    }
 579                }
 580            }
 581        }
 582
 583        Ok(())
 584    }
 585
 586    async fn upload_local_server_binary(
 587        &self,
 588        src_path: &Path,
 589        tmp_path_gz: &RelPath,
 590        delegate: &Arc<dyn RemoteClientDelegate>,
 591        cx: &mut AsyncApp,
 592    ) -> Result<()> {
 593        if let Some(parent) = tmp_path_gz.parent() {
 594            self.socket
 595                .run_command("mkdir", &["-p", parent.display(self.path_style()).as_ref()])
 596                .await?;
 597        }
 598
 599        let src_stat = fs::metadata(&src_path).await?;
 600        let size = src_stat.len();
 601
 602        let t0 = Instant::now();
 603        delegate.set_status(Some("Uploading remote development server"), cx);
 604        log::info!(
 605            "uploading remote development server to {:?} ({}kb)",
 606            tmp_path_gz,
 607            size / 1024
 608        );
 609        self.upload_file(src_path, tmp_path_gz)
 610            .await
 611            .context("failed to upload server binary")?;
 612        log::info!("uploaded remote development server in {:?}", t0.elapsed());
 613        Ok(())
 614    }
 615
 616    async fn extract_server_binary(
 617        &self,
 618        dst_path: &RelPath,
 619        tmp_path: &RelPath,
 620        delegate: &Arc<dyn RemoteClientDelegate>,
 621        cx: &mut AsyncApp,
 622    ) -> Result<()> {
 623        delegate.set_status(Some("Extracting remote development server"), cx);
 624        let server_mode = 0o755;
 625
 626        let orig_tmp_path = tmp_path.display(self.path_style());
 627        let script = if let Some(tmp_path) = orig_tmp_path.strip_suffix(".gz") {
 628            shell_script!(
 629                "gunzip -f {orig_tmp_path} && chmod {server_mode} {tmp_path} && mv {tmp_path} {dst_path}",
 630                server_mode = &format!("{:o}", server_mode),
 631                dst_path = &dst_path.display(self.path_style()),
 632            )
 633        } else {
 634            shell_script!(
 635                "chmod {server_mode} {orig_tmp_path} && mv {orig_tmp_path} {dst_path}",
 636                server_mode = &format!("{:o}", server_mode),
 637                dst_path = &dst_path.display(self.path_style())
 638            )
 639        };
 640        self.socket.run_command("sh", &["-c", &script]).await?;
 641        Ok(())
 642    }
 643
 644    async fn upload_file(&self, src_path: &Path, dest_path: &RelPath) -> Result<()> {
 645        log::debug!("uploading file {:?} to {:?}", src_path, dest_path);
 646        let mut command = util::command::new_smol_command("scp");
 647        let output = self
 648            .socket
 649            .ssh_options(&mut command)
 650            .args(
 651                self.socket
 652                    .connection_options
 653                    .port
 654                    .map(|port| vec!["-P".to_string(), port.to_string()])
 655                    .unwrap_or_default(),
 656            )
 657            .arg(src_path)
 658            .arg(format!(
 659                "{}:{}",
 660                self.socket.connection_options.scp_url(),
 661                dest_path.display(self.path_style())
 662            ))
 663            .output()
 664            .await?;
 665
 666        anyhow::ensure!(
 667            output.status.success(),
 668            "failed to upload file {} -> {}: {}",
 669            src_path.display(),
 670            dest_path.display(self.path_style()),
 671            String::from_utf8_lossy(&output.stderr)
 672        );
 673        Ok(())
 674    }
 675}
 676
 677impl SshSocket {
 678    #[cfg(not(target_os = "windows"))]
 679    async fn new(options: SshConnectionOptions, socket_path: PathBuf) -> Result<Self> {
 680        Ok(Self {
 681            connection_options: options,
 682            envs: HashMap::default(),
 683            socket_path,
 684        })
 685    }
 686
 687    #[cfg(target_os = "windows")]
 688    async fn new(
 689        options: SshConnectionOptions,
 690        password: askpass::EncryptedPassword,
 691        executor: gpui::BackgroundExecutor,
 692    ) -> Result<Self> {
 693        let mut envs = HashMap::default();
 694        let get_password =
 695            move |_| Task::ready(std::ops::ControlFlow::Continue(Ok(password.clone())));
 696
 697        let _proxy = askpass::PasswordProxy::new(get_password, executor).await?;
 698        envs.insert("SSH_ASKPASS_REQUIRE".into(), "force".into());
 699        envs.insert(
 700            "SSH_ASKPASS".into(),
 701            _proxy.script_path().as_ref().display().to_string(),
 702        );
 703
 704        Ok(Self {
 705            connection_options: options,
 706            envs,
 707            _proxy,
 708        })
 709    }
 710
 711    // :WARNING: ssh unquotes arguments when executing on the remote :WARNING:
 712    // e.g. $ ssh host sh -c 'ls -l' is equivalent to $ ssh host sh -c ls -l
 713    // and passes -l as an argument to sh, not to ls.
 714    // Furthermore, some setups (e.g. Coder) will change directory when SSH'ing
 715    // into a machine. You must use `cd` to get back to $HOME.
 716    // You need to do it like this: $ ssh host "cd; sh -c 'ls -l /tmp'"
 717    fn ssh_command(&self, program: &str, args: &[impl AsRef<str>]) -> process::Command {
 718        let mut command = util::command::new_smol_command("ssh");
 719        let mut to_run = shlex::try_quote(program).unwrap().into_owned();
 720        for arg in args {
 721            // We're trying to work with: sh, bash, zsh, fish, tcsh, ...?
 722            debug_assert!(
 723                !arg.as_ref().contains('\n'),
 724                "multiline arguments do not work in all shells"
 725            );
 726            to_run.push(' ');
 727            to_run.push_str(&shlex::try_quote(arg.as_ref()).unwrap());
 728        }
 729        let to_run = format!("cd; {to_run}");
 730        self.ssh_options(&mut command)
 731            .arg(self.connection_options.ssh_url())
 732            .arg("-T")
 733            .arg(to_run);
 734        log::debug!("ssh {:?}", command);
 735        command
 736    }
 737
 738    async fn run_command(&self, program: &str, args: &[&str]) -> Result<String> {
 739        let output = self.ssh_command(program, args).output().await?;
 740        anyhow::ensure!(
 741            output.status.success(),
 742            "failed to run command: {}",
 743            String::from_utf8_lossy(&output.stderr)
 744        );
 745        Ok(String::from_utf8_lossy(&output.stdout).to_string())
 746    }
 747
 748    #[cfg(not(target_os = "windows"))]
 749    fn ssh_options<'a>(&self, command: &'a mut process::Command) -> &'a mut process::Command {
 750        command
 751            .stdin(Stdio::piped())
 752            .stdout(Stdio::piped())
 753            .stderr(Stdio::piped())
 754            .args(self.connection_options.additional_args())
 755            .args(["-o", "ControlMaster=no", "-o"])
 756            .arg(format!("ControlPath={}", self.socket_path.display()))
 757    }
 758
 759    #[cfg(target_os = "windows")]
 760    fn ssh_options<'a>(&self, command: &'a mut process::Command) -> &'a mut process::Command {
 761        command
 762            .stdin(Stdio::piped())
 763            .stdout(Stdio::piped())
 764            .stderr(Stdio::piped())
 765            .args(self.connection_options.additional_args())
 766            .envs(self.envs.clone())
 767    }
 768
 769    // On Windows, we need to use `SSH_ASKPASS` to provide the password to ssh.
 770    // On Linux, we use the `ControlPath` option to create a socket file that ssh can use to
 771    #[cfg(not(target_os = "windows"))]
 772    fn ssh_args(&self) -> Vec<String> {
 773        let mut arguments = self.connection_options.additional_args();
 774        arguments.extend(vec![
 775            "-o".to_string(),
 776            "ControlMaster=no".to_string(),
 777            "-o".to_string(),
 778            format!("ControlPath={}", self.socket_path.display()),
 779            self.connection_options.ssh_url(),
 780        ]);
 781        arguments
 782    }
 783
 784    #[cfg(target_os = "windows")]
 785    fn ssh_args(&self) -> Vec<String> {
 786        let mut arguments = self.connection_options.additional_args();
 787        arguments.push(self.connection_options.ssh_url());
 788        arguments
 789    }
 790
 791    async fn platform(&self) -> Result<RemotePlatform> {
 792        let uname = self.run_command("uname", &["-sm"]).await?;
 793        let Some((os, arch)) = uname.split_once(" ") else {
 794            anyhow::bail!("unknown uname: {uname:?}")
 795        };
 796
 797        let os = match os.trim() {
 798            "Darwin" => "macos",
 799            "Linux" => "linux",
 800            _ => anyhow::bail!(
 801                "Prebuilt remote servers are not yet available for {os:?}. See https://zed.dev/docs/remote-development"
 802            ),
 803        };
 804        // exclude armv5,6,7 as they are 32-bit.
 805        let arch = if arch.starts_with("armv8")
 806            || arch.starts_with("armv9")
 807            || arch.starts_with("arm64")
 808            || arch.starts_with("aarch64")
 809        {
 810            "aarch64"
 811        } else if arch.starts_with("x86") {
 812            "x86_64"
 813        } else {
 814            anyhow::bail!(
 815                "Prebuilt remote servers are not yet available for {arch:?}. See https://zed.dev/docs/remote-development"
 816            )
 817        };
 818
 819        Ok(RemotePlatform { os, arch })
 820    }
 821
 822    async fn shell(&self) -> String {
 823        match self.run_command("sh", &["-c", "echo $SHELL"]).await {
 824            Ok(shell) => shell.trim().to_owned(),
 825            Err(e) => {
 826                log::error!("Failed to get shell: {e}");
 827                "sh".to_owned()
 828            }
 829        }
 830    }
 831}
 832
 833fn parse_port_number(port_str: &str) -> Result<u16> {
 834    port_str
 835        .parse()
 836        .with_context(|| format!("parsing port number: {port_str}"))
 837}
 838
 839fn parse_port_forward_spec(spec: &str) -> Result<SshPortForwardOption> {
 840    let parts: Vec<&str> = spec.split(':').collect();
 841
 842    match parts.len() {
 843        4 => {
 844            let local_port = parse_port_number(parts[1])?;
 845            let remote_port = parse_port_number(parts[3])?;
 846
 847            Ok(SshPortForwardOption {
 848                local_host: Some(parts[0].to_string()),
 849                local_port,
 850                remote_host: Some(parts[2].to_string()),
 851                remote_port,
 852            })
 853        }
 854        3 => {
 855            let local_port = parse_port_number(parts[0])?;
 856            let remote_port = parse_port_number(parts[2])?;
 857
 858            Ok(SshPortForwardOption {
 859                local_host: None,
 860                local_port,
 861                remote_host: Some(parts[1].to_string()),
 862                remote_port,
 863            })
 864        }
 865        _ => anyhow::bail!("Invalid port forward format"),
 866    }
 867}
 868
 869impl SshConnectionOptions {
 870    pub fn parse_command_line(input: &str) -> Result<Self> {
 871        let input = input.trim_start_matches("ssh ");
 872        let mut hostname: Option<String> = None;
 873        let mut username: Option<String> = None;
 874        let mut port: Option<u16> = None;
 875        let mut args = Vec::new();
 876        let mut port_forwards: Vec<SshPortForwardOption> = Vec::new();
 877
 878        // disallowed: -E, -e, -F, -f, -G, -g, -M, -N, -n, -O, -q, -S, -s, -T, -t, -V, -v, -W
 879        const ALLOWED_OPTS: &[&str] = &[
 880            "-4", "-6", "-A", "-a", "-C", "-K", "-k", "-X", "-x", "-Y", "-y",
 881        ];
 882        const ALLOWED_ARGS: &[&str] = &[
 883            "-B", "-b", "-c", "-D", "-F", "-I", "-i", "-J", "-l", "-m", "-o", "-P", "-p", "-R",
 884            "-w",
 885        ];
 886
 887        let mut tokens = shlex::split(input).context("invalid input")?.into_iter();
 888
 889        'outer: while let Some(arg) = tokens.next() {
 890            if ALLOWED_OPTS.contains(&(&arg as &str)) {
 891                args.push(arg.to_string());
 892                continue;
 893            }
 894            if arg == "-p" {
 895                port = tokens.next().and_then(|arg| arg.parse().ok());
 896                continue;
 897            } else if let Some(p) = arg.strip_prefix("-p") {
 898                port = p.parse().ok();
 899                continue;
 900            }
 901            if arg == "-l" {
 902                username = tokens.next();
 903                continue;
 904            } else if let Some(l) = arg.strip_prefix("-l") {
 905                username = Some(l.to_string());
 906                continue;
 907            }
 908            if arg == "-L" || arg.starts_with("-L") {
 909                let forward_spec = if arg == "-L" {
 910                    tokens.next()
 911                } else {
 912                    Some(arg.strip_prefix("-L").unwrap().to_string())
 913                };
 914
 915                if let Some(spec) = forward_spec {
 916                    port_forwards.push(parse_port_forward_spec(&spec)?);
 917                } else {
 918                    anyhow::bail!("Missing port forward format");
 919                }
 920            }
 921
 922            for a in ALLOWED_ARGS {
 923                if arg == *a {
 924                    args.push(arg);
 925                    if let Some(next) = tokens.next() {
 926                        args.push(next);
 927                    }
 928                    continue 'outer;
 929                } else if arg.starts_with(a) {
 930                    args.push(arg);
 931                    continue 'outer;
 932                }
 933            }
 934            if arg.starts_with("-") || hostname.is_some() {
 935                anyhow::bail!("unsupported argument: {:?}", arg);
 936            }
 937            let mut input = &arg as &str;
 938            // Destination might be: username1@username2@ip2@ip1
 939            if let Some((u, rest)) = input.rsplit_once('@') {
 940                input = rest;
 941                username = Some(u.to_string());
 942            }
 943            if let Some((rest, p)) = input.split_once(':') {
 944                input = rest;
 945                port = p.parse().ok()
 946            }
 947            hostname = Some(input.to_string())
 948        }
 949
 950        let Some(hostname) = hostname else {
 951            anyhow::bail!("missing hostname");
 952        };
 953
 954        let port_forwards = match port_forwards.len() {
 955            0 => None,
 956            _ => Some(port_forwards),
 957        };
 958
 959        Ok(Self {
 960            host: hostname,
 961            username,
 962            port,
 963            port_forwards,
 964            args: Some(args),
 965            password: None,
 966            nickname: None,
 967            upload_binary_over_ssh: false,
 968        })
 969    }
 970
 971    pub fn ssh_url(&self) -> String {
 972        let mut result = String::from("ssh://");
 973        if let Some(username) = &self.username {
 974            // Username might be: username1@username2@ip2
 975            let username = urlencoding::encode(username);
 976            result.push_str(&username);
 977            result.push('@');
 978        }
 979        result.push_str(&self.host);
 980        if let Some(port) = self.port {
 981            result.push(':');
 982            result.push_str(&port.to_string());
 983        }
 984        result
 985    }
 986
 987    pub fn additional_args(&self) -> Vec<String> {
 988        let mut args = self.args.iter().flatten().cloned().collect::<Vec<String>>();
 989
 990        if let Some(forwards) = &self.port_forwards {
 991            args.extend(forwards.iter().map(|pf| {
 992                let local_host = match &pf.local_host {
 993                    Some(host) => host,
 994                    None => "localhost",
 995                };
 996                let remote_host = match &pf.remote_host {
 997                    Some(host) => host,
 998                    None => "localhost",
 999                };
1000
1001                format!(
1002                    "-L{}:{}:{}:{}",
1003                    local_host, pf.local_port, remote_host, pf.remote_port
1004                )
1005            }));
1006        }
1007
1008        args
1009    }
1010
1011    fn scp_url(&self) -> String {
1012        if let Some(username) = &self.username {
1013            format!("{}@{}", username, self.host)
1014        } else {
1015            self.host.clone()
1016        }
1017    }
1018
1019    pub fn connection_string(&self) -> String {
1020        let host = if let Some(username) = &self.username {
1021            format!("{}@{}", username, self.host)
1022        } else {
1023            self.host.clone()
1024        };
1025        if let Some(port) = &self.port {
1026            format!("{}:{}", host, port)
1027        } else {
1028            host
1029        }
1030    }
1031}
1032
1033fn build_command(
1034    input_program: Option<String>,
1035    input_args: &[String],
1036    input_env: &HashMap<String, String>,
1037    working_dir: Option<String>,
1038    port_forward: Option<(u16, String, u16)>,
1039    ssh_env: HashMap<String, String>,
1040    ssh_path_style: PathStyle,
1041    ssh_shell: &str,
1042    ssh_args: Vec<String>,
1043) -> Result<CommandTemplate> {
1044    use std::fmt::Write as _;
1045
1046    let mut exec = String::new();
1047    if let Some(working_dir) = working_dir {
1048        let working_dir = RemotePathBuf::new(working_dir, ssh_path_style).to_string();
1049
1050        // shlex will wrap the command in single quotes (''), disabling ~ expansion,
1051        // replace with with something that works
1052        const TILDE_PREFIX: &'static str = "~/";
1053        if working_dir.starts_with(TILDE_PREFIX) {
1054            let working_dir = working_dir.trim_start_matches("~").trim_start_matches("/");
1055            write!(exec, "cd \"$HOME/{working_dir}\" && ",).unwrap();
1056        } else {
1057            write!(exec, "cd \"{working_dir}\" && ",).unwrap();
1058        }
1059    } else {
1060        write!(exec, "cd && ").unwrap();
1061    };
1062    write!(exec, "exec env ").unwrap();
1063
1064    for (k, v) in input_env.iter() {
1065        if let Some((k, v)) = shlex::try_quote(k).ok().zip(shlex::try_quote(v).ok()) {
1066            write!(exec, "{}={} ", k, v).unwrap();
1067        }
1068    }
1069
1070    if let Some(input_program) = input_program {
1071        write!(exec, "{}", shlex::try_quote(&input_program).unwrap()).unwrap();
1072        for arg in input_args {
1073            let arg = shlex::try_quote(&arg)?;
1074            write!(exec, " {}", &arg).unwrap();
1075        }
1076    } else {
1077        write!(exec, "{ssh_shell} -l").unwrap();
1078    };
1079
1080    let mut args = Vec::new();
1081    args.extend(ssh_args);
1082
1083    if let Some((local_port, host, remote_port)) = port_forward {
1084        args.push("-L".into());
1085        args.push(format!("{local_port}:{host}:{remote_port}"));
1086    }
1087
1088    args.push("-t".into());
1089    args.push(exec);
1090    Ok(CommandTemplate {
1091        program: "ssh".into(),
1092        args,
1093        env: ssh_env,
1094    })
1095}
1096
1097#[cfg(test)]
1098mod tests {
1099    use super::*;
1100
1101    #[test]
1102    fn test_build_command() -> Result<()> {
1103        let mut input_env = HashMap::default();
1104        input_env.insert("INPUT_VA".to_string(), "val".to_string());
1105        let mut env = HashMap::default();
1106        env.insert("SSH_VAR".to_string(), "ssh-val".to_string());
1107
1108        let command = build_command(
1109            Some("remote_program".to_string()),
1110            &["arg1".to_string(), "arg2".to_string()],
1111            &input_env,
1112            Some("~/work".to_string()),
1113            None,
1114            env.clone(),
1115            PathStyle::Posix,
1116            "/bin/fish",
1117            vec!["-p".to_string(), "2222".to_string()],
1118        )?;
1119
1120        assert_eq!(command.program, "ssh");
1121        assert_eq!(
1122            command.args.iter().map(String::as_str).collect::<Vec<_>>(),
1123            [
1124                "-p",
1125                "2222",
1126                "-t",
1127                "cd \"$HOME/work\" && exec env INPUT_VA=val remote_program arg1 arg2"
1128            ]
1129        );
1130        assert_eq!(command.env, env);
1131
1132        let mut input_env = HashMap::default();
1133        input_env.insert("INPUT_VA".to_string(), "val".to_string());
1134        let mut env = HashMap::default();
1135        env.insert("SSH_VAR".to_string(), "ssh-val".to_string());
1136
1137        let command = build_command(
1138            None,
1139            &["arg1".to_string(), "arg2".to_string()],
1140            &input_env,
1141            None,
1142            Some((1, "foo".to_owned(), 2)),
1143            env.clone(),
1144            PathStyle::Posix,
1145            "/bin/fish",
1146            vec!["-p".to_string(), "2222".to_string()],
1147        )?;
1148
1149        assert_eq!(command.program, "ssh");
1150        assert_eq!(
1151            command.args.iter().map(String::as_str).collect::<Vec<_>>(),
1152            [
1153                "-p",
1154                "2222",
1155                "-L",
1156                "1:foo:2",
1157                "-t",
1158                "cd && exec env INPUT_VA=val /bin/fish -l"
1159            ]
1160        );
1161        assert_eq!(command.env, env);
1162
1163        Ok(())
1164    }
1165}