diff --git a/Cargo.lock b/Cargo.lock index 334aa20..8066eb2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -447,6 +447,7 @@ dependencies = [ "crossterm", "dirs", "ed25519-dalek", + "filetime", "hkdf", "hmac", "libc", @@ -577,6 +578,16 @@ dependencies = [ "winapi", ] +[[package]] +name = "filetime" +version = "0.2.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" +dependencies = [ + "cfg-if", + "libc", +] + [[package]] name = "foldhash" version = "0.1.5" diff --git a/Cargo.toml b/Cargo.toml index 3d6e796..ebf1ead 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -14,6 +14,7 @@ clap = { version = "4.5", features = ["derive"] } crossterm = "0.28" dirs = "5.0" ed25519-dalek = "2.1" +filetime = "0.2" hkdf = "0.12" hmac = "0.12" libc = "0.2" diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index e4cacef..20d35aa 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -1,4 +1,4 @@ -use anyhow::{Context, Result, anyhow}; +use anyhow::{Context, Result, anyhow, bail}; use base64::Engine; use base64::engine::general_purpose::URL_SAFE_NO_PAD; use clap::{ArgAction, Parser}; @@ -8,6 +8,15 @@ use dosh::auth::{ }; use dosh::config::{expand_tilde, load_client_config, load_hosts_config, load_server_config}; use dosh::crypto; +use dosh::exec_service::{ + EXEC_STREAM_SENTINEL, ExecRequest, ExecResponse, FrameDecoder as ExecFrameDecoder, + decode_response as decode_exec_response, encode_request as encode_exec_request, +}; +use dosh::file_transfer::{ + CHUNK_SIZE, CopyEndpoint, FILE_STREAM_SENTINEL, FileKind, FileRequest, FileResponse, + FrameDecoder, decode_response, encode_request, ensure_relative_child, parse_copy_endpoint, + remote_basename, remote_join, +}; use dosh::native::{ EnvVar, ForwardingKind, ForwardingRequest, KnownHostStatus, TrustResult, derive_native_session_key, generate_native_ephemeral, host_fingerprint, load_native_identity, @@ -24,18 +33,19 @@ use dosh::protocol::{ }; use dosh::ssh_agent; use serde::{Deserialize, Serialize}; +use sha2::{Digest, Sha256}; use std::collections::BTreeMap; use std::collections::HashMap; use std::collections::HashSet; use std::collections::VecDeque; use std::fs; -use std::io::{Read, Write}; +use std::io::{Read, Seek, Write}; use std::net::{ Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener as StdTcpListener, TcpStream as StdTcpStream, ToSocketAddrs, }; use std::path::{Path, PathBuf}; -use std::process::{Command, Stdio}; +use std::process::{Child, Command, Stdio}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, Ordering}; use std::time::{Duration, Instant}; @@ -221,6 +231,24 @@ enum StartupGateMode { async fn main() -> Result<()> { let mut args = Args::parse(); let config = load_client_config(None).unwrap_or_default(); + if args.server.as_deref() == Some("cp") { + return run_cp_command(&config, &args); + } + if args.server.as_deref() == Some("exec") { + return run_exec_command(&args); + } + if args.server.as_deref() == Some("ls") { + return run_file_ls_command(&args); + } + if args.server.as_deref() == Some("mkdir") { + return run_file_mkdir_command(&args); + } + if args.server.as_deref() == Some("rm") { + return run_file_rm_command(&args); + } + if args.server.as_deref() == Some("cat") { + return run_file_cat_command(&args); + } if args.server.as_deref() == Some("update") { return run_update(&config, update_check_requested(&args.command)?); } @@ -955,10 +983,12 @@ async fn run_doctor_for_host( println!("[ok] known host: trusted"); println!("[ok] native auth: authorized as {}", check.requested_user); println!( - "[ok] forwarding policy: tcp={} remote={} agent={}", + "[ok] server policy: tcp_forward={} remote_forward={} agent_forward={} file_transfer={} exec_command={}", check.allow_tcp_forwarding, check.allow_remote_forwarding, - check.allow_agent_forwarding + check.allow_agent_forwarding, + check.allow_file_transfer, + check.allow_exec_command ); Ok(()) } @@ -1275,6 +1305,977 @@ fn signal_background_ready() { } } +#[derive(Debug, Clone)] +struct CpOptions { + recursive: bool, + resume: bool, + preserve: bool, + progress: bool, + overwrite: bool, +} + +impl Default for CpOptions { + fn default() -> Self { + Self { + recursive: false, + resume: false, + preserve: true, + progress: true, + overwrite: true, + } + } +} + +struct FileForwarder { + child: Child, +} + +impl Drop for FileForwarder { + fn drop(&mut self) { + let _ = self.child.kill(); + let _ = self.child.wait(); + } +} + +struct FileServiceClient { + stream: StdTcpStream, + decoder: FrameDecoder, + pending: VecDeque, +} + +struct ExecServiceClient { + stream: StdTcpStream, + decoder: ExecFrameDecoder, + pending: VecDeque, +} + +impl ExecServiceClient { + fn connect(port: u16) -> Result { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + match StdTcpStream::connect(("127.0.0.1", port)) { + Ok(stream) => { + stream + .set_nodelay(true) + .context("set exec service TCP_NODELAY")?; + return Ok(Self { + stream, + decoder: ExecFrameDecoder::default(), + pending: VecDeque::new(), + }); + } + Err(err) if Instant::now() < deadline => { + std::thread::sleep(Duration::from_millis(20)); + if err.kind() == std::io::ErrorKind::ConnectionRefused { + continue; + } + } + Err(err) => return Err(err).context("connect Dosh exec service"), + } + } + } + + fn send(&mut self, request: ExecRequest) -> Result<()> { + let bytes = encode_exec_request(&request)?; + self.stream.write_all(&bytes).context("send exec request") + } + + fn recv(&mut self) -> Result { + if let Some(response) = self.pending.pop_front() { + return Ok(response); + } + let mut buf = [0u8; 16 * 1024]; + loop { + let n = self.stream.read(&mut buf).context("read exec response")?; + if n == 0 { + bail!("exec service closed"); + } + let frames = self.decoder.push(&buf[..n])?; + for frame in frames { + self.pending.push_back(decode_exec_response(&frame)?); + } + if let Some(response) = self.pending.pop_front() { + return Ok(response); + } + } + } +} + +impl FileServiceClient { + fn connect(port: u16) -> Result { + let deadline = Instant::now() + Duration::from_secs(10); + loop { + match StdTcpStream::connect(("127.0.0.1", port)) { + Ok(stream) => { + stream + .set_nodelay(true) + .context("set file service TCP_NODELAY")?; + return Ok(Self { + stream, + decoder: FrameDecoder::default(), + pending: VecDeque::new(), + }); + } + Err(err) if Instant::now() < deadline => { + std::thread::sleep(Duration::from_millis(20)); + if err.kind() == std::io::ErrorKind::ConnectionRefused { + continue; + } + } + Err(err) => return Err(err).context("connect Dosh file service"), + } + } + } + + fn send(&mut self, request: FileRequest) -> Result<()> { + let bytes = encode_request(&request)?; + self.stream.write_all(&bytes).context("send file request") + } + + fn recv(&mut self) -> Result { + if let Some(response) = self.pending.pop_front() { + return Ok(response); + } + let mut buf = [0u8; 16 * 1024]; + loop { + let n = self.stream.read(&mut buf).context("read file response")?; + if n == 0 { + bail!("file service closed"); + } + let frames = self.decoder.push(&buf[..n])?; + for frame in frames { + self.pending.push_back(decode_response(&frame)?); + } + if let Some(response) = self.pending.pop_front() { + return Ok(response); + } + } + } + + fn stat_optional(&mut self, path: &str) -> Result> { + self.send(FileRequest::Stat { + path: path.to_string(), + })?; + match self.recv()? { + FileResponse::Stat { meta } => Ok(Some(meta)), + FileResponse::Error { .. } => Ok(None), + other => bail!("unexpected file stat response: {other:?}"), + } + } + + fn stat(&mut self, path: &str) -> Result { + self.send(FileRequest::Stat { + path: path.to_string(), + })?; + match self.recv()? { + FileResponse::Stat { meta } => Ok(meta), + FileResponse::Error { message } => Err(anyhow!(message)), + other => bail!("unexpected file stat response: {other:?}"), + } + } + + fn list(&mut self, path: &str) -> Result> { + self.send(FileRequest::List { + path: path.to_string(), + })?; + match self.recv()? { + FileResponse::List { entries } => Ok(entries), + FileResponse::Error { message } => Err(anyhow!(message)), + other => bail!("unexpected file list response: {other:?}"), + } + } + + fn mkdir(&mut self, path: &str, mode: Option) -> Result<()> { + self.send(FileRequest::Mkdir { + path: path.to_string(), + mode, + })?; + match self.recv()? { + FileResponse::Ok => Ok(()), + FileResponse::Error { message } => Err(anyhow!(message)), + other => bail!("unexpected mkdir response: {other:?}"), + } + } +} + +fn run_cp_command(_config: &dosh::config::ClientConfig, args: &Args) -> Result<()> { + let (opts, src, dst) = parse_cp_args(&args.command)?; + let src = parse_copy_endpoint(&src); + let dst = parse_copy_endpoint(&dst); + match (src, dst) { + (CopyEndpoint::Local(src), CopyEndpoint::Remote { host, path }) => { + let port = free_loopback_port()?; + let _forwarder = spawn_service_forwarder(args, &host, port, FILE_STREAM_SENTINEL)?; + let mut client = FileServiceClient::connect(port)?; + upload_path(&mut client, &src, &path, &opts) + } + (CopyEndpoint::Remote { host, path }, CopyEndpoint::Local(dst)) => { + let port = free_loopback_port()?; + let _forwarder = spawn_service_forwarder(args, &host, port, FILE_STREAM_SENTINEL)?; + let mut client = FileServiceClient::connect(port)?; + download_path(&mut client, &path, &dst, &opts) + } + ( + CopyEndpoint::Remote { + host: src_host, + path: src_path, + }, + CopyEndpoint::Remote { + host: dst_host, + path: dst_path, + }, + ) => { + let src_port = free_loopback_port()?; + let dst_port = free_loopback_port()?; + let _src_forwarder = + spawn_service_forwarder(args, &src_host, src_port, FILE_STREAM_SENTINEL)?; + let _dst_forwarder = + spawn_service_forwarder(args, &dst_host, dst_port, FILE_STREAM_SENTINEL)?; + let mut src_client = FileServiceClient::connect(src_port)?; + let mut dst_client = FileServiceClient::connect(dst_port)?; + copy_remote_path( + &mut src_client, + &mut dst_client, + &src_path, + &dst_path, + &opts, + ) + } + (CopyEndpoint::Local(_), CopyEndpoint::Local(_)) => Err(anyhow!( + "dosh cp needs one remote path, e.g. `dosh cp file host:path`" + )), + } +} + +fn run_exec_command(args: &Args) -> Result<()> { + let host = args + .command + .first() + .ok_or_else(|| anyhow!("usage: dosh exec HOST COMMAND"))?; + let command = args.command.get(1..).unwrap_or_default(); + anyhow::ensure!(!command.is_empty(), "usage: dosh exec HOST COMMAND"); + let command = command.join(" "); + let port = free_loopback_port()?; + let _forwarder = spawn_service_forwarder(args, host, port, EXEC_STREAM_SENTINEL)?; + let mut client = ExecServiceClient::connect(port)?; + client.send(ExecRequest { command })?; + loop { + match client.recv()? { + ExecResponse::Stdout(bytes) => { + let mut stdout = std::io::stdout(); + stdout.write_all(&bytes)?; + stdout.flush()?; + } + ExecResponse::Stderr(bytes) => { + let mut stderr = std::io::stderr(); + stderr.write_all(&bytes)?; + stderr.flush()?; + } + ExecResponse::Exit { code } => { + if code.unwrap_or(1) == 0 { + return Ok(()); + } + return Err(anyhow!("remote command exited with status {:?}", code)); + } + ExecResponse::Error { message } => return Err(anyhow!(message)), + } + } +} + +fn run_file_ls_command(args: &Args) -> Result<()> { + if args.command.len() != 1 { + return Err(anyhow!("usage: dosh ls host:path")); + } + let (host, path) = parse_single_remote_path(&args.command[0], "ls")?; + let (_forwarder, mut client) = connect_file_service(args, &host)?; + let meta = client.stat(&path)?; + if meta.kind == FileKind::Directory { + for entry in client.list(&path)? { + println!( + "{}\t{}\t{}", + file_kind_label(entry.meta.kind), + entry.meta.len, + entry.name + ); + } + } else { + println!("{}\t{}\t{}", file_kind_label(meta.kind), meta.len, path); + } + Ok(()) +} + +fn run_file_mkdir_command(args: &Args) -> Result<()> { + if args.command.len() != 1 { + return Err(anyhow!("usage: dosh mkdir host:path")); + } + let (host, path) = parse_single_remote_path(&args.command[0], "mkdir")?; + let (_forwarder, mut client) = connect_file_service(args, &host)?; + client.mkdir(&path, None) +} + +fn run_file_rm_command(args: &Args) -> Result<()> { + let mut recursive = false; + let mut paths = Vec::new(); + for arg in &args.command { + match arg.as_str() { + "-r" | "-R" | "--recursive" => recursive = true, + other if other.starts_with('-') => { + return Err(anyhow!("dosh rm: unknown option {other:?}")); + } + _ => paths.push(arg.clone()), + } + } + if paths.len() != 1 { + return Err(anyhow!("usage: dosh rm [-r] host:path")); + } + let (host, path) = parse_single_remote_path(&paths[0], "rm")?; + let (_forwarder, mut client) = connect_file_service(args, &host)?; + client.send(FileRequest::Remove { path, recursive })?; + match client.recv()? { + FileResponse::Ok => Ok(()), + FileResponse::Error { message } => Err(anyhow!(message)), + other => bail!("unexpected rm response: {other:?}"), + } +} + +fn run_file_cat_command(args: &Args) -> Result<()> { + if args.command.len() != 1 { + return Err(anyhow!("usage: dosh cat host:path")); + } + let (host, path) = parse_single_remote_path(&args.command[0], "cat")?; + let (_forwarder, mut client) = connect_file_service(args, &host)?; + let meta = client.stat(&path)?; + anyhow::ensure!(meta.kind == FileKind::File, "{path} is not a file"); + client.send(FileRequest::Get { path, offset: 0 })?; + match client.recv()? { + FileResponse::Start { .. } => {} + FileResponse::Error { message } => return Err(anyhow!(message)), + other => bail!("unexpected cat start response: {other:?}"), + } + let mut stdout = std::io::stdout(); + let mut hasher = Sha256::new(); + loop { + match client.recv()? { + FileResponse::Chunk { bytes, .. } => { + hasher.update(&bytes); + stdout.write_all(&bytes)?; + } + FileResponse::Done { sha256, bytes } => { + stdout.flush()?; + let digest: [u8; 32] = hasher.finalize().into(); + anyhow::ensure!(sha256 == digest, "cat checksum mismatch for remote file"); + anyhow::ensure!(bytes == meta.len, "cat size mismatch for remote file"); + return Ok(()); + } + FileResponse::Error { message } => return Err(anyhow!(message)), + other => bail!("unexpected cat response: {other:?}"), + } + } +} + +fn parse_single_remote_path(raw: &str, command: &str) -> Result<(String, String)> { + match parse_copy_endpoint(raw) { + CopyEndpoint::Remote { host, path } => Ok((host, path)), + CopyEndpoint::Local(_) => Err(anyhow!("dosh {command} expects a remote host:path")), + } +} + +fn connect_file_service(args: &Args, host: &str) -> Result<(FileForwarder, FileServiceClient)> { + let port = free_loopback_port()?; + let forwarder = spawn_service_forwarder(args, host, port, FILE_STREAM_SENTINEL)?; + let client = FileServiceClient::connect(port)?; + Ok((forwarder, client)) +} + +fn file_kind_label(kind: FileKind) -> &'static str { + match kind { + FileKind::File => "file", + FileKind::Directory => "dir", + FileKind::Symlink => "symlink", + FileKind::Other => "other", + } +} + +fn parse_cp_args(raw: &[String]) -> Result<(CpOptions, String, String)> { + let mut opts = CpOptions::default(); + let mut paths = Vec::new(); + let mut stop_options = false; + for arg in raw { + if stop_options { + paths.push(arg.clone()); + continue; + } + match arg.as_str() { + "--" => stop_options = true, + "-r" | "-R" | "--recursive" => opts.recursive = true, + "--resume" => opts.resume = true, + "-p" | "--preserve" => opts.preserve = true, + "--no-preserve" => opts.preserve = false, + "--no-progress" | "-q" | "--quiet" => opts.progress = false, + "--no-clobber" | "-n" => opts.overwrite = false, + other if other.starts_with('-') => { + return Err(anyhow!("dosh cp: unknown option {other:?}")); + } + _ => paths.push(arg.clone()), + } + } + anyhow::ensure!(paths.len() == 2, "usage: dosh cp [-r] [--resume] SRC DST"); + Ok((opts, paths.remove(0), paths.remove(0))) +} + +fn spawn_service_forwarder( + args: &Args, + host: &str, + port: u16, + target_host: &str, +) -> Result { + let listener = + StdTcpListener::bind("127.0.0.1:0").context("bind file forwarder readiness socket")?; + listener + .set_nonblocking(true) + .context("configure file forwarder readiness socket")?; + let readiness_addr = listener.local_addr()?; + let token = URL_SAFE_NO_PAD.encode(crypto::random_16()); + let exe = std::env::current_exe().context("locate current dosh-client executable")?; + let mut command = Command::new(exe); + command + .arg("--local-forward") + .arg(format!("127.0.0.1:{port}:{target_host}:0")) + .arg("--forward-only") + .arg("--auth") + .arg("native") + .env("DOSH_BACKGROUND_CHILD", "1") + .env("DOSH_BACKGROUND_READY_ADDR", readiness_addr.to_string()) + .env("DOSH_BACKGROUND_READY_TOKEN", &token) + .stdin(Stdio::null()) + .stdout(Stdio::null()) + .stderr(Stdio::piped()); + if args.local_auth { + command.arg("--local-auth"); + } + if let Some(value) = args.ssh_port { + command.arg("--ssh-port").arg(value.to_string()); + } + if let Some(value) = &args.ssh_auth_command { + command.arg("--ssh-auth-command").arg(value); + } + if let Some(value) = &args.ssh_key { + command.arg("--ssh-key").arg(value); + } + if let Some(value) = &args.ssh_known_hosts { + command.arg("--ssh-known-hosts").arg(value); + } + if let Some(value) = &args.ssh_control_path { + command.arg("--ssh-control-path").arg(value); + } + if let Some(value) = args.dosh_port { + command.arg("--dosh-port").arg(value.to_string()); + } + if let Some(value) = &args.dosh_host { + command.arg("--dosh-host").arg(value); + } + for _ in 0..args.verbose { + command.arg("-v"); + } + command.arg(host); + let mut child = command.spawn().context("spawn Dosh file forwarder")?; + let deadline = Instant::now() + Duration::from_secs(30); + loop { + match listener.accept() { + Ok((mut stream, _)) => { + let mut raw = String::new(); + stream.read_to_string(&mut raw).ok(); + if raw.trim() == token { + return Ok(FileForwarder { child }); + } + let _ = child.kill(); + let _ = child.wait(); + return Err(anyhow!("file forwarder sent invalid readiness token")); + } + Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {} + Err(err) => return Err(err).context("wait for file forwarder readiness"), + } + if let Some(status) = child.try_wait()? { + let mut stderr = String::new(); + if let Some(mut pipe) = child.stderr.take() { + let _ = pipe.read_to_string(&mut stderr); + } + return Err(anyhow!( + "file forwarder exited before readiness: {status}; stderr={stderr}" + )); + } + if Instant::now() >= deadline { + let _ = child.kill(); + let _ = child.wait(); + return Err(anyhow!("file forwarder did not become ready")); + } + std::thread::sleep(Duration::from_millis(20)); + } +} + +fn free_loopback_port() -> Result { + let listener = StdTcpListener::bind("127.0.0.1:0").context("reserve local file port")?; + Ok(listener.local_addr()?.port()) +} + +fn upload_path( + client: &mut FileServiceClient, + local: &Path, + remote: &str, + opts: &CpOptions, +) -> Result<()> { + let metadata = + fs::symlink_metadata(local).with_context(|| format!("stat {}", local.display()))?; + let remote = upload_destination(client, local, remote, metadata.is_dir())?; + if metadata.is_dir() { + anyhow::ensure!( + opts.recursive, + "{} is a directory; pass -r to copy directories", + local.display() + ); + client.mkdir(&remote, mode_for(local, opts.preserve)?)?; + let mut entries = fs::read_dir(local) + .with_context(|| format!("list {}", local.display()))? + .collect::>>()?; + entries.sort_by_key(|entry| entry.file_name()); + for entry in entries { + let name = ensure_relative_child(&entry.path())?; + upload_path(client, &entry.path(), &remote_join(&remote, &name), opts)?; + } + return Ok(()); + } + anyhow::ensure!( + metadata.is_file(), + "{} is not a regular file", + local.display() + ); + upload_file( + client, + local, + &remote, + metadata.len(), + mode_for(local, opts.preserve)?, + opts, + ) +} + +fn upload_destination( + client: &mut FileServiceClient, + local: &Path, + remote: &str, + local_is_dir: bool, +) -> Result { + let base = ensure_relative_child(local)?; + if remote.ends_with('/') { + return Ok(remote_join(remote, &base)); + } + if let Some(meta) = client.stat_optional(remote)? + && meta.kind == FileKind::Directory + { + return Ok(remote_join(remote, &base)); + } + if local_is_dir && remote == "." { + return Ok(base); + } + Ok(remote.to_string()) +} + +fn upload_file( + client: &mut FileServiceClient, + local: &Path, + remote: &str, + size: u64, + mode: Option, + opts: &CpOptions, +) -> Result<()> { + client.send(FileRequest::PutStart { + path: remote.to_string(), + size, + mode, + modified_secs: modified_secs(local, opts.preserve)?, + overwrite: opts.overwrite, + resume: opts.resume, + })?; + let offset = match client.recv()? { + FileResponse::Resume { offset } => offset.min(size), + FileResponse::Error { message } => return Err(anyhow!(message)), + other => bail!("unexpected upload start response: {other:?}"), + }; + let mut file = fs::File::open(local).with_context(|| format!("open {}", local.display()))?; + let mut hasher = Sha256::new(); + if offset > 0 { + hash_prefix(&mut file, offset, &mut hasher)?; + } + file.seek(std::io::SeekFrom::Start(offset))?; + let mut sent = offset; + let mut buf = vec![0u8; CHUNK_SIZE]; + loop { + let n = file.read(&mut buf)?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + client.send(FileRequest::PutChunk { + offset: sent, + bytes: buf[..n].to_vec(), + })?; + sent = sent.saturating_add(n as u64); + } + let digest: [u8; 32] = hasher.finalize().into(); + client.send(FileRequest::PutFinish { sha256: digest })?; + match client.recv()? { + FileResponse::Done { sha256, bytes } if sha256 == digest && bytes == size => { + if opts.progress { + eprintln!( + "uploaded {} -> {} ({} bytes)", + local.display(), + remote, + bytes + ); + } + Ok(()) + } + FileResponse::Done { .. } => Err(anyhow!("upload verification failed for {remote}")), + FileResponse::Error { message } => Err(anyhow!(message)), + other => bail!("unexpected upload finish response: {other:?}"), + } +} + +fn download_path( + client: &mut FileServiceClient, + remote: &str, + local: &Path, + opts: &CpOptions, +) -> Result<()> { + let meta = client.stat(remote)?; + let destination = download_destination(local, remote, meta.kind == FileKind::Directory); + match meta.kind { + FileKind::Directory => { + anyhow::ensure!( + opts.recursive, + "{remote} is a directory; pass -r to copy directories" + ); + fs::create_dir_all(&destination) + .with_context(|| format!("create {}", destination.display()))?; + set_local_mode(&destination, meta.mode, opts.preserve)?; + for entry in client.list(remote)? { + download_path( + client, + &remote_join(remote, &entry.name), + &destination.join(&entry.name), + opts, + )?; + } + Ok(()) + } + FileKind::File => download_file(client, remote, &destination, &meta, opts), + _ => Err(anyhow!( + "remote path is not a regular file or directory: {remote}" + )), + } +} + +fn copy_remote_path( + src_client: &mut FileServiceClient, + dst_client: &mut FileServiceClient, + src: &str, + dst: &str, + opts: &CpOptions, +) -> Result<()> { + let meta = src_client.stat(src)?; + let dst = remote_copy_destination(dst_client, src, dst, meta.kind == FileKind::Directory)?; + match meta.kind { + FileKind::Directory => { + anyhow::ensure!( + opts.recursive, + "{src} is a directory; pass -r to copy directories" + ); + dst_client.mkdir(&dst, preserve_remote_mode(meta.mode, opts.preserve))?; + for entry in src_client.list(src)? { + copy_remote_path( + src_client, + dst_client, + &remote_join(src, &entry.name), + &remote_join(&dst, &entry.name), + opts, + )?; + } + Ok(()) + } + FileKind::File => copy_remote_file(src_client, dst_client, src, &dst, &meta, opts), + _ => Err(anyhow!( + "remote path is not a regular file or directory: {src}" + )), + } +} + +fn remote_copy_destination( + dst_client: &mut FileServiceClient, + src: &str, + dst: &str, + src_is_dir: bool, +) -> Result { + if dst.ends_with('/') { + return Ok(remote_join(dst, &remote_basename(src))); + } + if let Some(meta) = dst_client.stat_optional(dst)? + && meta.kind == FileKind::Directory + { + return Ok(remote_join(dst, &remote_basename(src))); + } + if src_is_dir && dst == "." { + return Ok(remote_basename(src)); + } + Ok(dst.to_string()) +} + +fn copy_remote_file( + src_client: &mut FileServiceClient, + dst_client: &mut FileServiceClient, + src: &str, + dst: &str, + meta: &dosh::file_transfer::FileMeta, + opts: &CpOptions, +) -> Result<()> { + dst_client.send(FileRequest::PutStart { + path: dst.to_string(), + size: meta.len, + mode: preserve_remote_mode(meta.mode, opts.preserve), + modified_secs: preserve_remote_mtime(meta.modified_secs, opts.preserve), + overwrite: opts.overwrite, + resume: opts.resume, + })?; + let offset = match dst_client.recv()? { + FileResponse::Resume { offset } => offset.min(meta.len), + FileResponse::Error { message } => return Err(anyhow!(message)), + other => bail!("unexpected remote copy start response: {other:?}"), + }; + src_client.send(FileRequest::Get { + path: src.to_string(), + offset, + })?; + match src_client.recv()? { + FileResponse::Start { .. } => {} + FileResponse::Error { message } => return Err(anyhow!(message)), + other => bail!("unexpected remote copy source start response: {other:?}"), + } + let mut copied = offset; + let source_hash = loop { + match src_client.recv()? { + FileResponse::Chunk { offset, bytes } => { + anyhow::ensure!( + offset == copied, + "remote copy offset mismatch: got {offset}, expected {copied}" + ); + let len = bytes.len(); + dst_client.send(FileRequest::PutChunk { offset, bytes })?; + copied = copied.saturating_add(len as u64); + } + FileResponse::Done { sha256, bytes } => { + anyhow::ensure!( + bytes == meta.len, + "remote copy source size mismatch for {src}" + ); + break sha256; + } + FileResponse::Error { message } => return Err(anyhow!(message)), + other => bail!("unexpected remote copy source response: {other:?}"), + } + }; + dst_client.send(FileRequest::PutFinish { + sha256: source_hash, + })?; + match dst_client.recv()? { + FileResponse::Done { sha256, bytes } if sha256 == source_hash && bytes == meta.len => { + if opts.progress { + eprintln!("copied {src} -> {dst} ({bytes} bytes)"); + } + Ok(()) + } + FileResponse::Done { .. } => Err(anyhow!("remote copy verification failed for {dst}")), + FileResponse::Error { message } => Err(anyhow!(message)), + other => bail!("unexpected remote copy finish response: {other:?}"), + } +} + +fn preserve_remote_mode(mode: Option, preserve: bool) -> Option { + preserve.then_some(mode).flatten() +} + +fn preserve_remote_mtime(modified_secs: Option, preserve: bool) -> Option { + preserve.then_some(modified_secs).flatten() +} + +fn download_destination(local: &Path, remote: &str, remote_is_dir: bool) -> PathBuf { + if local.is_dir() || path_ends_separator(local) { + return local.join(remote_basename(remote)); + } + if remote_is_dir && !local.exists() { + return local.to_path_buf(); + } + local.to_path_buf() +} + +fn download_file( + client: &mut FileServiceClient, + remote: &str, + local: &Path, + meta: &dosh::file_transfer::FileMeta, + opts: &CpOptions, +) -> Result<()> { + if let Some(parent) = local.parent() + && !parent.as_os_str().is_empty() + { + fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?; + } + let offset = if opts.resume && local.exists() { + fs::metadata(local)?.len().min(meta.len) + } else { + 0 + }; + let mut file = if offset > 0 { + fs::OpenOptions::new() + .read(true) + .append(true) + .open(local) + .with_context(|| format!("open {}", local.display()))? + } else { + fs::OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .open(local) + .with_context(|| format!("create {}", local.display()))? + }; + let mut hasher = Sha256::new(); + if offset > 0 { + let mut prefix = fs::File::open(local)?; + hash_prefix(&mut prefix, offset, &mut hasher)?; + } + client.send(FileRequest::Get { + path: remote.to_string(), + offset, + })?; + match client.recv()? { + FileResponse::Start { .. } => {} + FileResponse::Error { message } => return Err(anyhow!(message)), + other => bail!("unexpected download start response: {other:?}"), + } + let mut received = offset; + loop { + match client.recv()? { + FileResponse::Chunk { offset, bytes } => { + anyhow::ensure!( + offset == received, + "download offset mismatch: got {offset}, expected {received}" + ); + file.write_all(&bytes)?; + hasher.update(&bytes); + received = received.saturating_add(bytes.len() as u64); + } + FileResponse::Done { sha256, bytes } => { + file.flush()?; + let digest: [u8; 32] = hasher.finalize().into(); + anyhow::ensure!(sha256 == digest, "download checksum mismatch for {remote}"); + anyhow::ensure!(bytes == meta.len, "download size mismatch for {remote}"); + set_local_mode(local, meta.mode, opts.preserve)?; + set_local_mtime(local, meta.modified_secs, opts.preserve)?; + if opts.progress { + eprintln!( + "downloaded {} -> {} ({} bytes)", + remote, + local.display(), + bytes + ); + } + return Ok(()); + } + FileResponse::Error { message } => return Err(anyhow!(message)), + other => bail!("unexpected download response: {other:?}"), + } + } +} + +fn hash_prefix(file: &mut fs::File, bytes: u64, hasher: &mut Sha256) -> Result<()> { + file.seek(std::io::SeekFrom::Start(0))?; + let mut remaining = bytes; + let mut buf = vec![0u8; CHUNK_SIZE]; + while remaining > 0 { + let want = remaining.min(buf.len() as u64) as usize; + let n = file.read(&mut buf[..want])?; + if n == 0 { + bail!("file ended while hashing resume prefix"); + } + hasher.update(&buf[..n]); + remaining -= n as u64; + } + Ok(()) +} + +fn path_ends_separator(path: &Path) -> bool { + let raw = path.as_os_str().to_string_lossy(); + raw.ends_with('/') || raw.ends_with('\\') +} + +fn modified_secs(path: &Path, preserve: bool) -> Result> { + if !preserve { + return Ok(None); + } + Ok(fs::metadata(path)? + .modified() + .ok() + .and_then(|mtime| mtime.duration_since(std::time::UNIX_EPOCH).ok()) + .map(|duration| duration.as_secs())) +} + +fn set_local_mtime(path: &Path, modified_secs: Option, preserve: bool) -> Result<()> { + if !preserve { + return Ok(()); + } + let Some(modified_secs) = modified_secs else { + return Ok(()); + }; + filetime::set_file_mtime( + path, + filetime::FileTime::from_unix_time(modified_secs as i64, 0), + ) + .with_context(|| format!("set mtime {}", path.display())) +} + +#[cfg(unix)] +fn mode_for(path: &Path, preserve: bool) -> Result> { + use std::os::unix::fs::PermissionsExt; + if preserve { + Ok(Some(fs::metadata(path)?.permissions().mode() & 0o7777)) + } else { + Ok(None) + } +} + +#[cfg(not(unix))] +fn mode_for(_path: &Path, _preserve: bool) -> Result> { + Ok(None) +} + +#[cfg(unix)] +fn set_local_mode(path: &Path, mode: Option, preserve: bool) -> Result<()> { + use std::os::unix::fs::PermissionsExt; + if !preserve { + return Ok(()); + } + let Some(mode) = mode else { + return Ok(()); + }; + let mut permissions = fs::metadata(path)?.permissions(); + permissions.set_mode(mode & 0o7777); + fs::set_permissions(path, permissions).with_context(|| format!("chmod {}", path.display())) +} + +#[cfg(not(unix))] +fn set_local_mode(_path: &Path, _mode: Option, _preserve: bool) -> Result<()> { + Ok(()) +} + fn parse_local_forwards(raw: &[String]) -> Result> { raw.iter().map(|value| parse_local_forward(value)).collect() } @@ -1326,7 +2327,7 @@ fn parse_local_forward(raw: &str) -> Result { "invalid -L {raw:?}; listen port cannot be 0" ); anyhow::ensure!( - target_port != 0, + target_port != 0 || matches!(target_host, FILE_STREAM_SENTINEL | EXEC_STREAM_SENTINEL), "invalid -L {raw:?}; target port cannot be 0" ); Ok(LocalForward { diff --git a/src/bin/dosh-server.rs b/src/bin/dosh-server.rs index d9ae160..7fa73f9 100644 --- a/src/bin/dosh-server.rs +++ b/src/bin/dosh-server.rs @@ -1,4 +1,4 @@ -use anyhow::{Context, Result, anyhow}; +use anyhow::{Context, Result, anyhow, bail}; use clap::{Parser, Subcommand}; use dosh::auth::{ build_attach_ticket, build_bootstrap, encode_bootstrap, load_or_create_server_secret, now_secs, @@ -6,6 +6,14 @@ use dosh::auth::{ }; use dosh::config::{ServerConfig, expand_tilde, load_server_config}; use dosh::crypto; +use dosh::exec_service::{ + EXEC_STREAM_SENTINEL, ExecResponse, FrameDecoder as ExecFrameDecoder, + encode_response as encode_exec_response, +}; +use dosh::file_transfer::{ + CHUNK_SIZE, FILE_STREAM_SENTINEL, FileEntry, FileKind, FileMeta, FileRequest, FileResponse, + FrameDecoder, clean_remote_path, encode_response, +}; use dosh::native::{ EnvVar, ForwardingKind, ForwardingRequest, NativeAuthOk, NativeServerHello, derive_native_session_key, generate_native_ephemeral, host_public_key, host_public_key_line, @@ -20,16 +28,21 @@ use dosh::protocol::{ TicketAttachBody, TicketAttachEnvelope, TicketAttachOkEnvelope, }; use dosh::pty::{PtyHandle, PtyOutput, adopt_pty_from_fd, spawn_pty_session}; +use sha2::{Digest, Sha256}; use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; +use std::fs; +use std::io::{Read, Seek, SeekFrom, Write}; use std::net::{IpAddr, SocketAddr}; use std::os::unix::fs::PermissionsExt; use std::os::unix::net::UnixStream as StdUnixStream; -use std::path::PathBuf; +use std::path::{Path, PathBuf}; +use std::process::Stdio; use std::sync::atomic::{AtomicU64, Ordering}; use std::sync::{Arc, Mutex}; use std::time::{Duration, Instant}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream, UdpSocket, UnixListener}; +use tokio::process::Command as TokioCommand; use tokio::sync::mpsc; const STREAM_INITIAL_WINDOW: usize = 1024 * 1024; @@ -1017,6 +1030,8 @@ async fn handle_native_user_auth( allow_tcp_forwarding: locked.config.allow_tcp_forwarding, allow_remote_forwarding: locked.config.allow_remote_forwarding, allow_agent_forwarding: locked.config.allow_agent_forwarding, + allow_file_transfer: locked.config.allow_file_transfer, + allow_exec_command: locked.config.allow_exec_command, policy_flags: Vec::new(), server_version: env!("CARGO_PKG_VERSION").to_string(), }) @@ -1801,6 +1816,88 @@ async fn handle_stream_open( } } + if open.target_host == FILE_STREAM_SENTINEL { + let (writer_tx, writer_rx) = mpsc::channel::>(1024); + register_opened_stream( + state, + &session_name, + packet.header.conn_id, + open.stream_id, + writer_tx, + )?; + 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 { + if let Err(err) = run_file_stream_service( + state.clone(), + socket.clone(), + client_id, + stream_id, + writer_rx, + ) + .await + { + let _ = send_file_response_to_client( + &state, + &socket, + client_id, + stream_id, + FileResponse::Error { + message: err.to_string(), + }, + ) + .await; + } + let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await; + }); + return Ok(()); + } + + if open.target_host == EXEC_STREAM_SENTINEL { + let (writer_tx, writer_rx) = mpsc::channel::>(1024); + register_opened_stream( + state, + &session_name, + packet.header.conn_id, + open.stream_id, + writer_tx, + )?; + 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 { + if let Err(err) = run_exec_stream_service( + state.clone(), + socket.clone(), + client_id, + stream_id, + writer_rx, + ) + .await + { + let _ = send_exec_response_to_client( + &state, + &socket, + client_id, + stream_id, + ExecResponse::Error { + message: err.to_string(), + }, + ) + .await; + } + let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await; + }); + return Ok(()); + } + 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, @@ -1817,30 +1914,13 @@ async fn handle_stream_open( }; let (mut reader, mut writer) = stream.into_split(); let (writer_tx, mut writer_rx) = mpsc::channel::>(1024); - { - 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); - client.opened_streams.insert(open.stream_id); - client - .stream_send_credit - .insert(open.stream_id, STREAM_INITIAL_WINDOW); - client - .stream_next_send_offset - .entry(open.stream_id) - .or_insert(0); - client - .stream_next_recv_offset - .entry(open.stream_id) - .or_insert(0); - } + register_opened_stream( + state, + &session_name, + packet.header.conn_id, + open.stream_id, + writer_tx, + )?; tokio::spawn(async move { while let Some(bytes) = writer_rx.recv().await { @@ -1884,6 +1964,32 @@ async fn handle_stream_open( Ok(()) } +fn register_opened_stream( + state: &Arc>, + session_name: &str, + client_id: [u8; 16], + stream_id: u64, + writer_tx: mpsc::Sender>, +) -> Result<()> { + 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(&client_id) + .ok_or_else(|| anyhow!("unknown client"))?; + client.stream_writers.insert(stream_id, writer_tx); + client.opened_streams.insert(stream_id); + client + .stream_send_credit + .insert(stream_id, STREAM_INITIAL_WINDOW); + client.stream_next_send_offset.entry(stream_id).or_insert(0); + client.stream_next_recv_offset.entry(stream_id).or_insert(0); + Ok(()) +} + async fn handle_stream_data( state: &Arc>, socket: &Arc, @@ -2320,6 +2426,26 @@ fn stream_open_allowed( client: &ClientState, open: &StreamOpen, ) -> Result<()> { + if open.target_host == FILE_STREAM_SENTINEL { + anyhow::ensure!(config.allow_file_transfer, "file transfer disabled"); + let allowed = client.allowed_forwardings.iter().any(|forward| { + forward.kind == ForwardingKind::Local + && forward.target_host.as_deref() == Some(FILE_STREAM_SENTINEL) + && forward.target_port == Some(0) + }); + anyhow::ensure!(allowed, "file transfer was not requested during auth"); + return Ok(()); + } + if open.target_host == EXEC_STREAM_SENTINEL { + anyhow::ensure!(config.allow_exec_command, "exec command disabled"); + let allowed = client.allowed_forwardings.iter().any(|forward| { + forward.kind == ForwardingKind::Local + && forward.target_host.as_deref() == Some(EXEC_STREAM_SENTINEL) + && forward.target_port == Some(0) + }); + anyhow::ensure!(allowed, "exec command was not requested during auth"); + return Ok(()); + } anyhow::ensure!(config.allow_tcp_forwarding, "TCP forwarding disabled"); let allowed = client .allowed_forwardings @@ -2343,6 +2469,512 @@ fn stream_open_allowed( Ok(()) } +struct UploadState { + final_path: PathBuf, + temp_path: Option, + file: fs::File, + hasher: Sha256, + written: u64, + expected_size: u64, + modified_secs: Option, + atomic: bool, +} + +async fn run_file_stream_service( + state: Arc>, + socket: Arc, + client_id: [u8; 16], + stream_id: u64, + mut rx: mpsc::Receiver>, +) -> Result<()> { + let home = std::env::var_os("HOME") + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(".")); + let mut decoder = FrameDecoder::default(); + let mut upload: Option = None; + while let Some(bytes) = rx.recv().await { + for frame in decoder.push(&bytes)? { + let request = match dosh::file_transfer::decode_request(&frame) { + Ok(request) => request, + Err(err) => { + send_file_response_to_client( + &state, + &socket, + client_id, + stream_id, + FileResponse::Error { + message: err.to_string(), + }, + ) + .await?; + continue; + } + }; + if let Err(err) = handle_file_request( + &state, + &socket, + client_id, + stream_id, + &home, + &mut upload, + request, + ) + .await + { + send_file_response_to_client( + &state, + &socket, + client_id, + stream_id, + FileResponse::Error { + message: err.to_string(), + }, + ) + .await?; + } + } + } + Ok(()) +} + +async fn handle_file_request( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, + home: &Path, + upload: &mut Option, + request: FileRequest, +) -> Result<()> { + match request { + FileRequest::Stat { path } => { + let path = clean_remote_path(&path, home)?; + let meta = file_meta(&path, home)?; + send_file_response_to_client( + state, + socket, + client_id, + stream_id, + FileResponse::Stat { meta }, + ) + .await + } + FileRequest::List { path } => { + let path = clean_remote_path(&path, home)?; + let mut entries = Vec::new(); + for entry in fs::read_dir(&path).with_context(|| format!("list {}", path.display()))? { + let entry = entry?; + let name = entry.file_name().to_string_lossy().to_string(); + entries.push(FileEntry { + name, + meta: file_meta(&entry.path(), home)?, + }); + } + entries.sort_by(|a, b| a.name.cmp(&b.name)); + send_file_response_to_client( + state, + socket, + client_id, + stream_id, + FileResponse::List { entries }, + ) + .await + } + FileRequest::Mkdir { path, mode } => { + let path = clean_remote_path(&path, home)?; + fs::create_dir_all(&path).with_context(|| format!("create {}", path.display()))?; + if let Some(mode) = mode { + set_mode(&path, mode)?; + } + send_file_response_to_client(state, socket, client_id, stream_id, FileResponse::Ok) + .await + } + FileRequest::Remove { path, recursive } => { + let path = clean_remote_path(&path, home)?; + let metadata = + fs::symlink_metadata(&path).with_context(|| format!("stat {}", path.display()))?; + if metadata.is_dir() { + anyhow::ensure!( + recursive, + "{} is a directory; pass -r to remove directories", + path.display() + ); + fs::remove_dir_all(&path).with_context(|| format!("remove {}", path.display()))?; + } else { + fs::remove_file(&path).with_context(|| format!("remove {}", path.display()))?; + } + send_file_response_to_client(state, socket, client_id, stream_id, FileResponse::Ok) + .await + } + FileRequest::PutStart { + path, + size, + mode, + modified_secs, + overwrite, + resume, + } => { + if upload.is_some() { + bail!("upload already in progress"); + } + let final_path = clean_remote_path(&path, home)?; + if let Some(parent) = final_path.parent() { + fs::create_dir_all(parent) + .with_context(|| format!("create {}", parent.display()))?; + } + if final_path.exists() && !overwrite && !resume { + bail!("destination exists: {}", final_path.display()); + } + let mut hasher = Sha256::new(); + let (file, temp_path, written, atomic) = if resume && final_path.exists() { + let mut file = fs::OpenOptions::new() + .read(true) + .append(true) + .open(&final_path) + .with_context(|| format!("open {}", final_path.display()))?; + let written = file.metadata()?.len().min(size); + if written > 0 { + let mut prefix = fs::File::open(&final_path)?; + hash_exact_prefix(&mut prefix, written, &mut hasher)?; + } + file.seek(SeekFrom::Start(written))?; + (file, None, written, false) + } else { + let temp_path = upload_temp_path(&final_path, stream_id); + let file = fs::OpenOptions::new() + .create(true) + .truncate(true) + .write(true) + .open(&temp_path) + .with_context(|| format!("create {}", temp_path.display()))?; + (file, Some(temp_path), 0, true) + }; + *upload = Some(UploadState { + final_path, + temp_path, + file, + hasher, + written, + expected_size: size, + modified_secs, + atomic, + }); + if let Some(mode) = mode { + let target = upload + .as_ref() + .and_then(|state| state.temp_path.as_ref()) + .unwrap_or_else(|| &upload.as_ref().unwrap().final_path) + .clone(); + set_mode(&target, mode)?; + } + send_file_response_to_client( + state, + socket, + client_id, + stream_id, + FileResponse::Resume { offset: written }, + ) + .await + } + FileRequest::PutChunk { offset, bytes } => { + let Some(active) = upload.as_mut() else { + bail!("no upload in progress"); + }; + anyhow::ensure!( + offset == active.written, + "upload offset mismatch: got {offset}, expected {}", + active.written + ); + active.file.write_all(&bytes)?; + active.hasher.update(&bytes); + active.written = active.written.saturating_add(bytes.len() as u64); + Ok(()) + } + FileRequest::PutFinish { sha256 } => { + let Some(mut active) = upload.take() else { + bail!("no upload in progress"); + }; + active.file.flush()?; + active.file.sync_all()?; + anyhow::ensure!( + active.written == active.expected_size, + "upload size mismatch: got {}, expected {}", + active.written, + active.expected_size + ); + let digest: [u8; 32] = active.hasher.finalize().into(); + if digest != sha256 { + if let Some(temp_path) = &active.temp_path { + let _ = fs::remove_file(temp_path); + } + bail!("upload checksum mismatch"); + } + drop(active.file); + if active.atomic + && let Some(temp_path) = active.temp_path.take() + { + fs::rename(&temp_path, &active.final_path).with_context(|| { + format!( + "rename {} to {}", + temp_path.display(), + active.final_path.display() + ) + })?; + } + if let Some(modified_secs) = active.modified_secs { + filetime::set_file_mtime( + &active.final_path, + filetime::FileTime::from_unix_time(modified_secs as i64, 0), + ) + .with_context(|| format!("set mtime {}", active.final_path.display()))?; + } + send_file_response_to_client( + state, + socket, + client_id, + stream_id, + FileResponse::Done { + sha256: digest, + bytes: active.written, + }, + ) + .await + } + FileRequest::Get { path, offset } => { + let path = clean_remote_path(&path, home)?; + let meta = file_meta(&path, home)?; + anyhow::ensure!( + meta.kind == FileKind::File, + "remote path is not a file: {}", + meta.path + ); + let mut file = + fs::File::open(&path).with_context(|| format!("open {}", path.display()))?; + file.seek(SeekFrom::Start(offset))?; + send_file_response_to_client( + state, + socket, + client_id, + stream_id, + FileResponse::Start { + meta: meta.clone(), + offset, + }, + ) + .await?; + let mut hasher = Sha256::new(); + if offset > 0 { + let mut prefix = fs::File::open(&path)?; + hash_exact_prefix(&mut prefix, offset, &mut hasher)?; + } + let mut sent = offset; + let mut buf = vec![0u8; CHUNK_SIZE]; + loop { + let n = file.read(&mut buf)?; + if n == 0 { + break; + } + hasher.update(&buf[..n]); + send_file_response_to_client( + state, + socket, + client_id, + stream_id, + FileResponse::Chunk { + offset: sent, + bytes: buf[..n].to_vec(), + }, + ) + .await?; + sent = sent.saturating_add(n as u64); + } + let digest: [u8; 32] = hasher.finalize().into(); + send_file_response_to_client( + state, + socket, + client_id, + stream_id, + FileResponse::Done { + sha256: digest, + bytes: sent, + }, + ) + .await + } + } +} + +async fn send_file_response_to_client( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, + response: FileResponse, +) -> Result<()> { + send_stream_data_to_client( + state, + socket, + client_id, + stream_id, + encode_response(&response)?, + ) + .await +} + +async fn run_exec_stream_service( + state: Arc>, + socket: Arc, + client_id: [u8; 16], + stream_id: u64, + mut rx: mpsc::Receiver>, +) -> Result<()> { + let mut decoder = ExecFrameDecoder::default(); + while let Some(bytes) = rx.recv().await { + for frame in decoder.push(&bytes)? { + let request = match dosh::exec_service::decode_request(&frame) { + Ok(request) => request, + Err(err) => { + send_exec_response_to_client( + &state, + &socket, + client_id, + stream_id, + ExecResponse::Error { + message: err.to_string(), + }, + ) + .await?; + continue; + } + }; + let output = TokioCommand::new("sh") + .arg("-lc") + .arg(&request.command) + .stdin(Stdio::null()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output() + .await + .with_context(|| format!("run command {:?}", request.command))?; + for chunk in output.stdout.chunks(CHUNK_SIZE) { + send_exec_response_to_client( + &state, + &socket, + client_id, + stream_id, + ExecResponse::Stdout(chunk.to_vec()), + ) + .await?; + } + for chunk in output.stderr.chunks(CHUNK_SIZE) { + send_exec_response_to_client( + &state, + &socket, + client_id, + stream_id, + ExecResponse::Stderr(chunk.to_vec()), + ) + .await?; + } + send_exec_response_to_client( + &state, + &socket, + client_id, + stream_id, + ExecResponse::Exit { + code: output.status.code(), + }, + ) + .await?; + } + } + Ok(()) +} + +async fn send_exec_response_to_client( + state: &Arc>, + socket: &Arc, + client_id: [u8; 16], + stream_id: u64, + response: ExecResponse, +) -> Result<()> { + send_stream_data_to_client( + state, + socket, + client_id, + stream_id, + encode_exec_response(&response)?, + ) + .await +} + +fn file_meta(path: &Path, home: &Path) -> Result { + let metadata = + fs::symlink_metadata(path).with_context(|| format!("stat {}", path.display()))?; + let file_type = metadata.file_type(); + let kind = if file_type.is_file() { + FileKind::File + } else if file_type.is_dir() { + FileKind::Directory + } else if file_type.is_symlink() { + FileKind::Symlink + } else { + FileKind::Other + }; + let modified_secs = metadata + .modified() + .ok() + .and_then(|mtime| mtime.duration_since(std::time::UNIX_EPOCH).ok()) + .map(|duration| duration.as_secs()); + let display_path = path + .strip_prefix(home) + .ok() + .and_then(|relative| relative.to_str()) + .filter(|relative| !relative.is_empty()) + .map(|relative| format!("~/{relative}")) + .unwrap_or_else(|| path.display().to_string()); + Ok(FileMeta { + path: display_path, + kind, + len: metadata.len(), + mode: Some(metadata.permissions().mode() & 0o7777), + modified_secs, + }) +} + +fn set_mode(path: &Path, mode: u32) -> Result<()> { + let mut permissions = fs::metadata(path)?.permissions(); + permissions.set_mode(mode & 0o7777); + fs::set_permissions(path, permissions).with_context(|| format!("chmod {}", path.display())) +} + +fn upload_temp_path(final_path: &Path, stream_id: u64) -> PathBuf { + let name = final_path + .file_name() + .and_then(|name| name.to_str()) + .unwrap_or("upload"); + final_path.with_file_name(format!( + ".{name}.dosh-upload-{}-{stream_id}", + std::process::id() + )) +} + +fn hash_exact_prefix(file: &mut fs::File, bytes: u64, hasher: &mut Sha256) -> Result<()> { + file.seek(SeekFrom::Start(0))?; + let mut remaining = bytes; + let mut buf = vec![0u8; CHUNK_SIZE]; + while remaining > 0 { + let want = remaining.min(buf.len() as u64) as usize; + let n = file.read(&mut buf[..want])?; + if n == 0 { + bail!("file ended while hashing resume prefix"); + } + hasher.update(&buf[..n]); + remaining -= n as u64; + } + Ok(()) +} + async fn send_stream_open_ok( state: &Arc>, socket: &Arc, diff --git a/src/client.rs b/src/client.rs new file mode 100644 index 0000000..7edce19 --- /dev/null +++ b/src/client.rs @@ -0,0 +1,442 @@ +use crate::config::{ + ClientConfig, HostConfig, expand_tilde, load_client_config, load_hosts_config, +}; +use crate::crypto; +use crate::native::{ + self, EnvVar, ForwardingKind, ForwardingRequest, KnownHostStatus, NativeClientHello, + derive_native_session_key, generate_native_ephemeral, sign_user_auth_with_private_key, + supported_user_key_algorithms, trust_host, verify_known_host, verify_server_hello, +}; +use crate::protocol::{ + self, AttachReject, CLIENT_TO_SERVER, NativeAuthOkBody, NativeClientHelloBody, + NativeServerHelloBody, NativeUserAuthBody, PacketKind, SERVER_TO_CLIENT, +}; +use crate::ssh_agent; +use crate::transport::{DoshTransport, SessionRole, SessionTransportConfig, TransportConfig}; +use anyhow::{Context, Result, anyhow, bail}; +use std::net::{SocketAddr, ToSocketAddrs}; +use std::path::PathBuf; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; +use tokio::net::UdpSocket; + +#[derive(Debug, Clone)] +pub struct DoshClient { + config: ClientConfig, + hosts: crate::config::HostsConfig, +} + +impl DoshClient { + pub fn load() -> Result { + Ok(Self { + config: load_client_config(None)?, + hosts: load_hosts_config(None)?, + }) + } + + pub fn load_from_paths( + client_config: Option, + hosts_config: Option, + ) -> Result { + Ok(Self { + config: load_client_config(client_config)?, + hosts: load_hosts_config(hosts_config)?, + }) + } + + pub fn with_config(config: ClientConfig, hosts: crate::config::HostsConfig) -> Self { + Self { config, hosts } + } + + pub fn connect(&self, host: impl Into) -> DoshClientBuilder { + DoshClientBuilder::new(self.clone(), host.into()) + } +} + +#[derive(Debug, Clone)] +pub struct DoshClientBuilder { + client: DoshClient, + host: String, + services: Vec, + identity_files: Vec, + session: Option, + user: Option, + udp_host: Option, + udp_port: Option, + trust_on_first_use: Option, + use_ssh_agent: Option, + timeout: Option, + env: Vec, +} + +impl DoshClientBuilder { + pub fn new(client: DoshClient, host: String) -> Self { + Self { + client, + host, + services: Vec::new(), + identity_files: Vec::new(), + session: None, + user: None, + udp_host: None, + udp_port: None, + trust_on_first_use: None, + use_ssh_agent: None, + timeout: None, + env: Vec::new(), + } + } + + pub fn service(mut self, name: impl Into) -> Self { + self.services.push(name.into()); + self + } + + pub fn services(mut self, names: impl IntoIterator>) -> Self { + self.services.extend(names.into_iter().map(Into::into)); + self + } + + pub fn identity_file(mut self, path: impl Into) -> Self { + self.identity_files.push(path.into()); + self + } + + pub fn session(mut self, session: impl Into) -> Self { + self.session = Some(session.into()); + self + } + + pub fn user(mut self, user: impl Into) -> Self { + self.user = Some(user.into()); + self + } + + pub fn udp_host(mut self, host: impl Into) -> Self { + self.udp_host = Some(host.into()); + self + } + + pub fn udp_port(mut self, port: u16) -> Self { + self.udp_port = Some(port); + self + } + + pub fn trust_on_first_use(mut self, trust: bool) -> Self { + self.trust_on_first_use = Some(trust); + self + } + + pub fn use_ssh_agent(mut self, value: bool) -> Self { + self.use_ssh_agent = Some(value); + self + } + + pub fn timeout(mut self, timeout: Duration) -> Self { + self.timeout = Some(timeout); + self + } + + pub fn env(mut self, name: impl Into, value: impl Into) -> Self { + self.env.push(EnvVar { + name: name.into(), + value: value.into(), + }); + self + } + + pub async fn connect(self) -> Result { + let host_config = self + .client + .hosts + .hosts + .get(&self.host) + .cloned() + .unwrap_or_default(); + let raw_server = host_config.ssh.clone().unwrap_or_else(|| self.host.clone()); + let requested_user = self + .user + .clone() + .or_else(|| host_config.user.clone()) + .or_else(|| user_from_destination(&raw_server)) + .or_else(|| std::env::var("USER").ok()) + .unwrap_or_else(|| "unknown".to_string()); + let udp_host = self + .udp_host + .clone() + .or_else(|| host_config.dosh_host.clone()) + .or_else(|| self.client.config.dosh_host.clone()) + .unwrap_or_else(|| destination_host(&raw_server)); + let udp_port = self + .udp_port + .or(host_config.port) + .unwrap_or(self.client.config.dosh_port); + let peer_addr = resolve_addr(&udp_host, udp_port)?; + let timeout = self.timeout.unwrap_or_else(|| { + Duration::from_millis(self.client.config.native_auth_timeout_ms.max(1)) + }); + let session = self.session.unwrap_or_else(default_sdk_session); + let requested_forwardings = self + .services + .iter() + .map(|service| { + Ok(ForwardingRequest { + kind: ForwardingKind::Local, + bind_host: None, + listen_port: 0, + target_host: Some(crate::transport::service_target(service)?), + target_port: Some(0), + }) + }) + .collect::>>()?; + let socket = UdpSocket::bind("0.0.0.0:0").await?; + let (client_secret, client_public) = generate_native_ephemeral(); + let hello = NativeClientHello { + protocol_version: native::NATIVE_PROTOCOL_VERSION, + client_random: crypto::random_32(), + client_ephemeral_public: client_public, + requested_host: self.host.clone(), + requested_user, + requested_session: session.clone(), + requested_mode: "forward-only".to_string(), + terminal_size: (80, 24), + supported_aead: vec!["chacha20poly1305".to_string()], + supported_user_key_algorithms: supported_user_key_algorithms(), + cached_host_key_fingerprint: None, + attach_ticket_envelope: None, + requested_env: self.env, + }; + let packet = protocol::encode_plain( + PacketKind::NativeClientHello, + [0u8; 16], + 1, + 0, + &protocol::to_body(&NativeClientHelloBody { + hello: hello.clone(), + })?, + )?; + socket.send_to(&packet, peer_addr).await?; + + let mut buf = vec![0u8; 65535]; + let (n, _) = tokio::time::timeout(timeout, socket.recv_from(&mut buf)).await??; + let packet = protocol::decode(&buf[..n])?; + if packet.header.kind != PacketKind::NativeServerHello { + if packet.header.kind == PacketKind::AttachReject { + let reject: AttachReject = protocol::from_body(&packet.body)?; + bail!("native auth rejected: {}", reject.reason); + } + bail!("native auth received unexpected server response"); + } + let server_hello: NativeServerHelloBody = protocol::from_body(&packet.body)?; + verify_server_hello(&hello, &server_hello.hello)?; + verify_or_trust_host( + &self.client.config, + &self.host, + &server_hello.hello.host_key, + self.trust_on_first_use, + )?; + + let session_key = derive_native_session_key( + &client_secret, + server_hello.hello.server_ephemeral_public, + &hello, + &server_hello.hello, + )?; + let auth = sign_auth( + &self.client.config, + &host_config, + &hello, + &server_hello.hello, + requested_forwardings, + self.identity_files, + self.use_ssh_agent, + )?; + let mut pending_id = [0u8; 16]; + pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]); + let auth_packet = protocol::encode_encrypted( + PacketKind::NativeUserAuth, + pending_id, + 2, + 1, + &session_key, + CLIENT_TO_SERVER, + &protocol::to_body(&NativeUserAuthBody { auth })?, + )?; + socket.send_to(&auth_packet, peer_addr).await?; + let (n, _) = tokio::time::timeout(timeout, socket.recv_from(&mut buf)).await??; + let packet = protocol::decode(&buf[..n])?; + if packet.header.kind != PacketKind::NativeAuthOk { + if packet.header.kind == PacketKind::AttachReject { + let reject: AttachReject = protocol::from_body(&packet.body)?; + bail!("native auth rejected: {}", reject.reason); + } + bail!("native auth received unexpected auth response"); + } + let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT)?; + let ok: NativeAuthOkBody = protocol::from_body(&plain)?; + let transport = DoshTransport::new_owned( + socket, + SessionTransportConfig { + role: SessionRole::Client, + conn_id: ok.ok.client_id, + session_key: ok.ok.session_key, + peer_addr, + initial_send_seq: 2, + initial_ack: ok.ok.initial_seq, + stream: TransportConfig::default(), + }, + ); + Ok(ConnectedDoshClient { + host: self.host, + session: ok.ok.session, + transport, + }) + } +} + +pub struct ConnectedDoshClient { + pub host: String, + pub session: String, + pub transport: DoshTransport, +} + +impl ConnectedDoshClient { + pub fn into_transport(self) -> DoshTransport { + self.transport + } +} + +fn verify_or_trust_host( + config: &ClientConfig, + host: &str, + host_key: &native::HostPublicKey, + trust_override: Option, +) -> Result<()> { + let known_hosts = expand_tilde(&config.known_hosts); + match verify_known_host(&known_hosts, host, host_key)? { + KnownHostStatus::Trusted => Ok(()), + KnownHostStatus::Unknown if trust_override.unwrap_or(config.trust_on_first_use) => { + trust_host(&known_hosts, host, host_key, "sdk-tofu", false)?; + Ok(()) + } + KnownHostStatus::Unknown => Err(anyhow!( + "Dosh host key for {host} is not trusted; run `dosh trust {host}` first or enable trust_on_first_use" + )), + KnownHostStatus::Mismatch { expected, actual } => Err(anyhow!( + "Dosh host key mismatch for {host}: expected {expected}, got {actual}" + )), + } +} + +fn sign_auth( + config: &ClientConfig, + host_config: &HostConfig, + hello: &NativeClientHello, + server_hello: &native::NativeServerHello, + requested_forwardings: Vec, + explicit_identity_files: Vec, + use_ssh_agent: Option, +) -> Result { + let use_agent = use_ssh_agent.unwrap_or(config.use_ssh_agent); + let mut errors = Vec::new(); + if use_agent { + 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:#}")), + } + } + + let mut paths = explicit_identity_files; + if paths.is_empty() { + paths.extend(config.identity_files.iter().map(|path| expand_tilde(path))); + } + if paths.is_empty() { + paths.extend(default_identity_paths()); + } + for path in paths { + match native::load_native_identity(&path).and_then(|identity| { + sign_user_auth_with_private_key( + &identity, + hello, + server_hello, + requested_forwardings.clone(), + ) + }) { + Ok(auth) => return Ok(auth), + Err(err) => errors.push(format!("{}: {err:#}", path.display())), + } + } + let _ = host_config; + Err(anyhow!( + "native auth found no usable identity: {}", + errors.join("; ") + )) +} + +fn default_identity_paths() -> Vec { + ["~/.ssh/id_ed25519", "~/.ssh/id_ecdsa", "~/.ssh/id_rsa"] + .into_iter() + .map(expand_tilde) + .collect() +} + +fn resolve_addr(host: &str, port: u16) -> Result { + (host, port) + .to_socket_addrs() + .with_context(|| format!("resolve UDP target {host}:{port}"))? + .next() + .ok_or_else(|| anyhow!("no UDP address resolved for {host}:{port}")) +} + +fn destination_host(destination: &str) -> String { + destination + .rsplit('@') + .next() + .unwrap_or(destination) + .split(':') + .next() + .unwrap_or(destination) + .to_string() +} + +fn user_from_destination(destination: &str) -> Option { + destination + .rsplit_once('@') + .map(|(user, _)| user.to_string()) + .filter(|user| !user.is_empty()) +} + +fn default_sdk_session() -> String { + let millis = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_millis(); + format!("sdk-{millis}-{}", std::process::id()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_user_and_host_from_destination() { + assert_eq!( + user_from_destination("palav@example.com").as_deref(), + Some("palav") + ); + assert_eq!(destination_host("palav@example.com"), "example.com"); + assert_eq!(destination_host("example.com:2222"), "example.com"); + } + + #[test] + fn default_identity_paths_are_expanded() { + assert!( + default_identity_paths() + .iter() + .any(|path| path.ends_with(".ssh/id_ed25519")) + ); + } +} diff --git a/src/config.rs b/src/config.rs index 8719f9d..7dbc998 100644 --- a/src/config.rs +++ b/src/config.rs @@ -46,6 +46,10 @@ pub struct ServerConfig { pub allow_remote_non_loopback_bind: bool, #[serde(default)] pub allow_agent_forwarding: bool, + #[serde(default = "default_true")] + pub allow_file_transfer: bool, + #[serde(default = "default_true")] + pub allow_exec_command: bool, #[serde(default = "default_accept_env")] pub accept_env: Vec, /// Run terminal sessions under per-session holder processes so shells @@ -82,6 +86,8 @@ impl Default for ServerConfig { allow_remote_forwarding: false, allow_remote_non_loopback_bind: false, allow_agent_forwarding: false, + allow_file_transfer: true, + allow_exec_command: true, accept_env: default_accept_env(), persist_sessions: true, } diff --git a/src/exec_service.rs b/src/exec_service.rs new file mode 100644 index 0000000..40f4f59 --- /dev/null +++ b/src/exec_service.rs @@ -0,0 +1,71 @@ +use anyhow::{Context, Result, bail}; +use serde::{Deserialize, Serialize}; + +pub const EXEC_STREAM_SENTINEL: &str = "@dosh-exec"; +pub const FRAME_MAX_LEN: usize = 1024 * 1024; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct ExecRequest { + pub command: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub enum ExecResponse { + Stdout(Vec), + Stderr(Vec), + Exit { code: Option }, + Error { message: String }, +} + +#[derive(Default)] +pub struct FrameDecoder { + buf: Vec, +} + +impl FrameDecoder { + pub fn push(&mut self, bytes: &[u8]) -> Result>> { + self.buf.extend_from_slice(bytes); + let mut frames = Vec::new(); + loop { + if self.buf.len() < 4 { + break; + } + let len = u32::from_be_bytes(self.buf[..4].try_into().unwrap()) as usize; + if len > FRAME_MAX_LEN { + bail!("exec protocol frame too large: {len} bytes"); + } + if self.buf.len() < 4 + len { + break; + } + frames.push(self.buf[4..4 + len].to_vec()); + self.buf.drain(..4 + len); + } + Ok(frames) + } +} + +pub fn encode_request(request: &ExecRequest) -> Result> { + encode_frame(&bincode::serialize(request).context("encode exec request")?) +} + +pub fn encode_response(response: &ExecResponse) -> Result> { + encode_frame(&bincode::serialize(response).context("encode exec response")?) +} + +pub fn decode_request(frame: &[u8]) -> Result { + bincode::deserialize(frame).context("decode exec request") +} + +pub fn decode_response(frame: &[u8]) -> Result { + bincode::deserialize(frame).context("decode exec response") +} + +fn encode_frame(payload: &[u8]) -> Result> { + if payload.len() > FRAME_MAX_LEN { + bail!("exec protocol frame too large: {} bytes", payload.len()); + } + let mut out = Vec::with_capacity(4 + payload.len()); + out.extend_from_slice(&(payload.len() as u32).to_be_bytes()); + out.extend_from_slice(payload); + Ok(out) +} diff --git a/src/file_transfer.rs b/src/file_transfer.rs new file mode 100644 index 0000000..b939585 --- /dev/null +++ b/src/file_transfer.rs @@ -0,0 +1,250 @@ +use anyhow::{Context, Result, anyhow, bail}; +use serde::{Deserialize, Serialize}; +use std::path::{Path, PathBuf}; + +pub const FILE_STREAM_SENTINEL: &str = "@dosh-file"; +pub const FRAME_MAX_LEN: usize = 1024 * 1024; +pub const CHUNK_SIZE: usize = 32 * 1024; + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub enum FileRequest { + Stat { + path: String, + }, + List { + path: String, + }, + Mkdir { + path: String, + mode: Option, + }, + Remove { + path: String, + recursive: bool, + }, + PutStart { + path: String, + size: u64, + mode: Option, + modified_secs: Option, + overwrite: bool, + resume: bool, + }, + PutChunk { + offset: u64, + bytes: Vec, + }, + PutFinish { + sha256: [u8; 32], + }, + Get { + path: String, + offset: u64, + }, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub enum FileResponse { + Ok, + Resume { offset: u64 }, + Stat { meta: FileMeta }, + List { entries: Vec }, + Start { meta: FileMeta, offset: u64 }, + Chunk { offset: u64, bytes: Vec }, + Done { sha256: [u8; 32], bytes: u64 }, + Error { message: String }, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct FileMeta { + pub path: String, + pub kind: FileKind, + pub len: u64, + pub mode: Option, + pub modified_secs: Option, +} + +#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)] +pub enum FileKind { + File, + Directory, + Symlink, + Other, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] +pub struct FileEntry { + pub name: String, + pub meta: FileMeta, +} + +#[derive(Default)] +pub struct FrameDecoder { + buf: Vec, +} + +impl FrameDecoder { + pub fn push(&mut self, bytes: &[u8]) -> Result>> { + self.buf.extend_from_slice(bytes); + let mut frames = Vec::new(); + loop { + if self.buf.len() < 4 { + break; + } + let len = u32::from_be_bytes(self.buf[..4].try_into().unwrap()) as usize; + if len > FRAME_MAX_LEN { + bail!("file protocol frame too large: {len} bytes"); + } + if self.buf.len() < 4 + len { + break; + } + frames.push(self.buf[4..4 + len].to_vec()); + self.buf.drain(..4 + len); + } + Ok(frames) + } +} + +pub fn encode_request(request: &FileRequest) -> Result> { + encode_frame(&bincode::serialize(request).context("encode file request")?) +} + +pub fn encode_response(response: &FileResponse) -> Result> { + encode_frame(&bincode::serialize(response).context("encode file response")?) +} + +pub fn decode_request(frame: &[u8]) -> Result { + bincode::deserialize(frame).context("decode file request") +} + +pub fn decode_response(frame: &[u8]) -> Result { + bincode::deserialize(frame).context("decode file response") +} + +pub fn encode_frame(payload: &[u8]) -> Result> { + if payload.len() > FRAME_MAX_LEN { + bail!("file protocol frame too large: {} bytes", payload.len()); + } + let mut out = Vec::with_capacity(4 + payload.len()); + out.extend_from_slice(&(payload.len() as u32).to_be_bytes()); + out.extend_from_slice(payload); + Ok(out) +} + +pub fn parse_copy_endpoint(raw: &str) -> CopyEndpoint { + if looks_like_windows_path(raw) { + return CopyEndpoint::Local(PathBuf::from(raw)); + } + let Some(index) = raw.find(':') else { + return CopyEndpoint::Local(PathBuf::from(raw)); + }; + let host = &raw[..index]; + let path = &raw[index + 1..]; + if host.is_empty() || host.contains('/') || host.contains('\\') { + return CopyEndpoint::Local(PathBuf::from(raw)); + } + CopyEndpoint::Remote { + host: host.to_string(), + path: if path.is_empty() { + ".".to_string() + } else { + path.to_string() + }, + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum CopyEndpoint { + Local(PathBuf), + Remote { host: String, path: String }, +} + +pub fn remote_join(parent: &str, child: &str) -> String { + if parent.is_empty() || parent == "." { + return child.to_string(); + } + format!("{}/{}", parent.trim_end_matches('/'), child) +} + +pub fn remote_basename(path: &str) -> String { + path.trim_end_matches('/') + .rsplit('/') + .next() + .filter(|value| !value.is_empty()) + .unwrap_or(path) + .to_string() +} + +pub fn clean_remote_path(path: &str, home: &Path) -> Result { + if path.as_bytes().contains(&0) { + bail!("remote path contains NUL"); + } + let expanded = if path == "~" { + home.to_path_buf() + } else if let Some(rest) = path.strip_prefix("~/") { + home.join(rest) + } else { + let raw = Path::new(path); + if raw.is_absolute() { + raw.to_path_buf() + } else { + home.join(raw) + } + }; + Ok(expanded) +} + +pub fn ensure_relative_child(path: &Path) -> Result { + let name = path + .file_name() + .and_then(|name| name.to_str()) + .ok_or_else(|| anyhow!("path has no final component: {}", path.display()))?; + if name.is_empty() || name == "." || name == ".." { + bail!("invalid path final component: {}", path.display()); + } + Ok(name.to_string()) +} + +fn looks_like_windows_path(raw: &str) -> bool { + let bytes = raw.as_bytes(); + bytes.len() >= 3 + && bytes[0].is_ascii_alphabetic() + && bytes[1] == b':' + && (bytes[2] == b'\\' || bytes[2] == b'/') +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn frames_round_trip_incrementally() { + let first = encode_response(&FileResponse::Ok).unwrap(); + let second = encode_response(&FileResponse::Error { + message: "nope".to_string(), + }) + .unwrap(); + let mut decoder = FrameDecoder::default(); + assert!(decoder.push(&first[..2]).unwrap().is_empty()); + let frames = decoder.push(&first[2..]).unwrap(); + assert_eq!(frames.len(), 1); + assert_eq!(decode_response(&frames[0]).unwrap(), FileResponse::Ok); + let frames = decoder.push(&second).unwrap(); + assert_eq!(frames.len(), 1); + } + + #[test] + fn copy_endpoint_parses_scp_style_but_not_windows_drive() { + assert_eq!( + parse_copy_endpoint("palav:tmp/file"), + CopyEndpoint::Remote { + host: "palav".to_string(), + path: "tmp/file".to_string() + } + ); + assert_eq!( + parse_copy_endpoint("C:\\Users\\palav\\x"), + CopyEndpoint::Local(PathBuf::from("C:\\Users\\palav\\x")) + ); + } +} diff --git a/src/lib.rs b/src/lib.rs index 08cb824..1cc5a62 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,11 +1,25 @@ +//! Dosh is an encrypted UDP transport for remote terminals and embeddable Rust +//! application streams. +//! +//! Use [`client::DoshClient`] and [`server::DoshServer`] when you want Dosh to +//! handle native authentication, host key checks, roaming, keepalives, reliable +//! streams, and service routing. Use [`transport::DoshTransport`] only when an +//! application already owns session establishment and wants direct access to the +//! encrypted stream transport. + pub mod auth; pub mod build_info; +pub mod client; pub mod config; pub mod crypto; +pub mod exec_service; +pub mod file_transfer; pub mod native; #[cfg(unix)] pub mod persist; pub mod protocol; #[cfg(unix)] pub mod pty; +pub mod server; pub mod ssh_agent; +pub mod transport; diff --git a/src/protocol.rs b/src/protocol.rs index 8e4ed82..437cb79 100644 --- a/src/protocol.rs +++ b/src/protocol.rs @@ -357,6 +357,10 @@ pub struct NativeAuthCheckOkBody { pub allow_tcp_forwarding: bool, pub allow_remote_forwarding: bool, pub allow_agent_forwarding: bool, + #[serde(default)] + pub allow_file_transfer: bool, + #[serde(default)] + pub allow_exec_command: bool, pub policy_flags: Vec, #[serde(default)] pub server_version: String, diff --git a/src/server.rs b/src/server.rs new file mode 100644 index 0000000..b295a96 --- /dev/null +++ b/src/server.rs @@ -0,0 +1,591 @@ +use crate::config::{ServerConfig, load_server_config}; +use crate::crypto; +use crate::native::{ + self, ForwardingKind, ForwardingRequest, NativeAuthOk, NativeServerHello, + derive_native_session_key, generate_native_ephemeral, host_public_key, load_or_create_host_key, + sign_server_hello, verify_native_user_auth_from_config, +}; +use crate::protocol::{ + self, AttachReject, CLIENT_TO_SERVER, NativeAuthOkBody, NativeClientHelloBody, + NativeServerHelloBody, NativeUserAuthBody, PacketKind, SERVER_TO_CLIENT, +}; +use crate::transport::{ + DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, TransportConfig, + service_name_from_target, +}; +use anyhow::{Context, Result, anyhow, bail}; +use ed25519_dalek::SigningKey; +use std::collections::{HashMap, HashSet}; +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::net::UdpSocket; + +#[derive(Debug, Clone)] +pub struct DoshServerConfig { + pub server: ServerConfig, + pub bind_addr: Option, + pub services: HashSet, + pub transport: TransportConfig, + pub require_current_user: bool, + pub auth_timeout: Duration, +} + +impl DoshServerConfig { + pub fn new(server: ServerConfig) -> Self { + Self { + server, + bind_addr: None, + services: HashSet::new(), + transport: TransportConfig::default(), + require_current_user: true, + auth_timeout: Duration::from_secs(30), + } + } + + pub fn bind_addr(mut self, addr: SocketAddr) -> Self { + self.bind_addr = Some(addr); + self + } + + pub fn service(mut self, name: impl Into) -> Result { + self.services.insert(validate_service_name(name.into())?); + Ok(self) + } + + pub fn services(mut self, names: impl IntoIterator>) -> Result { + for name in names { + self.services.insert(validate_service_name(name.into())?); + } + Ok(self) + } + + pub fn require_current_user(mut self, value: bool) -> Self { + self.require_current_user = value; + self + } + + pub fn transport(mut self, transport: TransportConfig) -> Self { + self.transport = transport; + self + } +} + +impl Default for DoshServerConfig { + fn default() -> Self { + Self::new(ServerConfig::default()) + } +} + +#[derive(Debug, Clone)] +pub struct DoshAccepted { + pub conn_id: [u8; 16], + pub user: String, + pub session: String, + pub services: Vec, + pub peer_addr: SocketAddr, +} + +#[derive(Debug, Clone)] +pub enum DoshServerEvent { + Accepted(DoshAccepted), + Session { + conn_id: [u8; 16], + event: SessionEvent, + }, + Ignored, +} + +struct PendingServerAuth { + client: native::NativeClientHello, + server: NativeServerHello, + session_key: [u8; 32], + peer: SocketAddr, + created_at: Instant, +} + +pub struct DoshServer { + socket: Arc, + config: DoshServerConfig, + host_signing: SigningKey, + pending: HashMap<[u8; 16], PendingServerAuth>, + transports: HashMap<[u8; 16], DoshTransport>, + accepted: HashMap<[u8; 16], DoshAccepted>, +} + +impl DoshServer { + pub async fn load() -> Result { + Self::bind(DoshServerConfig::new(load_server_config(None)?)).await + } + + pub async fn load_from_path(path: Option) -> Result { + Self::bind(DoshServerConfig::new(load_server_config(path)?)).await + } + + pub async fn bind(config: DoshServerConfig) -> Result { + let bind_addr = match config.bind_addr { + Some(addr) => addr, + None => format!("{}:{}", config.server.bind, config.server.port) + .parse() + .with_context(|| { + format!( + "parse Dosh bind address {}:{}", + config.server.bind, config.server.port + ) + })?, + }; + let host_signing = load_or_create_host_key(&config.server)?; + let socket = Arc::new(UdpSocket::bind(bind_addr).await?); + Ok(Self { + socket, + config, + host_signing, + pending: HashMap::new(), + transports: HashMap::new(), + accepted: HashMap::new(), + }) + } + + pub fn local_addr(&self) -> Result { + Ok(self.socket.local_addr()?) + } + + pub fn connection(&self, conn_id: &[u8; 16]) -> Option<&DoshAccepted> { + self.accepted.get(conn_id) + } + + pub fn transport(&self, conn_id: &[u8; 16]) -> Option<&DoshTransport> { + self.transports.get(conn_id) + } + + pub fn transport_mut(&mut self, conn_id: &[u8; 16]) -> Option<&mut DoshTransport> { + self.transports.get_mut(conn_id) + } + + pub async fn recv(&mut self) -> Result { + let mut buf = vec![0u8; 65535]; + loop { + self.expire_pending(); + let (n, peer) = self.socket.recv_from(&mut buf).await?; + let packet = match protocol::decode(&buf[..n]) { + Ok(packet) => packet, + Err(_) => continue, + }; + if packet.header.kind == PacketKind::NativeClientHello { + self.handle_client_hello(peer, packet.body).await?; + return Ok(DoshServerEvent::Ignored); + } + if packet.header.kind == PacketKind::NativeUserAuth { + return self.handle_user_auth(peer, &packet).await; + } + if let Some(transport) = self.transports.get_mut(&packet.header.conn_id) { + let event = transport.handle_datagram(&buf[..n], peer).await?; + return Ok(DoshServerEvent::Session { + conn_id: packet.header.conn_id, + event, + }); + } + self.send_reject(peer, packet.header.conn_id, "unknown Dosh connection") + .await?; + return Ok(DoshServerEvent::Ignored); + } + } + + pub async fn accept_stream(&mut self, conn_id: [u8; 16], stream_id: u64) -> Result<()> { + self.transport_mut(&conn_id) + .ok_or_else(|| anyhow!("unknown Dosh connection"))? + .accept_stream(stream_id) + .await + } + + pub async fn reject_stream( + &mut self, + conn_id: [u8; 16], + stream_id: u64, + reason: impl Into, + ) -> Result<()> { + self.transport_mut(&conn_id) + .ok_or_else(|| anyhow!("unknown Dosh connection"))? + .reject_stream(stream_id, reason) + .await + } + + pub async fn send( + &mut self, + conn_id: [u8; 16], + stream_id: u64, + bytes: impl Into>, + ) -> Result<()> { + self.transport_mut(&conn_id) + .ok_or_else(|| anyhow!("unknown Dosh connection"))? + .send(stream_id, bytes) + .await + } + + pub async fn close(&mut self, conn_id: [u8; 16], stream_id: u64) -> Result<()> { + self.transport_mut(&conn_id) + .ok_or_else(|| anyhow!("unknown Dosh connection"))? + .close(stream_id) + .await + } + + pub async fn maintenance(&mut self) -> Result<()> { + for transport in self.transports.values_mut() { + transport.maintenance().await?; + } + Ok(()) + } + + async fn handle_client_hello(&mut self, peer: SocketAddr, body: Vec) -> Result<()> { + let req: NativeClientHelloBody = protocol::from_body(&body)?; + if let Err(err) = + native::check_native_protocol_version(req.hello.protocol_version, "client") + { + self.send_reject(peer, [0u8; 16], &err.to_string()).await?; + return Ok(()); + } + let result = self.build_server_hello(req.hello, peer); + let (pending_id, hello) = match result { + Ok(value) => value, + Err(err) => { + self.send_reject(peer, [0u8; 16], &err.to_string()).await?; + return Ok(()); + } + }; + let body = protocol::to_body(&NativeServerHelloBody { hello })?; + let out = protocol::encode_plain(PacketKind::NativeServerHello, pending_id, 1, 0, &body)?; + self.socket.send_to(&out, peer).await?; + Ok(()) + } + + fn build_server_hello( + &mut self, + client: native::NativeClientHello, + peer: SocketAddr, + ) -> Result<([u8; 16], NativeServerHello)> { + if !self.config.server.native_auth { + bail!("native auth disabled"); + } + if !client + .supported_aead + .iter() + .any(|algorithm| algorithm == "chacha20poly1305") + { + bail!("native auth requires chacha20poly1305"); + } + if !client + .supported_user_key_algorithms + .iter() + .any(|algorithm| native::is_supported_user_signature_algorithm(algorithm)) + { + bail!("native auth requires a supported user key algorithm"); + } + if self.config.require_current_user { + let current_user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string()); + if client.requested_user != current_user { + bail!("native auth user mismatch"); + } + } + + let (server_secret, server_public) = generate_native_ephemeral(); + let mut server = NativeServerHello { + protocol_version: native::NATIVE_PROTOCOL_VERSION, + server_random: crypto::random_32(), + server_ephemeral_public: server_public, + host_key: host_public_key(&self.host_signing), + chosen_aead: "chacha20poly1305".to_string(), + server_key_epoch: 1, + auth_challenge: crypto::random_32(), + rate_limit_remaining: None, + host_signature: Vec::new(), + }; + sign_server_hello(&self.host_signing, &client, &mut server)?; + let session_key = derive_native_session_key( + &server_secret, + client.client_ephemeral_public, + &client, + &server, + )?; + let mut pending_id = [0u8; 16]; + pending_id.copy_from_slice(&server.auth_challenge[..16]); + self.pending.insert( + pending_id, + PendingServerAuth { + client, + server: server.clone(), + session_key, + peer, + created_at: Instant::now(), + }, + ); + Ok((pending_id, server)) + } + + async fn handle_user_auth( + &mut self, + peer: SocketAddr, + packet: &protocol::Packet, + ) -> Result { + let Some(pending) = self.pending.remove(&packet.header.conn_id) else { + self.send_reject(peer, packet.header.conn_id, "unknown native auth") + .await?; + return Ok(DoshServerEvent::Ignored); + }; + if pending.peer != peer { + self.send_reject(peer, packet.header.conn_id, "native auth peer changed") + .await?; + return Ok(DoshServerEvent::Ignored); + } + if pending.created_at.elapsed() > self.config.auth_timeout { + self.send_reject(peer, packet.header.conn_id, "native auth expired") + .await?; + return Ok(DoshServerEvent::Ignored); + } + + let body = protocol::decrypt_body(packet, &pending.session_key, CLIENT_TO_SERVER)?; + let req: NativeUserAuthBody = protocol::from_body(&body)?; + let services = match self.verify_auth_and_services(&pending, &req.auth) { + Ok(services) => services, + Err(err) => { + self.send_reject(peer, packet.header.conn_id, &format!("{err:#}")) + .await?; + return Ok(DoshServerEvent::Ignored); + } + }; + + let conn_id = crypto::random_16(); + let key_id = protocol::session_key_id(&pending.session_key); + let ok = NativeAuthOk { + client_id: conn_id, + session: pending.client.requested_session.clone(), + mode: pending.client.requested_mode.clone(), + session_key: pending.session_key, + session_key_id: key_id, + attach_ticket: Vec::new(), + attach_ticket_psk: crypto::random_32(), + initial_seq: 1, + snapshot: Vec::new(), + policy_flags: services + .iter() + .map(|service| format!("service:{service}")) + .collect(), + }; + let body = protocol::to_body(&NativeAuthOkBody { ok })?; + let out = protocol::encode_encrypted( + PacketKind::NativeAuthOk, + conn_id, + 1, + packet.header.seq, + &pending.session_key, + SERVER_TO_CLIENT, + &body, + )?; + self.socket.send_to(&out, peer).await?; + + let transport = DoshTransport::new( + Arc::clone(&self.socket), + SessionTransportConfig { + role: SessionRole::Server, + conn_id, + session_key: pending.session_key, + peer_addr: peer, + initial_send_seq: 2, + initial_ack: packet.header.seq, + stream: self.config.transport.clone(), + }, + ); + let accepted = DoshAccepted { + conn_id, + user: pending.client.requested_user, + session: pending.client.requested_session, + services, + peer_addr: peer, + }; + self.transports.insert(conn_id, transport); + self.accepted.insert(conn_id, accepted.clone()); + Ok(DoshServerEvent::Accepted(accepted)) + } + + fn verify_auth_and_services( + &self, + pending: &PendingServerAuth, + auth: &native::NativeUserAuth, + ) -> Result> { + verify_native_user_auth_from_config( + &self.config.server, + &pending.client, + &pending.server, + auth, + Some(pending.peer.ip()), + ) + .context("verify native user auth")?; + validate_requested_services(&self.config.services, &auth.requested_forwardings) + } + + async fn send_reject(&self, peer: SocketAddr, conn_id: [u8; 16], reason: &str) -> Result<()> { + let body = protocol::to_body(&AttachReject { + reason: reason.to_string(), + })?; + let out = protocol::encode_plain(PacketKind::AttachReject, conn_id, 0, 0, &body)?; + self.socket.send_to(&out, peer).await?; + Ok(()) + } + + fn expire_pending(&mut self) { + let timeout = self.config.auth_timeout; + self.pending + .retain(|_, pending| pending.created_at.elapsed() <= timeout); + } +} + +fn validate_requested_services( + allowed_services: &HashSet, + requests: &[ForwardingRequest], +) -> Result> { + let mut services = Vec::new(); + for request in requests { + if request.kind != ForwardingKind::Local { + bail!("embedded Dosh services only accept local service requests"); + } + if request.listen_port != 0 || request.target_port != Some(0) { + bail!("embedded Dosh service requests must use port 0"); + } + let target = request + .target_host + .as_deref() + .ok_or_else(|| anyhow!("embedded Dosh service request is missing target host"))?; + let service = service_name_from_target(target) + .ok_or_else(|| anyhow!("target {target:?} is not a Dosh service"))?; + if !allowed_services.contains(service) { + bail!("Dosh service {service:?} is not registered"); + } + if !services.iter().any(|existing| existing == service) { + services.push(service.to_string()); + } + } + Ok(services) +} + +fn validate_service_name(name: String) -> Result { + crate::transport::service_target(&name)?; + Ok(name) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::client::DoshClient; + use crate::config::{ClientConfig, HostsConfig}; + use crate::transport::TransportEvent; + use ed25519_dalek::SigningKey; + + #[tokio::test] + async fn sdk_client_and_server_exchange_service_stream() { + let dir = tempfile::tempdir().unwrap(); + let host_key = dir.path().join("host_key"); + let authorized_keys = dir.path().join("authorized_keys"); + let identity = dir.path().join("id_ed25519"); + let known_hosts = dir.path().join("known_hosts"); + + let signing = SigningKey::from_bytes(&[91u8; 32]); + let keypair = ssh_key::private::Ed25519Keypair::from(&signing); + let private = + ssh_key::PrivateKey::new(ssh_key::private::KeypairData::from(keypair), "").unwrap(); + private + .write_openssh_file(&identity, ssh_key::LineEnding::LF) + .unwrap(); + std::fs::write( + &authorized_keys, + format!("{}\n", private.public_key().to_openssh().unwrap()), + ) + .unwrap(); + + let server_config = ServerConfig { + host_key: host_key.to_string_lossy().to_string(), + authorized_keys: vec![authorized_keys.to_string_lossy().to_string()], + ..ServerConfig::default() + }; + let server_config = DoshServerConfig::new(server_config) + .bind_addr("127.0.0.1:0".parse().unwrap()) + .service("echo") + .unwrap() + .require_current_user(false); + let mut server = DoshServer::bind(server_config).await.unwrap(); + let server_port = server.local_addr().unwrap().port(); + + let client_config = ClientConfig { + dosh_port: server_port, + trust_on_first_use: true, + known_hosts: known_hosts.to_string_lossy().to_string(), + identity_files: vec![identity.to_string_lossy().to_string()], + use_ssh_agent: false, + ..ClientConfig::default() + }; + let client = DoshClient::with_config(client_config, HostsConfig::default()); + let connect = client + .connect("127.0.0.1") + .user("sdk-user") + .service("echo") + .connect(); + let accept = async { + loop { + if let DoshServerEvent::Accepted(accepted) = server.recv().await.unwrap() { + break accepted; + } + } + }; + let (connected, accepted) = tokio::join!(connect, accept); + let mut client_transport = connected.unwrap().into_transport(); + let conn_id = accepted.conn_id; + assert_eq!(accepted.services, vec!["echo".to_string()]); + + let stream_id = client_transport.open_service("echo").await.unwrap(); + match server.recv().await.unwrap() { + DoshServerEvent::Session { + event: SessionEvent::Stream(TransportEvent::Open(open)), + .. + } => { + assert_eq!(open.stream_id, stream_id); + server.accept_stream(conn_id, open.stream_id).await.unwrap(); + } + other => panic!("unexpected event {other:?}"), + } + assert!(matches!( + client_transport.recv().await.unwrap(), + SessionEvent::Stream(TransportEvent::OpenOk { .. }) + )); + + client_transport + .send(stream_id, b"ping".to_vec()) + .await + .unwrap(); + loop { + match server.recv().await.unwrap() { + DoshServerEvent::Session { + event: SessionEvent::Stream(TransportEvent::Data(data)), + .. + } => { + assert_eq!(data.chunks, vec![b"ping".to_vec()]); + server + .send(conn_id, data.stream_id, b"pong".to_vec()) + .await + .unwrap(); + break; + } + DoshServerEvent::Session { .. } | DoshServerEvent::Ignored => {} + other => panic!("unexpected event {other:?}"), + } + } + loop { + match client_transport.recv().await.unwrap() { + SessionEvent::Stream(TransportEvent::Data(data)) => { + assert_eq!(data.chunks, vec![b"pong".to_vec()]); + break; + } + SessionEvent::Stream(_) | SessionEvent::Ignored => {} + other => panic!("unexpected event {other:?}"), + } + } + } +} diff --git a/src/transport.rs b/src/transport.rs new file mode 100644 index 0000000..8cfd50d --- /dev/null +++ b/src/transport.rs @@ -0,0 +1,1099 @@ +//! Embeddable reliable stream machinery for Dosh-native application transports. +//! +//! The terminal client/server binaries use the same wire packets exposed here: +//! `StreamOpen`, `StreamData`, `StreamWindowAdjust`, and `StreamClose`. This +//! module keeps the stream state machine independent from UDP sockets, native +//! auth, terminal rendering, or any particular service protocol. An embedding +//! application is expected to: +//! +//! 1. authenticate and maintain a Dosh session using [`crate::protocol`], +//! 2. encrypt each [`OutgoingStreamPacket::body`] as the given packet kind, +//! 3. pass received stream packet bodies back into [`StreamMux`], +//! 4. call [`StreamMux::tick`] periodically to retransmit lost stream data. +//! +//! That split is intentional: Dosh roaming/reconnect stays at the authenticated +//! session layer, while this module provides the reusable fast persistent byte +//! streams that higher-level app protocols can embed. + +use crate::protocol::{ + self, CLIENT_TO_SERVER, PacketKind, ReplayWindow, SERVER_TO_CLIENT, StreamClose, StreamData, + StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust, +}; +use anyhow::{Result, anyhow, bail}; +use std::collections::{BTreeMap, HashMap, HashSet, VecDeque}; +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::net::UdpSocket; + +pub const DEFAULT_INITIAL_WINDOW: usize = 1024 * 1024; +pub const DEFAULT_RETRANSMIT_AFTER: Duration = Duration::from_millis(200); +pub const DEFAULT_KEEPALIVE_AFTER: Duration = Duration::from_secs(2); +pub const SERVICE_TARGET_PREFIX: &str = "@dosh-"; + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct TransportConfig { + pub initial_window: usize, + pub retransmit_after: Duration, + pub keepalive_after: Duration, +} + +impl Default for TransportConfig { + fn default() -> Self { + Self { + initial_window: DEFAULT_INITIAL_WINDOW, + retransmit_after: DEFAULT_RETRANSMIT_AFTER, + keepalive_after: DEFAULT_KEEPALIVE_AFTER, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SessionRole { + Client, + Server, +} + +impl SessionRole { + fn send_direction(self) -> u32 { + match self { + Self::Client => CLIENT_TO_SERVER, + Self::Server => SERVER_TO_CLIENT, + } + } + + fn recv_direction(self) -> u32 { + match self { + Self::Client => SERVER_TO_CLIENT, + Self::Server => CLIENT_TO_SERVER, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct SessionTransportConfig { + pub role: SessionRole, + pub conn_id: [u8; 16], + pub session_key: [u8; 32], + pub peer_addr: SocketAddr, + pub initial_send_seq: u64, + pub initial_ack: u64, + pub stream: TransportConfig, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum SessionEvent { + Stream(TransportEvent), + Ping, + Pong, + Ignored, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OutgoingStreamPacket { + pub kind: PacketKind, + pub body: Vec, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct IncomingStreamData { + pub stream_id: u64, + pub chunks: Vec>, + pub received_offset: u64, + pub consumed: usize, + pub window_adjust: OutgoingStreamPacket, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct IncomingStreamOpen { + pub stream_id: u64, + pub target_host: String, + pub target_port: u16, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct AcceptedService { + pub name: String, + pub stream_id: u64, +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TransportEvent { + Open(IncomingStreamOpen), + OpenOk { + stream_id: u64, + flushed: Vec, + }, + OpenReject { + stream_id: u64, + reason: String, + }, + Data(IncomingStreamData), + WindowAdjust { + stream_id: u64, + flushed: Vec, + }, + Close { + stream_id: u64, + }, +} + +#[derive(Debug, Clone)] +struct PendingStreamChunk { + offset: u64, + bytes: Vec, + last_sent: Instant, + attempts: u32, +} + +#[derive(Debug, Clone)] +struct PendingStreamOpen { + target_host: String, + target_port: u16, + last_sent: Instant, + attempts: u32, +} + +#[derive(Debug, Clone)] +pub struct StreamMux { + config: TransportConfig, + opened_streams: HashSet, + pending_opens: HashMap, + send_credit: HashMap, + pending_data: HashMap>>, + sent_data: HashMap>, + next_send_offset: HashMap, + next_recv_offset: HashMap, + recv_pending: HashMap>>, +} + +#[derive(Debug, Clone, Default)] +pub struct ServiceRegistry { + services: HashSet, +} + +impl ServiceRegistry { + pub fn new() -> Self { + Self::default() + } + + pub fn register(&mut self, name: impl Into) -> Result<()> { + let name = name.into(); + service_target(&name)?; + self.services.insert(name); + Ok(()) + } + + pub fn is_registered(&self, name: &str) -> bool { + self.services.contains(name) + } + + pub fn match_open(&self, open: &IncomingStreamOpen) -> Result { + let Some(name) = service_name_from_target(&open.target_host) else { + bail!("stream target {} is not a Dosh service", open.target_host); + }; + if !self.services.contains(name) { + bail!("service {name:?} is not registered"); + } + if open.target_port != 0 { + bail!( + "service stream {} used non-zero target port", + open.stream_id + ); + } + Ok(AcceptedService { + name: name.to_string(), + stream_id: open.stream_id, + }) + } +} + +impl StreamMux { + pub fn new(config: TransportConfig) -> Self { + Self { + config, + opened_streams: HashSet::new(), + pending_opens: HashMap::new(), + send_credit: HashMap::new(), + pending_data: HashMap::new(), + sent_data: HashMap::new(), + next_send_offset: HashMap::new(), + next_recv_offset: HashMap::new(), + recv_pending: HashMap::new(), + } + } + + pub fn open_stream( + &mut self, + stream_id: u64, + target_host: impl Into, + target_port: u16, + ) -> Result { + let target_host = target_host.into(); + if self.opened_streams.contains(&stream_id) || self.pending_opens.contains_key(&stream_id) { + bail!("stream {stream_id} already exists"); + } + self.send_credit + .insert(stream_id, self.config.initial_window); + self.next_send_offset.entry(stream_id).or_insert(0); + self.next_recv_offset.entry(stream_id).or_insert(0); + self.pending_opens.insert( + stream_id, + PendingStreamOpen { + target_host: target_host.clone(), + target_port, + last_sent: Instant::now(), + attempts: 1, + }, + ); + encode_packet( + PacketKind::StreamOpen, + &StreamOpen { + stream_id, + target_host, + target_port, + }, + ) + } + + pub fn accept_open(&mut self, open: StreamOpen) -> Result { + self.opened_streams.insert(open.stream_id); + self.send_credit + .entry(open.stream_id) + .or_insert(self.config.initial_window); + self.next_send_offset.entry(open.stream_id).or_insert(0); + self.next_recv_offset.entry(open.stream_id).or_insert(0); + Ok(IncomingStreamOpen { + stream_id: open.stream_id, + target_host: open.target_host, + target_port: open.target_port, + }) + } + + pub fn open_ok(stream_id: u64) -> Result { + encode_packet(PacketKind::StreamOpenOk, &StreamOpenOk { stream_id }) + } + + pub fn open_reject(stream_id: u64, reason: impl Into) -> Result { + encode_packet( + PacketKind::StreamOpenReject, + &StreamOpenReject { + stream_id, + reason: reason.into(), + }, + ) + } + + pub fn handle_open_ok(&mut self, ok: StreamOpenOk) -> Result> { + self.pending_opens.remove(&ok.stream_id); + self.opened_streams.insert(ok.stream_id); + self.send_credit + .entry(ok.stream_id) + .or_insert(self.config.initial_window); + self.next_send_offset.entry(ok.stream_id).or_insert(0); + self.next_recv_offset.entry(ok.stream_id).or_insert(0); + self.flush_pending_data(ok.stream_id) + } + + pub fn handle_open_reject(&mut self, reject: StreamOpenReject) { + self.cleanup_stream(reject.stream_id); + } + + pub fn send_data( + &mut self, + stream_id: u64, + bytes: impl Into>, + ) -> Result> { + let bytes = bytes.into(); + if bytes.is_empty() { + return Ok(Vec::new()); + } + if !self.opened_streams.contains(&stream_id) + || self.send_credit.get(&stream_id).copied().unwrap_or(0) < bytes.len() + || self + .pending_data + .get(&stream_id) + .is_some_and(|pending| !pending.is_empty()) + { + self.pending_data + .entry(stream_id) + .or_default() + .push_back(bytes); + return Ok(Vec::new()); + } + Ok(vec![self.send_data_now(stream_id, bytes)?]) + } + + pub fn handle_data(&mut self, data: StreamData) -> Result { + let stream_id = data.stream_id; + let (chunks, consumed, received_offset) = self.accept_data(data); + let window_adjust = encode_packet( + PacketKind::StreamWindowAdjust, + &StreamWindowAdjust { + stream_id, + received_offset, + bytes: consumed.min(u32::MAX as usize) as u32, + }, + )?; + Ok(IncomingStreamData { + stream_id, + chunks, + consumed, + received_offset, + window_adjust, + }) + } + + pub fn handle_window_adjust( + &mut self, + adjust: StreamWindowAdjust, + ) -> Result> { + self.ack_data(adjust.stream_id, adjust.received_offset); + self.add_credit(adjust.stream_id, adjust.bytes as usize); + self.flush_pending_data(adjust.stream_id) + } + + pub fn close_stream(&mut self, stream_id: u64) -> Result { + self.cleanup_stream(stream_id); + encode_packet(PacketKind::StreamClose, &StreamClose { stream_id }) + } + + pub fn handle_close(&mut self, close: StreamClose) { + self.cleanup_stream(close.stream_id); + } + + pub fn handle_packet(&mut self, kind: PacketKind, body: &[u8]) -> Result { + match kind { + PacketKind::StreamOpen => { + let open: StreamOpen = protocol::from_body(body)?; + self.accept_open(open).map(TransportEvent::Open) + } + PacketKind::StreamOpenOk => { + let ok: StreamOpenOk = protocol::from_body(body)?; + let stream_id = ok.stream_id; + let flushed = self.handle_open_ok(ok)?; + Ok(TransportEvent::OpenOk { stream_id, flushed }) + } + PacketKind::StreamOpenReject => { + let reject: StreamOpenReject = protocol::from_body(body)?; + self.handle_open_reject(reject.clone()); + Ok(TransportEvent::OpenReject { + stream_id: reject.stream_id, + reason: reject.reason, + }) + } + PacketKind::StreamData => { + let data: StreamData = protocol::from_body(body)?; + self.handle_data(data).map(TransportEvent::Data) + } + PacketKind::StreamWindowAdjust => { + let adjust: StreamWindowAdjust = protocol::from_body(body)?; + let stream_id = adjust.stream_id; + let flushed = self.handle_window_adjust(adjust)?; + Ok(TransportEvent::WindowAdjust { stream_id, flushed }) + } + PacketKind::StreamClose => { + let close: StreamClose = protocol::from_body(body)?; + let stream_id = close.stream_id; + self.handle_close(close); + Ok(TransportEvent::Close { stream_id }) + } + _ => bail!("packet kind {kind:?} is not a stream transport packet"), + } + } + + pub fn tick(&mut self) -> Result> { + let now = Instant::now(); + let mut out = Vec::new(); + let mut retransmit_opens = Vec::new(); + for (stream_id, pending) in self.pending_opens.iter_mut() { + if now.duration_since(pending.last_sent) < self.config.retransmit_after { + continue; + } + pending.last_sent = now; + pending.attempts = pending.attempts.saturating_add(1); + retransmit_opens.push((*stream_id, pending.target_host.clone(), pending.target_port)); + } + for (stream_id, target_host, target_port) in retransmit_opens { + out.push(encode_packet( + PacketKind::StreamOpen, + &StreamOpen { + stream_id, + target_host, + target_port, + }, + )?); + } + + let mut retransmit_data = Vec::new(); + for (stream_id, chunks) in self.sent_data.iter_mut() { + for chunk in chunks.values_mut() { + if now.duration_since(chunk.last_sent) < self.config.retransmit_after { + continue; + } + chunk.last_sent = now; + chunk.attempts = chunk.attempts.saturating_add(1); + retransmit_data.push((*stream_id, chunk.offset, chunk.bytes.clone())); + } + } + for (stream_id, offset, bytes) in retransmit_data { + out.push(encode_packet( + PacketKind::StreamData, + &StreamData { + stream_id, + offset, + bytes, + }, + )?); + } + Ok(out) + } + + pub fn is_open(&self, stream_id: u64) -> bool { + self.opened_streams.contains(&stream_id) + } + + pub fn pending_bytes(&self, stream_id: u64) -> usize { + self.pending_data + .get(&stream_id) + .map(|queue| queue.iter().map(Vec::len).sum()) + .unwrap_or(0) + } + + fn send_data_now(&mut self, stream_id: u64, bytes: Vec) -> Result { + let credit = self + .send_credit + .get_mut(&stream_id) + .ok_or_else(|| anyhow!("stream {stream_id} has no send window"))?; + if *credit < bytes.len() { + bail!("stream {stream_id} send window exhausted"); + } + *credit -= bytes.len(); + let offset = *self.next_send_offset.entry(stream_id).or_default(); + self.next_send_offset + .insert(stream_id, offset.saturating_add(bytes.len() as u64)); + self.sent_data.entry(stream_id).or_default().insert( + offset, + PendingStreamChunk { + offset, + bytes: bytes.clone(), + last_sent: Instant::now(), + attempts: 1, + }, + ); + encode_packet( + PacketKind::StreamData, + &StreamData { + stream_id, + offset, + bytes, + }, + ) + } + + fn flush_pending_data(&mut self, stream_id: u64) -> Result> { + let mut out = Vec::new(); + while let Some(front_len) = self + .pending_data + .get(&stream_id) + .and_then(|pending| pending.front().map(Vec::len)) + { + if self.send_credit.get(&stream_id).copied().unwrap_or(0) < front_len { + break; + } + let bytes = self + .pending_data + .get_mut(&stream_id) + .and_then(VecDeque::pop_front) + .expect("pending front exists"); + if self + .pending_data + .get(&stream_id) + .is_some_and(VecDeque::is_empty) + { + self.pending_data.remove(&stream_id); + } + out.push(self.send_data_now(stream_id, bytes)?); + } + Ok(out) + } + + fn accept_data(&mut self, data: StreamData) -> (Vec>, usize, u64) { + let stream_id = data.stream_id; + let expected = self.next_recv_offset.entry(stream_id).or_insert(0); + if data.offset < *expected { + return (Vec::new(), 0, *expected); + } + if data.offset > *expected { + self.recv_pending + .entry(stream_id) + .or_default() + .entry(data.offset) + .or_insert(data.bytes); + return (Vec::new(), 0, *expected); + } + + let mut chunks = vec![data.bytes]; + let mut consumed = chunks[0].len(); + *expected = expected.saturating_add(consumed as u64); + while let Some(bytes) = self + .recv_pending + .entry(stream_id) + .or_default() + .remove(expected) + { + consumed = consumed.saturating_add(bytes.len()); + *expected = expected.saturating_add(bytes.len() as u64); + chunks.push(bytes); + } + if self + .recv_pending + .get(&stream_id) + .is_some_and(BTreeMap::is_empty) + { + self.recv_pending.remove(&stream_id); + } + (chunks, consumed, *expected) + } + + fn ack_data(&mut self, stream_id: u64, received_offset: u64) { + let Some(sent) = self.sent_data.get_mut(&stream_id) else { + return; + }; + let acked_offsets = sent + .iter() + .filter_map(|(offset, chunk)| { + let end = chunk.offset.saturating_add(chunk.bytes.len() as u64); + (end <= received_offset).then_some(*offset) + }) + .collect::>(); + for offset in acked_offsets { + sent.remove(&offset); + } + if sent.is_empty() { + self.sent_data.remove(&stream_id); + } + } + + fn add_credit(&mut self, stream_id: u64, bytes: usize) { + let credit = self.send_credit.entry(stream_id).or_default(); + *credit = credit.saturating_add(bytes).min(self.config.initial_window); + } + + fn cleanup_stream(&mut self, stream_id: u64) { + self.opened_streams.remove(&stream_id); + self.pending_opens.remove(&stream_id); + self.send_credit.remove(&stream_id); + self.pending_data.remove(&stream_id); + self.sent_data.remove(&stream_id); + self.next_send_offset.remove(&stream_id); + self.next_recv_offset.remove(&stream_id); + self.recv_pending.remove(&stream_id); + } +} + +pub struct DoshTransport { + socket: Arc, + role: SessionRole, + conn_id: [u8; 16], + session_key: [u8; 32], + peer_addr: SocketAddr, + send_seq: u64, + ack_seq: u64, + replay: ReplayWindow, + mux: StreamMux, + next_stream_id: u64, + last_contact: Instant, + last_keepalive: Instant, +} + +impl DoshTransport { + pub fn new(socket: Arc, config: SessionTransportConfig) -> Self { + let now = Instant::now(); + Self { + socket, + role: config.role, + conn_id: config.conn_id, + session_key: config.session_key, + peer_addr: config.peer_addr, + send_seq: config.initial_send_seq, + ack_seq: config.initial_ack, + replay: ReplayWindow::default(), + mux: StreamMux::new(config.stream), + next_stream_id: 1, + last_contact: now, + last_keepalive: now, + } + } + + pub fn new_owned(socket: UdpSocket, config: SessionTransportConfig) -> Self { + Self::new(Arc::new(socket), config) + } + + pub fn local_addr(&self) -> Result { + Ok(self.socket.local_addr()?) + } + + pub fn conn_id(&self) -> [u8; 16] { + self.conn_id + } + + pub fn mux(&self) -> &StreamMux { + &self.mux + } + + pub fn mux_mut(&mut self) -> &mut StreamMux { + &mut self.mux + } + + pub fn peer_addr(&self) -> SocketAddr { + self.peer_addr + } + + pub fn set_peer_addr(&mut self, peer_addr: SocketAddr) { + self.peer_addr = peer_addr; + } + + pub fn last_contact(&self) -> Instant { + self.last_contact + } + + pub fn stale_for(&self) -> Duration { + self.last_contact.elapsed() + } + + pub fn allocate_stream_id(&mut self) -> u64 { + let id = self.next_stream_id; + self.next_stream_id = self.next_stream_id.saturating_add(1).max(1); + id + } + + pub async fn open_service(&mut self, name: &str) -> Result { + let stream_id = self.allocate_stream_id(); + let target = service_target(name)?; + let packet = self.mux.open_stream(stream_id, target, 0)?; + self.send_outgoing(packet).await?; + Ok(stream_id) + } + + pub async fn open_target( + &mut self, + target_host: impl Into, + target_port: u16, + ) -> Result { + let stream_id = self.allocate_stream_id(); + let packet = self + .mux + .open_stream(stream_id, target_host.into(), target_port)?; + self.send_outgoing(packet).await?; + Ok(stream_id) + } + + pub async fn accept_stream(&mut self, stream_id: u64) -> Result<()> { + self.send_outgoing(StreamMux::open_ok(stream_id)?).await + } + + pub async fn reject_stream(&mut self, stream_id: u64, reason: impl Into) -> Result<()> { + self.send_outgoing(StreamMux::open_reject(stream_id, reason)?) + .await + } + + pub async fn send(&mut self, stream_id: u64, bytes: impl Into>) -> Result<()> { + for packet in self.mux.send_data(stream_id, bytes)? { + self.send_outgoing(packet).await?; + } + Ok(()) + } + + pub async fn close(&mut self, stream_id: u64) -> Result<()> { + let packet = self.mux.close_stream(stream_id)?; + self.send_outgoing(packet).await + } + + pub async fn maintenance(&mut self) -> Result<()> { + for packet in self.mux.tick()? { + self.send_outgoing(packet).await?; + } + if self.last_keepalive.elapsed() >= self.mux.config.keepalive_after { + self.send_kind(PacketKind::Ping, b"").await?; + self.last_keepalive = Instant::now(); + } + Ok(()) + } + + pub async fn recv(&mut self) -> Result { + let mut buf = vec![0u8; 65535]; + loop { + let (n, peer) = self.socket.recv_from(&mut buf).await?; + let packet = match protocol::decode(&buf[..n]) { + Ok(packet) => packet, + Err(_) => continue, + }; + if packet.header.conn_id != self.conn_id { + continue; + } + if !self.replay.accept(packet.header.seq) { + continue; + } + let plain = + protocol::decrypt_body(&packet, &self.session_key, self.role.recv_direction())?; + self.peer_addr = peer; + self.ack_seq = self.ack_seq.max(packet.header.seq); + self.last_contact = Instant::now(); + return self + .handle_decrypted_packet(packet.header.kind, &plain) + .await; + } + } + + pub async fn handle_datagram( + &mut self, + datagram: &[u8], + peer: SocketAddr, + ) -> Result { + let packet = protocol::decode(datagram)?; + if packet.header.conn_id != self.conn_id { + return Ok(SessionEvent::Ignored); + } + if !self.replay.accept(packet.header.seq) { + return Ok(SessionEvent::Ignored); + } + let plain = protocol::decrypt_body(&packet, &self.session_key, self.role.recv_direction())?; + self.peer_addr = peer; + self.ack_seq = self.ack_seq.max(packet.header.seq); + self.last_contact = Instant::now(); + self.handle_decrypted_packet(packet.header.kind, &plain) + .await + } + + async fn handle_decrypted_packet( + &mut self, + kind: PacketKind, + plain: &[u8], + ) -> Result { + match kind { + PacketKind::Ping => { + self.send_kind(PacketKind::Pong, b"").await?; + Ok(SessionEvent::Ping) + } + PacketKind::Pong => Ok(SessionEvent::Pong), + PacketKind::StreamOpen + | PacketKind::StreamOpenOk + | PacketKind::StreamOpenReject + | PacketKind::StreamData + | PacketKind::StreamWindowAdjust + | PacketKind::StreamClose => { + let event = self.mux.handle_packet(kind, plain)?; + self.send_followups(&event).await?; + Ok(SessionEvent::Stream(event)) + } + _ => Ok(SessionEvent::Ignored), + } + } + + async fn send_followups(&mut self, event: &TransportEvent) -> Result<()> { + match event { + TransportEvent::OpenOk { flushed, .. } + | TransportEvent::WindowAdjust { flushed, .. } => { + for packet in flushed.clone() { + self.send_outgoing(packet).await?; + } + } + TransportEvent::Data(data) => { + self.send_outgoing(data.window_adjust.clone()).await?; + } + TransportEvent::Open(_) + | TransportEvent::OpenReject { .. } + | TransportEvent::Close { .. } => {} + } + Ok(()) + } + + async fn send_outgoing(&mut self, packet: OutgoingStreamPacket) -> Result<()> { + self.send_kind(packet.kind, &packet.body).await + } + + async fn send_kind(&mut self, kind: PacketKind, body: &[u8]) -> Result<()> { + let encoded = self.encode_kind(kind, body)?; + self.socket.send_to(&encoded, self.peer_addr).await?; + Ok(()) + } + + fn encode_kind(&mut self, kind: PacketKind, body: &[u8]) -> Result> { + let seq = self.send_seq; + self.send_seq = self.send_seq.saturating_add(1).max(1); + protocol::encode_encrypted( + kind, + self.conn_id, + seq, + self.ack_seq, + &self.session_key, + self.role.send_direction(), + body, + ) + } +} + +pub fn service_target(name: &str) -> Result { + if name.is_empty() + || !name + .bytes() + .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_')) + { + bail!("invalid Dosh service name {name:?}"); + } + Ok(format!("{SERVICE_TARGET_PREFIX}{name}")) +} + +pub fn is_service_target(target: &str) -> bool { + target.starts_with(SERVICE_TARGET_PREFIX) +} + +pub fn service_name_from_target(target: &str) -> Option<&str> { + target.strip_prefix(SERVICE_TARGET_PREFIX) +} + +fn encode_packet(kind: PacketKind, value: &T) -> Result { + Ok(OutgoingStreamPacket { + kind, + body: protocol::to_body(value)?, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn decode(packet: &OutgoingStreamPacket) -> T { + protocol::from_body(&packet.body).unwrap() + } + + #[test] + fn service_targets_are_reserved_and_validated() { + assert_eq!(service_target("chat_api").unwrap(), "@dosh-chat_api"); + assert!(is_service_target("@dosh-file")); + assert_eq!(service_name_from_target("@dosh-chat"), Some("chat")); + assert!(service_target("../bad").is_err()); + } + + #[test] + fn service_registry_matches_registered_service_opens() { + let mut registry = ServiceRegistry::new(); + registry.register("chat").unwrap(); + let accepted = registry + .match_open(&IncomingStreamOpen { + stream_id: 4, + target_host: "@dosh-chat".to_string(), + target_port: 0, + }) + .unwrap(); + assert_eq!( + accepted, + AcceptedService { + name: "chat".to_string(), + stream_id: 4, + } + ); + assert!( + registry + .match_open(&IncomingStreamOpen { + stream_id: 5, + target_host: "@dosh-other".to_string(), + target_port: 0, + }) + .is_err() + ); + } + + #[test] + fn opens_stream_and_flushes_queued_data_after_open_ok() { + let mut mux = StreamMux::new(TransportConfig::default()); + let open = mux + .open_stream(7, service_target("chat").unwrap(), 0) + .unwrap(); + assert_eq!(open.kind, PacketKind::StreamOpen); + assert!(mux.send_data(7, b"hello".to_vec()).unwrap().is_empty()); + assert_eq!(mux.pending_bytes(7), 5); + + let flushed = mux.handle_open_ok(StreamOpenOk { stream_id: 7 }).unwrap(); + assert_eq!(flushed.len(), 1); + assert_eq!(flushed[0].kind, PacketKind::StreamData); + let data: StreamData = decode(&flushed[0]); + assert_eq!(data.offset, 0); + assert_eq!(data.bytes, b"hello"); + } + + #[test] + fn receives_out_of_order_data_in_order_and_acks_offsets() { + let mut mux = StreamMux::new(TransportConfig::default()); + mux.accept_open(StreamOpen { + stream_id: 3, + target_host: "@dosh-test".to_string(), + target_port: 0, + }) + .unwrap(); + let first = mux + .handle_data(StreamData { + stream_id: 3, + offset: 5, + bytes: b"world".to_vec(), + }) + .unwrap(); + assert!(first.chunks.is_empty()); + assert_eq!(first.received_offset, 0); + + let second = mux + .handle_data(StreamData { + stream_id: 3, + offset: 0, + bytes: b"hello".to_vec(), + }) + .unwrap(); + assert_eq!(second.chunks, vec![b"hello".to_vec(), b"world".to_vec()]); + assert_eq!(second.consumed, 10); + assert_eq!(second.received_offset, 10); + } + + #[test] + fn window_adjust_acks_sent_data_and_releases_credit() { + let mut mux = StreamMux::new(TransportConfig { + initial_window: 5, + retransmit_after: Duration::from_secs(1), + ..TransportConfig::default() + }); + mux.open_stream(1, "@dosh-test", 0).unwrap(); + mux.handle_open_ok(StreamOpenOk { stream_id: 1 }).unwrap(); + let sent = mux.send_data(1, b"hello".to_vec()).unwrap(); + assert_eq!(sent.len(), 1); + assert!(mux.send_data(1, b"!".to_vec()).unwrap().is_empty()); + assert_eq!(mux.pending_bytes(1), 1); + let flushed = mux + .handle_window_adjust(StreamWindowAdjust { + stream_id: 1, + received_offset: 5, + bytes: 5, + }) + .unwrap(); + assert_eq!(flushed.len(), 1); + let data: StreamData = decode(&flushed[0]); + assert_eq!(data.offset, 5); + assert_eq!(data.bytes, b"!"); + } + + #[test] + fn handle_packet_decodes_and_dispatches_stream_packets() { + let mut mux = StreamMux::new(TransportConfig::default()); + let body = protocol::to_body(&StreamOpen { + stream_id: 9, + target_host: "@dosh-app".to_string(), + target_port: 0, + }) + .unwrap(); + let event = mux.handle_packet(PacketKind::StreamOpen, &body).unwrap(); + assert_eq!( + event, + TransportEvent::Open(IncomingStreamOpen { + stream_id: 9, + target_host: "@dosh-app".to_string(), + target_port: 0, + }) + ); + } + + #[tokio::test] + async fn session_transport_opens_service_and_moves_data_over_encrypted_udp() { + let client_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let client_addr = client_socket.local_addr().unwrap(); + let server_addr = server_socket.local_addr().unwrap(); + let key = [42u8; 32]; + let conn_id = [7u8; 16]; + let mut client = DoshTransport::new_owned( + client_socket, + SessionTransportConfig { + role: SessionRole::Client, + conn_id, + session_key: key, + peer_addr: server_addr, + initial_send_seq: 1, + initial_ack: 0, + stream: TransportConfig::default(), + }, + ); + let mut server = DoshTransport::new_owned( + server_socket, + SessionTransportConfig { + role: SessionRole::Server, + conn_id, + session_key: key, + peer_addr: client_addr, + initial_send_seq: 1, + initial_ack: 0, + stream: TransportConfig::default(), + }, + ); + + let stream_id = client.open_service("app").await.unwrap(); + let event = server.recv().await.unwrap(); + assert_eq!(server.peer_addr(), client_addr); + match event { + SessionEvent::Stream(TransportEvent::Open(open)) => { + assert_eq!(open.stream_id, stream_id); + assert_eq!(open.target_host, "@dosh-app"); + server.accept_stream(open.stream_id).await.unwrap(); + } + other => panic!("unexpected event {other:?}"), + } + + assert!(matches!( + client.recv().await.unwrap(), + SessionEvent::Stream(TransportEvent::OpenOk { .. }) + )); + client.send(stream_id, b"hello".to_vec()).await.unwrap(); + let event = server.recv().await.unwrap(); + match event { + SessionEvent::Stream(TransportEvent::Data(data)) => { + assert_eq!(data.stream_id, stream_id); + assert_eq!(data.chunks, vec![b"hello".to_vec()]); + } + other => panic!("unexpected event {other:?}"), + } + } + + #[tokio::test] + async fn session_transport_roams_to_authenticated_new_peer_addr() { + let server_socket = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let original_client = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let roaming_client = UdpSocket::bind("127.0.0.1:0").await.unwrap(); + let original_addr = original_client.local_addr().unwrap(); + let roaming_addr = roaming_client.local_addr().unwrap(); + let server_addr = server_socket.local_addr().unwrap(); + let key = [9u8; 32]; + let conn_id = [3u8; 16]; + let mut server = DoshTransport::new_owned( + server_socket, + SessionTransportConfig { + role: SessionRole::Server, + conn_id, + session_key: key, + peer_addr: original_addr, + initial_send_seq: 1, + initial_ack: 0, + stream: TransportConfig::default(), + }, + ); + let ping = protocol::encode_encrypted( + PacketKind::Ping, + conn_id, + 1, + 0, + &key, + CLIENT_TO_SERVER, + b"", + ) + .unwrap(); + roaming_client.send_to(&ping, server_addr).await.unwrap(); + assert!(matches!(server.recv().await.unwrap(), SessionEvent::Ping)); + assert_eq!(server.peer_addr(), roaming_addr); + } +} diff --git a/tests/integration_smoke.rs b/tests/integration_smoke.rs index bd19f0e..18e118f 100644 --- a/tests/integration_smoke.rs +++ b/tests/integration_smoke.rs @@ -760,7 +760,9 @@ fn native_doctor_checks_auth_without_opening_terminal() { assert!(output.status.success(), "stdout={stdout}\nstderr={stderr}"); assert!(stdout.contains("[ok] native udp"), "stdout={stdout}"); assert!(stdout.contains("[ok] native auth"), "stdout={stdout}"); - assert!(stdout.contains("[ok] forwarding policy"), "stdout={stdout}"); + assert!(stdout.contains("[ok] server policy"), "stdout={stdout}"); + assert!(stdout.contains("file_transfer=true"), "stdout={stdout}"); + assert!(stdout.contains("exec_command=true"), "stdout={stdout}"); } #[test] @@ -868,6 +870,177 @@ fn native_local_forward_background_smoke() { assert_eq!(&buf[..n], b"dosh-background-forward-ping"); } +#[test] +fn native_file_copy_recursive_round_trip() { + let dir = tempfile::tempdir().unwrap(); + let port = free_udp_port(); + let config = write_server_config(&dir, port); + write_native_client_auth(&dir, &config); + let mut server = start_server(&dir, &config); + let src = dir.path().join("srcdir"); + fs::create_dir_all(src.join("nested")).unwrap(); + fs::write(src.join("root.txt"), b"root file\n").unwrap(); + fs::write(src.join("nested/child.txt"), b"child file\n").unwrap(); + + let client_bin = env!("CARGO_BIN_EXE_dosh-client"); + let upload = Command::new(client_bin) + .arg("--dosh-host") + .arg("127.0.0.1") + .arg("--dosh-port") + .arg(port.to_string()) + .arg("cp") + .arg("-r") + .arg(&src) + .arg("local:remote-copy") + .env("HOME", dir.path()) + .output() + .unwrap(); + assert!( + upload.status.success(), + "upload failed: stdout={} stderr={}", + String::from_utf8_lossy(&upload.stdout), + String::from_utf8_lossy(&upload.stderr) + ); + + let remote_copy = Command::new(client_bin) + .arg("--dosh-host") + .arg("127.0.0.1") + .arg("--dosh-port") + .arg(port.to_string()) + .arg("cp") + .arg("-r") + .arg("local:remote-copy") + .arg("local:remote-copy-2") + .env("HOME", dir.path()) + .output() + .unwrap(); + assert!( + remote_copy.status.success(), + "remote copy failed: stdout={} stderr={}", + String::from_utf8_lossy(&remote_copy.stdout), + String::from_utf8_lossy(&remote_copy.stderr) + ); + + let list = Command::new(client_bin) + .arg("--dosh-host") + .arg("127.0.0.1") + .arg("--dosh-port") + .arg(port.to_string()) + .arg("ls") + .arg("local:remote-copy-2") + .env("HOME", dir.path()) + .output() + .unwrap(); + assert!( + list.status.success(), + "ls failed: stdout={} stderr={}", + String::from_utf8_lossy(&list.stdout), + String::from_utf8_lossy(&list.stderr) + ); + let list_stdout = String::from_utf8_lossy(&list.stdout); + assert!(list_stdout.contains("root.txt"), "stdout={list_stdout}"); + assert!(list_stdout.contains("nested"), "stdout={list_stdout}"); + + let cat = Command::new(client_bin) + .arg("--dosh-host") + .arg("127.0.0.1") + .arg("--dosh-port") + .arg(port.to_string()) + .arg("cat") + .arg("local:remote-copy-2/root.txt") + .env("HOME", dir.path()) + .output() + .unwrap(); + assert!( + cat.status.success(), + "cat failed: stdout={} stderr={}", + String::from_utf8_lossy(&cat.stdout), + String::from_utf8_lossy(&cat.stderr) + ); + assert_eq!(String::from_utf8_lossy(&cat.stdout), "root file\n"); + + let downloaded = dir.path().join("downloaded"); + let download = Command::new(client_bin) + .arg("--dosh-host") + .arg("127.0.0.1") + .arg("--dosh-port") + .arg(port.to_string()) + .arg("cp") + .arg("-r") + .arg("local:remote-copy-2") + .arg(&downloaded) + .env("HOME", dir.path()) + .output() + .unwrap(); + assert!( + download.status.success(), + "download failed: stdout={} stderr={}", + String::from_utf8_lossy(&download.stdout), + String::from_utf8_lossy(&download.stderr) + ); + assert_eq!( + fs::read_to_string(downloaded.join("root.txt")).unwrap(), + "root file\n" + ); + assert_eq!( + fs::read_to_string(downloaded.join("nested/child.txt")).unwrap(), + "child file\n" + ); + + let remove = Command::new(client_bin) + .arg("--dosh-host") + .arg("127.0.0.1") + .arg("--dosh-port") + .arg(port.to_string()) + .arg("rm") + .arg("-r") + .arg("local:remote-copy") + .env("HOME", dir.path()) + .output() + .unwrap(); + assert!( + remove.status.success(), + "rm failed: stdout={} stderr={}", + String::from_utf8_lossy(&remove.stdout), + String::from_utf8_lossy(&remove.stderr) + ); + + let _ = server.kill(); + let _ = server.wait(); +} + +#[test] +fn native_exec_command_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 mut server = start_server(&dir, &config); + let client_bin = env!("CARGO_BIN_EXE_dosh-client"); + let output = Command::new(client_bin) + .arg("--dosh-host") + .arg("127.0.0.1") + .arg("--dosh-port") + .arg(port.to_string()) + .arg("exec") + .arg("local") + .arg("printf out; printf err >&2") + .env("HOME", dir.path()) + .output() + .unwrap(); + let _ = server.kill(); + let _ = server.wait(); + + assert!( + output.status.success(), + "exec failed: stdout={} stderr={}", + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr) + ); + assert_eq!(String::from_utf8_lossy(&output.stdout), "out"); + assert_eq!(String::from_utf8_lossy(&output.stderr), "err"); +} + #[test] fn native_local_forward_bulk_load_does_not_delay_interactive_terminal() { use portable_pty::{CommandBuilder, NativePtySystem, PtySize, PtySystem};