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