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