From 6aa81d0ce3af6521bd3c4b55f04becee6d48c62c Mon Sep 17 00:00:00 2001 From: DuProcess <273172371+DuProcess@users.noreply.github.com> Date: Sat, 13 Jun 2026 15:36:10 -0400 Subject: [PATCH] Add native local port forwarding --- src/bin/dosh-client.rs | 441 +++++++++++++++++++++++++++++++++---- src/bin/dosh-server.rs | 301 ++++++++++++++++++++++++- src/protocol.rs | 54 +++++ tests/integration_smoke.rs | 103 ++++++++- 4 files changed, 844 insertions(+), 55 deletions(-) diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index 5a1a8e6..aeb9a46 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -9,27 +9,31 @@ use dosh::auth::{ use dosh::config::{expand_tilde, load_client_config, load_hosts_config, load_server_config}; use dosh::crypto; use dosh::native::{ - KnownHostStatus, TrustResult, derive_native_session_key, generate_native_ephemeral, - host_fingerprint, load_ed25519_identity, load_ed25519_identity_with_passphrase, - parse_host_public_key_line, remove_trusted_host, sign_user_auth, trust_host, verify_known_host, - verify_server_hello, + ForwardingKind, ForwardingRequest, KnownHostStatus, TrustResult, derive_native_session_key, + generate_native_ephemeral, host_fingerprint, load_ed25519_identity, + load_ed25519_identity_with_passphrase, parse_host_public_key_line, remove_trusted_host, + sign_user_auth, trust_host, verify_known_host, verify_server_hello, }; use dosh::protocol::{ self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody, NativeUserAuthBody, PacketKind, - Resize, ResumeRequest, SERVER_TO_CLIENT, TicketAttachBody, TicketAttachEnvelope, - TicketAttachOkEnvelope, + Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose, StreamData, StreamOpen, StreamOpenReject, + TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, }; use dosh::ssh_agent; use serde::{Deserialize, Serialize}; use std::collections::BTreeMap; +use std::collections::HashMap; use std::fs; use std::io::{Read, Write}; use std::net::{SocketAddr, ToSocketAddrs}; use std::path::{Path, PathBuf}; use std::process::{Command, Stdio}; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; -use tokio::net::UdpSocket; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpListener, UdpSocket}; use tokio::sync::mpsc; #[derive(Debug, Parser)] @@ -57,6 +61,10 @@ struct Args { replace: bool, #[arg(long)] attach_only: bool, + #[arg(short = 'L', long = "local-forward")] + local_forward: Vec, + #[arg(short = 'N', long)] + forward_only: bool, #[arg(long)] predict: bool, #[arg(long)] @@ -92,6 +100,30 @@ struct CachedCredential { last_rendered_seq: u64, } +#[derive(Debug, Clone, PartialEq, Eq)] +struct LocalForward { + bind_host: String, + listen_port: u16, + target_host: String, + target_port: u16, +} + +enum ForwardEvent { + Open { + stream_id: u64, + target_host: String, + target_port: u16, + writer: tokio::net::tcp::OwnedWriteHalf, + }, + Data { + stream_id: u64, + bytes: Vec, + }, + Close { + stream_id: u64, + }, +} + #[tokio::main(flavor = "current_thread")] async fn main() -> Result<()> { let args = Args::parse(); @@ -118,6 +150,8 @@ async fn main() -> Result<()> { .as_deref() .map(startup_command_from_string) }); + let local_forwards = parse_local_forwards(&args.local_forward)?; + let forwarding_requested = !local_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 { @@ -147,6 +181,19 @@ async fn main() -> Result<()> { let dosh_port = args.dosh_port.or(host.port).unwrap_or(config.dosh_port); let cache_path = cache_path(&config.credential_cache, &requested_server, &session, &mode); let (cols, rows) = size().unwrap_or((80, 24)); + let auth_preference = args + .auth + .clone() + .unwrap_or_else(|| config.auth_preference.clone()); + if forwarding_requested && args.local_auth { + return Err(anyhow!( + "port forwarding requires native auth; remove --local-auth" + )); + } + if forwarding_requested && !auth_allows(&auth_preference, "native") { + return Err(anyhow!("port forwarding requires --auth native or auto")); + } + let allow_cache = !args.no_cache && !forwarding_requested; let started = Instant::now(); let target_udp_host = if let Some(host) = args @@ -169,7 +216,7 @@ async fn main() -> Result<()> { }) }; - let credential = if !args.no_cache { + let credential = if allow_cache { load_cache(&cache_path).ok() } else { None @@ -181,7 +228,7 @@ async fn main() -> Result<()> { cached.udp_host = target_udp_host.clone(); cached.udp_port = dosh_port; log_timing(args.verbose, "credential_lookup_end", started.elapsed()); - if config.cache_attach_tickets && !args.no_cache { + if config.cache_attach_tickets && allow_cache { match try_ticket_attach(&socket, &cached, cols, rows).await { Ok((frame, cred)) => { log_timing(args.verbose, "udp_ticket_attach_ready", started.elapsed()); @@ -198,6 +245,8 @@ async fn main() -> Result<()> { predict, startup_command, config.reconnect_timeout_secs, + Vec::new(), + false, ) .await; } @@ -215,7 +264,7 @@ async fn main() -> Result<()> { match try_resume_fresh_process(&socket, &cached, cols, rows).await { Ok((frame, cred)) => { log_timing(args.verbose, "udp_resume_ready", started.elapsed()); - if !args.no_cache { + if allow_cache { save_cache(&cache_path, &cred)?; } if args.attach_only { @@ -230,6 +279,8 @@ async fn main() -> Result<()> { predict, startup_command, config.reconnect_timeout_secs, + Vec::new(), + false, ) .await; } @@ -244,10 +295,6 @@ async fn main() -> Result<()> { } } - let auth_preference = args - .auth - .clone() - .unwrap_or_else(|| config.auth_preference.clone()); if !args.local_auth && auth_allows(&auth_preference, "native") { let native_start = Instant::now(); match try_native_auth( @@ -263,12 +310,13 @@ async fn main() -> Result<()> { &mode, cols, rows, + forwarding_requests(&local_forwards), ) .await { Ok((frame, cred)) => { log_timing(args.verbose, "native_auth", native_start.elapsed()); - if !args.no_cache { + if allow_cache { save_cache(&cache_path, &cred)?; } log_timing(args.verbose, "terminal_ready", started.elapsed()); @@ -284,10 +332,12 @@ async fn main() -> Result<()> { predict, startup_command, config.reconnect_timeout_secs, + local_forwards, + args.forward_only, ) .await; } - Err(err) if auth_allows(&auth_preference, "ssh") => { + Err(err) if !forwarding_requested && auth_allows(&auth_preference, "ssh") => { log_debug( args.verbose, 2, @@ -328,7 +378,7 @@ async fn main() -> Result<()> { let (ok, mut cred) = bootstrap_attach(&socket, &server, &bootstrap, cols, rows).await?; cred.last_rendered_seq = ok.initial_seq; - if !args.no_cache { + if allow_cache { save_cache(&cache_path, &cred)?; } log_timing(args.verbose, "terminal_ready", started.elapsed()); @@ -351,6 +401,8 @@ async fn main() -> Result<()> { predict, startup_command, config.reconnect_timeout_secs, + Vec::new(), + false, ) .await } @@ -460,6 +512,61 @@ fn startup_command_from_string(command: &str) -> Vec { command.into_bytes() } +fn parse_local_forwards(raw: &[String]) -> Result> { + raw.iter().map(|value| parse_local_forward(value)).collect() +} + +fn parse_local_forward(raw: &str) -> Result { + let parts = raw.split(':').collect::>(); + let (bind_host, listen, target_host, target_port) = match parts.as_slice() { + [listen, target_host, target_port] => { + ("127.0.0.1".to_string(), *listen, *target_host, *target_port) + } + [bind_host, listen, target_host, target_port] => ( + (*bind_host).to_string(), + *listen, + *target_host, + *target_port, + ), + _ => { + return Err(anyhow!( + "invalid -L {raw:?}; expected listen_port:target_host:target_port or bind_host:listen_port:target_host:target_port" + )); + } + }; + let listen_port = listen + .parse::() + .with_context(|| format!("invalid local forward listen port in {raw:?}"))?; + let target_port = target_port + .parse::() + .with_context(|| format!("invalid local forward target port in {raw:?}"))?; + if target_host.is_empty() { + return Err(anyhow!("invalid -L {raw:?}; target host cannot be empty")); + } + if bind_host.is_empty() { + return Err(anyhow!("invalid -L {raw:?}; bind host cannot be empty")); + } + Ok(LocalForward { + bind_host, + listen_port, + target_host: target_host.to_string(), + target_port, + }) +} + +fn forwarding_requests(local_forwards: &[LocalForward]) -> Vec { + local_forwards + .iter() + .map(|forward| ForwardingRequest { + kind: ForwardingKind::Local, + 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), + }) + .collect() +} + fn run_sessions_command( server: &str, ssh_port: Option, @@ -822,6 +929,7 @@ async fn try_native_auth( mode: &str, cols: u16, rows: u16, + requested_forwardings: Vec, ) -> Result<(Frame, CachedCredential)> { let addr = resolve_addr(udp_host, udp_port)?; let ssh_config = ssh_config(server, ssh_port).unwrap_or_default(); @@ -907,6 +1015,7 @@ async fn try_native_auth( &ssh_config, &hello, &server_hello.hello, + requested_forwardings, )?; let mut pending_id = [0u8; 16]; pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]); @@ -966,17 +1075,22 @@ fn sign_native_user_auth( ssh_config: &SshConfig, hello: &dosh::native::NativeClientHello, server_hello: &dosh::native::NativeServerHello, + requested_forwardings: Vec, ) -> Result { let mut errors = Vec::new(); if config.use_ssh_agent { - match ssh_agent::sign_user_auth_with_agent(hello, server_hello, Vec::new()) { + match ssh_agent::sign_user_auth_with_agent( + hello, + server_hello, + requested_forwardings.clone(), + ) { Ok(auth) => return Ok(auth), Err(err) => errors.push(format!("ssh-agent: {err:#}")), } } match load_first_native_identity(config, cli_identity, ssh_config) { - Ok(identity) => sign_user_auth(&identity, hello, server_hello, Vec::new()), + Ok(identity) => sign_user_auth(&identity, hello, server_hello, requested_forwardings), Err(err) => { errors.push(format!("identity files: {err:#}")); Err(anyhow!( @@ -1281,8 +1395,14 @@ async fn run_terminal( predict: bool, startup_command: Option>, reconnect_timeout_secs: u64, + local_forwards: Vec, + forward_only: bool, ) -> Result<()> { - let _raw = RawMode::enter()?; + let _raw = if forward_only { + None + } else { + Some(RawMode::enter()?) + }; let addr = resolve_addr(&cred.udp_host, cred.udp_port)?; let mut send_seq = 2u64; let mut last_packet_at = Instant::now(); @@ -1290,10 +1410,21 @@ async fn run_terminal( let mut resize_tick = tokio::time::interval(Duration::from_millis(250)); let mut last_size = size().unwrap_or((80, 24)); let mut frame_buffer = FrameBuffer::default(); - let mut predictor = Predictor::new(predict && cred.mode != "view-only"); + let mut predictor = Predictor::new(predict && cred.mode != "view-only" && !forward_only); + let (forward_tx, mut forward_rx) = mpsc::unbounded_channel::(); + let _forward_keepalive = if local_forwards.is_empty() { + Some(forward_tx.clone()) + } else { + None + }; + let stream_ids = Arc::new(AtomicU64::new(1)); + start_local_forwards(&local_forwards, forward_tx, stream_ids).await?; + let mut stream_writers: HashMap = HashMap::new(); if let Some(frame) = first_frame { - render_frame(&frame)?; - predictor.observe_output(&frame.bytes); + if !forward_only { + render_frame(&frame)?; + predictor.observe_output(&frame.bytes); + } if frame.closed { return Ok(()); } @@ -1302,26 +1433,32 @@ async fn run_terminal( } if let Some(bytes) = startup_command && cred.mode != "view-only" + && !forward_only { send_input(&socket, addr, &cred, &mut send_seq, bytes).await?; } let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::>(); - std::thread::Builder::new() - .name("dosh-stdin".to_string()) - .spawn(move || { - let mut stdin = std::io::stdin(); - let mut buf = [0u8; 4096]; - loop { - match stdin.read(&mut buf) { - Ok(0) => break, - Ok(n) => { - let _ = stdin_tx.send(buf[..n].to_vec()); + let _stdin_keepalive = if forward_only { + Some(stdin_tx) + } else { + std::thread::Builder::new() + .name("dosh-stdin".to_string()) + .spawn(move || { + let mut stdin = std::io::stdin(); + let mut buf = [0u8; 4096]; + loop { + match stdin.read(&mut buf) { + Ok(0) => break, + Ok(n) => { + let _ = stdin_tx.send(buf[..n].to_vec()); + } + Err(_) => break, } - Err(_) => break, } - } - })?; + })?; + None + }; let mut recv_buf = vec![0u8; 65535]; loop { @@ -1371,8 +1508,10 @@ async fn run_terminal( ) .await? { - render_frame(&frame)?; - predictor.observe_output(&frame.bytes); + if !forward_only { + render_frame(&frame)?; + predictor.observe_output(&frame.bytes); + } last_packet_at = Instant::now(); } continue; @@ -1384,8 +1523,10 @@ async fn run_terminal( let frames = frame_buffer.accept(frame, &mut cred.last_rendered_seq); for frame in frames { predictor.clear_pending()?; - render_frame(&frame)?; - predictor.observe_output(&frame.bytes); + if !forward_only { + render_frame(&frame)?; + predictor.observe_output(&frame.bytes); + } if frame.closed { send_ack(&socket, addr, &cred, &mut send_seq).await?; return Ok(()); @@ -1409,17 +1550,78 @@ async fn run_terminal( ) .await? { - render_frame(&frame)?; - predictor.observe_output(&frame.bytes); + if !forward_only { + render_frame(&frame)?; + predictor.observe_output(&frame.bytes); + } last_packet_at = Instant::now(); } } else { return Err(anyhow!("server rejected session: {}", reject.reason)); } } + PacketKind::StreamOpenOk => {} + PacketKind::StreamOpenReject => { + let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { + continue; + }; + let Ok(reject) = protocol::from_body::(&plain) else { + continue; + }; + last_packet_at = Instant::now(); + stream_writers.remove(&reject.stream_id); + } + PacketKind::StreamData => { + let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { + continue; + }; + let Ok(data) = protocol::from_body::(&plain) else { + continue; + }; + last_packet_at = Instant::now(); + if let Some(writer) = stream_writers.get_mut(&data.stream_id) { + let _ = writer.write_all(&data.bytes).await; + } + } + PacketKind::StreamClose => { + let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else { + continue; + }; + let Ok(close) = protocol::from_body::(&plain) else { + continue; + }; + last_packet_at = Instant::now(); + stream_writers.remove(&close.stream_id); + } _ => {} } } + forward_event = forward_rx.recv() => { + match forward_event { + Some(ForwardEvent::Open { stream_id, target_host, target_port, writer }) => { + stream_writers.insert(stream_id, writer); + send_stream_open( + &socket, + addr, + &cred, + &mut send_seq, + stream_id, + target_host, + target_port, + ) + .await?; + } + Some(ForwardEvent::Data { stream_id, bytes }) => { + send_stream_data(&socket, addr, &cred, &mut send_seq, stream_id, bytes).await?; + } + Some(ForwardEvent::Close { stream_id }) => { + stream_writers.remove(&stream_id); + send_stream_close(&socket, addr, &cred, &mut send_seq, stream_id).await?; + } + None if forward_only => break, + None => {} + } + } _ = status_tick.tick() => { let stale = last_packet_at.elapsed(); if stale >= Duration::from_secs(reconnect_timeout_secs.max(1)) { @@ -1433,8 +1635,10 @@ async fn run_terminal( ) .await? { - render_frame(&frame)?; - predictor.observe_output(&frame.bytes); + if !forward_only { + render_frame(&frame)?; + predictor.observe_output(&frame.bytes); + } last_packet_at = Instant::now(); } } else if stale >= Duration::from_secs(2) { @@ -1457,6 +1661,60 @@ async fn run_terminal( Ok(()) } +async fn start_local_forwards( + local_forwards: &[LocalForward], + forward_tx: mpsc::UnboundedSender, + stream_ids: Arc, +) -> Result<()> { + for forward in local_forwards { + let bind = format!("{}:{}", forward.bind_host, forward.listen_port); + let listener = TcpListener::bind(&bind) + .await + .with_context(|| format!("bind local forward {bind}"))?; + let forward = forward.clone(); + 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 stream_id = stream_ids.fetch_add(1, Ordering::Relaxed); + let (mut reader, writer) = stream.into_split(); + let _ = forward_tx.send(ForwardEvent::Open { + stream_id, + target_host: forward.target_host.clone(), + target_port: forward.target_port, + writer, + }); + let forward_tx = 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, + bytes: buf[..n].to_vec(), + }) + .is_err() + { + return; + } + } + Err(_) => break, + } + } + let _ = forward_tx.send(ForwardEvent::Close { stream_id }); + }); + } + }); + } + Ok(()) +} + async fn reconnect( socket: &UdpSocket, cred: &mut CachedCredential, @@ -1520,6 +1778,68 @@ async fn send_input( Ok(()) } +async fn send_stream_open( + socket: &UdpSocket, + addr: SocketAddr, + cred: &CachedCredential, + send_seq: &mut u64, + stream_id: u64, + target_host: String, + target_port: u16, +) -> Result<()> { + let body = protocol::to_body(&StreamOpen { + stream_id, + target_host, + target_port, + })?; + send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamOpen, &body).await +} + +async fn send_stream_data( + socket: &UdpSocket, + addr: SocketAddr, + cred: &CachedCredential, + send_seq: &mut u64, + stream_id: u64, + bytes: Vec, +) -> Result<()> { + let body = protocol::to_body(&StreamData { stream_id, bytes })?; + send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamData, &body).await +} + +async fn send_stream_close( + socket: &UdpSocket, + addr: SocketAddr, + cred: &CachedCredential, + send_seq: &mut u64, + stream_id: u64, +) -> Result<()> { + let body = protocol::to_body(&StreamClose { stream_id })?; + send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamClose, &body).await +} + +async fn send_stream_packet( + socket: &UdpSocket, + addr: SocketAddr, + cred: &CachedCredential, + send_seq: &mut u64, + kind: PacketKind, + body: &[u8], +) -> Result<()> { + let packet = protocol::encode_encrypted( + kind, + cred.client_id, + *send_seq, + cred.last_rendered_seq, + &cred.session_key, + CLIENT_TO_SERVER, + body, + )?; + *send_seq += 1; + socket.send_to(&packet, addr).await?; + Ok(()) +} + struct Predictor { enabled: bool, alternate_screen: bool, @@ -1763,9 +2083,10 @@ const TERMINAL_CLEANUP: &[u8] = concat!( #[cfg(test)] mod tests { use super::{ - FrameBuffer, Predictor, SshConfig, auth_allows, load_first_native_identity_with_prompt, - parse_ssh_config, raw_contains_host_table, ssh_destination_host, ssh_username, - startup_command, toml_bare_key_or_quoted, + 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, }; use dosh::config::ClientConfig; use dosh::protocol::Frame; @@ -1976,4 +2297,30 @@ mod tests { Some(b"tm work\n".to_vec()) ); } + + #[test] + fn parses_local_forward_with_default_bind() { + assert_eq!( + parse_local_forward("8080:127.0.0.1:80").unwrap(), + LocalForward { + bind_host: "127.0.0.1".to_string(), + listen_port: 8080, + target_host: "127.0.0.1".to_string(), + target_port: 80, + } + ); + } + + #[test] + fn parses_local_forward_with_explicit_bind() { + assert_eq!( + parse_local_forward("0.0.0.0:8080:localhost:8000").unwrap(), + LocalForward { + bind_host: "0.0.0.0".to_string(), + listen_port: 8080, + target_host: "localhost".to_string(), + target_port: 8000, + } + ); + } } diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index 00f438d..824c455 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -7,22 +7,23 @@ use dosh::auth::{ use dosh::config::{ServerConfig, load_server_config}; use dosh::crypto; use dosh::native::{ - NativeAuthOk, NativeServerHello, derive_native_session_key, generate_native_ephemeral, - host_public_key, host_public_key_line, load_or_create_host_key, sign_server_hello, - verify_native_user_auth_from_config, + ForwardingKind, ForwardingRequest, NativeAuthOk, NativeServerHello, derive_native_session_key, + generate_native_ephemeral, host_public_key, host_public_key_line, load_or_create_host_key, + sign_server_hello, verify_native_user_auth_from_config, }; use dosh::protocol::{ self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody, NativeUserAuthBody, PacketKind, - ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT, TicketAttachBody, TicketAttachEnvelope, - TicketAttachOkEnvelope, + ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose, StreamData, StreamOpen, + StreamOpenOk, StreamOpenReject, TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, }; use dosh::pty::{PtyHandle, PtyOutput, spawn_pty_session}; use std::collections::{HashMap, VecDeque}; use std::net::SocketAddr; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; -use tokio::net::UdpSocket; +use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tokio::net::{TcpStream, UdpSocket}; use tokio::sync::mpsc; #[derive(Debug, Parser)] @@ -186,6 +187,8 @@ struct ClientState { last_seen: Instant, pending: VecDeque, last_screen: Option, + allowed_forwardings: Vec, + stream_writers: HashMap>>, } #[derive(Clone)] @@ -263,7 +266,15 @@ async fn handle_packet( handle_ticket_attach(state, socket, peer, packet.body).await } PacketKind::ResumeRequest => handle_resume(state, socket, peer, &packet).await, - PacketKind::Input | PacketKind::Resize | PacketKind::Ping | PacketKind::Ack + PacketKind::Input + | PacketKind::Resize + | PacketKind::Ping + | PacketKind::Ack + | PacketKind::StreamOpen + | PacketKind::StreamData + | PacketKind::StreamClose + | PacketKind::StreamEof + | PacketKind::StreamWindowAdjust if find_client_key(state, &packet.header.conn_id).is_err() => { send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client").await @@ -273,6 +284,9 @@ async fn handle_packet( PacketKind::Ping => handle_ping(state, socket, peer, &packet).await, PacketKind::Ack => handle_ack(state, &packet).await, PacketKind::Detach => handle_detach(state, &packet).await, + PacketKind::StreamOpen => handle_stream_open(state, socket, peer, &packet).await, + PacketKind::StreamData => handle_stream_data(state, peer, &packet).await, + PacketKind::StreamClose => handle_stream_close(state, &packet).await, _ => Ok(()), } } @@ -459,6 +473,8 @@ async fn handle_native_user_auth( last_seen: Instant::now(), pending: VecDeque::new(), last_screen: Some(screen), + allowed_forwardings: req.auth.requested_forwardings.clone(), + stream_writers: HashMap::new(), }, ); Ok(( @@ -562,6 +578,8 @@ async fn handle_bootstrap_attach( last_seen: Instant::now(), pending: VecDeque::new(), last_screen: Some(screen), + allowed_forwardings: Vec::new(), + stream_writers: HashMap::new(), }, ); ( @@ -670,6 +688,8 @@ async fn handle_ticket_attach( last_seen: Instant::now(), pending: VecDeque::new(), last_screen: Some(screen), + allowed_forwardings: Vec::new(), + stream_writers: HashMap::new(), }, ); (client_id, output_seq, snapshot) @@ -913,6 +933,273 @@ async fn handle_ack(state: &Arc>, packet: &protocol::Packet) Ok(()) } +async fn handle_stream_open( + state: &Arc>, + socket: &Arc, + peer: SocketAddr, + packet: &protocol::Packet, +) -> Result<()> { + let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; + let open: StreamOpen = protocol::from_body(&body)?; + let allowed = { + let mut locked = state.lock().expect("server state poisoned"); + let config = locked.config.clone(); + let session = locked + .sessions + .get_mut(&session_name) + .ok_or_else(|| anyhow!("unknown session"))?; + let client = session + .clients + .get_mut(&packet.header.conn_id) + .ok_or_else(|| anyhow!("unknown client"))?; + if !client.replay.accept(packet.header.seq) { + return Ok(()); + } + client.endpoint = peer; + client.last_seen = Instant::now(); + stream_open_allowed(&config, client, &open) + }; + if let Err(err) = allowed { + return send_stream_open_reject( + state, + socket, + packet.header.conn_id, + open.stream_id, + err.to_string(), + ) + .await; + } + + let target = format!("{}:{}", open.target_host, open.target_port); + let stream = match TcpStream::connect((open.target_host.as_str(), open.target_port)).await { + Ok(stream) => stream, + Err(err) => { + return send_stream_open_reject( + state, + socket, + packet.header.conn_id, + open.stream_id, + format!("connect {target}: {err}"), + ) + .await; + } + }; + let (mut reader, mut writer) = stream.into_split(); + let (writer_tx, mut writer_rx) = mpsc::unbounded_channel::>(); + { + let mut locked = state.lock().expect("server state poisoned"); + let session = locked + .sessions + .get_mut(&session_name) + .ok_or_else(|| anyhow!("unknown session"))?; + let client = session + .clients + .get_mut(&packet.header.conn_id) + .ok_or_else(|| anyhow!("unknown client"))?; + client.stream_writers.insert(open.stream_id, writer_tx); + } + + tokio::spawn(async move { + while let Some(bytes) = writer_rx.recv().await { + if writer.write_all(&bytes).await.is_err() { + break; + } + } + }); + + send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?; + + let state = Arc::clone(state); + let socket = Arc::clone(socket); + let client_id = packet.header.conn_id; + let stream_id = open.stream_id; + tokio::spawn(async move { + let mut buf = [0u8; 16 * 1024]; + loop { + match reader.read(&mut buf).await { + Ok(0) => break, + Ok(n) => { + if let Err(err) = send_stream_data_to_client( + &state, + &socket, + client_id, + stream_id, + buf[..n].to_vec(), + ) + .await + { + eprintln!("stream send error: {err:#}"); + break; + } + } + Err(_) => break, + } + } + let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await; + }); + + Ok(()) +} + +async fn handle_stream_data( + state: &Arc>, + peer: SocketAddr, + packet: &protocol::Packet, +) -> Result<()> { + let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; + let data: StreamData = protocol::from_body(&body)?; + let writer = { + let mut locked = state.lock().expect("server state poisoned"); + let session = locked + .sessions + .get_mut(&session_name) + .ok_or_else(|| anyhow!("unknown session"))?; + let client = session + .clients + .get_mut(&packet.header.conn_id) + .ok_or_else(|| anyhow!("unknown client"))?; + if !client.replay.accept(packet.header.seq) { + return Ok(()); + } + client.endpoint = peer; + client.last_seen = Instant::now(); + client.stream_writers.get(&data.stream_id).cloned() + }; + if let Some(writer) = writer { + let _ = writer.send(data.bytes); + } + Ok(()) +} + +async fn handle_stream_close( + state: &Arc>, + packet: &protocol::Packet, +) -> Result<()> { + let (key, session_name) = find_client_key(state, &packet.header.conn_id)?; + let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?; + let close: StreamClose = protocol::from_body(&body)?; + let mut locked = state.lock().expect("server state poisoned"); + let Some(session) = locked.sessions.get_mut(&session_name) else { + return Ok(()); + }; + let Some(client) = session.clients.get_mut(&packet.header.conn_id) else { + return Ok(()); + }; + if !client.replay.accept(packet.header.seq) { + return Ok(()); + } + client.last_seen = Instant::now(); + client.stream_writers.remove(&close.stream_id); + Ok(()) +} + +fn stream_open_allowed( + config: &ServerConfig, + client: &ClientState, + open: &StreamOpen, +) -> Result<()> { + anyhow::ensure!(config.allow_tcp_forwarding, "TCP forwarding disabled"); + let allowed = client.allowed_forwardings.iter().any(|forward| { + forward.kind == ForwardingKind::Local + && forward.target_host.as_deref() == Some(open.target_host.as_str()) + && forward.target_port == Some(open.target_port) + }); + anyhow::ensure!( + allowed, + "stream target {}:{} was not requested during auth", + open.target_host, + open.target_port + ); + Ok(()) +} + +async fn send_stream_open_ok( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, +) -> Result<()> { + let body = protocol::to_body(&StreamOpenOk { stream_id })?; + send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenOk, body).await +} + +async fn send_stream_open_reject( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, + reason: String, +) -> Result<()> { + let body = protocol::to_body(&StreamOpenReject { stream_id, reason })?; + send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenReject, body).await +} + +async fn send_stream_data_to_client( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, + bytes: Vec, +) -> Result<()> { + let body = protocol::to_body(&StreamData { stream_id, bytes })?; + send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamData, body).await +} + +async fn send_stream_close_to_client( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, +) -> Result<()> { + { + let mut locked = state.lock().expect("server state poisoned"); + for session in locked.sessions.values_mut() { + if let Some(client) = session.clients.get_mut(&client_id) { + client.stream_writers.remove(&stream_id); + break; + } + } + } + let body = protocol::to_body(&StreamClose { stream_id })?; + send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamClose, body).await +} + +async fn send_stream_packet_to_client( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + kind: PacketKind, + body: Vec, +) -> Result<()> { + let send = { + let mut locked = state.lock().expect("server state poisoned"); + let mut found = None; + for session in locked.sessions.values_mut() { + if let Some(client) = session.clients.get_mut(&client_id) { + client.send_seq += 1; + let packet = protocol::encode_encrypted( + kind, + client_id, + client.send_seq, + client.last_acked, + &client.session_key, + SERVER_TO_CLIENT, + &body, + )?; + found = Some((client.endpoint, packet)); + break; + } + } + found + }; + if let Some((endpoint, packet)) = send { + socket.send_to(&packet, endpoint).await?; + } + Ok(()) +} + async fn broadcast_output( state: &Arc>, socket: &Arc, diff --git a/src/protocol.rs b/src/protocol.rs index a65c2da..8f67d41 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -30,6 +30,13 @@ pub enum PacketKind { NativeServerHello = 15, NativeUserAuth = 16, NativeAuthOk = 17, + StreamOpen = 18, + StreamOpenOk = 19, + StreamOpenReject = 20, + StreamData = 21, + StreamWindowAdjust = 22, + StreamEof = 23, + StreamClose = 24, } impl TryFrom for PacketKind { @@ -54,6 +61,13 @@ impl TryFrom for PacketKind { 15 => Self::NativeServerHello, 16 => Self::NativeUserAuth, 17 => Self::NativeAuthOk, + 18 => Self::StreamOpen, + 19 => Self::StreamOpenOk, + 20 => Self::StreamOpenReject, + 21 => Self::StreamData, + 22 => Self::StreamWindowAdjust, + 23 => Self::StreamEof, + 24 => Self::StreamClose, _ => bail!("unknown packet kind {value}"), }) } @@ -278,6 +292,46 @@ pub struct Resize { pub rows: u16, } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StreamOpen { + pub stream_id: u64, + pub target_host: String, + pub target_port: u16, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StreamOpenOk { + pub stream_id: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StreamOpenReject { + pub stream_id: u64, + pub reason: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StreamData { + pub stream_id: u64, + pub bytes: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StreamWindowAdjust { + pub stream_id: u64, + pub bytes: u32, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StreamEof { + pub stream_id: u64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StreamClose { + pub stream_id: u64, +} + #[derive(Debug, Clone, Serialize, Deserialize)] pub struct Frame { pub session: String, diff --git a/tests/integration_smoke.rs b/tests/integration_smoke.rs index 8b6c045..ad23b66 100644 --- a/tests/integration_smoke.rs +++ b/tests/integration_smoke.rs @@ -1,6 +1,6 @@ use std::fs; use std::io::{Read, Write}; -use std::net::UdpSocket; +use std::net::{TcpListener, TcpStream, UdpSocket}; use std::os::unix::net::{UnixListener, UnixStream}; use std::path::{Path, PathBuf}; use std::process::{Child, Command, Stdio}; @@ -24,6 +24,11 @@ fn free_udp_port() -> u16 { socket.local_addr().unwrap().port() } +fn free_tcp_port() -> u16 { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap().port() +} + fn write_server_config(dir: &tempfile::TempDir, port: u16) -> std::path::PathBuf { write_server_config_with_timeout(dir, port, 30) } @@ -143,6 +148,50 @@ fn native_attach_once_with_agent( .unwrap() } +fn connect_with_retry(port: u16, timeout: Duration, child: &mut Child) -> TcpStream { + let deadline = std::time::Instant::now() + timeout; + loop { + if let Some(status) = child.try_wait().unwrap() { + let mut stderr = String::new(); + if let Some(mut pipe) = child.stderr.take() { + let _ = pipe.read_to_string(&mut stderr); + } + panic!( + "client exited before local forward was ready: status={status}; stderr={stderr}" + ); + } + match TcpStream::connect(("127.0.0.1", port)) { + Ok(stream) => return stream, + Err(err) if std::time::Instant::now() < deadline => { + let _ = err; + thread::sleep(Duration::from_millis(50)); + } + Err(err) => panic!("connect to local forward 127.0.0.1:{port}: {err}"), + } + } +} + +fn start_echo_server() -> u16 { + let listener = TcpListener::bind("127.0.0.1:0").unwrap(); + let port = listener.local_addr().unwrap().port(); + thread::spawn(move || { + if let Ok((mut stream, _)) = listener.accept() { + let mut buf = [0u8; 1024]; + loop { + match stream.read(&mut buf) { + Ok(0) | Err(_) => break, + Ok(n) => { + if stream.write_all(&buf[..n]).is_err() { + break; + } + } + } + } + } + }); + port +} + fn write_native_client_auth(dir: &tempfile::TempDir, config_path: &std::path::Path) { let ssh_dir = dir.path().join(".ssh"); fs::create_dir_all(&ssh_dir).unwrap(); @@ -495,6 +544,58 @@ fn native_attach_only_smoke_uses_ssh_agent() { assert!(stderr.contains("terminal_ready"), "stderr={stderr}"); } +#[test] +fn native_local_forward_echo_smoke() { + let dir = tempfile::tempdir().unwrap(); + let port = free_udp_port(); + let config = write_server_config(&dir, port); + write_native_client_auth(&dir, &config); + let echo_port = start_echo_server(); + let local_port = free_tcp_port(); + let mut server = start_server(&dir, &config); + let client_bin = env!("CARGO_BIN_EXE_dosh-client"); + let mut client = Command::new(client_bin) + .arg("--auth") + .arg("native") + .arg("--no-cache") + .arg("-N") + .arg("-L") + .arg(format!("{local_port}:127.0.0.1:{echo_port}")) + .arg("--session") + .arg("default") + .arg("--dosh-host") + .arg("127.0.0.1") + .arg("--dosh-port") + .arg(port.to_string()) + .arg("local") + .env("HOME", dir.path()) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()) + .spawn() + .unwrap(); + + let mut stream = connect_with_retry(local_port, Duration::from_secs(5), &mut client); + stream + .set_read_timeout(Some(Duration::from_secs(3))) + .unwrap(); + stream.write_all(b"dosh-forward-ping").unwrap(); + let mut buf = [0u8; 64]; + let n = stream.read(&mut buf).unwrap(); + + let _ = client.kill(); + let client_output = client.wait_with_output().unwrap(); + let _ = server.kill(); + let _ = server.wait(); + + assert_eq!(&buf[..n], b"dosh-forward-ping"); + assert!( + client_output.status.success() || client_output.status.code().is_none(), + "stderr={}", + String::from_utf8_lossy(&client_output.stderr) + ); +} + #[test] fn ticket_attach_after_server_restart_smoke() { let dir = tempfile::tempdir().unwrap();