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