Files
zed/crates/remote/src/remote_client.rs
T
Lukas Wirth f5ba029313 remote: Implement client side connection support for windows remotes (#45084)
Obviously this doesn't do too much without having an actual windows
server binary for the remote side, but it does at least improve the
error message as right now we will complain about `uname` not being a
valid powershell command.

Release Notes:

- N/A *or* Added/Fixed/Improved ...
2025-12-17 11:31:18 +01:00

1766 lines
59 KiB
Rust

use crate::{
SshConnectionOptions,
protocol::MessageId,
proxy::ProxyLaunchError,
transport::{
docker::{DockerConnectionOptions, DockerExecConnection},
ssh::SshRemoteConnection,
wsl::{WslConnectionOptions, WslRemoteConnection},
},
};
use anyhow::{Context as _, Result, anyhow};
use askpass::EncryptedPassword;
use async_trait::async_trait;
use collections::HashMap;
use futures::{
Future, FutureExt as _, StreamExt as _,
channel::{
mpsc::{self, Sender, UnboundedReceiver, UnboundedSender},
oneshot,
},
future::{BoxFuture, Shared},
select, select_biased,
};
use gpui::{
App, AppContext as _, AsyncApp, BackgroundExecutor, BorrowAppContext, Context, Entity,
EventEmitter, FutureExt, Global, Task, WeakEntity,
};
use parking_lot::Mutex;
use release_channel::ReleaseChannel;
use rpc::{
AnyProtoClient, ErrorExt, ProtoClient, ProtoMessageHandlerSet, RpcError,
proto::{self, Envelope, EnvelopedMessage, PeerId, RequestMessage, build_typed_envelope},
};
use semver::Version;
use std::{
collections::VecDeque,
fmt,
ops::ControlFlow,
path::PathBuf,
sync::{
Arc, Weak,
atomic::{AtomicU32, AtomicU64, Ordering::SeqCst},
},
time::{Duration, Instant},
};
use util::{
ResultExt,
paths::{PathStyle, RemotePathBuf},
};
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum RemoteOs {
Linux,
MacOs,
Windows,
}
impl RemoteOs {
pub fn as_str(&self) -> &'static str {
match self {
RemoteOs::Linux => "linux",
RemoteOs::MacOs => "macos",
RemoteOs::Windows => "windows",
}
}
pub fn is_windows(&self) -> bool {
matches!(self, RemoteOs::Windows)
}
}
impl std::fmt::Display for RemoteOs {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Copy, Clone, Debug, PartialEq, Eq)]
pub enum RemoteArch {
X86_64,
Aarch64,
}
impl RemoteArch {
pub fn as_str(&self) -> &'static str {
match self {
RemoteArch::X86_64 => "x86_64",
RemoteArch::Aarch64 => "aarch64",
}
}
}
impl std::fmt::Display for RemoteArch {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Copy, Clone, Debug)]
pub struct RemotePlatform {
pub os: RemoteOs,
pub arch: RemoteArch,
}
#[derive(Clone, Debug)]
pub struct CommandTemplate {
pub program: String,
pub args: Vec<String>,
pub env: HashMap<String, String>,
}
pub trait RemoteClientDelegate: Send + Sync {
fn ask_password(
&self,
prompt: String,
tx: oneshot::Sender<EncryptedPassword>,
cx: &mut AsyncApp,
);
fn get_download_url(
&self,
platform: RemotePlatform,
release_channel: ReleaseChannel,
version: Option<Version>,
cx: &mut AsyncApp,
) -> Task<Result<Option<String>>>;
fn download_server_binary_locally(
&self,
platform: RemotePlatform,
release_channel: ReleaseChannel,
version: Option<Version>,
cx: &mut AsyncApp,
) -> Task<Result<PathBuf>>;
fn set_status(&self, status: Option<&str>, cx: &mut AsyncApp);
}
const MAX_MISSED_HEARTBEATS: usize = 5;
const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5);
const INITIAL_CONNECTION_TIMEOUT: Duration =
Duration::from_secs(if cfg!(debug_assertions) { 5 } else { 60 });
const MAX_RECONNECT_ATTEMPTS: usize = 3;
enum State {
Connecting,
Connected {
remote_connection: Arc<dyn RemoteConnection>,
delegate: Arc<dyn RemoteClientDelegate>,
multiplex_task: Task<Result<()>>,
heartbeat_task: Task<Result<()>>,
},
HeartbeatMissed {
missed_heartbeats: usize,
ssh_connection: Arc<dyn RemoteConnection>,
delegate: Arc<dyn RemoteClientDelegate>,
multiplex_task: Task<Result<()>>,
heartbeat_task: Task<Result<()>>,
},
Reconnecting,
ReconnectFailed {
ssh_connection: Arc<dyn RemoteConnection>,
delegate: Arc<dyn RemoteClientDelegate>,
error: anyhow::Error,
attempts: usize,
},
ReconnectExhausted,
ServerNotRunning,
}
impl fmt::Display for State {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Connecting => write!(f, "connecting"),
Self::Connected { .. } => write!(f, "connected"),
Self::Reconnecting => write!(f, "reconnecting"),
Self::ReconnectFailed { .. } => write!(f, "reconnect failed"),
Self::ReconnectExhausted => write!(f, "reconnect exhausted"),
Self::HeartbeatMissed { .. } => write!(f, "heartbeat missed"),
Self::ServerNotRunning { .. } => write!(f, "server not running"),
}
}
}
impl State {
fn remote_connection(&self) -> Option<Arc<dyn RemoteConnection>> {
match self {
Self::Connected {
remote_connection: ssh_connection,
..
} => Some(ssh_connection.clone()),
Self::HeartbeatMissed { ssh_connection, .. } => Some(ssh_connection.clone()),
Self::ReconnectFailed { ssh_connection, .. } => Some(ssh_connection.clone()),
_ => None,
}
}
fn can_reconnect(&self) -> bool {
match self {
Self::Connected { .. }
| Self::HeartbeatMissed { .. }
| Self::ReconnectFailed { .. } => true,
State::Connecting
| State::Reconnecting
| State::ReconnectExhausted
| State::ServerNotRunning => false,
}
}
fn is_reconnect_failed(&self) -> bool {
matches!(self, Self::ReconnectFailed { .. })
}
fn is_reconnect_exhausted(&self) -> bool {
matches!(self, Self::ReconnectExhausted { .. })
}
fn is_server_not_running(&self) -> bool {
matches!(self, Self::ServerNotRunning)
}
fn is_reconnecting(&self) -> bool {
matches!(self, Self::Reconnecting { .. })
}
fn heartbeat_recovered(self) -> Self {
match self {
Self::HeartbeatMissed {
ssh_connection,
delegate,
multiplex_task,
heartbeat_task,
..
} => Self::Connected {
remote_connection: ssh_connection,
delegate,
multiplex_task,
heartbeat_task,
},
_ => self,
}
}
fn heartbeat_missed(self) -> Self {
match self {
Self::Connected {
remote_connection: ssh_connection,
delegate,
multiplex_task,
heartbeat_task,
} => Self::HeartbeatMissed {
missed_heartbeats: 1,
ssh_connection,
delegate,
multiplex_task,
heartbeat_task,
},
Self::HeartbeatMissed {
missed_heartbeats,
ssh_connection,
delegate,
multiplex_task,
heartbeat_task,
} => Self::HeartbeatMissed {
missed_heartbeats: missed_heartbeats + 1,
ssh_connection,
delegate,
multiplex_task,
heartbeat_task,
},
_ => self,
}
}
}
/// The state of the ssh connection.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ConnectionState {
Connecting,
Connected,
HeartbeatMissed,
Reconnecting,
Disconnected,
}
impl From<&State> for ConnectionState {
fn from(value: &State) -> Self {
match value {
State::Connecting => Self::Connecting,
State::Connected { .. } => Self::Connected,
State::Reconnecting | State::ReconnectFailed { .. } => Self::Reconnecting,
State::HeartbeatMissed { .. } => Self::HeartbeatMissed,
State::ReconnectExhausted => Self::Disconnected,
State::ServerNotRunning => Self::Disconnected,
}
}
}
pub struct RemoteClient {
client: Arc<ChannelClient>,
unique_identifier: String,
connection_options: RemoteConnectionOptions,
path_style: PathStyle,
state: Option<State>,
}
#[derive(Debug)]
pub enum RemoteClientEvent {
Disconnected,
}
impl EventEmitter<RemoteClientEvent> for RemoteClient {}
/// Identifies the socket on the remote server so that reconnects
/// can re-join the same project.
pub enum ConnectionIdentifier {
Setup(u64),
Workspace(i64),
}
static NEXT_ID: AtomicU64 = AtomicU64::new(1);
impl ConnectionIdentifier {
pub fn setup() -> Self {
Self::Setup(NEXT_ID.fetch_add(1, SeqCst))
}
// This string gets used in a socket name, and so must be relatively short.
// The total length of:
// /home/{username}/.local/share/zed/server_state/{name}/stdout.sock
// Must be less than about 100 characters
// https://unix.stackexchange.com/questions/367008/why-is-socket-path-length-limited-to-a-hundred-chars
// So our strings should be at most 20 characters or so.
fn to_string(&self, cx: &App) -> String {
let identifier_prefix = match ReleaseChannel::global(cx) {
ReleaseChannel::Stable => "".to_string(),
release_channel => format!("{}-", release_channel.dev_name()),
};
match self {
Self::Setup(setup_id) => format!("{identifier_prefix}setup-{setup_id}"),
Self::Workspace(workspace_id) => {
format!("{identifier_prefix}workspace-{workspace_id}",)
}
}
}
}
pub async fn connect(
connection_options: RemoteConnectionOptions,
delegate: Arc<dyn RemoteClientDelegate>,
cx: &mut AsyncApp,
) -> Result<Arc<dyn RemoteConnection>> {
cx.update(|cx| {
cx.update_default_global(|pool: &mut ConnectionPool, cx| {
pool.connect(connection_options.clone(), delegate.clone(), cx)
})
})?
.await
.map_err(|e| e.cloned())
}
impl RemoteClient {
pub fn new(
unique_identifier: ConnectionIdentifier,
remote_connection: Arc<dyn RemoteConnection>,
cancellation: oneshot::Receiver<()>,
delegate: Arc<dyn RemoteClientDelegate>,
cx: &mut App,
) -> Task<Result<Option<Entity<Self>>>> {
let unique_identifier = unique_identifier.to_string(cx);
cx.spawn(async move |cx| {
let success = Box::pin(async move {
let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
let client = cx.update(|cx| {
ChannelClient::new(
incoming_rx,
outgoing_tx,
cx,
"client",
remote_connection.has_wsl_interop(),
)
})?;
let path_style = remote_connection.path_style();
let this = cx.new(|_| Self {
client: client.clone(),
unique_identifier: unique_identifier.clone(),
connection_options: remote_connection.connection_options(),
path_style,
state: Some(State::Connecting),
})?;
let io_task = remote_connection.start_proxy(
unique_identifier,
false,
incoming_tx,
outgoing_rx,
connection_activity_tx,
delegate.clone(),
cx,
);
let ready = client
.wait_for_remote_started()
.with_timeout(INITIAL_CONNECTION_TIMEOUT, cx.background_executor())
.await;
match ready {
Ok(Some(_)) => {}
Ok(None) => {
let mut error = "remote client exited before becoming ready".to_owned();
if let Some(status) = io_task.now_or_never() {
match status {
Ok(exit_code) => {
error.push_str(&format!(", exit_code={exit_code:?}"))
}
Err(e) => error.push_str(&format!(", error={e:?}")),
}
}
let error = anyhow::anyhow!("{error}");
log::error!("failed to establish connection: {}", error);
return Err(error);
}
Err(_) => {
let mut error =
"remote client did not become ready within the timeout".to_owned();
if let Some(status) = io_task.now_or_never() {
match status {
Ok(exit_code) => {
error.push_str(&format!(", exit_code={exit_code:?}"))
}
Err(e) => error.push_str(&format!(", error={e:?}")),
}
}
let error = anyhow::anyhow!("{error}");
log::error!("failed to establish connection: {}", error);
return Err(error);
}
}
let multiplex_task = Self::monitor(this.downgrade(), io_task, cx);
if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
log::error!("failed to establish connection: {}", error);
return Err(error);
}
let heartbeat_task = Self::heartbeat(this.downgrade(), connection_activity_rx, cx);
this.update(cx, |this, _| {
this.state = Some(State::Connected {
remote_connection,
delegate,
multiplex_task,
heartbeat_task,
});
})?;
Ok(Some(this))
});
select! {
_ = cancellation.fuse() => {
Ok(None)
}
result = success.fuse() => result
}
})
}
pub fn proto_client_from_channels(
incoming_rx: mpsc::UnboundedReceiver<Envelope>,
outgoing_tx: mpsc::UnboundedSender<Envelope>,
cx: &App,
name: &'static str,
has_wsl_interop: bool,
) -> AnyProtoClient {
ChannelClient::new(incoming_rx, outgoing_tx, cx, name, has_wsl_interop).into()
}
pub fn shutdown_processes<T: RequestMessage>(
&mut self,
shutdown_request: Option<T>,
executor: BackgroundExecutor,
) -> Option<impl Future<Output = ()> + use<T>> {
let state = self.state.take()?;
log::info!("shutting down ssh processes");
let State::Connected {
multiplex_task,
heartbeat_task,
remote_connection: ssh_connection,
delegate,
} = state
else {
return None;
};
let client = self.client.clone();
Some(async move {
if let Some(shutdown_request) = shutdown_request {
client.send(shutdown_request).log_err();
// We wait 50ms instead of waiting for a response, because
// waiting for a response would require us to wait on the main thread
// which we want to avoid in an `on_app_quit` callback.
executor.timer(Duration::from_millis(50)).await;
}
// Drop `multiplex_task` because it owns our ssh_proxy_process, which is a
// child of master_process.
drop(multiplex_task);
// Now drop the rest of state, which kills master process.
drop(heartbeat_task);
drop(ssh_connection);
drop(delegate);
})
}
fn reconnect(&mut self, cx: &mut Context<Self>) -> Result<()> {
let can_reconnect = self
.state
.as_ref()
.map(|state| state.can_reconnect())
.unwrap_or(false);
if !can_reconnect {
log::info!("aborting reconnect, because not in state that allows reconnecting");
let error = if let Some(state) = self.state.as_ref() {
format!("invalid state, cannot reconnect while in state {state}")
} else {
"no state set".to_string()
};
anyhow::bail!(error);
}
let state = self.state.take().unwrap();
let (attempts, remote_connection, delegate) = match state {
State::Connected {
remote_connection: ssh_connection,
delegate,
multiplex_task,
heartbeat_task,
}
| State::HeartbeatMissed {
ssh_connection,
delegate,
multiplex_task,
heartbeat_task,
..
} => {
drop(multiplex_task);
drop(heartbeat_task);
(0, ssh_connection, delegate)
}
State::ReconnectFailed {
attempts,
ssh_connection,
delegate,
..
} => (attempts, ssh_connection, delegate),
State::Connecting
| State::Reconnecting
| State::ReconnectExhausted
| State::ServerNotRunning => unreachable!(),
};
let attempts = attempts + 1;
if attempts > MAX_RECONNECT_ATTEMPTS {
log::error!(
"Failed to reconnect to after {} attempts, giving up",
MAX_RECONNECT_ATTEMPTS
);
self.set_state(State::ReconnectExhausted, cx);
return Ok(());
}
self.set_state(State::Reconnecting, cx);
log::info!("Trying to reconnect to ssh server... Attempt {}", attempts);
let unique_identifier = self.unique_identifier.clone();
let client = self.client.clone();
let reconnect_task = cx.spawn(async move |this, cx| {
macro_rules! failed {
($error:expr, $attempts:expr, $ssh_connection:expr, $delegate:expr) => {
delegate.set_status(Some(&format!("{error:#}", error = $error)), cx);
return State::ReconnectFailed {
error: anyhow!($error),
attempts: $attempts,
ssh_connection: $ssh_connection,
delegate: $delegate,
};
};
}
if let Err(error) = remote_connection
.kill()
.await
.context("Failed to kill ssh process")
{
failed!(error, attempts, remote_connection, delegate);
};
let connection_options = remote_connection.connection_options();
let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
let (ssh_connection, io_task) = match async {
let ssh_connection = cx
.update_global(|pool: &mut ConnectionPool, cx| {
pool.connect(connection_options, delegate.clone(), cx)
})?
.await
.map_err(|error| error.cloned())?;
let io_task = ssh_connection.start_proxy(
unique_identifier,
true,
incoming_tx,
outgoing_rx,
connection_activity_tx,
delegate.clone(),
cx,
);
anyhow::Ok((ssh_connection, io_task))
}
.await
{
Ok((ssh_connection, io_task)) => (ssh_connection, io_task),
Err(error) => {
failed!(error, attempts, remote_connection, delegate);
}
};
let multiplex_task = Self::monitor(this.clone(), io_task, cx);
client.reconnect(incoming_rx, outgoing_tx, cx);
if let Err(error) = client.resync(HEARTBEAT_TIMEOUT).await {
failed!(error, attempts, ssh_connection, delegate);
};
State::Connected {
remote_connection: ssh_connection,
delegate,
multiplex_task,
heartbeat_task: Self::heartbeat(this.clone(), connection_activity_rx, cx),
}
});
cx.spawn(async move |this, cx| {
let new_state = reconnect_task.await;
this.update(cx, |this, cx| {
this.try_set_state(cx, |old_state| {
if old_state.is_reconnecting() {
match &new_state {
State::Connecting
| State::Reconnecting
| State::HeartbeatMissed { .. }
| State::ServerNotRunning => {}
State::Connected { .. } => {
log::info!("Successfully reconnected");
}
State::ReconnectFailed {
error, attempts, ..
} => {
log::error!(
"Reconnect attempt {} failed: {:?}. Starting new attempt...",
attempts,
error
);
}
State::ReconnectExhausted => {
log::error!("Reconnect attempt failed and all attempts exhausted");
}
}
Some(new_state)
} else {
None
}
});
if this.state_is(State::is_reconnect_failed) {
this.reconnect(cx)
} else if this.state_is(State::is_reconnect_exhausted) {
Ok(())
} else {
log::debug!("State has transition from Reconnecting into new state while attempting reconnect.");
Ok(())
}
})
})
.detach_and_log_err(cx);
Ok(())
}
fn heartbeat(
this: WeakEntity<Self>,
mut connection_activity_rx: mpsc::Receiver<()>,
cx: &mut AsyncApp,
) -> Task<Result<()>> {
let Ok(client) = this.read_with(cx, |this, _| this.client.clone()) else {
return Task::ready(Err(anyhow!("SshRemoteClient lost")));
};
cx.spawn(async move |cx| {
let mut missed_heartbeats = 0;
let keepalive_timer = cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse();
futures::pin_mut!(keepalive_timer);
loop {
select_biased! {
result = connection_activity_rx.next().fuse() => {
if result.is_none() {
log::warn!("ssh heartbeat: connection activity channel has been dropped. stopping.");
return Ok(());
}
if missed_heartbeats != 0 {
missed_heartbeats = 0;
let _ =this.update(cx, |this, cx| {
this.handle_heartbeat_result(missed_heartbeats, cx)
})?;
}
}
_ = keepalive_timer => {
log::debug!("Sending heartbeat to server...");
let result = select_biased! {
_ = connection_activity_rx.next().fuse() => {
Ok(())
}
ping_result = client.ping(HEARTBEAT_TIMEOUT).fuse() => {
ping_result
}
};
if result.is_err() {
missed_heartbeats += 1;
log::warn!(
"No heartbeat from server after {:?}. Missed heartbeat {} out of {}.",
HEARTBEAT_TIMEOUT,
missed_heartbeats,
MAX_MISSED_HEARTBEATS
);
} else if missed_heartbeats != 0 {
missed_heartbeats = 0;
} else {
continue;
}
let result = this.update(cx, |this, cx| {
this.handle_heartbeat_result(missed_heartbeats, cx)
})?;
if result.is_break() {
return Ok(());
}
}
}
keepalive_timer.set(cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse());
}
})
}
fn handle_heartbeat_result(
&mut self,
missed_heartbeats: usize,
cx: &mut Context<Self>,
) -> ControlFlow<()> {
let state = self.state.take().unwrap();
let next_state = if missed_heartbeats > 0 {
state.heartbeat_missed()
} else {
state.heartbeat_recovered()
};
self.set_state(next_state, cx);
if missed_heartbeats >= MAX_MISSED_HEARTBEATS {
log::error!(
"Missed last {} heartbeats. Reconnecting...",
missed_heartbeats
);
self.reconnect(cx)
.context("failed to start reconnect process after missing heartbeats")
.log_err();
ControlFlow::Break(())
} else {
ControlFlow::Continue(())
}
}
fn monitor(
this: WeakEntity<Self>,
io_task: Task<Result<i32>>,
cx: &AsyncApp,
) -> Task<Result<()>> {
cx.spawn(async move |cx| {
let result = io_task.await;
match result {
Ok(exit_code) => {
if let Some(error) = ProxyLaunchError::from_exit_code(exit_code) {
match error {
ProxyLaunchError::ServerNotRunning => {
log::error!("failed to reconnect because server is not running");
this.update(cx, |this, cx| {
this.set_state(State::ServerNotRunning, cx);
})?;
}
}
} else if exit_code > 0 {
log::error!("proxy process terminated unexpectedly");
this.update(cx, |this, cx| {
this.reconnect(cx).ok();
})?;
}
}
Err(error) => {
log::warn!("ssh io task died with error: {:?}. reconnecting...", error);
this.update(cx, |this, cx| {
this.reconnect(cx).ok();
})?;
}
}
Ok(())
})
}
fn state_is(&self, check: impl FnOnce(&State) -> bool) -> bool {
self.state.as_ref().is_some_and(check)
}
fn try_set_state(&mut self, cx: &mut Context<Self>, map: impl FnOnce(&State) -> Option<State>) {
let new_state = self.state.as_ref().and_then(map);
if let Some(new_state) = new_state {
self.state.replace(new_state);
cx.notify();
}
}
fn set_state(&mut self, state: State, cx: &mut Context<Self>) {
log::info!("setting state to '{}'", &state);
let is_reconnect_exhausted = state.is_reconnect_exhausted();
let is_server_not_running = state.is_server_not_running();
self.state.replace(state);
if is_reconnect_exhausted || is_server_not_running {
cx.emit(RemoteClientEvent::Disconnected);
}
cx.notify();
}
pub fn shell(&self) -> Option<String> {
Some(self.remote_connection()?.shell())
}
pub fn default_system_shell(&self) -> Option<String> {
Some(self.remote_connection()?.default_system_shell())
}
pub fn shares_network_interface(&self) -> bool {
self.remote_connection()
.map_or(false, |connection| connection.shares_network_interface())
}
pub fn build_command(
&self,
program: Option<String>,
args: &[String],
env: &HashMap<String, String>,
working_dir: Option<String>,
port_forward: Option<(u16, String, u16)>,
) -> Result<CommandTemplate> {
let Some(connection) = self.remote_connection() else {
return Err(anyhow!("no ssh connection"));
};
connection.build_command(program, args, env, working_dir, port_forward)
}
pub fn build_forward_ports_command(
&self,
forwards: Vec<(u16, String, u16)>,
) -> Result<CommandTemplate> {
let Some(connection) = self.remote_connection() else {
return Err(anyhow!("no ssh connection"));
};
connection.build_forward_ports_command(forwards)
}
pub fn upload_directory(
&self,
src_path: PathBuf,
dest_path: RemotePathBuf,
cx: &App,
) -> Task<Result<()>> {
let Some(connection) = self.remote_connection() else {
return Task::ready(Err(anyhow!("no ssh connection")));
};
connection.upload_directory(src_path, dest_path, cx)
}
pub fn proto_client(&self) -> AnyProtoClient {
self.client.clone().into()
}
pub fn connection_options(&self) -> RemoteConnectionOptions {
self.connection_options.clone()
}
pub fn connection(&self) -> Option<Arc<dyn RemoteConnection>> {
if let State::Connected {
remote_connection, ..
} = self.state.as_ref()?
{
Some(remote_connection.clone())
} else {
None
}
}
pub fn connection_state(&self) -> ConnectionState {
self.state
.as_ref()
.map(ConnectionState::from)
.unwrap_or(ConnectionState::Disconnected)
}
pub fn is_disconnected(&self) -> bool {
self.connection_state() == ConnectionState::Disconnected
}
pub fn path_style(&self) -> PathStyle {
self.path_style
}
#[cfg(any(test, feature = "test-support"))]
pub fn simulate_disconnect(&self, client_cx: &mut App) -> Task<()> {
let opts = self.connection_options();
client_cx.spawn(async move |cx| {
let connection = cx
.update_global(|c: &mut ConnectionPool, _| {
if let Some(ConnectionPoolEntry::Connecting(c)) = c.connections.get(&opts) {
c.clone()
} else {
panic!("missing test connection")
}
})
.unwrap()
.await
.unwrap();
connection.simulate_disconnect(cx);
})
}
#[cfg(any(test, feature = "test-support"))]
pub fn fake_server(
client_cx: &mut gpui::TestAppContext,
server_cx: &mut gpui::TestAppContext,
) -> (RemoteConnectionOptions, AnyProtoClient) {
use crate::transport::ssh::SshConnectionHost;
let port = client_cx
.update(|cx| cx.default_global::<ConnectionPool>().connections.len() as u16 + 1);
let opts = RemoteConnectionOptions::Ssh(SshConnectionOptions {
host: SshConnectionHost::from("<fake>".to_string()),
port: Some(port),
..Default::default()
});
let (outgoing_tx, _) = mpsc::unbounded::<Envelope>();
let (_, incoming_rx) = mpsc::unbounded::<Envelope>();
let server_client = server_cx
.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "fake-server", false));
let connection: Arc<dyn RemoteConnection> = Arc::new(fake::FakeRemoteConnection {
connection_options: opts.clone(),
server_cx: fake::SendableCx::new(server_cx),
server_channel: server_client.clone(),
});
client_cx.update(|cx| {
cx.update_default_global(|c: &mut ConnectionPool, cx| {
c.connections.insert(
opts.clone(),
ConnectionPoolEntry::Connecting(
cx.background_spawn({
let connection = connection.clone();
async move { Ok(connection.clone()) }
})
.shared(),
),
);
})
});
(opts, server_client.into())
}
#[cfg(any(test, feature = "test-support"))]
pub async fn fake_client(
opts: RemoteConnectionOptions,
client_cx: &mut gpui::TestAppContext,
) -> Entity<Self> {
let (_tx, rx) = oneshot::channel();
let mut cx = client_cx.to_async();
let connection = connect(opts, Arc::new(fake::Delegate), &mut cx)
.await
.unwrap();
client_cx
.update(|cx| {
Self::new(
ConnectionIdentifier::setup(),
connection,
rx,
Arc::new(fake::Delegate),
cx,
)
})
.await
.unwrap()
.unwrap()
}
fn remote_connection(&self) -> Option<Arc<dyn RemoteConnection>> {
self.state
.as_ref()
.and_then(|state| state.remote_connection())
}
}
enum ConnectionPoolEntry {
Connecting(Shared<Task<Result<Arc<dyn RemoteConnection>, Arc<anyhow::Error>>>>),
Connected(Weak<dyn RemoteConnection>),
}
#[derive(Default)]
struct ConnectionPool {
connections: HashMap<RemoteConnectionOptions, ConnectionPoolEntry>,
}
impl Global for ConnectionPool {}
impl ConnectionPool {
pub fn connect(
&mut self,
opts: RemoteConnectionOptions,
delegate: Arc<dyn RemoteClientDelegate>,
cx: &mut App,
) -> Shared<Task<Result<Arc<dyn RemoteConnection>, Arc<anyhow::Error>>>> {
let connection = self.connections.get(&opts);
match connection {
Some(ConnectionPoolEntry::Connecting(task)) => {
delegate.set_status(
Some("Waiting for existing connection attempt"),
&mut cx.to_async(),
);
return task.clone();
}
Some(ConnectionPoolEntry::Connected(ssh)) => {
if let Some(ssh) = ssh.upgrade()
&& !ssh.has_been_killed()
{
return Task::ready(Ok(ssh)).shared();
}
self.connections.remove(&opts);
}
None => {}
}
let task = cx
.spawn({
let opts = opts.clone();
let delegate = delegate.clone();
async move |cx| {
let connection = match opts.clone() {
RemoteConnectionOptions::Ssh(opts) => {
SshRemoteConnection::new(opts, delegate, cx)
.await
.map(|connection| Arc::new(connection) as Arc<dyn RemoteConnection>)
}
RemoteConnectionOptions::Wsl(opts) => {
WslRemoteConnection::new(opts, delegate, cx)
.await
.map(|connection| Arc::new(connection) as Arc<dyn RemoteConnection>)
}
RemoteConnectionOptions::Docker(opts) => {
DockerExecConnection::new(opts, delegate, cx)
.await
.map(|connection| Arc::new(connection) as Arc<dyn RemoteConnection>)
}
};
cx.update_global(|pool: &mut Self, _| {
debug_assert!(matches!(
pool.connections.get(&opts),
Some(ConnectionPoolEntry::Connecting(_))
));
match connection {
Ok(connection) => {
pool.connections.insert(
opts.clone(),
ConnectionPoolEntry::Connected(Arc::downgrade(&connection)),
);
Ok(connection)
}
Err(error) => {
pool.connections.remove(&opts);
Err(Arc::new(error))
}
}
})?
}
})
.shared();
self.connections
.insert(opts.clone(), ConnectionPoolEntry::Connecting(task.clone()));
task
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum RemoteConnectionOptions {
Ssh(SshConnectionOptions),
Wsl(WslConnectionOptions),
Docker(DockerConnectionOptions),
}
impl RemoteConnectionOptions {
pub fn display_name(&self) -> String {
match self {
RemoteConnectionOptions::Ssh(opts) => opts.host.to_string(),
RemoteConnectionOptions::Wsl(opts) => opts.distro_name.clone(),
RemoteConnectionOptions::Docker(opts) => opts.name.clone(),
}
}
}
impl From<SshConnectionOptions> for RemoteConnectionOptions {
fn from(opts: SshConnectionOptions) -> Self {
RemoteConnectionOptions::Ssh(opts)
}
}
impl From<WslConnectionOptions> for RemoteConnectionOptions {
fn from(opts: WslConnectionOptions) -> Self {
RemoteConnectionOptions::Wsl(opts)
}
}
#[cfg(target_os = "windows")]
/// Open a wsl path (\\wsl.localhost\<distro>\path)
#[derive(Debug, Clone, PartialEq, Eq, gpui::Action)]
#[action(namespace = workspace, no_json, no_register)]
pub struct OpenWslPath {
pub distro: WslConnectionOptions,
pub paths: Vec<PathBuf>,
}
#[async_trait(?Send)]
pub trait RemoteConnection: Send + Sync {
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>>;
fn upload_directory(
&self,
src_path: PathBuf,
dest_path: RemotePathBuf,
cx: &App,
) -> Task<Result<()>>;
async fn kill(&self) -> Result<()>;
fn has_been_killed(&self) -> bool;
fn shares_network_interface(&self) -> bool {
false
}
fn build_command(
&self,
program: Option<String>,
args: &[String],
env: &HashMap<String, String>,
working_dir: Option<String>,
port_forward: Option<(u16, String, u16)>,
) -> Result<CommandTemplate>;
fn build_forward_ports_command(
&self,
forwards: Vec<(u16, String, u16)>,
) -> Result<CommandTemplate>;
fn connection_options(&self) -> RemoteConnectionOptions;
fn path_style(&self) -> PathStyle;
fn shell(&self) -> String;
fn default_system_shell(&self) -> String;
fn has_wsl_interop(&self) -> bool;
#[cfg(any(test, feature = "test-support"))]
fn simulate_disconnect(&self, _: &AsyncApp) {}
}
type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>;
struct Signal<T> {
tx: Mutex<Option<oneshot::Sender<T>>>,
rx: Shared<Task<Option<T>>>,
}
impl<T: Send + Clone + 'static> Signal<T> {
pub fn new(cx: &App) -> Self {
let (tx, rx) = oneshot::channel();
let task = cx
.background_executor()
.spawn(async move { rx.await.ok() })
.shared();
Self {
tx: Mutex::new(Some(tx)),
rx: task,
}
}
fn set(&self, value: T) {
if let Some(tx) = self.tx.lock().take() {
let _ = tx.send(value);
}
}
fn wait(&self) -> Shared<Task<Option<T>>> {
self.rx.clone()
}
}
struct ChannelClient {
next_message_id: AtomicU32,
outgoing_tx: Mutex<mpsc::UnboundedSender<Envelope>>,
buffer: Mutex<VecDeque<Envelope>>,
response_channels: ResponseChannels,
message_handlers: Mutex<ProtoMessageHandlerSet>,
max_received: AtomicU32,
name: &'static str,
task: Mutex<Task<Result<()>>>,
remote_started: Signal<()>,
has_wsl_interop: bool,
}
impl ChannelClient {
fn new(
incoming_rx: mpsc::UnboundedReceiver<Envelope>,
outgoing_tx: mpsc::UnboundedSender<Envelope>,
cx: &App,
name: &'static str,
has_wsl_interop: bool,
) -> Arc<Self> {
Arc::new_cyclic(|this| Self {
outgoing_tx: Mutex::new(outgoing_tx),
next_message_id: AtomicU32::new(0),
max_received: AtomicU32::new(0),
response_channels: ResponseChannels::default(),
message_handlers: Default::default(),
buffer: Mutex::new(VecDeque::new()),
name,
task: Mutex::new(Self::start_handling_messages(
this.clone(),
incoming_rx,
&cx.to_async(),
)),
remote_started: Signal::new(cx),
has_wsl_interop,
})
}
fn wait_for_remote_started(&self) -> Shared<Task<Option<()>>> {
self.remote_started.wait()
}
fn start_handling_messages(
this: Weak<Self>,
mut incoming_rx: mpsc::UnboundedReceiver<Envelope>,
cx: &AsyncApp,
) -> Task<Result<()>> {
cx.spawn(async move |cx| {
if let Some(this) = this.upgrade() {
let envelope = proto::RemoteStarted {}.into_envelope(0, None, None);
this.outgoing_tx.lock().unbounded_send(envelope).ok();
};
let peer_id = PeerId { owner_id: 0, id: 0 };
while let Some(incoming) = incoming_rx.next().await {
let Some(this) = this.upgrade() else {
return anyhow::Ok(());
};
if let Some(ack_id) = incoming.ack_id {
let mut buffer = this.buffer.lock();
while buffer.front().is_some_and(|msg| msg.id <= ack_id) {
buffer.pop_front();
}
}
if let Some(proto::envelope::Payload::FlushBufferedMessages(_)) = &incoming.payload
{
log::debug!(
"{}:ssh message received. name:FlushBufferedMessages",
this.name
);
{
let buffer = this.buffer.lock();
for envelope in buffer.iter() {
this.outgoing_tx
.lock()
.unbounded_send(envelope.clone())
.ok();
}
}
let mut envelope = proto::Ack {}.into_envelope(0, Some(incoming.id), None);
envelope.id = this.next_message_id.fetch_add(1, SeqCst);
this.outgoing_tx.lock().unbounded_send(envelope).ok();
continue;
}
if let Some(proto::envelope::Payload::RemoteStarted(_)) = &incoming.payload {
this.remote_started.set(());
let mut envelope = proto::Ack {}.into_envelope(0, Some(incoming.id), None);
envelope.id = this.next_message_id.fetch_add(1, SeqCst);
this.outgoing_tx.lock().unbounded_send(envelope).ok();
continue;
}
this.max_received.store(incoming.id, SeqCst);
if let Some(request_id) = incoming.responding_to {
let request_id = MessageId(request_id);
let sender = this.response_channels.lock().remove(&request_id);
if let Some(sender) = sender {
let (tx, rx) = oneshot::channel();
if incoming.payload.is_some() {
sender.send((incoming, tx)).ok();
}
rx.await.ok();
}
} else if let Some(envelope) =
build_typed_envelope(peer_id, Instant::now(), incoming)
{
let type_name = envelope.payload_type_name();
let message_id = envelope.message_id();
if let Some(future) = ProtoMessageHandlerSet::handle_message(
&this.message_handlers,
envelope,
this.clone().into(),
cx.clone(),
) {
log::debug!("{}:ssh message received. name:{type_name}", this.name);
cx.foreground_executor()
.spawn(async move {
match future.await {
Ok(_) => {
log::debug!(
"{}:ssh message handled. name:{type_name}",
this.name
);
}
Err(error) => {
log::error!(
"{}:error handling message. type:{}, error:{:#}",
this.name,
type_name,
format!("{error:#}").lines().fold(
String::new(),
|mut message, line| {
if !message.is_empty() {
message.push(' ');
}
message.push_str(line);
message
}
)
);
}
}
})
.detach()
} else {
log::error!("{}:unhandled ssh message name:{type_name}", this.name);
if let Err(e) = AnyProtoClient::from(this.clone()).send_response(
message_id,
anyhow::anyhow!("no handler registered for {type_name}").to_proto(),
) {
log::error!(
"{}:error sending error response for {type_name}:{e:#}",
this.name
);
}
}
}
}
anyhow::Ok(())
})
}
fn reconnect(
self: &Arc<Self>,
incoming_rx: UnboundedReceiver<Envelope>,
outgoing_tx: UnboundedSender<Envelope>,
cx: &AsyncApp,
) {
*self.outgoing_tx.lock() = outgoing_tx;
*self.task.lock() = Self::start_handling_messages(Arc::downgrade(self), incoming_rx, cx);
}
fn request<T: RequestMessage>(
&self,
payload: T,
) -> impl 'static + Future<Output = Result<T::Response>> {
self.request_internal(payload, true)
}
fn request_internal<T: RequestMessage>(
&self,
payload: T,
use_buffer: bool,
) -> impl 'static + Future<Output = Result<T::Response>> {
log::debug!("ssh request start. name:{}", T::NAME);
let response =
self.request_dynamic(payload.into_envelope(0, None, None), T::NAME, use_buffer);
async move {
let response = response.await?;
log::debug!("ssh request finish. name:{}", T::NAME);
T::Response::from_envelope(response).context("received a response of the wrong type")
}
}
async fn resync(&self, timeout: Duration) -> Result<()> {
smol::future::or(
async {
self.request_internal(proto::FlushBufferedMessages {}, false)
.await?;
for envelope in self.buffer.lock().iter() {
self.outgoing_tx
.lock()
.unbounded_send(envelope.clone())
.ok();
}
Ok(())
},
async {
smol::Timer::after(timeout).await;
anyhow::bail!("Timed out resyncing remote client")
},
)
.await
}
async fn ping(&self, timeout: Duration) -> Result<()> {
smol::future::or(
async {
self.request(proto::Ping {}).await?;
Ok(())
},
async {
smol::Timer::after(timeout).await;
anyhow::bail!("Timed out pinging remote client")
},
)
.await
}
fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {
log::debug!("ssh send name:{}", T::NAME);
self.send_dynamic(payload.into_envelope(0, None, None))
}
fn request_dynamic(
&self,
mut envelope: proto::Envelope,
type_name: &'static str,
use_buffer: bool,
) -> impl 'static + Future<Output = Result<proto::Envelope>> {
envelope.id = self.next_message_id.fetch_add(1, SeqCst);
let (tx, rx) = oneshot::channel();
let mut response_channels_lock = self.response_channels.lock();
response_channels_lock.insert(MessageId(envelope.id), tx);
drop(response_channels_lock);
let result = if use_buffer {
self.send_buffered(envelope)
} else {
self.send_unbuffered(envelope)
};
async move {
if let Err(error) = &result {
log::error!("failed to send message: {error}");
anyhow::bail!("failed to send message: {error}");
}
let response = rx.await.context("connection lost")?.0;
if let Some(proto::envelope::Payload::Error(error)) = &response.payload {
return Err(RpcError::from_proto(error, type_name));
}
Ok(response)
}
}
pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {
envelope.id = self.next_message_id.fetch_add(1, SeqCst);
self.send_buffered(envelope)
}
fn send_buffered(&self, mut envelope: proto::Envelope) -> Result<()> {
envelope.ack_id = Some(self.max_received.load(SeqCst));
self.buffer.lock().push_back(envelope.clone());
// ignore errors on send (happen while we're reconnecting)
// assume that the global "disconnected" overlay is sufficient.
self.outgoing_tx.lock().unbounded_send(envelope).ok();
Ok(())
}
fn send_unbuffered(&self, mut envelope: proto::Envelope) -> Result<()> {
envelope.ack_id = Some(self.max_received.load(SeqCst));
self.outgoing_tx.lock().unbounded_send(envelope).ok();
Ok(())
}
}
impl ProtoClient for ChannelClient {
fn request(
&self,
envelope: proto::Envelope,
request_type: &'static str,
) -> BoxFuture<'static, Result<proto::Envelope>> {
self.request_dynamic(envelope, request_type, true).boxed()
}
fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
self.send_dynamic(envelope)
}
fn send_response(&self, envelope: Envelope, _message_type: &'static str) -> anyhow::Result<()> {
self.send_dynamic(envelope)
}
fn message_handler_set(&self) -> &Mutex<ProtoMessageHandlerSet> {
&self.message_handlers
}
fn is_via_collab(&self) -> bool {
false
}
fn has_wsl_interop(&self) -> bool {
self.has_wsl_interop
}
}
#[cfg(any(test, feature = "test-support"))]
mod fake {
use super::{ChannelClient, RemoteClientDelegate, RemoteConnection, RemotePlatform};
use crate::remote_client::{CommandTemplate, RemoteConnectionOptions};
use anyhow::Result;
use askpass::EncryptedPassword;
use async_trait::async_trait;
use collections::HashMap;
use futures::{
FutureExt, SinkExt, StreamExt,
channel::{
mpsc::{self, Sender},
oneshot,
},
select_biased,
};
use gpui::{App, AppContext as _, AsyncApp, Task, TestAppContext};
use release_channel::ReleaseChannel;
use rpc::proto::Envelope;
use semver::Version;
use std::{path::PathBuf, sync::Arc};
use util::paths::{PathStyle, RemotePathBuf};
pub(super) struct FakeRemoteConnection {
pub(super) connection_options: RemoteConnectionOptions,
pub(super) server_channel: Arc<ChannelClient>,
pub(super) server_cx: SendableCx,
}
pub(super) struct SendableCx(AsyncApp);
impl SendableCx {
// SAFETY: When run in test mode, GPUI is always single threaded.
pub(super) fn new(cx: &TestAppContext) -> Self {
Self(cx.to_async())
}
// SAFETY: Enforce that we're on the main thread by requiring a valid AsyncApp
fn get(&self, _: &AsyncApp) -> AsyncApp {
self.0.clone()
}
}
// SAFETY: There is no way to access a SendableCx from a different thread, see [`SendableCx::new`] and [`SendableCx::get`]
unsafe impl Send for SendableCx {}
unsafe impl Sync for SendableCx {}
#[async_trait(?Send)]
impl RemoteConnection for FakeRemoteConnection {
async fn kill(&self) -> Result<()> {
Ok(())
}
fn has_been_killed(&self) -> bool {
false
}
fn build_command(
&self,
program: Option<String>,
args: &[String],
env: &HashMap<String, String>,
_: Option<String>,
_: Option<(u16, String, u16)>,
) -> Result<CommandTemplate> {
let ssh_program = program.unwrap_or_else(|| "sh".to_string());
let mut ssh_args = Vec::new();
ssh_args.push(ssh_program);
ssh_args.extend(args.iter().cloned());
Ok(CommandTemplate {
program: "ssh".into(),
args: ssh_args,
env: env.clone(),
})
}
fn build_forward_ports_command(
&self,
forwards: Vec<(u16, String, u16)>,
) -> anyhow::Result<CommandTemplate> {
Ok(CommandTemplate {
program: "ssh".into(),
args: std::iter::once("-N".to_owned())
.chain(forwards.into_iter().map(|(local_port, host, remote_port)| {
format!("{local_port}:{host}:{remote_port}")
}))
.collect(),
env: Default::default(),
})
}
fn upload_directory(
&self,
_src_path: PathBuf,
_dest_path: RemotePathBuf,
_cx: &App,
) -> Task<Result<()>> {
unreachable!()
}
fn connection_options(&self) -> RemoteConnectionOptions {
self.connection_options.clone()
}
fn simulate_disconnect(&self, cx: &AsyncApp) {
let (outgoing_tx, _) = mpsc::unbounded::<Envelope>();
let (_, incoming_rx) = mpsc::unbounded::<Envelope>();
self.server_channel
.reconnect(incoming_rx, outgoing_tx, &self.server_cx.get(cx));
}
fn start_proxy(
&self,
_unique_identifier: String,
_reconnect: bool,
mut client_incoming_tx: mpsc::UnboundedSender<Envelope>,
mut client_outgoing_rx: mpsc::UnboundedReceiver<Envelope>,
mut connection_activity_tx: Sender<()>,
_delegate: Arc<dyn RemoteClientDelegate>,
cx: &mut AsyncApp,
) -> Task<Result<i32>> {
let (mut server_incoming_tx, server_incoming_rx) = mpsc::unbounded::<Envelope>();
let (server_outgoing_tx, mut server_outgoing_rx) = mpsc::unbounded::<Envelope>();
self.server_channel.reconnect(
server_incoming_rx,
server_outgoing_tx,
&self.server_cx.get(cx),
);
cx.background_spawn(async move {
loop {
select_biased! {
server_to_client = server_outgoing_rx.next().fuse() => {
let Some(server_to_client) = server_to_client else {
return Ok(1)
};
connection_activity_tx.try_send(()).ok();
client_incoming_tx.send(server_to_client).await.ok();
}
client_to_server = client_outgoing_rx.next().fuse() => {
let Some(client_to_server) = client_to_server else {
return Ok(1)
};
server_incoming_tx.send(client_to_server).await.ok();
}
}
}
})
}
fn path_style(&self) -> PathStyle {
PathStyle::local()
}
fn shell(&self) -> String {
"sh".to_owned()
}
fn default_system_shell(&self) -> String {
"sh".to_owned()
}
fn has_wsl_interop(&self) -> bool {
false
}
}
pub(super) struct Delegate;
impl RemoteClientDelegate for Delegate {
fn ask_password(&self, _: String, _: oneshot::Sender<EncryptedPassword>, _: &mut AsyncApp) {
unreachable!()
}
fn download_server_binary_locally(
&self,
_: RemotePlatform,
_: ReleaseChannel,
_: Option<Version>,
_: &mut AsyncApp,
) -> Task<Result<PathBuf>> {
unreachable!()
}
fn get_download_url(
&self,
_platform: RemotePlatform,
_release_channel: ReleaseChannel,
_version: Option<Version>,
_cx: &mut AsyncApp,
) -> Task<Result<Option<String>>> {
unreachable!()
}
fn set_status(&self, _: Option<&str>, _: &mut AsyncApp) {}
}
}