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