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