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