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