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