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