Add native remote port forwarding

2026-06-13 15:41:48 -04:00
parent 6aa81d0ce3
commit 6a0f12c59e
3 changed files with 451 additions and 15 deletions
@@ -17,8 +17,8 @@ use dosh::native::{
 use dosh::protocol::{
    self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
    NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody, NativeUserAuthBody, PacketKind,
-    Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose, StreamData, StreamOpen, StreamOpenReject,
-    TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
+    Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose, StreamData, StreamOpen, StreamOpenOk,
+    StreamOpenReject, TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope,
 };
 use dosh::ssh_agent;
 use serde::{Deserialize, Serialize};
@@ -33,7 +33,7 @@ use std::sync::Arc;
 use std::sync::atomic::{AtomicU64, Ordering};
 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
 use tokio::io::{AsyncReadExt, AsyncWriteExt};
-use tokio::net::{TcpListener, UdpSocket};
+use tokio::net::{TcpListener, TcpStream, UdpSocket};
 use tokio::sync::mpsc;

 #[derive(Debug, Parser)]
@@ -63,6 +63,8 @@ struct Args {
    attach_only: bool,
    #[arg(short = 'L', long = "local-forward")]
    local_forward: Vec<String>,
+    #[arg(short = 'R', long = "remote-forward")]
+    remote_forward: Vec<String>,
    #[arg(short = 'N', long)]
    forward_only: bool,
    #[arg(long)]
@@ -108,6 +110,14 @@ struct LocalForward {
    target_port: u16,
 }

+#[derive(Debug, Clone, PartialEq, Eq)]
+struct RemoteForward {
+    bind_host: String,
+    listen_port: u16,
+    target_host: String,
+    target_port: u16,
+}
+
 enum ForwardEvent {
    Open {
        stream_id: u64,
@@ -151,7 +161,8 @@ async fn main() -> Result<()> {
            .map(startup_command_from_string)
    });
    let local_forwards = parse_local_forwards(&args.local_forward)?;
-    let forwarding_requested = !local_forwards.is_empty();
+    let remote_forwards = parse_remote_forwards(&args.remote_forward)?;
+    let forwarding_requested = !local_forwards.is_empty() || !remote_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 {
@@ -310,7 +321,7 @@ async fn main() -> Result<()> {
            &mode,
            cols,
            rows,
-            forwarding_requests(&local_forwards),
+            forwarding_requests(&local_forwards, &remote_forwards),
        )
        .await
        {
@@ -516,6 +527,12 @@ fn parse_local_forwards(raw: &[String]) -> Result<Vec<LocalForward>> {
    raw.iter().map(|value| parse_local_forward(value)).collect()
 }

+fn parse_remote_forwards(raw: &[String]) -> Result<Vec<RemoteForward>> {
+    raw.iter()
+        .map(|value| parse_remote_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() {
@@ -554,8 +571,21 @@ fn parse_local_forward(raw: &str) -> Result<LocalForward> {
    })
 }

-fn forwarding_requests(local_forwards: &[LocalForward]) -> Vec<ForwardingRequest> {
-    local_forwards
+fn parse_remote_forward(raw: &str) -> Result<RemoteForward> {
+    let local = parse_local_forward(raw)?;
+    Ok(RemoteForward {
+        bind_host: local.bind_host,
+        listen_port: local.listen_port,
+        target_host: local.target_host,
+        target_port: local.target_port,
+    })
+}
+
+fn forwarding_requests(
+    local_forwards: &[LocalForward],
+    remote_forwards: &[RemoteForward],
+) -> Vec<ForwardingRequest> {
+    let mut requests = local_forwards
        .iter()
        .map(|forward| ForwardingRequest {
            kind: ForwardingKind::Local,
@@ -564,7 +594,15 @@ fn forwarding_requests(local_forwards: &[LocalForward]) -> Vec<ForwardingRequest
            target_host: Some(forward.target_host.clone()),
            target_port: Some(forward.target_port),
        })
-        .collect()
+        .collect::<Vec<_>>();
+    requests.extend(remote_forwards.iter().map(|forward| ForwardingRequest {
+        kind: ForwardingKind::Remote,
+        bind_host: Some(forward.bind_host.clone()),
+        listen_port: forward.listen_port,
+        target_host: Some(forward.target_host.clone()),
+        target_port: Some(forward.target_port),
+    }));
+    requests
 }

 fn run_sessions_command(
@@ -1417,8 +1455,9 @@ async fn run_terminal(
    } else {
        None
    };
+    let remote_forward_tx = forward_tx.clone();
    let stream_ids = Arc::new(AtomicU64::new(1));
-    start_local_forwards(&local_forwards, forward_tx, stream_ids).await?;
+    start_local_forwards(&local_forwards, forward_tx.clone(), stream_ids).await?;
    let mut stream_writers: HashMap<u64, tokio::net::tcp::OwnedWriteHalf> = HashMap::new();
    if let Some(frame) = first_frame {
        if !forward_only {
@@ -1560,7 +1599,66 @@ async fn run_terminal(
                            return Err(anyhow!("server rejected session: {}", reject.reason));
                        }
                    }
-                    PacketKind::StreamOpenOk => {}
+                    PacketKind::StreamOpen => {
+                        let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
+                            continue;
+                        };
+                        let Ok(open) = protocol::from_body::<StreamOpen>(&plain) else {
+                            continue;
+                        };
+                        last_packet_at = Instant::now();
+                        match TcpStream::connect((open.target_host.as_str(), open.target_port)).await {
+                            Ok(stream) => {
+                                let (mut reader, writer) = stream.into_split();
+                                stream_writers.insert(open.stream_id, writer);
+                                send_stream_open_ok(&socket, addr, &cred, &mut send_seq, open.stream_id).await?;
+                                let forward_tx = remote_forward_tx.clone();
+                                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: open.stream_id,
+                                                        bytes: buf[..n].to_vec(),
+                                                    })
+                                                    .is_err()
+                                                {
+                                                    return;
+                                                }
+                                            }
+                                            Err(_) => break,
+                                        }
+                                    }
+                                    let _ = forward_tx.send(ForwardEvent::Close {
+                                        stream_id: open.stream_id,
+                                    });
+                                });
+                            }
+                            Err(err) => {
+                                send_stream_open_reject(
+                                    &socket,
+                                    addr,
+                                    &cred,
+                                    &mut send_seq,
+                                    open.stream_id,
+                                    format!("connect {}:{}: {err}", open.target_host, open.target_port),
+                                )
+                                .await?;
+                            }
+                        }
+                    }
+                    PacketKind::StreamOpenOk => {
+                        let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
+                            continue;
+                        };
+                        let Ok(_ok) = protocol::from_body::<StreamOpenOk>(&plain) else {
+                            continue;
+                        };
+                        last_packet_at = Instant::now();
+                    }
                    PacketKind::StreamOpenReject => {
                        let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
                            continue;
@@ -1795,6 +1893,45 @@ async fn send_stream_open(
    send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamOpen, &body).await
 }

+async fn send_stream_open_ok(
+    socket: &UdpSocket,
+    addr: SocketAddr,
+    cred: &CachedCredential,
+    send_seq: &mut u64,
+    stream_id: u64,
+) -> Result<()> {
+    let body = protocol::to_body(&StreamOpenOk { stream_id })?;
+    send_stream_packet(
+        socket,
+        addr,
+        cred,
+        send_seq,
+        PacketKind::StreamOpenOk,
+        &body,
+    )
+    .await
+}
+
+async fn send_stream_open_reject(
+    socket: &UdpSocket,
+    addr: SocketAddr,
+    cred: &CachedCredential,
+    send_seq: &mut u64,
+    stream_id: u64,
+    reason: String,
+) -> Result<()> {
+    let body = protocol::to_body(&StreamOpenReject { stream_id, reason })?;
+    send_stream_packet(
+        socket,
+        addr,
+        cred,
+        send_seq,
+        PacketKind::StreamOpenReject,
+        &body,
+    )
+    .await
+}
+
 async fn send_stream_data(
    socket: &UdpSocket,
    addr: SocketAddr,
@@ -2083,10 +2220,10 @@ const TERMINAL_CLEANUP: &[u8] = concat!(
 #[cfg(test)]
 mod tests {
    use super::{
-        FrameBuffer, LocalForward, Predictor, SshConfig, auth_allows,
-        load_first_native_identity_with_prompt, parse_local_forward, parse_ssh_config,
-        raw_contains_host_table, ssh_destination_host, ssh_username, startup_command,
-        toml_bare_key_or_quoted,
+        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,
    };
    use dosh::config::ClientConfig;
    use dosh::protocol::Frame;
@@ -2323,4 +2460,17 @@ mod tests {
            }
        );
    }
+
+    #[test]
+    fn parses_remote_forward_with_default_bind() {
+        assert_eq!(
+            parse_remote_forward("9090:127.0.0.1:5432").unwrap(),
+            RemoteForward {
+                bind_host: "127.0.0.1".to_string(),
+                listen_port: 9090,
+                target_host: "127.0.0.1".to_string(),
+                target_port: 5432,
+            }
+        );
+    }
 }