Add native remote port forwarding
ci / test (push) Has been cancelled
ci / remote-bench (push) Has been cancelled

This commit is contained in:
DuProcess
2026-06-13 15:41:48 -04:00
parent 6aa81d0ce3
commit 6a0f12c59e
3 changed files with 451 additions and 15 deletions
+164 -14
View File
@@ -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,
}
);
}
}
+224 -1
View File
@@ -23,7 +23,7 @@ use std::net::SocketAddr;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpStream, UdpSocket};
use tokio::net::{TcpListener, TcpStream, UdpSocket};
use tokio::sync::mpsc;
#[derive(Debug, Parser)]
@@ -164,6 +164,7 @@ struct ServerState {
pty_tx: mpsc::UnboundedSender<PtyOutput>,
sessions: HashMap<String, Session>,
pending_native: HashMap<[u8; 16], PendingNativeAuth>,
next_server_stream_id: u64,
}
struct Session {
@@ -219,6 +220,7 @@ impl ServerState {
pty_tx,
sessions: HashMap::new(),
pending_native: HashMap::new(),
next_server_stream_id: 1u64 << 63,
}
}
@@ -271,6 +273,8 @@ async fn handle_packet(
| PacketKind::Ping
| PacketKind::Ack
| PacketKind::StreamOpen
| PacketKind::StreamOpenOk
| PacketKind::StreamOpenReject
| PacketKind::StreamData
| PacketKind::StreamClose
| PacketKind::StreamEof
@@ -285,6 +289,8 @@ async fn handle_packet(
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::StreamOpenOk => handle_stream_open_ok(state, peer, &packet).await,
PacketKind::StreamOpenReject => handle_stream_open_reject(state, peer, &packet).await,
PacketKind::StreamData => handle_stream_data(state, peer, &packet).await,
PacketKind::StreamClose => handle_stream_close(state, &packet).await,
_ => Ok(()),
@@ -406,6 +412,7 @@ async fn handle_native_user_auth(
String,
u64,
Vec<u8>,
Vec<ForwardingRequest>,
)> = {
let mut locked = state.lock().expect("server state poisoned");
verify_native_user_auth_from_config(
@@ -486,6 +493,7 @@ async fn handle_native_user_auth(
mode,
output_seq,
snapshot,
req.auth.requested_forwardings.clone(),
))
})
};
@@ -499,6 +507,7 @@ async fn handle_native_user_auth(
mode,
output_seq,
snapshot,
requested_forwardings,
) = match attached {
Ok(attached) => attached,
Err(err) => {
@@ -506,6 +515,18 @@ async fn handle_native_user_auth(
.await;
}
};
if let Err(err) = start_remote_forwards(
state,
socket,
client_id,
requested_forwardings.clone(),
pending.client.requested_session.clone(),
)
.await
{
remove_client(state, client_id);
return send_reject_to_client(socket, peer, packet.header.conn_id, &err.to_string()).await;
}
let ok = NativeAuthOk {
client_id,
@@ -909,6 +930,13 @@ async fn handle_detach(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packe
Ok(())
}
fn remove_client(state: &Arc<Mutex<ServerState>>, client_id: [u8; 16]) {
let mut locked = state.lock().expect("server state poisoned");
for session in locked.sessions.values_mut() {
session.clients.remove(&client_id);
}
}
async fn handle_ack(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet) -> Result<()> {
let (key, _) = find_client_key(state, &packet.header.conn_id)?;
let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
@@ -1073,6 +1101,57 @@ async fn handle_stream_data(
Ok(())
}
async fn handle_stream_open_ok(
state: &Arc<Mutex<ServerState>>,
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 _ok: StreamOpenOk = protocol::from_body(&body)?;
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();
Ok(())
}
async fn handle_stream_open_reject(
state: &Arc<Mutex<ServerState>>,
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 reject: StreamOpenReject = protocol::from_body(&body)?;
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.remove(&reject.stream_id);
Ok(())
}
async fn handle_stream_close(
state: &Arc<Mutex<ServerState>>,
packet: &protocol::Packet,
@@ -1095,6 +1174,134 @@ async fn handle_stream_close(
Ok(())
}
async fn start_remote_forwards(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
requested_forwardings: Vec<ForwardingRequest>,
session_name: String,
) -> Result<()> {
let remote_forwards = requested_forwardings
.into_iter()
.filter(|forward| forward.kind == ForwardingKind::Remote)
.collect::<Vec<_>>();
if remote_forwards.is_empty() {
return Ok(());
}
let config = {
let locked = state.lock().expect("server state poisoned");
locked.config.clone()
};
anyhow::ensure!(
config.allow_remote_forwarding,
"remote forwarding disabled on server"
);
for forward in remote_forwards {
let bind_host = forward
.bind_host
.clone()
.unwrap_or_else(|| "127.0.0.1".to_string());
anyhow::ensure!(
is_loopback_bind(&bind_host),
"remote forwarding may only bind loopback in this server build"
);
let listen_port = forward.listen_port;
let target_host = forward
.target_host
.clone()
.ok_or_else(|| anyhow!("remote forward target host is required"))?;
let target_port = forward
.target_port
.ok_or_else(|| anyhow!("remote forward target port is required"))?;
let bind = format!("{bind_host}:{listen_port}");
let listener = TcpListener::bind(&bind)
.await
.with_context(|| format!("bind remote forward {bind}"))?;
let state = Arc::clone(state);
let socket = Arc::clone(socket);
let session_name = session_name.clone();
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
break;
};
let stream_id = allocate_server_stream_id(&state);
let (mut reader, mut writer) = stream.into_split();
let (writer_tx, mut writer_rx) = mpsc::unbounded_channel::<Vec<u8>>();
{
let mut locked = state.lock().expect("server state poisoned");
if let Some(session) = locked.sessions.get_mut(&session_name)
&& let Some(client) = session.clients.get_mut(&client_id)
{
client.stream_writers.insert(stream_id, writer_tx);
} else {
break;
}
}
tokio::spawn(async move {
while let Some(bytes) = writer_rx.recv().await {
if writer.write_all(&bytes).await.is_err() {
break;
}
}
});
if let Err(err) = send_stream_open_to_client(
&state,
&socket,
client_id,
stream_id,
target_host.clone(),
target_port,
)
.await
{
eprintln!("remote forward open error: {err:#}");
break;
}
let state = Arc::clone(&state);
let socket = Arc::clone(&socket);
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!("remote forward data error: {err:#}");
break;
}
}
Err(_) => break,
}
}
let _ =
send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
});
}
});
}
Ok(())
}
fn allocate_server_stream_id(state: &Arc<Mutex<ServerState>>) -> u64 {
let mut locked = state.lock().expect("server state poisoned");
let stream_id = locked.next_server_stream_id;
locked.next_server_stream_id = locked.next_server_stream_id.wrapping_add(1).max(1u64 << 63);
stream_id
}
fn is_loopback_bind(host: &str) -> bool {
matches!(host, "127.0.0.1" | "localhost" | "::1")
}
fn stream_open_allowed(
config: &ServerConfig,
client: &ClientState,
@@ -1125,6 +1332,22 @@ async fn send_stream_open_ok(
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenOk, body).await
}
async fn send_stream_open_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
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_to_client(state, socket, client_id, PacketKind::StreamOpen, body).await
}
async fn send_stream_open_reject(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
+63
View File
@@ -33,6 +33,17 @@ fn write_server_config(dir: &tempfile::TempDir, port: u16) -> std::path::PathBuf
write_server_config_with_timeout(dir, port, 30)
}
fn write_server_config_with_remote_forwarding(
dir: &tempfile::TempDir,
port: u16,
) -> std::path::PathBuf {
let config = write_server_config(dir, port);
let mut raw = fs::read_to_string(&config).unwrap();
raw.push_str("allow_remote_forwarding = true\n");
fs::write(&config, raw).unwrap();
config
}
fn write_server_config_with_timeout(
dir: &tempfile::TempDir,
port: u16,
@@ -596,6 +607,58 @@ fn native_local_forward_echo_smoke() {
);
}
#[test]
fn native_remote_forward_echo_smoke() {
let dir = tempfile::tempdir().unwrap();
let port = free_udp_port();
let config = write_server_config_with_remote_forwarding(&dir, port);
write_native_client_auth(&dir, &config);
let echo_port = start_echo_server();
let remote_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("-R")
.arg(format!("{remote_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(remote_port, Duration::from_secs(5), &mut client);
stream
.set_read_timeout(Some(Duration::from_secs(3)))
.unwrap();
stream.write_all(b"dosh-remote-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-remote-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();