2111 lines · 69.0 KB · Rust
RawDownload
use crate::{    RemoteArch, RemoteClientDelegate, RemoteOs, RemotePlatform,    remote_client::{CommandTemplate, Interactive, RemoteConnection, RemoteConnectionOptions},    transport::{parse_platform, parse_shell},};use anyhow::{Context as _, Result, anyhow};use async_trait::async_trait;use collections::HashMap;use futures::{    AsyncReadExt as _, FutureExt as _,    channel::mpsc::{Sender, UnboundedReceiver, UnboundedSender},    select_biased,};use gpui::{App, AppContext as _, AsyncApp, Task};use parking_lot::Mutex;use paths::remote_server_dir_relative;use release_channel::{AppVersion, ReleaseChannel};use rpc::proto::Envelope;use semver::Version;pub use settings::SshPortForwardOption;use smol::fs;use std::{    net::IpAddr,    path::{Path, PathBuf},    sync::Arc,    time::Instant,};use tempfile::TempDir;use util::command::{Child, Stdio};use util::{    paths::{PathStyle, RemotePathBuf},    rel_path::RelPath,    shell::ShellKind,}; pub(crate) struct SshRemoteConnection {    socket: SshSocket,    master_process: Mutex<Option<MasterProcess>>,    remote_binary_path: Option<Arc<RelPath>>,    ssh_platform: RemotePlatform,    ssh_path_style: PathStyle,    ssh_shell: String,    ssh_shell_kind: ShellKind,    ssh_default_system_shell: String,    _temp_dir: TempDir,} #[derive(Debug, Clone, PartialEq, Eq, Hash)]pub enum SshConnectionHost {    IpAddr(IpAddr),    Hostname(String),} impl SshConnectionHost {    pub fn to_bracketed_string(&self) -> String {        match self {            Self::IpAddr(IpAddr::V4(ip)) => ip.to_string(),            Self::IpAddr(IpAddr::V6(ip)) => format!("[{}]", ip),            Self::Hostname(hostname) => hostname.clone(),        }    }     pub fn to_string(&self) -> String {        match self {            Self::IpAddr(ip) => ip.to_string(),            Self::Hostname(hostname) => hostname.clone(),        }    }} impl From<&str> for SshConnectionHost {    fn from(value: &str) -> Self {        if let Ok(address) = value.parse() {            Self::IpAddr(address)        } else {            Self::Hostname(value.to_string())        }    }} impl From<String> for SshConnectionHost {    fn from(value: String) -> Self {        if let Ok(address) = value.parse() {            Self::IpAddr(address)        } else {            Self::Hostname(value)        }    }} impl Default for SshConnectionHost {    fn default() -> Self {        Self::Hostname(Default::default())    }} fn bracket_ipv6(host: &str) -> String {    if host.contains(':') && !host.starts_with('[') {        format!("[{}]", host)    } else {        host.to_string()    }} #[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]pub struct SshConnectionOptions {    pub host: SshConnectionHost,    pub username: Option<String>,    pub port: Option<u16>,    pub password: Option<String>,    pub args: Option<Vec<String>>,    pub port_forwards: Option<Vec<SshPortForwardOption>>,    pub connection_timeout: Option<u16>,     pub nickname: Option<String>,    pub upload_binary_over_ssh: bool,} impl From<settings::SshConnection> for SshConnectionOptions {    fn from(val: settings::SshConnection) -> Self {        SshConnectionOptions {            host: val.host.to_string().into(),            username: val.username,            port: val.port,            password: None,            args: Some(val.args),            nickname: val.nickname,            upload_binary_over_ssh: val.upload_binary_over_ssh.unwrap_or_default(),            port_forwards: val.port_forwards,            connection_timeout: val.connection_timeout,        }    }} struct SshSocket {    connection_options: SshConnectionOptions,    #[cfg(not(windows))]    socket_path: std::path::PathBuf,    /// Extra environment variables needed for the ssh process    envs: HashMap<String, String>,    #[cfg(windows)]    _proxy: askpass::PasswordProxy,} struct MasterProcess {    process: Child,} #[cfg(not(windows))]impl MasterProcess {    pub fn new(        askpass_script_path: &std::ffi::OsStr,        additional_args: Vec<String>,        socket_path: &std::path::Path,        destination: &str,    ) -> Result<Self> {        let args = [            "-N",            "-o",            "ControlPersist=no",            "-o",            "ControlMaster=yes",            "-o",        ];         let mut master_process = util::command::new_command("ssh");        master_process            .kill_on_drop(true)            .stdin(Stdio::null())            .stdout(Stdio::piped())            .stderr(Stdio::piped())            .env("SSH_ASKPASS_REQUIRE", "force")            .env("SSH_ASKPASS", askpass_script_path)            .args(additional_args)            .args(args);         master_process.arg(format!("ControlPath={}", socket_path.display()));         let process = master_process.arg(&destination).spawn()?;         Ok(MasterProcess { process })    }     pub async fn wait_connected(&mut self) -> Result<()> {        let Some(mut stdout) = self.process.stdout.take() else {            anyhow::bail!("ssh process stdout capture failed");        };         let mut output = Vec::new();        stdout.read_to_end(&mut output).await?;        Ok(())    }} #[cfg(windows)]impl MasterProcess {    const CONNECTION_ESTABLISHED_MAGIC: &str = "ZED_SSH_CONNECTION_ESTABLISHED";     pub fn new(        askpass_script_path: &std::ffi::OsStr,        additional_args: Vec<String>,        destination: &str,    ) -> Result<Self> {        // On Windows, `ControlMaster` and `ControlPath` are not supported:        // https://github.com/PowerShell/Win32-OpenSSH/issues/405        // https://github.com/PowerShell/Win32-OpenSSH/wiki/Project-Scope        //        // Using an ugly workaround to detect connection establishment        // -N doesn't work with JumpHosts as windows openssh never closes stdin in that case        let args = [            "-t",            &format!("echo '{}'; exec $0", Self::CONNECTION_ESTABLISHED_MAGIC),        ];         let mut master_process = util::command::new_command("ssh");        master_process            .kill_on_drop(true)            .stdin(Stdio::null())            .stdout(Stdio::piped())            .stderr(Stdio::piped())            .env("SSH_ASKPASS_REQUIRE", "force")            .env("SSH_ASKPASS", askpass_script_path)            .args(additional_args)            .arg(destination)            .args(args);         let process = master_process.spawn()?;         Ok(MasterProcess { process })    }     pub async fn wait_connected(&mut self) -> Result<()> {        use smol::io::AsyncBufReadExt;         let Some(stdout) = self.process.stdout.take() else {            anyhow::bail!("ssh process stdout capture failed");        };         let mut reader = smol::io::BufReader::new(stdout);         let mut line = String::new();         loop {            let n = reader.read_line(&mut line).await?;            if n == 0 {                anyhow::bail!("ssh process exited before connection established");            }             if line.contains(Self::CONNECTION_ESTABLISHED_MAGIC) {                return Ok(());            }        }    }} impl AsRef<Child> for MasterProcess {    fn as_ref(&self) -> &Child {        &self.process    }} impl AsMut<Child> for MasterProcess {    fn as_mut(&mut self) -> &mut Child {        &mut self.process    }} #[async_trait(?Send)]impl RemoteConnection for SshRemoteConnection {    async fn kill(&self) -> Result<()> {        let Some(mut process) = self.master_process.lock().take() else {            return Ok(());        };        process.as_mut().kill().ok();        process.as_mut().status().await?;        Ok(())    }     fn has_been_killed(&self) -> bool {        self.master_process.lock().is_none()    }     fn connection_options(&self) -> RemoteConnectionOptions {        RemoteConnectionOptions::Ssh(self.socket.connection_options.clone())    }     fn shell(&self) -> String {        self.ssh_shell.clone()    }     fn default_system_shell(&self) -> String {        self.ssh_default_system_shell.clone()    }     fn build_command(        &self,        input_program: Option<String>,        input_args: &[String],        input_env: &HashMap<String, String>,        working_dir: Option<String>,        port_forward: Option<(u16, String, u16)>,        interactive: Interactive,    ) -> Result<CommandTemplate> {        let Self {            ssh_path_style,            socket,            ssh_shell_kind,            ssh_shell,            ..        } = self;        let env = socket.envs.clone();         if self.ssh_platform.os.is_windows() {            build_command_windows(                input_program,                input_args,                input_env,                working_dir,                port_forward,                env,                *ssh_path_style,                ssh_shell,                *ssh_shell_kind,                socket.ssh_command_options(),                &socket.connection_options.ssh_destination(),                interactive,            )        } else {            build_command_posix(                input_program,                input_args,                input_env,                working_dir,                port_forward,                env,                *ssh_path_style,                ssh_shell,                *ssh_shell_kind,                socket.ssh_command_options(),                &socket.connection_options.ssh_destination(),                interactive,            )        }    }     fn build_forward_ports_command(        &self,        forwards: Vec<(u16, String, u16)>,    ) -> Result<CommandTemplate> {        let Self { socket, .. } = self;        let mut args = socket.ssh_command_options();        args.push("-N".into());        for (local_port, host, remote_port) in forwards {            args.push("-L".into());            args.push(format!(                "{}:{}:{}",                local_port,                bracket_ipv6(&host),                remote_port            ));        }        args.push(socket.connection_options.ssh_destination());        Ok(CommandTemplate {            program: "ssh".into(),            args,            env: Default::default(),        })    }     fn upload_directory(        &self,        src_path: PathBuf,        dest_path: RemotePathBuf,        cx: &App,    ) -> Task<Result<()>> {        let dest_path_str = dest_path.to_string();        let src_path_display = src_path.display().to_string();         let mut sftp_command = self.build_sftp_command();        let mut scp_command =            self.build_scp_command(&src_path, &dest_path_str, Some(&["-C", "-r"]));         cx.background_spawn(async move {            // We will try SFTP first, and if that fails, we will fall back to SCP.            // If SCP fails also, we give up and return an error.            // The reason we allow a fallback from SFTP to SCP is that if the user has to specify a password,            // depending on the implementation of SSH stack, SFTP may disable interactive password prompts in batch mode.            // This is for example the case on Windows as evidenced by this implementation snippet:            // https://github.com/PowerShell/openssh-portable/blob/b8c08ef9da9450a94a9c5ef717d96a7bd83f3332/sshconnect2.c#L417            if Self::is_sftp_available().await {                log::debug!("using SFTP for directory upload");                let mut child = sftp_command.spawn()?;                if let Some(mut stdin) = child.stdin.take() {                    use futures::AsyncWriteExt;                    let sftp_batch = format!("put -r \"{src_path_display}\" \"{dest_path_str}\"\n");                    stdin.write_all(sftp_batch.as_bytes()).await?;                    stdin.flush().await?;                }                 let output = child.output().await?;                if output.status.success() {                    return Ok(());                }                 let stderr = String::from_utf8_lossy(&output.stderr);                log::debug!("failed to upload directory via SFTP {src_path_display} -> {dest_path_str}: {stderr}");            }             log::debug!("using SCP for directory upload");            let output = scp_command.output().await?;             if output.status.success() {                return Ok(());            }             let stderr = String::from_utf8_lossy(&output.stderr);            log::debug!("failed to upload directory via SCP {src_path_display} -> {dest_path_str}: {stderr}");             anyhow::bail!(                "failed to upload directory via SFTP/SCP {} -> {}: {}",                src_path_display,                dest_path_str,                stderr,            );        })    }     fn start_proxy(        &self,        unique_identifier: String,        reconnect: bool,        incoming_tx: UnboundedSender<Envelope>,        outgoing_rx: UnboundedReceiver<Envelope>,        connection_activity_tx: Sender<()>,        delegate: Arc<dyn RemoteClientDelegate>,        cx: &mut AsyncApp,    ) -> Task<Result<i32>> {        const VARS: [&str; 3] = ["RUST_LOG", "RUST_BACKTRACE", "ZED_GENERATE_MINIDUMPS"];        delegate.set_status(Some("Starting proxy"), cx);         let Some(remote_binary_path) = self.remote_binary_path.clone() else {            return Task::ready(Err(anyhow!("Remote binary path not set")));        };         let mut ssh_command = if self.ssh_platform.os.is_windows() {            // TODO: Set the `VARS` environment variables, we do not have `env` on windows            // so this needs a different approach            let mut proxy_args = vec![];            proxy_args.push("proxy".to_owned());            proxy_args.push("--identifier".to_owned());            proxy_args.push(unique_identifier);             if reconnect {                proxy_args.push("--reconnect".to_owned());            }            self.socket.ssh_command(                self.ssh_shell_kind,                &remote_binary_path.display(self.path_style()),                &proxy_args,                false,            )        } else {            let mut proxy_args = vec![];            for env_var in VARS {                if let Some(value) = std::env::var(env_var).ok() {                    proxy_args.push(format!("{env_var}={value}"));                }            }            proxy_args.push(remote_binary_path.display(self.path_style()).into_owned());            proxy_args.push("proxy".to_owned());            proxy_args.push("--identifier".to_owned());            proxy_args.push(unique_identifier);             if reconnect {                proxy_args.push("--reconnect".to_owned());            }            self.socket                .ssh_command(self.ssh_shell_kind, "env", &proxy_args, false)        };         let ssh_proxy_process = match ssh_command            // IMPORTANT: we kill this process when we drop the task that uses it.            .kill_on_drop(true)            .spawn()        {            Ok(process) => process,            Err(error) => {                return Task::ready(Err(anyhow!("failed to spawn remote server: {}", error)));            }        };         super::handle_rpc_messages_over_child_process_stdio(            ssh_proxy_process,            incoming_tx,            outgoing_rx,            connection_activity_tx,            cx,        )    } 

Lines 1–500 of 2111.