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 crate::{Builder, Inner, IrohSsh};
11use std::{pin::Pin, process::Stdio};
12
13use ed25519_dalek::SECRET_KEY_LENGTH;
14use iroh::{
15    Endpoint, NodeId, SecretKey,
16    endpoint::Connection,
17    protocol::{ProtocolHandler, Router},
18};
19use tokio::{
20    net::{TcpListener, TcpStream},
21    process::{Child, Command},
22};
23use homedir::my_home;
24
25
26impl Builder {
27    pub fn new() -> Self {
28        Self {
29            secret_key: SecretKey::generate(rand::rngs::OsRng).to_bytes(),
30            accept_incoming: false,
31            accept_port: None,
32        }
33    }
34
35    pub fn accept_incoming(mut self, accept_incoming: bool) -> Self {
36        self.accept_incoming = accept_incoming;
37        self
38    }
39
40    pub fn accept_port(mut self, accept_port: u16) -> Self {
41        self.accept_port = Some(accept_port);
42        self
43    }
44
45    pub fn secret_key(mut self, secret_key: &[u8; SECRET_KEY_LENGTH]) -> Self {
46        self.secret_key = *secret_key;
47        self
48    }
49
50    pub fn dot_ssh_integration(mut self) -> Self {
51        if let Ok(secret_key) = dot_ssh(&SecretKey::from_bytes(&self.secret_key)) {
52            self.secret_key = secret_key.to_bytes();
53        }
54        self
55    }
56
57    pub async fn build(self: &mut Self) -> anyhow::Result<IrohSsh> {
58        // Iroh setup
59        let secret_key = SecretKey::from_bytes(&self.secret_key);
60        let endpoint = Endpoint::builder()
61            .secret_key(secret_key)
62            .discovery_n0()
63            .bind()
64            .await?;
65
66        let mut iroh_ssh = IrohSsh {
67            public_key: endpoint.node_id().as_bytes().clone(),
68            secret_key: self.secret_key,
69            inner: None,
70        };
71
72        let router = if self.accept_incoming {
73            Router::builder(endpoint.clone()).accept(&IrohSsh::ALPN(), iroh_ssh.clone())
74        } else {
75            Router::builder(endpoint.clone())
76        }
77        .spawn();
78
79        iroh_ssh.add_inner(endpoint, router);
80
81        if self.accept_incoming && self.accept_port.is_some() {
82            tokio::spawn({
83                let iroh_ssh = iroh_ssh.clone();
84                let accept_port = self.accept_port.expect("accept_port not set");
85                async move {
86                    iroh_ssh._spawn(accept_port).await.expect("spawn failed");
87                }
88            });
89        }
90
91        Ok(iroh_ssh)
92    }
93}
94
95impl IrohSsh {
96    pub fn new() -> Builder {
97        Builder::new()
98    }
99
100    #[allow(non_snake_case)]
101    pub fn ALPN() -> Vec<u8> {
102        format!("/iroh/ssh").into_bytes()
103    }
104
105    fn add_inner(&mut self, endpoint: Endpoint, router: Router) {
106        self.inner = Some(Inner { endpoint, router });
107    }
108
109    pub async fn connect(&self, ssh_user: &str, node_id: NodeId) -> anyhow::Result<Child> {
110        let inner = self.inner.as_ref().expect("inner not set");
111        let conn = inner.endpoint.connect(node_id, &IrohSsh::ALPN()).await?;
112        let listener = TcpListener::bind("127.0.0.1:0").await?;
113        let port = listener.local_addr()?.port();
114
115        tokio::spawn(async move {
116            loop {
117                match listener.accept().await {
118                    Ok((mut stream, _)) => match conn.open_bi().await {
119                        Ok((mut iroh_send, mut iroh_recv)) => {
120                            tokio::spawn(async move {
121                                let (mut local_read, mut local_write) = stream.split();
122                                let a_to_b = async move {
123                                    tokio::io::copy(&mut local_read, &mut iroh_send).await
124                                };
125                                let b_to_a = async move {
126                                    tokio::io::copy(&mut iroh_recv, &mut local_write).await
127                                };
128
129                                tokio::select! {
130                                    result = a_to_b => {
131                                        let _ = result;
132                                    },
133                                    result = b_to_a => {
134                                        let _ = result;
135                                    },
136                                };
137                            });
138                        }
139                        Err(_) => break,
140                    },
141                    Err(_) => break,
142                }
143            }
144        });
145        tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
146        let ssh_process = Command::new("ssh")
147            .arg("-tt") // Force pseudo-terminal allocation
148            .arg(format!("{}@127.0.0.1", ssh_user))
149            .arg("-p")
150            .arg(port.to_string())
151            .arg("-o")
152            .arg("StrictHostKeyChecking=no")
153            .arg("-o")
154            .arg("UserKnownHostsFile=/dev/null")
155            .arg("-o")
156            .arg("LogLevel=ERROR") // Reduce SSH debug output
157            .stdin(Stdio::inherit())
158            .stdout(Stdio::inherit())
159            .stderr(Stdio::inherit())
160            .spawn()?;
161
162        Ok(ssh_process)
163    }
164
165    pub fn node_id(&self) -> NodeId {
166        self.inner
167            .as_ref()
168            .expect("inner not set")
169            .endpoint
170            .node_id()
171    }
172
173    async fn _spawn(self, port: u16) -> anyhow::Result<()> {
174        println!("Server listening for iroh connections...");
175
176        while let Some(incoming) = self
177            .inner
178            .clone()
179            .expect("inner not set")
180            .endpoint
181            .accept()
182            .await
183        {
184            let mut connecting = match incoming.accept() {
185                Ok(connecting) => connecting,
186                Err(err) => {
187                    println!("Incoming connection failure: {err:#}");
188                    continue;
189                }
190            };
191
192            let alpn = ok_or_continue!(connecting.alpn().await);
193            let conn = ok_or_continue!(connecting.await);
194            let node_id = ok_or_continue!(conn.remote_node_id());
195
196            println!("{}: {node_id} connected", String::from_utf8_lossy(&alpn));
197
198            tokio::spawn(async move {
199                match conn.accept_bi().await {
200                    Ok((mut iroh_send, mut iroh_recv)) => {
201                        println!("Accepted bidirectional stream from {}", node_id);
202
203                        match TcpStream::connect(format!("127.0.0.1:{}", port)).await {
204                            Ok(mut ssh_stream) => {
205                                println!("Connected to local SSH server on port {}", port);
206
207                                let (mut local_read, mut local_write) = ssh_stream.split();
208
209                                let a_to_b = async move {
210                                    tokio::io::copy(&mut local_read, &mut iroh_send).await
211                                };
212                                let b_to_a = async move {
213                                    tokio::io::copy(&mut iroh_recv, &mut local_write).await
214                                };
215
216                                tokio::select! {
217                                    result = a_to_b => {
218                                        println!("SSH->Iroh stream ended: {:?}", result);
219                                    },
220                                    result = b_to_a => {
221                                        println!("Iroh->SSH stream ended: {:?}", result);
222                                    },
223                                };
224                            }
225                            Err(e) => {
226                                println!("Failed to connect to SSH server: {}", e);
227                            }
228                        }
229                    }
230                    Err(e) => {
231                        println!("Failed to accept bidirectional stream: {}", e);
232                    }
233                }
234            });
235        }
236        Ok(())
237    }
238}
239
240impl ProtocolHandler for IrohSsh {
241    fn accept(
242        &self,
243        conn: Connection,
244    ) -> Pin<Box<dyn Future<Output = anyhow::Result<()>> + Send + 'static>> {
245        let iroh_ssh = self.clone();
246
247        Box::pin(async move {
248            iroh_ssh.accept(conn).await?;
249            Ok(())
250        })
251    }
252}
253
254pub fn dot_ssh(default_secret_key: &SecretKey) -> anyhow::Result<SecretKey> {
255    let distro_home = my_home()?.ok_or_else(|| anyhow::anyhow!("home directory not found"))?;
256    let ssh_dir = distro_home.join(".ssh");
257    let pub_key = ssh_dir.join("irohssh_ed25519.pub");
258    let priv_key = ssh_dir.join("irohssh_ed25519");
259
260    // check pub and priv key already exists
261    if pub_key.exists() && priv_key.exists() {
262        // read secret key
263        if let Ok(secret_key) = std::fs::read(priv_key.clone()) {
264            let mut sk_bytes = [0u8; SECRET_KEY_LENGTH];
265            sk_bytes.copy_from_slice(z32::decode(secret_key.as_slice())?.as_slice());
266            return Ok(SecretKey::from_bytes(&sk_bytes));
267        }
268    }
269    
270    let key = default_secret_key.clone();
271    let secret_key = key.secret();
272    let public_key = key.public();
273
274    std::fs::write(pub_key, z32::encode(public_key.as_bytes()))?;
275    std::fs::write(priv_key, z32::encode(secret_key.as_bytes()))?;
276
277    Ok(key)
278}
279
iroh_ssh/ssh.rs

iroh_ssh/
ssh.rs
