Add native dynamic port forwarding

src/bin/dosh-client.rs

@@ -24,9 +24,10 @@ use dosh::ssh_agent;
 use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 use std::collections::HashMap;
+use std::collections::HashSet;
 use std::fs;
 use std::io::{Read, Write};
-use std::net::{SocketAddr, ToSocketAddrs};
+use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs};
 use std::path::{Path, PathBuf};
 use std::process::{Command, Stdio};
 use std::sync::Arc;
@@ -65,6 +66,8 @@ struct Args {
     local_forward: Vec<String>,
     #[arg(short = 'R', long = "remote-forward")]
     remote_forward: Vec<String>,
+    #[arg(short = 'D', long = "dynamic-forward")]
+    dynamic_forward: Vec<String>,
     #[arg(short = 'N', long)]
     forward_only: bool,
     #[arg(long)]
@@ -118,12 +121,19 @@ struct RemoteForward {
     target_port: u16,
 }
 
+#[derive(Debug, Clone, PartialEq, Eq)]
+struct DynamicForward {
+    bind_host: String,
+    listen_port: u16,
+}
+
 enum ForwardEvent {
     Open {
         stream_id: u64,
         target_host: String,
         target_port: u16,
         writer: tokio::net::tcp::OwnedWriteHalf,
+        socks5_reply: bool,
     },
     Data {
         stream_id: u64,
@@ -162,7 +172,9 @@ async fn main() -> Result<()> {
     });
     let local_forwards = parse_local_forwards(&args.local_forward)?;
     let remote_forwards = parse_remote_forwards(&args.remote_forward)?;
-    let forwarding_requested = !local_forwards.is_empty() || !remote_forwards.is_empty();
+    let dynamic_forwards = parse_dynamic_forwards(&args.dynamic_forward)?;
+    let forwarding_requested =
+        !local_forwards.is_empty() || !remote_forwards.is_empty() || !dynamic_forwards.is_empty();
     let predict = args.predict || host.predict.unwrap_or(config.predict);
     let session = select_session(args.session.as_deref(), args.new);
     let mode = if args.view_only || config.view_only {
@@ -257,6 +269,7 @@ async fn main() -> Result<()> {
                         startup_command,
                         config.reconnect_timeout_secs,
                         Vec::new(),
+                        Vec::new(),
                         false,
                     )
                     .await;
@@ -291,6 +304,7 @@ async fn main() -> Result<()> {
                         startup_command,
                         config.reconnect_timeout_secs,
                         Vec::new(),
+                        Vec::new(),
                         false,
                     )
                     .await;
@@ -321,7 +335,7 @@ async fn main() -> Result<()> {
             &mode,
             cols,
             rows,
-            forwarding_requests(&local_forwards, &remote_forwards),
+            forwarding_requests(&local_forwards, &remote_forwards, &dynamic_forwards),
         )
         .await
         {
@@ -344,6 +358,7 @@ async fn main() -> Result<()> {
                     startup_command,
                     config.reconnect_timeout_secs,
                     local_forwards,
+                    dynamic_forwards,
                     args.forward_only,
                 )
                 .await;
@@ -413,6 +428,7 @@ async fn main() -> Result<()> {
         startup_command,
         config.reconnect_timeout_secs,
         Vec::new(),
+        Vec::new(),
         false,
     )
     .await
@@ -533,6 +549,12 @@ fn parse_remote_forwards(raw: &[String]) -> Result<Vec<RemoteForward>> {
         .collect()
 }
 
+fn parse_dynamic_forwards(raw: &[String]) -> Result<Vec<DynamicForward>> {
+    raw.iter()
+        .map(|value| parse_dynamic_forward(value))
+        .collect()
+}
+
 fn parse_local_forward(raw: &str) -> Result<LocalForward> {
     let parts = raw.split(':').collect::<Vec<_>>();
     let (bind_host, listen, target_host, target_port) = match parts.as_slice() {
@@ -581,9 +603,33 @@ fn parse_remote_forward(raw: &str) -> Result<RemoteForward> {
     })
 }
 
+fn parse_dynamic_forward(raw: &str) -> Result<DynamicForward> {
+    let parts = raw.split(':').collect::<Vec<_>>();
+    let (bind_host, listen) = match parts.as_slice() {
+        [listen] => ("127.0.0.1".to_string(), *listen),
+        [bind_host, listen] => ((*bind_host).to_string(), *listen),
+        _ => {
+            return Err(anyhow!(
+                "invalid -D {raw:?}; expected listen_port or bind_host:listen_port"
+            ));
+        }
+    };
+    let listen_port = listen
+        .parse::<u16>()
+        .with_context(|| format!("invalid dynamic forward listen port in {raw:?}"))?;
+    if bind_host.is_empty() {
+        return Err(anyhow!("invalid -D {raw:?}; bind host cannot be empty"));
+    }
+    Ok(DynamicForward {
+        bind_host,
+        listen_port,
+    })
+}
+
 fn forwarding_requests(
     local_forwards: &[LocalForward],
     remote_forwards: &[RemoteForward],
+    dynamic_forwards: &[DynamicForward],
 ) -> Vec<ForwardingRequest> {
     let mut requests = local_forwards
         .iter()
@@ -602,6 +648,13 @@ fn forwarding_requests(
         target_host: Some(forward.target_host.clone()),
         target_port: Some(forward.target_port),
     }));
+    requests.extend(dynamic_forwards.iter().map(|forward| ForwardingRequest {
+        kind: ForwardingKind::Dynamic,
+        bind_host: Some(forward.bind_host.clone()),
+        listen_port: forward.listen_port,
+        target_host: None,
+        target_port: None,
+    }));
     requests
 }
 
@@ -1434,6 +1487,7 @@ async fn run_terminal(
     startup_command: Option<Vec<u8>>,
     reconnect_timeout_secs: u64,
     local_forwards: Vec<LocalForward>,
+    dynamic_forwards: Vec<DynamicForward>,
     forward_only: bool,
 ) -> Result<()> {
     let _raw = if forward_only {
@@ -1457,8 +1511,10 @@ async fn run_terminal(
     };
     let remote_forward_tx = forward_tx.clone();
     let stream_ids = Arc::new(AtomicU64::new(1));
-    start_local_forwards(&local_forwards, forward_tx.clone(), stream_ids).await?;
+    start_local_forwards(&local_forwards, forward_tx.clone(), Arc::clone(&stream_ids)).await?;
+    start_dynamic_forwards(&dynamic_forwards, forward_tx.clone(), stream_ids).await?;
     let mut stream_writers: HashMap<u64, tokio::net::tcp::OwnedWriteHalf> = HashMap::new();
+    let mut pending_socks_replies: HashSet<u64> = HashSet::new();
     if let Some(frame) = first_frame {
         if !forward_only {
             render_frame(&frame)?;
@@ -1656,10 +1712,15 @@ async fn run_terminal(
                         let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
                             continue;
                         };
-                        let Ok(_ok) = protocol::from_body::<StreamOpenOk>(&plain) else {
+                        let Ok(ok) = protocol::from_body::<StreamOpenOk>(&plain) else {
                             continue;
                         };
                         last_packet_at = Instant::now();
+                        if pending_socks_replies.remove(&ok.stream_id)
+                            && let Some(writer) = stream_writers.get_mut(&ok.stream_id)
+                        {
+                            let _ = writer.write_all(&[5, 0, 0, 1, 0, 0, 0, 0, 0, 0]).await;
+                        }
                     }
                     PacketKind::StreamOpenReject => {
                         let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
@@ -1669,6 +1730,11 @@ async fn run_terminal(
                             continue;
                         };
                         last_packet_at = Instant::now();
+                        if pending_socks_replies.remove(&reject.stream_id)
+                            && let Some(writer) = stream_writers.get_mut(&reject.stream_id)
+                        {
+                            let _ = writer.write_all(&[5, 5, 0, 1, 0, 0, 0, 0, 0, 0]).await;
+                        }
                         stream_writers.remove(&reject.stream_id);
                     }
                     PacketKind::StreamData => {
@@ -1691,6 +1757,7 @@ async fn run_terminal(
                             continue;
                         };
                         last_packet_at = Instant::now();
+                        pending_socks_replies.remove(&close.stream_id);
                         stream_writers.remove(&close.stream_id);
                     }
                     _ => {}
@@ -1698,8 +1765,11 @@ async fn run_terminal(
             }
             forward_event = forward_rx.recv() => {
                 match forward_event {
-                    Some(ForwardEvent::Open { stream_id, target_host, target_port, writer }) => {
+                    Some(ForwardEvent::Open { stream_id, target_host, target_port, writer, socks5_reply }) => {
                         stream_writers.insert(stream_id, writer);
+                        if socks5_reply {
+                            pending_socks_replies.insert(stream_id);
+                        }
                         send_stream_open(
                             &socket,
                             addr,
@@ -1715,6 +1785,7 @@ async fn run_terminal(
                         send_stream_data(&socket, addr, &cred, &mut send_seq, stream_id, bytes).await?;
                     }
                     Some(ForwardEvent::Close { stream_id }) => {
+                        pending_socks_replies.remove(&stream_id);
                         stream_writers.remove(&stream_id);
                         send_stream_close(&socket, addr, &cred, &mut send_seq, stream_id).await?;
                     }
@@ -1787,6 +1858,7 @@ async fn start_local_forwards(
                         target_host: forward.target_host.clone(),
                         target_port: forward.target_port,
                         writer,
+                        socks5_reply: false,
                     })
                     .await;
                 let forward_tx = forward_tx.clone();
@@ -1818,6 +1890,124 @@ async fn start_local_forwards(
     Ok(())
 }
 
+async fn start_dynamic_forwards(
+    dynamic_forwards: &[DynamicForward],
+    forward_tx: mpsc::Sender<ForwardEvent>,
+    stream_ids: Arc<AtomicU64>,
+) -> Result<()> {
+    for forward in dynamic_forwards {
+        let bind = format!("{}:{}", forward.bind_host, forward.listen_port);
+        let listener = TcpListener::bind(&bind)
+            .await
+            .with_context(|| format!("bind dynamic forward {bind}"))?;
+        let forward_tx = forward_tx.clone();
+        let stream_ids = Arc::clone(&stream_ids);
+        tokio::spawn(async move {
+            loop {
+                let Ok((stream, _)) = listener.accept().await else {
+                    break;
+                };
+                let forward_tx = forward_tx.clone();
+                let stream_ids = Arc::clone(&stream_ids);
+                tokio::spawn(async move {
+                    let _ = handle_socks5_connect(stream, forward_tx, stream_ids).await;
+                });
+            }
+        });
+    }
+    Ok(())
+}
+
+async fn handle_socks5_connect(
+    mut stream: TcpStream,
+    forward_tx: mpsc::Sender<ForwardEvent>,
+    stream_ids: Arc<AtomicU64>,
+) -> Result<()> {
+    let mut hello = [0u8; 2];
+    stream.read_exact(&mut hello).await?;
+    anyhow::ensure!(hello[0] == 5, "unsupported SOCKS version {}", hello[0]);
+    let mut methods = vec![0u8; hello[1] as usize];
+    stream.read_exact(&mut methods).await?;
+    if !methods.contains(&0) {
+        stream.write_all(&[5, 0xff]).await?;
+        return Ok(());
+    }
+    stream.write_all(&[5, 0]).await?;
+
+    let mut header = [0u8; 4];
+    stream.read_exact(&mut header).await?;
+    anyhow::ensure!(
+        header[0] == 5,
+        "unsupported SOCKS request version {}",
+        header[0]
+    );
+    if header[1] != 1 {
+        stream.write_all(&[5, 7, 0, 1, 0, 0, 0, 0, 0, 0]).await?;
+        return Ok(());
+    }
+    let target_host = match header[3] {
+        1 => {
+            let mut raw = [0u8; 4];
+            stream.read_exact(&mut raw).await?;
+            Ipv4Addr::from(raw).to_string()
+        }
+        3 => {
+            let mut len = [0u8; 1];
+            stream.read_exact(&mut len).await?;
+            let mut raw = vec![0u8; len[0] as usize];
+            stream.read_exact(&mut raw).await?;
+            String::from_utf8(raw).context("SOCKS domain is not valid UTF-8")?
+        }
+        4 => {
+            let mut raw = [0u8; 16];
+            stream.read_exact(&mut raw).await?;
+            Ipv6Addr::from(raw).to_string()
+        }
+        atyp => {
+            stream.write_all(&[5, 8, 0, 1, 0, 0, 0, 0, 0, 0]).await?;
+            return Err(anyhow!("unsupported SOCKS address type {atyp}"));
+        }
+    };
+    let mut port = [0u8; 2];
+    stream.read_exact(&mut port).await?;
+    let target_port = u16::from_be_bytes(port);
+
+    let stream_id = stream_ids.fetch_add(1, Ordering::Relaxed);
+    let (mut reader, writer) = stream.into_split();
+    forward_tx
+        .send(ForwardEvent::Open {
+            stream_id,
+            target_host,
+            target_port,
+            writer,
+            socks5_reply: true,
+        })
+        .await?;
+    tokio::spawn(async move {
+        let mut buf = [0u8; 16 * 1024];
+        loop {
+            match reader.read(&mut buf).await {
+                Ok(0) => break,
+                Ok(n) => {
+                    if forward_tx
+                        .send(ForwardEvent::Data {
+                            stream_id,
+                            bytes: buf[..n].to_vec(),
+                        })
+                        .await
+                        .is_err()
+                    {
+                        return;
+                    }
+                }
+                Err(_) => break,
+            }
+        }
+        let _ = forward_tx.send(ForwardEvent::Close { stream_id }).await;
+    });
+    Ok(())
+}
+
 async fn reconnect(
     socket: &UdpSocket,
     cred: &mut CachedCredential,
@@ -2225,10 +2415,10 @@ const TERMINAL_CLEANUP: &[u8] = concat!(
 #[cfg(test)]
 mod tests {
     use super::{
-        FrameBuffer, LocalForward, Predictor, RemoteForward, SshConfig, auth_allows,
-        load_first_native_identity_with_prompt, parse_local_forward, parse_remote_forward,
-        parse_ssh_config, raw_contains_host_table, ssh_destination_host, ssh_username,
-        startup_command, toml_bare_key_or_quoted,
+        DynamicForward, FrameBuffer, LocalForward, Predictor, RemoteForward, SshConfig,
+        auth_allows, load_first_native_identity_with_prompt, parse_dynamic_forward,
+        parse_local_forward, parse_remote_forward, parse_ssh_config, raw_contains_host_table,
+        ssh_destination_host, ssh_username, startup_command, toml_bare_key_or_quoted,
     };
     use dosh::config::ClientConfig;
     use dosh::protocol::Frame;
@@ -2478,4 +2668,15 @@ mod tests {
             }
         );
     }
+
+    #[test]
+    fn parses_dynamic_forward_with_default_bind() {
+        assert_eq!(
+            parse_dynamic_forward("1080").unwrap(),
+            DynamicForward {
+                bind_host: "127.0.0.1".to_string(),
+                listen_port: 1080,
+            }
+        );
+    }
 }