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