use crate::{    json_log::LogRecord,    protocol::{        message_len_from_buffer, read_message_with_len, write_message, MessageId, MESSAGE_LEN_SIZE,    },};use anyhow::{anyhow, Context as _, Result};use collections::HashMap;use futures::{    channel::{mpsc, oneshot},    future::{BoxFuture, LocalBoxFuture},    select_biased, AsyncReadExt as _, AsyncWriteExt as _, Future, FutureExt as _, StreamExt as _,};use gpui::{AppContext, AsyncAppContext, Model, SemanticVersion, WeakModel};use parking_lot::Mutex;use rpc::{    proto::{        self, build_typed_envelope, AnyTypedEnvelope, Envelope, EnvelopedMessage, PeerId,        ProtoClient, RequestMessage,    },    TypedEnvelope,};use smol::{    fs,    process::{self, Stdio},};use std::{    any::TypeId,    ffi::OsStr,    path::{Path, PathBuf},    sync::{        atomic::{AtomicU32, Ordering::SeqCst},        Arc,    },    time::Instant,};use tempfile::TempDir; #[derive(Clone)]pub struct SshSocket {    connection_options: SshConnectionOptions,    socket_path: PathBuf,} pub struct SshSession {    next_message_id: AtomicU32,    response_channels: ResponseChannels,    outgoing_tx: mpsc::UnboundedSender<Envelope>,    spawn_process_tx: mpsc::UnboundedSender<SpawnRequest>,    client_socket: Option<SshSocket>,    message_handlers: Mutex<        HashMap<            TypeId,            Arc<                dyn Send                    + Sync                    + Fn(                        Box<dyn AnyTypedEnvelope>,                        Arc<SshSession>,                        AsyncAppContext,                    ) -> Option<LocalBoxFuture<'static, Result<()>>>,            >,        >,    >,} struct SshClientState {    socket: SshSocket,    _master_process: process::Child,    _temp_dir: TempDir,} #[derive(Debug, Clone, PartialEq, Eq)]pub struct SshConnectionOptions {    pub host: String,    pub username: Option<String>,    pub port: Option<u16>,    pub password: Option<String>,} impl SshConnectionOptions {    pub fn ssh_url(&self) -> String {        let mut result = String::from("ssh://");        if let Some(username) = &self.username {            result.push_str(username);            result.push('@');        }        result.push_str(&self.host);        if let Some(port) = self.port {            result.push(':');            result.push_str(&port.to_string());        }        result    }     fn scp_url(&self) -> String {        if let Some(username) = &self.username {            format!("{}@{}", username, self.host)        } else {            self.host.clone()        }    }     pub fn connection_string(&self) -> String {        let host = if let Some(username) = &self.username {            format!("{}@{}", username, self.host)        } else {            self.host.clone()        };        if let Some(port) = &self.port {            format!("{}:{}", host, port)        } else {            host        }    }} struct SpawnRequest {    command: String,    process_tx: oneshot::Sender<process::Child>,} #[derive(Copy, Clone, Debug)]pub struct SshPlatform {    pub os: &'static str,    pub arch: &'static str,} pub trait SshClientDelegate {    fn ask_password(        &self,        prompt: String,        cx: &mut AsyncAppContext,    ) -> oneshot::Receiver<Result<String>>;    fn remote_server_binary_path(&self, cx: &mut AsyncAppContext) -> Result<PathBuf>;    fn get_server_binary(        &self,        platform: SshPlatform,        cx: &mut AsyncAppContext,    ) -> oneshot::Receiver<Result<(PathBuf, SemanticVersion)>>;    fn set_status(&self, status: Option<&str>, cx: &mut AsyncAppContext);} type ResponseChannels = Mutex<HashMap<MessageId, oneshot::Sender<(Envelope, oneshot::Sender<()>)>>>; impl SshSession {    pub async fn client(        connection_options: SshConnectionOptions,        delegate: Arc<dyn SshClientDelegate>,        cx: &mut AsyncAppContext,    ) -> Result<Arc<Self>> {        let client_state = SshClientState::new(connection_options, delegate.clone(), cx).await?;         let platform = client_state.query_platform().await?;        let (local_binary_path, version) = delegate.get_server_binary(platform, cx).await??;        let remote_binary_path = delegate.remote_server_binary_path(cx)?;        client_state            .ensure_server_binary(                &delegate,                &local_binary_path,                &remote_binary_path,                version,                cx,            )            .await?;         let (spawn_process_tx, mut spawn_process_rx) = mpsc::unbounded::<SpawnRequest>();        let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded::<Envelope>();        let (incoming_tx, incoming_rx) = mpsc::unbounded::<Envelope>();         let socket = client_state.socket.clone();        run_cmd(socket.ssh_command(&remote_binary_path).arg("version")).await?;         let mut remote_server_child = socket            .ssh_command(&format!(                "RUST_LOG={} {:?} run",                std::env::var("RUST_LOG").unwrap_or(String::new()),                remote_binary_path,            ))            .spawn()            .context("failed to spawn remote server")?;        let mut child_stderr = remote_server_child.stderr.take().unwrap();        let mut child_stdout = remote_server_child.stdout.take().unwrap();        let mut child_stdin = remote_server_child.stdin.take().unwrap();         let executor = cx.background_executor().clone();        executor.clone().spawn(async move {            let mut stdin_buffer = Vec::new();            let mut stdout_buffer = Vec::new();            let mut stderr_buffer = Vec::new();            let mut stderr_offset = 0;             loop {                stdout_buffer.resize(MESSAGE_LEN_SIZE, 0);                stderr_buffer.resize(stderr_offset + 1024, 0);                 select_biased! {                    outgoing = outgoing_rx.next().fuse() => {                        let Some(outgoing) = outgoing else {                            return anyhow::Ok(());                        };                         write_message(&mut child_stdin, &mut stdin_buffer, outgoing).await?;                    }                     request = spawn_process_rx.next().fuse() => {                        let Some(request) = request else {                            return Ok(());                        };                         log::info!("spawn process: {:?}", request.command);                        let child = client_state.socket                            .ssh_command(&request.command)                            .spawn()                            .context("failed to create channel")?;                        request.process_tx.send(child).ok();                    }                     result = child_stdout.read(&mut stdout_buffer).fuse() => {                        match result {                            Ok(len) => {                                if len == 0 {                                    child_stdin.close().await?;                                    let status = remote_server_child.status().await?;                                    if !status.success() {                                        log::info!("channel exited with status: {status:?}");                                    }                                    return Ok(());                                }                                 if len < stdout_buffer.len() {                                    child_stdout.read_exact(&mut stdout_buffer[len..]).await?;                                }                                 let message_len = message_len_from_buffer(&stdout_buffer);                                match read_message_with_len(&mut child_stdout, &mut stdout_buffer, message_len).await {                                    Ok(envelope) => {                                        incoming_tx.unbounded_send(envelope).ok();                                    }                                    Err(error) => {                                        log::error!("error decoding message {error:?}");                                    }                                }                            }                            Err(error) => {                                Err(anyhow!("error reading stdout: {error:?}"))?;                            }                        }                    }                     result = child_stderr.read(&mut stderr_buffer[stderr_offset..]).fuse() => {                        match result {                            Ok(len) => {                                stderr_offset += len;                                let mut start_ix = 0;                                while let Some(ix) = stderr_buffer[start_ix..stderr_offset].iter().position(|b| b == &b'\n') {                                    let line_ix = start_ix + ix;                                    let content = &stderr_buffer[start_ix..line_ix];                                    start_ix = line_ix + 1;                                    if let Ok(record) = serde_json::from_slice::<LogRecord>(&content) {                                        record.log(log::logger())                                    } else {                                        eprintln!("(remote) {}", String::from_utf8_lossy(content));                                    }                                }                                stderr_buffer.drain(0..start_ix);                                stderr_offset -= start_ix;                            }                            Err(error) => {                                Err(anyhow!("error reading stderr: {error:?}"))?;                            }                        }                    }                }            }        }).detach();         cx.update(|cx| Self::new(incoming_rx, outgoing_tx, spawn_process_tx, Some(socket), cx))    }     pub fn server(        incoming_rx: mpsc::UnboundedReceiver<Envelope>,        outgoing_tx: mpsc::UnboundedSender<Envelope>,        cx: &AppContext,    ) -> Arc<SshSession> {        let (tx, _rx) = mpsc::unbounded();        Self::new(incoming_rx, outgoing_tx, tx, None, cx)    }     #[cfg(any(test, feature = "test-support"))]    pub fn fake(        client_cx: &mut gpui::TestAppContext,        server_cx: &mut gpui::TestAppContext,    ) -> (Arc<Self>, Arc<Self>) {        let (server_to_client_tx, server_to_client_rx) = mpsc::unbounded();        let (client_to_server_tx, client_to_server_rx) = mpsc::unbounded();        let (tx, _rx) = mpsc::unbounded();        (            client_cx.update(|cx| {                Self::new(                    server_to_client_rx,                    client_to_server_tx,                    tx.clone(),                    None, // todo()                    cx,                )            }),            server_cx.update(|cx| {                Self::new(                    client_to_server_rx,                    server_to_client_tx,                    tx.clone(),                    None,                    cx,                )            }),        )    }     fn new(        mut incoming_rx: mpsc::UnboundedReceiver<Envelope>,        outgoing_tx: mpsc::UnboundedSender<Envelope>,        spawn_process_tx: mpsc::UnboundedSender<SpawnRequest>,        client_socket: Option<SshSocket>,        cx: &AppContext,    ) -> Arc<SshSession> {        let this = Arc::new(Self {            next_message_id: AtomicU32::new(0),            response_channels: ResponseChannels::default(),            outgoing_tx,            spawn_process_tx,            client_socket,            message_handlers: Default::default(),        });         cx.spawn(|cx| {            let this = this.clone();            async move {                let peer_id = PeerId { owner_id: 0, id: 0 };                while let Some(incoming) = incoming_rx.next().await {                    if let Some(request_id) = incoming.responding_to {                        let request_id = MessageId(request_id);                        let sender = this.response_channels.lock().remove(&request_id);                        if let Some(sender) = sender {                            let (tx, rx) = oneshot::channel();                            if incoming.payload.is_some() {                                sender.send((incoming, tx)).ok();                            }                            rx.await.ok();                        }                    } else if let Some(envelope) =                        build_typed_envelope(peer_id, Instant::now(), incoming)                    {                        log::debug!(                            "ssh message received. name:{}",                            envelope.payload_type_name()                        );                        let type_id = envelope.payload_type_id();                        let handler = this.message_handlers.lock().get(&type_id).cloned();                        if let Some(handler) = handler {                            if let Some(future) = handler(envelope, this.clone(), cx.clone()) {                                future.await.ok();                            } else {                                this.message_handlers.lock().remove(&type_id);                            }                        }                    }                }                anyhow::Ok(())            }        })        .detach();         this    }     pub fn request<T: RequestMessage>(        &self,        payload: T,    ) -> impl 'static + Future<Output = Result<T::Response>> {        log::debug!("ssh request start. name:{}", T::NAME);        let response = self.request_dynamic(payload.into_envelope(0, None, None), "");        async move {            let response = response.await?;            log::debug!("ssh request finish. name:{}", T::NAME);            T::Response::from_envelope(response)                .ok_or_else(|| anyhow!("received a response of the wrong type"))        }    }     pub fn send<T: EnvelopedMessage>(&self, payload: T) -> Result<()> {        self.send_dynamic(payload.into_envelope(0, None, None))    }     pub fn request_dynamic(        &self,        mut envelope: proto::Envelope,        _request_type: &'static str,    ) -> impl 'static + Future<Output = Result<proto::Envelope>> {        envelope.id = self.next_message_id.fetch_add(1, SeqCst);        let (tx, rx) = oneshot::channel();        self.response_channels            .lock()            .insert(MessageId(envelope.id), tx);        self.outgoing_tx.unbounded_send(envelope).ok();        async move { Ok(rx.await.context("connection lost")?.0) }    }     pub fn send_dynamic(&self, mut envelope: proto::Envelope) -> Result<()> {        envelope.id = self.next_message_id.fetch_add(1, SeqCst);        self.outgoing_tx.unbounded_send(envelope)?;        Ok(())    }     pub async fn spawn_process(&self, command: String) -> process::Child {        let (process_tx, process_rx) = oneshot::channel();        self.spawn_process_tx            .unbounded_send(SpawnRequest {                command,                process_tx,            })            .ok();        process_rx.await.unwrap()    }     pub fn ssh_args(&self) -> Vec<String> {        self.client_socket.as_ref().unwrap().ssh_args()    }     pub fn add_message_handler<M, E, H, F>(&self, entity: WeakModel<E>, handler: H)    where        M: EnvelopedMessage,        E: 'static,        H: 'static + Sync + Send + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F,        F: 'static + Future<Output = Result<()>>,    {        let message_type_id = TypeId::of::<M>();        self.message_handlers.lock().insert(            message_type_id,            Arc::new(move |envelope, _, cx| {                let entity = entity.upgrade()?;                let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();                Some(handler(entity, *envelope, cx).boxed_local())            }),        );    }     pub fn add_request_handler<M, E, H, F>(&self, entity: WeakModel<E>, handler: H)    where        M: EnvelopedMessage + RequestMessage,        E: 'static,        H: 'static + Sync + Send + Fn(Model<E>, TypedEnvelope<M>, AsyncAppContext) -> F,        F: 'static + Future<Output = Result<M::Response>>,    {        let message_type_id = TypeId::of::<M>();        self.message_handlers.lock().insert(            message_type_id,            Arc::new(move |envelope, this, cx| {                let entity = entity.upgrade()?;                let envelope = envelope.into_any().downcast::<TypedEnvelope<M>>().unwrap();                let request_id = envelope.message_id();                Some(                    handler(entity, *envelope, cx)                        .then(move |result| async move {                            this.outgoing_tx.unbounded_send(result?.into_envelope(                                this.next_message_id.fetch_add(1, SeqCst),                                Some(request_id),                                None,                            ))?;                            Ok(())                        })                        .boxed_local(),                )            }),        );    }} impl ProtoClient for SshSession {    fn request(        &self,        envelope: proto::Envelope,        request_type: &'static str,    ) -> BoxFuture<'static, Result<proto::Envelope>> {        self.request_dynamic(envelope, request_type).boxed()    }     fn send(&self, envelope: proto::Envelope) -> Result<()> {        self.send_dynamic(envelope)    }} impl SshClientState {    #[cfg(not(unix))]    async fn new(        _connection_options: SshConnectionOptions,        _delegate: Arc<dyn SshClientDelegate>,        _cx: &mut AsyncAppContext,    ) -> Result<Self> {        Err(anyhow!("ssh is not supported on this platform"))

Lines 1–500 of 755. Load more