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,
17 select_biased, AsyncReadExt as _, Future, FutureExt as _, StreamExt as _,
18};
19use gpui::{
20 AppContext, AsyncAppContext, Context, EventEmitter, Model, ModelContext, SemanticVersion, Task,
21 WeakModel,
22};
23use parking_lot::Mutex;
24use rpc::{
25 proto::{self, build_typed_envelope, Envelope, EnvelopedMessage, PeerId, RequestMessage},
26 AnyProtoClient, EntityMessageSubscriber, ProtoClient, ProtoMessageHandlerSet, RpcError,
27};
28use smol::{
29 fs,
30 process::{self, Child, Stdio},
31};
32use std::{
33 any::TypeId,
34 collections::VecDeque,
35 ffi::OsStr,
36 fmt,
37 ops::ControlFlow,
38 path::{Path, PathBuf},
39 sync::{
40 atomic::{AtomicU32, Ordering::SeqCst},
41 Arc, Weak,
42 },
43 time::{Duration, Instant},
44};
45use tempfile::TempDir;
46use util::ResultExt;
47
48#[derive(
49 Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, serde::Serialize, serde::Deserialize,
50)]
51pub struct SshProjectId(pub u64);
52
53#[derive(Clone)]
54pub struct SshSocket {
55 connection_options: SshConnectionOptions,
56 socket_path: PathBuf,
57}
58
59#[derive(Debug, Default, Clone, PartialEq, Eq)]
60pub struct SshConnectionOptions {
61 pub host: String,
62 pub username: Option<String>,
63 pub port: Option<u16>,
64 pub password: Option<String>,
65 pub args: Option<Vec<String>>,
66}
67
68impl SshConnectionOptions {
69 pub fn parse_command_line(input: &str) -> Result<Self> {
70 let input = input.trim_start_matches("ssh ");
71 let mut hostname: Option<String> = None;
72 let mut username: Option<String> = None;
73 let mut port: Option<u16> = None;
74 let mut args = Vec::new();
75
76 // disallowed: -E, -e, -F, -f, -G, -g, -M, -N, -n, -O, -q, -S, -s, -T, -t, -V, -v, -W
77 const ALLOWED_OPTS: &[&str] = &[
78 "-4", "-6", "-A", "-a", "-C", "-K", "-k", "-X", "-x", "-Y", "-y",
79 ];
80 const ALLOWED_ARGS: &[&str] = &[
81 "-B", "-b", "-c", "-D", "-I", "-i", "-J", "-L", "-l", "-m", "-o", "-P", "-p", "-R",
82 "-w",
83 ];
84
85 let mut tokens = shlex::split(input)
86 .ok_or_else(|| anyhow!("invalid input"))?
87 .into_iter();
88
89 'outer: while let Some(arg) = tokens.next() {
90 if ALLOWED_OPTS.contains(&(&arg as &str)) {
91 args.push(arg.to_string());
92 continue;
93 }
94 if arg == "-p" {
95 port = tokens.next().and_then(|arg| arg.parse().ok());
96 continue;
97 } else if let Some(p) = arg.strip_prefix("-p") {
98 port = p.parse().ok();
99 continue;
100 }
101 if arg == "-l" {
102 username = tokens.next();
103 continue;
104 } else if let Some(l) = arg.strip_prefix("-l") {
105 username = Some(l.to_string());
106 continue;
107 }
108 for a in ALLOWED_ARGS {
109 if arg == *a {
110 args.push(arg);
111 if let Some(next) = tokens.next() {
112 args.push(next);
113 }
114 continue 'outer;
115 } else if arg.starts_with(a) {
116 args.push(arg);
117 continue 'outer;
118 }
119 }
120 if arg.starts_with("-") || hostname.is_some() {
121 anyhow::bail!("unsupported argument: {:?}", arg);
122 }
123 let mut input = &arg as &str;
124 if let Some((u, rest)) = input.split_once('@') {
125 input = rest;
126 username = Some(u.to_string());
127 }
128 if let Some((rest, p)) = input.split_once(':') {
129 input = rest;
130 port = p.parse().ok()
131 }
132 hostname = Some(input.to_string())
133 }
134
135 let Some(hostname) = hostname else {
136 anyhow::bail!("missing hostname");
137 };
138
139 Ok(Self {
140 host: hostname.to_string(),
141 username: username.clone(),
142 port,
143 password: None,
144 args: Some(args),
145 })
146 }
147
148 pub fn ssh_url(&self) -> String {
149 let mut result = String::from("ssh://");
150 if let Some(username) = &self.username {
151 result.push_str(username);
152 result.push('@');
153 }
154 result.push_str(&self.host);
155 if let Some(port) = self.port {
156 result.push(':');
157 result.push_str(&port.to_string());
158 }
159 result
160 }
161
162 pub fn additional_args(&self) -> Option<&Vec<String>> {
163 self.args.as_ref()
164 }
165
166 fn scp_url(&self) -> String {
167 if let Some(username) = &self.username {
168 format!("{}@{}", username, self.host)
169 } else {
170 self.host.clone()
171 }
172 }
173
174 pub fn connection_string(&self) -> String {
175 let host = if let Some(username) = &self.username {
176 format!("{}@{}", username, self.host)
177 } else {
178 self.host.clone()
179 };
180 if let Some(port) = &self.port {
181 format!("{}:{}", host, port)
182 } else {
183 host
184 }
185 }
186
187 // Uniquely identifies dev server projects on a remote host. Needs to be
188 // stable for the same dev server project.
189 pub fn remote_server_identifier(&self) -> String {
190 let mut identifier = format!("dev-server-{:?}", self.host);
191 if let Some(username) = self.username.as_ref() {
192 identifier.push('-');
193 identifier.push_str(&username);
194 }
195 identifier
196 }
197}
198
199#[derive(Copy, Clone, Debug)]
200pub struct SshPlatform {
201 pub os: &'static str,
202 pub arch: &'static str,
203}
204
205impl SshPlatform {
206 pub fn triple(&self) -> Option<String> {
207 Some(format!(
208 "{}-{}",
209 self.arch,
210 match self.os {
211 "linux" => "unknown-linux-gnu",
212 "macos" => "apple-darwin",
213 _ => return None,
214 }
215 ))
216 }
217}
218
219pub trait SshClientDelegate: Send + Sync {
220 fn ask_password(
221 &self,
222 prompt: String,
223 cx: &mut AsyncAppContext,
224 ) -> oneshot::Receiver<Result<String>>;
225 fn remote_server_binary_path(
226 &self,
227 platform: SshPlatform,
228 cx: &mut AsyncAppContext,
229 ) -> Result<PathBuf>;
230 fn get_server_binary(
231 &self,
232 platform: SshPlatform,
233 cx: &mut AsyncAppContext,
234 ) -> oneshot::Receiver<Result<(PathBuf, SemanticVersion)>>;
235 fn set_status(&self, status: Option<&str>, cx: &mut AsyncAppContext);
236}
237
238impl SshSocket {
239 fn ssh_command<S: AsRef<OsStr>>(&self, program: S) -> process::Command {
240 let mut command = process::Command::new("ssh");
241 self.ssh_options(&mut command)
242 .arg(self.connection_options.ssh_url())
243 .arg(program);
244 command
245 }
246
247 fn ssh_options<'a>(&self, command: &'a mut process::Command) -> &'a mut process::Command {
248 command
249 .stdin(Stdio::piped())
250 .stdout(Stdio::piped())
251 .stderr(Stdio::piped())
252 .args(["-o", "ControlMaster=no", "-o"])
253 .arg(format!("ControlPath={}", self.socket_path.display()))
254 }
255
256 fn ssh_args(&self) -> Vec<String> {
257 vec![
258 "-o".to_string(),
259 "ControlMaster=no".to_string(),
260 "-o".to_string(),
261 format!("ControlPath={}", self.socket_path.display()),
262 self.connection_options.ssh_url(),
263 ]
264 }
265}
266
267async fn run_cmd(command: &mut process::Command) -> Result<String> {
268 let output = command.output().await?;
269 if output.status.success() {
270 Ok(String::from_utf8_lossy(&output.stdout).to_string())
271 } else {
272 Err(anyhow!(
273 "failed to run command: {}",
274 String::from_utf8_lossy(&output.stderr)
275 ))
276 }
277}
278
279const MAX_MISSED_HEARTBEATS: usize = 5;
280const HEARTBEAT_INTERVAL: Duration = Duration::from_secs(5);
281const HEARTBEAT_TIMEOUT: Duration = Duration::from_secs(5);
282
283const MAX_RECONNECT_ATTEMPTS: usize = 3;
284
285enum State {
286 Connecting,
287 Connected {
288 ssh_connection: Box<dyn SshRemoteProcess>,
289 delegate: Arc<dyn SshClientDelegate>,
290
291 multiplex_task: Task<Result<()>>,
292 heartbeat_task: Task<Result<()>>,
293 },
294 HeartbeatMissed {
295 missed_heartbeats: usize,
296
297 ssh_connection: Box<dyn SshRemoteProcess>,
298 delegate: Arc<dyn SshClientDelegate>,
299
300 multiplex_task: Task<Result<()>>,
301 heartbeat_task: Task<Result<()>>,
302 },
303 Reconnecting,
304 ReconnectFailed {
305 ssh_connection: Box<dyn SshRemoteProcess>,
306 delegate: Arc<dyn SshClientDelegate>,
307
308 error: anyhow::Error,
309 attempts: usize,
310 },
311 ReconnectExhausted,
312 ServerNotRunning,
313}
314
315impl fmt::Display for State {
316 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
317 match self {
318 Self::Connecting => write!(f, "connecting"),
319 Self::Connected { .. } => write!(f, "connected"),
320 Self::Reconnecting => write!(f, "reconnecting"),
321 Self::ReconnectFailed { .. } => write!(f, "reconnect failed"),
322 Self::ReconnectExhausted => write!(f, "reconnect exhausted"),
323 Self::HeartbeatMissed { .. } => write!(f, "heartbeat missed"),
324 Self::ServerNotRunning { .. } => write!(f, "server not running"),
325 }
326 }
327}
328
329impl State {
330 fn ssh_connection(&self) -> Option<&dyn SshRemoteProcess> {
331 match self {
332 Self::Connected { ssh_connection, .. } => Some(ssh_connection.as_ref()),
333 Self::HeartbeatMissed { ssh_connection, .. } => Some(ssh_connection.as_ref()),
334 Self::ReconnectFailed { ssh_connection, .. } => Some(ssh_connection.as_ref()),
335 _ => None,
336 }
337 }
338
339 fn can_reconnect(&self) -> bool {
340 match self {
341 Self::Connected { .. }
342 | Self::HeartbeatMissed { .. }
343 | Self::ReconnectFailed { .. } => true,
344 State::Connecting
345 | State::Reconnecting
346 | State::ReconnectExhausted
347 | State::ServerNotRunning => false,
348 }
349 }
350
351 fn is_reconnect_failed(&self) -> bool {
352 matches!(self, Self::ReconnectFailed { .. })
353 }
354
355 fn is_reconnect_exhausted(&self) -> bool {
356 matches!(self, Self::ReconnectExhausted { .. })
357 }
358
359 fn is_server_not_running(&self) -> bool {
360 matches!(self, Self::ServerNotRunning)
361 }
362
363 fn is_reconnecting(&self) -> bool {
364 matches!(self, Self::Reconnecting { .. })
365 }
366
367 fn heartbeat_recovered(self) -> Self {
368 match self {
369 Self::HeartbeatMissed {
370 ssh_connection,
371 delegate,
372 multiplex_task,
373 heartbeat_task,
374 ..
375 } => Self::Connected {
376 ssh_connection,
377 delegate,
378 multiplex_task,
379 heartbeat_task,
380 },
381 _ => self,
382 }
383 }
384
385 fn heartbeat_missed(self) -> Self {
386 match self {
387 Self::Connected {
388 ssh_connection,
389 delegate,
390 multiplex_task,
391 heartbeat_task,
392 } => Self::HeartbeatMissed {
393 missed_heartbeats: 1,
394 ssh_connection,
395 delegate,
396 multiplex_task,
397 heartbeat_task,
398 },
399 Self::HeartbeatMissed {
400 missed_heartbeats,
401 ssh_connection,
402 delegate,
403 multiplex_task,
404 heartbeat_task,
405 } => Self::HeartbeatMissed {
406 missed_heartbeats: missed_heartbeats + 1,
407 ssh_connection,
408 delegate,
409 multiplex_task,
410 heartbeat_task,
411 },
412 _ => self,
413 }
414 }
415}
416
417/// The state of the ssh connection.
418#[derive(Clone, Copy, Debug, PartialEq, Eq)]
419pub enum ConnectionState {
420 Connecting,
421 Connected,
422 HeartbeatMissed,
423 Reconnecting,
424 Disconnected,
425}
426
427impl From<&State> for ConnectionState {
428 fn from(value: &State) -> Self {
429 match value {
430 State::Connecting => Self::Connecting,
431 State::Connected { .. } => Self::Connected,
432 State::Reconnecting | State::ReconnectFailed { .. } => Self::Reconnecting,
433 State::HeartbeatMissed { .. } => Self::HeartbeatMissed,
434 State::ReconnectExhausted => Self::Disconnected,
435 State::ServerNotRunning => Self::Disconnected,
436 }
437 }
438}
439
440pub struct SshRemoteClient {
441 client: Arc<ChannelClient>,
442 unique_identifier: String,
443 connection_options: SshConnectionOptions,
444 state: Arc<Mutex<Option<State>>>,
445}
446
447#[derive(Debug)]
448pub enum SshRemoteEvent {
449 Disconnected,
450}
451
452impl EventEmitter<SshRemoteEvent> for SshRemoteClient {}
453
454impl SshRemoteClient {
455 pub fn new(
456 unique_identifier: String,
457 connection_options: SshConnectionOptions,
458 delegate: Arc<dyn SshClientDelegate>,
459 cx: &AppContext,
460 ) -> Task<Result<Model<Self>>> {
461 cx.spawn(|mut cx| async move {
462 let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
463 let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
464 let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
465
466 let client =
467 cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "client"))?;
468 let this = cx.new_model(|_| Self {
469 client: client.clone(),
470 unique_identifier: unique_identifier.clone(),
471 connection_options: connection_options.clone(),
472 state: Arc::new(Mutex::new(Some(State::Connecting))),
473 })?;
474
475 let (ssh_connection, io_task) = Self::establish_connection(
476 unique_identifier,
477 false,
478 connection_options,
479 incoming_tx,
480 outgoing_rx,
481 connection_activity_tx,
482 delegate.clone(),
483 &mut cx,
484 )
485 .await?;
486
487 let multiplex_task = Self::monitor(this.downgrade(), io_task, &cx);
488
489 if let Err(error) = client.ping(HEARTBEAT_TIMEOUT).await {
490 log::error!("failed to establish connection: {}", error);
491 return Err(error);
492 }
493
494 let heartbeat_task = Self::heartbeat(this.downgrade(), connection_activity_rx, &mut cx);
495
496 this.update(&mut cx, |this, _| {
497 *this.state.lock() = Some(State::Connected {
498 ssh_connection,
499 delegate,
500 multiplex_task,
501 heartbeat_task,
502 });
503 })?;
504
505 Ok(this)
506 })
507 }
508
509 pub fn shutdown_processes<T: RequestMessage>(
510 &self,
511 shutdown_request: Option<T>,
512 ) -> Option<impl Future<Output = ()>> {
513 let state = self.state.lock().take()?;
514 log::info!("shutting down ssh processes");
515
516 let State::Connected {
517 multiplex_task,
518 heartbeat_task,
519 ssh_connection,
520 delegate,
521 } = state
522 else {
523 return None;
524 };
525
526 let client = self.client.clone();
527
528 Some(async move {
529 if let Some(shutdown_request) = shutdown_request {
530 client.send(shutdown_request).log_err();
531 // We wait 50ms instead of waiting for a response, because
532 // waiting for a response would require us to wait on the main thread
533 // which we want to avoid in an `on_app_quit` callback.
534 smol::Timer::after(Duration::from_millis(50)).await;
535 }
536
537 // Drop `multiplex_task` because it owns our ssh_proxy_process, which is a
538 // child of master_process.
539 drop(multiplex_task);
540 // Now drop the rest of state, which kills master process.
541 drop(heartbeat_task);
542 drop(ssh_connection);
543 drop(delegate);
544 })
545 }
546
547 fn reconnect(&mut self, cx: &mut ModelContext<Self>) -> Result<()> {
548 let mut lock = self.state.lock();
549
550 let can_reconnect = lock
551 .as_ref()
552 .map(|state| state.can_reconnect())
553 .unwrap_or(false);
554 if !can_reconnect {
555 let error = if let Some(state) = lock.as_ref() {
556 format!("invalid state, cannot reconnect while in state {state}")
557 } else {
558 "no state set".to_string()
559 };
560 log::info!("aborting reconnect, because not in state that allows reconnecting");
561 return Err(anyhow!(error));
562 }
563
564 let state = lock.take().unwrap();
565 let (attempts, mut ssh_connection, delegate) = match state {
566 State::Connected {
567 ssh_connection,
568 delegate,
569 multiplex_task,
570 heartbeat_task,
571 }
572 | State::HeartbeatMissed {
573 ssh_connection,
574 delegate,
575 multiplex_task,
576 heartbeat_task,
577 ..
578 } => {
579 drop(multiplex_task);
580 drop(heartbeat_task);
581 (0, ssh_connection, delegate)
582 }
583 State::ReconnectFailed {
584 attempts,
585 ssh_connection,
586 delegate,
587 ..
588 } => (attempts, ssh_connection, delegate),
589 State::Connecting
590 | State::Reconnecting
591 | State::ReconnectExhausted
592 | State::ServerNotRunning => unreachable!(),
593 };
594
595 let attempts = attempts + 1;
596 if attempts > MAX_RECONNECT_ATTEMPTS {
597 log::error!(
598 "Failed to reconnect to after {} attempts, giving up",
599 MAX_RECONNECT_ATTEMPTS
600 );
601 drop(lock);
602 self.set_state(State::ReconnectExhausted, cx);
603 return Ok(());
604 }
605 drop(lock);
606
607 self.set_state(State::Reconnecting, cx);
608
609 log::info!("Trying to reconnect to ssh server... Attempt {}", attempts);
610
611 let identifier = self.unique_identifier.clone();
612 let client = self.client.clone();
613 let reconnect_task = cx.spawn(|this, mut cx| async move {
614 macro_rules! failed {
615 ($error:expr, $attempts:expr, $ssh_connection:expr, $delegate:expr) => {
616 return State::ReconnectFailed {
617 error: anyhow!($error),
618 attempts: $attempts,
619 ssh_connection: $ssh_connection,
620 delegate: $delegate,
621 };
622 };
623 }
624
625 if let Err(error) = ssh_connection
626 .kill()
627 .await
628 .context("Failed to kill ssh process")
629 {
630 failed!(error, attempts, ssh_connection, delegate);
631 };
632
633 let connection_options = ssh_connection.connection_options();
634
635 let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
636 let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
637 let (connection_activity_tx, connection_activity_rx) = mpsc::channel::<()>(1);
638
639 let (ssh_connection, io_task) = match Self::establish_connection(
640 identifier,
641 true,
642 connection_options,
643 incoming_tx,
644 outgoing_rx,
645 connection_activity_tx,
646 delegate.clone(),
647 &mut cx,
648 )
649 .await
650 {
651 Ok((ssh_connection, ssh_process)) => (ssh_connection, ssh_process),
652 Err(error) => {
653 failed!(error, attempts, ssh_connection, delegate);
654 }
655 };
656
657 let multiplex_task = Self::monitor(this.clone(), io_task, &cx);
658 client.reconnect(incoming_rx, outgoing_tx, &cx);
659
660 if let Err(error) = client.resync(HEARTBEAT_TIMEOUT).await {
661 failed!(error, attempts, ssh_connection, delegate);
662 };
663
664 State::Connected {
665 ssh_connection,
666 delegate,
667 multiplex_task,
668 heartbeat_task: Self::heartbeat(this.clone(), connection_activity_rx, &mut cx),
669 }
670 });
671
672 cx.spawn(|this, mut cx| async move {
673 let new_state = reconnect_task.await;
674 this.update(&mut cx, |this, cx| {
675 this.try_set_state(cx, |old_state| {
676 if old_state.is_reconnecting() {
677 match &new_state {
678 State::Connecting
679 | State::Reconnecting { .. }
680 | State::HeartbeatMissed { .. }
681 | State::ServerNotRunning => {}
682 State::Connected { .. } => {
683 log::info!("Successfully reconnected");
684 }
685 State::ReconnectFailed {
686 error, attempts, ..
687 } => {
688 log::error!(
689 "Reconnect attempt {} failed: {:?}. Starting new attempt...",
690 attempts,
691 error
692 );
693 }
694 State::ReconnectExhausted => {
695 log::error!("Reconnect attempt failed and all attempts exhausted");
696 }
697 }
698 Some(new_state)
699 } else {
700 None
701 }
702 });
703
704 if this.state_is(State::is_reconnect_failed) {
705 this.reconnect(cx)
706 } else if this.state_is(State::is_reconnect_exhausted) {
707 Ok(())
708 } else {
709 log::debug!("State has transition from Reconnecting into new state while attempting reconnect.");
710 Ok(())
711 }
712 })
713 })
714 .detach_and_log_err(cx);
715
716 Ok(())
717 }
718
719 fn heartbeat(
720 this: WeakModel<Self>,
721 mut connection_activity_rx: mpsc::Receiver<()>,
722 cx: &mut AsyncAppContext,
723 ) -> Task<Result<()>> {
724 let Ok(client) = this.update(cx, |this, _| this.client.clone()) else {
725 return Task::ready(Err(anyhow!("SshRemoteClient lost")));
726 };
727
728 cx.spawn(|mut cx| {
729 let this = this.clone();
730 async move {
731 let mut missed_heartbeats = 0;
732
733 let keepalive_timer = cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse();
734 futures::pin_mut!(keepalive_timer);
735
736 loop {
737 select_biased! {
738 result = connection_activity_rx.next().fuse() => {
739 if result.is_none() {
740 log::warn!("ssh heartbeat: connection activity channel has been dropped. stopping.");
741 return Ok(());
742 }
743
744 if missed_heartbeats != 0 {
745 missed_heartbeats = 0;
746 this.update(&mut cx, |this, mut cx| {
747 this.handle_heartbeat_result(missed_heartbeats, &mut cx)
748 })?;
749 }
750 }
751 _ = keepalive_timer => {
752 log::debug!("Sending heartbeat to server...");
753
754 let result = select_biased! {
755 _ = connection_activity_rx.next().fuse() => {
756 Ok(())
757 }
758 ping_result = client.ping(HEARTBEAT_TIMEOUT).fuse() => {
759 ping_result
760 }
761 };
762
763 if result.is_err() {
764 missed_heartbeats += 1;
765 log::warn!(
766 "No heartbeat from server after {:?}. Missed heartbeat {} out of {}.",
767 HEARTBEAT_TIMEOUT,
768 missed_heartbeats,
769 MAX_MISSED_HEARTBEATS
770 );
771 } else if missed_heartbeats != 0 {
772 missed_heartbeats = 0;
773 } else {
774 continue;
775 }
776
777 let result = this.update(&mut cx, |this, mut cx| {
778 this.handle_heartbeat_result(missed_heartbeats, &mut cx)
779 })?;
780 if result.is_break() {
781 return Ok(());
782 }
783 }
784 }
785
786 keepalive_timer.set(cx.background_executor().timer(HEARTBEAT_INTERVAL).fuse());
787 }
788 }
789 })
790 }
791
792 fn handle_heartbeat_result(
793 &mut self,
794 missed_heartbeats: usize,
795 cx: &mut ModelContext<Self>,
796 ) -> ControlFlow<()> {
797 let state = self.state.lock().take().unwrap();
798 let next_state = if missed_heartbeats > 0 {
799 state.heartbeat_missed()
800 } else {
801 state.heartbeat_recovered()
802 };
803
804 self.set_state(next_state, cx);
805
806 if missed_heartbeats >= MAX_MISSED_HEARTBEATS {
807 log::error!(
808 "Missed last {} heartbeats. Reconnecting...",
809 missed_heartbeats
810 );
811
812 self.reconnect(cx)
813 .context("failed to start reconnect process after missing heartbeats")
814 .log_err();
815 ControlFlow::Break(())
816 } else {
817 ControlFlow::Continue(())
818 }
819 }
820
821 fn multiplex(
822 mut ssh_proxy_process: Child,
823 incoming_tx: UnboundedSender<Envelope>,
824 mut outgoing_rx: UnboundedReceiver<Envelope>,
825 mut connection_activity_tx: Sender<()>,
826 cx: &AsyncAppContext,
827 ) -> Task<Result<i32>> {
828 let mut child_stderr = ssh_proxy_process.stderr.take().unwrap();
829 let mut child_stdout = ssh_proxy_process.stdout.take().unwrap();
830 let mut child_stdin = ssh_proxy_process.stdin.take().unwrap();
831
832 let mut stdin_buffer = Vec::new();
833 let mut stdout_buffer = Vec::new();
834 let mut stderr_buffer = Vec::new();
835 let mut stderr_offset = 0;
836
837 let stdin_task = cx.background_executor().spawn(async move {
838 while let Some(outgoing) = outgoing_rx.next().await {
839 write_message(&mut child_stdin, &mut stdin_buffer, outgoing).await?;
840 }
841 anyhow::Ok(())
842 });
843
844 let stdout_task = cx.background_executor().spawn({
845 let mut connection_activity_tx = connection_activity_tx.clone();
846 async move {
847 loop {
848 stdout_buffer.resize(MESSAGE_LEN_SIZE, 0);
849 let len = child_stdout.read(&mut stdout_buffer).await?;
850
851 if len == 0 {
852 return anyhow::Ok(());
853 }
854
855 if len < MESSAGE_LEN_SIZE {
856 child_stdout.read_exact(&mut stdout_buffer[len..]).await?;
857 }
858
859 let message_len = message_len_from_buffer(&stdout_buffer);
860 let envelope =
861 read_message_with_len(&mut child_stdout, &mut stdout_buffer, message_len)
862 .await?;
863 connection_activity_tx.try_send(()).ok();
864 incoming_tx.unbounded_send(envelope).ok();
865 }
866 }
867 });
868
869 let stderr_task: Task<anyhow::Result<()>> = cx.background_executor().spawn(async move {
870 loop {
871 stderr_buffer.resize(stderr_offset + 1024, 0);
872
873 let len = child_stderr
874 .read(&mut stderr_buffer[stderr_offset..])
875 .await?;
876 if len == 0 {
877 return anyhow::Ok(());
878 }
879
880 stderr_offset += len;
881 let mut start_ix = 0;
882 while let Some(ix) = stderr_buffer[start_ix..stderr_offset]
883 .iter()
884 .position(|b| b == &b'\n')
885 {
886 let line_ix = start_ix + ix;
887 let content = &stderr_buffer[start_ix..line_ix];
888 start_ix = line_ix + 1;
889 if let Ok(record) = serde_json::from_slice::<LogRecord>(content) {
890 record.log(log::logger())
891 } else {
892 eprintln!("(remote) {}", String::from_utf8_lossy(content));
893 }
894 }
895 stderr_buffer.drain(0..start_ix);
896 stderr_offset -= start_ix;
897
898 connection_activity_tx.try_send(()).ok();
899 }
900 });
901
902 cx.spawn(|_| async move {
903 let result = futures::select! {
904 result = stdin_task.fuse() => {
905 result.context("stdin")
906 }
907 result = stdout_task.fuse() => {
908 result.context("stdout")
909 }
910 result = stderr_task.fuse() => {
911 result.context("stderr")
912 }
913 };
914
915 let status = ssh_proxy_process.status().await?.code().unwrap_or(1);
916 match result {
917 Ok(_) => Ok(status),
918 Err(error) => Err(error),
919 }
920 })
921 }
922
923 fn monitor(
924 this: WeakModel<Self>,
925 io_task: Task<Result<i32>>,
926 cx: &AsyncAppContext,
927 ) -> Task<Result<()>> {
928 cx.spawn(|mut cx| async move {
929 let result = io_task.await;
930
931 match result {
932 Ok(exit_code) => {
933 if let Some(error) = ProxyLaunchError::from_exit_code(exit_code) {
934 match error {
935 ProxyLaunchError::ServerNotRunning => {
936 log::error!("failed to reconnect because server is not running");
937 this.update(&mut cx, |this, cx| {
938 this.set_state(State::ServerNotRunning, cx);
939 })?;
940 }
941 }
942 } else if exit_code > 0 {
943 log::error!("proxy process terminated unexpectedly");
944 this.update(&mut cx, |this, cx| {
945 this.reconnect(cx).ok();
946 })?;
947 }
948 }
949 Err(error) => {
950 log::warn!("ssh io task died with error: {:?}. reconnecting...", error);
951 this.update(&mut cx, |this, cx| {
952 this.reconnect(cx).ok();
953 })?;
954 }
955 }
956
957 Ok(())
958 })
959 }
960
961 fn state_is(&self, check: impl FnOnce(&State) -> bool) -> bool {
962 self.state.lock().as_ref().map_or(false, check)
963 }
964
965 fn try_set_state(
966 &self,
967 cx: &mut ModelContext<Self>,
968 map: impl FnOnce(&State) -> Option<State>,
969 ) {
970 let mut lock = self.state.lock();
971 let new_state = lock.as_ref().and_then(map);
972
973 if let Some(new_state) = new_state {
974 lock.replace(new_state);
975 cx.notify();
976 }
977 }
978
979 fn set_state(&self, state: State, cx: &mut ModelContext<Self>) {
980 log::info!("setting state to '{}'", &state);
981
982 let is_reconnect_exhausted = state.is_reconnect_exhausted();
983 let is_server_not_running = state.is_server_not_running();
984 self.state.lock().replace(state);
985
986 if is_reconnect_exhausted || is_server_not_running {
987 cx.emit(SshRemoteEvent::Disconnected);
988 }
989 cx.notify();
990 }
991
992 #[allow(clippy::too_many_arguments)]
993 async fn establish_connection(
994 unique_identifier: String,
995 reconnect: bool,
996 connection_options: SshConnectionOptions,
997 incoming_tx: UnboundedSender<Envelope>,
998 outgoing_rx: UnboundedReceiver<Envelope>,
999 connection_activity_tx: Sender<()>,
1000 delegate: Arc<dyn SshClientDelegate>,
1001 cx: &mut AsyncAppContext,
1002 ) -> Result<(Box<dyn SshRemoteProcess>, Task<Result<i32>>)> {
1003 #[cfg(any(test, feature = "test-support"))]
1004 if let Some(fake) = fake::SshRemoteConnection::new(&connection_options) {
1005 let io_task = fake::SshRemoteConnection::multiplex(
1006 fake.connection_options(),
1007 incoming_tx,
1008 outgoing_rx,
1009 connection_activity_tx,
1010 cx,
1011 )
1012 .await;
1013 return Ok((fake, io_task));
1014 }
1015
1016 let ssh_connection =
1017 SshRemoteConnection::new(connection_options, delegate.clone(), cx).await?;
1018
1019 let platform = ssh_connection.query_platform().await?;
1020 let remote_binary_path = delegate.remote_server_binary_path(platform, cx)?;
1021 if !reconnect {
1022 ssh_connection
1023 .ensure_server_binary(&delegate, &remote_binary_path, platform, cx)
1024 .await?;
1025 }
1026
1027 let socket = ssh_connection.socket.clone();
1028 run_cmd(socket.ssh_command(&remote_binary_path).arg("version")).await?;
1029
1030 delegate.set_status(Some("Starting proxy"), cx);
1031
1032 let mut start_proxy_command = format!(
1033 "RUST_LOG={} RUST_BACKTRACE={} {:?} proxy --identifier {}",
1034 std::env::var("RUST_LOG").unwrap_or_default(),
1035 std::env::var("RUST_BACKTRACE").unwrap_or_default(),
1036 remote_binary_path,
1037 unique_identifier,
1038 );
1039 if reconnect {
1040 start_proxy_command.push_str(" --reconnect");
1041 }
1042
1043 let ssh_proxy_process = socket
1044 .ssh_command(start_proxy_command)
1045 // IMPORTANT: we kill this process when we drop the task that uses it.
1046 .kill_on_drop(true)
1047 .spawn()
1048 .context("failed to spawn remote server")?;
1049
1050 let io_task = Self::multiplex(
1051 ssh_proxy_process,
1052 incoming_tx,
1053 outgoing_rx,
1054 connection_activity_tx,
1055 &cx,
1056 );
1057
1058 Ok((Box::new(ssh_connection), io_task))
1059 }
1060
1061 pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Model<E>) {
1062 self.client.subscribe_to_entity(remote_id, entity);
1063 }
1064
1065 pub fn ssh_args(&self) -> Option<Vec<String>> {
1066 self.state
1067 .lock()
1068 .as_ref()
1069 .and_then(|state| state.ssh_connection())
1070 .map(|ssh_connection| ssh_connection.ssh_args())
1071 }
1072
1073 pub fn proto_client(&self) -> AnyProtoClient {
1074 self.client.clone().into()
1075 }
1076
1077 pub fn connection_string(&self) -> String {
1078 self.connection_options.connection_string()
1079 }
1080
1081 pub fn connection_options(&self) -> SshConnectionOptions {
1082 self.connection_options.clone()
1083 }
1084
1085 pub fn connection_state(&self) -> ConnectionState {
1086 self.state
1087 .lock()
1088 .as_ref()
1089 .map(ConnectionState::from)
1090 .unwrap_or(ConnectionState::Disconnected)
1091 }
1092
1093 pub fn is_disconnected(&self) -> bool {
1094 self.connection_state() == ConnectionState::Disconnected
1095 }
1096
1097 #[cfg(any(test, feature = "test-support"))]
1098 pub fn simulate_disconnect(&self, client_cx: &mut AppContext) -> Task<()> {
1099 let port = self.connection_options().port.unwrap();
1100 client_cx.spawn(|cx| async move {
1101 let (channel, server_cx) = cx
1102 .update_global(|c: &mut fake::ServerConnections, _| c.get(port))
1103 .unwrap();
1104
1105 let (outgoing_tx, _) = mpsc::unbounded::<Envelope>();
1106 let (_, incoming_rx) = mpsc::unbounded::<Envelope>();
1107 channel.reconnect(incoming_rx, outgoing_tx, &server_cx);
1108 })
1109 }
1110
1111 #[cfg(any(test, feature = "test-support"))]
1112 pub fn fake_server(
1113 client_cx: &mut gpui::TestAppContext,
1114 server_cx: &mut gpui::TestAppContext,
1115 ) -> (u16, Arc<ChannelClient>) {
1116 use gpui::BorrowAppContext;
1117 let (outgoing_tx, _) = mpsc::unbounded::<Envelope>();
1118 let (_, incoming_rx) = mpsc::unbounded::<Envelope>();
1119 let server_client =
1120 server_cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx, "fake-server"));
1121 let port = client_cx.update(|cx| {
1122 cx.update_default_global(|c: &mut fake::ServerConnections, _| {
1123 c.push(server_client.clone(), server_cx.to_async())
1124 })
1125 });
1126 (port, server_client)
1127 }
1128
1129 #[cfg(any(test, feature = "test-support"))]
1130 pub async fn fake_client(port: u16, client_cx: &mut gpui::TestAppContext) -> Model<Self> {
1131 client_cx
1132 .update(|cx| {
1133 Self::new(
1134 "fake".to_string(),
1135 SshConnectionOptions {
1136 host: "<fake>".to_string(),
1137 port: Some(port),
1138 ..Default::default()
1139 },
1140 Arc::new(fake::Delegate),
1141 cx,
1142 )
1143 })
1144 .await
1145 .unwrap()
1146 }
1147}
1148
1149impl From<SshRemoteClient> for AnyProtoClient {
1150 fn from(client: SshRemoteClient) -> Self {
1151 AnyProtoClient::new(client.client.clone())
1152 }
1153}
1154
1155#[async_trait]
1156trait SshRemoteProcess: Send + Sync {
1157 async fn kill(&mut self) -> Result<()>;
1158 fn ssh_args(&self) -> Vec<String>;
1159 fn connection_options(&self) -> SshConnectionOptions;
1160}
1161
1162struct SshRemoteConnection {
1163 socket: SshSocket,
1164 master_process: process::Child,
1165 _temp_dir: TempDir,
1166}
1167
1168impl Drop for SshRemoteConnection {
1169 fn drop(&mut self) {
1170 if let Err(error) = self.master_process.kill() {
1171 log::error!("failed to kill SSH master process: {}", error);
1172 }
1173 }
1174}
1175
1176#[async_trait]
1177impl SshRemoteProcess for SshRemoteConnection {
1178 async fn kill(&mut self) -> Result<()> {
1179 self.master_process.kill()?;
1180
1181 self.master_process.status().await?;
1182
1183 Ok(())
1184 }
1185
1186 fn ssh_args(&self) -> Vec<String> {
1187 self.socket.ssh_args()
1188 }
1189
1190 fn connection_options(&self) -> SshConnectionOptions {
1191 self.socket.connection_options.clone()
1192 }
1193}
1194
1195impl SshRemoteConnection {
1196 #[cfg(not(unix))]
1197 async fn new(
1198 _connection_options: SshConnectionOptions,
1199 _delegate: Arc<dyn SshClientDelegate>,
1200 _cx: &mut AsyncAppContext,
1201 ) -> Result<Self> {
1202 Err(anyhow!("ssh is not supported on this platform"))
1203 }
1204
1205 #[cfg(unix)]
1206 async fn new(
1207 connection_options: SshConnectionOptions,
1208 delegate: Arc<dyn SshClientDelegate>,
1209 cx: &mut AsyncAppContext,
1210 ) -> Result<Self> {
1211 use futures::AsyncWriteExt as _;
1212 use futures::{io::BufReader, AsyncBufReadExt as _};
1213 use smol::{fs::unix::PermissionsExt as _, net::unix::UnixListener};
1214 use util::ResultExt as _;
1215
1216 delegate.set_status(Some("Connecting"), cx);
1217
1218 let url = connection_options.ssh_url();
1219 let temp_dir = tempfile::Builder::new()
1220 .prefix("zed-ssh-session")
1221 .tempdir()?;
1222
1223 // Create a domain socket listener to handle requests from the askpass program.
1224 let askpass_socket = temp_dir.path().join("askpass.sock");
1225 let (askpass_opened_tx, askpass_opened_rx) = oneshot::channel::<()>();
1226 let listener =
1227 UnixListener::bind(&askpass_socket).context("failed to create askpass socket")?;
1228
1229 let askpass_task = cx.spawn({
1230 let delegate = delegate.clone();
1231 |mut cx| async move {
1232 let mut askpass_opened_tx = Some(askpass_opened_tx);
1233
1234 while let Ok((mut stream, _)) = listener.accept().await {
1235 if let Some(askpass_opened_tx) = askpass_opened_tx.take() {
1236 askpass_opened_tx.send(()).ok();
1237 }
1238 let mut buffer = Vec::new();
1239 let mut reader = BufReader::new(&mut stream);
1240 if reader.read_until(b'\0', &mut buffer).await.is_err() {
1241 buffer.clear();
1242 }
1243 let password_prompt = String::from_utf8_lossy(&buffer);
1244 if let Some(password) = delegate
1245 .ask_password(password_prompt.to_string(), &mut cx)
1246 .await
1247 .context("failed to get ssh password")
1248 .and_then(|p| p)
1249 .log_err()
1250 {
1251 stream.write_all(password.as_bytes()).await.log_err();
1252 }
1253 }
1254 }
1255 });
1256
1257 // Create an askpass script that communicates back to this process.
1258 let askpass_script = format!(
1259 "{shebang}\n{print_args} | nc -U {askpass_socket} 2> /dev/null \n",
1260 askpass_socket = askpass_socket.display(),
1261 print_args = "printf '%s\\0' \"$@\"",
1262 shebang = "#!/bin/sh",
1263 );
1264 let askpass_script_path = temp_dir.path().join("askpass.sh");
1265 fs::write(&askpass_script_path, askpass_script).await?;
1266 fs::set_permissions(&askpass_script_path, std::fs::Permissions::from_mode(0o755)).await?;
1267
1268 // Start the master SSH process, which does not do anything except for establish
1269 // the connection and keep it open, allowing other ssh commands to reuse it
1270 // via a control socket.
1271 let socket_path = temp_dir.path().join("ssh.sock");
1272 let mut master_process = process::Command::new("ssh")
1273 .stdin(Stdio::null())
1274 .stdout(Stdio::piped())
1275 .stderr(Stdio::piped())
1276 .env("SSH_ASKPASS_REQUIRE", "force")
1277 .env("SSH_ASKPASS", &askpass_script_path)
1278 .args(connection_options.additional_args().unwrap_or(&Vec::new()))
1279 .args([
1280 "-N",
1281 "-o",
1282 "ControlPersist=no",
1283 "-o",
1284 "ControlMaster=yes",
1285 "-o",
1286 ])
1287 .arg(format!("ControlPath={}", socket_path.display()))
1288 .arg(&url)
1289 .spawn()?;
1290
1291 // Wait for this ssh process to close its stdout, indicating that authentication
1292 // has completed.
1293 let stdout = master_process.stdout.as_mut().unwrap();
1294 let mut output = Vec::new();
1295 let connection_timeout = Duration::from_secs(10);
1296
1297 let result = select_biased! {
1298 _ = askpass_opened_rx.fuse() => {
1299 // If the askpass script has opened, that means the user is typing
1300 // their password, in which case we don't want to timeout anymore,
1301 // since we know a connection has been established.
1302 stdout.read_to_end(&mut output).await?;
1303 Ok(())
1304 }
1305 result = stdout.read_to_end(&mut output).fuse() => {
1306 result?;
1307 Ok(())
1308 }
1309 _ = futures::FutureExt::fuse(smol::Timer::after(connection_timeout)) => {
1310 Err(anyhow!("Exceeded {:?} timeout trying to connect to host", connection_timeout))
1311 }
1312 };
1313
1314 if let Err(e) = result {
1315 return Err(e.context("Failed to connect to host"));
1316 }
1317
1318 drop(askpass_task);
1319
1320 if master_process.try_status()?.is_some() {
1321 output.clear();
1322 let mut stderr = master_process.stderr.take().unwrap();
1323 stderr.read_to_end(&mut output).await?;
1324
1325 let error_message = format!(
1326 "failed to connect: {}",
1327 String::from_utf8_lossy(&output).trim()
1328 );
1329 Err(anyhow!(error_message))?;
1330 }
1331
1332 Ok(Self {
1333 socket: SshSocket {
1334 connection_options,
1335 socket_path,
1336 },
1337 master_process,
1338 _temp_dir: temp_dir,
1339 })
1340 }
1341
1342 async fn ensure_server_binary(
1343 &self,
1344 delegate: &Arc<dyn SshClientDelegate>,
1345 dst_path: &Path,
1346 platform: SshPlatform,
1347 cx: &mut AsyncAppContext,
1348 ) -> Result<()> {
1349 if std::env::var("ZED_USE_CACHED_REMOTE_SERVER").is_ok() {
1350 if let Ok(installed_version) =
1351 run_cmd(self.socket.ssh_command(dst_path).arg("version")).await
1352 {
1353 log::info!("using cached server binary version {}", installed_version);
1354 return Ok(());
1355 }
1356 }
1357
1358 let mut dst_path_gz = dst_path.to_path_buf();
1359 dst_path_gz.set_extension("gz");
1360
1361 if let Some(parent) = dst_path.parent() {
1362 run_cmd(self.socket.ssh_command("mkdir").arg("-p").arg(parent)).await?;
1363 }
1364
1365 let (src_path, version) = delegate.get_server_binary(platform, cx).await??;
1366
1367 let mut server_binary_exists = false;
1368 if !server_binary_exists && cfg!(not(debug_assertions)) {
1369 if let Ok(installed_version) =
1370 run_cmd(self.socket.ssh_command(dst_path).arg("version")).await
1371 {
1372 if installed_version.trim() == version.to_string() {
1373 server_binary_exists = true;
1374 }
1375 }
1376 }
1377
1378 if server_binary_exists {
1379 log::info!("remote development server already present",);
1380 return Ok(());
1381 }
1382
1383 let src_stat = fs::metadata(&src_path).await?;
1384 let size = src_stat.len();
1385 let server_mode = 0o755;
1386
1387 let t0 = Instant::now();
1388 delegate.set_status(Some("Uploading remote development server"), cx);
1389 log::info!("uploading remote development server ({}kb)", size / 1024);
1390 self.upload_file(&src_path, &dst_path_gz)
1391 .await
1392 .context("failed to upload server binary")?;
1393 log::info!("uploaded remote development server in {:?}", t0.elapsed());
1394
1395 delegate.set_status(Some("Extracting remote development server"), cx);
1396 run_cmd(
1397 self.socket
1398 .ssh_command("gunzip")
1399 .arg("--force")
1400 .arg(&dst_path_gz),
1401 )
1402 .await?;
1403
1404 delegate.set_status(Some("Marking remote development server executable"), cx);
1405 run_cmd(
1406 self.socket
1407 .ssh_command("chmod")
1408 .arg(format!("{:o}", server_mode))
1409 .arg(dst_path),
1410 )
1411 .await?;
1412
1413 Ok(())
1414 }
1415
1416 async fn query_platform(&self) -> Result<SshPlatform> {
1417 let os = run_cmd(self.socket.ssh_command("uname").arg("-s")).await?;
1418 let arch = run_cmd(self.socket.ssh_command("uname").arg("-m")).await?;
1419
1420 let os = match os.trim() {
1421 "Darwin" => "macos",
1422 "Linux" => "linux",
1423 _ => Err(anyhow!("unknown uname os {os:?}"))?,
1424 };
1425 let arch = if arch.starts_with("arm") || arch.starts_with("aarch64") {
1426 "aarch64"
1427 } else if arch.starts_with("x86") || arch.starts_with("i686") {
1428 "x86_64"
1429 } else {
1430 Err(anyhow!("unknown uname architecture {arch:?}"))?
1431 };
1432
1433 Ok(SshPlatform { os, arch })
1434 }
1435
1436 async fn upload_file(&self, src_path: &Path, dest_path: &Path) -> Result<()> {
1437 let mut command = process::Command::new("scp");
1438 let output = self
1439 .socket
1440 .ssh_options(&mut command)
1441 .args(
1442 self.socket
1443 .connection_options
1444 .port
1445 .map(|port| vec!["-P".to_string(), port.to_string()])
1446 .unwrap_or_default(),
1447 )
1448 .arg(src_path)
1449 .arg(format!(
1450 "{}:{}",
1451 self.socket.connection_options.scp_url(),
1452 dest_path.display()
1453 ))
1454 .output()
1455 .await?;
1456
1457 if output.status.success() {
1458 Ok(())
1459 } else {
1460 Err(anyhow!(
1461 "failed to upload file {} -> {}: {}",
1462 src_path.display(),
1463 dest_path.display(),
1464 String::from_utf8_lossy(&output.stderr)
1465 ))
1466 }
1467 }
1468}
1469
1470type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>;
1471
1472pub struct ChannelClient {
1473 next_message_id: AtomicU32,
1474 outgoing_tx: Mutex<mpsc::UnboundedSender<Envelope>>,
1475 buffer: Mutex<VecDeque<Envelope>>,
1476 response_channels: ResponseChannels,
1477 message_handlers: Mutex<ProtoMessageHandlerSet>,
1478 max_received: AtomicU32,
1479 name: &'static str,
1480 task: Mutex<Task<Result<()>>>,
1481}
1482
1483impl ChannelClient {
1484 pub fn new(
1485 incoming_rx: mpsc::UnboundedReceiver<Envelope>,
1486 outgoing_tx: mpsc::UnboundedSender<Envelope>,
1487 cx: &AppContext,
1488 name: &'static str,
1489 ) -> Arc<Self> {
1490 Arc::new_cyclic(|this| Self {
1491 outgoing_tx: Mutex::new(outgoing_tx),
1492 next_message_id: AtomicU32::new(0),
1493 max_received: AtomicU32::new(0),
1494 response_channels: ResponseChannels::default(),
1495 message_handlers: Default::default(),
1496 buffer: Mutex::new(VecDeque::new()),
1497 name,
1498 task: Mutex::new(Self::start_handling_messages(
1499 this.clone(),
1500 incoming_rx,
1501 &cx.to_async(),
1502 )),
1503 })
1504 }
1505
1506 fn start_handling_messages(
1507 this: Weak<Self>,
1508 mut incoming_rx: mpsc::UnboundedReceiver<Envelope>,
1509 cx: &AsyncAppContext,
1510 ) -> Task<Result<()>> {
1511 cx.spawn(|cx| {
1512 async move {
1513 let peer_id = PeerId { owner_id: 0, id: 0 };
1514 while let Some(incoming) = incoming_rx.next().await {
1515 let Some(this) = this.upgrade() else {
1516 return anyhow::Ok(());
1517 };
1518 if let Some(ack_id) = incoming.ack_id {
1519 let mut buffer = this.buffer.lock();
1520 while buffer.front().is_some_and(|msg| msg.id <= ack_id) {
1521 buffer.pop_front();
1522 }
1523 }
1524 if let Some(proto::envelope::Payload::FlushBufferedMessages(_)) =
1525 &incoming.payload
1526 {
1527 log::debug!("{}:ssh message received. name:FlushBufferedMessages", this.name);
1528 {
1529 let buffer = this.buffer.lock();
1530 for envelope in buffer.iter() {
1531 this.outgoing_tx.lock().unbounded_send(envelope.clone()).ok();
1532 }
1533 }
1534 let mut envelope = proto::Ack{}.into_envelope(0, Some(incoming.id), None);
1535 envelope.id = this.next_message_id.fetch_add(1, SeqCst);
1536 this.outgoing_tx.lock().unbounded_send(envelope).ok();
1537 continue;
1538 }
1539
1540 this.max_received.store(incoming.id, SeqCst);
1541
1542 if let Some(request_id) = incoming.responding_to {
1543 let request_id = MessageId(request_id);
1544 let sender = this.response_channels.lock().remove(&request_id);
1545 if let Some(sender) = sender {
1546 let (tx, rx) = oneshot::channel();
1547 if incoming.payload.is_some() {
1548 sender.send((incoming, tx)).ok();
1549 }
1550 rx.await.ok();
1551 }
1552 } else if let Some(envelope) =
1553 build_typed_envelope(peer_id, Instant::now(), incoming)
1554 {
1555 let type_name = envelope.payload_type_name();
1556 if let Some(future) = ProtoMessageHandlerSet::handle_message(
1557 &this.message_handlers,
1558 envelope,
1559 this.clone().into(),
1560 cx.clone(),
1561 ) {
1562 log::debug!("{}:ssh message received. name:{type_name}", this.name);
1563 cx.foreground_executor().spawn(async move {
1564 match future.await {
1565 Ok(_) => {
1566 log::debug!("{}:ssh message handled. name:{type_name}", this.name);
1567 }
1568 Err(error) => {
1569 log::error!(
1570 "{}:error handling message. type:{type_name}, error:{error}", this.name,
1571 );
1572 }
1573 }
1574 }).detach()
1575 } else {
1576 log::error!("{}:unhandled ssh message name:{type_name}", this.name);
1577 }
1578 }
1579 }
1580 anyhow::Ok(())
1581 }
1582 })
1583 }
1584
1585 pub fn reconnect(
1586 self: &Arc<Self>,
1587 incoming_rx: UnboundedReceiver<Envelope>,
1588 outgoing_tx: UnboundedSender<Envelope>,
1589 cx: &AsyncAppContext,
1590 ) {
1591 *self.outgoing_tx.lock() = outgoing_tx;
1592 *self.task.lock() = Self::start_handling_messages(Arc::downgrade(self), incoming_rx, cx);
1593 }
1594
1595 pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Model<E>) {
1596 let id = (TypeId::of::<E>(), remote_id);
1597
1598 let mut message_handlers = self.message_handlers.lock();
1599 if message_handlers
1600 .entities_by_type_and_remote_id
1601 .contains_key(&id)
1602 {
1603 panic!("already subscribed to entity");
1604 }
1605
1606 message_handlers.entities_by_type_and_remote_id.insert(
1607 id,
1608 EntityMessageSubscriber::Entity {
1609 handle: entity.downgrade().into(),
1610 },
1611 );
1612 }
1613
1614 pub fn request<T: RequestMessage>(
1615 &self,
1616 payload: T,
1617 ) -> impl 'static + Future<Output = Result<T::Response>> {
1618 self.request_internal(payload, true)
1619 }
1620
1621 fn request_internal<T: RequestMessage>(
1622 &self,
1623 payload: T,
1624 use_buffer: bool,
1625 ) -> impl 'static + Future<Output = Result<T::Response>> {
1626 log::debug!("ssh request start. name:{}", T::NAME);
1627 let response =
1628 self.request_dynamic(payload.into_envelope(0, None, None), T::NAME, use_buffer);
1629 async move {
1630 let response = response.await?;
1631 log::debug!("ssh request finish. name:{}", T::NAME);
1632 T::Response::from_envelope(response)
1633 .ok_or_else(|| anyhow!("received a response of the wrong type"))
1634 }
1635 }
1636
1637 pub async fn resync(&self, timeout: Duration) -> Result<()> {
1638 smol::future::or(
1639 async {
1640 self.request_internal(proto::FlushBufferedMessages {}, false)
1641 .await?;
1642
1643 for envelope in self.buffer.lock().iter() {
1644 self.outgoing_tx
1645 .lock()
1646 .unbounded_send(envelope.clone())
1647 .ok();
1648 }
1649 Ok(())
1650 },
1651 async {
1652 smol::Timer::after(timeout).await;
1653 Err(anyhow!("Timeout detected"))
1654 },
1655 )
1656 .await
1657 }
1658
1659 pub async fn ping(&self, timeout: Duration) -> Result<()> {
1660 smol::future::or(
1661 async {
1662 self.request(proto::Ping {}).await?;
1663 Ok(())
1664 },
1665 async {
1666 smol::Timer::after(timeout).await;
1667 Err(anyhow!("Timeout detected"))
1668 },
1669 )
1670 .await
1671 }
1672
1673 pub fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {
1674 log::debug!("ssh send name:{}", T::NAME);
1675 self.send_dynamic(payload.into_envelope(0, None, None))
1676 }
1677
1678 fn request_dynamic(
1679 &self,
1680 mut envelope: proto::Envelope,
1681 type_name: &'static str,
1682 use_buffer: bool,
1683 ) -> impl 'static + Future<Output = Result<proto::Envelope>> {
1684 envelope.id = self.next_message_id.fetch_add(1, SeqCst);
1685 let (tx, rx) = oneshot::channel();
1686 let mut response_channels_lock = self.response_channels.lock();
1687 response_channels_lock.insert(MessageId(envelope.id), tx);
1688 drop(response_channels_lock);
1689
1690 let result = if use_buffer {
1691 self.send_buffered(envelope)
1692 } else {
1693 self.send_unbuffered(envelope)
1694 };
1695 async move {
1696 if let Err(error) = &result {
1697 log::error!("failed to send message: {}", error);
1698 return Err(anyhow!("failed to send message: {}", error));
1699 }
1700
1701 let response = rx.await.context("connection lost")?.0;
1702 if let Some(proto::envelope::Payload::Error(error)) = &response.payload {
1703 return Err(RpcError::from_proto(error, type_name));
1704 }
1705 Ok(response)
1706 }
1707 }
1708
1709 pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {
1710 envelope.id = self.next_message_id.fetch_add(1, SeqCst);
1711 self.send_buffered(envelope)
1712 }
1713
1714 fn send_buffered(&self, mut envelope: proto::Envelope) -> Result<()> {
1715 envelope.ack_id = Some(self.max_received.load(SeqCst));
1716 self.buffer.lock().push_back(envelope.clone());
1717 // ignore errors on send (happen while we're reconnecting)
1718 // assume that the global "disconnected" overlay is sufficient.
1719 self.outgoing_tx.lock().unbounded_send(envelope).ok();
1720 Ok(())
1721 }
1722
1723 fn send_unbuffered(&self, mut envelope: proto::Envelope) -> Result<()> {
1724 envelope.ack_id = Some(self.max_received.load(SeqCst));
1725 self.outgoing_tx.lock().unbounded_send(envelope).ok();
1726 Ok(())
1727 }
1728}
1729
1730impl ProtoClient for ChannelClient {
1731 fn request(
1732 &self,
1733 envelope: proto::Envelope,
1734 request_type: &'static str,
1735 ) -> BoxFuture<'static, Result<proto::Envelope>> {
1736 self.request_dynamic(envelope, request_type, true).boxed()
1737 }
1738
1739 fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
1740 self.send_dynamic(envelope)
1741 }
1742
1743 fn send_response(&self, envelope: Envelope, _message_type: &'static str) -> anyhow::Result<()> {
1744 self.send_dynamic(envelope)
1745 }
1746
1747 fn message_handler_set(&self) -> &Mutex<ProtoMessageHandlerSet> {
1748 &self.message_handlers
1749 }
1750
1751 fn is_via_collab(&self) -> bool {
1752 false
1753 }
1754}
1755
1756#[cfg(any(test, feature = "test-support"))]
1757mod fake {
1758 use std::{path::PathBuf, sync::Arc};
1759
1760 use anyhow::Result;
1761 use async_trait::async_trait;
1762 use futures::{
1763 channel::{
1764 mpsc::{self, Sender},
1765 oneshot,
1766 },
1767 select_biased, FutureExt, SinkExt, StreamExt,
1768 };
1769 use gpui::{AsyncAppContext, BorrowAppContext, Global, SemanticVersion, Task};
1770 use rpc::proto::Envelope;
1771
1772 use super::{
1773 ChannelClient, SshClientDelegate, SshConnectionOptions, SshPlatform, SshRemoteProcess,
1774 };
1775
1776 pub(super) struct SshRemoteConnection {
1777 connection_options: SshConnectionOptions,
1778 }
1779
1780 impl SshRemoteConnection {
1781 pub(super) fn new(
1782 connection_options: &SshConnectionOptions,
1783 ) -> Option<Box<dyn SshRemoteProcess>> {
1784 if connection_options.host == "<fake>" {
1785 return Some(Box::new(Self {
1786 connection_options: connection_options.clone(),
1787 }));
1788 }
1789 return None;
1790 }
1791 pub(super) async fn multiplex(
1792 connection_options: SshConnectionOptions,
1793 mut client_incoming_tx: mpsc::UnboundedSender<Envelope>,
1794 mut client_outgoing_rx: mpsc::UnboundedReceiver<Envelope>,
1795 mut connection_activity_tx: Sender<()>,
1796 cx: &mut AsyncAppContext,
1797 ) -> Task<Result<i32>> {
1798 let (mut server_incoming_tx, server_incoming_rx) = mpsc::unbounded::<Envelope>();
1799 let (server_outgoing_tx, mut server_outgoing_rx) = mpsc::unbounded::<Envelope>();
1800
1801 let (channel, server_cx) = cx
1802 .update(|cx| {
1803 cx.update_global(|conns: &mut ServerConnections, _| {
1804 conns.get(connection_options.port.unwrap())
1805 })
1806 })
1807 .unwrap();
1808 channel.reconnect(server_incoming_rx, server_outgoing_tx, &server_cx);
1809
1810 // send to proxy_tx to get to the server.
1811 // receive from
1812
1813 cx.background_executor().spawn(async move {
1814 loop {
1815 select_biased! {
1816 server_to_client = server_outgoing_rx.next().fuse() => {
1817 let Some(server_to_client) = server_to_client else {
1818 return Ok(1)
1819 };
1820 connection_activity_tx.try_send(()).ok();
1821 client_incoming_tx.send(server_to_client).await.ok();
1822 }
1823 client_to_server = client_outgoing_rx.next().fuse() => {
1824 let Some(client_to_server) = client_to_server else {
1825 return Ok(1)
1826 };
1827 server_incoming_tx.send(client_to_server).await.ok();
1828 }
1829 }
1830 }
1831 })
1832 }
1833 }
1834
1835 #[async_trait]
1836 impl SshRemoteProcess for SshRemoteConnection {
1837 async fn kill(&mut self) -> Result<()> {
1838 Ok(())
1839 }
1840
1841 fn ssh_args(&self) -> Vec<String> {
1842 Vec::new()
1843 }
1844
1845 fn connection_options(&self) -> SshConnectionOptions {
1846 self.connection_options.clone()
1847 }
1848 }
1849
1850 #[derive(Default)]
1851 pub(super) struct ServerConnections(Vec<(Arc<ChannelClient>, AsyncAppContext)>);
1852 impl Global for ServerConnections {}
1853
1854 impl ServerConnections {
1855 pub(super) fn push(&mut self, server: Arc<ChannelClient>, cx: AsyncAppContext) -> u16 {
1856 self.0.push((server.clone(), cx));
1857 self.0.len() as u16 - 1
1858 }
1859
1860 pub(super) fn get(&mut self, port: u16) -> (Arc<ChannelClient>, AsyncAppContext) {
1861 self.0
1862 .get(port as usize)
1863 .expect("no fake server for port")
1864 .clone()
1865 }
1866 }
1867
1868 pub(super) struct Delegate;
1869
1870 impl SshClientDelegate for Delegate {
1871 fn ask_password(
1872 &self,
1873 _: String,
1874 _: &mut AsyncAppContext,
1875 ) -> oneshot::Receiver<Result<String>> {
1876 unreachable!()
1877 }
1878 fn remote_server_binary_path(
1879 &self,
1880 _: SshPlatform,
1881 _: &mut AsyncAppContext,
1882 ) -> Result<PathBuf> {
1883 unreachable!()
1884 }
1885 fn get_server_binary(
1886 &self,
1887 _: SshPlatform,
1888 _: &mut AsyncAppContext,
1889 ) -> oneshot::Receiver<Result<(PathBuf, SemanticVersion)>> {
1890 unreachable!()
1891 }
1892 fn set_status(&self, _: Option<&str>, _: &mut AsyncAppContext) {
1893 unreachable!()
1894 }
1895 }
1896}