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