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