iroh_ssh/
ssh.rs

1macro_rules! ok_or_continue {
2    ($expr:expr) => {
3        match $expr {
4            Ok(val) => val,
5            Err(_) => continue,
6        }
7    };
8}
9
10use std::{pin::Pin};
11
12use crate::{Builder, Inner, IrohSsh};
13
14use ed25519_dalek::{SECRET_KEY_LENGTH};
15use iroh::{
16    endpoint::{Connection}, protocol::{ProtocolHandler, Router}, Endpoint, NodeId, SecretKey
17};
18use tokio::{
19    
20    net::TcpStream,
21};
22
23impl Builder {
24    pub fn new() -> Self {
25        Self {
26            secret_key: SecretKey::generate(rand::rngs::OsRng).to_bytes(),
27            accept_incoming: false,
28            accept_port: None,
29        }
30    }
31
32    pub fn accept_incoming(mut self, accept_incoming: bool) -> Self {
33        self.accept_incoming = accept_incoming;
34        self
35    }
36
37    pub fn accept_port(mut self, accept_port: u16) -> Self {
38        self.accept_port = Some(accept_port);
39        self
40    }
41
42    pub fn secret_key(mut self, secret_key: &[u8; SECRET_KEY_LENGTH]) -> Self {
43        self.secret_key = *secret_key;
44        self
45    }
46
47    pub async fn build(self: &mut Self) -> anyhow::Result<IrohSsh> {
48        // Iroh setup
49        let secret_key = SecretKey::from_bytes(&self.secret_key);
50        let endpoint = Endpoint::builder()
51            .secret_key(secret_key)
52            .discovery_n0()
53            .bind()
54            .await?;
55
56        let mut iroh_ssh = IrohSsh {
57            public_key: endpoint.node_id().as_bytes().clone(),
58            secret_key: self.secret_key,
59            inner: None,
60        };
61
62        let router = if self.accept_incoming {
63            Router::builder(endpoint.clone()).accept(&IrohSsh::ALPN(), iroh_ssh.clone())
64        } else {
65            Router::builder(endpoint.clone())
66        }
67        .spawn();
68
69        iroh_ssh.add_inner(endpoint, router);
70
71        if self.accept_incoming && self.accept_port.is_some() {
72            tokio::spawn({
73                let iroh_ssh = iroh_ssh.clone();
74                let accept_port = self.accept_port.expect("accept_port not set");
75                async move {
76                    iroh_ssh._spawn(accept_port).await.expect("spawn failed");
77                }
78            });
79        }
80
81        Ok(iroh_ssh)
82    }
83}
84
85impl IrohSsh {
86    pub fn new() -> Builder {
87        Builder::new()
88    }
89
90    #[allow(non_snake_case)]
91    pub fn ALPN() -> Vec<u8> {
92        format!("/iroh/ssh").into_bytes()
93    }
94
95    fn add_inner(&mut self, endpoint: Endpoint, router: Router) {
96        self.inner = Some(Inner { endpoint, router });
97    }
98
99    pub async fn connect(&self, node_id: NodeId) -> anyhow::Result<Connection> {
100        let inner = self.inner.as_ref().expect("inner not set");
101        inner.endpoint.connect(node_id, &IrohSsh::ALPN()).await
102    }
103
104    pub fn node_id(&self) -> NodeId {
105        self.inner.as_ref().expect("inner not set").endpoint.node_id()
106    }
107
108    async fn _spawn(self, port: u16) -> anyhow::Result<()> {
109        while let Some(incoming) = self
110            .inner
111            .clone()
112            .expect("inner not set")
113            .endpoint
114            .accept()
115            .await
116        {
117            let mut connecting = match incoming.accept() {
118                Ok(connecting) => connecting,
119                Err(err) => {
120                    println!("incomming connection failure: {err:#}");
121                    continue;
122                }
123            };
124            let alpn = ok_or_continue!(connecting.alpn().await);
125
126            let conn = ok_or_continue!(connecting.await);
127            let node_id = ok_or_continue!(conn.remote_node_id());
128            println!("{}: {node_id} incoming...", String::from_utf8_lossy(&alpn),);
129
130            tokio::spawn(async move {
131                if let Ok((mut send, mut recv)) = conn.accept_bi().await {
132                    if let Ok(mut ssh_stream) = TcpStream::connect(format!("127.0.0.1:{}", port)).await {
133                    
134                    let (mut local_read, mut local_write) = ssh_stream.split();
135                    let a_to_b = async move { tokio::io::copy(&mut local_read, &mut send).await };
136                    let b_to_a = async move { tokio::io::copy(&mut recv, &mut local_write).await };
137
138                    tokio::select! {
139                        result = a_to_b => { let _ = result; },
140                        result = b_to_a => { let _ = result; },
141                    };
142                }};
143            });
144        }
145        Ok(())
146    }
147}
148
149impl ProtocolHandler for IrohSsh {
150    fn accept(
151        &self,
152        conn: Connection,
153    ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + 'static>> {
154        let iroh_ssh = self.clone();
155
156        Box::pin(async move {
157            iroh_ssh.accept(conn).await?;
158            Ok(())
159        })
160    }
161}
iroh_ssh/ssh.rs

iroh_ssh/
ssh.rs
