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