1use std::io;2use std::pin::Pin;3use std::task::{Context, Poll};45use anyhow::{anyhow, Result};6use camino::Utf8PathBuf;7use remowt_link_shared::iroh_tunnel::IrohBiStream;8use russh::client::Msg;9use russh::{Channel, ChannelStream};10use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};11use tokio::net::{UnixListener, UnixStream};12use tokio::sync::oneshot;1314pub enum RemowtListener {15 Ssh(oneshot::Receiver<Channel<Msg>>),16 Local(UnixListener, Utf8PathBuf),17 Iroh(oneshot::Receiver<RemowtStream>),18}1920impl RemowtListener {21 pub async fn accept(self) -> Result<RemowtStream> {22 match self {23 RemowtListener::Ssh(rx) => {24 let ch = rx25 .await26 .map_err(|_| anyhow!("agent never connected the forwarded socket"))?;27 Ok(RemowtStream::Ssh(ch.into_stream()))28 }29 RemowtListener::Local(listener, path) => {30 let (stream, _) = listener.accept().await?;31 let _ = std::fs::remove_file(&path);32 Ok(RemowtStream::Local(stream))33 }34 RemowtListener::Iroh(rx) => rx35 .await36 .map_err(|_| anyhow!("agent never opened the iroh tunnel")),37 }38 }39}4041pub enum RemowtStream {42 Ssh(ChannelStream<Msg>),43 Local(UnixStream),44 Iroh(IrohBiStream),45}4647impl AsyncRead for RemowtStream {48 fn poll_read(49 self: Pin<&mut Self>,50 cx: &mut Context<'_>,51 buf: &mut ReadBuf<'_>,52 ) -> Poll<io::Result<()>> {53 match self.get_mut() {54 RemowtStream::Ssh(s) => Pin::new(s).poll_read(cx, buf),55 RemowtStream::Local(s) => Pin::new(s).poll_read(cx, buf),56 RemowtStream::Iroh(s) => Pin::new(s).poll_read(cx, buf),57 }58 }59}6061impl AsyncWrite for RemowtStream {62 fn poll_write(63 self: Pin<&mut Self>,64 cx: &mut Context<'_>,65 buf: &[u8],66 ) -> Poll<io::Result<usize>> {67 match self.get_mut() {68 RemowtStream::Ssh(s) => Pin::new(s).poll_write(cx, buf),69 RemowtStream::Local(s) => Pin::new(s).poll_write(cx, buf),70 RemowtStream::Iroh(s) => Pin::new(s).poll_write(cx, buf),71 }72 }7374 fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {75 match self.get_mut() {76 RemowtStream::Ssh(s) => Pin::new(s).poll_flush(cx),77 RemowtStream::Local(s) => Pin::new(s).poll_flush(cx),78 RemowtStream::Iroh(s) => Pin::new(s).poll_flush(cx),79 }80 }8182 fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> {83 match self.get_mut() {84 RemowtStream::Ssh(s) => Pin::new(s).poll_shutdown(cx),85 RemowtStream::Local(s) => Pin::new(s).poll_shutdown(cx),86 RemowtStream::Iroh(s) => Pin::new(s).poll_shutdown(cx),87 }88 }89}