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