Files
oak-gpui/crates/remote/src/remote_client.rs
T
Piotr OsiewiczandConrad Irwin 2d55c088cc releases: Add build number to Nightly builds (#42990)
- **Remove semantic_version crate and use semver instead**
- **Update upload-nightly**


Release Notes:

- N/A

---------

Co-authored-by: Conrad Irwin <conrad.irwin@gmail.com>
2025-11-24 13:34:04 +01:00

1687 lines
57 KiB
Rust

use crate::{
SshConnectionOptions,
protocol::MessageId,
proxy::ProxyLaunchError,
transport::{
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)]
pub struct RemotePlatform {
pub os: &'static str,
pub arch: &'static str,
}
#[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(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"))?;
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,
) -> AnyProtoClient {
ChannelClient::new(incoming_rx, outgoing_tx, cx, name).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) {
let port = client_cx
.update(|cx| cx.default_global::<ConnectionPool>().connections.len() as u16 + 1);
let opts = RemoteConnectionOptions::Ssh(SshConnectionOptions {
host: "<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"));
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>)
}
};
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),
}
impl RemoteConnectionOptions {
pub fn display_name(&self) -> String {
match self {
RemoteConnectionOptions::Ssh(opts) => opts.host.clone(),
RemoteConnectionOptions::Wsl(opts) => opts.distro_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;
#[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<()>,
}
impl ChannelClient {
fn new(
incoming_rx: mpsc::UnboundedReceiver<Envelope>,
outgoing_tx: mpsc::UnboundedSender<Envelope>,
cx: &App,
name: &'static str,
) -> 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),
})
}
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
}
}
#[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()
}
}
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) {}
}
}