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};
7use anyhow::{anyhow, Context as _, Result};
8use collections::HashMap;
9use futures::{
10 channel::{
11 mpsc::{self, UnboundedReceiver, UnboundedSender},
12 oneshot,
13 },
14 future::BoxFuture,
15 select_biased, AsyncReadExt as _, AsyncWriteExt as _, Future, FutureExt as _, SinkExt,
16 StreamExt as _,
17};
18use gpui::{AppContext, AsyncAppContext, Model, SemanticVersion, Task};
19use parking_lot::Mutex;
20use rpc::{
21 proto::{self, build_typed_envelope, Envelope, EnvelopedMessage, PeerId, RequestMessage},
22 AnyProtoClient, EntityMessageSubscriber, ProtoClient, ProtoMessageHandlerSet, RpcError,
23};
24use smol::{
25 fs,
26 process::{self, Child, Stdio},
27};
28use std::{
29 any::TypeId,
30 ffi::OsStr,
31 path::{Path, PathBuf},
32 sync::{
33 atomic::{AtomicU32, Ordering::SeqCst},
34 Arc, Weak,
35 },
36 time::Instant,
37};
38use tempfile::TempDir;
39use util::maybe;
40
41#[derive(
42 Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Copy, serde::Serialize, serde::Deserialize,
43)]
44pub struct SshProjectId(pub u64);
45
46#[derive(Clone)]
47pub struct SshSocket {
48 connection_options: SshConnectionOptions,
49 socket_path: PathBuf,
50}
51
52#[derive(Debug, Default, Clone, PartialEq, Eq)]
53pub struct SshConnectionOptions {
54 pub host: String,
55 pub username: Option<String>,
56 pub port: Option<u16>,
57 pub password: Option<String>,
58}
59
60impl SshConnectionOptions {
61 pub fn ssh_url(&self) -> String {
62 let mut result = String::from("ssh://");
63 if let Some(username) = &self.username {
64 result.push_str(username);
65 result.push('@');
66 }
67 result.push_str(&self.host);
68 if let Some(port) = self.port {
69 result.push(':');
70 result.push_str(&port.to_string());
71 }
72 result
73 }
74
75 fn scp_url(&self) -> String {
76 if let Some(username) = &self.username {
77 format!("{}@{}", username, self.host)
78 } else {
79 self.host.clone()
80 }
81 }
82
83 pub fn connection_string(&self) -> String {
84 let host = if let Some(username) = &self.username {
85 format!("{}@{}", username, self.host)
86 } else {
87 self.host.clone()
88 };
89 if let Some(port) = &self.port {
90 format!("{}:{}", host, port)
91 } else {
92 host
93 }
94 }
95}
96
97#[derive(Copy, Clone, Debug)]
98pub struct SshPlatform {
99 pub os: &'static str,
100 pub arch: &'static str,
101}
102
103pub trait SshClientDelegate: Send + Sync {
104 fn ask_password(
105 &self,
106 prompt: String,
107 cx: &mut AsyncAppContext,
108 ) -> oneshot::Receiver<Result<String>>;
109 fn remote_server_binary_path(&self, cx: &mut AsyncAppContext) -> Result<PathBuf>;
110 fn get_server_binary(
111 &self,
112 platform: SshPlatform,
113 cx: &mut AsyncAppContext,
114 ) -> oneshot::Receiver<Result<(PathBuf, SemanticVersion)>>;
115 fn set_status(&self, status: Option<&str>, cx: &mut AsyncAppContext);
116 fn set_error(&self, error_message: String, cx: &mut AsyncAppContext);
117}
118
119impl SshSocket {
120 fn ssh_command<S: AsRef<OsStr>>(&self, program: S) -> process::Command {
121 let mut command = process::Command::new("ssh");
122 self.ssh_options(&mut command)
123 .arg(self.connection_options.ssh_url())
124 .arg(program);
125 command
126 }
127
128 fn ssh_options<'a>(&self, command: &'a mut process::Command) -> &'a mut process::Command {
129 command
130 .stdin(Stdio::piped())
131 .stdout(Stdio::piped())
132 .stderr(Stdio::piped())
133 .args(["-o", "ControlMaster=no", "-o"])
134 .arg(format!("ControlPath={}", self.socket_path.display()))
135 }
136
137 fn ssh_args(&self) -> Vec<String> {
138 vec![
139 "-o".to_string(),
140 "ControlMaster=no".to_string(),
141 "-o".to_string(),
142 format!("ControlPath={}", self.socket_path.display()),
143 self.connection_options.ssh_url(),
144 ]
145 }
146}
147
148async fn run_cmd(command: &mut process::Command) -> Result<String> {
149 let output = command.output().await?;
150 if output.status.success() {
151 Ok(String::from_utf8_lossy(&output.stdout).to_string())
152 } else {
153 Err(anyhow!(
154 "failed to run command: {}",
155 String::from_utf8_lossy(&output.stderr)
156 ))
157 }
158}
159#[cfg(unix)]
160async fn read_with_timeout(
161 stdout: &mut process::ChildStdout,
162 timeout: std::time::Duration,
163 output: &mut Vec<u8>,
164) -> Result<(), std::io::Error> {
165 smol::future::or(
166 async {
167 stdout.read_to_end(output).await?;
168 Ok::<_, std::io::Error>(())
169 },
170 async {
171 smol::Timer::after(timeout).await;
172
173 Err(std::io::Error::new(
174 std::io::ErrorKind::TimedOut,
175 "Read operation timed out",
176 ))
177 },
178 )
179 .await
180}
181
182struct ChannelForwarder {
183 quit_tx: UnboundedSender<()>,
184 forwarding_task: Task<(UnboundedSender<Envelope>, UnboundedReceiver<Envelope>)>,
185}
186
187impl ChannelForwarder {
188 fn new(
189 mut incoming_tx: UnboundedSender<Envelope>,
190 mut outgoing_rx: UnboundedReceiver<Envelope>,
191 cx: &mut AsyncAppContext,
192 ) -> (Self, UnboundedSender<Envelope>, UnboundedReceiver<Envelope>) {
193 let (quit_tx, mut quit_rx) = mpsc::unbounded::<()>();
194
195 let (proxy_incoming_tx, mut proxy_incoming_rx) = mpsc::unbounded::<Envelope>();
196 let (mut proxy_outgoing_tx, proxy_outgoing_rx) = mpsc::unbounded::<Envelope>();
197
198 let forwarding_task = cx.background_executor().spawn(async move {
199 loop {
200 select_biased! {
201 _ = quit_rx.next().fuse() => {
202 break;
203 },
204 incoming_envelope = proxy_incoming_rx.next().fuse() => {
205 if let Some(envelope) = incoming_envelope {
206 if incoming_tx.send(envelope).await.is_err() {
207 break;
208 }
209 } else {
210 break;
211 }
212 }
213 outgoing_envelope = outgoing_rx.next().fuse() => {
214 if let Some(envelope) = outgoing_envelope {
215 if proxy_outgoing_tx.send(envelope).await.is_err() {
216 break;
217 }
218 } else {
219 break;
220 }
221 }
222 }
223 }
224
225 (incoming_tx, outgoing_rx)
226 });
227
228 (
229 Self {
230 forwarding_task,
231 quit_tx,
232 },
233 proxy_incoming_tx,
234 proxy_outgoing_rx,
235 )
236 }
237
238 async fn into_channels(mut self) -> (UnboundedSender<Envelope>, UnboundedReceiver<Envelope>) {
239 let _ = self.quit_tx.send(()).await;
240 self.forwarding_task.await
241 }
242}
243
244struct SshRemoteClientState {
245 ssh_connection: SshRemoteConnection,
246 delegate: Arc<dyn SshClientDelegate>,
247 forwarder: ChannelForwarder,
248 multiplex_task: Task<Result<()>>,
249}
250
251pub struct SshRemoteClient {
252 client: Arc<ChannelClient>,
253 inner_state: Mutex<Option<SshRemoteClientState>>,
254 connection_options: SshConnectionOptions,
255}
256
257impl SshRemoteClient {
258 pub async fn new(
259 connection_options: SshConnectionOptions,
260 delegate: Arc<dyn SshClientDelegate>,
261 cx: &mut AsyncAppContext,
262 ) -> Result<Arc<Self>> {
263 let (outgoing_tx, outgoing_rx) = mpsc::unbounded::<Envelope>();
264 let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();
265
266 let client = cx.update(|cx| ChannelClient::new(incoming_rx, outgoing_tx, cx))?;
267 let this = Arc::new(Self {
268 client,
269 inner_state: Mutex::new(None),
270 connection_options: connection_options.clone(),
271 });
272
273 let inner_state = {
274 let (proxy, proxy_incoming_tx, proxy_outgoing_rx) =
275 ChannelForwarder::new(incoming_tx, outgoing_rx, cx);
276
277 let (ssh_connection, ssh_process) =
278 Self::establish_connection(connection_options, delegate.clone(), cx).await?;
279
280 let multiplex_task = Self::multiplex(
281 Arc::downgrade(&this),
282 ssh_process,
283 proxy_incoming_tx,
284 proxy_outgoing_rx,
285 cx,
286 );
287
288 SshRemoteClientState {
289 ssh_connection,
290 delegate,
291 forwarder: proxy,
292 multiplex_task,
293 }
294 };
295
296 this.inner_state.lock().replace(inner_state);
297
298 Ok(this)
299 }
300
301 fn reconnect(this: Arc<Self>, cx: &mut AsyncAppContext) -> Result<()> {
302 let Some(state) = this.inner_state.lock().take() else {
303 return Err(anyhow!("reconnect is already in progress"));
304 };
305
306 let SshRemoteClientState {
307 mut ssh_connection,
308 delegate,
309 forwarder: proxy,
310 multiplex_task,
311 } = state;
312 drop(multiplex_task);
313
314 cx.spawn(|mut cx| async move {
315 let (incoming_tx, outgoing_rx) = proxy.into_channels().await;
316
317 ssh_connection.master_process.kill()?;
318 ssh_connection
319 .master_process
320 .status()
321 .await
322 .context("Failed to kill ssh process")?;
323
324 let connection_options = ssh_connection.socket.connection_options.clone();
325
326 let (ssh_connection, ssh_process) =
327 Self::establish_connection(connection_options, delegate.clone(), &mut cx).await?;
328
329 let (proxy, proxy_incoming_tx, proxy_outgoing_rx) =
330 ChannelForwarder::new(incoming_tx, outgoing_rx, &mut cx);
331
332 let inner_state = SshRemoteClientState {
333 ssh_connection,
334 delegate,
335 forwarder: proxy,
336 multiplex_task: Self::multiplex(
337 Arc::downgrade(&this),
338 ssh_process,
339 proxy_incoming_tx,
340 proxy_outgoing_rx,
341 &mut cx,
342 ),
343 };
344 this.inner_state.lock().replace(inner_state);
345
346 anyhow::Ok(())
347 })
348 .detach();
349
350 anyhow::Ok(())
351 }
352
353 fn multiplex(
354 this: Weak<Self>,
355 mut ssh_process: Child,
356 incoming_tx: UnboundedSender<Envelope>,
357 mut outgoing_rx: UnboundedReceiver<Envelope>,
358 cx: &mut AsyncAppContext,
359 ) -> Task<Result<()>> {
360 let mut child_stderr = ssh_process.stderr.take().unwrap();
361 let mut child_stdout = ssh_process.stdout.take().unwrap();
362 let mut child_stdin = ssh_process.stdin.take().unwrap();
363
364 let io_task = cx.background_executor().spawn(async move {
365 let mut stdin_buffer = Vec::new();
366 let mut stdout_buffer = Vec::new();
367 let mut stderr_buffer = Vec::new();
368 let mut stderr_offset = 0;
369
370 loop {
371 stdout_buffer.resize(MESSAGE_LEN_SIZE, 0);
372 stderr_buffer.resize(stderr_offset + 1024, 0);
373
374 select_biased! {
375 outgoing = outgoing_rx.next().fuse() => {
376 let Some(outgoing) = outgoing else {
377 return anyhow::Ok(());
378 };
379
380 write_message(&mut child_stdin, &mut stdin_buffer, outgoing).await?;
381 }
382
383 result = child_stdout.read(&mut stdout_buffer).fuse() => {
384 match result {
385 Ok(0) => {
386 child_stdin.close().await?;
387 outgoing_rx.close();
388 let status = ssh_process.status().await?;
389 if !status.success() {
390 log::error!("ssh process exited with status: {status:?}");
391 return Err(anyhow!("ssh process exited with non-zero status code: {:?}", status.code()));
392 }
393 return Ok(());
394 }
395 Ok(len) => {
396 if len < stdout_buffer.len() {
397 child_stdout.read_exact(&mut stdout_buffer[len..]).await?;
398 }
399
400 let message_len = message_len_from_buffer(&stdout_buffer);
401 match read_message_with_len(&mut child_stdout, &mut stdout_buffer, message_len).await {
402 Ok(envelope) => {
403 incoming_tx.unbounded_send(envelope).ok();
404 }
405 Err(error) => {
406 log::error!("error decoding message {error:?}");
407 }
408 }
409 }
410 Err(error) => {
411 Err(anyhow!("error reading stdout: {error:?}"))?;
412 }
413 }
414 }
415
416 result = child_stderr.read(&mut stderr_buffer[stderr_offset..]).fuse() => {
417 match result {
418 Ok(len) => {
419 stderr_offset += len;
420 let mut start_ix = 0;
421 while let Some(ix) = stderr_buffer[start_ix..stderr_offset].iter().position(|b| b == &b'\n') {
422 let line_ix = start_ix + ix;
423 let content = &stderr_buffer[start_ix..line_ix];
424 start_ix = line_ix + 1;
425 if let Ok(mut record) = serde_json::from_slice::<LogRecord>(content) {
426 record.message = format!("(remote) {}", record.message);
427 record.log(log::logger())
428 } else {
429 eprintln!("(remote) {}", String::from_utf8_lossy(content));
430 }
431 }
432 stderr_buffer.drain(0..start_ix);
433 stderr_offset -= start_ix;
434 }
435 Err(error) => {
436 Err(anyhow!("error reading stderr: {error:?}"))?;
437 }
438 }
439 }
440 }
441 }
442 });
443
444 cx.spawn(|mut cx| async move {
445 let result = io_task.await;
446
447 if let Err(error) = result {
448 log::warn!("ssh io task died with error: {:?}. reconnecting...", error);
449 if let Some(this) = this.upgrade() {
450 Self::reconnect(this, &mut cx).ok();
451 }
452 }
453
454 Ok(())
455 })
456 }
457
458 async fn establish_connection(
459 connection_options: SshConnectionOptions,
460 delegate: Arc<dyn SshClientDelegate>,
461 cx: &mut AsyncAppContext,
462 ) -> Result<(SshRemoteConnection, Child)> {
463 let ssh_connection =
464 SshRemoteConnection::new(connection_options, delegate.clone(), cx).await?;
465
466 let platform = ssh_connection.query_platform().await?;
467 let (local_binary_path, version) = delegate.get_server_binary(platform, cx).await??;
468 let remote_binary_path = delegate.remote_server_binary_path(cx)?;
469 ssh_connection
470 .ensure_server_binary(
471 &delegate,
472 &local_binary_path,
473 &remote_binary_path,
474 version,
475 cx,
476 )
477 .await?;
478
479 let socket = ssh_connection.socket.clone();
480 run_cmd(socket.ssh_command(&remote_binary_path).arg("version")).await?;
481
482 let ssh_process = socket
483 .ssh_command(format!(
484 "RUST_LOG={} RUST_BACKTRACE={} {:?} run",
485 std::env::var("RUST_LOG").unwrap_or_default(),
486 std::env::var("RUST_BACKTRACE").unwrap_or_default(),
487 remote_binary_path,
488 ))
489 .spawn()
490 .context("failed to spawn remote server")?;
491
492 Ok((ssh_connection, ssh_process))
493 }
494
495 pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Model<E>) {
496 self.client.subscribe_to_entity(remote_id, entity);
497 }
498
499 pub fn ssh_args(&self) -> Option<Vec<String>> {
500 let state = self.inner_state.lock();
501 state
502 .as_ref()
503 .map(|state| state.ssh_connection.socket.ssh_args())
504 }
505
506 pub fn to_proto_client(&self) -> AnyProtoClient {
507 self.client.clone().into()
508 }
509
510 pub fn connection_string(&self) -> String {
511 self.connection_options.connection_string()
512 }
513
514 pub fn is_reconnect_underway(&self) -> bool {
515 maybe!({ Some(self.inner_state.try_lock()?.is_none()) }).unwrap_or_default()
516 }
517 #[cfg(any(test, feature = "test-support"))]
518 pub fn fake(
519 client_cx: &mut gpui::TestAppContext,
520 server_cx: &mut gpui::TestAppContext,
521 ) -> (Arc<Self>, Arc<ChannelClient>) {
522 let (server_to_client_tx, server_to_client_rx) = mpsc::unbounded();
523 let (client_to_server_tx, client_to_server_rx) = mpsc::unbounded();
524
525 (
526 client_cx.update(|cx| {
527 let client = ChannelClient::new(server_to_client_rx, client_to_server_tx, cx);
528 Arc::new(Self {
529 client,
530 inner_state: Mutex::new(None),
531 connection_options: SshConnectionOptions::default(),
532 })
533 }),
534 server_cx.update(|cx| ChannelClient::new(client_to_server_rx, server_to_client_tx, cx)),
535 )
536 }
537}
538
539impl From<SshRemoteClient> for AnyProtoClient {
540 fn from(client: SshRemoteClient) -> Self {
541 AnyProtoClient::new(client.client.clone())
542 }
543}
544
545struct SshRemoteConnection {
546 socket: SshSocket,
547 master_process: process::Child,
548 _temp_dir: TempDir,
549}
550
551impl Drop for SshRemoteConnection {
552 fn drop(&mut self) {
553 if let Err(error) = self.master_process.kill() {
554 log::error!("failed to kill SSH master process: {}", error);
555 }
556 }
557}
558
559impl SshRemoteConnection {
560 #[cfg(not(unix))]
561 async fn new(
562 _connection_options: SshConnectionOptions,
563 _delegate: Arc<dyn SshClientDelegate>,
564 _cx: &mut AsyncAppContext,
565 ) -> Result<Self> {
566 Err(anyhow!("ssh is not supported on this platform"))
567 }
568
569 #[cfg(unix)]
570 async fn new(
571 connection_options: SshConnectionOptions,
572 delegate: Arc<dyn SshClientDelegate>,
573 cx: &mut AsyncAppContext,
574 ) -> Result<Self> {
575 use futures::{io::BufReader, AsyncBufReadExt as _};
576 use smol::{fs::unix::PermissionsExt as _, net::unix::UnixListener};
577 use util::ResultExt as _;
578
579 delegate.set_status(Some("connecting"), cx);
580
581 let url = connection_options.ssh_url();
582 let temp_dir = tempfile::Builder::new()
583 .prefix("zed-ssh-session")
584 .tempdir()?;
585
586 // Create a domain socket listener to handle requests from the askpass program.
587 let askpass_socket = temp_dir.path().join("askpass.sock");
588 let listener =
589 UnixListener::bind(&askpass_socket).context("failed to create askpass socket")?;
590
591 let askpass_task = cx.spawn({
592 let delegate = delegate.clone();
593 |mut cx| async move {
594 while let Ok((mut stream, _)) = listener.accept().await {
595 let mut buffer = Vec::new();
596 let mut reader = BufReader::new(&mut stream);
597 if reader.read_until(b'\0', &mut buffer).await.is_err() {
598 buffer.clear();
599 }
600 let password_prompt = String::from_utf8_lossy(&buffer);
601 if let Some(password) = delegate
602 .ask_password(password_prompt.to_string(), &mut cx)
603 .await
604 .context("failed to get ssh password")
605 .and_then(|p| p)
606 .log_err()
607 {
608 stream.write_all(password.as_bytes()).await.log_err();
609 }
610 }
611 }
612 });
613
614 // Create an askpass script that communicates back to this process.
615 let askpass_script = format!(
616 "{shebang}\n{print_args} | nc -U {askpass_socket} 2> /dev/null \n",
617 askpass_socket = askpass_socket.display(),
618 print_args = "printf '%s\\0' \"$@\"",
619 shebang = "#!/bin/sh",
620 );
621 let askpass_script_path = temp_dir.path().join("askpass.sh");
622 fs::write(&askpass_script_path, askpass_script).await?;
623 fs::set_permissions(&askpass_script_path, std::fs::Permissions::from_mode(0o755)).await?;
624
625 // Start the master SSH process, which does not do anything except for establish
626 // the connection and keep it open, allowing other ssh commands to reuse it
627 // via a control socket.
628 let socket_path = temp_dir.path().join("ssh.sock");
629 let mut master_process = process::Command::new("ssh")
630 .stdin(Stdio::null())
631 .stdout(Stdio::piped())
632 .stderr(Stdio::piped())
633 .env("SSH_ASKPASS_REQUIRE", "force")
634 .env("SSH_ASKPASS", &askpass_script_path)
635 .args(["-N", "-o", "ControlMaster=yes", "-o"])
636 .arg(format!("ControlPath={}", socket_path.display()))
637 .arg(&url)
638 .spawn()?;
639
640 // Wait for this ssh process to close its stdout, indicating that authentication
641 // has completed.
642 let stdout = master_process.stdout.as_mut().unwrap();
643 let mut output = Vec::new();
644 let connection_timeout = std::time::Duration::from_secs(10);
645 let result = read_with_timeout(stdout, connection_timeout, &mut output).await;
646 if let Err(e) = result {
647 let error_message = if e.kind() == std::io::ErrorKind::TimedOut {
648 format!(
649 "Failed to connect to host. Timed out after {:?}.",
650 connection_timeout
651 )
652 } else {
653 format!("Failed to connect to host: {}.", e)
654 };
655
656 delegate.set_error(error_message, cx);
657 return Err(e.into());
658 }
659
660 drop(askpass_task);
661
662 if master_process.try_status()?.is_some() {
663 output.clear();
664 let mut stderr = master_process.stderr.take().unwrap();
665 stderr.read_to_end(&mut output).await?;
666 Err(anyhow!(
667 "failed to connect: {}",
668 String::from_utf8_lossy(&output)
669 ))?;
670 }
671
672 Ok(Self {
673 socket: SshSocket {
674 connection_options,
675 socket_path,
676 },
677 master_process,
678 _temp_dir: temp_dir,
679 })
680 }
681
682 async fn ensure_server_binary(
683 &self,
684 delegate: &Arc<dyn SshClientDelegate>,
685 src_path: &Path,
686 dst_path: &Path,
687 version: SemanticVersion,
688 cx: &mut AsyncAppContext,
689 ) -> Result<()> {
690 let mut dst_path_gz = dst_path.to_path_buf();
691 dst_path_gz.set_extension("gz");
692
693 if let Some(parent) = dst_path.parent() {
694 run_cmd(self.socket.ssh_command("mkdir").arg("-p").arg(parent)).await?;
695 }
696
697 let mut server_binary_exists = false;
698 if cfg!(not(debug_assertions)) {
699 if let Ok(installed_version) =
700 run_cmd(self.socket.ssh_command(dst_path).arg("version")).await
701 {
702 if installed_version.trim() == version.to_string() {
703 server_binary_exists = true;
704 }
705 }
706 }
707
708 if server_binary_exists {
709 log::info!("remote development server already present",);
710 return Ok(());
711 }
712
713 let src_stat = fs::metadata(src_path).await?;
714 let size = src_stat.len();
715 let server_mode = 0o755;
716
717 let t0 = Instant::now();
718 delegate.set_status(Some("uploading remote development server"), cx);
719 log::info!("uploading remote development server ({}kb)", size / 1024);
720 self.upload_file(src_path, &dst_path_gz)
721 .await
722 .context("failed to upload server binary")?;
723 log::info!("uploaded remote development server in {:?}", t0.elapsed());
724
725 delegate.set_status(Some("extracting remote development server"), cx);
726 run_cmd(
727 self.socket
728 .ssh_command("gunzip")
729 .arg("--force")
730 .arg(&dst_path_gz),
731 )
732 .await?;
733
734 delegate.set_status(Some("unzipping remote development server"), cx);
735 run_cmd(
736 self.socket
737 .ssh_command("chmod")
738 .arg(format!("{:o}", server_mode))
739 .arg(dst_path),
740 )
741 .await?;
742
743 Ok(())
744 }
745
746 async fn query_platform(&self) -> Result<SshPlatform> {
747 let os = run_cmd(self.socket.ssh_command("uname").arg("-s")).await?;
748 let arch = run_cmd(self.socket.ssh_command("uname").arg("-m")).await?;
749
750 let os = match os.trim() {
751 "Darwin" => "macos",
752 "Linux" => "linux",
753 _ => Err(anyhow!("unknown uname os {os:?}"))?,
754 };
755 let arch = if arch.starts_with("arm") || arch.starts_with("aarch64") {
756 "aarch64"
757 } else if arch.starts_with("x86") || arch.starts_with("i686") {
758 "x86_64"
759 } else {
760 Err(anyhow!("unknown uname architecture {arch:?}"))?
761 };
762
763 Ok(SshPlatform { os, arch })
764 }
765
766 async fn upload_file(&self, src_path: &Path, dest_path: &Path) -> Result<()> {
767 let mut command = process::Command::new("scp");
768 let output = self
769 .socket
770 .ssh_options(&mut command)
771 .args(
772 self.socket
773 .connection_options
774 .port
775 .map(|port| vec!["-P".to_string(), port.to_string()])
776 .unwrap_or_default(),
777 )
778 .arg(src_path)
779 .arg(format!(
780 "{}:{}",
781 self.socket.connection_options.scp_url(),
782 dest_path.display()
783 ))
784 .output()
785 .await?;
786
787 if output.status.success() {
788 Ok(())
789 } else {
790 Err(anyhow!(
791 "failed to upload file {} -> {}: {}",
792 src_path.display(),
793 dest_path.display(),
794 String::from_utf8_lossy(&output.stderr)
795 ))
796 }
797 }
798}
799
800type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>;
801
802pub struct ChannelClient {
803 next_message_id: AtomicU32,
804 outgoing_tx: mpsc::UnboundedSender<Envelope>,
805 response_channels: ResponseChannels, // Lock
806 message_handlers: Mutex<ProtoMessageHandlerSet>, // Lock
807}
808
809impl ChannelClient {
810 pub fn new(
811 incoming_rx: mpsc::UnboundedReceiver<Envelope>,
812 outgoing_tx: mpsc::UnboundedSender<Envelope>,
813 cx: &AppContext,
814 ) -> Arc<Self> {
815 let this = Arc::new(Self {
816 outgoing_tx,
817 next_message_id: AtomicU32::new(0),
818 response_channels: ResponseChannels::default(),
819 message_handlers: Default::default(),
820 });
821
822 Self::start_handling_messages(this.clone(), incoming_rx, cx);
823
824 this
825 }
826
827 fn start_handling_messages(
828 this: Arc<Self>,
829 mut incoming_rx: mpsc::UnboundedReceiver<Envelope>,
830 cx: &AppContext,
831 ) {
832 cx.spawn(|cx| {
833 let this = Arc::downgrade(&this);
834 async move {
835 let peer_id = PeerId { owner_id: 0, id: 0 };
836 while let Some(incoming) = incoming_rx.next().await {
837 let Some(this) = this.upgrade() else {
838 return anyhow::Ok(());
839 };
840
841 if let Some(request_id) = incoming.responding_to {
842 let request_id = MessageId(request_id);
843 let sender = this.response_channels.lock().remove(&request_id);
844 if let Some(sender) = sender {
845 let (tx, rx) = oneshot::channel();
846 if incoming.payload.is_some() {
847 sender.send((incoming, tx)).ok();
848 }
849 rx.await.ok();
850 }
851 } else if let Some(envelope) =
852 build_typed_envelope(peer_id, Instant::now(), incoming)
853 {
854 let type_name = envelope.payload_type_name();
855 if let Some(future) = ProtoMessageHandlerSet::handle_message(
856 &this.message_handlers,
857 envelope,
858 this.clone().into(),
859 cx.clone(),
860 ) {
861 log::debug!("ssh message received. name:{type_name}");
862 match future.await {
863 Ok(_) => {
864 log::debug!("ssh message handled. name:{type_name}");
865 }
866 Err(error) => {
867 log::error!(
868 "error handling message. type:{type_name}, error:{error}",
869 );
870 }
871 }
872 } else {
873 log::error!("unhandled ssh message name:{type_name}");
874 }
875 }
876 }
877 anyhow::Ok(())
878 }
879 })
880 .detach();
881 }
882
883 pub fn subscribe_to_entity<E: 'static>(&self, remote_id: u64, entity: &Model<E>) {
884 let id = (TypeId::of::<E>(), remote_id);
885
886 let mut message_handlers = self.message_handlers.lock();
887 if message_handlers
888 .entities_by_type_and_remote_id
889 .contains_key(&id)
890 {
891 panic!("already subscribed to entity");
892 }
893
894 message_handlers.entities_by_type_and_remote_id.insert(
895 id,
896 EntityMessageSubscriber::Entity {
897 handle: entity.downgrade().into(),
898 },
899 );
900 }
901
902 pub fn request<T: RequestMessage>(
903 &self,
904 payload: T,
905 ) -> impl 'static + Future<Output = Result<T::Response>> {
906 log::debug!("ssh request start. name:{}", T::NAME);
907 let response = self.request_dynamic(payload.into_envelope(0, None, None), T::NAME);
908 async move {
909 let response = response.await?;
910 log::debug!("ssh request finish. name:{}", T::NAME);
911 T::Response::from_envelope(response)
912 .ok_or_else(|| anyhow!("received a response of the wrong type"))
913 }
914 }
915
916 pub fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {
917 log::debug!("ssh send name:{}", T::NAME);
918 self.send_dynamic(payload.into_envelope(0, None, None))
919 }
920
921 pub fn request_dynamic(
922 &self,
923 mut envelope: proto::Envelope,
924 type_name: &'static str,
925 ) -> impl 'static + Future<Output = Result<proto::Envelope>> {
926 envelope.id = self.next_message_id.fetch_add(1, SeqCst);
927 let (tx, rx) = oneshot::channel();
928 let mut response_channels_lock = self.response_channels.lock();
929 response_channels_lock.insert(MessageId(envelope.id), tx);
930 drop(response_channels_lock);
931 let result = self.outgoing_tx.unbounded_send(envelope);
932 async move {
933 if let Err(error) = &result {
934 log::error!("failed to send message: {}", error);
935 return Err(anyhow!("failed to send message: {}", error));
936 }
937
938 let response = rx.await.context("connection lost")?.0;
939 if let Some(proto::envelope::Payload::Error(error)) = &response.payload {
940 return Err(RpcError::from_proto(error, type_name));
941 }
942 Ok(response)
943 }
944 }
945
946 pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {
947 envelope.id = self.next_message_id.fetch_add(1, SeqCst);
948 self.outgoing_tx.unbounded_send(envelope)?;
949 Ok(())
950 }
951}
952
953impl ProtoClient for ChannelClient {
954 fn request(
955 &self,
956 envelope: proto::Envelope,
957 request_type: &'static str,
958 ) -> BoxFuture<'static, Result<proto::Envelope>> {
959 self.request_dynamic(envelope, request_type).boxed()
960 }
961
962 fn send(&self, envelope: proto::Envelope, _message_type: &'static str) -> Result<()> {
963 self.send_dynamic(envelope)
964 }
965
966 fn send_response(&self, envelope: Envelope, _message_type: &'static str) -> anyhow::Result<()> {
967 self.send_dynamic(envelope)
968 }
969
970 fn message_handler_set(&self) -> &Mutex<ProtoMessageHandlerSet> {
971 &self.message_handlers
972 }
973
974 fn is_via_collab(&self) -> bool {
975 false
976 }
977}