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