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