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