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