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