1use crate::{
2 json_log::LogRecord,
3 protocol::{
4 MESSAGE_LEN_SIZE, MessageId, message_len_from_buffer, read_message_with_len, write_message,
5 },
6 proxy::ProxyLaunchError,
7};
8use anyhow::{Context as _, Result, anyhow};
9use async_trait::async_trait;
10use collections::HashMap;
11use futures::{
12 AsyncReadExt as _, Future, FutureExt as _, StreamExt as _,
13 channel::{
14 mpsc::{self, Sender, UnboundedReceiver, UnboundedSender},
15 oneshot,
16 },
17 future::{BoxFuture, Shared},
18 select, select_biased,
19};
20use gpui::{
21 App, AppContext as _, AsyncApp, BackgroundExecutor, BorrowAppContext, Context, Entity,
22 EventEmitter, Global, SemanticVersion, Task, WeakEntity,
23};
24use itertools::Itertools;
25use parking_lot::Mutex;
26use paths;
27use release_channel::{AppCommitSha, AppVersion, ReleaseChannel};
28use rpc::{
29 AnyProtoClient, EntityMessageSubscriber, ErrorExt, ProtoClient, ProtoMessageHandlerSet,
30 RpcError,
31 proto::{self, Envelope, EnvelopedMessage, PeerId, RequestMessage, build_typed_envelope},
32};
33use schemars::JsonSchema;
34use serde::{Deserialize, Serialize};
35use smol::{
36 fs,
37 process::{self, Child, Stdio},
38};
39use std::{
40 any::TypeId,
41 collections::VecDeque,
42 fmt, iter,
43 ops::ControlFlow,
44 path::{Path, PathBuf},
45 sync::{
46 Arc, Weak,
47 atomic::{AtomicU32, AtomicU64, Ordering::SeqCst},
48 },
49 time::{Duration, Instant},
50};
51use tempfile::TempDir;
52use util::ResultExt;
53
54#[derive(
55 Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, serde::Serialize, serde::Deserialize,
56)]
57pub struct SshProjectId(pub u64);
58
59#[derive(Clone)]
60pub struct SshSocket {
61 connection_options: SshConnectionOptions,
62 socket_path: PathBuf,
63}
64
65#[derive(Debug, Clone, PartialEq, Eq, Hash, Deserialize, Serialize, JsonSchema)]
66pub struct SshPortForwardOption {
67 #[serde(skip_serializing_if = "Option::is_none")]
68 pub local_host: Option<String>,
69 pub local_port: u16,
70 #[serde(skip_serializing_if = "Option::is_none")]
71 pub remote_host: Option<String>,
72 pub remote_port: u16,
73}
74
75#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)]
76pub struct SshConnectionOptions {
77 pub host: String,
78 pub username: Option<String>,
79 pub port: Option<u16>,
80 pub password: Option<String>,
81 pub args: Option<Vec<String>>,
82 pub port_forwards: Option<Vec<SshPortForwardOption>>,
83
84 pub nickname: Option<String>,
85 pub upload_binary_over_ssh: bool,
86}
87
88#[macro_export]
89macro_rules! shell_script {
90 ($fmt:expr, $($name:ident = $arg:expr),+ $(,)?) => {{
91 format!(
92 $fmt,
93 $(
94 $name = shlex::try_quote($arg).unwrap()
95 ),+
96 )
97 }};
98}
99
100fn parse_port_number(port_str: &str) -> Result<u16> {
101 port_str
102 .parse()
103 .map_err(|e| anyhow!("Invalid port number: {}: {}", port_str, e))
104}
105
106fn parse_port_forward_spec(spec: &str) -> Result<SshPortForwardOption> {
107 let parts: Vec<&str> = spec.split(':').collect();
108
109 match parts.len() {
110 4 => {
111 let local_port = parse_port_number(parts[1])?;
112 let remote_port = parse_port_number(parts[3])?;
113
114 Ok(SshPortForwardOption {
115 local_host: Some(parts[0].to_string()),
116 local_port,
117 remote_host: Some(parts[2].to_string()),
118 remote_port,
119 })
120 }
121 3 => {
122 let local_port = parse_port_number(parts[0])?;
123 let remote_port = parse_port_number(parts[2])?;
124
125 Ok(SshPortForwardOption {
126 local_host: None,
127 local_port,
128 remote_host: Some(parts[1].to_string()),
129 remote_port,
130 })
131 }
132 _ => anyhow::bail!("Invalid port forward format"),
133 }
134}
135
136impl SshConnectionOptions {
137 pub fn parse_command_line(input: &str) -> Result<Self> {
138 let input = input.trim_start_matches("ssh ");
139 let mut hostname: Option<String> = None;
140 let mut username: Option<String> = None;
141 let mut port: Option<u16> = None;
142 let mut args = Vec::new();
143 let mut port_forwards: Vec<SshPortForwardOption> = Vec::new();
144
145 // disallowed: -E, -e, -F, -f, -G, -g, -M, -N, -n, -O, -q, -S, -s, -T, -t, -V, -v, -W
146 const ALLOWED_OPTS: &[&str] = &[
147 "-4", "-6", "-A", "-a", "-C", "-K", "-k", "-X", "-x", "-Y", "-y",
148 ];
149 const ALLOWED_ARGS: &[&str] = &[
150 "-B", "-b", "-c", "-D", "-F", "-I", "-i", "-J", "-l", "-m", "-o", "-P", "-p", "-R",
151 "-w",
152 ];
153
154 let mut tokens = shlex::split(input)
155 .ok_or_else(|| anyhow!("invalid input"))?
156 .into_iter();
157
158 'outer: while let Some(arg) = tokens.next() {
159 if ALLOWED_OPTS.contains(&(&arg as &str)) {
160 args.push(arg.to_string());
161 continue;
162 }
163 if arg == "-p" {
164 port = tokens.next().and_then(|arg| arg.parse().ok());
165 continue;
166 } else if let Some(p) = arg.strip_prefix("-p") {
167 port = p.parse().ok();
168 continue;
169 }
170 if arg == "-l" {
171 username = tokens.next();
172 continue;
173 } else if let Some(l) = arg.strip_prefix("-l") {
174 username = Some(l.to_string());
175 continue;
176 }
177 if arg == "-L" || arg.starts_with("-L") {
178 let forward_spec = if arg == "-L" {
179 tokens.next()
180 } else {
181 Some(arg.strip_prefix("-L").unwrap().to_string())
182 };
183
184 if let Some(spec) = forward_spec {
185 port_forwards.push(parse_port_forward_spec(&spec)?);
186 } else {
187 anyhow::bail!("Missing port forward format");
188 }
189 }
190
191 for a in ALLOWED_ARGS {
192 if arg == *a {
193 args.push(arg);
194 if let Some(next) = tokens.next() {
195 args.push(next);
196 }
197 continue 'outer;
198 } else if arg.starts_with(a) {
199 args.push(arg);
200 continue 'outer;
201 }
202 }
203 if arg.starts_with("-") || hostname.is_some() {
204 anyhow::bail!("unsupported argument: {:?}", arg);
205 }
206 let mut input = &arg as &str;
207 // Destination might be: username1@username2@ip2@ip1
208 if let Some((u, rest)) = input.rsplit_once('@') {
209 input = rest;
210 username = Some(u.to_string());
211 }
212 if let Some((rest, p)) = input.split_once(':') {
213 input = rest;
214 port = p.parse().ok()
215 }
216 hostname = Some(input.to_string())
217 }
218
219 let Some(hostname) = hostname else {
220 anyhow::bail!("missing hostname");
221 };
222
223 let port_forwards = match port_forwards.len() {
224 0 => None,
225 _ => Some(port_forwards),
226 };
227
228 Ok(Self {
229 host: hostname.to_string(),
230 username: username.clone(),
231 port,
232 port_forwards,
233 args: Some(args),
234 password: None,
235 nickname: None,
236 upload_binary_over_ssh: false,
237 })
238 }
239
240 pub fn ssh_url(&self) -> String {
241 let mut result = String::from("ssh://");
242 if let Some(username) = &self.username {
243 // Username might be: username1@username2@ip2
244 let username = urlencoding::encode(username);
245 result.push_str(&username);
246 result.push('@');
247 }
248 result.push_str(&self.host);
249 if let Some(port) = self.port {
250 result.push(':');
251 result.push_str(&port.to_string());
252 }
253 result
254 }
255
256 pub fn additional_args(&self) -> Vec<String> {
257 let mut args = self.args.iter().flatten().cloned().collect::<Vec<String>>();
258
259 if let Some(forwards) = &self.port_forwards {
260 args.extend(forwards.iter().map(|pf| {
261 let local_host = match &pf.local_host {
262 Some(host) => host,
263 None => "localhost",
264 };
265 let remote_host = match &pf.remote_host {
266 Some(host) => host,
267 None => "localhost",
268 };
269
270 format!(
271 "-L{}:{}:{}:{}",
272 local_host, pf.local_port, remote_host, pf.remote_port
273 )
274 }));
275 }
276
277 args
278 }
279
280 fn scp_url(&self) -> String {
281 if let Some(username) = &self.username {
282 format!("{}@{}", username, self.host)
283 } else {
284 self.host.clone()
285 }
286 }
287
288 pub fn connection_string(&self) -> String {
289 let host = if let Some(username) = &self.username {
290 format!("{}@{}", username, self.host)
291 } else {
292 self.host.clone()
293 };
294 if let Some(port) = &self.port {
295 format!("{}:{}", host, port)
296 } else {
297 host
298 }
299 }
300}
301
302#[derive(Copy, Clone, Debug)]
303pub struct SshPlatform {
304 pub os: &'static str,
305 pub arch: &'static str,
306}
307
308impl SshPlatform {
309 pub fn triple(&self) -> Option<String> {
310 Some(format!(
311 "{}-{}",
312 self.arch,
313 match self.os {
314 "linux" => "unknown-linux-gnu",
315 "macos" => "apple-darwin",
316 _ => return None,
317 }
318 ))
319 }
320}
321
322pub trait SshClientDelegate: Send + Sync {
323 fn ask_password(&self, prompt: String, tx: oneshot::Sender<String>, cx: &mut AsyncApp);
324 fn get_download_params(
325 &self,
326 platform: SshPlatform,
327 release_channel: ReleaseChannel,
328 version: Option<SemanticVersion>,
329 cx: &mut AsyncApp,
330 ) -> Task<Result<Option<(String, String)>>>;
331
332 fn download_server_binary_locally(
333 &self,
334 platform: SshPlatform,
335 release_channel: ReleaseChannel,
336 version: Option<SemanticVersion>,
337 cx: &mut AsyncApp,
338 ) -> Task<Result<PathBuf>>;
339 fn set_status(&self, status: Option<&str>, cx: &mut AsyncApp);
340}
341
342impl SshSocket {
343 // :WARNING: ssh unquotes arguments when executing on the remote :WARNING:
344 // e.g. $ ssh host sh -c 'ls -l' is equivalent to $ ssh host sh -c ls -l
345 // and passes -l as an argument to sh, not to ls.
346 // Furthermore, some setups (e.g. Coder) will change directory when SSH'ing
347 // into a machine. You must use `cd` to get back to $HOME.
348 // You need to do it like this: $ ssh host "cd; sh -c 'ls -l /tmp'"
349 fn ssh_command(&self, program: &str, args: &[&str]) -> process::Command {
350 let mut command = util::command::new_smol_command("ssh");
351 let to_run = iter::once(&program)
352 .chain(args.iter())
353 .map(|token| {
354 // We're trying to work with: sh, bash, zsh, fish, tcsh, ...?
355 debug_assert!(
356 !token.contains('\n'),
357 "multiline arguments do not work in all shells"
358 );
359 shlex::try_quote(token).unwrap()
360 })
361 .join(" ");
362 let to_run = format!("cd; {to_run}");
363 log::debug!("ssh {} {:?}", self.connection_options.ssh_url(), to_run);
364 self.ssh_options(&mut command)
365 .arg(self.connection_options.ssh_url())
366 .arg(to_run);
367 command
368 }
369
370 async fn run_command(&self, program: &str, args: &[&str]) -> Result<String> {
371 let output = self.ssh_command(program, args).output().await?;
372 if output.status.success() {
373 Ok(String::from_utf8_lossy(&output.stdout).to_string())
374 } else {
375 Err(anyhow!(
376 "failed to run command: {}",
377 String::from_utf8_lossy(&output.stderr)
378 ))
379 }
380 }
381
382 fn ssh_options<'a>(&self, command: &'a mut process::Command) -> &'a mut process::Command {
383 command
384 .stdin(Stdio::piped())
385 .stdout(Stdio::piped())
386 .stderr(Stdio::piped())
387 .args(["-o", "ControlMaster=no", "-o"])
388 .arg(format!("ControlPath={}", self.socket_path.display()))
389 }
390
391 fn ssh_args(&self) -> Vec<String> {
392 vec![
393 "-o".to_string(),
394 "ControlMaster=no".to_string(),
395 "-o".to_string(),
396 format!("ControlPath={}", self.socket_path.display()),
397 self.connection_options.ssh_url(),
398 ]
399 }
400}
401
402const MAX_MISSED_HEARTBEATS: usize = 5;
403const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
404const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5);
405
406const MAX_RECONNECT_ATTEMPTS: usize = 3;
407
408enum State {
409 Connecting,
410 Connected {
411 ssh_connection: Arc<dyn RemoteConnection>,
412 delegate: Arc<dyn SshClientDelegate>,
413
414 multiplex_task: Task<Result<()>>,
415 heartbeat_task: Task<Result<()>>,
416 },
417 HeartbeatMissed {
418 missed_heartbeats: usize,
419
420 ssh_connection: Arc<dyn RemoteConnection>,
421 delegate: Arc<dyn SshClientDelegate>,
422
423 multiplex_task: Task<Result<()>>,
424 heartbeat_task: Task<Result<()>>,
425 },
426 Reconnecting,
427 ReconnectFailed {
428 ssh_connection: Arc<dyn RemoteConnection>,
429 delegate: Arc<dyn SshClientDelegate>,
430
431 error: anyhow::Error,
432 attempts: usize,
433 },
434 ReconnectExhausted,
435 ServerNotRunning,
436}
437
438impl fmt::Display for State {
439 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
440 match self {
441 Self::Connecting => write!(f, "connecting"),
442 Self::Connected { .. } => write!(f, "connected"),
443 Self::Reconnecting => write!(f, "reconnecting"),
444 Self::ReconnectFailed { .. } => write!(f, "reconnect failed"),
445 Self::ReconnectExhausted => write!(f, "reconnect exhausted"),
446 Self::HeartbeatMissed { .. } => write!(f, "heartbeat missed"),
447 Self::ServerNotRunning { .. } => write!(f, "server not running"),
448 }
449 }
450}
451
452impl State {
453 fn ssh_connection(&self) -> Option<&dyn RemoteConnection> {
454 match self {
455 Self::Connected { ssh_connection, .. } => Some(ssh_connection.as_ref()),
456 Self::HeartbeatMissed { ssh_connection, .. } => Some(ssh_connection.as_ref()),
457 Self::ReconnectFailed { ssh_connection, .. } => Some(ssh_connection.as_ref()),
458 _ => None,
459 }
460 }
461
462 fn can_reconnect(&self) -> bool {
463 match self {
464 Self::Connected { .. }
465 | Self::HeartbeatMissed { .. }
466 | Self::ReconnectFailed { .. } => true,
467 State::Connecting
468 | State::Reconnecting
469 | State::ReconnectExhausted
470 | State::ServerNotRunning => false,
471 }
472 }
473
474 fn is_reconnect_failed(&self) -> bool {
475 matches!(self, Self::ReconnectFailed { .. })
476 }
477
478 fn is_reconnect_exhausted(&self) -> bool {
479 matches!(self, Self::ReconnectExhausted { .. })
480 }
481
482 fn is_server_not_running(&self) -> bool {
483 matches!(self, Self::ServerNotRunning)
484 }
485
486 fn is_reconnecting(&self) -> bool {
487 matches!(self, Self::Reconnecting { .. })
488 }
489
490 fn heartbeat_recovered(self) -> Self {
491 match self {
492 Self::HeartbeatMissed {
493 ssh_connection,
494 delegate,
495 multiplex_task,
496 heartbeat_task,
497 ..
498 } => Self::Connected {
499 ssh_connection,
500 delegate,
501 multiplex_task,
502 heartbeat_task,
503 },
504 _ => self,
505 }
506 }
507
508 fn heartbeat_missed(self) -> Self {
509 match self {
510 Self::Connected {
511 ssh_connection,
512 delegate,
513 multiplex_task,
514 heartbeat_task,
515 } => Self::HeartbeatMissed {
516 missed_heartbeats: 1,
517 ssh_connection,
518 delegate,
519 multiplex_task,
520 heartbeat_task,
521 },
522 Self::HeartbeatMissed {
523 missed_heartbeats,
524 ssh_connection,
525 delegate,
526 multiplex_task,
527 heartbeat_task,
528 } => Self::HeartbeatMissed {
529 missed_heartbeats: missed_heartbeats + 1,
530 ssh_connection,
531 delegate,
532 multiplex_task,
533 heartbeat_task,
534 },
535 _ => self,
536 }
537 }
538}
539
540/// The state of the ssh connection.
541#[derive(Clone, Copy, Debug, PartialEq, Eq)]
542pub enum ConnectionState {
543 Connecting,
544 Connected,
545 HeartbeatMissed,
546 Reconnecting,
547 Disconnected,
548}
549
550impl From<&State> for ConnectionState {
551 fn from(value: &State) -> Self {
552 match value {
553 State::Connecting => Self::Connecting,
554 State::Connected { .. } => Self::Connected,
555 State::Reconnecting | State::ReconnectFailed { .. } => Self::Reconnecting,
556 State::HeartbeatMissed { .. } => Self::HeartbeatMissed,
557 State::ReconnectExhausted => Self::Disconnected,
558 State::ServerNotRunning => Self::Disconnected,
559 }
560 }
561}
562
563pub struct SshRemoteClient {
564 client: Arc<ChannelClient>,
565 unique_identifier: String,
566 connection_options: SshConnectionOptions,
567 state: Arc<Mutex<Option<State>>>,
568}
569
570#[derive(Debug)]
571pub enum SshRemoteEvent {
572 Disconnected,
573}
574
575impl EventEmitter<SshRemoteEvent> for SshRemoteClient {}
576
577// Identifies the socket on the remote server so that reconnects
578// can re-join the same project.
579pub enum ConnectionIdentifier {
580 Setup(u64),
581 Workspace(i64),
582}
583
584static NEXT_ID: AtomicU64 = AtomicU64::new(1);
585
586impl ConnectionIdentifier {
587 pub fn setup() -> Self {
588 Self::Setup(NEXT_ID.fetch_add(1, SeqCst))
589 }
590 // This string gets used in a socket name, and so must be relatively short.
591 // The total length of:
592 // /home/{username}/.local/share/zed/server_state/{name}/stdout.sock
593 // Must be less than about 100 characters
594 // https://unix.stackexchange.com/questions/367008/why-is-socket-path-length-limited-to-a-hundred-chars
595 // So our strings should be at most 20 characters or so.
596 fn to_string(&self, cx: &App) -> String {
597 let identifier_prefix = match ReleaseChannel::global(cx) {
598 ReleaseChannel::Stable => "".to_string(),
599 release_channel => format!("{}-", release_channel.dev_name()),
600 };
601 match self {
602 Self::Setup(setup_id) => format!("{identifier_prefix}setup-{setup_id}"),
603 Self::Workspace(workspace_id) => {
604 format!("{identifier_prefix}workspace-{workspace_id}",)
605 }
606 }
607 }
608}
609
610impl SshRemoteClient {
611 pub fn new(
612 unique_identifier: ConnectionIdentifier,
613 connection_options: SshConnectionOptions,
614 cancellation: oneshot::Receiver<()>,
615 delegate: Arc<dyn SshClientDelegate>,
616 cx: &mut App,
617 ) -> Task<Result<Option<Entity<Self>>>> {
618 let unique_identifier = unique_identifier.to_string(cx);
619 cx.spawn(async move |cx| {
620 let success = Box::pin(async move {
621 let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
622 let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
623 let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
624
625 let client =
626 cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "client"))?;
627 let this = cx.new(|_| Self {
628 client: client.clone(),
629 unique_identifier: unique_identifier.clone(),
630 connection_options: connection_options.clone(),
631 state: Arc::new(Mutex::new(Some(State::Connecting))),
632 })?;
633
634 let ssh_connection = cx
635 .update(|cx| {
636 cx.update_default_global(|pool: &mut ConnectionPool, cx| {
637 pool.connect(connection_options, &delegate, cx)
638 })
639 })?
640 .await
641 .map_err(|e| e.cloned())?;
642
643 let io_task = ssh_connection.start_proxy(
644 unique_identifier,
645 false,
646 incoming_tx,
647 outgoing_rx,
648 connection_activity_tx,
649 delegate.clone(),
650 cx,
651 );
652
653 let multiplex_task = Self::monitor(this.downgrade(), io_task, &cx);
654
655 if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
656 log::error!("failed to establish connection: {}", error);
657 return Err(error);
658 }
659
660 let heartbeat_task = Self::heartbeat(this.downgrade(), connection_activity_rx, cx);
661
662 this.update(cx, |this, _| {
663 *this.state.lock() = Some(State::Connected {
664 ssh_connection,
665 delegate,
666 multiplex_task,
667 heartbeat_task,
668 });
669 })?;
670
671 Ok(Some(this))
672 });
673
674 select! {
675 _ = cancellation.fuse() => {
676 Ok(None)
677 }
678 result = success.fuse() => result
679 }
680 })
681 }
682
683 pub fn shutdown_processes<T: RequestMessage>(
684 &self,
685 shutdown_request: Option<T>,
686 executor: BackgroundExecutor,
687 ) -> Option<impl Future<Output = ()> + use<T>> {
688 let state = self.state.lock().take()?;
689 log::info!("shutting down ssh processes");
690
691 let State::Connected {
692 multiplex_task,
693 heartbeat_task,
694 ssh_connection,
695 delegate,
696 } = state
697 else {
698 return None;
699 };
700
701 let client = self.client.clone();
702
703 Some(async move {
704 if let Some(shutdown_request) = shutdown_request {
705 client.send(shutdown_request).log_err();
706 // We wait 50ms instead of waiting for a response, because
707 // waiting for a response would require us to wait on the main thread
708 // which we want to avoid in an `on_app_quit` callback.
709 executor.timer(Duration::from_millis(50)).await;
710 }
711
712 // Drop `multiplex_task` because it owns our ssh_proxy_process, which is a
713 // child of master_process.
714 drop(multiplex_task);
715 // Now drop the rest of state, which kills master process.
716 drop(heartbeat_task);
717 drop(ssh_connection);
718 drop(delegate);
719 })
720 }
721
722 fn reconnect(&mut self, cx: &mut Context<Self>) -> Result<()> {
723 let mut lock = self.state.lock();
724
725 let can_reconnect = lock
726 .as_ref()
727 .map(|state| state.can_reconnect())
728 .unwrap_or(false);
729 if !can_reconnect {
730 let error = if let Some(state) = lock.as_ref() {
731 format!("invalid state, cannot reconnect while in state {state}")
732 } else {
733 "no state set".to_string()
734 };
735 log::info!("aborting reconnect, because not in state that allows reconnecting");
736 return Err(anyhow!(error));
737 }
738
739 let state = lock.take().unwrap();
740 let (attempts, ssh_connection, delegate) = match state {
741 State::Connected {
742 ssh_connection,
743 delegate,
744 multiplex_task,
745 heartbeat_task,
746 }
747 | State::HeartbeatMissed {
748 ssh_connection,
749 delegate,
750 multiplex_task,
751 heartbeat_task,
752 ..
753 } => {
754 drop(multiplex_task);
755 drop(heartbeat_task);
756 (0, ssh_connection, delegate)
757 }
758 State::ReconnectFailed {
759 attempts,
760 ssh_connection,
761 delegate,
762 ..
763 } => (attempts, ssh_connection, delegate),
764 State::Connecting
765 | State::Reconnecting
766 | State::ReconnectExhausted
767 | State::ServerNotRunning => unreachable!(),
768 };
769
770 let attempts = attempts + 1;
771 if attempts > MAX_RECONNECT_ATTEMPTS {
772 log::error!(
773 "Failed to reconnect to after {} attempts, giving up",
774 MAX_RECONNECT_ATTEMPTS
775 );
776 drop(lock);
777 self.set_state(State::ReconnectExhausted, cx);
778 return Ok(());
779 }
780 drop(lock);
781
782 self.set_state(State::Reconnecting, cx);
783
784 log::info!("Trying to reconnect to ssh server... Attempt {}", attempts);
785
786 let unique_identifier = self.unique_identifier.clone();
787 let client = self.client.clone();
788 let reconnect_task = cx.spawn(async move |this, cx| {
789 macro_rules! failed {
790 ($error:expr, $attempts:expr, $ssh_connection:expr, $delegate:expr) => {
791 return State::ReconnectFailed {
792 error: anyhow!($error),
793 attempts: $attempts,
794 ssh_connection: $ssh_connection,
795 delegate: $delegate,
796 };
797 };
798 }
799
800 if let Err(error) = ssh_connection
801 .kill()
802 .await
803 .context("Failed to kill ssh process")
804 {
805 failed!(error, attempts, ssh_connection, delegate);
806 };
807
808 let connection_options = ssh_connection.connection_options();
809
810 let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
811 let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
812 let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
813
814 let (ssh_connection, io_task) = match async {
815 let ssh_connection = cx
816 .update_global(|pool: &mut ConnectionPool, cx| {
817 pool.connect(connection_options, &delegate, cx)
818 })?
819 .await
820 .map_err(|error| error.cloned())?;
821
822 let io_task = ssh_connection.start_proxy(
823 unique_identifier,
824 true,
825 incoming_tx,
826 outgoing_rx,
827 connection_activity_tx,
828 delegate.clone(),
829 cx,
830 );
831 anyhow::Ok((ssh_connection, io_task))
832 }
833 .await
834 {
835 Ok((ssh_connection, io_task)) => (ssh_connection, io_task),
836 Err(error) => {
837 failed!(error, attempts, ssh_connection, delegate);
838 }
839 };
840
841 let multiplex_task = Self::monitor(this.clone(), io_task, &cx);
842 client.reconnect(incoming_rx, outgoing_tx, &cx);
843
844 if let Err(error) = client.resync(HEARTBEAT_TIMEOUT).await {
845 failed!(error, attempts, ssh_connection, delegate);
846 };
847
848 State::Connected {
849 ssh_connection,
850 delegate,
851 multiplex_task,
852 heartbeat_task: Self::heartbeat(this.clone(), connection_activity_rx, cx),
853 }
854 });
855
856 cx.spawn(async move |this, cx| {
857 let new_state = reconnect_task.await;
858 this.update(cx, |this, cx| {
859 this.try_set_state(cx, |old_state| {
860 if old_state.is_reconnecting() {
861 match &new_state {
862 State::Connecting
863 | State::Reconnecting { .. }
864 | State::HeartbeatMissed { .. }
865 | State::ServerNotRunning => {}
866 State::Connected { .. } => {
867 log::info!("Successfully reconnected");
868 }
869 State::ReconnectFailed {
870 error, attempts, ..
871 } => {
872 log::error!(
873 "Reconnect attempt {} failed: {:?}. Starting new attempt...",
874 attempts,
875 error
876 );
877 }
878 State::ReconnectExhausted => {
879 log::error!("Reconnect attempt failed and all attempts exhausted");
880 }
881 }
882 Some(new_state)
883 } else {
884 None
885 }
886 });
887
888 if this.state_is(State::is_reconnect_failed) {
889 this.reconnect(cx)
890 } else if this.state_is(State::is_reconnect_exhausted) {
891 Ok(())
892 } else {
893 log::debug!("State has transition from Reconnecting into new state while attempting reconnect.");
894 Ok(())
895 }
896 })
897 })
898 .detach_and_log_err(cx);
899
900 Ok(())
901 }
902
903 fn heartbeat(
904 this: WeakEntity<Self>,
905 mut connection_activity_rx: mpsc::Receiver<()>,
906 cx: &mut AsyncApp,
907 ) -> Task<Result<()>> {
908 let Ok(client) = this.update(cx, |this, _| this.client.clone()) else {
909 return Task::ready(Err(anyhow!("SshRemoteClient lost")));
910 };
911
912 cx.spawn(async move |cx| {
913 let mut missed_heartbeats = 0;
914
915 let keepalive_timer = cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse();
916 futures::pin_mut!(keepalive_timer);
917
918 loop {
919 select_biased! {
920 result = connection_activity_rx.next().fuse() => {
921 if result.is_none() {
922 log::warn!("ssh heartbeat: connection activity channel has been dropped. stopping.");
923 return Ok(());
924 }
925
926 if missed_heartbeats != 0 {
927 missed_heartbeats = 0;
928 let _ =this.update(cx, |this, mut cx| {
929 this.handle_heartbeat_result(missed_heartbeats, &mut cx)
930 })?;
931 }
932 }
933 _ = keepalive_timer => {
934 log::debug!("Sending heartbeat to server...");
935
936 let result = select_biased! {
937 _ = connection_activity_rx.next().fuse() => {
938 Ok(())
939 }
940 ping_result = client.ping(HEARTBEAT_TIMEOUT).fuse() => {
941 ping_result
942 }
943 };
944
945 if result.is_err() {
946 missed_heartbeats += 1;
947 log::warn!(
948 "No heartbeat from server after {:?}. Missed heartbeat {} out of {}.",
949 HEARTBEAT_TIMEOUT,
950 missed_heartbeats,
951 MAX_MISSED_HEARTBEATS
952 );
953 } else if missed_heartbeats != 0 {
954 missed_heartbeats = 0;
955 } else {
956 continue;
957 }
958
959 let result = this.update(cx, |this, mut cx| {
960 this.handle_heartbeat_result(missed_heartbeats, &mut cx)
961 })?;
962 if result.is_break() {
963 return Ok(());
964 }
965 }
966 }
967
968 keepalive_timer.set(cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse());
969 }
970
971 })
972 }
973
974 fn handle_heartbeat_result(
975 &mut self,
976 missed_heartbeats: usize,
977 cx: &mut Context<Self>,
978 ) -> ControlFlow<()> {
979 let state = self.state.lock().take().unwrap();
980 let next_state = if missed_heartbeats > 0 {
981 state.heartbeat_missed()
982 } else {
983 state.heartbeat_recovered()
984 };
985
986 self.set_state(next_state, cx);
987
988 if missed_heartbeats >= MAX_MISSED_HEARTBEATS {
989 log::error!(
990 "Missed last {} heartbeats. Reconnecting...",
991 missed_heartbeats
992 );
993
994 self.reconnect(cx)
995 .context("failed to start reconnect process after missing heartbeats")
996 .log_err();
997 ControlFlow::Break(())
998 } else {
999 ControlFlow::Continue(())
1000 }
1001 }
1002
1003 fn monitor(
1004 this: WeakEntity<Self>,
1005 io_task: Task<Result<i32>>,
1006 cx: &AsyncApp,
1007 ) -> Task<Result<()>> {
1008 cx.spawn(async move |cx| {
1009 let result = io_task.await;
1010
1011 match result {
1012 Ok(exit_code) => {
1013 if let Some(error) = ProxyLaunchError::from_exit_code(exit_code) {
1014 match error {
1015 ProxyLaunchError::ServerNotRunning => {
1016 log::error!("failed to reconnect because server is not running");
1017 this.update(cx, |this, cx| {
1018 this.set_state(State::ServerNotRunning, cx);
1019 })?;
1020 }
1021 }
1022 } else if exit_code > 0 {
1023 log::error!("proxy process terminated unexpectedly");
1024 this.update(cx, |this, cx| {
1025 this.reconnect(cx).ok();
1026 })?;
1027 }
1028 }
1029 Err(error) => {
1030 log::warn!("ssh io task died with error: {:?}. reconnecting...", error);
1031 this.update(cx, |this, cx| {
1032 this.reconnect(cx).ok();
1033 })?;
1034 }
1035 }
1036
1037 Ok(())
1038 })
1039 }
1040
1041 fn state_is(&self, check: impl FnOnce(&State) -> bool) -> bool {
1042 self.state.lock().as_ref().map_or(false, check)
1043 }
1044
1045 fn try_set_state(&self, cx: &mut Context<Self>, map: impl FnOnce(&State) -> Option<State>) {
1046 let mut lock = self.state.lock();
1047 let new_state = lock.as_ref().and_then(map);
1048
1049 if let Some(new_state) = new_state {
1050 lock.replace(new_state);
1051 cx.notify();
1052 }
1053 }
1054
1055 fn set_state(&self, state: State, cx: &mut Context<Self>) {
1056 log::info!("setting state to '{}'", &state);
1057
1058 let is_reconnect_exhausted = state.is_reconnect_exhausted();
1059 let is_server_not_running = state.is_server_not_running();
1060 self.state.lock().replace(state);
1061
1062 if is_reconnect_exhausted || is_server_not_running {
1063 cx.emit(SshRemoteEvent::Disconnected);
1064 }
1065 cx.notify();
1066 }
1067
1068 pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Entity<E>) {
1069 self.client.subscribe_to_entity(remote_id, entity);
1070 }
1071
1072 pub fn ssh_args(&self) -> Option<Vec<String>> {
1073 self.state
1074 .lock()
1075 .as_ref()
1076 .and_then(|state| state.ssh_connection())
1077 .map(|ssh_connection| ssh_connection.ssh_args())
1078 }
1079
1080 pub fn upload_directory(
1081 &self,
1082 src_path: PathBuf,
1083 dest_path: PathBuf,
1084 cx: &App,
1085 ) -> Task<Result<()>> {
1086 let state = self.state.lock();
1087 let Some(connection) = state.as_ref().and_then(|state| state.ssh_connection()) else {
1088 return Task::ready(Err(anyhow!("no ssh connection")));
1089 };
1090 connection.upload_directory(src_path, dest_path, cx)
1091 }
1092
1093 pub fn proto_client(&self) -> AnyProtoClient {
1094 self.client.clone().into()
1095 }
1096
1097 pub fn connection_string(&self) -> String {
1098 self.connection_options.connection_string()
1099 }
1100
1101 pub fn connection_options(&self) -> SshConnectionOptions {
1102 self.connection_options.clone()
1103 }
1104
1105 pub fn connection_state(&self) -> ConnectionState {
1106 self.state
1107 .lock()
1108 .as_ref()
1109 .map(ConnectionState::from)
1110 .unwrap_or(ConnectionState::Disconnected)
1111 }
1112
1113 pub fn is_disconnected(&self) -> bool {
1114 self.connection_state() == ConnectionState::Disconnected
1115 }
1116
1117 #[cfg(any(test, feature = "test-support"))]
1118 pub fn simulate_disconnect(&self, client_cx: &mut App) -> Task<()> {
1119 let opts = self.connection_options();
1120 client_cx.spawn(async move |cx| {
1121 let connection = cx
1122 .update_global(|c: &mut ConnectionPool, _| {
1123 if let Some(ConnectionPoolEntry::Connecting(c)) = c.connections.get(&opts) {
1124 c.clone()
1125 } else {
1126 panic!("missing test connection")
1127 }
1128 })
1129 .unwrap()
1130 .await
1131 .unwrap();
1132
1133 connection.simulate_disconnect(&cx);
1134 })
1135 }
1136
1137 #[cfg(any(test, feature = "test-support"))]
1138 pub fn fake_server(
1139 client_cx: &mut gpui::TestAppContext,
1140 server_cx: &mut gpui::TestAppContext,
1141 ) -> (SshConnectionOptions, Arc<ChannelClient>) {
1142 let port = client_cx
1143 .update(|cx| cx.default_global::<ConnectionPool>().connections.len() as u16 + 1);
1144 let opts = SshConnectionOptions {
1145 host: "<fake>".to_string(),
1146 port: Some(port),
1147 ..Default::default()
1148 };
1149 let (outgoing_tx, _) = mpsc::unbounded::<Envelope>();
1150 let (_, incoming_rx) = mpsc::unbounded::<Envelope>();
1151 let server_client =
1152 server_cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "fake-server"));
1153 let connection: Arc<dyn RemoteConnection> = Arc::new(fake::FakeRemoteConnection {
1154 connection_options: opts.clone(),
1155 server_cx: fake::SendableCx::new(server_cx),
1156 server_channel: server_client.clone(),
1157 });
1158
1159 client_cx.update(|cx| {
1160 cx.update_default_global(|c: &mut ConnectionPool, cx| {
1161 c.connections.insert(
1162 opts.clone(),
1163 ConnectionPoolEntry::Connecting(
1164 cx.background_spawn({
1165 let connection = connection.clone();
1166 async move { Ok(connection.clone()) }
1167 })
1168 .shared(),
1169 ),
1170 );
1171 })
1172 });
1173
1174 (opts, server_client)
1175 }
1176
1177 #[cfg(any(test, feature = "test-support"))]
1178 pub async fn fake_client(
1179 opts: SshConnectionOptions,
1180 client_cx: &mut gpui::TestAppContext,
1181 ) -> Entity<Self> {
1182 let (_tx, rx) = oneshot::channel();
1183 client_cx
1184 .update(|cx| {
1185 Self::new(
1186 ConnectionIdentifier::setup(),
1187 opts,
1188 rx,
1189 Arc::new(fake::Delegate),
1190 cx,
1191 )
1192 })
1193 .await
1194 .unwrap()
1195 .unwrap()
1196 }
1197}
1198
1199enum ConnectionPoolEntry {
1200 Connecting(Shared<Task<Result<Arc<dyn RemoteConnection>, Arc<anyhow::Error>>>>),
1201 Connected(Weak<dyn RemoteConnection>),
1202}
1203
1204#[derive(Default)]
1205struct ConnectionPool {
1206 connections: HashMap<SshConnectionOptions, ConnectionPoolEntry>,
1207}
1208
1209impl Global for ConnectionPool {}
1210
1211impl ConnectionPool {
1212 pub fn connect(
1213 &mut self,
1214 opts: SshConnectionOptions,
1215 delegate: &Arc<dyn SshClientDelegate>,
1216 cx: &mut App,
1217 ) -> Shared<Task<Result<Arc<dyn RemoteConnection>, Arc<anyhow::Error>>>> {
1218 let connection = self.connections.get(&opts);
1219 match connection {
1220 Some(ConnectionPoolEntry::Connecting(task)) => {
1221 let delegate = delegate.clone();
1222 cx.spawn(async move |cx| {
1223 delegate.set_status(Some("Waiting for existing connection attempt"), cx);
1224 })
1225 .detach();
1226 return task.clone();
1227 }
1228 Some(ConnectionPoolEntry::Connected(ssh)) => {
1229 if let Some(ssh) = ssh.upgrade() {
1230 if !ssh.has_been_killed() {
1231 return Task::ready(Ok(ssh)).shared();
1232 }
1233 }
1234 self.connections.remove(&opts);
1235 }
1236 None => {}
1237 }
1238
1239 let task = cx
1240 .spawn({
1241 let opts = opts.clone();
1242 let delegate = delegate.clone();
1243 async move |cx| {
1244 let connection = SshRemoteConnection::new(opts.clone(), delegate, cx)
1245 .await
1246 .map(|connection| Arc::new(connection) as Arc<dyn RemoteConnection>);
1247
1248 cx.update_global(|pool: &mut Self, _| {
1249 debug_assert!(matches!(
1250 pool.connections.get(&opts),
1251 Some(ConnectionPoolEntry::Connecting(_))
1252 ));
1253 match connection {
1254 Ok(connection) => {
1255 pool.connections.insert(
1256 opts.clone(),
1257 ConnectionPoolEntry::Connected(Arc::downgrade(&connection)),
1258 );
1259 Ok(connection)
1260 }
1261 Err(error) => {
1262 pool.connections.remove(&opts);
1263 Err(Arc::new(error))
1264 }
1265 }
1266 })?
1267 }
1268 })
1269 .shared();
1270
1271 self.connections
1272 .insert(opts.clone(), ConnectionPoolEntry::Connecting(task.clone()));
1273 task
1274 }
1275}
1276
1277impl From<SshRemoteClient> for AnyProtoClient {
1278 fn from(client: SshRemoteClient) -> Self {
1279 AnyProtoClient::new(client.client.clone())
1280 }
1281}
1282
1283#[async_trait(?Send)]
1284trait RemoteConnection: Send + Sync {
1285 fn start_proxy(
1286 &self,
1287 unique_identifier: String,
1288 reconnect: bool,
1289 incoming_tx: UnboundedSender<Envelope>,
1290 outgoing_rx: UnboundedReceiver<Envelope>,
1291 connection_activity_tx: Sender<()>,
1292 delegate: Arc<dyn SshClientDelegate>,
1293 cx: &mut AsyncApp,
1294 ) -> Task<Result<i32>>;
1295 fn upload_directory(&self, src_path: PathBuf, dest_path: PathBuf, cx: &App)
1296 -> Task<Result<()>>;
1297 async fn kill(&self) -> Result<()>;
1298 fn has_been_killed(&self) -> bool;
1299 fn ssh_args(&self) -> Vec<String>;
1300 fn connection_options(&self) -> SshConnectionOptions;
1301
1302 #[cfg(any(test, feature = "test-support"))]
1303 fn simulate_disconnect(&self, _: &AsyncApp) {}
1304}
1305
1306struct SshRemoteConnection {
1307 socket: SshSocket,
1308 master_process: Mutex<Option<Child>>,
1309 remote_binary_path: Option<PathBuf>,
1310 _temp_dir: TempDir,
1311}
1312
1313#[async_trait(?Send)]
1314impl RemoteConnection for SshRemoteConnection {
1315 async fn kill(&self) -> Result<()> {
1316 let Some(mut process) = self.master_process.lock().take() else {
1317 return Ok(());
1318 };
1319 process.kill().ok();
1320 process.status().await?;
1321 Ok(())
1322 }
1323
1324 fn has_been_killed(&self) -> bool {
1325 self.master_process.lock().is_none()
1326 }
1327
1328 fn ssh_args(&self) -> Vec<String> {
1329 self.socket.ssh_args()
1330 }
1331
1332 fn connection_options(&self) -> SshConnectionOptions {
1333 self.socket.connection_options.clone()
1334 }
1335
1336 fn upload_directory(
1337 &self,
1338 src_path: PathBuf,
1339 dest_path: PathBuf,
1340 cx: &App,
1341 ) -> Task<Result<()>> {
1342 let mut command = util::command::new_smol_command("scp");
1343 let output = self
1344 .socket
1345 .ssh_options(&mut command)
1346 .args(
1347 self.socket
1348 .connection_options
1349 .port
1350 .map(|port| vec!["-P".to_string(), port.to_string()])
1351 .unwrap_or_default(),
1352 )
1353 .arg("-C")
1354 .arg("-r")
1355 .arg(&src_path)
1356 .arg(format!(
1357 "{}:{}",
1358 self.socket.connection_options.scp_url(),
1359 dest_path.display()
1360 ))
1361 .output();
1362
1363 cx.background_spawn(async move {
1364 let output = output.await?;
1365
1366 if !output.status.success() {
1367 return Err(anyhow!(
1368 "failed to upload directory {} -> {}: {}",
1369 src_path.display(),
1370 dest_path.display(),
1371 String::from_utf8_lossy(&output.stderr)
1372 ));
1373 }
1374
1375 Ok(())
1376 })
1377 }
1378
1379 fn start_proxy(
1380 &self,
1381 unique_identifier: String,
1382 reconnect: bool,
1383 incoming_tx: UnboundedSender<Envelope>,
1384 outgoing_rx: UnboundedReceiver<Envelope>,
1385 connection_activity_tx: Sender<()>,
1386 delegate: Arc<dyn SshClientDelegate>,
1387 cx: &mut AsyncApp,
1388 ) -> Task<Result<i32>> {
1389 delegate.set_status(Some("Starting proxy"), cx);
1390
1391 let Some(remote_binary_path) = self.remote_binary_path.clone() else {
1392 return Task::ready(Err(anyhow!("Remote binary path not set")));
1393 };
1394
1395 let mut start_proxy_command = shell_script!(
1396 "exec {binary_path} proxy --identifier {identifier}",
1397 binary_path = &remote_binary_path.to_string_lossy(),
1398 identifier = &unique_identifier,
1399 );
1400
1401 if let Some(rust_log) = std::env::var("RUST_LOG").ok() {
1402 start_proxy_command = format!(
1403 "RUST_LOG={} {}",
1404 shlex::try_quote(&rust_log).unwrap(),
1405 start_proxy_command
1406 )
1407 }
1408 if let Some(rust_backtrace) = std::env::var("RUST_BACKTRACE").ok() {
1409 start_proxy_command = format!(
1410 "RUST_BACKTRACE={} {}",
1411 shlex::try_quote(&rust_backtrace).unwrap(),
1412 start_proxy_command
1413 )
1414 }
1415 if reconnect {
1416 start_proxy_command.push_str(" --reconnect");
1417 }
1418
1419 let ssh_proxy_process = match self
1420 .socket
1421 .ssh_command("sh", &["-c", &start_proxy_command])
1422 // IMPORTANT: we kill this process when we drop the task that uses it.
1423 .kill_on_drop(true)
1424 .spawn()
1425 {
1426 Ok(process) => process,
1427 Err(error) => {
1428 return Task::ready(Err(anyhow!("failed to spawn remote server: {}", error)));
1429 }
1430 };
1431
1432 Self::multiplex(
1433 ssh_proxy_process,
1434 incoming_tx,
1435 outgoing_rx,
1436 connection_activity_tx,
1437 &cx,
1438 )
1439 }
1440}
1441
1442impl SshRemoteConnection {
1443 #[cfg(not(unix))]
1444 async fn new(
1445 _connection_options: SshConnectionOptions,
1446 _delegate: Arc<dyn SshClientDelegate>,
1447 _cx: &mut AsyncApp,
1448 ) -> Result<Self> {
1449 Err(anyhow!("ssh is not supported on this platform"))
1450 }
1451
1452 #[cfg(unix)]
1453 async fn new(
1454 connection_options: SshConnectionOptions,
1455 delegate: Arc<dyn SshClientDelegate>,
1456 cx: &mut AsyncApp,
1457 ) -> Result<Self> {
1458 use askpass::AskPassResult;
1459
1460 delegate.set_status(Some("Connecting"), cx);
1461
1462 let url = connection_options.ssh_url();
1463
1464 let temp_dir = tempfile::Builder::new()
1465 .prefix("zed-ssh-session")
1466 .tempdir()?;
1467 let askpass_delegate = askpass::AskPassDelegate::new(cx, {
1468 let delegate = delegate.clone();
1469 move |prompt, tx, cx| delegate.ask_password(prompt, tx, cx)
1470 });
1471
1472 let mut askpass =
1473 askpass::AskPassSession::new(cx.background_executor(), askpass_delegate).await?;
1474
1475 // Start the master SSH process, which does not do anything except for establish
1476 // the connection and keep it open, allowing other ssh commands to reuse it
1477 // via a control socket.
1478 let socket_path = temp_dir.path().join("ssh.sock");
1479
1480 let mut master_process = process::Command::new("ssh")
1481 .stdin(Stdio::null())
1482 .stdout(Stdio::piped())
1483 .stderr(Stdio::piped())
1484 .env("SSH_ASKPASS_REQUIRE", "force")
1485 .env("SSH_ASKPASS", &askpass.script_path())
1486 .args(connection_options.additional_args())
1487 .args([
1488 "-N",
1489 "-o",
1490 "ControlPersist=no",
1491 "-o",
1492 "ControlMaster=yes",
1493 "-o",
1494 ])
1495 .arg(format!("ControlPath={}", socket_path.display()))
1496 .arg(&url)
1497 .kill_on_drop(true)
1498 .spawn()?;
1499 // Wait for this ssh process to close its stdout, indicating that authentication
1500 // has completed.
1501 let mut stdout = master_process.stdout.take().unwrap();
1502 let mut output = Vec::new();
1503
1504 let result = select_biased! {
1505 result = askpass.run().fuse() => {
1506 match result {
1507 AskPassResult::CancelledByUser => {
1508 master_process.kill().ok();
1509 Err(anyhow!("SSH connection canceled"))?
1510 }
1511 AskPassResult::Timedout => {
1512 Err(anyhow!("connecting to host timed out"))?
1513 }
1514 }
1515 }
1516 _ = stdout.read_to_end(&mut output).fuse() => {
1517 anyhow::Ok(())
1518 }
1519 };
1520
1521 if let Err(e) = result {
1522 return Err(e.context("Failed to connect to host"));
1523 }
1524
1525 if master_process.try_status()?.is_some() {
1526 output.clear();
1527 let mut stderr = master_process.stderr.take().unwrap();
1528 stderr.read_to_end(&mut output).await?;
1529
1530 let error_message = format!(
1531 "failed to connect: {}",
1532 String::from_utf8_lossy(&output).trim()
1533 );
1534 Err(anyhow!(error_message))?;
1535 }
1536
1537 drop(askpass);
1538
1539 let socket = SshSocket {
1540 connection_options,
1541 socket_path,
1542 };
1543
1544 let mut this = Self {
1545 socket,
1546 master_process: Mutex::new(Some(master_process)),
1547 _temp_dir: temp_dir,
1548 remote_binary_path: None,
1549 };
1550
1551 let (release_channel, version, commit) = cx.update(|cx| {
1552 (
1553 ReleaseChannel::global(cx),
1554 AppVersion::global(cx),
1555 AppCommitSha::try_global(cx),
1556 )
1557 })?;
1558 this.remote_binary_path = Some(
1559 this.ensure_server_binary(&delegate, release_channel, version, commit, cx)
1560 .await?,
1561 );
1562
1563 Ok(this)
1564 }
1565
1566 async fn platform(&self) -> Result<SshPlatform> {
1567 let uname = self.socket.run_command("sh", &["-c", "uname -sm"]).await?;
1568 let Some((os, arch)) = uname.split_once(" ") else {
1569 Err(anyhow!("unknown uname: {uname:?}"))?
1570 };
1571
1572 let os = match os.trim() {
1573 "Darwin" => "macos",
1574 "Linux" => "linux",
1575 _ => Err(anyhow!(
1576 "Prebuilt remote servers are not yet available for {os:?}. See https://zed.dev/docs/remote-development"
1577 ))?,
1578 };
1579 // exclude armv5,6,7 as they are 32-bit.
1580 let arch = if arch.starts_with("armv8")
1581 || arch.starts_with("armv9")
1582 || arch.starts_with("arm64")
1583 || arch.starts_with("aarch64")
1584 {
1585 "aarch64"
1586 } else if arch.starts_with("x86") {
1587 "x86_64"
1588 } else {
1589 Err(anyhow!(
1590 "Prebuilt remote servers are not yet available for {arch:?}. See https://zed.dev/docs/remote-development"
1591 ))?
1592 };
1593
1594 Ok(SshPlatform { os, arch })
1595 }
1596
1597 fn multiplex(
1598 mut ssh_proxy_process: Child,
1599 incoming_tx: UnboundedSender<Envelope>,
1600 mut outgoing_rx: UnboundedReceiver<Envelope>,
1601 mut connection_activity_tx: Sender<()>,
1602 cx: &AsyncApp,
1603 ) -> Task<Result<i32>> {
1604 let mut child_stderr = ssh_proxy_process.stderr.take().unwrap();
1605 let mut child_stdout = ssh_proxy_process.stdout.take().unwrap();
1606 let mut child_stdin = ssh_proxy_process.stdin.take().unwrap();
1607
1608 let mut stdin_buffer = Vec::new();
1609 let mut stdout_buffer = Vec::new();
1610 let mut stderr_buffer = Vec::new();
1611 let mut stderr_offset = 0;
1612
1613 let stdin_task = cx.background_spawn(async move {
1614 while let Some(outgoing) = outgoing_rx.next().await {
1615 write_message(&mut child_stdin, &mut stdin_buffer, outgoing).await?;
1616 }
1617 anyhow::Ok(())
1618 });
1619
1620 let stdout_task = cx.background_spawn({
1621 let mut connection_activity_tx = connection_activity_tx.clone();
1622 async move {
1623 loop {
1624 stdout_buffer.resize(MESSAGE_LEN_SIZE, 0);
1625 let len = child_stdout.read(&mut stdout_buffer).await?;
1626
1627 if len == 0 {
1628 return anyhow::Ok(());
1629 }
1630
1631 if len < MESSAGE_LEN_SIZE {
1632 child_stdout.read_exact(&mut stdout_buffer[len..]).await?;
1633 }
1634
1635 let message_len = message_len_from_buffer(&stdout_buffer);
1636 let envelope =
1637 read_message_with_len(&mut child_stdout, &mut stdout_buffer, message_len)
1638 .await?;
1639 connection_activity_tx.try_send(()).ok();
1640 incoming_tx.unbounded_send(envelope).ok();
1641 }
1642 }
1643 });
1644
1645 let stderr_task: Task<anyhow::Result<()>> = cx.background_spawn(async move {
1646 loop {
1647 stderr_buffer.resize(stderr_offset + 1024, 0);
1648
1649 let len = child_stderr
1650 .read(&mut stderr_buffer[stderr_offset..])
1651 .await?;
1652 if len == 0 {
1653 return anyhow::Ok(());
1654 }
1655
1656 stderr_offset += len;
1657 let mut start_ix = 0;
1658 while let Some(ix) = stderr_buffer[start_ix..stderr_offset]
1659 .iter()
1660 .position(|b| b == &b'\n')
1661 {
1662 let line_ix = start_ix + ix;
1663 let content = &stderr_buffer[start_ix..line_ix];
1664 start_ix = line_ix + 1;
1665 if let Ok(record) = serde_json::from_slice::<LogRecord>(content) {
1666 record.log(log::logger())
1667 } else {
1668 eprintln!("(remote) {}", String::from_utf8_lossy(content));
1669 }
1670 }
1671 stderr_buffer.drain(0..start_ix);
1672 stderr_offset -= start_ix;
1673
1674 connection_activity_tx.try_send(()).ok();
1675 }
1676 });
1677
1678 cx.spawn(async move |_| {
1679 let result = futures::select! {
1680 result = stdin_task.fuse() => {
1681 result.context("stdin")
1682 }
1683 result = stdout_task.fuse() => {
1684 result.context("stdout")
1685 }
1686 result = stderr_task.fuse() => {
1687 result.context("stderr")
1688 }
1689 };
1690
1691 let status = ssh_proxy_process.status().await?.code().unwrap_or(1);
1692 match result {
1693 Ok(_) => Ok(status),
1694 Err(error) => Err(error),
1695 }
1696 })
1697 }
1698
1699 #[allow(unused)]
1700 async fn ensure_server_binary(
1701 &self,
1702 delegate: &Arc<dyn SshClientDelegate>,
1703 release_channel: ReleaseChannel,
1704 version: SemanticVersion,
1705 commit: Option<AppCommitSha>,
1706 cx: &mut AsyncApp,
1707 ) -> Result<PathBuf> {
1708 let version_str = match release_channel {
1709 ReleaseChannel::Nightly => {
1710 let commit = commit.map(|s| s.0.to_string()).unwrap_or_default();
1711
1712 format!("{}-{}", version, commit)
1713 }
1714 ReleaseChannel::Dev => "build".to_string(),
1715 _ => version.to_string(),
1716 };
1717 let binary_name = format!(
1718 "zed-remote-server-{}-{}",
1719 release_channel.dev_name(),
1720 version_str
1721 );
1722 let dst_path = paths::remote_server_dir_relative().join(binary_name);
1723 let tmp_path_gz = PathBuf::from(format!(
1724 "{}-download-{}.gz",
1725 dst_path.to_string_lossy(),
1726 std::process::id()
1727 ));
1728
1729 #[cfg(debug_assertions)]
1730 if std::env::var("ZED_BUILD_REMOTE_SERVER").is_ok() {
1731 let src_path = self
1732 .build_local(self.platform().await?, delegate, cx)
1733 .await?;
1734 self.upload_local_server_binary(&src_path, &tmp_path_gz, delegate, cx)
1735 .await?;
1736 self.extract_server_binary(&dst_path, &tmp_path_gz, delegate, cx)
1737 .await?;
1738 return Ok(dst_path);
1739 }
1740
1741 if self
1742 .socket
1743 .run_command(&dst_path.to_string_lossy(), &["version"])
1744 .await
1745 .is_ok()
1746 {
1747 return Ok(dst_path);
1748 }
1749
1750 let wanted_version = cx.update(|cx| match release_channel {
1751 ReleaseChannel::Nightly => Ok(None),
1752 ReleaseChannel::Dev => {
1753 anyhow::bail!(
1754 "ZED_BUILD_REMOTE_SERVER is not set and no remote server exists at ({:?})",
1755 dst_path
1756 )
1757 }
1758 _ => Ok(Some(AppVersion::global(cx))),
1759 })??;
1760
1761 let platform = self.platform().await?;
1762
1763 if !self.socket.connection_options.upload_binary_over_ssh {
1764 if let Some((url, body)) = delegate
1765 .get_download_params(platform, release_channel, wanted_version, cx)
1766 .await?
1767 {
1768 match self
1769 .download_binary_on_server(&url, &body, &tmp_path_gz, delegate, cx)
1770 .await
1771 {
1772 Ok(_) => {
1773 self.extract_server_binary(&dst_path, &tmp_path_gz, delegate, cx)
1774 .await?;
1775 return Ok(dst_path);
1776 }
1777 Err(e) => {
1778 log::error!(
1779 "Failed to download binary on server, attempting to upload server: {}",
1780 e
1781 )
1782 }
1783 }
1784 }
1785 }
1786
1787 let src_path = delegate
1788 .download_server_binary_locally(platform, release_channel, wanted_version, cx)
1789 .await?;
1790 self.upload_local_server_binary(&src_path, &tmp_path_gz, delegate, cx)
1791 .await?;
1792 self.extract_server_binary(&dst_path, &tmp_path_gz, delegate, cx)
1793 .await?;
1794 return Ok(dst_path);
1795 }
1796
1797 async fn download_binary_on_server(
1798 &self,
1799 url: &str,
1800 body: &str,
1801 tmp_path_gz: &Path,
1802 delegate: &Arc<dyn SshClientDelegate>,
1803 cx: &mut AsyncApp,
1804 ) -> Result<()> {
1805 if let Some(parent) = tmp_path_gz.parent() {
1806 self.socket
1807 .run_command("mkdir", &["-p", &parent.to_string_lossy()])
1808 .await?;
1809 }
1810
1811 delegate.set_status(Some("Downloading remote development server on host"), cx);
1812
1813 match self
1814 .socket
1815 .run_command(
1816 "curl",
1817 &[
1818 "-f",
1819 "-L",
1820 "-X",
1821 "GET",
1822 "-H",
1823 "Content-Type: application/json",
1824 "-d",
1825 &body,
1826 &url,
1827 "-o",
1828 &tmp_path_gz.to_string_lossy(),
1829 ],
1830 )
1831 .await
1832 {
1833 Ok(_) => {}
1834 Err(e) => {
1835 if self.socket.run_command("which", &["curl"]).await.is_ok() {
1836 return Err(e);
1837 }
1838
1839 match self
1840 .socket
1841 .run_command(
1842 "wget",
1843 &[
1844 "--method=GET",
1845 "--header=Content-Type: application/json",
1846 "--body-data",
1847 &body,
1848 &url,
1849 "-O",
1850 &tmp_path_gz.to_string_lossy(),
1851 ],
1852 )
1853 .await
1854 {
1855 Ok(_) => {}
1856 Err(e) => {
1857 if self.socket.run_command("which", &["wget"]).await.is_ok() {
1858 return Err(e);
1859 } else {
1860 anyhow::bail!("Neither curl nor wget is available");
1861 }
1862 }
1863 }
1864 }
1865 }
1866
1867 Ok(())
1868 }
1869
1870 async fn upload_local_server_binary(
1871 &self,
1872 src_path: &Path,
1873 tmp_path_gz: &Path,
1874 delegate: &Arc<dyn SshClientDelegate>,
1875 cx: &mut AsyncApp,
1876 ) -> Result<()> {
1877 if let Some(parent) = tmp_path_gz.parent() {
1878 self.socket
1879 .run_command("mkdir", &["-p", &parent.to_string_lossy()])
1880 .await?;
1881 }
1882
1883 let src_stat = fs::metadata(&src_path).await?;
1884 let size = src_stat.len();
1885
1886 let t0 = Instant::now();
1887 delegate.set_status(Some("Uploading remote development server"), cx);
1888 log::info!(
1889 "uploading remote development server to {:?} ({}kb)",
1890 tmp_path_gz,
1891 size / 1024
1892 );
1893 self.upload_file(&src_path, &tmp_path_gz)
1894 .await
1895 .context("failed to upload server binary")?;
1896 log::info!("uploaded remote development server in {:?}", t0.elapsed());
1897 Ok(())
1898 }
1899
1900 async fn extract_server_binary(
1901 &self,
1902 dst_path: &Path,
1903 tmp_path_gz: &Path,
1904 delegate: &Arc<dyn SshClientDelegate>,
1905 cx: &mut AsyncApp,
1906 ) -> Result<()> {
1907 delegate.set_status(Some("Extracting remote development server"), cx);
1908 let server_mode = 0o755;
1909
1910 let script = shell_script!(
1911 "gunzip -f {tmp_path_gz} && chmod {server_mode} {tmp_path} && mv {tmp_path} {dst_path}",
1912 tmp_path_gz = &tmp_path_gz.to_string_lossy(),
1913 tmp_path = &tmp_path_gz.to_string_lossy().strip_suffix(".gz").unwrap(),
1914 server_mode = &format!("{:o}", server_mode),
1915 dst_path = &dst_path.to_string_lossy()
1916 );
1917 self.socket.run_command("sh", &["-c", &script]).await?;
1918 Ok(())
1919 }
1920
1921 async fn upload_file(&self, src_path: &Path, dest_path: &Path) -> Result<()> {
1922 log::debug!("uploading file {:?} to {:?}", src_path, dest_path);
1923 let mut command = util::command::new_smol_command("scp");
1924 let output = self
1925 .socket
1926 .ssh_options(&mut command)
1927 .args(
1928 self.socket
1929 .connection_options
1930 .port
1931 .map(|port| vec!["-P".to_string(), port.to_string()])
1932 .unwrap_or_default(),
1933 )
1934 .arg(src_path)
1935 .arg(format!(
1936 "{}:{}",
1937 self.socket.connection_options.scp_url(),
1938 dest_path.display()
1939 ))
1940 .output()
1941 .await?;
1942
1943 if output.status.success() {
1944 Ok(())
1945 } else {
1946 Err(anyhow!(
1947 "failed to upload file {} -> {}: {}",
1948 src_path.display(),
1949 dest_path.display(),
1950 String::from_utf8_lossy(&output.stderr)
1951 ))
1952 }
1953 }
1954
1955 #[cfg(debug_assertions)]
1956 async fn build_local(
1957 &self,
1958 platform: SshPlatform,
1959 delegate: &Arc<dyn SshClientDelegate>,
1960 cx: &mut AsyncApp,
1961 ) -> Result<PathBuf> {
1962 use smol::process::{Command, Stdio};
1963
1964 async fn run_cmd(command: &mut Command) -> Result<()> {
1965 let output = command
1966 .kill_on_drop(true)
1967 .stderr(Stdio::inherit())
1968 .output()
1969 .await?;
1970 if !output.status.success() {
1971 Err(anyhow!("Failed to run command: {:?}", command))?;
1972 }
1973 Ok(())
1974 }
1975
1976 if platform.arch == std::env::consts::ARCH && platform.os == std::env::consts::OS {
1977 delegate.set_status(Some("Building remote server binary from source"), cx);
1978 log::info!("building remote server binary from source");
1979 run_cmd(Command::new("cargo").args([
1980 "build",
1981 "--package",
1982 "remote_server",
1983 "--features",
1984 "debug-embed",
1985 "--target-dir",
1986 "target/remote_server",
1987 ]))
1988 .await?;
1989
1990 delegate.set_status(Some("Compressing binary"), cx);
1991
1992 run_cmd(Command::new("gzip").args([
1993 "-9",
1994 "-f",
1995 "target/remote_server/debug/remote_server",
1996 ]))
1997 .await?;
1998
1999 let path = std::env::current_dir()?.join("target/remote_server/debug/remote_server.gz");
2000 return Ok(path);
2001 }
2002 let Some(triple) = platform.triple() else {
2003 anyhow::bail!("can't cross compile for: {:?}", platform);
2004 };
2005 smol::fs::create_dir_all("target/remote_server").await?;
2006
2007 delegate.set_status(Some("Installing cross.rs for cross-compilation"), cx);
2008 log::info!("installing cross");
2009 run_cmd(Command::new("cargo").args([
2010 "install",
2011 "cross",
2012 "--git",
2013 "https://github.com/cross-rs/cross",
2014 ]))
2015 .await?;
2016
2017 delegate.set_status(
2018 Some(&format!(
2019 "Building remote server binary from source for {} with Docker",
2020 &triple
2021 )),
2022 cx,
2023 );
2024 log::info!("building remote server binary from source for {}", &triple);
2025 run_cmd(
2026 Command::new("cross")
2027 .args([
2028 "build",
2029 "--package",
2030 "remote_server",
2031 "--features",
2032 "debug-embed",
2033 "--target-dir",
2034 "target/remote_server",
2035 "--target",
2036 &triple,
2037 ])
2038 .env(
2039 "CROSS_CONTAINER_OPTS",
2040 "--mount type=bind,src=./target,dst=/app/target",
2041 ),
2042 )
2043 .await?;
2044
2045 delegate.set_status(Some("Compressing binary"), cx);
2046
2047 run_cmd(Command::new("gzip").args([
2048 "-9",
2049 "-f",
2050 &format!("target/remote_server/{}/debug/remote_server", triple),
2051 ]))
2052 .await?;
2053
2054 let path = std::env::current_dir()?.join(format!(
2055 "target/remote_server/{}/debug/remote_server.gz",
2056 triple
2057 ));
2058
2059 return Ok(path);
2060 }
2061}
2062
2063type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>;
2064
2065pub struct ChannelClient {
2066 next_message_id: AtomicU32,
2067 outgoing_tx: Mutex<mpsc::UnboundedSender<Envelope>>,
2068 buffer: Mutex<VecDeque<Envelope>>,
2069 response_channels: ResponseChannels,
2070 message_handlers: Mutex<ProtoMessageHandlerSet>,
2071 max_received: AtomicU32,
2072 name: &'static str,
2073 task: Mutex<Task<Result<()>>>,
2074}
2075
2076impl ChannelClient {
2077 pub fn new(
2078 incoming_rx: mpsc::UnboundedReceiver<Envelope>,
2079 outgoing_tx: mpsc::UnboundedSender<Envelope>,
2080 cx: &App,
2081 name: &'static str,
2082 ) -> Arc<Self> {
2083 Arc::new_cyclic(|this| Self {
2084 outgoing_tx: Mutex::new(outgoing_tx),
2085 next_message_id: AtomicU32::new(0),
2086 max_received: AtomicU32::new(0),
2087 response_channels: ResponseChannels::default(),
2088 message_handlers: Default::default(),
2089 buffer: Mutex::new(VecDeque::new()),
2090 name,
2091 task: Mutex::new(Self::start_handling_messages(
2092 this.clone(),
2093 incoming_rx,
2094 &cx.to_async(),
2095 )),
2096 })
2097 }
2098
2099 fn start_handling_messages(
2100 this: Weak<Self>,
2101 mut incoming_rx: mpsc::UnboundedReceiver<Envelope>,
2102 cx: &AsyncApp,
2103 ) -> Task<Result<()>> {
2104 cx.spawn(async move |cx| {
2105 let peer_id = PeerId { owner_id: 0, id: 0 };
2106 while let Some(incoming) = incoming_rx.next().await {
2107 let Some(this) = this.upgrade() else {
2108 return anyhow::Ok(());
2109 };
2110 if let Some(ack_id) = incoming.ack_id {
2111 let mut buffer = this.buffer.lock();
2112 while buffer.front().is_some_and(|msg| msg.id <= ack_id) {
2113 buffer.pop_front();
2114 }
2115 }
2116 if let Some(proto::envelope::Payload::FlushBufferedMessages(_)) = &incoming.payload
2117 {
2118 log::debug!(
2119 "{}:ssh message received. name:FlushBufferedMessages",
2120 this.name
2121 );
2122 {
2123 let buffer = this.buffer.lock();
2124 for envelope in buffer.iter() {
2125 this.outgoing_tx
2126 .lock()
2127 .unbounded_send(envelope.clone())
2128 .ok();
2129 }
2130 }
2131 let mut envelope = proto::Ack {}.into_envelope(0, Some(incoming.id), None);
2132 envelope.id = this.next_message_id.fetch_add(1, SeqCst);
2133 this.outgoing_tx.lock().unbounded_send(envelope).ok();
2134 continue;
2135 }
2136
2137 this.max_received.store(incoming.id, SeqCst);
2138
2139 if let Some(request_id) = incoming.responding_to {
2140 let request_id = MessageId(request_id);
2141 let sender = this.response_channels.lock().remove(&request_id);
2142 if let Some(sender) = sender {
2143 let (tx, rx) = oneshot::channel();
2144 if incoming.payload.is_some() {
2145 sender.send((incoming, tx)).ok();
2146 }
2147 rx.await.ok();
2148 }
2149 } else if let Some(envelope) =
2150 build_typed_envelope(peer_id, Instant::now(), incoming)
2151 {
2152 let type_name = envelope.payload_type_name();
2153 if let Some(future) = ProtoMessageHandlerSet::handle_message(
2154 &this.message_handlers,
2155 envelope,
2156 this.clone().into(),
2157 cx.clone(),
2158 ) {
2159 log::debug!("{}:ssh message received. name:{type_name}", this.name);
2160 cx.foreground_executor()
2161 .spawn(async move {
2162 match future.await {
2163 Ok(_) => {
2164 log::debug!(
2165 "{}:ssh message handled. name:{type_name}",
2166 this.name
2167 );
2168 }
2169 Err(error) => {
2170 log::error!(
2171 "{}:error handling message. type:{}, error:{}",
2172 this.name,
2173 type_name,
2174 format!("{error:#}").lines().fold(
2175 String::new(),
2176 |mut message, line| {
2177 if !message.is_empty() {
2178 message.push(' ');
2179 }
2180 message.push_str(line);
2181 message
2182 }
2183 )
2184 );
2185 }
2186 }
2187 })
2188 .detach()
2189 } else {
2190 log::error!("{}:unhandled ssh message name:{type_name}", this.name);
2191 }
2192 }
2193 }
2194 anyhow::Ok(())
2195 })
2196 }
2197
2198 pub fn reconnect(
2199 self: &Arc<Self>,
2200 incoming_rx: UnboundedReceiver<Envelope>,
2201 outgoing_tx: UnboundedSender<Envelope>,
2202 cx: &AsyncApp,
2203 ) {
2204 *self.outgoing_tx.lock() = outgoing_tx;
2205 *self.task.lock() = Self::start_handling_messages(Arc::downgrade(self), incoming_rx, cx);
2206 }
2207
2208 pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Entity<E>) {
2209 let id = (TypeId::of::<E>(), remote_id);
2210
2211 let mut message_handlers = self.message_handlers.lock();
2212 if message_handlers
2213 .entities_by_type_and_remote_id
2214 .contains_key(&id)
2215 {
2216 panic!("already subscribed to entity");
2217 }
2218
2219 message_handlers.entities_by_type_and_remote_id.insert(
2220 id,
2221 EntityMessageSubscriber::Entity {
2222 handle: entity.downgrade().into(),
2223 },
2224 );
2225 }
2226
2227 pub fn request<T: RequestMessage>(
2228 &self,
2229 payload: T,
2230 ) -> impl 'static + Future<Output = Result<T::Response>> {
2231 self.request_internal(payload, true)
2232 }
2233
2234 fn request_internal<T: RequestMessage>(
2235 &self,
2236 payload: T,
2237 use_buffer: bool,
2238 ) -> impl 'static + Future<Output = Result<T::Response>> {
2239 log::debug!("ssh request start. name:{}", T::NAME);
2240 let response =
2241 self.request_dynamic(payload.into_envelope(0, None, None), T::NAME, use_buffer);
2242 async move {
2243 let response = response.await?;
2244 log::debug!("ssh request finish. name:{}", T::NAME);
2245 T::Response::from_envelope(response)
2246 .ok_or_else(|| anyhow!("received a response of the wrong type"))
2247 }
2248 }
2249
2250 pub async fn resync(&self, timeout: Duration) -> Result<()> {
2251 smol::future::or(
2252 async {
2253 self.request_internal(proto::FlushBufferedMessages {}, false)
2254 .await?;
2255
2256 for envelope in self.buffer.lock().iter() {
2257 self.outgoing_tx
2258 .lock()
2259 .unbounded_send(envelope.clone())
2260 .ok();
2261 }
2262 Ok(())
2263 },
2264 async {
2265 smol::Timer::after(timeout).await;
2266 Err(anyhow!("Timeout detected"))
2267 },
2268 )
2269 .await
2270 }
2271
2272 pub async fn ping(&self, timeout: Duration) -> Result<()> {
2273 smol::future::or(
2274 async {
2275 self.request(proto::Ping {}).await?;
2276 Ok(())
2277 },
2278 async {
2279 smol::Timer::after(timeout).await;
2280 Err(anyhow!("Timeout detected"))
2281 },
2282 )
2283 .await
2284 }
2285
2286 pub fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {
2287 log::debug!("ssh send name:{}", T::NAME);
2288 self.send_dynamic(payload.into_envelope(0, None, None))
2289 }
2290
2291 fn request_dynamic(
2292 &self,
2293 mut envelope: proto::Envelope,
2294 type_name: &'static str,
2295 use_buffer: bool,
2296 ) -> impl 'static + Future<Output = Result<proto::Envelope>> {
2297 envelope.id = self.next_message_id.fetch_add(1, SeqCst);
2298 let (tx, rx) = oneshot::channel();
2299 let mut response_channels_lock = self.response_channels.lock();
2300 response_channels_lock.insert(MessageId(envelope.id), tx);
2301 drop(response_channels_lock);
2302
2303 let result = if use_buffer {
2304 self.send_buffered(envelope)
2305 } else {
2306 self.send_unbuffered(envelope)
2307 };
2308 async move {
2309 if let Err(error) = &result {
2310 log::error!("failed to send message: {}", error);
2311 return Err(anyhow!("failed to send message: {}", error));
2312 }
2313
2314 let response = rx.await.context("connection lost")?.0;
2315 if let Some(proto::envelope::Payload::Error(error)) = &response.payload {
2316 return Err(RpcError::from_proto(error, type_name));
2317 }
2318 Ok(response)
2319 }
2320 }
2321
2322 pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {
2323 envelope.id = self.next_message_id.fetch_add(1, SeqCst);
2324 self.send_buffered(envelope)
2325 }
2326
2327 fn send_buffered(&self, mut envelope: proto::Envelope) -> Result<()> {
2328 envelope.ack_id = Some(self.max_received.load(SeqCst));
2329 self.buffer.lock().push_back(envelope.clone());
2330 // ignore errors on send (happen while we're reconnecting)
2331 // assume that the global "disconnected" overlay is sufficient.
2332 self.outgoing_tx.lock().unbounded_send(envelope).ok();
2333 Ok(())
2334 }
2335
2336 fn send_unbuffered(&self, mut envelope: proto::Envelope) -> Result<()> {
2337 envelope.ack_id = Some(self.max_received.load(SeqCst));
2338 self.outgoing_tx.lock().unbounded_send(envelope).ok();
2339 Ok(())
2340 }
2341}
2342
2343impl ProtoClient for ChannelClient {
2344 fn request(
2345 &self,
2346 envelope: proto::Envelope,
2347 request_type: &'static str,
2348 ) -> BoxFuture<'static, Result<proto::Envelope>> {
2349 self.request_dynamic(envelope, request_type, true).boxed()
2350 }
2351
2352 fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
2353 self.send_dynamic(envelope)
2354 }
2355
2356 fn send_response(&self, envelope: Envelope, _message_type: &'static str) -> anyhow::Result<()> {
2357 self.send_dynamic(envelope)
2358 }
2359
2360 fn message_handler_set(&self) -> &Mutex<ProtoMessageHandlerSet> {
2361 &self.message_handlers
2362 }
2363
2364 fn is_via_collab(&self) -> bool {
2365 false
2366 }
2367}
2368
2369#[cfg(any(test, feature = "test-support"))]
2370mod fake {
2371 use std::{path::PathBuf, sync::Arc};
2372
2373 use anyhow::Result;
2374 use async_trait::async_trait;
2375 use futures::{
2376 FutureExt, SinkExt, StreamExt,
2377 channel::{
2378 mpsc::{self, Sender},
2379 oneshot,
2380 },
2381 select_biased,
2382 };
2383 use gpui::{App, AppContext as _, AsyncApp, SemanticVersion, Task, TestAppContext};
2384 use release_channel::ReleaseChannel;
2385 use rpc::proto::Envelope;
2386
2387 use super::{
2388 ChannelClient, RemoteConnection, SshClientDelegate, SshConnectionOptions, SshPlatform,
2389 };
2390
2391 pub(super) struct FakeRemoteConnection {
2392 pub(super) connection_options: SshConnectionOptions,
2393 pub(super) server_channel: Arc<ChannelClient>,
2394 pub(super) server_cx: SendableCx,
2395 }
2396
2397 pub(super) struct SendableCx(AsyncApp);
2398 impl SendableCx {
2399 // SAFETY: When run in test mode, GPUI is always single threaded.
2400 pub(super) fn new(cx: &TestAppContext) -> Self {
2401 Self(cx.to_async())
2402 }
2403
2404 // SAFETY: Enforce that we're on the main thread by requiring a valid AsyncApp
2405 fn get(&self, _: &AsyncApp) -> AsyncApp {
2406 self.0.clone()
2407 }
2408 }
2409
2410 // SAFETY: There is no way to access a SendableCx from a different thread, see [`SendableCx::new`] and [`SendableCx::get`]
2411 unsafe impl Send for SendableCx {}
2412 unsafe impl Sync for SendableCx {}
2413
2414 #[async_trait(?Send)]
2415 impl RemoteConnection for FakeRemoteConnection {
2416 async fn kill(&self) -> Result<()> {
2417 Ok(())
2418 }
2419
2420 fn has_been_killed(&self) -> bool {
2421 false
2422 }
2423
2424 fn ssh_args(&self) -> Vec<String> {
2425 Vec::new()
2426 }
2427 fn upload_directory(
2428 &self,
2429 _src_path: PathBuf,
2430 _dest_path: PathBuf,
2431 _cx: &App,
2432 ) -> Task<Result<()>> {
2433 unreachable!()
2434 }
2435
2436 fn connection_options(&self) -> SshConnectionOptions {
2437 self.connection_options.clone()
2438 }
2439
2440 fn simulate_disconnect(&self, cx: &AsyncApp) {
2441 let (outgoing_tx, _) = mpsc::unbounded::<Envelope>();
2442 let (_, incoming_rx) = mpsc::unbounded::<Envelope>();
2443 self.server_channel
2444 .reconnect(incoming_rx, outgoing_tx, &self.server_cx.get(&cx));
2445 }
2446
2447 fn start_proxy(
2448 &self,
2449
2450 _unique_identifier: String,
2451 _reconnect: bool,
2452 mut client_incoming_tx: mpsc::UnboundedSender<Envelope>,
2453 mut client_outgoing_rx: mpsc::UnboundedReceiver<Envelope>,
2454 mut connection_activity_tx: Sender<()>,
2455 _delegate: Arc<dyn SshClientDelegate>,
2456 cx: &mut AsyncApp,
2457 ) -> Task<Result<i32>> {
2458 let (mut server_incoming_tx, server_incoming_rx) = mpsc::unbounded::<Envelope>();
2459 let (server_outgoing_tx, mut server_outgoing_rx) = mpsc::unbounded::<Envelope>();
2460
2461 self.server_channel.reconnect(
2462 server_incoming_rx,
2463 server_outgoing_tx,
2464 &self.server_cx.get(cx),
2465 );
2466
2467 cx.background_spawn(async move {
2468 loop {
2469 select_biased! {
2470 server_to_client = server_outgoing_rx.next().fuse() => {
2471 let Some(server_to_client) = server_to_client else {
2472 return Ok(1)
2473 };
2474 connection_activity_tx.try_send(()).ok();
2475 client_incoming_tx.send(server_to_client).await.ok();
2476 }
2477 client_to_server = client_outgoing_rx.next().fuse() => {
2478 let Some(client_to_server) = client_to_server else {
2479 return Ok(1)
2480 };
2481 server_incoming_tx.send(client_to_server).await.ok();
2482 }
2483 }
2484 }
2485 })
2486 }
2487 }
2488
2489 pub(super) struct Delegate;
2490
2491 impl SshClientDelegate for Delegate {
2492 fn ask_password(&self, _: String, _: oneshot::Sender<String>, _: &mut AsyncApp) {
2493 unreachable!()
2494 }
2495
2496 fn download_server_binary_locally(
2497 &self,
2498 _: SshPlatform,
2499 _: ReleaseChannel,
2500 _: Option<SemanticVersion>,
2501 _: &mut AsyncApp,
2502 ) -> Task<Result<PathBuf>> {
2503 unreachable!()
2504 }
2505
2506 fn get_download_params(
2507 &self,
2508 _platform: SshPlatform,
2509 _release_channel: ReleaseChannel,
2510 _version: Option<SemanticVersion>,
2511 _cx: &mut AsyncApp,
2512 ) -> Task<Result<Option<(String, String)>>> {
2513 unreachable!()
2514 }
2515
2516 fn set_status(&self, _: Option<&str>, _: &mut AsyncApp) {}
2517 }
2518}