1use std::collections::HashMap;2use std::pin::Pin;3use std::sync::{Arc, Mutex};4use std::task::{Context, Poll};5use std::time::Duration;6use std::{fmt, io};78use anyhow::Result;9use bytes::{Bytes, BytesMut};10use camino::Utf8PathBuf;11use futures::{SinkExt as _, Stream};12use iroh::endpoint::transports::{13 CustomEndpoint, CustomSender, CustomTransport, RecvInfo, Transmit,14};15use iroh::endpoint::{presets, Connection, RecvStream, SendStream};16use iroh::{Endpoint, EndpointId, SecretKey};17use iroh_base::CustomAddr;18use serde::{Deserialize, Serialize};19use tokio::io::{AsyncRead, AsyncWrite, ReadBuf, ReadHalf, WriteHalf};20use tokio::net::UnixStream;21use tokio::sync::{mpsc, watch};22use tokio_util::codec::{FramedRead, FramedWrite, LengthDelimitedCodec};23use tracing::{debug, warn};242526pub const REMOWT_SSH_TRANSPORT_ID: u64 = 0x_72_6d_77_74_73_73_68;2728pub const REMOWT_ALPN: &[u8] = b"remowt/tunnel/0";2930#[derive(Clone, Debug, Serialize, Deserialize)]31pub enum TunnelAddr {32 Unix(Utf8PathBuf),33 Iroh { token: u64 },34}3536pub fn ssh_custom_addr(id: EndpointId) -> CustomAddr {37 CustomAddr::from((REMOWT_SSH_TRANSPORT_ID, &id.as_bytes()[..]))38}3940pub async fn build_endpoint<S>(41 secret: SecretKey,42 stream: S,43 remote: EndpointId,44 accept: bool,45) -> Result<Endpoint>46where47 S: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static,48{49 let local = secret.public();50 let transport = Arc::new(SshTransport::new(stream, local, remote));51 let mut builder = Endpoint::builder(presets::N0)52 .secret_key(secret)53 .add_custom_transport(transport);54 if accept {55 builder = builder.alpns(vec![REMOWT_ALPN.to_vec()]);56 }57 Ok(builder.bind().await?)58}5960struct SshTransport {61 local: CustomAddr,62 remote: CustomAddr,63 addrs: n0_watcher::Watchable<Vec<CustomAddr>>,64 endpoint: Mutex<Option<SshEndpointParts>>,65 sender: Arc<SshSender>,66}6768struct SshEndpointParts {69 reader: Pin<Box<dyn FramedDatagrams>>,70}7172trait FramedDatagrams: Stream<Item = io::Result<BytesMut>> + Send + Sync {}73impl<T: Stream<Item = io::Result<BytesMut>> + Send + Sync> FramedDatagrams for T {}7475impl fmt::Debug for SshTransport {76 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {77 f.debug_struct("SshTransport").finish_non_exhaustive()78 }79}8081impl SshTransport {82 fn new<S>(stream: S, local: EndpointId, remote: EndpointId) -> Self83 where84 S: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static,85 {86 let (read, write): (ReadHalf<S>, WriteHalf<S>) = tokio::io::split(stream);87 let reader = FramedRead::new(read, LengthDelimitedCodec::new());8889 let (tx, rx) = mpsc::unbounded_channel::<Bytes>();90 tokio::spawn(writer_task(rx, write));9192 let local = ssh_custom_addr(local);93 Self {94 addrs: n0_watcher::Watchable::new(vec![local.clone()]),95 local,96 remote: ssh_custom_addr(remote),97 endpoint: Mutex::new(Some(SshEndpointParts {98 reader: Box::pin(reader),99 })),100 sender: Arc::new(SshSender { tx }),101 }102 }103}104105async fn writer_task<W: AsyncWrite + Unpin>(mut rx: mpsc::UnboundedReceiver<Bytes>, write: W) {106 let mut framed = FramedWrite::new(write, LengthDelimitedCodec::new());107 while let Some(datagram) = rx.recv().await {108 if let Err(e) = framed.send(datagram).await {109 debug!("ssh transport writer ended: {e}");110 break;111 }112 }113}114115impl CustomTransport for SshTransport {116 fn bind(&self) -> io::Result<Box<dyn CustomEndpoint>> {117 let parts = self118 .endpoint119 .lock()120 .expect("not poisoned")121 .take()122 .ok_or_else(|| io::Error::other("ssh transport already bound"))?;123 Ok(Box::new(SshEndpoint {124 local: self.local.clone(),125 remote: self.remote.clone(),126 addrs: self.addrs.clone(),127 sender: self.sender.clone(),128 reader: parts.reader,129 }))130 }131}132133struct SshEndpoint {134 local: CustomAddr,135 remote: CustomAddr,136 addrs: n0_watcher::Watchable<Vec<CustomAddr>>,137 sender: Arc<SshSender>,138 reader: Pin<Box<dyn FramedDatagrams>>,139}140141impl fmt::Debug for SshEndpoint {142 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {143 f.debug_struct("SshEndpoint").finish_non_exhaustive()144 }145}146147impl CustomEndpoint for SshEndpoint {148 fn watch_local_addrs(&self) -> n0_watcher::Direct<Vec<CustomAddr>> {149 self.addrs.watch()150 }151152 fn create_sender(&self) -> Arc<dyn CustomSender> {153 self.sender.clone()154 }155156 fn poll_recv(157 &mut self,158 cx: &mut Context,159 bufs: &mut [io::IoSliceMut<'_>],160 metas: &mut [noq_udp::RecvMeta],161 recv_infos: &mut [RecvInfo],162 ) -> Poll<io::Result<usize>> {163 assert_eq!(bufs.len(), metas.len());164 assert_eq!(bufs.len(), recv_infos.len());165 if bufs.is_empty() {166 return Poll::Ready(Ok(0));167 }168 let mut count = 0;169 while count < bufs.len() {170 match self.reader.as_mut().poll_next(cx) {171 Poll::Ready(Some(Ok(frame))) => {172 let buf = &mut bufs[count];173 if buf.len() < frame.len() {174 warn!("ssh transport datagram {} > buf {}", frame.len(), buf.len());175 if count > 0 {176 break;177 }178 return Poll::Ready(Err(io::Error::other("datagram exceeds recv buffer")));179 }180 buf[..frame.len()].copy_from_slice(&frame);181 metas[count].len = frame.len();182 metas[count].stride = frame.len();183 recv_infos[count] =184 RecvInfo::new(self.remote.clone(), Some(self.local.clone()));185 count += 1;186 }187 Poll::Ready(Some(Err(e))) => {188 if count > 0 {189 break;190 }191 return Poll::Ready(Err(e));192 }193 Poll::Ready(None) => {194 if count > 0 {195 break;196 }197 return Poll::Ready(Err(io::Error::other("ssh transport closed")));198 }199 Poll::Pending => {200 if count > 0 {201 break;202 }203 return Poll::Pending;204 }205 }206 }207 Poll::Ready(Ok(count))208 }209}210211#[derive(Debug)]212struct SshSender {213 tx: mpsc::UnboundedSender<Bytes>,214}215216impl CustomSender for SshSender {217 fn is_valid_send_addr(&self, addr: &CustomAddr) -> bool {218 addr.id() == REMOWT_SSH_TRANSPORT_ID219 }220221 fn poll_send(222 &self,223 _cx: &mut Context,224 _dst: &CustomAddr,225 _src: Option<&CustomAddr>,226 transmit: &Transmit<'_>,227 ) -> Poll<io::Result<()>> {228 let segment = transmit.segment_size.unwrap_or(transmit.contents.len());229 if segment == 0 || transmit.contents.is_empty() {230 return Poll::Ready(Ok(()));231 }232 for chunk in transmit.contents.chunks(segment) {233 if self.tx.send(Bytes::copy_from_slice(chunk)).is_err() {234 return Poll::Ready(Err(io::Error::other("ssh transport writer gone")));235 }236 }237 Poll::Ready(Ok(()))238 }239}240241pub struct IrohBiStream {242 send: SendStream,243 recv: RecvStream,244}245246impl IrohBiStream {247 pub fn new(send: SendStream, recv: RecvStream) -> Self {248 Self { send, recv }249 }250}251252impl AsyncRead for IrohBiStream {253 fn poll_read(254 self: Pin<&mut Self>,255 cx: &mut Context<'_>,256 buf: &mut ReadBuf<'_>,257 ) -> Poll<io::Result<()>> {258 AsyncRead::poll_read(Pin::new(&mut self.get_mut().recv), cx, buf)259 }260}261262impl AsyncWrite for IrohBiStream {263 fn poll_write(264 self: Pin<&mut Self>,265 cx: &mut Context<'_>,266 buf: &[u8],267 ) -> Poll<io::Result<usize>> {268 AsyncWrite::poll_write(Pin::new(&mut self.get_mut().send), cx, buf)269 }270 fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {271 AsyncWrite::poll_flush(Pin::new(&mut self.get_mut().send), cx)272 }273 fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {274 AsyncWrite::poll_shutdown(Pin::new(&mut self.get_mut().send), cx)275 }276}277278pub enum TunnelStream {279 Unix(UnixStream),280 Iroh(IrohBiStream),281}282283impl AsyncRead for TunnelStream {284 fn poll_read(285 self: Pin<&mut Self>,286 cx: &mut Context<'_>,287 buf: &mut ReadBuf<'_>,288 ) -> Poll<io::Result<()>> {289 match self.get_mut() {290 TunnelStream::Unix(s) => Pin::new(s).poll_read(cx, buf),291 TunnelStream::Iroh(s) => Pin::new(s).poll_read(cx, buf),292 }293 }294}295296impl AsyncWrite for TunnelStream {297 fn poll_write(298 self: Pin<&mut Self>,299 cx: &mut Context<'_>,300 buf: &[u8],301 ) -> Poll<io::Result<usize>> {302 match self.get_mut() {303 TunnelStream::Unix(s) => Pin::new(s).poll_write(cx, buf),304 TunnelStream::Iroh(s) => Pin::new(s).poll_write(cx, buf),305 }306 }307 fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {308 match self.get_mut() {309 TunnelStream::Unix(s) => Pin::new(s).poll_flush(cx),310 TunnelStream::Iroh(s) => Pin::new(s).poll_flush(cx),311 }312 }313 fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {314 match self.get_mut() {315 TunnelStream::Unix(s) => Pin::new(s).poll_shutdown(cx),316 TunnelStream::Iroh(s) => Pin::new(s).poll_shutdown(cx),317 }318 }319}320321#[derive(Clone)]322pub struct TunnelDialer {323 conn: watch::Sender<Option<Connection>>,324 router: Arc<Mutex<Option<Arc<DatagramRouter>>>>,325}326327impl Default for TunnelDialer {328 fn default() -> Self {329 Self::new()330 }331}332333impl TunnelDialer {334 pub fn new() -> Self {335 let (conn, _) = watch::channel(None);336 Self {337 conn,338 router: Arc::new(Mutex::new(None)),339 }340 }341342 pub fn set_conn(&self, conn: Connection) {343 *self.router.lock().expect("lock") = Some(DatagramRouter::spawn(conn.clone()));344 self.conn.send_replace(Some(conn));345 }346347 pub fn router(&self) -> Option<Arc<DatagramRouter>> {348 self.router.lock().expect("lock").clone()349 }350351 pub async fn connect_tunnel(&self, addr: &TunnelAddr) -> io::Result<TunnelStream> {352 match addr {353 TunnelAddr::Unix(path) => Ok(TunnelStream::Unix(UnixStream::connect(path).await?)),354 TunnelAddr::Iroh { token } => {355 let mut rx = self.conn.subscribe();356 let conn =357 tokio::time::timeout(Duration::from_secs(5), rx.wait_for(|c| c.is_some()))358 .await359 .map_err(|_| io::Error::other("timed out waiting for iroh connection"))?360 .map_err(|_| io::Error::other("iroh connection channel closed"))?361 .clone()362 .expect("is_some");363 let (mut send, recv) = conn.open_bi().await.map_err(io::Error::other)?;364 send.write_all(&token.to_be_bytes())365 .await366 .map_err(io::Error::other)?;367 Ok(TunnelStream::Iroh(IrohBiStream::new(send, recv)))368 }369 }370 }371}372373const DGRAM_HEADER: usize = 16;374375type DatagramRoutes = Arc<Mutex<HashMap<u64, mpsc::UnboundedSender<(u64, Bytes)>>>>;376377pub struct DatagramRouter {378 conn: Connection,379 routes: DatagramRoutes,380}381382impl DatagramRouter {383 pub fn spawn(conn: Connection) -> Arc<Self> {384 let routes: DatagramRoutes = Arc::new(Mutex::new(HashMap::new()));385 let read_routes = routes.clone();386 let read_conn = conn.clone();387 tokio::spawn(async move {388 loop {389 match read_conn.read_datagram().await {390 Ok(dg) if dg.len() >= DGRAM_HEADER => {391 let session = u64::from_be_bytes(dg[0..8].try_into().expect("8 bytes"));392 let sub = u64::from_be_bytes(dg[8..16].try_into().expect("8 bytes"));393 let payload = dg.slice(DGRAM_HEADER..);394 let tx = read_routes.lock().expect("lock").get(&session).cloned();395 if let Some(tx) = tx {396 let _ = tx.send((sub, payload));397 }398 }399 Ok(_) => {}400 Err(e) => {401 debug!("datagram read loop ended: {e}");402 break;403 }404 }405 }406 });407 Arc::new(Self { conn, routes })408 }409410 pub fn register(&self, session: u64) -> mpsc::UnboundedReceiver<(u64, Bytes)> {411 let (tx, rx) = mpsc::unbounded_channel();412 self.routes.lock().expect("lock").insert(session, tx);413 rx414 }415416 pub fn unregister(&self, session: u64) {417 self.routes.lock().expect("lock").remove(&session);418 }419420 pub fn send(&self, session: u64, sub: u64, payload: &[u8]) -> io::Result<()> {421 let mut buf = BytesMut::with_capacity(DGRAM_HEADER + payload.len());422 buf.extend_from_slice(&session.to_be_bytes());423 buf.extend_from_slice(&sub.to_be_bytes());424 buf.extend_from_slice(payload);425 self.conn426 .send_datagram(buf.freeze())427 .map_err(io::Error::other)428 }429}430431#[cfg(test)]432mod tests {433 use super::*;434 use iroh::endpoint::presets;435 use iroh::{EndpointAddr, RelayMode, TransportAddr};436437 async fn endpoint<S>(secret: SecretKey, stream: S, remote: EndpointId, accept: bool) -> Endpoint438 where439 S: AsyncRead + AsyncWrite + Unpin + Send + Sync + 'static,440 {441 let local = secret.public();442 let transport = Arc::new(SshTransport::new(stream, local, remote));443 let mut builder = Endpoint::builder(presets::N0)444 .secret_key(secret)445 .relay_mode(RelayMode::Disabled)446 .clear_ip_transports()447 .add_custom_transport(transport);448 if accept {449 builder = builder.alpns(vec![REMOWT_ALPN.to_vec()]);450 }451 builder.bind().await.expect("bind")452 }453454 #[tokio::test]455 async fn echo_over_ssh_transport() {456 let (client_pipe, agent_pipe) = tokio::io::duplex(256 * 1024);457 let client_secret = SecretKey::generate();458 let agent_secret = SecretKey::generate();459 let client_id = client_secret.public();460 let agent_id = agent_secret.public();461462 let client = endpoint(client_secret, client_pipe, agent_id, false).await;463 let agent = endpoint(agent_secret, agent_pipe, client_id, true).await;464465 let server = tokio::spawn(async move {466 let incoming = agent.accept().await.expect("incoming");467 let conn = incoming468 .accept()469 .expect("accept")470 .await471 .expect("connecting");472 assert_eq!(conn.remote_id(), client_id);473 let (mut send, mut recv) = conn.accept_bi().await.expect("accept_bi");474 let mut buf = [0u8; 5];475 recv.read_exact(&mut buf).await.expect("read");476 send.write_all(&buf).await.expect("write");477 send.finish().expect("finish");478 conn.closed().await;479 });480481 let addr =482 EndpointAddr::from_parts(agent_id, [TransportAddr::Custom(ssh_custom_addr(agent_id))]);483 let conn = client.connect(addr, REMOWT_ALPN).await.expect("connect");484 assert_eq!(conn.remote_id(), agent_id);485 let (mut send, mut recv) = conn.open_bi().await.expect("open_bi");486 send.write_all(b"hello").await.expect("write");487 send.finish().expect("finish");488 let mut buf = [0u8; 5];489 recv.read_exact(&mut buf).await.expect("read");490 assert_eq!(&buf, b"hello");491492 conn.close(0u32.into(), b"done");493 let _ = server.await;494 }495496 #[tokio::test]497 async fn dgram_round_trip() {498 let (client_pipe, agent_pipe) = tokio::io::duplex(256 * 1024);499 let client_secret = SecretKey::generate();500 let agent_secret = SecretKey::generate();501 let client_id = client_secret.public();502 let agent_id = agent_secret.public();503504 let client = endpoint(client_secret, client_pipe, agent_id, false).await;505 let agent = endpoint(agent_secret, agent_pipe, client_id, true).await;506507 let server = tokio::spawn(async move {508 let incoming = agent.accept().await.expect("incoming");509 let conn = incoming510 .accept()511 .expect("accept")512 .await513 .expect("connecting");514 let router = DatagramRouter::spawn(conn);515 let mut rx = router.register(7);516 if let Some((sub, payload)) = rx.recv().await {517 router.send(7, sub, &payload).expect("send reply");518 }519 tokio::time::sleep(Duration::from_millis(200)).await;520 });521522 let addr =523 EndpointAddr::from_parts(agent_id, [TransportAddr::Custom(ssh_custom_addr(agent_id))]);524 let conn = client.connect(addr, REMOWT_ALPN).await.expect("connect");525 let router = DatagramRouter::spawn(conn);526 let mut rx = router.register(7);527 router.send(7, 42, b"ping").expect("send");528 let (sub, payload) = tokio::time::timeout(Duration::from_secs(5), rx.recv())529 .await530 .expect("datagram timed out")531 .expect("router closed");532 assert_eq!(sub, 42);533 assert_eq!(&payload[..], b"ping");534 let _ = server.await;535 }536}