ssh.rs

   1use crate::{
   2    RemoteArch, RemoteClientDelegate, RemoteOs, RemotePlatform,
   3    remote_client::{CommandTemplate, Interactive, RemoteConnection, RemoteConnectionOptions},
   4    transport::{parse_platform, parse_shell},
   5};
   6use anyhow::{Context as _, Result, anyhow};
   7use async_trait::async_trait;
   8use collections::HashMap;
   9use futures::{
  10    AsyncReadExt as _, FutureExt as _,
  11    channel::mpsc::{Sender, UnboundedReceiver, UnboundedSender},
  12    select_biased,
  13};
  14use gpui::{App, AppContext as _, AsyncApp, Task};
  15use parking_lot::Mutex;
  16use paths::remote_server_dir_relative;
  17use release_channel::{AppVersion, ReleaseChannel};
  18use rpc::proto::Envelope;
  19use semver::Version;
  20pub use settings::SshPortForwardOption;
  21use smol::{
  22    fs,
  23    process::{self, Child, Stdio},
  24};
  25use std::{
  26    net::IpAddr,
  27    path::{Path, PathBuf},
  28    sync::Arc,
  29    time::Instant,
  30};
  31use tempfile::TempDir;
  32use util::{
  33    paths::{PathStyle, RemotePathBuf},
  34    rel_path::RelPath,
  35    shell::ShellKind,
  36};
  37
  38pub(crate) struct SshRemoteConnection {
  39    socket: SshSocket,
  40    master_process: Mutex<Option<MasterProcess>>,
  41    remote_binary_path: Option<Arc<RelPath>>,
  42    ssh_platform: RemotePlatform,
  43    ssh_path_style: PathStyle,
  44    ssh_shell: String,
  45    ssh_shell_kind: ShellKind,
  46    ssh_default_system_shell: String,
  47    _temp_dir: TempDir,
  48}
  49
  50#[derive(Debug, Clone, PartialEq, Eq, Hash)]
  51pub enum SshConnectionHost {
  52    IpAddr(IpAddr),
  53    Hostname(String),
  54}
  55
  56impl SshConnectionHost {
  57    pub fn to_bracketed_string(&self) -> String {
  58        match self {
  59            Self::IpAddr(IpAddr::V4(ip)) => ip.to_string(),
  60            Self::IpAddr(IpAddr::V6(ip)) => format!("[{}]", ip),
  61            Self::Hostname(hostname) => hostname.clone(),
  62        }
  63    }
  64
  65    pub fn to_string(&self) -> String {
  66        match self {
  67            Self::IpAddr(ip) => ip.to_string(),
  68            Self::Hostname(hostname) => hostname.clone(),
  69        }
  70    }
  71}
  72
  73impl From<&str> for SshConnectionHost {
  74    fn from(value: &str) -> Self {
  75        if let Ok(address) = value.parse() {
  76            Self::IpAddr(address)
  77        } else {
  78            Self::Hostname(value.to_string())
  79        }
  80    }
  81}
  82
  83impl From<String> for SshConnectionHost {
  84    fn from(value: String) -> Self {
  85        if let Ok(address) = value.parse() {
  86            Self::IpAddr(address)
  87        } else {
  88            Self::Hostname(value)
  89        }
  90    }
  91}
  92
  93impl Default for SshConnectionHost {
  94    fn default() -> Self {
  95        Self::Hostname(Default::default())
  96    }
  97}
  98
  99#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
 100pub struct SshConnectionOptions {
 101    pub host: SshConnectionHost,
 102    pub username: Option<String>,
 103    pub port: Option<u16>,
 104    pub password: Option<String>,
 105    pub args: Option<Vec<String>>,
 106    pub port_forwards: Option<Vec<SshPortForwardOption>>,
 107    pub connection_timeout: Option<u16>,
 108
 109    pub nickname: Option<String>,
 110    pub upload_binary_over_ssh: bool,
 111}
 112
 113impl From<settings::SshConnection> for SshConnectionOptions {
 114    fn from(val: settings::SshConnection) -> Self {
 115        SshConnectionOptions {
 116            host: val.host.to_string().into(),
 117            username: val.username,
 118            port: val.port,
 119            password: None,
 120            args: Some(val.args),
 121            nickname: val.nickname,
 122            upload_binary_over_ssh: val.upload_binary_over_ssh.unwrap_or_default(),
 123            port_forwards: val.port_forwards,
 124            connection_timeout: val.connection_timeout,
 125        }
 126    }
 127}
 128
 129struct SshSocket {
 130    connection_options: SshConnectionOptions,
 131    #[cfg(not(target_os = "windows"))]
 132    socket_path: std::path::PathBuf,
 133    envs: HashMap<String, String>,
 134    #[cfg(target_os = "windows")]
 135    _proxy: askpass::PasswordProxy,
 136}
 137
 138struct MasterProcess {
 139    process: Child,
 140}
 141
 142#[cfg(not(target_os = "windows"))]
 143impl MasterProcess {
 144    pub fn new(
 145        askpass_script_path: &std::ffi::OsStr,
 146        additional_args: Vec<String>,
 147        socket_path: &std::path::Path,
 148        destination: &str,
 149    ) -> Result<Self> {
 150        let args = [
 151            "-N",
 152            "-o",
 153            "ControlPersist=no",
 154            "-o",
 155            "ControlMaster=yes",
 156            "-o",
 157        ];
 158
 159        let mut master_process = util::command::new_smol_command("ssh");
 160        master_process
 161            .kill_on_drop(true)
 162            .stdin(Stdio::null())
 163            .stdout(Stdio::piped())
 164            .stderr(Stdio::piped())
 165            .env("SSH_ASKPASS_REQUIRE", "force")
 166            .env("SSH_ASKPASS", askpass_script_path)
 167            .args(additional_args)
 168            .args(args);
 169
 170        master_process.arg(format!("ControlPath={}", socket_path.display()));
 171
 172        let process = master_process.arg(&destination).spawn()?;
 173
 174        Ok(MasterProcess { process })
 175    }
 176
 177    pub async fn wait_connected(&mut self) -> Result<()> {
 178        let Some(mut stdout) = self.process.stdout.take() else {
 179            anyhow::bail!("ssh process stdout capture failed");
 180        };
 181
 182        let mut output = Vec::new();
 183        stdout.read_to_end(&mut output).await?;
 184        Ok(())
 185    }
 186}
 187
 188#[cfg(target_os = "windows")]
 189impl MasterProcess {
 190    const CONNECTION_ESTABLISHED_MAGIC: &str = "ZED_SSH_CONNECTION_ESTABLISHED";
 191
 192    pub fn new(
 193        askpass_script_path: &std::ffi::OsStr,
 194        additional_args: Vec<String>,
 195        destination: &str,
 196    ) -> Result<Self> {
 197        // On Windows, `ControlMaster` and `ControlPath` are not supported:
 198        // https://github.com/PowerShell/Win32-OpenSSH/issues/405
 199        // https://github.com/PowerShell/Win32-OpenSSH/wiki/Project-Scope
 200        //
 201        // Using an ugly workaround to detect connection establishment
 202        // -N doesn't work with JumpHosts as windows openssh never closes stdin in that case
 203        let args = [
 204            "-t",
 205            &format!("echo '{}'; exec $0", Self::CONNECTION_ESTABLISHED_MAGIC),
 206        ];
 207
 208        let mut master_process = util::command::new_smol_command("ssh");
 209        master_process
 210            .kill_on_drop(true)
 211            .stdin(Stdio::null())
 212            .stdout(Stdio::piped())
 213            .stderr(Stdio::piped())
 214            .env("SSH_ASKPASS_REQUIRE", "force")
 215            .env("SSH_ASKPASS", askpass_script_path)
 216            .args(additional_args)
 217            .arg(destination)
 218            .args(args);
 219
 220        let process = master_process.spawn()?;
 221
 222        Ok(MasterProcess { process })
 223    }
 224
 225    pub async fn wait_connected(&mut self) -> Result<()> {
 226        use smol::io::AsyncBufReadExt;
 227
 228        let Some(stdout) = self.process.stdout.take() else {
 229            anyhow::bail!("ssh process stdout capture failed");
 230        };
 231
 232        let mut reader = smol::io::BufReader::new(stdout);
 233
 234        let mut line = String::new();
 235
 236        loop {
 237            let n = reader.read_line(&mut line).await?;
 238            if n == 0 {
 239                anyhow::bail!("ssh process exited before connection established");
 240            }
 241
 242            if line.contains(Self::CONNECTION_ESTABLISHED_MAGIC) {
 243                return Ok(());
 244            }
 245        }
 246    }
 247}
 248
 249impl AsRef<Child> for MasterProcess {
 250    fn as_ref(&self) -> &Child {
 251        &self.process
 252    }
 253}
 254
 255impl AsMut<Child> for MasterProcess {
 256    fn as_mut(&mut self) -> &mut Child {
 257        &mut self.process
 258    }
 259}
 260
 261#[async_trait(?Send)]
 262impl RemoteConnection for SshRemoteConnection {
 263    async fn kill(&self) -> Result<()> {
 264        let Some(mut process) = self.master_process.lock().take() else {
 265            return Ok(());
 266        };
 267        process.as_mut().kill().ok();
 268        process.as_mut().status().await?;
 269        Ok(())
 270    }
 271
 272    fn has_been_killed(&self) -> bool {
 273        self.master_process.lock().is_none()
 274    }
 275
 276    fn connection_options(&self) -> RemoteConnectionOptions {
 277        RemoteConnectionOptions::Ssh(self.socket.connection_options.clone())
 278    }
 279
 280    fn shell(&self) -> String {
 281        self.ssh_shell.clone()
 282    }
 283
 284    fn default_system_shell(&self) -> String {
 285        self.ssh_default_system_shell.clone()
 286    }
 287
 288    fn build_command(
 289        &self,
 290        input_program: Option<String>,
 291        input_args: &[String],
 292        input_env: &HashMap<String, String>,
 293        working_dir: Option<String>,
 294        port_forward: Option<(u16, String, u16)>,
 295        interactive: Interactive,
 296    ) -> Result<CommandTemplate> {
 297        let Self {
 298            ssh_path_style,
 299            socket,
 300            ssh_shell_kind,
 301            ssh_shell,
 302            ..
 303        } = self;
 304        let env = socket.envs.clone();
 305        build_command(
 306            input_program,
 307            input_args,
 308            input_env,
 309            working_dir,
 310            port_forward,
 311            env,
 312            *ssh_path_style,
 313            ssh_shell,
 314            *ssh_shell_kind,
 315            socket.ssh_args(),
 316            interactive,
 317        )
 318    }
 319
 320    fn build_forward_ports_command(
 321        &self,
 322        forwards: Vec<(u16, String, u16)>,
 323    ) -> Result<CommandTemplate> {
 324        let Self { socket, .. } = self;
 325        let mut args = socket.ssh_args();
 326        args.push("-N".into());
 327        for (local_port, host, remote_port) in forwards {
 328            args.push("-L".into());
 329            args.push(format!("{local_port}:{host}:{remote_port}"));
 330        }
 331        Ok(CommandTemplate {
 332            program: "ssh".into(),
 333            args,
 334            env: Default::default(),
 335        })
 336    }
 337
 338    fn upload_directory(
 339        &self,
 340        src_path: PathBuf,
 341        dest_path: RemotePathBuf,
 342        cx: &App,
 343    ) -> Task<Result<()>> {
 344        let dest_path_str = dest_path.to_string();
 345        let src_path_display = src_path.display().to_string();
 346
 347        let mut sftp_command = self.build_sftp_command();
 348        let mut scp_command =
 349            self.build_scp_command(&src_path, &dest_path_str, Some(&["-C", "-r"]));
 350
 351        cx.background_spawn(async move {
 352            // We will try SFTP first, and if that fails, we will fall back to SCP.
 353            // If SCP fails also, we give up and return an error.
 354            // The reason we allow a fallback from SFTP to SCP is that if the user has to specify a password,
 355            // depending on the implementation of SSH stack, SFTP may disable interactive password prompts in batch mode.
 356            // This is for example the case on Windows as evidenced by this implementation snippet:
 357            // https://github.com/PowerShell/openssh-portable/blob/b8c08ef9da9450a94a9c5ef717d96a7bd83f3332/sshconnect2.c#L417
 358            if Self::is_sftp_available().await {
 359                log::debug!("using SFTP for directory upload");
 360                let mut child = sftp_command.spawn()?;
 361                if let Some(mut stdin) = child.stdin.take() {
 362                    use futures::AsyncWriteExt;
 363                    let sftp_batch = format!("put -r \"{src_path_display}\" \"{dest_path_str}\"\n");
 364                    stdin.write_all(sftp_batch.as_bytes()).await?;
 365                    stdin.flush().await?;
 366                }
 367
 368                let output = child.output().await?;
 369                if output.status.success() {
 370                    return Ok(());
 371                }
 372
 373                let stderr = String::from_utf8_lossy(&output.stderr);
 374                log::debug!("failed to upload directory via SFTP {src_path_display} -> {dest_path_str}: {stderr}");
 375            }
 376
 377            log::debug!("using SCP for directory upload");
 378            let output = scp_command.output().await?;
 379
 380            if output.status.success() {
 381                return Ok(());
 382            }
 383
 384            let stderr = String::from_utf8_lossy(&output.stderr);
 385            log::debug!("failed to upload directory via SCP {src_path_display} -> {dest_path_str}: {stderr}");
 386
 387            anyhow::bail!(
 388                "failed to upload directory via SFTP/SCP {} -> {}: {}",
 389                src_path_display,
 390                dest_path_str,
 391                stderr,
 392            );
 393        })
 394    }
 395
 396    fn start_proxy(
 397        &self,
 398        unique_identifier: String,
 399        reconnect: bool,
 400        incoming_tx: UnboundedSender<Envelope>,
 401        outgoing_rx: UnboundedReceiver<Envelope>,
 402        connection_activity_tx: Sender<()>,
 403        delegate: Arc<dyn RemoteClientDelegate>,
 404        cx: &mut AsyncApp,
 405    ) -> Task<Result<i32>> {
 406        const VARS: [&str; 3] = ["RUST_LOG", "RUST_BACKTRACE", "ZED_GENERATE_MINIDUMPS"];
 407        delegate.set_status(Some("Starting proxy"), cx);
 408
 409        let Some(remote_binary_path) = self.remote_binary_path.clone() else {
 410            return Task::ready(Err(anyhow!("Remote binary path not set")));
 411        };
 412
 413        let mut ssh_command = if self.ssh_platform.os.is_windows() {
 414            // TODO: Set the `VARS` environment variables, we do not have `env` on windows
 415            // so this needs a different approach
 416            let mut proxy_args = vec![];
 417            proxy_args.push("proxy".to_owned());
 418            proxy_args.push("--identifier".to_owned());
 419            proxy_args.push(unique_identifier);
 420
 421            if reconnect {
 422                proxy_args.push("--reconnect".to_owned());
 423            }
 424            self.socket.ssh_command(
 425                self.ssh_shell_kind,
 426                &remote_binary_path.display(self.path_style()),
 427                &proxy_args,
 428                false,
 429            )
 430        } else {
 431            let mut proxy_args = vec![];
 432            for env_var in VARS {
 433                if let Some(value) = std::env::var(env_var).ok() {
 434                    proxy_args.push(format!("{}='{}'", env_var, value));
 435                }
 436            }
 437            proxy_args.push(remote_binary_path.display(self.path_style()).into_owned());
 438            proxy_args.push("proxy".to_owned());
 439            proxy_args.push("--identifier".to_owned());
 440            proxy_args.push(unique_identifier);
 441
 442            if reconnect {
 443                proxy_args.push("--reconnect".to_owned());
 444            }
 445            self.socket
 446                .ssh_command(self.ssh_shell_kind, "env", &proxy_args, false)
 447        };
 448
 449        let ssh_proxy_process = match ssh_command
 450            // IMPORTANT: we kill this process when we drop the task that uses it.
 451            .kill_on_drop(true)
 452            .spawn()
 453        {
 454            Ok(process) => process,
 455            Err(error) => {
 456                return Task::ready(Err(anyhow!("failed to spawn remote server: {}", error)));
 457            }
 458        };
 459
 460        super::handle_rpc_messages_over_child_process_stdio(
 461            ssh_proxy_process,
 462            incoming_tx,
 463            outgoing_rx,
 464            connection_activity_tx,
 465            cx,
 466        )
 467    }
 468
 469    fn path_style(&self) -> PathStyle {
 470        self.ssh_path_style
 471    }
 472
 473    fn has_wsl_interop(&self) -> bool {
 474        false
 475    }
 476}
 477
 478impl SshRemoteConnection {
 479    pub(crate) async fn new(
 480        connection_options: SshConnectionOptions,
 481        delegate: Arc<dyn RemoteClientDelegate>,
 482        cx: &mut AsyncApp,
 483    ) -> Result<Self> {
 484        use askpass::AskPassResult;
 485
 486        let destination = connection_options.ssh_destination();
 487
 488        let temp_dir = tempfile::Builder::new()
 489            .prefix("zed-ssh-session")
 490            .tempdir()?;
 491        let askpass_delegate = askpass::AskPassDelegate::new(cx, {
 492            let delegate = delegate.clone();
 493            move |prompt, tx, cx| delegate.ask_password(prompt, tx, cx)
 494        });
 495
 496        let mut askpass =
 497            askpass::AskPassSession::new(cx.background_executor().clone(), askpass_delegate)
 498                .await?;
 499
 500        delegate.set_status(Some("Connecting"), cx);
 501
 502        // Start the master SSH process, which does not do anything except for establish
 503        // the connection and keep it open, allowing other ssh commands to reuse it
 504        // via a control socket.
 505        #[cfg(not(target_os = "windows"))]
 506        let socket_path = temp_dir.path().join("ssh.sock");
 507
 508        #[cfg(target_os = "windows")]
 509        let mut master_process = MasterProcess::new(
 510            askpass.script_path().as_ref(),
 511            connection_options.additional_args(),
 512            &destination,
 513        )?;
 514        #[cfg(not(target_os = "windows"))]
 515        let mut master_process = MasterProcess::new(
 516            askpass.script_path().as_ref(),
 517            connection_options.additional_args(),
 518            &socket_path,
 519            &destination,
 520        )?;
 521
 522        let result = select_biased! {
 523            result = askpass.run().fuse() => {
 524                match result {
 525                    AskPassResult::CancelledByUser => {
 526                        master_process.as_mut().kill().ok();
 527                        anyhow::bail!("SSH connection canceled")
 528                    }
 529                    AskPassResult::Timedout => {
 530                        anyhow::bail!("connecting to host timed out")
 531                    }
 532                }
 533            }
 534            _ = master_process.wait_connected().fuse() => {
 535                anyhow::Ok(())
 536            }
 537        };
 538
 539        if let Err(e) = result {
 540            return Err(e.context("Failed to connect to host"));
 541        }
 542
 543        if master_process.as_mut().try_status()?.is_some() {
 544            let mut output = Vec::new();
 545            output.clear();
 546            let mut stderr = master_process.as_mut().stderr.take().unwrap();
 547            stderr.read_to_end(&mut output).await?;
 548
 549            let error_message = format!(
 550                "failed to connect: {}",
 551                String::from_utf8_lossy(&output).trim()
 552            );
 553            anyhow::bail!(error_message);
 554        }
 555
 556        #[cfg(not(target_os = "windows"))]
 557        let socket = SshSocket::new(connection_options, socket_path).await?;
 558        #[cfg(target_os = "windows")]
 559        let socket = SshSocket::new(
 560            connection_options,
 561            askpass
 562                .get_password()
 563                .or_else(|| askpass::EncryptedPassword::try_from("").ok())
 564                .context("Failed to fetch askpass password")?,
 565            cx.background_executor().clone(),
 566        )
 567        .await?;
 568        drop(askpass);
 569
 570        let is_windows = socket.probe_is_windows().await;
 571        log::info!("Remote is windows: {}", is_windows);
 572
 573        let ssh_shell = socket.shell(is_windows).await;
 574        log::info!("Remote shell discovered: {}", ssh_shell);
 575
 576        let ssh_shell_kind = ShellKind::new(&ssh_shell, is_windows);
 577        let ssh_platform = socket.platform(ssh_shell_kind, is_windows).await?;
 578        log::info!("Remote platform discovered: {:?}", ssh_platform);
 579
 580        let (ssh_path_style, ssh_default_system_shell) = match ssh_platform.os {
 581            RemoteOs::Windows => (PathStyle::Windows, ssh_shell.clone()),
 582            _ => (PathStyle::Posix, String::from("/bin/sh")),
 583        };
 584
 585        let mut this = Self {
 586            socket,
 587            master_process: Mutex::new(Some(master_process)),
 588            _temp_dir: temp_dir,
 589            remote_binary_path: None,
 590            ssh_path_style,
 591            ssh_platform,
 592            ssh_shell,
 593            ssh_shell_kind,
 594            ssh_default_system_shell,
 595        };
 596
 597        let (release_channel, version) =
 598            cx.update(|cx| (ReleaseChannel::global(cx), AppVersion::global(cx)));
 599        this.remote_binary_path = Some(
 600            this.ensure_server_binary(&delegate, release_channel, version, cx)
 601                .await?,
 602        );
 603
 604        Ok(this)
 605    }
 606
 607    async fn ensure_server_binary(
 608        &self,
 609        delegate: &Arc<dyn RemoteClientDelegate>,
 610        release_channel: ReleaseChannel,
 611        version: Version,
 612        cx: &mut AsyncApp,
 613    ) -> Result<Arc<RelPath>> {
 614        let version_str = match release_channel {
 615            ReleaseChannel::Dev => "build".to_string(),
 616            _ => version.to_string(),
 617        };
 618        let binary_name = format!(
 619            "zed-remote-server-{}-{}{}",
 620            release_channel.dev_name(),
 621            version_str,
 622            if self.ssh_platform.os.is_windows() {
 623                ".exe"
 624            } else {
 625                ""
 626            }
 627        );
 628        let dst_path =
 629            paths::remote_server_dir_relative().join(RelPath::unix(&binary_name).unwrap());
 630
 631        let binary_exists_on_server = self
 632            .socket
 633            .run_command(
 634                self.ssh_shell_kind,
 635                &dst_path.display(self.path_style()),
 636                &["version"],
 637                true,
 638            )
 639            .await
 640            .is_ok();
 641
 642        #[cfg(any(debug_assertions, feature = "build-remote-server-binary"))]
 643        if let Some(remote_server_path) = super::build_remote_server_from_source(
 644            &self.ssh_platform,
 645            delegate.as_ref(),
 646            binary_exists_on_server,
 647            cx,
 648        )
 649        .await?
 650        {
 651            let tmp_path = paths::remote_server_dir_relative().join(
 652                RelPath::unix(&format!(
 653                    "download-{}-{}",
 654                    std::process::id(),
 655                    remote_server_path.file_name().unwrap().to_string_lossy()
 656                ))
 657                .unwrap(),
 658            );
 659            self.upload_local_server_binary(&remote_server_path, &tmp_path, delegate, cx)
 660                .await?;
 661            self.extract_server_binary(&dst_path, &tmp_path, delegate, cx)
 662                .await?;
 663            return Ok(dst_path);
 664        }
 665
 666        if binary_exists_on_server {
 667            return Ok(dst_path);
 668        }
 669
 670        let wanted_version = cx.update(|cx| match release_channel {
 671            ReleaseChannel::Nightly => Ok(None),
 672            ReleaseChannel::Dev => {
 673                anyhow::bail!(
 674                    "ZED_BUILD_REMOTE_SERVER is not set and no remote server exists at ({:?})",
 675                    dst_path
 676                )
 677            }
 678            _ => Ok(Some(AppVersion::global(cx))),
 679        })?;
 680
 681        let tmp_path_compressed = remote_server_dir_relative().join(
 682            RelPath::unix(&format!(
 683                "{}-download-{}.{}",
 684                binary_name,
 685                std::process::id(),
 686                if self.ssh_platform.os.is_windows() {
 687                    "zip"
 688                } else {
 689                    "gz"
 690                }
 691            ))
 692            .unwrap(),
 693        );
 694        if !self.socket.connection_options.upload_binary_over_ssh
 695            && let Some(url) = delegate
 696                .get_download_url(
 697                    self.ssh_platform,
 698                    release_channel,
 699                    wanted_version.clone(),
 700                    cx,
 701                )
 702                .await?
 703        {
 704            match self
 705                .download_binary_on_server(&url, &tmp_path_compressed, delegate, cx)
 706                .await
 707            {
 708                Ok(_) => {
 709                    self.extract_server_binary(&dst_path, &tmp_path_compressed, delegate, cx)
 710                        .await
 711                        .context("extracting server binary")?;
 712                    return Ok(dst_path);
 713                }
 714                Err(e) => {
 715                    log::error!(
 716                        "Failed to download binary on server, attempting to download locally and then upload it the server: {e:#}",
 717                    )
 718                }
 719            }
 720        }
 721
 722        let src_path = delegate
 723            .download_server_binary_locally(
 724                self.ssh_platform,
 725                release_channel,
 726                wanted_version.clone(),
 727                cx,
 728            )
 729            .await
 730            .context("downloading server binary locally")?;
 731        self.upload_local_server_binary(&src_path, &tmp_path_compressed, delegate, cx)
 732            .await
 733            .context("uploading server binary")?;
 734        self.extract_server_binary(&dst_path, &tmp_path_compressed, delegate, cx)
 735            .await
 736            .context("extracting server binary")?;
 737        Ok(dst_path)
 738    }
 739
 740    async fn download_binary_on_server(
 741        &self,
 742        url: &str,
 743        tmp_path: &RelPath,
 744        delegate: &Arc<dyn RemoteClientDelegate>,
 745        cx: &mut AsyncApp,
 746    ) -> Result<()> {
 747        if let Some(parent) = tmp_path.parent() {
 748            let res = self
 749                .socket
 750                .run_command(
 751                    self.ssh_shell_kind,
 752                    "mkdir",
 753                    &["-p", parent.display(self.path_style()).as_ref()],
 754                    true,
 755                )
 756                .await;
 757            if !self.ssh_platform.os.is_windows() {
 758                // mkdir fails on windows if the path already exists ...
 759                res?;
 760            }
 761        }
 762
 763        delegate.set_status(Some("Downloading remote development server on host"), cx);
 764
 765        let connection_timeout = self
 766            .socket
 767            .connection_options
 768            .connection_timeout
 769            .unwrap_or(10)
 770            .to_string();
 771
 772        match self
 773            .socket
 774            .run_command(
 775                self.ssh_shell_kind,
 776                "curl",
 777                &[
 778                    "-f",
 779                    "-L",
 780                    "--connect-timeout",
 781                    &connection_timeout,
 782                    url,
 783                    "-o",
 784                    &tmp_path.display(self.path_style()),
 785                ],
 786                true,
 787            )
 788            .await
 789        {
 790            Ok(_) => {}
 791            Err(e) => {
 792                if self
 793                    .socket
 794                    .run_command(self.ssh_shell_kind, "which", &["curl"], true)
 795                    .await
 796                    .is_ok()
 797                {
 798                    return Err(e);
 799                }
 800
 801                log::info!("curl is not available, trying wget");
 802                match self
 803                    .socket
 804                    .run_command(
 805                        self.ssh_shell_kind,
 806                        "wget",
 807                        &[
 808                            "--connect-timeout",
 809                            &connection_timeout,
 810                            "--tries",
 811                            "1",
 812                            url,
 813                            "-O",
 814                            &tmp_path.display(self.path_style()),
 815                        ],
 816                        true,
 817                    )
 818                    .await
 819                {
 820                    Ok(_) => {}
 821                    Err(e) => {
 822                        if self
 823                            .socket
 824                            .run_command(self.ssh_shell_kind, "which", &["wget"], true)
 825                            .await
 826                            .is_ok()
 827                        {
 828                            return Err(e);
 829                        } else {
 830                            anyhow::bail!("Neither curl nor wget is available");
 831                        }
 832                    }
 833                }
 834            }
 835        }
 836
 837        Ok(())
 838    }
 839
 840    async fn upload_local_server_binary(
 841        &self,
 842        src_path: &Path,
 843        tmp_path: &RelPath,
 844        delegate: &Arc<dyn RemoteClientDelegate>,
 845        cx: &mut AsyncApp,
 846    ) -> Result<()> {
 847        if let Some(parent) = tmp_path.parent() {
 848            let res = self
 849                .socket
 850                .run_command(
 851                    self.ssh_shell_kind,
 852                    "mkdir",
 853                    &["-p", parent.display(self.path_style()).as_ref()],
 854                    true,
 855                )
 856                .await;
 857            if !self.ssh_platform.os.is_windows() {
 858                // mkdir fails on windows if the path already exists ...
 859                res?;
 860            }
 861        }
 862
 863        let src_stat = fs::metadata(&src_path)
 864            .await
 865            .with_context(|| format!("failed to get metadata for {:?}", src_path))?;
 866        let size = src_stat.len();
 867
 868        let t0 = Instant::now();
 869        delegate.set_status(Some("Uploading remote development server"), cx);
 870        log::info!(
 871            "uploading remote development server to {:?} ({}kb)",
 872            tmp_path,
 873            size / 1024
 874        );
 875        self.upload_file(src_path, tmp_path)
 876            .await
 877            .context("failed to upload server binary")?;
 878        log::info!("uploaded remote development server in {:?}", t0.elapsed());
 879        Ok(())
 880    }
 881
 882    async fn extract_server_binary(
 883        &self,
 884        dst_path: &RelPath,
 885        tmp_path: &RelPath,
 886        delegate: &Arc<dyn RemoteClientDelegate>,
 887        cx: &mut AsyncApp,
 888    ) -> Result<()> {
 889        delegate.set_status(Some("Extracting remote development server"), cx);
 890
 891        if self.ssh_platform.os.is_windows() {
 892            self.extract_server_binary_windows(dst_path, tmp_path).await
 893        } else {
 894            self.extract_server_binary_posix(dst_path, tmp_path).await
 895        }
 896    }
 897
 898    async fn extract_server_binary_posix(
 899        &self,
 900        dst_path: &RelPath,
 901        tmp_path: &RelPath,
 902    ) -> Result<()> {
 903        let shell_kind = ShellKind::Posix;
 904        let server_mode = 0o755;
 905        let orig_tmp_path = tmp_path.display(self.path_style());
 906        let server_mode = format!("{:o}", server_mode);
 907        let server_mode = shell_kind
 908            .try_quote(&server_mode)
 909            .context("shell quoting")?;
 910        let dst_path = dst_path.display(self.path_style());
 911        let dst_path = shell_kind.try_quote(&dst_path).context("shell quoting")?;
 912        let script = if let Some(tmp_path) = orig_tmp_path.strip_suffix(".gz") {
 913            let orig_tmp_path = shell_kind
 914                .try_quote(&orig_tmp_path)
 915                .context("shell quoting")?;
 916            let tmp_path = shell_kind.try_quote(&tmp_path).context("shell quoting")?;
 917            format!(
 918                "gunzip -f {orig_tmp_path} && chmod {server_mode} {tmp_path} && mv {tmp_path} {dst_path}",
 919            )
 920        } else {
 921            let orig_tmp_path = shell_kind
 922                .try_quote(&orig_tmp_path)
 923                .context("shell quoting")?;
 924            format!("chmod {server_mode} {orig_tmp_path} && mv {orig_tmp_path} {dst_path}",)
 925        };
 926        let args = shell_kind.args_for_shell(false, script.to_string());
 927        self.socket
 928            .run_command(self.ssh_shell_kind, "sh", &args, true)
 929            .await?;
 930        Ok(())
 931    }
 932
 933    async fn extract_server_binary_windows(
 934        &self,
 935        dst_path: &RelPath,
 936        tmp_path: &RelPath,
 937    ) -> Result<()> {
 938        let shell_kind = ShellKind::Pwsh;
 939        let orig_tmp_path = tmp_path.display(self.path_style());
 940        let dst_path = dst_path.display(self.path_style());
 941        let dst_path = shell_kind.try_quote(&dst_path).context("shell quoting")?;
 942
 943        let script = if let Some(tmp_path) = orig_tmp_path.strip_suffix(".zip") {
 944            let orig_tmp_path = shell_kind
 945                .try_quote(&orig_tmp_path)
 946                .context("shell quoting")?;
 947            let tmp_path = shell_kind.try_quote(tmp_path).context("shell quoting")?;
 948            format!(
 949                "Expand-Archive -Force -Path {orig_tmp_path} -DestinationPath {tmp_path} -ErrorAction Stop;
 950                 Move-Item -Force {tmp_path} {dst_path}",
 951            )
 952        } else {
 953            let orig_tmp_path = shell_kind
 954                .try_quote(&orig_tmp_path)
 955                .context("shell quoting")?;
 956            format!("Move-Item -Force {orig_tmp_path} {dst_path}")
 957        };
 958
 959        let args = shell_kind.args_for_shell(false, script);
 960        self.socket
 961            .run_command(self.ssh_shell_kind, "powershell", &args, true)
 962            .await?;
 963        Ok(())
 964    }
 965
 966    fn build_scp_command(
 967        &self,
 968        src_path: &Path,
 969        dest_path_str: &str,
 970        args: Option<&[&str]>,
 971    ) -> process::Command {
 972        let mut command = util::command::new_smol_command("scp");
 973        self.socket.ssh_options(&mut command, false).args(
 974            self.socket
 975                .connection_options
 976                .port
 977                .map(|port| vec!["-P".to_string(), port.to_string()])
 978                .unwrap_or_default(),
 979        );
 980        if let Some(args) = args {
 981            command.args(args);
 982        }
 983        command.arg(src_path).arg(format!(
 984            "{}:{}",
 985            self.socket.connection_options.scp_destination(),
 986            dest_path_str
 987        ));
 988        command
 989    }
 990
 991    fn build_sftp_command(&self) -> process::Command {
 992        let mut command = util::command::new_smol_command("sftp");
 993        self.socket.ssh_options(&mut command, false).args(
 994            self.socket
 995                .connection_options
 996                .port
 997                .map(|port| vec!["-P".to_string(), port.to_string()])
 998                .unwrap_or_default(),
 999        );
1000        command.arg("-b").arg("-");
1001        command.arg(self.socket.connection_options.scp_destination());
1002        command.stdin(Stdio::piped());
1003        command
1004    }
1005
1006    async fn upload_file(&self, src_path: &Path, dest_path: &RelPath) -> Result<()> {
1007        log::debug!("uploading file {:?} to {:?}", src_path, dest_path);
1008
1009        let src_path_display = src_path.display().to_string();
1010        let dest_path_str = dest_path.display(self.path_style());
1011
1012        // We will try SFTP first, and if that fails, we will fall back to SCP.
1013        // If SCP fails also, we give up and return an error.
1014        // The reason we allow a fallback from SFTP to SCP is that if the user has to specify a password,
1015        // depending on the implementation of SSH stack, SFTP may disable interactive password prompts in batch mode.
1016        // This is for example the case on Windows as evidenced by this implementation snippet:
1017        // https://github.com/PowerShell/openssh-portable/blob/b8c08ef9da9450a94a9c5ef717d96a7bd83f3332/sshconnect2.c#L417
1018        if Self::is_sftp_available().await {
1019            log::debug!("using SFTP for file upload");
1020            let mut command = self.build_sftp_command();
1021            let sftp_batch = format!("put {src_path_display} {dest_path_str}\n");
1022
1023            let mut child = command.spawn()?;
1024            if let Some(mut stdin) = child.stdin.take() {
1025                use futures::AsyncWriteExt;
1026                stdin.write_all(sftp_batch.as_bytes()).await?;
1027                stdin.flush().await?;
1028            }
1029
1030            let output = child.output().await?;
1031            if output.status.success() {
1032                return Ok(());
1033            }
1034
1035            let stderr = String::from_utf8_lossy(&output.stderr);
1036            log::debug!(
1037                "failed to upload file via SFTP {src_path_display} -> {dest_path_str}: {stderr}"
1038            );
1039        }
1040
1041        log::debug!("using SCP for file upload");
1042        let mut command = self.build_scp_command(src_path, &dest_path_str, None);
1043        let output = command.output().await?;
1044
1045        if output.status.success() {
1046            return Ok(());
1047        }
1048
1049        let stderr = String::from_utf8_lossy(&output.stderr);
1050        log::debug!(
1051            "failed to upload file via SCP {src_path_display} -> {dest_path_str}: {stderr}",
1052        );
1053        anyhow::bail!(
1054            "failed to upload file via STFP/SCP {} -> {}: {}",
1055            src_path_display,
1056            dest_path_str,
1057            stderr,
1058        );
1059    }
1060
1061    async fn is_sftp_available() -> bool {
1062        which::which("sftp").is_ok()
1063    }
1064}
1065
1066impl SshSocket {
1067    #[cfg(not(target_os = "windows"))]
1068    async fn new(options: SshConnectionOptions, socket_path: PathBuf) -> Result<Self> {
1069        Ok(Self {
1070            connection_options: options,
1071            envs: HashMap::default(),
1072            socket_path,
1073        })
1074    }
1075
1076    #[cfg(target_os = "windows")]
1077    async fn new(
1078        options: SshConnectionOptions,
1079        password: askpass::EncryptedPassword,
1080        executor: gpui::BackgroundExecutor,
1081    ) -> Result<Self> {
1082        let mut envs = HashMap::default();
1083        let get_password =
1084            move |_| Task::ready(std::ops::ControlFlow::Continue(Ok(password.clone())));
1085
1086        let _proxy = askpass::PasswordProxy::new(get_password, executor).await?;
1087        envs.insert("SSH_ASKPASS_REQUIRE".into(), "force".into());
1088        envs.insert(
1089            "SSH_ASKPASS".into(),
1090            _proxy.script_path().as_ref().display().to_string(),
1091        );
1092
1093        Ok(Self {
1094            connection_options: options,
1095            envs,
1096            _proxy,
1097        })
1098    }
1099
1100    // :WARNING: ssh unquotes arguments when executing on the remote :WARNING:
1101    // e.g. $ ssh host sh -c 'ls -l' is equivalent to $ ssh host sh -c ls -l
1102    // and passes -l as an argument to sh, not to ls.
1103    // Furthermore, some setups (e.g. Coder) will change directory when SSH'ing
1104    // into a machine. You must use `cd` to get back to $HOME.
1105    // You need to do it like this: $ ssh host "cd; sh -c 'ls -l /tmp'"
1106    fn ssh_command(
1107        &self,
1108        shell_kind: ShellKind,
1109        program: &str,
1110        args: &[impl AsRef<str>],
1111        allow_pseudo_tty: bool,
1112    ) -> process::Command {
1113        let mut command = util::command::new_smol_command("ssh");
1114        let program = shell_kind.prepend_command_prefix(program);
1115        let mut to_run = shell_kind
1116            .try_quote_prefix_aware(&program)
1117            .expect("shell quoting")
1118            .into_owned();
1119        for arg in args {
1120            // We're trying to work with: sh, bash, zsh, fish, tcsh, ...?
1121            debug_assert!(
1122                !arg.as_ref().contains('\n'),
1123                "multiline arguments do not work in all shells"
1124            );
1125            to_run.push(' ');
1126            to_run.push_str(&shell_kind.try_quote(arg.as_ref()).expect("shell quoting"));
1127        }
1128        let to_run = if shell_kind == ShellKind::Cmd {
1129            to_run // 'cd' prints the current directory in CMD
1130        } else {
1131            let separator = shell_kind.sequential_commands_separator();
1132            format!("cd{separator} {to_run}")
1133        };
1134        self.ssh_options(&mut command, true)
1135            .arg(self.connection_options.ssh_destination());
1136        if !allow_pseudo_tty {
1137            command.arg("-T");
1138        }
1139        command.arg(to_run);
1140        log::debug!("ssh {:?}", command);
1141        command
1142    }
1143
1144    async fn run_command(
1145        &self,
1146        shell_kind: ShellKind,
1147        program: &str,
1148        args: &[impl AsRef<str>],
1149        allow_pseudo_tty: bool,
1150    ) -> Result<String> {
1151        let mut command = self.ssh_command(shell_kind, program, args, allow_pseudo_tty);
1152        let output = command.output().await?;
1153        log::debug!("{:?}: {:?}", command, output);
1154        anyhow::ensure!(
1155            output.status.success(),
1156            "failed to run command {command:?}: {}",
1157            String::from_utf8_lossy(&output.stderr)
1158        );
1159        Ok(String::from_utf8_lossy(&output.stdout).to_string())
1160    }
1161
1162    #[cfg(not(target_os = "windows"))]
1163    fn ssh_options<'a>(
1164        &self,
1165        command: &'a mut process::Command,
1166        include_port_forwards: bool,
1167    ) -> &'a mut process::Command {
1168        let args = if include_port_forwards {
1169            self.connection_options.additional_args()
1170        } else {
1171            self.connection_options.additional_args_for_scp()
1172        };
1173
1174        command
1175            .stdin(Stdio::piped())
1176            .stdout(Stdio::piped())
1177            .stderr(Stdio::piped())
1178            .args(args)
1179            .args(["-o", "ControlMaster=no", "-o"])
1180            .arg(format!("ControlPath={}", self.socket_path.display()))
1181    }
1182
1183    #[cfg(target_os = "windows")]
1184    fn ssh_options<'a>(
1185        &self,
1186        command: &'a mut process::Command,
1187        include_port_forwards: bool,
1188    ) -> &'a mut process::Command {
1189        let args = if include_port_forwards {
1190            self.connection_options.additional_args()
1191        } else {
1192            self.connection_options.additional_args_for_scp()
1193        };
1194
1195        command
1196            .stdin(Stdio::piped())
1197            .stdout(Stdio::piped())
1198            .stderr(Stdio::piped())
1199            .args(args)
1200            .envs(self.envs.clone())
1201    }
1202
1203    // On Windows, we need to use `SSH_ASKPASS` to provide the password to ssh.
1204    // On Linux, we use the `ControlPath` option to create a socket file that ssh can use to
1205    #[cfg(not(target_os = "windows"))]
1206    fn ssh_args(&self) -> Vec<String> {
1207        let mut arguments = self.connection_options.additional_args();
1208        arguments.extend(vec![
1209            "-o".to_string(),
1210            "ControlMaster=no".to_string(),
1211            "-o".to_string(),
1212            format!("ControlPath={}", self.socket_path.display()),
1213            self.connection_options.ssh_destination(),
1214        ]);
1215        arguments
1216    }
1217
1218    #[cfg(target_os = "windows")]
1219    fn ssh_args(&self) -> Vec<String> {
1220        let mut arguments = self.connection_options.additional_args();
1221        arguments.push(self.connection_options.ssh_destination());
1222        arguments
1223    }
1224
1225    async fn platform(&self, shell: ShellKind, is_windows: bool) -> Result<RemotePlatform> {
1226        if is_windows {
1227            self.platform_windows(shell).await
1228        } else {
1229            self.platform_posix(shell).await
1230        }
1231    }
1232
1233    async fn platform_posix(&self, shell: ShellKind) -> Result<RemotePlatform> {
1234        let output = self
1235            .run_command(shell, "uname", &["-sm"], false)
1236            .await
1237            .context("Failed to run 'uname -sm' to determine platform")?;
1238        parse_platform(&output)
1239    }
1240
1241    async fn platform_windows(&self, shell: ShellKind) -> Result<RemotePlatform> {
1242        let output = self
1243            .run_command(
1244                shell,
1245                "cmd.exe",
1246                &["/c", "echo", "%PROCESSOR_ARCHITECTURE%"],
1247                false,
1248            )
1249            .await
1250            .context(
1251                "Failed to run 'echo %PROCESSOR_ARCHITECTURE%' to determine Windows architecture",
1252            )?;
1253
1254        Ok(RemotePlatform {
1255            os: RemoteOs::Windows,
1256            arch: match output.trim() {
1257                "AMD64" => RemoteArch::X86_64,
1258                "ARM64" => RemoteArch::Aarch64,
1259                arch => anyhow::bail!(
1260                    "Prebuilt remote servers are not yet available for windows-{arch}. See https://zed.dev/docs/remote-development"
1261                ),
1262            },
1263        })
1264    }
1265
1266    /// Probes whether the remote host is running Windows.
1267    ///
1268    /// This is done by attempting to run a simple Windows-specific command.
1269    /// If it succeeds and returns Windows-like output, we assume it's Windows.
1270    async fn probe_is_windows(&self) -> bool {
1271        match self
1272            .run_command(ShellKind::Cmd, "cmd.exe", &["/c", "ver"], false)
1273            .await
1274        {
1275            // Windows 'ver' command outputs something like "Microsoft Windows [Version 10.0.19045.5011]"
1276            Ok(output) => output.trim().contains("indows"),
1277            Err(_) => false,
1278        }
1279    }
1280
1281    async fn shell(&self, is_windows: bool) -> String {
1282        if is_windows {
1283            self.shell_windows().await
1284        } else {
1285            self.shell_posix().await
1286        }
1287    }
1288
1289    async fn shell_posix(&self) -> String {
1290        const DEFAULT_SHELL: &str = "sh";
1291        match self
1292            .run_command(ShellKind::Posix, "sh", &["-c", "echo $SHELL"], false)
1293            .await
1294        {
1295            Ok(output) => parse_shell(&output, DEFAULT_SHELL),
1296            Err(e) => {
1297                log::error!("Failed to detect remote shell: {e}");
1298                DEFAULT_SHELL.to_owned()
1299            }
1300        }
1301    }
1302
1303    async fn shell_windows(&self) -> String {
1304        const DEFAULT_SHELL: &str = "cmd.exe";
1305
1306        // We detect the shell used by the SSH session by running the following command in PowerShell:
1307        // (Get-CimInstance Win32_Process -Filter "ProcessId = $((Get-CimInstance Win32_Process -Filter ProcessId=$PID).ParentProcessId)").Name
1308        // This prints the name of PowerShell's parent process (which will be the shell that SSH launched).
1309        // We pass it as a Base64 encoded string since we don't yet know how to correctly quote that command.
1310        // (We'd need to know what the shell is to do that...)
1311        match self
1312            .run_command(
1313                ShellKind::Cmd,
1314                "powershell",
1315                &[
1316                    "-E",
1317                    "KABHAGUAdAAtAEMAaQBtAEkAbgBzAHQAYQBuAGMAZQAgAFcAaQBuADMAMgBfAFAAcgBvAGMAZQBzAHMAIAAtAEYAaQBsAHQAZQByACAAIgBQAHIAbwBjAGUAcwBzAEkAZAAgAD0AIAAkACgAKABHAGUAdAAtAEMAaQBtAEkAbgBzAHQAYQBuAGMAZQAgAFcAaQBuADMAMgBfAFAAcgBvAGMAZQBzAHMAIAAtAEYAaQBsAHQAZQByACAAUAByAG8AYwBlAHMAcwBJAGQAPQAkAFAASQBEACkALgBQAGEAcgBlAG4AdABQAHIAbwBjAGUAcwBzAEkAZAApACIAKQAuAE4AYQBtAGUA",
1318                ],
1319                false,
1320            )
1321            .await
1322        {
1323            Ok(output) => parse_shell(&output, DEFAULT_SHELL),
1324            Err(e) => {
1325                log::error!("Failed to detect remote shell: {e}");
1326                DEFAULT_SHELL.to_owned()
1327            }
1328        }
1329    }
1330}
1331
1332fn parse_port_number(port_str: &str) -> Result<u16> {
1333    port_str
1334        .parse()
1335        .with_context(|| format!("parsing port number: {port_str}"))
1336}
1337
1338fn parse_port_forward_spec(spec: &str) -> Result<SshPortForwardOption> {
1339    let parts: Vec<&str> = spec.split(':').collect();
1340
1341    match parts.len() {
1342        4 => {
1343            let local_port = parse_port_number(parts[1])?;
1344            let remote_port = parse_port_number(parts[3])?;
1345
1346            Ok(SshPortForwardOption {
1347                local_host: Some(parts[0].to_string()),
1348                local_port,
1349                remote_host: Some(parts[2].to_string()),
1350                remote_port,
1351            })
1352        }
1353        3 => {
1354            let local_port = parse_port_number(parts[0])?;
1355            let remote_port = parse_port_number(parts[2])?;
1356
1357            Ok(SshPortForwardOption {
1358                local_host: None,
1359                local_port,
1360                remote_host: Some(parts[1].to_string()),
1361                remote_port,
1362            })
1363        }
1364        _ => anyhow::bail!("Invalid port forward format"),
1365    }
1366}
1367
1368impl SshConnectionOptions {
1369    pub fn parse_command_line(input: &str) -> Result<Self> {
1370        let input = input.trim_start_matches("ssh ");
1371        let mut hostname: Option<String> = None;
1372        let mut username: Option<String> = None;
1373        let mut port: Option<u16> = None;
1374        let mut args = Vec::new();
1375        let mut port_forwards: Vec<SshPortForwardOption> = Vec::new();
1376
1377        // disallowed: -E, -e, -F, -f, -G, -g, -M, -N, -n, -O, -q, -S, -s, -T, -t, -V, -v, -W
1378        const ALLOWED_OPTS: &[&str] = &[
1379            "-4", "-6", "-A", "-a", "-C", "-K", "-k", "-X", "-x", "-Y", "-y",
1380        ];
1381        const ALLOWED_ARGS: &[&str] = &[
1382            "-B", "-b", "-c", "-D", "-F", "-I", "-i", "-J", "-l", "-m", "-o", "-P", "-p", "-R",
1383            "-w",
1384        ];
1385
1386        let mut tokens = ShellKind::Posix
1387            .split(input)
1388            .context("invalid input")?
1389            .into_iter();
1390
1391        'outer: while let Some(arg) = tokens.next() {
1392            if ALLOWED_OPTS.contains(&(&arg as &str)) {
1393                args.push(arg.to_string());
1394                continue;
1395            }
1396            if arg == "-p" {
1397                port = tokens.next().and_then(|arg| arg.parse().ok());
1398                continue;
1399            } else if let Some(p) = arg.strip_prefix("-p") {
1400                port = p.parse().ok();
1401                continue;
1402            }
1403            if arg == "-l" {
1404                username = tokens.next();
1405                continue;
1406            } else if let Some(l) = arg.strip_prefix("-l") {
1407                username = Some(l.to_string());
1408                continue;
1409            }
1410            if arg == "-L" || arg.starts_with("-L") {
1411                let forward_spec = if arg == "-L" {
1412                    tokens.next()
1413                } else {
1414                    Some(arg.strip_prefix("-L").unwrap().to_string())
1415                };
1416
1417                if let Some(spec) = forward_spec {
1418                    port_forwards.push(parse_port_forward_spec(&spec)?);
1419                } else {
1420                    anyhow::bail!("Missing port forward format");
1421                }
1422            }
1423
1424            for a in ALLOWED_ARGS {
1425                if arg == *a {
1426                    args.push(arg);
1427                    if let Some(next) = tokens.next() {
1428                        args.push(next);
1429                    }
1430                    continue 'outer;
1431                } else if arg.starts_with(a) {
1432                    args.push(arg);
1433                    continue 'outer;
1434                }
1435            }
1436            if arg.starts_with("-") || hostname.is_some() {
1437                anyhow::bail!("unsupported argument: {:?}", arg);
1438            }
1439            let mut input = &arg as &str;
1440            // Destination might be: username1@username2@ip2@ip1
1441            if let Some((u, rest)) = input.rsplit_once('@') {
1442                input = rest;
1443                username = Some(u.to_string());
1444            }
1445
1446            // Handle port parsing, accounting for IPv6 addresses
1447            // IPv6 addresses can be: 2001:db8::1 or [2001:db8::1]:22
1448            if input.starts_with('[') {
1449                if let Some((rest, p)) = input.rsplit_once("]:") {
1450                    input = rest.strip_prefix('[').unwrap_or(rest);
1451                    port = p.parse().ok();
1452                } else if input.ends_with(']') {
1453                    input = input.strip_prefix('[').unwrap_or(input);
1454                    input = input.strip_suffix(']').unwrap_or(input);
1455                }
1456            } else if let Some((rest, p)) = input.rsplit_once(':')
1457                && !rest.contains(":")
1458            {
1459                input = rest;
1460                port = p.parse().ok();
1461            }
1462
1463            hostname = Some(input.to_string())
1464        }
1465
1466        let Some(hostname) = hostname else {
1467            anyhow::bail!("missing hostname");
1468        };
1469
1470        let port_forwards = match port_forwards.len() {
1471            0 => None,
1472            _ => Some(port_forwards),
1473        };
1474
1475        Ok(Self {
1476            host: hostname.into(),
1477            username,
1478            port,
1479            port_forwards,
1480            args: Some(args),
1481            password: None,
1482            nickname: None,
1483            upload_binary_over_ssh: false,
1484            connection_timeout: None,
1485        })
1486    }
1487
1488    pub fn ssh_destination(&self) -> String {
1489        let mut result = String::default();
1490        if let Some(username) = &self.username {
1491            // Username might be: username1@username2@ip2
1492            let username = urlencoding::encode(username);
1493            result.push_str(&username);
1494            result.push('@');
1495        }
1496
1497        result.push_str(&self.host.to_string());
1498        result
1499    }
1500
1501    pub fn additional_args_for_scp(&self) -> Vec<String> {
1502        self.args.iter().flatten().cloned().collect::<Vec<String>>()
1503    }
1504
1505    pub fn additional_args(&self) -> Vec<String> {
1506        let mut args = self.additional_args_for_scp();
1507
1508        if let Some(timeout) = self.connection_timeout {
1509            args.extend(["-o".to_string(), format!("ConnectTimeout={}", timeout)]);
1510        }
1511
1512        if let Some(port) = self.port {
1513            args.push("-p".to_string());
1514            args.push(port.to_string());
1515        }
1516
1517        if let Some(forwards) = &self.port_forwards {
1518            args.extend(forwards.iter().map(|pf| {
1519                let local_host = match &pf.local_host {
1520                    Some(host) => host,
1521                    None => "localhost",
1522                };
1523                let remote_host = match &pf.remote_host {
1524                    Some(host) => host,
1525                    None => "localhost",
1526                };
1527
1528                format!(
1529                    "-L{}:{}:{}:{}",
1530                    local_host, pf.local_port, remote_host, pf.remote_port
1531                )
1532            }));
1533        }
1534
1535        args
1536    }
1537
1538    fn scp_destination(&self) -> String {
1539        if let Some(username) = &self.username {
1540            format!("{}@{}", username, self.host.to_bracketed_string())
1541        } else {
1542            self.host.to_string()
1543        }
1544    }
1545
1546    pub fn connection_string(&self) -> String {
1547        let host = if let Some(port) = &self.port {
1548            format!("{}:{}", self.host.to_bracketed_string(), port)
1549        } else {
1550            self.host.to_string()
1551        };
1552
1553        if let Some(username) = &self.username {
1554            format!("{}@{}", username, host)
1555        } else {
1556            host
1557        }
1558    }
1559}
1560
1561fn build_command(
1562    input_program: Option<String>,
1563    input_args: &[String],
1564    input_env: &HashMap<String, String>,
1565    working_dir: Option<String>,
1566    port_forward: Option<(u16, String, u16)>,
1567    ssh_env: HashMap<String, String>,
1568    ssh_path_style: PathStyle,
1569    ssh_shell: &str,
1570    ssh_shell_kind: ShellKind,
1571    ssh_args: Vec<String>,
1572    interactive: Interactive,
1573) -> Result<CommandTemplate> {
1574    use std::fmt::Write as _;
1575
1576    let mut exec = String::new();
1577    if let Some(working_dir) = working_dir {
1578        let working_dir = RemotePathBuf::new(working_dir, ssh_path_style).to_string();
1579
1580        // shlex will wrap the command in single quotes (''), disabling ~ expansion,
1581        // replace with something that works
1582        const TILDE_PREFIX: &'static str = "~/";
1583        if working_dir.starts_with(TILDE_PREFIX) {
1584            let working_dir = working_dir.trim_start_matches("~").trim_start_matches("/");
1585            write!(
1586                exec,
1587                "cd \"$HOME/{working_dir}\" {} ",
1588                ssh_shell_kind.sequential_and_commands_separator()
1589            )?;
1590        } else {
1591            write!(
1592                exec,
1593                "cd \"{working_dir}\" {} ",
1594                ssh_shell_kind.sequential_and_commands_separator()
1595            )?;
1596        }
1597    } else {
1598        write!(
1599            exec,
1600            "cd {} ",
1601            ssh_shell_kind.sequential_and_commands_separator()
1602        )?;
1603    };
1604    write!(exec, "exec env ")?;
1605
1606    for (k, v) in input_env.iter() {
1607        write!(
1608            exec,
1609            "{}={} ",
1610            k,
1611            ssh_shell_kind.try_quote(v).context("shell quoting")?
1612        )?;
1613    }
1614
1615    if let Some(input_program) = input_program {
1616        write!(
1617            exec,
1618            "{}",
1619            ssh_shell_kind
1620                .try_quote_prefix_aware(&input_program)
1621                .context("shell quoting")?
1622        )?;
1623        for arg in input_args {
1624            let arg = ssh_shell_kind.try_quote(&arg).context("shell quoting")?;
1625            write!(exec, " {}", &arg)?;
1626        }
1627    } else {
1628        write!(exec, "{ssh_shell} -l")?;
1629    };
1630
1631    let mut args = Vec::new();
1632    args.extend(ssh_args);
1633
1634    if let Some((local_port, host, remote_port)) = port_forward {
1635        args.push("-L".into());
1636        args.push(format!("{local_port}:{host}:{remote_port}"));
1637    }
1638
1639    // -q suppresses the "Connection to ... closed." message that SSH prints when
1640    // the connection terminates with -t (pseudo-terminal allocation)
1641    args.push("-q".into());
1642    match interactive {
1643        // -t forces pseudo-TTY allocation (for interactive use)
1644        Interactive::Yes => args.push("-t".into()),
1645        // -T disables pseudo-TTY allocation (for non-interactive piped stdio)
1646        Interactive::No => args.push("-T".into()),
1647    }
1648    args.push(exec);
1649
1650    Ok(CommandTemplate {
1651        program: "ssh".into(),
1652        args,
1653        env: ssh_env,
1654    })
1655}
1656
1657#[cfg(test)]
1658mod tests {
1659    use super::*;
1660
1661    #[test]
1662    fn test_build_command() -> Result<()> {
1663        let mut input_env = HashMap::default();
1664        input_env.insert("INPUT_VA".to_string(), "val".to_string());
1665        let mut env = HashMap::default();
1666        env.insert("SSH_VAR".to_string(), "ssh-val".to_string());
1667
1668        // Test non-interactive command (interactive=false should use -T)
1669        let command = build_command(
1670            Some("remote_program".to_string()),
1671            &["arg1".to_string(), "arg2".to_string()],
1672            &input_env,
1673            Some("~/work".to_string()),
1674            None,
1675            env.clone(),
1676            PathStyle::Posix,
1677            "/bin/bash",
1678            ShellKind::Posix,
1679            vec!["-o".to_string(), "ControlMaster=auto".to_string()],
1680            Interactive::No,
1681        )?;
1682        assert_eq!(command.program, "ssh");
1683        // Should contain -T for non-interactive
1684        assert!(command.args.iter().any(|arg| arg == "-T"));
1685        assert!(!command.args.iter().any(|arg| arg == "-t"));
1686
1687        // Test interactive command (interactive=true should use -t)
1688        let command = build_command(
1689            Some("remote_program".to_string()),
1690            &["arg1".to_string(), "arg2".to_string()],
1691            &input_env,
1692            Some("~/work".to_string()),
1693            None,
1694            env.clone(),
1695            PathStyle::Posix,
1696            "/bin/fish",
1697            ShellKind::Fish,
1698            vec!["-p".to_string(), "2222".to_string()],
1699            Interactive::Yes,
1700        )?;
1701
1702        assert_eq!(command.program, "ssh");
1703        assert_eq!(
1704            command.args.iter().map(String::as_str).collect::<Vec<_>>(),
1705            [
1706                "-p",
1707                "2222",
1708                "-q",
1709                "-t",
1710                "cd \"$HOME/work\" && exec env INPUT_VA=val remote_program arg1 arg2"
1711            ]
1712        );
1713        assert_eq!(command.env, env);
1714
1715        let mut input_env = HashMap::default();
1716        input_env.insert("INPUT_VA".to_string(), "val".to_string());
1717        let mut env = HashMap::default();
1718        env.insert("SSH_VAR".to_string(), "ssh-val".to_string());
1719
1720        let command = build_command(
1721            None,
1722            &[],
1723            &input_env,
1724            None,
1725            Some((1, "foo".to_owned(), 2)),
1726            env.clone(),
1727            PathStyle::Posix,
1728            "/bin/fish",
1729            ShellKind::Fish,
1730            vec!["-p".to_string(), "2222".to_string()],
1731            Interactive::Yes,
1732        )?;
1733
1734        assert_eq!(command.program, "ssh");
1735        assert_eq!(
1736            command.args.iter().map(String::as_str).collect::<Vec<_>>(),
1737            [
1738                "-p",
1739                "2222",
1740                "-L",
1741                "1:foo:2",
1742                "-q",
1743                "-t",
1744                "cd && exec env INPUT_VA=val /bin/fish -l"
1745            ]
1746        );
1747        assert_eq!(command.env, env);
1748
1749        Ok(())
1750    }
1751
1752    #[test]
1753    fn scp_args_exclude_port_forward_flags() {
1754        let options = SshConnectionOptions {
1755            host: "example.com".into(),
1756            args: Some(vec![
1757                "-p".to_string(),
1758                "2222".to_string(),
1759                "-o".to_string(),
1760                "StrictHostKeyChecking=no".to_string(),
1761            ]),
1762            port_forwards: Some(vec![SshPortForwardOption {
1763                local_host: Some("127.0.0.1".to_string()),
1764                local_port: 8080,
1765                remote_host: Some("127.0.0.1".to_string()),
1766                remote_port: 80,
1767            }]),
1768            ..Default::default()
1769        };
1770
1771        let ssh_args = options.additional_args();
1772        assert!(
1773            ssh_args.iter().any(|arg| arg.starts_with("-L")),
1774            "expected ssh args to include port-forward: {ssh_args:?}"
1775        );
1776
1777        let scp_args = options.additional_args_for_scp();
1778        assert_eq!(
1779            scp_args,
1780            vec![
1781                "-p".to_string(),
1782                "2222".to_string(),
1783                "-o".to_string(),
1784                "StrictHostKeyChecking=no".to_string(),
1785            ]
1786        );
1787    }
1788
1789    #[test]
1790    fn test_host_parsing() -> Result<()> {
1791        let opts = SshConnectionOptions::parse_command_line("user@2001:db8::1")?;
1792        assert_eq!(opts.host, "2001:db8::1".into());
1793        assert_eq!(opts.username, Some("user".to_string()));
1794        assert_eq!(opts.port, None);
1795
1796        let opts = SshConnectionOptions::parse_command_line("user@[2001:db8::1]:2222")?;
1797        assert_eq!(opts.host, "2001:db8::1".into());
1798        assert_eq!(opts.username, Some("user".to_string()));
1799        assert_eq!(opts.port, Some(2222));
1800
1801        let opts = SshConnectionOptions::parse_command_line("user@[2001:db8::1]")?;
1802        assert_eq!(opts.host, "2001:db8::1".into());
1803        assert_eq!(opts.username, Some("user".to_string()));
1804        assert_eq!(opts.port, None);
1805
1806        let opts = SshConnectionOptions::parse_command_line("2001:db8::1")?;
1807        assert_eq!(opts.host, "2001:db8::1".into());
1808        assert_eq!(opts.username, None);
1809        assert_eq!(opts.port, None);
1810
1811        let opts = SshConnectionOptions::parse_command_line("[2001:db8::1]:2222")?;
1812        assert_eq!(opts.host, "2001:db8::1".into());
1813        assert_eq!(opts.username, None);
1814        assert_eq!(opts.port, Some(2222));
1815
1816        let opts = SshConnectionOptions::parse_command_line("user@example.com:2222")?;
1817        assert_eq!(opts.host, "example.com".into());
1818        assert_eq!(opts.username, Some("user".to_string()));
1819        assert_eq!(opts.port, Some(2222));
1820
1821        let opts = SshConnectionOptions::parse_command_line("user@192.168.1.1:2222")?;
1822        assert_eq!(opts.host, "192.168.1.1".into());
1823        assert_eq!(opts.username, Some("user".to_string()));
1824        assert_eq!(opts.port, Some(2222));
1825
1826        Ok(())
1827    }
1828}