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