Palav/dosh
1
0
Fork 0
Code Issues Pull Requests Actions 58 Packages Projects Releases 51 Wiki Activity

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

Browse Source
This commit is contained in:
DuProcess
2026-06-13 15:36:10 -04:00
parent f324a7627f
commit 6aa81d0ce3
4 changed files with 844 additions and 55 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/bin/dosh-client.rs
+394 -47
View File
@@ -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<String>,
#[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<u8>,
},
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<u8> {
command.into_bytes()
}
fn parse_local_forwards(raw: &[String]) -> Result<Vec<LocalForward>> {
raw.iter().map(|value| parse_local_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() {
[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::<u16>()
.with_context(|| format!("invalid local forward listen port in {raw:?}"))?;
let target_port = target_port
.parse::<u16>()
.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<ForwardingRequest> {
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<u16>,
@@ -822,6 +929,7 @@ async fn try_native_auth(
mode: &str,
cols: u16,
rows: u16,
requested_forwardings: Vec<ForwardingRequest>,
) -> 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<ForwardingRequest>,
) -> Result<dosh::native::NativeUserAuth> {
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<Vec<u8>>,
reconnect_timeout_secs: u64,
local_forwards: Vec<LocalForward>,
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::<ForwardEvent>();
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<u64, tokio::net::tcp::OwnedWriteHalf> = 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::<Vec<u8>>();
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::<StreamOpenReject>(&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::<StreamData>(&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::<StreamClose>(&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<ForwardEvent>,
stream_ids: Arc<AtomicU64>,
) -> 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<u8>,
) -> 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,
}
);
}
}
src/bin/dosh-server.rs
+294 -7
View File
@@ -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<PendingFrame>,
last_screen: Option<vt100::Screen>,
allowed_forwardings: Vec<ForwardingRequest>,
stream_writers: HashMap<u64, mpsc::UnboundedSender<Vec<u8>>>,
}
#[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<Mutex<ServerState>>, packet: &protocol::Packet)
Ok(())
}
async fn handle_stream_open(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
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::<Vec<u8>>();
{
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<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 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<Mutex<ServerState>>,
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<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
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<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
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<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
bytes: Vec<u8>,
) -> 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<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
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<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
kind: PacketKind,
body: Vec<u8>,
) -> 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<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
src/protocol.rs
+54
View File
@@ -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<u8> for PacketKind {
@@ -54,6 +61,13 @@ impl TryFrom<u8> 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<u8>,
}
#[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,
tests/integration_smoke.rs
+102 -1
View File
@@ -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();