ssh_session.rs

  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}