1use std::collections::HashMap;2use std::net::SocketAddr;3use std::sync::atomic::{AtomicU64, Ordering};4use std::sync::{Arc, Mutex};56use remowt_endpoints::forward::ForwardClient;7use remowt_link_shared::editor::{EditorBackend, Error};8use remowt_link_shared::BifConfig;9use russh::client::Handle;10use tokio::net::{TcpListener, UdpSocket, UnixListener};11use tracing::error;1213use crate::{Remowt, SshHandler};1415pub struct SshEditor {16 pub sess: Arc<Handle<SshHandler>>,17 pub conn: Remowt,18}19impl EditorBackend for SshEditor {20 async fn open_editor(&self, socket_path: String) -> Result<(), Error> {21 let local = std::env::temp_dir().join(format!("remowt-nvim-{}.sock", uuid::Uuid::new_v4()));22 let _ = std::fs::remove_file(&local);23 let listener = UnixListener::bind(&local).map_err(|e| Error::Failed(e.to_string()))?;2425 let sess = self.sess.clone();26 let forward = tokio::spawn(async move {27 loop {28 let Ok((mut stream, _)) = listener.accept().await else {29 break;30 };31 let sess = sess.clone();32 let remote = socket_path.clone();33 tokio::spawn(async move {34 match sess.channel_open_direct_streamlocal(remote).await {35 Ok(ch) => {36 let mut remote = ch.into_stream();37 let _ = tokio::io::copy_bidirectional(&mut stream, &mut remote).await;38 }39 Err(e) => error!("opening direct-streamlocal to nvim failed: {e}"),40 }41 });42 }43 });4445 let status = tokio::process::Command::new("neovide")46 .arg("--no-fork")47 .arg("--server")48 .arg(&local)49 .status()50 .await51 .map_err(|e| Error::Failed(format!("spawning neovide: {e}")));5253 forward.abort();54 let _ = std::fs::remove_file(&local);5556 match status? {57 s if s.success() => Ok(()),58 s => Err(Error::Failed(format!("neovide exited with {s}"))),59 }60 }6162 async fn expose_tcp(&self, addr: String) -> Result<u16, Error> {63 let listener = TcpListener::bind(("127.0.0.1", 0))64 .await65 .map_err(|e| Error::Failed(e.to_string()))?;66 let local = listener67 .local_addr()68 .map_err(|e| Error::Failed(e.to_string()))?69 .port();7071 let conn = self.conn.clone();72 tokio::spawn(async move {73 loop {74 let Ok((mut tcp, _)) = listener.accept().await else {75 break;76 };77 let conn = conn.clone();78 let addr = addr.clone();79 tokio::spawn(async move {80 let (forwarded, tunnel) = match conn.bind_fast_tunnel("forward", false).await {81 Ok(v) => v,82 Err(e) => {83 error!("forward: bind tunnel failed: {e}");84 return;85 }86 };87 let fclient: ForwardClient<BifConfig> = conn.endpoints();88 match fclient.connect_tcp(tunnel, addr).await {89 Ok(Ok(())) => {}90 Ok(Err(e)) => {91 error!("forward: agent connect_tcp failed: {e}");92 return;93 }94 Err(e) => {95 error!("forward: connect_tcp rpc failed: {e}");96 return;97 }98 }99 match forwarded.accept().await {100 Ok(mut stream) => {101 let _ = tokio::io::copy_bidirectional(&mut tcp, &mut stream).await;102 }103 Err(e) => error!("forward: accept tunnel failed: {e}"),104 }105 });106 }107 });108109 Ok(local)110 }111112 async fn expose_udp(&self, addr: String) -> Result<u16, Error> {113 let router = self.conn.datagram_router().ok_or_else(|| {114 Error::Failed(115 "udp forward requires the iroh fast path, which is not established".into(),116 )117 })?;118119 let fclient: ForwardClient<BifConfig> = self.conn.endpoints();120 let session = fclient121 .open_udp(addr)122 .await123 .map_err(|e| Error::Failed(format!("open_udp rpc: {e}")))?124 .map_err(|e| Error::Failed(format!("agent open_udp: {e}")))?;125126 let sock = Arc::new(127 UdpSocket::bind(("127.0.0.1", 0))128 .await129 .map_err(|e| Error::Failed(e.to_string()))?,130 );131 let local = sock132 .local_addr()133 .map_err(|e| Error::Failed(e.to_string()))?134 .port();135136 let sub_for_source: Arc<Mutex<HashMap<SocketAddr, u64>>> =137 Arc::new(Mutex::new(HashMap::new()));138 let source_for_sub: Arc<Mutex<HashMap<u64, SocketAddr>>> =139 Arc::new(Mutex::new(HashMap::new()));140 let next_sub = Arc::new(AtomicU64::new(0));141 let mut rx = router.register(session);142143 let up_sock = sock.clone();144 let up_router = router.clone();145 let down_source_for_sub = source_for_sub.clone();146 tokio::spawn(async move {147 let mut buf = vec![0u8; 65535];148 loop {149 let (n, src) = match up_sock.recv_from(&mut buf).await {150 Ok(v) => v,151 Err(_) => break,152 };153 let sub = {154 let mut by_src = sub_for_source.lock().expect("lock");155 if let Some(&sub) = by_src.get(&src) {156 sub157 } else {158 let sub = next_sub.fetch_add(1, Ordering::Relaxed);159 by_src.insert(src, sub);160 source_for_sub.lock().expect("lock").insert(sub, src);161 sub162 }163 };164 if up_router.send(session, sub, &buf[..n]).is_err() {165 break;166 }167 }168 up_router.unregister(session);169 });170171 let down_sock = sock.clone();172 tokio::spawn(async move {173 while let Some((sub, payload)) = rx.recv().await {174 let dst = down_source_for_sub.lock().expect("lock").get(&sub).copied();175 if let Some(dst) = dst {176 let _ = down_sock.send_to(&payload, dst).await;177 }178 }179 });180181 Ok(local)182 }183}