Palav/dosh
1
0
Fork 0
Code Issues Pull Requests Actions 46 Packages Projects Releases 51 Wiki Activity
Files
e11ba1a65245a3ebbbc90ff5cde65c3264ca7b6e
dosh/src/bin/dosh-client.rs
T

9330 lines
332 KiB
Rust
Raw Blame History

use anyhow::{Context, Result, anyhow, bail};
use base64::Engine;
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
use clap::{ArgAction, Parser};
use crossterm::terminal::{disable_raw_mode, enable_raw_mode, size};
use dosh::auth::{
BootstrapResponse, build_bootstrap, decode_bootstrap, load_or_create_server_secret,
};
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,
load_native_identity_with_passphrase, parse_host_public_key_line, remove_trusted_host,
sign_user_auth_with_private_key, supported_user_key_algorithms, trust_host, verify_known_host,
verify_server_hello,
};
use dosh::protocol::{
self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
NativeAuthCheckOkBody, NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody,
NativeUserAuthBody, PacketKind, Rekey, Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose,
StreamData, StreamOpen, StreamOpenOk, StreamOpenReject, StreamWindowAdjust, TicketAttachBody,
TicketAttachEnvelope, TicketAttachOkEnvelope,
};
use dosh::ssh_agent;
use dosh::transport::{
DoshTransport, SessionEvent, SessionRole, SessionTransportConfig, TransportConfig,
TransportEvent,
};
use dosh::udp::is_transient_udp_send_error;
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, Seek, Write};
use std::net::{
Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener as StdTcpListener, TcpStream as StdTcpStream,
ToSocketAddrs,
};
use std::path::{Path, PathBuf};
use std::process::{Child, Command, Stdio};
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{Duration, Instant};
use tokio::io::{AsyncReadExt, AsyncWriteExt};
#[cfg(unix)]
use tokio::net::UnixStream;
use tokio::net::{TcpListener, TcpStream, UdpSocket};
use tokio::sync::mpsc;
const STREAM_INITIAL_WINDOW: usize = 1024 * 1024;
const STREAM_RETIRED_TOMBSTONES: usize = 16 * 1024;
const MAX_PENDING_USER_INPUT_BYTES: usize = 1024 * 1024;
const STARTUP_INPUT_HOLD: Duration = Duration::from_millis(750);
const POST_SUBMIT_ALL_INPUT_HOLD: Duration = Duration::from_millis(120);
const STALE_TERMINAL_INPUT_AFTER: Duration = Duration::from_secs(2);
const POST_RECONNECT_STALE_INPUT_GRACE: Duration = Duration::from_secs(2);
const FOCUS_REPAINT_COOLDOWN: Duration = Duration::from_secs(1);
const ALT_SCREEN_IDLE_REPAINT_AFTER: Duration = Duration::from_secs(15);
/// Sentinel `target_host` the server uses on a server-initiated `StreamOpen` that
/// represents an SSH-agent connection (rather than a TCP target). The client
/// splices these into its local `SSH_AUTH_SOCK` instead of dialing TCP. The
/// `:0` port and this reserved name are never valid TCP forward targets.
const AGENT_STREAM_SENTINEL: &str = "@dosh-agent";
/// Current terminal size, with a sane fallback.
///
/// `crossterm::size()` returns `Err` when stdout is not a TTY (piped), but it
/// can also return `Ok((0, 0))` for a real PTY that has not been sized yet
/// (e.g. launched under `script` with no controlling window). Sending `0` to
/// the server is meaningless and used to crash older servers, so treat any
/// zero dimension the same as a missing size and fall back to 80x24.
fn terminal_size() -> (u16, u16) {
match size() {
Ok((cols, rows)) if cols > 0 && rows > 0 => (cols, rows),
_ => (80, 24),
}
}
#[derive(Debug, Clone, Parser)]
#[command(
name = "dosh-client",
version = dosh::build_info::VERSION,
long_version = dosh::build_info::LONG_VERSION
)]
struct Args {
#[arg()]
server: Option<String>,
#[arg(trailing_var_arg = true, allow_hyphen_values = true)]
command: Vec<String>,
#[arg(long)]
session: Option<String>,
#[arg(long)]
new: bool,
#[arg(long)]
view_only: bool,
#[arg(long)]
local_auth: bool,
#[arg(long)]
auth: Option<String>,
#[arg(long)]
no_cache: bool,
#[arg(long)]
remove: bool,
#[arg(long)]
replace: bool,
#[arg(long)]
attach_only: bool,
#[arg(short = 'L', long = "local-forward")]
local_forward: Vec<String>,
#[arg(short = 'R', long = "remote-forward")]
remote_forward: Vec<String>,
#[arg(short = 'D', long = "dynamic-forward")]
dynamic_forward: Vec<String>,
/// Forward the local ssh-agent (SSH_AUTH_SOCK) to the remote session.
/// SECURITY: opt-in only; the server must also set allow_agent_forwarding.
#[arg(short = 'A', long = "forward-agent")]
forward_agent: bool,
#[arg(short = 'N', long)]
forward_only: bool,
#[arg(short = 'f', long)]
background: bool,
#[arg(long)]
predict: bool,
#[arg(long)]
predict_mode: Option<String>,
#[arg(short = 'a', long = "predict-always")]
predict_always: bool,
#[arg(short = 'n', long = "predict-never")]
predict_never: bool,
#[arg(long)]
escape_key: Option<String>,
#[arg(long)]
ssh_port: Option<u16>,
#[arg(long)]
ssh_auth_command: Option<String>,
#[arg(long)]
ssh_key: Option<PathBuf>,
#[arg(long)]
ssh_known_hosts: Option<PathBuf>,
#[arg(long)]
ssh_control_path: Option<PathBuf>,
#[arg(long)]
dosh_port: Option<u16>,
#[arg(long)]
dosh_host: Option<String>,
#[arg(short, long, action = ArgAction::Count)]
verbose: u8,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
struct CachedCredential {
server: String,
session: String,
mode: String,
udp_host: String,
udp_port: u16,
client_id: [u8; 16],
session_key: [u8; 32],
session_key_id: [u8; 16],
attach_ticket: Vec<u8>,
attach_ticket_psk: [u8; 32],
last_rendered_seq: u64,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct LocalForward {
bind_host: String,
listen_port: u16,
target_host: String,
target_port: u16,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct RemoteForward {
bind_host: String,
listen_port: u16,
target_host: String,
target_port: u16,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct DynamicForward {
bind_host: String,
listen_port: u16,
}
enum ForwardEvent {
Open {
stream_id: u64,
target_host: String,
target_port: u16,
writer: tokio::net::tcp::OwnedWriteHalf,
socks5_reply: bool,
},
Data {
stream_id: u64,
bytes: Vec<u8>,
},
Close {
stream_id: u64,
},
}
#[derive(Clone)]
struct PendingStreamChunk {
offset: u64,
bytes: Vec<u8>,
last_sent: Instant,
attempts: u32,
}
#[derive(Clone)]
struct PendingStreamOpen {
target_host: String,
target_port: u16,
last_sent: Instant,
attempts: u32,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum StartupGateMode {
HoldAll,
HoldControl,
}
#[tokio::main(flavor = "current_thread")]
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)?);
}
if args.server.as_deref() == Some("setup") {
return run_setup_command(&config, &args).await;
}
if args.server.as_deref() == Some("selftest") {
return run_selftest_command(&config, &args).await;
}
if args.server.as_deref() == Some("proxy-stdio") {
return run_proxy_stdio_command(&config, &args).await;
}
if args.server.as_deref() == Some("vscode") {
return run_vscode_command(&config, &args);
}
if matches!(args.server.as_deref(), Some("recover" | "repair")) {
return run_recover_command(&config, &args).await;
}
if matches!(args.server.as_deref(), Some("status" | "sessions")) {
return run_status_command(&config, &args);
}
if args.server.as_deref() == Some("restart") {
return run_restart_command(&config, &args);
}
if args.server.as_deref() == Some("forward") {
args = rewrite_forward_command(args)?;
}
if args.server.as_deref() == Some("import-ssh") {
return run_import_ssh(&config, &args.command);
}
if args.server.as_deref() == Some("trust") {
return run_trust_command(&config, &args);
}
if args.server.as_deref() == Some("doctor") {
return run_doctor_command(&config, &args).await;
}
let requested_server = args.server.clone().unwrap_or_else(|| config.server.clone());
let hosts = load_hosts_config(None).unwrap_or_default();
let host = hosts
.hosts
.get(&requested_server)
.cloned()
.unwrap_or_default();
let raw_server = host.ssh.clone().unwrap_or_else(|| requested_server.clone());
let server = ssh_with_user(&raw_server, host.user.as_deref());
let startup_command = resolved_startup_command(&args.command, &config, &host);
let local_forwards = parse_local_forwards(&args.local_forward)?;
let remote_forwards = parse_remote_forwards(&args.remote_forward)?;
let dynamic_forwards = parse_dynamic_forwards(&args.dynamic_forward)?;
// SSH-agent forwarding (opt-in via -A / forward_agent). SECURITY: this exposes
// your local agent to the remote host for the session's lifetime; only ever
// active on explicit opt-in AND when the server's allow_agent_forwarding is
// set. We capture the LOCAL agent socket path now so the run loop can splice
// each server-initiated agent stream into it.
let forward_agent = args.forward_agent || config.forward_agent;
let agent_sock = if forward_agent {
match std::env::var_os("SSH_AUTH_SOCK") {
Some(path) if !path.is_empty() => Some(PathBuf::from(path)),
_ => {
return Err(anyhow!(
"agent forwarding requested but SSH_AUTH_SOCK is not set"
));
}
}
} else {
None
};
let forwarding_requested = !local_forwards.is_empty()
|| !remote_forwards.is_empty()
|| !dynamic_forwards.is_empty()
|| forward_agent;
let predict = args.predict || host.predict.unwrap_or(config.predict);
let predict_mode = selected_predict_mode(&args, &config)?;
let escape_key = selected_escape_key(&args, &config)?;
let session = select_session(args.session.as_deref(), args.new);
let mode = if args.forward_only {
"forward-only"
} else if args.view_only || config.view_only {
"view-only"
} else {
"read-write"
}
.to_string();
let ssh_port = args.ssh_port.or(host.ssh_port).or(config.ssh_port);
let ssh_auth_command = args
.ssh_auth_command
.clone()
.or_else(|| config.ssh_auth_command.clone())
.unwrap_or_else(|| "~/.local/bin/dosh-auth".to_string());
if args
.command
.first()
.is_some_and(|command| command == "sessions")
{
return run_sessions_command(
&server,
ssh_port,
args.ssh_key.as_deref(),
args.ssh_known_hosts.as_deref(),
);
}
let dosh_port = args.dosh_port.or(host.port).unwrap_or(config.dosh_port);
let cache_path = cache_path(&config.credential_cache, &requested_server, &session, &mode);
let (cols, rows) = terminal_size();
let auth_preference = args
.auth
.clone()
.unwrap_or_else(|| config.auth_preference.clone());
if forwarding_requested && args.local_auth {
return Err(anyhow!(
"port forwarding requires native auth; remove --local-auth"
));
}
if forwarding_requested && !auth_allows(&auth_preference, "native") {
return Err(anyhow!("port forwarding requires --auth native or auto"));
}
if args.background && std::env::var_os("DOSH_BACKGROUND_CHILD").is_none() {
return spawn_background_forwarder(&args, forwarding_requested);
}
let allow_cache = !args.no_cache && !forwarding_requested;
let started = Instant::now();
let mut resolved_ssh_config = None;
let requested_udp_host = args
.dosh_host
.clone()
.or_else(|| host.dosh_host.clone())
.or_else(|| config.dosh_host.clone());
let target_udp_host = if args.local_auth {
requested_udp_host.unwrap_or_else(|| "127.0.0.1".to_string())
} else {
match ssh_config(&server, ssh_port) {
Ok(parsed) => {
let hostname = selected_udp_host(requested_udp_host.as_deref(), &server, &parsed)?;
resolved_ssh_config = Some(parsed);
hostname
}
Err(err) => {
log_debug(
args.verbose,
2,
&format!("dosh could not resolve SSH config hostname, using {server}: {err:#}"),
);
resolved_ssh_config = Some(SshConfig::default());
selected_udp_host(
requested_udp_host.as_deref(),
&server,
&SshConfig::default(),
)?
}
}
};
let credential = if allow_cache {
load_cache(&cache_path).ok()
} else {
None
};
let cache_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref());
let socket = UdpSocket::bind("0.0.0.0:0").await?;
if let Some(mut cached) = credential.clone() {
cached.udp_host = target_udp_host.clone();
cached.udp_port = dosh_port;
log_timing(args.verbose, "credential_lookup_end", started.elapsed());
if config.cache_attach_tickets && allow_cache {
match try_ticket_attach(&socket, &cached, cols, rows, cache_requested_env.clone()).await
{
Ok((frame, cred)) => {
log_timing(args.verbose, "udp_ticket_attach_ready", started.elapsed());
save_cache(&cache_path, &cred)?;
if args.attach_only {
render_frame(&frame)?;
detach_once(&socket, &cred, 2).await?;
return Ok(());
}
return run_terminal(
socket,
cred,
Some(frame),
predict,
predict_mode,
escape_key,
startup_command,
config.reconnect_timeout_secs,
Vec::new(),
Vec::new(),
false,
None,
)
.await;
}
Err(err) => {
log_debug(
args.verbose,
2,
&format!(
"dosh ticket attach failed, falling back to SSH bootstrap: {err:#}"
),
);
}
}
} else {
match try_resume_fresh_process(&socket, &cached, cols, rows).await {
Ok((frame, cred)) => {
log_timing(args.verbose, "udp_resume_ready", started.elapsed());
if allow_cache {
save_cache(&cache_path, &cred)?;
}
if args.attach_only {
render_frame(&frame)?;
detach_once(&socket, &cred, 2).await?;
return Ok(());
}
return run_terminal(
socket,
cred,
Some(frame),
predict,
predict_mode,
escape_key,
startup_command,
config.reconnect_timeout_secs,
Vec::new(),
Vec::new(),
false,
None,
)
.await;
}
Err(err) => {
log_debug(
args.verbose,
2,
&format!("dosh resume failed, falling back to SSH bootstrap: {err:#}"),
);
}
}
}
}
if !args.local_auth && auth_allows(&auth_preference, "native") {
if resolved_ssh_config.is_none() {
resolved_ssh_config = Some(ssh_config(&server, ssh_port).unwrap_or_default());
}
let cold_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref());
let native_start = Instant::now();
match try_native_auth(
&socket,
&config,
&requested_server,
&server,
ssh_port,
&target_udp_host,
dosh_port,
args.ssh_key.as_deref(),
&session,
&mode,
cols,
rows,
forwarding_requests(
&local_forwards,
&remote_forwards,
&dynamic_forwards,
forward_agent,
),
resolved_ssh_config.as_ref(),
cold_requested_env,
true,
)
.await
{
Ok((frame, cred)) => {
log_timing(args.verbose, "native_auth", native_start.elapsed());
if allow_cache {
save_cache(&cache_path, &cred)?;
}
log_timing(args.verbose, "terminal_ready", started.elapsed());
if args.attach_only {
render_frame(&frame)?;
detach_once(&socket, &cred, 2).await?;
return Ok(());
}
return run_terminal(
socket,
cred,
Some(frame),
predict,
predict_mode,
escape_key,
startup_command,
config.reconnect_timeout_secs,
local_forwards,
dynamic_forwards,
args.forward_only,
agent_sock,
)
.await;
}
Err(err) if !forwarding_requested && auth_allows(&auth_preference, "ssh") => {
log_debug(
args.verbose,
2,
&format!("dosh native auth failed, falling back to SSH bootstrap: {err:#}"),
);
}
Err(err) => return Err(err.context("native auth failed")),
}
}
if !args.local_auth && !auth_allows(&auth_preference, "ssh") {
return Err(anyhow!(
"auth preference {auth_preference:?} does not allow SSH fallback"
));
}
let bootstrap_start = Instant::now();
if !args.local_auth && resolved_ssh_config.is_none() {
resolved_ssh_config = Some(ssh_config(&server, ssh_port).unwrap_or_default());
}
let cold_requested_env = requested_env(&config, &host, resolved_ssh_config.as_ref());
let bootstrap = if args.local_auth {
local_bootstrap(&session, &mode, cols, rows, dosh_port, target_udp_host)?
} else {
let mut bootstrap = ssh_bootstrap(
&server,
ssh_port,
&ssh_auth_command,
args.ssh_key.as_deref(),
args.ssh_known_hosts.as_deref(),
args.ssh_control_path.as_deref(),
&session,
&mode,
cols,
rows,
)?;
bootstrap.udp_host = target_udp_host;
bootstrap.udp_port = dosh_port;
bootstrap
};
log_timing(args.verbose, "ssh_bootstrap", bootstrap_start.elapsed());
let (ok, mut cred) =
bootstrap_attach(&socket, &server, &bootstrap, cols, rows, cold_requested_env).await?;
cred.last_rendered_seq = ok.initial_seq;
if allow_cache {
save_cache(&cache_path, &cred)?;
}
log_timing(args.verbose, "terminal_ready", started.elapsed());
let first = Frame {
session: ok.session,
output_seq: ok.initial_seq,
bytes: ok.snapshot,
snapshot: true,
closed: false,
};
if args.attach_only {
render_frame(&first)?;
detach_once(&socket, &cred, 2).await?;
return Ok(());
}
run_terminal(
socket,
cred,
Some(first),
predict,
predict_mode,
escape_key,
startup_command,
config.reconnect_timeout_secs,
Vec::new(),
Vec::new(),
false,
None,
)
.await
}
fn run_trust_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
let host = args
.command
.first()
.ok_or_else(|| anyhow!("usage: dosh trust [--remove] [--replace] <host>"))?;
let path = expand_tilde(&config.known_hosts);
if args.remove {
if remove_trusted_host(&path, host)? {
eprintln!("dosh trust: removed {host} from {}", path.display());
} else {
eprintln!("dosh trust: {host} was not trusted in {}", path.display());
}
return Ok(());
}
let hosts = load_hosts_config(None).unwrap_or_default();
let host_config = hosts.hosts.get(host).cloned().unwrap_or_default();
let raw_server = host_config.ssh.clone().unwrap_or_else(|| host.to_string());
let server = ssh_with_user(&raw_server, host_config.user.as_deref());
let ssh_port = args.ssh_port.or(host_config.ssh_port).or(config.ssh_port);
let ssh_auth_command = args
.ssh_auth_command
.clone()
.or_else(|| config.ssh_auth_command.clone())
.unwrap_or_else(|| "~/.local/bin/dosh-auth".to_string());
let public_key = fetch_host_key_over_ssh(
&server,
ssh_port,
&ssh_auth_command,
args.ssh_key.as_deref(),
args.ssh_known_hosts.as_deref(),
args.ssh_control_path.as_deref(),
)?;
match trust_host(&path, host, &public_key, "ssh", args.replace)? {
TrustResult::AlreadyTrusted => {
eprintln!(
"dosh trust: {host} already trusted ({})",
host_fingerprint(&public_key)
);
}
TrustResult::Trusted => {
eprintln!(
"dosh trust: trusted {host} {} in {}",
host_fingerprint(&public_key),
path.display()
);
}
}
Ok(())
}
async fn run_setup_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
let host = args
.command
.first()
.cloned()
.ok_or_else(|| anyhow!("usage: dosh setup <ssh-host>"))?;
println!("Dosh setup for {host}");
let imported = import_ssh_aliases(config, std::slice::from_ref(&host))?;
for entry in imported {
match entry.action {
ImportAction::Added => println!(
"[ok] host config: [{}] ssh={} dosh_host={} port={}",
entry.alias, entry.ssh, entry.dosh_host, entry.port
),
ImportAction::Skipped => println!("[ok] host config: [{}] already exists", entry.alias),
}
}
let hosts = load_hosts_config(None).unwrap_or_default();
let host_config = hosts.hosts.get(&host).cloned().unwrap_or_default();
let raw_server = host_config.ssh.clone().unwrap_or_else(|| host.clone());
let server = ssh_with_user(&raw_server, host_config.user.as_deref());
let ssh_port = args.ssh_port.or(host_config.ssh_port).or(config.ssh_port);
let ssh_auth_command = args
.ssh_auth_command
.clone()
.or_else(|| config.ssh_auth_command.clone())
.unwrap_or_else(|| "~/.local/bin/dosh-auth".to_string());
let known_hosts = expand_tilde(&config.known_hosts);
let public_key = fetch_host_key_over_ssh(
&server,
ssh_port,
&ssh_auth_command,
args.ssh_key.as_deref(),
args.ssh_known_hosts.as_deref(),
args.ssh_control_path.as_deref(),
)
.with_context(|| {
format!("fetch Dosh host key over SSH; make sure dosh-auth is installed on {server}")
})?;
match trust_host(&known_hosts, &host, &public_key, "ssh", args.replace)? {
TrustResult::AlreadyTrusted => println!(
"[ok] trust: already trusted {}",
host_fingerprint(&public_key)
),
TrustResult::Trusted => println!(
"[ok] trust: pinned {} in {}",
host_fingerprint(&public_key),
known_hosts.display()
),
}
println!();
run_doctor_for_host(config, args, &host).await?;
println!();
println!("Next:");
println!(" dosh {host}");
println!(" dosh update --check");
Ok(())
}
async fn run_selftest_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
let host = args
.command
.first()
.cloned()
.ok_or_else(|| anyhow!("usage: dosh selftest <host>"))?;
println!("Dosh selftest for {host}");
println!("[ok] terminal size fallback: {:?}", terminal_size());
let overlay = render_status_overlay("[dosh] reconnecting — last contact 3s ago", 24);
ensure_tui_safe_status_overlay(&overlay)?;
println!("[ok] disconnect UI: save/restore top status bar");
parse_local_forward("18080:127.0.0.1:80")?;
parse_remote_forward("19090:127.0.0.1:5432")?;
parse_dynamic_forward("11080")?;
println!("[ok] forwarding parser: -L, -R, -D");
println!();
run_doctor_for_host(config, args, &host).await?;
println!();
println!("[ok] selftest complete");
Ok(())
}
async fn run_doctor_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
let requested = args
.command
.first()
.cloned()
.ok_or_else(|| anyhow!("usage: dosh doctor <host>"))?;
run_doctor_for_host(config, args, &requested).await
}
async fn run_recover_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
if args.command.len() != 1 {
return Err(anyhow!("usage: dosh recover <host>"));
}
let requested = args.command[0].clone();
println!("Dosh recovery for {requested}");
let removed = clear_cached_credentials(&config.credential_cache, &requested)
.context("clear cached attach credentials")?;
println!("[ok] removed {removed} cached credential(s)");
println!("[info] host trust was left unchanged");
run_doctor_for_host(config, args, &requested).await
}
fn run_status_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
if args.command.len() != 1 {
return Err(anyhow!("usage: dosh status <host>"));
}
let requested = args.command[0].clone();
let hosts = load_hosts_config(None).unwrap_or_default();
let host_config = hosts.hosts.get(&requested).cloned().unwrap_or_default();
let server = status_ssh_target(&requested, &host_config);
let ssh_port = args.ssh_port.or(host_config.ssh_port).or(config.ssh_port);
println!("Dosh status for {requested}");
if is_local_status_target(&server) {
return run_local_sessions_command();
}
run_sessions_command(
&server,
ssh_port,
args.ssh_key.as_deref(),
args.ssh_known_hosts.as_deref(),
)
}
fn run_restart_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
if args.command.len() != 1 {
return Err(anyhow!("usage: dosh restart <host>"));
}
let requested = args.command[0].clone();
let hosts = load_hosts_config(None).unwrap_or_default();
let host_config = hosts.hosts.get(&requested).cloned().unwrap_or_default();
let server = status_ssh_target(&requested, &host_config);
let ssh_port = args.ssh_port.or(host_config.ssh_port).or(config.ssh_port);
println!("Restarting Dosh server for {requested}");
if is_local_status_target(&server) {
return run_local_script(RESTART_STATUS_SCRIPT, "restart local Dosh server");
}
run_remote_script(
&server,
ssh_port,
args.ssh_key.as_deref(),
args.ssh_known_hosts.as_deref(),
RESTART_STATUS_SCRIPT,
"restart remote Dosh server",
)
}
fn status_ssh_target(requested: &str, host_config: &dosh::config::HostConfig) -> String {
let raw_server = host_config
.ssh
.clone()
.unwrap_or_else(|| requested.to_string());
ssh_with_user(&raw_server, host_config.user.as_deref())
}
fn is_local_status_target(server: &str) -> bool {
let host = server.rsplit('@').next().unwrap_or(server);
matches!(host, "localhost" | "127.0.0.1" | "::1")
}
const SESSIONS_STATUS_SCRIPT: &str = r#"printf 'tmux sessions:\n'; tmux ls 2>/dev/null || printf ' none\n'; printf '\nDosh service:\n'; systemctl --user --no-pager --plain status dosh-server.service 2>/dev/null | sed -n '1,8p' || pgrep -af dosh-server || printf ' not running\n'"#;
const RESTART_STATUS_SCRIPT: &str = r#"if command -v systemctl >/dev/null 2>&1 && systemctl --user list-unit-files dosh-server.service >/dev/null 2>&1; then systemctl --user restart dosh-server.service; else pkill -f 'dosh-server serve' 2>/dev/null || true; nohup "$HOME/.local/bin/dosh-server" serve >/tmp/dosh-server.log 2>&1 & fi; sleep 1; systemctl --user --no-pager --plain status dosh-server.service 2>/dev/null | sed -n '1,8p' || pgrep -af dosh-server || printf ' not running\n'"#;
fn run_local_sessions_command() -> Result<()> {
run_local_script(SESSIONS_STATUS_SCRIPT, "run local status command")
}
fn run_local_script(script: &str, context: &str) -> Result<()> {
let status = Command::new("sh")
.arg("-c")
.arg(script)
.status()
.with_context(|| context.to_string())?;
if !status.success() {
return Err(anyhow!("{context} failed with status {status}"));
}
Ok(())
}
async fn run_doctor_for_host(
config: &dosh::config::ClientConfig,
args: &Args,
requested: &str,
) -> Result<()> {
let hosts = load_hosts_config(None).unwrap_or_default();
let host_config = hosts.hosts.get(requested).cloned().unwrap_or_default();
let raw_server = host_config
.ssh
.clone()
.unwrap_or_else(|| requested.to_string());
let server = ssh_with_user(&raw_server, host_config.user.as_deref());
let ssh_port = args.ssh_port.or(host_config.ssh_port).or(config.ssh_port);
let dosh_port = args
.dosh_port
.or(host_config.port)
.unwrap_or(config.dosh_port);
println!("Dosh doctor for {requested}");
println!(
"[info] config: client={} hosts={} known_hosts={}",
expand_tilde("~/.config/dosh/client.toml").display(),
expand_tilde("~/.config/dosh/hosts.toml").display(),
expand_tilde(&config.known_hosts).display()
);
println!(
"[info] client: auth={} predict={}({}) disconnect_status={} cache_tickets={}",
config.auth_preference,
config.predict,
config.predict_mode,
config.disconnect_status,
config.cache_attach_tickets
);
let parsed_ssh_config = match ssh_config(&server, ssh_port) {
Ok(parsed) => {
let hostname = parsed
.hostname
.clone()
.unwrap_or_else(|| ssh_destination_host(&server));
println!(
"[ok] ssh config: host={} user={} port={} identities_only={} proxy={}",
hostname,
parsed.user.as_deref().unwrap_or("<default>"),
parsed.port.unwrap_or(ssh_port.unwrap_or(22)),
parsed.identities_only,
parsed
.proxy_jump
.as_deref()
.or(parsed.proxy_command.as_deref())
.unwrap_or("<none>")
);
parsed
}
Err(err) => {
println!("[fail] ssh config: {err:#}");
SshConfig::default()
}
};
match ssh_fallback_probe(
&server,
ssh_port,
args.ssh_key.as_deref(),
args.ssh_known_hosts.as_deref(),
) {
Ok(()) => println!("[ok] ssh fallback: reachable"),
Err(err) => println!("[warn] ssh fallback: {err:#}"),
}
let requested_udp_host = args
.dosh_host
.clone()
.or_else(|| host_config.dosh_host.clone())
.or_else(|| config.dosh_host.clone());
let udp_host =
match selected_udp_host(requested_udp_host.as_deref(), &server, &parsed_ssh_config) {
Ok(host) => host,
Err(err) => {
println!("[fail] native endpoint: {err:#}");
ssh_destination_host(&server)
}
};
println!("[info] native endpoint: {udp_host}:{dosh_port}");
match resolve_addr(&udp_host, dosh_port) {
Ok(addr) => println!("[ok] udp resolve: {addr}"),
Err(err) => println!("[fail] udp resolve: {err:#}"),
}
println!(
"[info] forwarding requested by config: agent={} default_session={}",
config.forward_agent, config.default_session
);
let socket = UdpSocket::bind("0.0.0.0:0").await?;
match try_native_auth_check(
&socket,
config,
requested,
&server,
ssh_port,
&udp_host,
dosh_port,
args.ssh_key.as_deref(),
Some(&parsed_ssh_config),
)
.await
{
Ok(check) => {
println!("[ok] native udp: reachable");
println!(
"[ok] server version: {}",
if check.server_version.is_empty() {
"<unknown>"
} else {
&check.server_version
}
);
if server_version_mismatch(&check.server_version) {
println!(
"[warn] version mismatch: client={} server={}",
env!("CARGO_PKG_VERSION"),
check.server_version
);
}
println!("[ok] known host: trusted");
println!("[ok] native auth: authorized as {}", check.requested_user);
println!(
"[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_file_transfer,
check.allow_exec_command
);
Ok(())
}
Err(err) => {
println!("[fail] native/auth check: {err:#}");
Err(err.context("dosh doctor failed"))
}
}
}
fn server_version_mismatch(server_version: &str) -> bool {
!server_version.is_empty() && server_version != env!("CARGO_PKG_VERSION")
}
fn ensure_tui_safe_status_overlay(bytes: &[u8]) -> Result<()> {
anyhow::ensure!(
bytes.starts_with(b"\x1b7"),
"status overlay must save cursor"
);
anyhow::ensure!(
bytes.ends_with(b"\x1b8"),
"status overlay must restore cursor"
);
anyhow::ensure!(
!contains_bytes(bytes, b"\n") && !contains_bytes(bytes, b"\r"),
"status overlay must not scroll"
);
anyhow::ensure!(
!contains_bytes(bytes, b"\x1b[?25l") && !contains_bytes(bytes, b"\x1b[?25h"),
"status overlay must not toggle cursor visibility"
);
Ok(())
}
fn ssh_fallback_probe(
server: &str,
ssh_port: Option<u16>,
ssh_key: Option<&Path>,
ssh_known_hosts: Option<&Path>,
) -> Result<()> {
let mut command = Command::new("ssh");
command.arg("-o").arg("BatchMode=yes");
command.arg("-o").arg("ConnectTimeout=5");
command.arg("-T");
if let Some(ssh_port) = ssh_port {
command.arg("-p").arg(ssh_port.to_string());
}
if let Some(key) = ssh_key {
command.arg("-i").arg(key);
}
if let Some(known_hosts) = ssh_known_hosts {
command
.arg("-o")
.arg(format!("UserKnownHostsFile={}", known_hosts.display()));
}
let output = command.arg(server).arg("true").output()?;
anyhow::ensure!(
output.status.success(),
"{}",
String::from_utf8_lossy(&output.stderr)
);
Ok(())
}
fn fetch_host_key_over_ssh(
server: &str,
ssh_port: Option<u16>,
ssh_auth_command: &str,
ssh_key: Option<&Path>,
ssh_known_hosts: Option<&Path>,
ssh_control_path: Option<&Path>,
) -> Result<dosh::native::HostPublicKey> {
let mut command = Command::new("ssh");
if let Some(ssh_port) = ssh_port {
command.arg("-p").arg(ssh_port.to_string());
}
command.arg("-T");
if let Some(key) = ssh_key {
command.arg("-i").arg(key);
}
if let Some(known_hosts) = ssh_known_hosts {
command
.arg("-o")
.arg(format!("UserKnownHostsFile={}", known_hosts.display()));
}
if let Some(control_path) = ssh_control_path {
command.arg("-S").arg(control_path);
}
let output = command
.arg(server)
.arg(ssh_auth_command)
.arg("--host-key")
.output()
.context("fetch Dosh host key over SSH")?;
if !output.status.success() {
return Err(anyhow!(
"Dosh host-key fetch failed: {}",
String::from_utf8_lossy(&output.stderr)
));
}
let raw = String::from_utf8(output.stdout)?;
parse_host_public_key_line(&raw)
}
fn startup_command(args: &[String]) -> Option<Vec<u8>> {
if args.is_empty() {
return None;
}
let mut command = args.join(" ");
command.push('\r');
Some(command.into_bytes())
}
fn resolved_startup_command(
args: &[String],
config: &dosh::config::ClientConfig,
host: &dosh::config::HostConfig,
) -> Option<Vec<u8>> {
if args.is_empty() {
return host
.default_command
.as_deref()
.map(startup_command_from_string);
}
match resolve_command_extension(args, config, host) {
ExtensionResolution::Command(command) => Some(startup_command_from_string(&command)),
ExtensionResolution::Disabled => None,
ExtensionResolution::NoMatch => startup_command(args),
}
}
enum ExtensionResolution {
Command(String),
Disabled,
NoMatch,
}
fn resolve_command_extension(
args: &[String],
config: &dosh::config::ClientConfig,
host: &dosh::config::HostConfig,
) -> ExtensionResolution {
let Some(name) = args.first() else {
return ExtensionResolution::NoMatch;
};
if let Some(extension) = host.extensions.get(name) {
if extension.disabled {
return ExtensionResolution::Disabled;
}
if let Some(command) = extension.command.as_deref() {
return ExtensionResolution::Command(expand_command_extension(command, &args[1..]));
}
}
if let Some(extension) = config.extensions.get(name) {
if extension.disabled {
return ExtensionResolution::Disabled;
}
if let Some(command) = extension.command.as_deref() {
return ExtensionResolution::Command(expand_command_extension(command, &args[1..]));
}
}
ExtensionResolution::NoMatch
}
fn expand_command_extension(template: &str, args: &[String]) -> String {
let quoted_args = args
.iter()
.map(|arg| shell_word(arg))
.collect::<Vec<_>>()
.join(" ");
if template.contains("{args}") {
return template.replace("{args}", &quoted_args);
}
if args.is_empty() {
template.to_string()
} else {
format!("{template} {quoted_args}")
}
}
fn startup_command_from_string(command: &str) -> Vec<u8> {
let mut command = command.to_string();
command.push('\r');
command.into_bytes()
}
fn spawn_background_forwarder(args: &Args, forwarding_requested: bool) -> Result<()> {
if !forwarding_requested {
return Err(anyhow!(
"-f requires at least one -L, -R, or -D forwarding request"
));
}
if !args.forward_only {
return Err(anyhow!("-f currently requires -N for forward-only mode"));
}
let listener =
StdTcpListener::bind("127.0.0.1:0").context("bind background readiness socket")?;
listener
.set_nonblocking(true)
.context("configure background 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 child = Command::new(exe)
.args(std::env::args_os().skip(1))
.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::null())
.spawn()
.context("spawn background dosh forwarder")?;
if let Some(path) = std::env::var_os("DOSH_BACKGROUND_PID_FILE") {
fs::write(path, child.id().to_string()).context("write background pid file")?;
}
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 {
eprintln!("dosh background forwarding ready");
return Ok(());
}
return Err(anyhow!("background forwarder sent invalid readiness token"));
}
Err(err) if err.kind() == std::io::ErrorKind::WouldBlock => {}
Err(err) => return Err(err).context("wait for background readiness"),
}
if let Some(status) = child.try_wait()? {
return Err(anyhow!(
"background dosh forwarder exited before readiness: {status}"
));
}
if Instant::now() >= deadline {
let _ = child.kill();
let _ = child.wait();
return Err(anyhow!("background dosh forwarder did not become ready"));
}
std::thread::sleep(Duration::from_millis(20));
}
}
fn rewrite_forward_command(mut args: Args) -> Result<Args> {
let mut tokens = std::mem::take(&mut args.command).into_iter();
let host = tokens.next().ok_or_else(|| {
anyhow!("usage: dosh forward <host> [-L spec] [-R spec] [-D spec] [-A] [-f]")
})?;
let mut trailing_command = Vec::new();
while let Some(token) = tokens.next() {
match token.as_str() {
"-L" | "--local-forward" => {
let spec = tokens
.next()
.ok_or_else(|| anyhow!("dosh forward: {token} requires a forwarding spec"))?;
args.local_forward.push(spec);
}
"-R" | "--remote-forward" => {
let spec = tokens
.next()
.ok_or_else(|| anyhow!("dosh forward: {token} requires a forwarding spec"))?;
args.remote_forward.push(spec);
}
"-D" | "--dynamic-forward" => {
let spec = tokens
.next()
.ok_or_else(|| anyhow!("dosh forward: {token} requires a forwarding spec"))?;
args.dynamic_forward.push(spec);
}
"-A" | "--forward-agent" => args.forward_agent = true,
"-f" | "--background" => args.background = true,
"-N" | "--forward-only" => args.forward_only = true,
"--" => {
trailing_command.extend(tokens);
break;
}
other => return Err(anyhow!("dosh forward: unknown option {other:?}")),
}
}
if !trailing_command.is_empty() {
return Err(anyhow!(
"dosh forward does not run remote commands; remove {:?}",
trailing_command
));
}
if args.local_forward.is_empty()
&& args.remote_forward.is_empty()
&& args.dynamic_forward.is_empty()
&& !args.forward_agent
{
return Err(anyhow!(
"dosh forward requires at least one -L, -R, -D, or -A request"
));
}
args.server = Some(host);
args.command = Vec::new();
args.forward_only = true;
Ok(args)
}
fn signal_background_ready() {
let Some(addr) = std::env::var_os("DOSH_BACKGROUND_READY_ADDR") else {
return;
};
let Some(token) = std::env::var_os("DOSH_BACKGROUND_READY_TOKEN") else {
return;
};
if let Ok(mut stream) = StdTcpStream::connect(addr.to_string_lossy().as_ref()) {
let _ = stream.write_all(token.to_string_lossy().as_bytes());
let _ = stream.shutdown(std::net::Shutdown::Write);
}
}
#[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<FileResponse>,
}
struct ExecServiceClient {
stream: StdTcpStream,
decoder: ExecFrameDecoder,
pending: VecDeque<ExecResponse>,
}
async fn run_proxy_stdio_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
if args.command.len() != 3 {
return Err(anyhow!(
"usage: dosh proxy-stdio HOST TARGET_HOST TARGET_PORT"
));
}
anyhow::ensure!(
!args.local_auth,
"proxy-stdio requires native auth; local auth cannot authenticate SSH ProxyCommand streams"
);
let requested_server = args.command[0].clone();
let target_host = args.command[1].clone();
let target_port = args.command[2]
.parse::<u16>()
.with_context(|| format!("invalid proxy target port {:?}", args.command[2]))?;
anyhow::ensure!(
valid_forward_host(&target_host),
"invalid proxy target host"
);
anyhow::ensure!(target_port != 0, "proxy target port cannot be 0");
let hosts = load_hosts_config(None).unwrap_or_default();
let host = hosts
.hosts
.get(&requested_server)
.cloned()
.unwrap_or_default();
let raw_server = host.ssh.clone().unwrap_or_else(|| requested_server.clone());
let server = ssh_with_user(&raw_server, host.user.as_deref());
let ssh_port = args.ssh_port.or(host.ssh_port).or(config.ssh_port);
let resolved_ssh_config = ssh_config(&server, ssh_port).unwrap_or_default();
let requested_udp_host = args
.dosh_host
.clone()
.or_else(|| host.dosh_host.clone())
.or_else(|| config.dosh_host.clone());
let target_udp_host =
selected_udp_host(requested_udp_host.as_deref(), &server, &resolved_ssh_config)?;
let dosh_port = args.dosh_port.or(host.port).unwrap_or(config.dosh_port);
let socket = UdpSocket::bind("0.0.0.0:0").await?;
let session = args
.session
.clone()
.unwrap_or_else(protocol::generate_implicit_session_name);
let requested_forwardings = vec![ForwardingRequest {
kind: ForwardingKind::Local,
bind_host: None,
listen_port: 0,
target_host: Some(target_host.clone()),
target_port: Some(target_port),
}];
let (_frame, cred) = try_native_auth(
&socket,
config,
&requested_server,
&server,
ssh_port,
&target_udp_host,
dosh_port,
args.ssh_key.as_deref(),
&session,
"forward-only",
80,
24,
requested_forwardings,
Some(&resolved_ssh_config),
Vec::new(),
false,
)
.await?;
let peer_addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
let mut transport = DoshTransport::new_owned(
socket,
SessionTransportConfig {
role: SessionRole::Client,
conn_id: cred.client_id,
session_key: cred.session_key,
peer_addr,
initial_send_seq: 2,
initial_ack: cred.last_rendered_seq,
stream: TransportConfig::default(),
},
);
let stream_id = transport.open_target(target_host, target_port).await?;
proxy_stdio_loop(transport, stream_id).await
}
async fn proxy_stdio_loop(mut transport: DoshTransport, stream_id: u64) -> Result<()> {
let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::<Vec<u8>>();
std::thread::Builder::new()
.name("dosh-proxy-stdin".to_string())
.spawn(move || {
let mut stdin = std::io::stdin();
let mut buf = [0u8; 16 * 1024];
loop {
match stdin.read(&mut buf) {
Ok(0) => break,
Ok(n) => {
if stdin_tx.send(buf[..n].to_vec()).is_err() {
break;
}
}
Err(_) => break,
}
}
})?;
let mut stdout = tokio::io::stdout();
let mut maintenance = tokio::time::interval(Duration::from_millis(50));
let mut stdin_closed = false;
loop {
tokio::select! {
stdin_msg = stdin_rx.recv(), if !stdin_closed => {
match stdin_msg {
Some(bytes) => transport.send(stream_id, bytes).await?,
None => {
stdin_closed = true;
let _ = transport.close(stream_id).await;
}
}
}
event = transport.recv() => {
match event? {
SessionEvent::Stream(TransportEvent::OpenOk { stream_id: opened, .. }) if opened == stream_id => {}
SessionEvent::Stream(TransportEvent::OpenReject { stream_id: rejected, reason }) if rejected == stream_id => {
return Err(anyhow!("Dosh proxy stream rejected: {reason}"));
}
SessionEvent::Stream(TransportEvent::Data(data)) if data.stream_id == stream_id => {
for chunk in data.chunks {
stdout.write_all(&chunk).await?;
}
stdout.flush().await?;
}
SessionEvent::Stream(TransportEvent::Close { stream_id: closed }) if closed == stream_id => {
stdout.flush().await?;
return Ok(());
}
SessionEvent::Stream(_) | SessionEvent::Ping | SessionEvent::Pong | SessionEvent::Ignored => {}
}
}
_ = maintenance.tick() => {
transport.maintenance().await?;
}
}
}
}
fn run_vscode_command(config: &dosh::config::ClientConfig, args: &Args) -> Result<()> {
let mut tokens = args.command.iter();
let first = tokens
.next()
.ok_or_else(|| anyhow!("usage: dosh vscode [setup|open] HOST [REMOTE_PATH]"))?;
let (launch, host, remote_path) = match first.as_str() {
"setup" => {
let host = tokens
.next()
.ok_or_else(|| anyhow!("usage: dosh vscode setup HOST"))?;
(false, host.as_str(), None)
}
"open" => {
let host = tokens
.next()
.ok_or_else(|| anyhow!("usage: dosh vscode open HOST [REMOTE_PATH]"))?;
(true, host.as_str(), tokens.next().map(String::as_str))
}
host => (true, host, tokens.next().map(String::as_str)),
};
if tokens.next().is_some() {
return Err(anyhow!("dosh vscode accepts at most one remote path"));
}
let generated = ensure_vscode_ssh_host(config, args, host)?;
println!("[ok] VS Code SSH host: {}", generated.alias);
println!("[ok] SSH target: 127.0.0.1:{}", generated.target_ssh_port);
println!("[ok] Dosh proxy: {}", generated.proxy_command);
if launch {
launch_vscode_remote(&generated.alias, remote_path)?;
} else {
println!("Open in VS Code Remote-SSH as: {}", generated.alias);
}
Ok(())
}
struct GeneratedVscodeHost {
alias: String,
proxy_command: String,
target_ssh_port: u16,
}
fn ensure_vscode_ssh_host(
config: &dosh::config::ClientConfig,
args: &Args,
requested: &str,
) -> Result<GeneratedVscodeHost> {
let hosts = load_hosts_config(None).unwrap_or_default();
let host_config = hosts.hosts.get(requested).cloned().unwrap_or_default();
let raw_server = host_config
.ssh
.clone()
.unwrap_or_else(|| requested.to_string());
let server = ssh_with_user(&raw_server, host_config.user.as_deref());
let ssh_port = args.ssh_port.or(host_config.ssh_port).or(config.ssh_port);
let ssh_config = ssh_config(&server, ssh_port).unwrap_or_default();
let target_ssh_port = ssh_config.port.or(ssh_port).unwrap_or(22);
let user = host_config
.user
.or(ssh_username(&server))
.or(ssh_config.user);
let alias = format!("dosh-{}", vscode_safe_alias(requested));
let exe = std::env::current_exe()
.ok()
.map(|path| path.display().to_string())
.unwrap_or_else(|| "dosh".to_string());
let mut proxy_command = ssh_config_word(&exe);
if let Some(host) = args
.dosh_host
.clone()
.or_else(|| host_config.dosh_host.clone())
.or_else(|| config.dosh_host.clone())
{
proxy_command.push_str(&format!(" --dosh-host {}", ssh_config_word(&host)));
}
if let Some(port) = args.dosh_port.or(host_config.port) {
proxy_command.push_str(&format!(" --dosh-port {port}"));
}
proxy_command.push_str(&format!(
" proxy-stdio {} %h %p",
ssh_config_word(requested)
));
let ssh_dir = expand_tilde("~/.ssh");
fs::create_dir_all(&ssh_dir).with_context(|| format!("create {}", ssh_dir.display()))?;
let include_path = ssh_dir.join("config.dosh");
let main_config = ssh_dir.join("config");
ensure_ssh_include(&main_config, "config.dosh")?;
let mut block = String::new();
block.push_str(&format!("# BEGIN DOSH {alias}\n"));
block.push_str(&format!("Host {alias}\n"));
block.push_str(" HostName 127.0.0.1\n");
block.push_str(&format!(" Port {target_ssh_port}\n"));
block.push_str(&format!(" HostKeyAlias {requested}\n"));
if let Some(user) = user {
block.push_str(&format!(" User {user}\n"));
}
block.push_str(" ClearAllForwardings yes\n");
block.push_str(" ServerAliveInterval 15\n");
block.push_str(" ServerAliveCountMax 3\n");
block.push_str(&format!(" ProxyCommand {proxy_command}\n"));
block.push_str(&format!("# END DOSH {alias}\n"));
upsert_managed_block(&include_path, &alias, &block)?;
Ok(GeneratedVscodeHost {
alias,
proxy_command,
target_ssh_port,
})
}
fn ensure_ssh_include(path: &Path, include: &str) -> Result<()> {
let raw = fs::read_to_string(path).unwrap_or_default();
if raw.lines().any(|line| {
line.trim()
.eq_ignore_ascii_case(&format!("include {include}"))
}) {
return Ok(());
}
let mut next = String::new();
next.push_str(&format!("Include {include}\n"));
if !raw.is_empty() {
next.push('\n');
next.push_str(&raw);
}
fs::write(path, next).with_context(|| format!("write {}", path.display()))
}
fn upsert_managed_block(path: &Path, alias: &str, block: &str) -> Result<()> {
let raw = fs::read_to_string(path).unwrap_or_default();
let begin = format!("# BEGIN DOSH {alias}");
let end = format!("# END DOSH {alias}");
let mut out = String::new();
let mut skipping = false;
let mut replaced = false;
for line in raw.lines() {
if line.trim() == begin {
if !out.ends_with('\n') && !out.is_empty() {
out.push('\n');
}
out.push_str(block);
replaced = true;
skipping = true;
continue;
}
if skipping {
if line.trim() == end {
skipping = false;
}
continue;
}
out.push_str(line);
out.push('\n');
}
if !replaced {
if !out.trim().is_empty() && !out.ends_with("\n\n") {
out.push('\n');
}
out.push_str(block);
}
fs::write(path, out).with_context(|| format!("write {}", path.display()))
}
fn vscode_safe_alias(value: &str) -> String {
let mut out = String::new();
for byte in value.bytes() {
if byte.is_ascii_alphanumeric() || byte == b'-' || byte == b'_' {
out.push(byte as char);
} else {
out.push('-');
}
}
let out = out.trim_matches('-').to_string();
if out.is_empty() {
"host".to_string()
} else {
out
}
}
fn launch_vscode_remote(alias: &str, remote_path: Option<&str>) -> Result<()> {
let remote = format!("ssh-remote+{alias}");
let mut command = Command::new("code");
command.arg("--remote").arg(&remote);
if let Some(path) = remote_path {
command.arg(path);
}
match command.status() {
Ok(status) if status.success() => Ok(()),
Ok(status) => Err(anyhow!("code exited with status {status}")),
Err(err) if err.kind() == std::io::ErrorKind::NotFound => {
println!(
"Run: code --remote {}{}",
shell_word(&remote),
remote_path
.map(|path| format!(" {}", shell_word(path)))
.unwrap_or_default()
);
Ok(())
}
Err(err) => Err(err).context("launch VS Code"),
}
}
impl ExecServiceClient {
fn connect(port: u16) -> Result<Self> {
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<ExecResponse> {
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<Self> {
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<FileResponse> {
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<Option<dosh::file_transfer::FileMeta>> {
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<dosh::file_transfer::FileMeta> {
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<Vec<dosh::file_transfer::FileEntry>> {
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<u32>) -> 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<FileForwarder> {
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<u16> {
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::<std::io::Result<Vec<_>>>()?;
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<String> {
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<u32>,
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<String> {
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<u32>, preserve: bool) -> Option<u32> {
preserve.then_some(mode).flatten()
}
fn preserve_remote_mtime(modified_secs: Option<u64>, preserve: bool) -> Option<u64> {
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 {
let file = fs::OpenOptions::new()
.read(true)
.write(true)
.append(true)
.open(local)
.with_context(|| format!("open {}", local.display()))?;
file.set_len(offset)
.with_context(|| format!("truncate {}", local.display()))?;
file
} 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<Option<u64>> {
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<u64>, 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<Option<u32>> {
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<Option<u32>> {
Ok(None)
}
#[cfg(unix)]
fn set_local_mode(path: &Path, mode: Option<u32>, 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<u32>, _preserve: bool) -> Result<()> {
Ok(())
}
fn parse_local_forwards(raw: &[String]) -> Result<Vec<LocalForward>> {
raw.iter().map(|value| parse_local_forward(value)).collect()
}
fn parse_remote_forwards(raw: &[String]) -> Result<Vec<RemoteForward>> {
raw.iter()
.map(|value| parse_remote_forward(value))
.collect()
}
fn parse_dynamic_forwards(raw: &[String]) -> Result<Vec<DynamicForward>> {
raw.iter()
.map(|value| parse_dynamic_forward(value))
.collect()
}
fn parse_local_forward(raw: &str) -> Result<LocalForward> {
let parts = raw.split(':').collect::<Vec<_>>();
let (bind_host, listen, target_host, target_port) = match parts.as_slice() {
[listen, target_host, target_port] => {
("127.0.0.1".to_string(), *listen, *target_host, *target_port)
}
[bind_host, listen, target_host, target_port] => (
(*bind_host).to_string(),
*listen,
*target_host,
*target_port,
),
_ => {
return Err(anyhow!(
"invalid -L {raw:?}; expected listen_port:target_host:target_port or bind_host:listen_port:target_host:target_port"
));
}
};
let listen_port = listen
.parse::<u16>()
.with_context(|| format!("invalid local forward listen port in {raw:?}"))?;
let target_port = target_port
.parse::<u16>()
.with_context(|| format!("invalid local forward target port in {raw:?}"))?;
if !valid_forward_host(target_host) {
return Err(anyhow!("invalid -L {raw:?}; target host cannot be empty"));
}
if !valid_forward_host(&bind_host) {
return Err(anyhow!("invalid -L {raw:?}; bind host cannot be empty"));
}
anyhow::ensure!(
listen_port != 0,
"invalid -L {raw:?}; listen port cannot be 0"
);
anyhow::ensure!(
target_port != 0 || matches!(target_host, FILE_STREAM_SENTINEL | EXEC_STREAM_SENTINEL),
"invalid -L {raw:?}; target port cannot be 0"
);
Ok(LocalForward {
bind_host,
listen_port,
target_host: target_host.to_string(),
target_port,
})
}
fn parse_remote_forward(raw: &str) -> Result<RemoteForward> {
let local = parse_local_forward(raw)?;
Ok(RemoteForward {
bind_host: local.bind_host,
listen_port: local.listen_port,
target_host: local.target_host,
target_port: local.target_port,
})
}
fn parse_dynamic_forward(raw: &str) -> Result<DynamicForward> {
let parts = raw.split(':').collect::<Vec<_>>();
let (bind_host, listen) = match parts.as_slice() {
[listen] => ("127.0.0.1".to_string(), *listen),
[bind_host, listen] => ((*bind_host).to_string(), *listen),
_ => {
return Err(anyhow!(
"invalid -D {raw:?}; expected listen_port or bind_host:listen_port"
));
}
};
let listen_port = listen
.parse::<u16>()
.with_context(|| format!("invalid dynamic forward listen port in {raw:?}"))?;
if !valid_forward_host(&bind_host) {
return Err(anyhow!("invalid -D {raw:?}; bind host cannot be empty"));
}
anyhow::ensure!(
listen_port != 0,
"invalid -D {raw:?}; listen port cannot be 0"
);
Ok(DynamicForward {
bind_host,
listen_port,
})
}
fn valid_forward_host(host: &str) -> bool {
!host.is_empty()
&& !host
.bytes()
.any(|byte| byte == 0 || byte == b'\n' || byte == b'\r' || byte == b'\t')
}
fn forwarding_requests(
local_forwards: &[LocalForward],
remote_forwards: &[RemoteForward],
dynamic_forwards: &[DynamicForward],
forward_agent: bool,
) -> Vec<ForwardingRequest> {
let mut requests = local_forwards
.iter()
.map(|forward| ForwardingRequest {
kind: ForwardingKind::Local,
bind_host: Some(forward.bind_host.clone()),
listen_port: forward.listen_port,
target_host: Some(forward.target_host.clone()),
target_port: Some(forward.target_port),
})
.collect::<Vec<_>>();
requests.extend(remote_forwards.iter().map(|forward| ForwardingRequest {
kind: ForwardingKind::Remote,
bind_host: Some(forward.bind_host.clone()),
listen_port: forward.listen_port,
target_host: Some(forward.target_host.clone()),
target_port: Some(forward.target_port),
}));
requests.extend(dynamic_forwards.iter().map(|forward| ForwardingRequest {
kind: ForwardingKind::Dynamic,
bind_host: Some(forward.bind_host.clone()),
listen_port: forward.listen_port,
target_host: None,
target_port: None,
}));
if forward_agent {
requests.push(ForwardingRequest {
kind: ForwardingKind::Agent,
bind_host: None,
listen_port: 0,
target_host: None,
target_port: None,
});
}
requests
}
fn run_sessions_command(
server: &str,
ssh_port: Option<u16>,
ssh_key: Option<&std::path::Path>,
ssh_known_hosts: Option<&std::path::Path>,
) -> Result<()> {
run_remote_script(
server,
ssh_port,
ssh_key,
ssh_known_hosts,
SESSIONS_STATUS_SCRIPT,
"run remote status command",
)
}
fn run_remote_script(
server: &str,
ssh_port: Option<u16>,
ssh_key: Option<&std::path::Path>,
ssh_known_hosts: Option<&std::path::Path>,
script: &str,
context: &str,
) -> Result<()> {
let mut command = Command::new("ssh");
if let Some(ssh_port) = ssh_port {
command.arg("-p").arg(ssh_port.to_string());
}
command.arg("-T");
if let Some(key) = ssh_key {
command.arg("-i").arg(key);
}
if let Some(known_hosts) = ssh_known_hosts {
command
.arg("-o")
.arg(format!("UserKnownHostsFile={}", known_hosts.display()));
}
let remote_command = format!("sh -c {}", shell_word(script));
let status = command
.arg(server)
.arg(remote_command)
.status()
.with_context(|| context.to_string())?;
if !status.success() {
return Err(anyhow!("{context} failed with status {status}"));
}
Ok(())
}
fn update_check_requested(command: &[String]) -> Result<bool> {
match command {
[] => Ok(false),
[flag] if flag == "--check" || flag == "-n" => Ok(true),
_ => Err(anyhow!("usage: dosh update [--check]")),
}
}
fn release_artifact_name() -> &'static str {
match (std::env::consts::OS, std::env::consts::ARCH) {
("macos", "aarch64") => "dosh-macos-aarch64.tar.gz",
("macos", "x86_64") => "dosh-macos-x86_64.tar.gz",
("linux", "aarch64") => "dosh-linux-aarch64.tar.gz",
("linux", "x86_64") => "dosh-linux-x86_64.tar.gz",
("windows", "x86_64") => "dosh-windows-x86_64.zip",
("windows", "aarch64") => "dosh-windows-aarch64.zip",
_ => "dosh-unknown.tar.gz",
}
}
fn latest_release_download_url(repo: &str, artifact: &str) -> Option<String> {
let web = repo
.strip_suffix(".git")
.unwrap_or(repo)
.trim_end_matches('/');
if !web.starts_with("http://") && !web.starts_with("https://") {
return None;
}
Some(format!("{web}/releases/latest/download/{artifact}"))
}
fn release_tag_from_effective_url(web: &str, effective: &str) -> Option<String> {
let prefix = format!("{web}/releases/tag/");
let tag = effective.strip_prefix(&prefix)?;
if tag.is_empty() {
None
} else {
Some(tag.to_string())
}
}
fn release_version_from_tag(tag: &str) -> &str {
tag.strip_prefix('v').unwrap_or(tag)
}
fn update_version_status(local: &str, latest: &str) -> &'static str {
match compare_dotted_versions(latest, local) {
std::cmp::Ordering::Equal => "current",
std::cmp::Ordering::Greater => "update available",
std::cmp::Ordering::Less => "different",
}
}
fn compare_dotted_versions(left: &str, right: &str) -> std::cmp::Ordering {
let left = parse_dotted_version(left);
let right = parse_dotted_version(right);
let width = left.len().max(right.len());
for index in 0..width {
let ordering = left
.get(index)
.copied()
.unwrap_or(0)
.cmp(&right.get(index).copied().unwrap_or(0));
if ordering != std::cmp::Ordering::Equal {
return ordering;
}
}
std::cmp::Ordering::Equal
}
fn parse_dotted_version(version: &str) -> Vec<u64> {
version
.split(|byte: char| !byte.is_ascii_digit())
.filter(|part| !part.is_empty())
.map(|part| part.parse::<u64>().unwrap_or(0))
.collect()
}
fn latest_release_tag(repo: &str) -> Result<Option<String>> {
let web = repo
.strip_suffix(".git")
.unwrap_or(repo)
.trim_end_matches('/');
if !web.starts_with("http://") && !web.starts_with("https://") {
return Ok(None);
}
let output = Command::new("curl")
.arg("-fsSL")
.arg("-o")
.arg("/dev/null")
.arg("-w")
.arg("%{url_effective}")
.arg(format!("{web}/releases/latest"))
.stdin(Stdio::null())
.output()
.with_context(|| format!("resolve latest release for {web}"))?;
if !output.status.success() {
return Ok(None);
}
let effective = String::from_utf8_lossy(&output.stdout);
Ok(release_tag_from_effective_url(web, &effective))
}
fn release_tag_download_url(repo: &str, tag: &str, artifact: &str) -> Option<String> {
let web = repo
.strip_suffix(".git")
.unwrap_or(repo)
.trim_end_matches('/');
if !web.starts_with("http://") && !web.starts_with("https://") {
return None;
}
Some(format!("{web}/releases/download/{tag}/{artifact}"))
}
fn url_reachable(url: &str) -> Result<bool> {
let status = Command::new("curl")
.arg("-fsIL")
.arg(url)
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status()
.with_context(|| format!("check {url}"))?;
Ok(status.success())
}
fn run_update(config: &dosh::config::ClientConfig, check_only: bool) -> Result<()> {
let repo = config
.update_repo
.clone()
.unwrap_or_else(|| "https://git.palav.dev/Palav/dosh.git".to_string());
let raw_base = raw_base_from_repo(&repo);
let installer = format!("{raw_base}/raw/branch/main/install.sh");
let artifact = release_artifact_name();
if check_only {
let local_version = env!("CARGO_PKG_VERSION");
let latest_tag = latest_release_tag(&repo)?;
println!("dosh {local_version}");
println!("repo: {repo}");
println!("installer: {installer}");
match latest_tag.as_deref() {
Some(tag) => {
let latest_version = release_version_from_tag(tag);
println!(
"latest: {latest_version} ({})",
update_version_status(local_version, latest_version)
);
}
None => println!("latest: unknown"),
}
if let Some(latest_url) = latest_release_download_url(&repo, artifact) {
let mut status = "missing";
let mut display_url = latest_url.clone();
if let Some(tag_url) = latest_tag
.as_deref()
.and_then(|tag| release_tag_download_url(&repo, tag, artifact))
&& url_reachable(&tag_url)?
{
status = "available";
display_url = tag_url;
} else if url_reachable(&latest_url)? {
status = "available";
}
println!("prebuilt: {status} ({artifact})");
println!("prebuilt_url: {display_url}");
} else {
println!("prebuilt: unavailable for non-HTTP repo");
}
return Ok(());
}
let update_cache = dirs::cache_dir()
.unwrap_or_else(|| expand_tilde("~/.cache"))
.join("dosh")
.join("source");
let use_prebuilt = std::env::var("DOSH_USE_PREBUILT").unwrap_or_else(|_| "1".to_string());
let mut script = format!(
"curl -fsSL {} | DOSH_REPO={} DOSH_PORT={} DOSH_UPDATE_CACHE={} DOSH_UPDATE_QUIET=1 DOSH_USE_PREBUILT={} sh -s -- client",
shell_word(&installer),
shell_word(&repo),
config.update_port.unwrap_or(config.dosh_port),
shell_word(&update_cache.display().to_string()),
shell_word(&use_prebuilt)
);
if config.server != "user@example.com" {
script.push_str(&format!(" --server {}", shell_word(&config.server)));
}
if let Some(host) = &config.dosh_host {
script.push_str(&format!(" --dosh-host {}", shell_word(host)));
}
let status = Command::new("sh")
.arg("-c")
.arg(script)
.stdin(Stdio::null())
.status()
.context("run dosh update installer")?;
if !status.success() {
return Err(anyhow!("dosh update failed with status {status}"));
}
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
enum ImportAction {
Added,
Skipped,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct ImportedHost {
alias: String,
ssh: String,
dosh_host: String,
port: u16,
action: ImportAction,
}
fn run_import_ssh(config: &dosh::config::ClientConfig, aliases: &[String]) -> Result<()> {
if aliases.is_empty() {
return Err(anyhow!("usage: dosh import-ssh <ssh-host> [ssh-host...]"));
}
let imported = import_ssh_aliases(config, aliases)?;
for entry in imported {
match entry.action {
ImportAction::Added => eprintln!(
"dosh import-ssh: added [{}] ssh={} dosh_host={}",
entry.alias, entry.ssh, entry.dosh_host
),
ImportAction::Skipped => {
eprintln!(
"dosh import-ssh: [{}] already exists, skipping",
entry.alias
)
}
}
}
eprintln!(
"dosh import-ssh: wrote {}",
expand_tilde("~/.config/dosh/hosts.toml").display()
);
Ok(())
}
fn import_ssh_aliases(
config: &dosh::config::ClientConfig,
aliases: &[String],
) -> Result<Vec<ImportedHost>> {
let path = expand_tilde("~/.config/dosh/hosts.toml");
let mut raw = if path.exists() {
fs::read_to_string(&path).with_context(|| format!("read {}", path.display()))?
} else {
String::new()
};
if !raw.ends_with('\n') && !raw.is_empty() {
raw.push('\n');
}
let mut imported = Vec::new();
for alias in aliases {
let ssh_config = ssh_config(alias, None)
.with_context(|| format!("read SSH config for {alias} with ssh -G"))?;
let udp_host = ssh_config
.hostname
.clone()
.unwrap_or_else(|| ssh_destination_host(alias));
if raw_contains_host_table(&raw, alias) {
imported.push(ImportedHost {
alias: alias.clone(),
ssh: alias.clone(),
dosh_host: udp_host,
port: config.dosh_port,
action: ImportAction::Skipped,
});
continue;
}
raw.push('\n');
raw.push_str(&format!("[{}]\n", toml_bare_key_or_quoted(alias)));
raw.push_str(&format!("ssh = {}\n", toml_string(alias)));
raw.push_str(&format!("dosh_host = {}\n", toml_string(&udp_host)));
raw.push_str(&format!("port = {}\n", config.dosh_port));
if let Some(user) = &ssh_config.user {
raw.push_str(&format!("user = {}\n", toml_string(user)));
}
if let Some(port) = ssh_config.port
&& port != 22
{
raw.push_str(&format!("ssh_port = {port}\n"));
}
raw.push_str("predict = true\n");
imported.push(ImportedHost {
alias: alias.clone(),
ssh: alias.clone(),
dosh_host: udp_host,
port: config.dosh_port,
action: ImportAction::Added,
});
}
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
fs::write(&path, raw).with_context(|| format!("write {}", path.display()))?;
Ok(imported)
}
fn raw_contains_host_table(raw: &str, alias: &str) -> bool {
let quoted = toml_string(alias);
raw.lines().any(|line| {
let line = line.trim();
line == format!("[{alias}]") || line == format!("[{quoted}]")
})
}
fn toml_string(value: &str) -> String {
let escaped = value
.replace('\\', "\\\\")
.replace('"', "\\\"")
.replace('\n', "\\n");
format!("\"{escaped}\"")
}
fn toml_bare_key_or_quoted(value: &str) -> String {
if value
.chars()
.all(|c| c.is_ascii_alphanumeric() || c == '_' || c == '-')
{
value.to_string()
} else {
toml_string(value)
}
}
fn raw_base_from_repo(repo: &str) -> String {
repo.trim_end_matches(".git").to_string()
}
fn shell_word(value: &str) -> String {
format!("'{}'", value.replace('\'', "'\\''"))
}
fn ssh_config_word(value: &str) -> String {
if cfg!(windows) {
format!("\"{}\"", value.replace('"', "\\\""))
} else {
shell_word(value)
}
}
fn local_bootstrap(
session: &str,
mode: &str,
cols: u16,
rows: u16,
port: u16,
host: String,
) -> Result<BootstrapResponse> {
let mut server_config = load_server_config(None)?;
server_config.port = port;
let secret = load_or_create_server_secret(&server_config)?;
let user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
let nonce = crypto::random_12();
build_bootstrap(
&server_config,
&secret,
user,
session.to_string(),
mode.to_string(),
(cols, rows),
nonce,
host,
)
}
fn auth_allows(preference: &str, method: &str) -> bool {
preference
.split(',')
.map(str::trim)
.any(|entry| entry == "auto" || entry == method)
}
fn log_timing(verbose: u8, name: &str, elapsed: Duration) {
if verbose >= 1 {
eprintln!("dosh timing {name}={}ms", elapsed.as_millis());
}
}
fn log_debug(verbose: u8, level: u8, message: &str) {
if verbose >= level {
eprintln!("{message}");
}
}
fn select_session(requested: Option<&str>, _force_new: bool) -> String {
if let Some(session) = requested {
return session.to_string();
}
// Implicit session: this gives plain `dosh host` a fresh terminal by
// default. The server can still persist it across restarts because the
// reconnect ticket cache keeps the generated name.
protocol::generate_implicit_session_name()
}
#[allow(clippy::too_many_arguments)]
fn ssh_bootstrap(
server: &str,
ssh_port: Option<u16>,
ssh_auth_command: &str,
ssh_key: Option<&Path>,
ssh_known_hosts: Option<&Path>,
ssh_control_path: Option<&Path>,
session: &str,
mode: &str,
cols: u16,
rows: u16,
) -> Result<BootstrapResponse> {
let nonce = crypto::random_12();
let nonce_b64 = URL_SAFE_NO_PAD.encode(nonce);
let size = format!("{cols}x{rows}");
let mut command = Command::new("ssh");
if let Some(ssh_port) = ssh_port {
command.arg("-p").arg(ssh_port.to_string());
}
command.arg("-T");
if let Some(key) = ssh_key {
command.arg("-i").arg(key);
}
if let Some(known_hosts) = ssh_known_hosts {
command
.arg("-o")
.arg(format!("UserKnownHostsFile={}", known_hosts.display()));
}
if let Some(control_path) = ssh_control_path {
command.arg("-S").arg(control_path);
}
let output = command
.arg(server)
.arg(ssh_auth_command)
.arg("--protocol")
.arg("1")
.arg(format!("--nonce={nonce_b64}"))
.arg(format!("--session={session}"))
.arg(format!("--mode={mode}"))
.arg(format!("--size={size}"))
.output()
.context("run ssh dosh-auth")?;
if !output.status.success() {
return Err(anyhow!(
"ssh bootstrap failed: {}",
String::from_utf8_lossy(&output.stderr)
));
}
let raw = String::from_utf8(output.stdout)?;
decode_bootstrap(&raw)
}
fn ssh_destination_host(server: &str) -> String {
let without_user = server.rsplit_once('@').map_or(server, |(_, host)| host);
let without_path = without_user
.strip_prefix("ssh://")
.unwrap_or(without_user)
.split('/')
.next()
.unwrap_or(without_user);
if let Some(stripped) = without_path.strip_prefix('[')
&& let Some((host, _)) = stripped.split_once(']')
{
return host.to_string();
}
without_path
.split_once(':')
.map_or(without_path, |(host, _)| host)
.to_string()
}
fn selected_udp_host(
explicit: Option<&str>,
server: &str,
ssh_config: &SshConfig,
) -> Result<String> {
let ssh_host = || {
ssh_config
.hostname
.clone()
.unwrap_or_else(|| ssh_destination_host(server))
};
let Some(raw) = explicit.map(str::trim).filter(|value| !value.is_empty()) else {
return Ok(ssh_host());
};
match raw.to_ascii_lowercase().as_str() {
"ssh" | "auto" => Ok(ssh_host()),
"localhost" => Ok("127.0.0.1".to_string()),
"any" => Err(anyhow!(
"dosh_host={raw:?} is a server bind policy, not a client destination; use ssh/auto or a host/IP"
)),
_ => Ok(raw.to_string()),
}
}
fn ssh_username(server: &str) -> Option<String> {
let (user, _) = server.rsplit_once('@')?;
if user.is_empty() {
None
} else {
Some(user.to_string())
}
}
fn ssh_with_user(server: &str, user: Option<&str>) -> String {
let Some(user) = user else {
return server.to_string();
};
if user.is_empty() || ssh_username(server).is_some() || server.starts_with("ssh://") {
server.to_string()
} else {
format!("{user}@{server}")
}
}
#[derive(Debug, Clone, Default)]
struct SshConfig {
hostname: Option<String>,
port: Option<u16>,
user: Option<String>,
identity_files: Vec<String>,
user_known_hosts_file: Option<String>,
identities_only: bool,
proxy_jump: Option<String>,
proxy_command: Option<String>,
send_env: Vec<String>,
set_env: Vec<EnvVar>,
}
fn ssh_config(server: &str, ssh_port: Option<u16>) -> Result<SshConfig> {
let mut command = Command::new("ssh");
command.arg("-G");
if let Some(ssh_port) = ssh_port {
command.arg("-p").arg(ssh_port.to_string());
}
let output = command
.arg(server)
.output()
.with_context(|| format!("run ssh -G {server}"))?;
if !output.status.success() {
return Err(anyhow!(
"ssh -G failed: {}",
String::from_utf8_lossy(&output.stderr)
));
}
let raw = String::from_utf8(output.stdout)?;
Ok(parse_ssh_config(&raw))
}
fn parse_ssh_config(raw: &str) -> SshConfig {
let mut config = SshConfig::default();
for line in raw.lines() {
let line = line.trim();
let Some((key, value)) = line.split_once(char::is_whitespace) else {
continue;
};
let value = value.trim();
if key.eq_ignore_ascii_case("hostname") && !empty_or_none(value) {
config.hostname = Some(value.to_string());
} else if key.eq_ignore_ascii_case("port") {
config.port = value.parse().ok();
} else if key.eq_ignore_ascii_case("user") && !empty_or_none(value) {
config.user = Some(value.to_string());
} else if key.eq_ignore_ascii_case("identityfile") && !empty_or_none(value) {
config.identity_files.push(value.to_string());
} else if key.eq_ignore_ascii_case("userknownhostsfile") && !empty_or_none(value) {
config.user_known_hosts_file = Some(value.to_string());
} else if key.eq_ignore_ascii_case("identitiesonly") {
config.identities_only = value.eq_ignore_ascii_case("yes");
} else if key.eq_ignore_ascii_case("proxyjump") && !empty_or_none(value) {
config.proxy_jump = Some(value.to_string());
} else if key.eq_ignore_ascii_case("proxycommand") && !empty_or_none(value) {
config.proxy_command = Some(value.to_string());
} else if key.eq_ignore_ascii_case("sendenv") && !empty_or_none(value) {
config
.send_env
.extend(value.split_whitespace().map(ToString::to_string));
} else if key.eq_ignore_ascii_case("setenv") && !empty_or_none(value) {
config.set_env.extend(parse_set_env_values(value));
}
}
config
}
fn parse_set_env_values(value: &str) -> Vec<EnvVar> {
value
.split_whitespace()
.filter_map(|entry| {
let (name, value) = entry.split_once('=')?;
if valid_env_name(name) && !value.as_bytes().contains(&0) {
Some(EnvVar {
name: name.to_string(),
value: value.to_string(),
})
} else {
None
}
})
.collect()
}
fn empty_or_none(value: &str) -> bool {
value.is_empty() || value.eq_ignore_ascii_case("none")
}
fn requested_env(
config: &dosh::config::ClientConfig,
host: &dosh::config::HostConfig,
ssh_config: Option<&SshConfig>,
) -> Vec<EnvVar> {
let mut values = BTreeMap::new();
let mut patterns = host
.send_env
.clone()
.unwrap_or_else(|| config.send_env.clone());
if let Some(ssh_config) = ssh_config {
patterns.extend(ssh_config.send_env.clone());
}
for (name, value) in std::env::vars() {
if valid_env_name(&name) && patterns.iter().any(|pattern| glob_matches(pattern, &name)) {
values.insert(name, value);
}
}
for (name, value) in &config.set_env {
if valid_env_name(name) && !value.as_bytes().contains(&0) {
values.insert(name.clone(), value.clone());
}
}
for (name, value) in &host.set_env {
if valid_env_name(name) && !value.as_bytes().contains(&0) {
values.insert(name.clone(), value.clone());
}
}
if let Some(ssh_config) = ssh_config {
for env in &ssh_config.set_env {
values.insert(env.name.clone(), env.value.clone());
}
}
values
.into_iter()
.map(|(name, value)| EnvVar { name, value })
.collect()
}
fn valid_env_name(name: &str) -> bool {
let mut chars = name.chars();
let Some(first) = chars.next() else {
return false;
};
if !(first == '_' || first.is_ascii_alphabetic()) {
return false;
}
chars.all(|ch| ch == '_' || ch.is_ascii_alphanumeric())
}
fn glob_matches(pattern: &str, value: &str) -> bool {
let pattern = pattern.as_bytes();
let value = value.as_bytes();
let (mut p, mut v) = (0, 0);
let mut star = None;
let mut star_value = 0;
while v < value.len() {
if p < pattern.len() && (pattern[p] == b'?' || pattern[p] == value[v]) {
p += 1;
v += 1;
} else if p < pattern.len() && pattern[p] == b'*' {
star = Some(p);
p += 1;
star_value = v;
} else if let Some(star_pos) = star {
p = star_pos + 1;
star_value += 1;
v = star_value;
} else {
return false;
}
}
while p < pattern.len() && pattern[p] == b'*' {
p += 1;
}
p == pattern.len()
}
fn resolve_addr(host: &str, port: u16) -> Result<SocketAddr> {
(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}"))
}
async fn send_terminal_udp(socket: &UdpSocket, packet: &[u8], addr: SocketAddr) -> Result<bool> {
match socket.send_to(packet, addr).await {
Ok(_) => Ok(true),
Err(err) if is_transient_udp_send_error(&err) => Ok(false),
Err(err) => Err(err).with_context(|| format!("send UDP packet to {addr}")),
}
}
#[allow(clippy::too_many_arguments)]
async fn try_native_auth(
socket: &UdpSocket,
config: &dosh::config::ClientConfig,
requested_server: &str,
server: &str,
ssh_port: Option<u16>,
udp_host: &str,
udp_port: u16,
cli_identity: Option<&Path>,
session: &str,
mode: &str,
cols: u16,
rows: u16,
requested_forwardings: Vec<ForwardingRequest>,
ssh_config_hint: Option<&SshConfig>,
requested_env: Vec<EnvVar>,
allow_passphrase_prompt: bool,
) -> Result<(Frame, CachedCredential)> {
let addr = resolve_addr(udp_host, udp_port)?;
let owned_ssh_config;
let ssh_config = if let Some(ssh_config) = ssh_config_hint {
ssh_config
} else {
owned_ssh_config = ssh_config(server, ssh_port).unwrap_or_default();
&owned_ssh_config
};
let requested_user = ssh_username(server)
.or(ssh_config.user.clone())
.unwrap_or_else(|| std::env::var("USER").unwrap_or_else(|_| "unknown".to_string()));
let (client_secret, client_public) = generate_native_ephemeral();
let hello = dosh::native::NativeClientHello {
protocol_version: dosh::native::NATIVE_PROTOCOL_VERSION,
client_random: crypto::random_32(),
client_ephemeral_public: client_public,
requested_host: requested_server.to_string(),
requested_user,
requested_session: session.to_string(),
requested_mode: mode.to_string(),
terminal_size: (cols, rows),
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,
};
let packet = protocol::encode_plain(
PacketKind::NativeClientHello,
[0u8; 16],
1,
0,
&protocol::to_body(&NativeClientHelloBody {
hello: hello.clone(),
})?,
)?;
socket.send_to(&packet, addr).await?;
let mut buf = vec![0u8; 65535];
let (n, _) = tokio::time::timeout(
Duration::from_millis(config.native_auth_timeout_ms.max(1)),
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)?;
return Err(anyhow!("native auth rejected: {}", reject.reason));
}
return Err(anyhow!("native auth received unexpected server response"));
}
let server_hello: NativeServerHelloBody = protocol::from_body(&packet.body)?;
verify_server_hello(&hello, &server_hello.hello)?;
let known_hosts = expand_tilde(&config.known_hosts);
match verify_known_host(&known_hosts, requested_server, &server_hello.hello.host_key)? {
KnownHostStatus::Trusted => {}
KnownHostStatus::Unknown if config.trust_on_first_use => {
trust_host(
&known_hosts,
requested_server,
&server_hello.hello.host_key,
"tofu",
false,
)?;
}
KnownHostStatus::Unknown => {
return Err(anyhow!(
"Dosh host key for {requested_server} is not trusted; run `dosh trust {requested_server}` first"
));
}
KnownHostStatus::Mismatch { expected, actual } => {
return Err(anyhow!(
"Dosh host key mismatch for {requested_server}: expected {expected}, got {actual}"
));
}
}
let session_key = derive_native_session_key(
&client_secret,
server_hello.hello.server_ephemeral_public,
&hello,
&server_hello.hello,
)?;
let auth = sign_native_user_auth(
config,
cli_identity,
ssh_config,
&hello,
&server_hello.hello,
requested_forwardings,
allow_passphrase_prompt,
)?;
let mut pending_id = [0u8; 16];
pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]);
let auth_body = protocol::to_body(&NativeUserAuthBody { auth })?;
let packet = protocol::encode_encrypted(
PacketKind::NativeUserAuth,
pending_id,
2,
1,
&session_key,
CLIENT_TO_SERVER,
&auth_body,
)?;
socket.send_to(&packet, addr).await?;
let (n, _) = tokio::time::timeout(
Duration::from_millis(config.native_auth_timeout_ms.max(1)),
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)?;
return Err(anyhow!("native auth rejected: {}", reject.reason));
}
return Err(anyhow!("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 frame = Frame {
session: ok.ok.session.clone(),
output_seq: ok.ok.initial_seq,
bytes: ok.ok.snapshot.clone(),
snapshot: true,
closed: false,
};
let cred = CachedCredential {
server: requested_server.to_string(),
session: ok.ok.session,
mode: ok.ok.mode,
udp_host: udp_host.to_string(),
udp_port,
client_id: ok.ok.client_id,
session_key: ok.ok.session_key,
session_key_id: ok.ok.session_key_id,
attach_ticket: ok.ok.attach_ticket,
attach_ticket_psk: ok.ok.attach_ticket_psk,
last_rendered_seq: ok.ok.initial_seq,
};
Ok((frame, cred))
}
#[allow(clippy::too_many_arguments)]
async fn try_native_auth_check(
socket: &UdpSocket,
config: &dosh::config::ClientConfig,
requested_server: &str,
server: &str,
ssh_port: Option<u16>,
udp_host: &str,
udp_port: u16,
cli_identity: Option<&Path>,
ssh_config_hint: Option<&SshConfig>,
) -> Result<NativeAuthCheckOkBody> {
let addr = resolve_addr(udp_host, udp_port)?;
let owned_ssh_config;
let ssh_config = if let Some(ssh_config) = ssh_config_hint {
ssh_config
} else {
owned_ssh_config = ssh_config(server, ssh_port).unwrap_or_default();
&owned_ssh_config
};
let requested_user = ssh_username(server)
.or(ssh_config.user.clone())
.unwrap_or_else(|| std::env::var("USER").unwrap_or_else(|_| "unknown".to_string()));
let (client_secret, client_public) = generate_native_ephemeral();
let hello = dosh::native::NativeClientHello {
protocol_version: dosh::native::NATIVE_PROTOCOL_VERSION,
client_random: crypto::random_32(),
client_ephemeral_public: client_public,
requested_host: requested_server.to_string(),
requested_user,
requested_session: "doctor".to_string(),
requested_mode: "doctor".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: Vec::new(),
};
let packet = protocol::encode_plain(
PacketKind::NativeClientHello,
[0u8; 16],
1,
0,
&protocol::to_body(&NativeClientHelloBody {
hello: hello.clone(),
})?,
)?;
socket.send_to(&packet, addr).await?;
let mut buf = vec![0u8; 65535];
let (n, _) = tokio::time::timeout(
Duration::from_millis(config.native_auth_timeout_ms.max(1)),
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)?;
return Err(anyhow!("native doctor rejected: {}", reject.reason));
}
return Err(anyhow!("native doctor received unexpected server response"));
}
let server_hello: NativeServerHelloBody = protocol::from_body(&packet.body)?;
verify_server_hello(&hello, &server_hello.hello)?;
let known_hosts = expand_tilde(&config.known_hosts);
match verify_known_host(&known_hosts, requested_server, &server_hello.hello.host_key)? {
KnownHostStatus::Trusted => {}
KnownHostStatus::Unknown if config.trust_on_first_use => {
trust_host(
&known_hosts,
requested_server,
&server_hello.hello.host_key,
"tofu",
false,
)?;
}
KnownHostStatus::Unknown => {
return Err(anyhow!(
"Dosh host key for {requested_server} is not trusted; run `dosh trust {requested_server}` first"
));
}
KnownHostStatus::Mismatch { expected, actual } => {
return Err(anyhow!(
"Dosh host key mismatch for {requested_server}: expected {expected}, got {actual}"
));
}
}
let session_key = derive_native_session_key(
&client_secret,
server_hello.hello.server_ephemeral_public,
&hello,
&server_hello.hello,
)?;
let auth = sign_native_user_auth(
config,
cli_identity,
ssh_config,
&hello,
&server_hello.hello,
Vec::new(),
true,
)?;
let mut pending_id = [0u8; 16];
pending_id.copy_from_slice(&server_hello.hello.auth_challenge[..16]);
let packet = protocol::encode_encrypted(
PacketKind::NativeUserAuth,
pending_id,
2,
1,
&session_key,
CLIENT_TO_SERVER,
&protocol::to_body(&NativeUserAuthBody { auth })?,
)?;
socket.send_to(&packet, addr).await?;
let (n, _) = tokio::time::timeout(
Duration::from_millis(config.native_auth_timeout_ms.max(1)),
socket.recv_from(&mut buf),
)
.await??;
let packet = protocol::decode(&buf[..n])?;
if packet.header.kind != PacketKind::NativeAuthCheckOk {
if packet.header.kind == PacketKind::AttachReject {
let reject: AttachReject = protocol::from_body(&packet.body)?;
return Err(anyhow!("native doctor rejected: {}", reject.reason));
}
return Err(anyhow!("native doctor received unexpected auth response"));
}
let plain = protocol::decrypt_body(&packet, &session_key, SERVER_TO_CLIENT)?;
protocol::from_body(&plain)
}
fn sign_native_user_auth(
config: &dosh::config::ClientConfig,
cli_identity: Option<&Path>,
ssh_config: &SshConfig,
hello: &dosh::native::NativeClientHello,
server_hello: &dosh::native::NativeServerHello,
requested_forwardings: Vec<ForwardingRequest>,
allow_passphrase_prompt: bool,
) -> Result<dosh::native::NativeUserAuth> {
let mut errors = Vec::new();
if config.use_ssh_agent && !ssh_config.identities_only {
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:#}")),
}
} else if config.use_ssh_agent && ssh_config.identities_only {
errors.push("ssh-agent: skipped because SSH config sets IdentitiesOnly=yes".to_string());
}
let identity = if allow_passphrase_prompt {
load_first_native_identity(config, cli_identity, ssh_config)
} else {
load_first_native_identity_noninteractive(config, cli_identity, ssh_config)
};
match identity {
Ok(identity) => {
sign_user_auth_with_private_key(&identity, hello, server_hello, requested_forwardings)
}
Err(err) => {
errors.push(format!("identity files: {err:#}"));
Err(anyhow!(
"native auth found no usable key: {}",
errors.join("; ")
))
}
}
}
fn load_first_native_identity(
config: &dosh::config::ClientConfig,
cli_identity: Option<&Path>,
ssh_config: &SshConfig,
) -> Result<ssh_key::PrivateKey> {
load_first_native_identity_with_prompt(
config,
cli_identity,
ssh_config,
prompt_identity_passphrase,
)
}
fn load_first_native_identity_noninteractive(
config: &dosh::config::ClientConfig,
cli_identity: Option<&Path>,
ssh_config: &SshConfig,
) -> Result<ssh_key::PrivateKey> {
load_first_native_identity_with_prompt(config, cli_identity, ssh_config, |path| {
Err(anyhow!(
"{} is encrypted; unlock it in ssh-agent or use an unencrypted key for proxy-stdio",
path.display()
))
})
}
fn load_first_native_identity_with_prompt<F>(
config: &dosh::config::ClientConfig,
cli_identity: Option<&Path>,
ssh_config: &SshConfig,
mut prompt: F,
) -> Result<ssh_key::PrivateKey>
where
F: FnMut(&Path) -> Result<String>,
{
let mut paths = Vec::new();
if let Some(path) = cli_identity {
push_identity_path(&mut paths, path.to_path_buf());
}
for path in &ssh_config.identity_files {
push_identity_path(&mut paths, expand_tilde(path));
}
if !ssh_config.identities_only {
for path in &config.identity_files {
push_identity_path(&mut paths, expand_tilde(path));
}
}
let mut errors = Vec::new();
for path in paths {
if !path.exists() {
continue;
}
match load_native_identity(&path) {
Ok(identity) => return Ok(identity),
Err(err) if err.to_string().contains(" is encrypted") => {
match prompt(&path).and_then(|passphrase| {
load_native_identity_with_passphrase(&path, Some(&passphrase))
}) {
Ok(identity) => return Ok(identity),
Err(err) => errors.push(format!("{}: {err:#}", path.display())),
}
}
Err(err) => errors.push(format!("{}: {err:#}", path.display())),
}
}
if errors.is_empty() {
Err(anyhow!("no usable native identity found"))
} else {
Err(anyhow!(
"no usable native identity found: {}",
errors.join("; ")
))
}
}
fn push_identity_path(paths: &mut Vec<PathBuf>, path: PathBuf) {
if !paths.iter().any(|existing| existing == &path) {
paths.push(path);
}
}
fn prompt_identity_passphrase(path: &Path) -> Result<String> {
rpassword::prompt_password(format!("Enter passphrase for {}: ", path.display()))
.with_context(|| format!("read passphrase for {}", path.display()))
}
async fn bootstrap_attach(
socket: &UdpSocket,
server_name: &str,
bootstrap: &BootstrapResponse,
cols: u16,
rows: u16,
requested_env: Vec<EnvVar>,
) -> Result<(AttachOk, CachedCredential)> {
let addr = resolve_addr(&bootstrap.udp_host, bootstrap.udp_port)?;
let req = BootstrapAttachRequest {
bootstrap: bootstrap.clone(),
cols,
rows,
requested_env,
};
let body = protocol::to_body(&req)?;
let packet =
protocol::encode_plain(PacketKind::BootstrapAttachRequest, [0u8; 16], 1, 0, &body)?;
socket.send_to(&packet, addr).await?;
let expected_key_id = protocol::session_key_id(&bootstrap.session_key);
let ok = recv_response_until(socket, Duration::from_secs(5), |packet| {
if packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16] {
let reject: AttachReject = protocol::from_body(&packet.body)?;
return Err(anyhow!("attach rejected: {}", reject.reason));
}
if packet.header.kind != PacketKind::AttachOk
|| packet.header.session_key_id != expected_key_id
{
return Ok(None);
}
let Ok(plain) = protocol::decrypt_body(&packet, &bootstrap.session_key, SERVER_TO_CLIENT)
else {
return Ok(None);
};
Ok(protocol::from_body::<AttachOk>(&plain).ok())
})
.await
.with_context(|| {
format!(
"no Dosh bootstrap response from {}:{}; check UDP reachability and server firewall",
bootstrap.udp_host, bootstrap.udp_port
)
})?;
let cred = CachedCredential {
server: server_name.to_string(),
session: ok.session.clone(),
mode: ok.mode.clone(),
udp_host: bootstrap.udp_host.clone(),
udp_port: bootstrap.udp_port,
client_id: ok.client_id,
session_key: ok.session_key,
session_key_id: ok.session_key_id,
attach_ticket: bootstrap.attach_ticket.clone(),
attach_ticket_psk: bootstrap.attach_ticket_psk,
last_rendered_seq: ok.initial_seq,
};
Ok((ok, cred))
}
async fn try_ticket_attach(
socket: &UdpSocket,
cached: &CachedCredential,
cols: u16,
rows: u16,
requested_env: Vec<EnvVar>,
) -> Result<(Frame, CachedCredential)> {
let addr = resolve_addr(&cached.udp_host, cached.udp_port)?;
let client_nonce = crypto::random_12();
let request_key = crypto::hkdf32(
&cached.attach_ticket_psk,
&client_nonce,
b"dosh/ticket-attach-request/v1",
)?;
let body = protocol::to_body(&TicketAttachBody {
session: cached.session.clone(),
mode: cached.mode.clone(),
cols,
rows,
requested_env,
})?;
let ciphertext = crypto::seal(
&request_key,
&client_nonce,
b"dosh-ticket-attach-request-v1",
&body,
)?;
let envelope = TicketAttachEnvelope {
ticket: cached.attach_ticket.clone(),
client_nonce,
ciphertext,
};
let packet = protocol::encode_plain(
PacketKind::TicketAttachRequest,
[0u8; 16],
1,
0,
&protocol::to_body(&envelope)?,
)?;
socket.send_to(&packet, addr).await?;
let ok = recv_response_until(socket, Duration::from_millis(700), |packet| {
if packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16] {
let reject: AttachReject = protocol::from_body(&packet.body)?;
return Err(anyhow!("ticket attach rejected: {}", reject.reason));
}
if packet.header.kind != PacketKind::AttachOk {
return Ok(None);
}
let Ok(envelope) = protocol::from_body::<TicketAttachOkEnvelope>(&packet.body) else {
return Ok(None);
};
let mut salt = Vec::with_capacity(24);
salt.extend_from_slice(&client_nonce);
salt.extend_from_slice(&envelope.server_nonce);
let response_key = crypto::hkdf32(
&cached.attach_ticket_psk,
&salt,
b"dosh/ticket-attach-ok/v1",
)?;
let Ok(plain) = crypto::open(
&response_key,
&envelope.server_nonce,
b"dosh-ticket-attach-ok-v1",
&envelope.ciphertext,
) else {
return Ok(None);
};
Ok(protocol::from_body::<AttachOk>(&plain).ok())
})
.await
.with_context(|| {
format!(
"cached ticket attach to {}:{} timed out; check UDP reachability or run `dosh recover {}`",
cached.udp_host, cached.udp_port, cached.server
)
})?;
let frame = Frame {
session: ok.session.clone(),
output_seq: ok.initial_seq,
bytes: ok.snapshot.clone(),
snapshot: true,
closed: false,
};
let mut next = cached.clone();
next.client_id = ok.client_id;
next.session_key = ok.session_key;
next.session_key_id = ok.session_key_id;
next.last_rendered_seq = ok.initial_seq;
Ok((frame, next))
}
async fn try_resume_fresh_process(
socket: &UdpSocket,
cached: &CachedCredential,
cols: u16,
rows: u16,
) -> Result<(Frame, CachedCredential)> {
let addr = resolve_addr(&cached.udp_host, cached.udp_port)?;
let req = ResumeRequest {
session: cached.session.clone(),
last_rendered_seq: cached.last_rendered_seq,
cols,
rows,
};
let body = protocol::to_body(&req)?;
let packet = protocol::encode_encrypted(
PacketKind::ResumeRequest,
cached.client_id,
1,
0,
&cached.session_key,
CLIENT_TO_SERVER,
&body,
)?;
socket.send_to(&packet, addr).await?;
let frame = recv_response_until(socket, Duration::from_millis(700), |packet| {
if !is_resume_response_for_client(&packet, cached.client_id) {
return Ok(None);
}
if packet.header.kind == PacketKind::AttachReject {
let reject: AttachReject = protocol::from_body(&packet.body)?;
return Err(anyhow!("resume rejected: {}", reject.reason));
}
if packet.header.kind != PacketKind::ResumeOk {
return Ok(None);
}
let Ok(plain) = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT)
else {
return Ok(None);
};
Ok(protocol::from_body::<Frame>(&plain).ok())
})
.await
.with_context(|| {
format!(
"cached resume to {}:{} timed out; check UDP reachability or run `dosh recover {}`",
cached.udp_host, cached.udp_port, cached.server
)
})?;
let mut next = cached.clone();
next.last_rendered_seq = frame.output_seq;
Ok((frame, next))
}
async fn try_live_resume(
socket: &UdpSocket,
cached: &CachedCredential,
send_seq: &mut u64,
cols: u16,
rows: u16,
) -> Result<(Frame, CachedCredential)> {
let addr = resolve_addr(&cached.udp_host, cached.udp_port)?;
let req = ResumeRequest {
session: cached.session.clone(),
last_rendered_seq: cached.last_rendered_seq,
cols,
rows,
};
let body = protocol::to_body(&req)?;
let packet = protocol::encode_encrypted(
PacketKind::ResumeRequest,
cached.client_id,
*send_seq,
cached.last_rendered_seq,
&cached.session_key,
CLIENT_TO_SERVER,
&body,
)?;
*send_seq += 1;
socket.send_to(&packet, addr).await?;
let frame = recv_response_until(socket, Duration::from_millis(700), |packet| {
if !is_resume_response_for_client(&packet, cached.client_id) {
return Ok(None);
}
if packet.header.kind == PacketKind::AttachReject {
let reject: AttachReject = protocol::from_body(&packet.body)?;
return Err(anyhow!("resume rejected: {}", reject.reason));
}
if packet.header.kind != PacketKind::ResumeOk {
return Ok(None);
}
let Ok(plain) = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT)
else {
return Ok(None);
};
Ok(protocol::from_body::<Frame>(&plain).ok())
})
.await
.with_context(|| {
format!(
"live resume to {}:{} timed out; still waiting for the server to become reachable",
cached.udp_host, cached.udp_port
)
})?;
let mut next = cached.clone();
next.last_rendered_seq = frame.output_seq;
Ok((frame, next))
}
fn is_resume_response_for_client(packet: &protocol::Packet, client_id: [u8; 16]) -> bool {
packet.header.conn_id == client_id
|| (packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16])
}
async fn recv_response_until<T, F>(socket: &UdpSocket, wait: Duration, mut accept: F) -> Result<T>
where
F: FnMut(protocol::Packet) -> Result<Option<T>>,
{
let started = Instant::now();
let mut buf = vec![0u8; 65535];
loop {
let elapsed = started.elapsed();
if elapsed >= wait {
return Err(anyhow!("timed out waiting for server response"));
}
let remaining = wait - elapsed;
let (n, _) = tokio::time::timeout(remaining, socket.recv_from(&mut buf))
.await
.context("timed out waiting for server response")??;
let Ok(packet) = protocol::decode(&buf[..n]) else {
continue;
};
if let Some(response) = accept(packet)? {
return Ok(response);
}
}
}
#[allow(clippy::too_many_arguments)]
async fn run_terminal(
socket: UdpSocket,
mut cred: CachedCredential,
first_frame: Option<Frame>,
predict: bool,
predict_mode: PredictMode,
escape_key: Option<Vec<u8>>,
startup_command: Option<Vec<u8>>,
reconnect_timeout_secs: u64,
local_forwards: Vec<LocalForward>,
dynamic_forwards: Vec<DynamicForward>,
forward_only: bool,
agent_sock: Option<PathBuf>,
) -> Result<()> {
#[cfg(not(unix))]
let _ = &agent_sock;
let _raw = if forward_only {
None
} else {
Some(RawMode::enter()?)
};
let mut addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
let mut send_seq = 2u64;
let mut last_packet_at = Instant::now();
let mut status_tick = tokio::time::interval(Duration::from_secs(1));
let mut resize_tick = tokio::time::interval(Duration::from_millis(250));
let mut stream_retransmit_tick = tokio::time::interval(Duration::from_millis(200));
let mut frame_gap_tick = tokio::time::interval(Duration::from_millis(250));
let mut last_size = terminal_size();
// React to terminal resize the instant it happens via SIGWINCH (mosh-style),
// instead of waiting up to one `resize_tick`. The 250ms poll below stays as a
// fallback for environments where the signal doesn't fire. `signal()` can fail
// (rare), in which case we rely on the poll.
#[cfg(unix)]
let mut winch =
tokio::signal::unix::signal(tokio::signal::unix::SignalKind::window_change()).ok();
let mut frame_buffer = FrameBuffer::default();
// Resolve the prediction display policy (off / experimental / always). An
// env var wins for ad-hoc tuning; otherwise the client config's
// `predict_mode` provides the persistent default. Predictions only run in a
// read-write interactive session.
let mut predictor = Predictor::with_mode(
predict && cred.mode != "view-only" && !forward_only,
predict_mode,
);
// Non-destructive disconnect status line. Off in forward-only mode (no TTY to
// draw on) and respecting the client config (default on, env override).
let mut disconnect_status = DisconnectStatus::new(resolve_disconnect_status() && !forward_only);
let (forward_tx, mut forward_rx) = mpsc::channel::<ForwardEvent>(1024);
let _forward_keepalive = if local_forwards.is_empty() {
Some(forward_tx.clone())
} else {
None
};
let remote_forward_tx = forward_tx.clone();
let stream_ids = Arc::new(AtomicU64::new(1));
start_local_forwards(&local_forwards, forward_tx.clone(), Arc::clone(&stream_ids)).await?;
start_dynamic_forwards(&dynamic_forwards, forward_tx.clone(), stream_ids).await?;
signal_background_ready();
let mut stream_writers: HashMap<u64, tokio::net::tcp::OwnedWriteHalf> = HashMap::new();
// Write halves for server-initiated SSH-agent streams, spliced into the local
// unix agent socket. Kept separate from the TCP `stream_writers` so the
// existing TCP forwarding paths are untouched.
#[cfg(unix)]
let mut agent_writers: HashMap<u64, tokio::net::unix::OwnedWriteHalf> = HashMap::new();
let mut pending_socks_replies: HashSet<u64> = HashSet::new();
let mut opened_streams: HashSet<u64> = HashSet::new();
let mut stream_send_credit: HashMap<u64, usize> = HashMap::new();
let mut stream_pending_data: HashMap<u64, VecDeque<Vec<u8>>> = HashMap::new();
let mut stream_pending_opens: HashMap<u64, PendingStreamOpen> = HashMap::new();
let mut stream_next_send_offset: HashMap<u64, u64> = HashMap::new();
let mut stream_sent_data: HashMap<u64, BTreeMap<u64, PendingStreamChunk>> = HashMap::new();
let mut stream_next_recv_offset: HashMap<u64, u64> = HashMap::new();
let mut stream_recv_pending: HashMap<u64, BTreeMap<u64, Vec<u8>>> = HashMap::new();
let mut stream_retired: HashSet<u64> = HashSet::new();
let mut stream_retired_order: VecDeque<u64> = VecDeque::new();
let mut pending_user_input: VecDeque<PendingUserInput> = VecDeque::new();
let mut pending_user_input_bytes = 0usize;
let mut startup_input_hold_until: Option<Instant> = None;
let mut startup_gate_mode = StartupGateMode::HoldControl;
let mut stale_terminal_input_suppress_until: Option<Instant> = None;
let mut last_terminal_frame_at = Instant::now();
let mut last_focus_repaint_at = Instant::now() - FOCUS_REPAINT_COOLDOWN;
let mut last_idle_repaint_attempt_at = Instant::now() - ALT_SCREEN_IDLE_REPAINT_AFTER;
if let Some(frame) = first_frame {
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
last_terminal_frame_at = Instant::now();
}
if frame.closed {
return Ok(());
}
cred.last_rendered_seq = frame.output_seq;
send_ack(&socket, addr, &cred, &mut send_seq).await?;
}
if let Some(bytes) = startup_command
&& cred.mode != "view-only"
&& !forward_only
{
if !send_input(&socket, addr, &cred, &mut send_seq, bytes.clone()).await? {
queue_stale_pending_user_input(
&mut pending_user_input,
&mut pending_user_input_bytes,
bytes,
)?;
}
startup_input_hold_until = Some(Instant::now() + STARTUP_INPUT_HOLD);
startup_gate_mode = StartupGateMode::HoldAll;
}
let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::<Vec<u8>>();
let _stdin_keepalive = if forward_only {
Some(stdin_tx)
} else {
std::thread::Builder::new()
.name("dosh-stdin".to_string())
.spawn(move || {
let mut stdin = std::io::stdin();
let mut buf = [0u8; 4096];
loop {
match stdin.read(&mut buf) {
Ok(0) => break,
Ok(n) => {
let _ = stdin_tx.send(buf[..n].to_vec());
}
Err(_) => break,
}
}
})?;
None
};
// Retain the previous epoch's key briefly after a rekey so any in-flight
// pre-rekey frame still decrypts instead of triggering a needless reconnect.
let mut previous_session_key: Option<[u8; 32]> = None;
let mut recv_buf = vec![0u8; 65535];
loop {
tokio::select! {
biased;
stdin_msg = stdin_rx.recv() => {
match stdin_msg {
Some(mut bytes) => {
if input_matches_escape(&bytes, escape_key.as_deref()) {
break;
}
let saw_focus_in = input_contains_focus_in(&bytes);
if !forward_only
&& saw_focus_in
&& last_focus_repaint_at.elapsed() >= FOCUS_REPAINT_COOLDOWN
{
last_focus_repaint_at = Instant::now();
if let Some(frame) = reconnect(
&socket,
&mut cred,
&mut send_seq,
last_size,
&mut frame_buffer,
&mut predictor,
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
arm_stale_terminal_input_suppression(
&mut stale_terminal_input_suppress_until,
);
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
last_terminal_frame_at = Instant::now();
last_packet_at = Instant::now();
flush_pending_user_input(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
&mut pending_user_input,
&mut pending_user_input_bytes,
)
.await?;
}
}
bytes = strip_terminal_focus_reports(&bytes);
if bytes.is_empty() {
continue;
}
if should_strip_stale_terminal_reports(
last_packet_at.elapsed(),
stale_terminal_input_suppress_until,
) {
bytes = strip_stale_mouse_reports(&bytes);
if bytes.is_empty() {
continue;
}
}
if cred.mode != "view-only" && cred.mode != "forward-only" {
if let Some(deadline) = startup_input_hold_until {
if !predictor.alternate_screen
&& Instant::now() < deadline
&& should_hold_during_startup_gate(
startup_gate_mode,
!pending_user_input.is_empty(),
&bytes,
)
{
queue_pending_user_input(
&mut pending_user_input,
&mut pending_user_input_bytes,
bytes,
)?;
continue;
}
startup_input_hold_until = None;
startup_gate_mode = StartupGateMode::HoldControl;
flush_pending_user_input(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
&mut pending_user_input,
&mut pending_user_input_bytes,
)
.await?;
}
if !predictor.alternate_screen
&& let Some((send_now, mut hold_for_later)) =
split_after_command_submit(&bytes)
{
if send_input(&socket, addr, &cred, &mut send_seq, send_now.clone()).await? {
predictor.observe_input(&send_now)?;
} else {
queue_stale_pending_user_input(
&mut pending_user_input,
&mut pending_user_input_bytes,
send_now,
)?;
}
startup_input_hold_until = Some(
Instant::now()
+ post_submit_hold_duration(&hold_for_later),
);
startup_gate_mode = if should_hold_post_submit_input(&hold_for_later)
{
StartupGateMode::HoldControl
} else {
StartupGateMode::HoldAll
};
if should_hold_post_submit_input(&hold_for_later) {
queue_pending_user_input(
&mut pending_user_input,
&mut pending_user_input_bytes,
hold_for_later,
)?;
continue;
} else if !hold_for_later.is_empty() {
let sent = send_input(
&socket,
addr,
&cred,
&mut send_seq,
hold_for_later.clone(),
)
.await?;
if sent {
predictor.observe_input(&hold_for_later)?;
} else {
queue_stale_pending_user_input(
&mut pending_user_input,
&mut pending_user_input_bytes,
std::mem::take(&mut hold_for_later),
)?;
}
}
continue;
}
if last_packet_at.elapsed() >= Duration::from_secs(2) {
queue_stale_pending_user_input(
&mut pending_user_input,
&mut pending_user_input_bytes,
bytes,
)?;
if let Some(frame) = reconnect(
&socket,
&mut cred,
&mut send_seq,
last_size,
&mut frame_buffer,
&mut predictor,
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
arm_stale_terminal_input_suppression(
&mut stale_terminal_input_suppress_until,
);
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
last_terminal_frame_at = Instant::now();
}
last_packet_at = Instant::now();
flush_pending_user_input(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
&mut pending_user_input,
&mut pending_user_input_bytes,
)
.await?;
}
} else {
if send_input(&socket, addr, &cred, &mut send_seq, bytes.clone()).await? {
predictor.observe_input(&bytes)?;
} else {
queue_stale_pending_user_input(
&mut pending_user_input,
&mut pending_user_input_bytes,
bytes,
)?;
}
}
}
}
None => break,
}
}
_ = async {
#[cfg(unix)]
match winch.as_mut() {
Some(w) => { w.recv().await; }
None => std::future::pending::<()>().await,
}
#[cfg(not(unix))]
std::future::pending::<()>().await
} => {
maybe_send_resize(&socket, addr, &cred, &mut send_seq, &mut last_size).await?;
}
_ = resize_tick.tick() => {
maybe_send_resize(&socket, addr, &cred, &mut send_seq, &mut last_size).await?;
}
_ = frame_gap_tick.tick() => {
if frame_buffer.resync_due()
&& let Some(frame) = reconnect(
&socket,
&mut cred,
&mut send_seq,
last_size,
&mut frame_buffer,
&mut predictor,
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
arm_stale_terminal_input_suppression(
&mut stale_terminal_input_suppress_until,
);
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
last_terminal_frame_at = Instant::now();
}
last_packet_at = Instant::now();
flush_pending_user_input(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
&mut pending_user_input,
&mut pending_user_input_bytes,
)
.await?;
}
}
recv = socket.recv_from(&mut recv_buf) => {
let (n, _) = recv?;
let Ok(packet) = protocol::decode(&recv_buf[..n]) else {
continue;
};
if packet.header.conn_id != [0u8; 16] && packet.header.conn_id != cred.client_id {
continue;
}
// Contact resumed: clear any disconnect status line immediately.
disconnect_status.apply(disconnect_status.on_contact())?;
match packet.header.kind {
PacketKind::Frame | PacketKind::ResumeOk => {
let decrypted = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT)
.or_else(|err| {
// Fall back to the previous epoch key for in-flight
// pre-rekey frames before declaring the link stale.
match previous_session_key {
Some(prev) => protocol::decrypt_body(&packet, &prev, SERVER_TO_CLIENT),
None => Err(err),
}
});
let Ok(plain) = decrypted else {
if let Some(frame) = reconnect(
&socket,
&mut cred,
&mut send_seq,
last_size,
&mut frame_buffer,
&mut predictor,
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
arm_stale_terminal_input_suppression(
&mut stale_terminal_input_suppress_until,
);
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
last_terminal_frame_at = Instant::now();
}
last_packet_at = Instant::now();
flush_pending_user_input(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
&mut pending_user_input,
&mut pending_user_input_bytes,
)
.await?;
}
continue;
};
let Ok(frame) = protocol::from_body::<Frame>(&plain) else {
continue;
};
last_packet_at = Instant::now();
let frames = frame_buffer.accept(frame, &mut cred.last_rendered_seq);
for frame in frames {
predictor.clear_pending()?;
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
last_terminal_frame_at = Instant::now();
flush_startup_input_if_ready(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
(&mut pending_user_input, &mut pending_user_input_bytes),
(&mut startup_input_hold_until, &mut startup_gate_mode),
)
.await?;
}
if frame.closed {
send_ack(&socket, addr, &cred, &mut send_seq).await?;
return Ok(());
}
}
send_ack(&socket, addr, &cred, &mut send_seq).await?;
}
PacketKind::Pong => {
last_packet_at = Instant::now();
if should_flush_terminal_input_after_contact(
forward_only,
&cred.mode,
pending_user_input.is_empty(),
) {
flush_pending_user_input(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
&mut pending_user_input,
&mut pending_user_input_bytes,
)
.await?;
}
}
PacketKind::Rekey => {
// Server-initiated transport rekey (spec §11). The Rekey is
// sealed under the current key; once decrypted we derive the
// next key from the shipped fresh material + current key,
// switch atomically (keeping the old key for grace), and
// confirm with a RekeyAck under the NEW key.
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
continue;
};
let Ok(rekey) = protocol::from_body::<Rekey>(&plain) else {
continue;
};
let previous_key = cred.session_key;
let previous_key_id = cred.session_key_id;
let Ok(new_key) = dosh::native::derive_rekey_session_key(
&previous_key,
&rekey.rekey_material,
&previous_key_id,
rekey.epoch,
) else {
continue;
};
// Only accept if our derivation matches the server's id.
if protocol::session_key_id(&new_key) != rekey.new_session_key_id {
continue;
}
previous_session_key = Some(previous_key);
cred.session_key = new_key;
cred.session_key_id = rekey.new_session_key_id;
last_packet_at = Instant::now();
send_seq += 1;
let ack = protocol::encode_encrypted(
PacketKind::RekeyAck,
cred.client_id,
send_seq,
cred.last_rendered_seq,
&cred.session_key,
CLIENT_TO_SERVER,
b"",
)?;
let _ = send_terminal_udp(&socket, &ack, addr).await?;
}
PacketKind::AttachReject => {
let reject: AttachReject = protocol::from_body(&packet.body)?;
if reject.reason == "unknown client" {
if let Some(frame) = reconnect(
&socket,
&mut cred,
&mut send_seq,
last_size,
&mut frame_buffer,
&mut predictor,
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
arm_stale_terminal_input_suppression(
&mut stale_terminal_input_suppress_until,
);
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
last_terminal_frame_at = Instant::now();
}
last_packet_at = Instant::now();
flush_pending_user_input(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
&mut pending_user_input,
&mut pending_user_input_bytes,
)
.await?;
}
} else {
return Err(anyhow!("server rejected session: {}", reject.reason));
}
}
PacketKind::StreamOpen => {
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
continue;
};
let Ok(open) = protocol::from_body::<StreamOpen>(&plain) else {
continue;
};
last_packet_at = Instant::now();
if stream_retired.contains(&open.stream_id) {
send_stream_open_reject(
&socket,
addr,
&cred,
&mut send_seq,
open.stream_id,
"stream closed".to_string(),
)
.await?;
continue;
}
if opened_streams.contains(&open.stream_id) {
send_stream_open_ok(
&socket,
addr,
&cred,
&mut send_seq,
open.stream_id,
)
.await?;
continue;
}
// SSH-agent stream: the server is asking us to splice this
// stream into our LOCAL ssh-agent. Only honor it when we
// actually opted in (agent_sock is Some); otherwise reject,
// so a server cannot reach our agent without consent.
if open.target_host == AGENT_STREAM_SENTINEL {
#[cfg(not(unix))]
{
send_stream_open_reject(
&socket,
addr,
&cred,
&mut send_seq,
open.stream_id,
"agent forwarding is not supported on this client platform"
.to_string(),
)
.await?;
}
#[cfg(unix)]
{
let Some(sock_path) = agent_sock.clone() else {
send_stream_open_reject(
&socket,
addr,
&cred,
&mut send_seq,
open.stream_id,
"agent forwarding not enabled by client".to_string(),
)
.await?;
continue;
};
match UnixStream::connect(&sock_path).await {
Ok(stream) => {
let (mut reader, writer) = stream.into_split();
agent_writers.insert(open.stream_id, writer);
opened_streams.insert(open.stream_id);
stream_send_credit
.insert(open.stream_id, STREAM_INITIAL_WINDOW);
stream_next_send_offset.entry(open.stream_id).or_insert(0);
stream_next_recv_offset.entry(open.stream_id).or_insert(0);
send_stream_open_ok(
&socket,
addr,
&cred,
&mut send_seq,
open.stream_id,
)
.await?;
let forward_tx = remote_forward_tx.clone();
tokio::spawn(async move {
let mut buf = [0u8; 16 * 1024];
loop {
match reader.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if forward_tx
.send(ForwardEvent::Data {
stream_id: open.stream_id,
bytes: buf[..n].to_vec(),
})
.await
.is_err()
{
return;
}
}
Err(_) => break,
}
}
let _ = forward_tx
.send(ForwardEvent::Close {
stream_id: open.stream_id,
})
.await;
});
}
Err(err) => {
send_stream_open_reject(
&socket,
addr,
&cred,
&mut send_seq,
open.stream_id,
format!("connect local agent: {err}"),
)
.await?;
}
}
}
continue;
}
match TcpStream::connect((open.target_host.as_str(), open.target_port)).await {
Ok(stream) => {
let (mut reader, writer) = stream.into_split();
stream_writers.insert(open.stream_id, writer);
opened_streams.insert(open.stream_id);
stream_send_credit.insert(open.stream_id, STREAM_INITIAL_WINDOW);
stream_next_send_offset.entry(open.stream_id).or_insert(0);
stream_next_recv_offset.entry(open.stream_id).or_insert(0);
send_stream_open_ok(&socket, addr, &cred, &mut send_seq, open.stream_id).await?;
let forward_tx = remote_forward_tx.clone();
tokio::spawn(async move {
let mut buf = [0u8; 16 * 1024];
loop {
match reader.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if forward_tx
.send(ForwardEvent::Data {
stream_id: open.stream_id,
bytes: buf[..n].to_vec(),
})
.await
.is_err()
{
return;
}
}
Err(_) => break,
}
}
let _ = forward_tx.send(ForwardEvent::Close {
stream_id: open.stream_id,
}).await;
});
}
Err(err) => {
send_stream_open_reject(
&socket,
addr,
&cred,
&mut send_seq,
open.stream_id,
format!("connect {}:{}: {err}", open.target_host, open.target_port),
)
.await?;
}
}
}
PacketKind::StreamOpenOk => {
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
continue;
};
let Ok(ok) = protocol::from_body::<StreamOpenOk>(&plain) else {
continue;
};
last_packet_at = Instant::now();
if stream_pending_opens.remove(&ok.stream_id).is_none() {
continue;
}
opened_streams.insert(ok.stream_id);
stream_send_credit.entry(ok.stream_id).or_insert(STREAM_INITIAL_WINDOW);
stream_next_send_offset.entry(ok.stream_id).or_insert(0);
stream_next_recv_offset.entry(ok.stream_id).or_insert(0);
if pending_socks_replies.remove(&ok.stream_id)
&& let Some(writer) = stream_writers.get_mut(&ok.stream_id)
{
let _ = writer.write_all(&[5, 0, 0, 1, 0, 0, 0, 0, 0, 0]).await;
}
flush_stream_pending_data(
&socket,
addr,
&cred,
&mut send_seq,
ok.stream_id,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_next_send_offset,
&mut stream_sent_data,
).await?;
}
PacketKind::StreamOpenReject => {
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
continue;
};
let Ok(reject) = protocol::from_body::<StreamOpenReject>(&plain) else {
continue;
};
last_packet_at = Instant::now();
stream_pending_opens.remove(&reject.stream_id);
if pending_socks_replies.remove(&reject.stream_id)
&& let Some(writer) = stream_writers.get_mut(&reject.stream_id)
{
let _ = writer.write_all(&[5, 5, 0, 1, 0, 0, 0, 0, 0, 0]).await;
}
retire_stream_state(
reject.stream_id,
&mut opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_pending_opens,
&mut stream_next_send_offset,
&mut stream_sent_data,
&mut stream_next_recv_offset,
&mut stream_recv_pending,
&mut stream_retired,
&mut stream_retired_order,
);
stream_writers.remove(&reject.stream_id);
}
PacketKind::StreamData => {
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
continue;
};
let Ok(data) = protocol::from_body::<StreamData>(&plain) else {
continue;
};
last_packet_at = Instant::now();
let stream_id = data.stream_id;
if !opened_streams.contains(&stream_id) {
continue;
}
let (writes, consumed, received_offset) =
accept_stream_data(&mut stream_next_recv_offset, &mut stream_recv_pending, data);
if let Some(writer) = stream_writers.get_mut(&stream_id) {
for bytes in &writes {
let _ = writer.write_all(bytes).await;
}
} else {
#[cfg(unix)]
if let Some(writer) = agent_writers.get_mut(&stream_id) {
for bytes in &writes {
let _ = writer.write_all(bytes).await;
}
}
}
// Always return flow-control credit so a stream whose local
// writer is already gone cannot wedge the server's send window.
send_stream_window_adjust(
&socket,
addr,
&cred,
&mut send_seq,
stream_id,
consumed,
received_offset,
).await?;
}
PacketKind::StreamWindowAdjust => {
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
continue;
};
let Ok(adjust) = protocol::from_body::<StreamWindowAdjust>(&plain) else {
continue;
};
last_packet_at = Instant::now();
if !opened_streams.contains(&adjust.stream_id) {
continue;
}
ack_stream_data(
&mut stream_sent_data,
&mut stream_send_credit,
adjust.stream_id,
adjust.received_offset,
);
flush_stream_pending_data(
&socket,
addr,
&cred,
&mut send_seq,
adjust.stream_id,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_next_send_offset,
&mut stream_sent_data,
).await?;
}
PacketKind::StreamClose => {
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
continue;
};
let Ok(close) = protocol::from_body::<StreamClose>(&plain) else {
continue;
};
last_packet_at = Instant::now();
pending_socks_replies.remove(&close.stream_id);
retire_stream_state(
close.stream_id,
&mut opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_pending_opens,
&mut stream_next_send_offset,
&mut stream_sent_data,
&mut stream_next_recv_offset,
&mut stream_recv_pending,
&mut stream_retired,
&mut stream_retired_order,
);
stream_writers.remove(&close.stream_id);
#[cfg(unix)]
agent_writers.remove(&close.stream_id);
}
_ => {}
}
}
forward_event = forward_rx.recv() => {
match forward_event {
Some(ForwardEvent::Open { stream_id, target_host, target_port, writer, socks5_reply }) => {
stream_writers.insert(stream_id, writer);
stream_send_credit.insert(stream_id, STREAM_INITIAL_WINDOW);
stream_next_send_offset.entry(stream_id).or_insert(0);
stream_next_recv_offset.entry(stream_id).or_insert(0);
if socks5_reply {
pending_socks_replies.insert(stream_id);
}
stream_pending_opens.insert(stream_id, PendingStreamOpen {
target_host: target_host.clone(),
target_port,
last_sent: Instant::now(),
attempts: 1,
});
send_stream_open(
&socket,
addr,
&cred,
&mut send_seq,
stream_id,
target_host,
target_port,
)
.await?;
}
Some(ForwardEvent::Data { stream_id, bytes }) => {
queue_or_send_stream_data(
&socket,
addr,
&cred,
&mut send_seq,
stream_id,
bytes,
&opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_next_send_offset,
&mut stream_sent_data,
).await?;
}
Some(ForwardEvent::Close { stream_id }) => {
pending_socks_replies.remove(&stream_id);
retire_stream_state(
stream_id,
&mut opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_pending_opens,
&mut stream_next_send_offset,
&mut stream_sent_data,
&mut stream_next_recv_offset,
&mut stream_recv_pending,
&mut stream_retired,
&mut stream_retired_order,
);
stream_writers.remove(&stream_id);
#[cfg(unix)]
agent_writers.remove(&stream_id);
send_stream_close(&socket, addr, &cred, &mut send_seq, stream_id).await?;
}
None if forward_only => break,
None => {}
}
}
_ = status_tick.tick() => {
let stale = last_packet_at.elapsed();
if stale >= Duration::from_secs(reconnect_timeout_secs.max(1)) {
if let Some(frame) = reconnect(
&socket,
&mut cred,
&mut send_seq,
last_size,
&mut frame_buffer,
&mut predictor,
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
arm_stale_terminal_input_suppression(
&mut stale_terminal_input_suppress_until,
);
if !forward_only {
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
last_terminal_frame_at = Instant::now();
}
last_packet_at = Instant::now();
flush_pending_user_input(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
&mut pending_user_input,
&mut pending_user_input_bytes,
)
.await?;
}
} else if stale >= Duration::from_secs(2) {
let packet = protocol::encode_encrypted(
PacketKind::Ping,
cred.client_id,
send_seq,
cred.last_rendered_seq,
&cred.session_key,
CLIENT_TO_SERVER,
b"",
)?;
send_seq += 1;
let _ = send_terminal_udp(&socket, &packet, addr).await?;
}
// Re-evaluate the prediction display policy on every timer tick so
// a latency spike (or recovery) flips speculation on/off promptly
// without waiting for the next keystroke to drive `redraw`.
if !forward_only {
predictor.refresh_policy()?;
flush_startup_input_if_ready(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
(&mut pending_user_input, &mut pending_user_input_bytes),
(&mut startup_input_hold_until, &mut startup_gate_mode),
)
.await?;
let now = Instant::now();
if should_repaint_idle_alternate_screen(
predictor.alternate_screen,
last_terminal_frame_at,
last_idle_repaint_attempt_at,
now,
) {
last_idle_repaint_attempt_at = now;
if let Some(frame) = reconnect(
&socket,
&mut cred,
&mut send_seq,
last_size,
&mut frame_buffer,
&mut predictor,
)
.await?
{
refresh_live_addr(&mut addr, &cred)?;
arm_stale_terminal_input_suppression(
&mut stale_terminal_input_suppress_until,
);
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
last_terminal_frame_at = Instant::now();
last_packet_at = Instant::now();
flush_pending_user_input(
&socket,
addr,
&cred,
&mut send_seq,
&mut predictor,
&mut pending_user_input,
&mut pending_user_input_bytes,
)
.await?;
}
}
}
// Refresh / clear the non-destructive disconnect status line based
// on how long the link has been silent (recomputed after any
// reconnect attempt above may have reset `last_packet_at`).
if !forward_only {
let action = if predictor.alternate_screen {
disconnect_status.on_suppressed()
} else {
disconnect_status.on_tick(last_packet_at.elapsed())
};
disconnect_status.apply(action)?;
}
}
_ = stream_retransmit_tick.tick() => {
retransmit_stream_data(
&socket,
addr,
&cred,
&mut send_seq,
&mut stream_sent_data,
).await?;
retransmit_stream_opens(
&socket,
addr,
&cred,
&mut send_seq,
&mut stream_pending_opens,
).await?;
}
}
}
// Leave the terminal clean on exit: erase any lingering status line.
let _ = disconnect_status.apply(StatusAction::Clear);
let _ = detach_once(&socket, &cred, send_seq).await;
Ok(())
}
async fn start_local_forwards(
local_forwards: &[LocalForward],
forward_tx: mpsc::Sender<ForwardEvent>,
stream_ids: Arc<AtomicU64>,
) -> Result<()> {
for forward in local_forwards {
let bind = format!("{}:{}", forward.bind_host, forward.listen_port);
let listener = TcpListener::bind(&bind)
.await
.with_context(|| format!("bind local forward {bind}"))?;
let forward = forward.clone();
let forward_tx = forward_tx.clone();
let stream_ids = Arc::clone(&stream_ids);
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
break;
};
let stream_id = stream_ids.fetch_add(1, Ordering::Relaxed);
let (mut reader, writer) = stream.into_split();
let _ = forward_tx
.send(ForwardEvent::Open {
stream_id,
target_host: forward.target_host.clone(),
target_port: forward.target_port,
writer,
socks5_reply: false,
})
.await;
let forward_tx = forward_tx.clone();
tokio::spawn(async move {
let mut buf = [0u8; 16 * 1024];
loop {
match reader.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if forward_tx
.send(ForwardEvent::Data {
stream_id,
bytes: buf[..n].to_vec(),
})
.await
.is_err()
{
return;
}
}
Err(_) => break,
}
}
let _ = forward_tx.send(ForwardEvent::Close { stream_id }).await;
});
}
});
}
Ok(())
}
async fn start_dynamic_forwards(
dynamic_forwards: &[DynamicForward],
forward_tx: mpsc::Sender<ForwardEvent>,
stream_ids: Arc<AtomicU64>,
) -> Result<()> {
for forward in dynamic_forwards {
let bind = format!("{}:{}", forward.bind_host, forward.listen_port);
let listener = TcpListener::bind(&bind)
.await
.with_context(|| format!("bind dynamic forward {bind}"))?;
let forward_tx = forward_tx.clone();
let stream_ids = Arc::clone(&stream_ids);
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
break;
};
let forward_tx = forward_tx.clone();
let stream_ids = Arc::clone(&stream_ids);
tokio::spawn(async move {
let _ = handle_socks5_connect(stream, forward_tx, stream_ids).await;
});
}
});
}
Ok(())
}
async fn handle_socks5_connect(
mut stream: TcpStream,
forward_tx: mpsc::Sender<ForwardEvent>,
stream_ids: Arc<AtomicU64>,
) -> Result<()> {
let mut hello = [0u8; 2];
stream.read_exact(&mut hello).await?;
anyhow::ensure!(hello[0] == 5, "unsupported SOCKS version {}", hello[0]);
let mut methods = vec![0u8; hello[1] as usize];
stream.read_exact(&mut methods).await?;
if !methods.contains(&0) {
stream.write_all(&[5, 0xff]).await?;
return Ok(());
}
stream.write_all(&[5, 0]).await?;
let mut header = [0u8; 4];
stream.read_exact(&mut header).await?;
anyhow::ensure!(
header[0] == 5,
"unsupported SOCKS request version {}",
header[0]
);
if header[1] != 1 {
stream.write_all(&[5, 7, 0, 1, 0, 0, 0, 0, 0, 0]).await?;
return Ok(());
}
let target_host = match header[3] {
1 => {
let mut raw = [0u8; 4];
stream.read_exact(&mut raw).await?;
Ipv4Addr::from(raw).to_string()
}
3 => {
let mut len = [0u8; 1];
stream.read_exact(&mut len).await?;
let mut raw = vec![0u8; len[0] as usize];
stream.read_exact(&mut raw).await?;
String::from_utf8(raw).context("SOCKS domain is not valid UTF-8")?
}
4 => {
let mut raw = [0u8; 16];
stream.read_exact(&mut raw).await?;
Ipv6Addr::from(raw).to_string()
}
atyp => {
stream.write_all(&[5, 8, 0, 1, 0, 0, 0, 0, 0, 0]).await?;
return Err(anyhow!("unsupported SOCKS address type {atyp}"));
}
};
let mut port = [0u8; 2];
stream.read_exact(&mut port).await?;
let target_port = u16::from_be_bytes(port);
let stream_id = stream_ids.fetch_add(1, Ordering::Relaxed);
let (mut reader, writer) = stream.into_split();
forward_tx
.send(ForwardEvent::Open {
stream_id,
target_host,
target_port,
writer,
socks5_reply: true,
})
.await?;
tokio::spawn(async move {
let mut buf = [0u8; 16 * 1024];
loop {
match reader.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if forward_tx
.send(ForwardEvent::Data {
stream_id,
bytes: buf[..n].to_vec(),
})
.await
.is_err()
{
return;
}
}
Err(_) => break,
}
}
let _ = forward_tx.send(ForwardEvent::Close { stream_id }).await;
});
Ok(())
}
async fn reconnect(
socket: &UdpSocket,
cred: &mut CachedCredential,
send_seq: &mut u64,
size: (u16, u16),
frame_buffer: &mut FrameBuffer,
predictor: &mut Predictor,
) -> Result<Option<Frame>> {
if let Ok((frame, next)) = try_live_resume(socket, cred, send_seq, size.0, size.1).await {
*cred = next;
frame_buffer.clear();
predictor.reset();
cred.last_rendered_seq = frame.output_seq;
send_ack(
socket,
resolve_addr(&cred.udp_host, cred.udp_port)?,
cred,
send_seq,
)
.await?;
return Ok(Some(frame));
}
let Ok((frame, next)) = try_ticket_attach(socket, cred, size.0, size.1, Vec::new()).await
else {
return Ok(None);
};
*cred = next;
*send_seq = 2;
frame_buffer.clear();
predictor.reset();
cred.last_rendered_seq = frame.output_seq;
send_ack(
socket,
resolve_addr(&cred.udp_host, cred.udp_port)?,
cred,
send_seq,
)
.await?;
Ok(Some(frame))
}
fn refresh_live_addr(addr: &mut SocketAddr, cred: &CachedCredential) -> Result<()> {
*addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
Ok(())
}
fn queue_pending_user_input(
pending: &mut VecDeque<PendingUserInput>,
pending_bytes: &mut usize,
bytes: Vec<u8>,
) -> Result<()> {
queue_pending_user_input_with_filter(pending, pending_bytes, bytes, false)
}
fn queue_stale_pending_user_input(
pending: &mut VecDeque<PendingUserInput>,
pending_bytes: &mut usize,
bytes: Vec<u8>,
) -> Result<()> {
queue_pending_user_input_with_filter(pending, pending_bytes, bytes, true)
}
fn queue_pending_user_input_with_filter(
pending: &mut VecDeque<PendingUserInput>,
pending_bytes: &mut usize,
bytes: Vec<u8>,
strip_mouse_reports: bool,
) -> Result<()> {
let next = pending_bytes
.checked_add(bytes.len())
.ok_or_else(|| anyhow!("pending input buffer overflow"))?;
anyhow::ensure!(
next <= MAX_PENDING_USER_INPUT_BYTES,
"pending input buffer full while disconnected"
);
*pending_bytes = next;
pending.push_back(PendingUserInput {
bytes,
strip_mouse_reports,
});
Ok(())
}
fn requeue_pending_user_input_front(
pending: &mut VecDeque<PendingUserInput>,
pending_bytes: &mut usize,
input: PendingUserInput,
) -> Result<()> {
let next = pending_bytes
.checked_add(input.bytes.len())
.ok_or_else(|| anyhow!("pending input buffer overflow"))?;
anyhow::ensure!(
next <= MAX_PENDING_USER_INPUT_BYTES,
"pending input buffer full while disconnected"
);
*pending_bytes = next;
pending.push_front(input);
Ok(())
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct PendingUserInput {
bytes: Vec<u8>,
strip_mouse_reports: bool,
}
fn split_after_command_submit(bytes: &[u8]) -> Option<(Vec<u8>, Vec<u8>)> {
let split_at = bytes
.iter()
.position(|byte| matches!(byte, b'\r' | b'\n'))?
+ 1;
if split_at == bytes.len() {
return Some((bytes.to_vec(), Vec::new()));
}
let later = bytes[split_at..].to_vec();
let now = bytes[..split_at].to_vec();
Some((now, later))
}
fn should_hold_post_submit_input(bytes: &[u8]) -> bool {
matches!(bytes.first(), Some(0x1b | 0x9b))
}
fn post_submit_hold_duration(followup: &[u8]) -> Duration {
if should_hold_post_submit_input(followup) {
STARTUP_INPUT_HOLD
} else {
POST_SUBMIT_ALL_INPUT_HOLD
}
}
fn should_hold_during_startup_gate(
mode: StartupGateMode,
already_queued: bool,
bytes: &[u8],
) -> bool {
mode == StartupGateMode::HoldAll || already_queued || should_hold_post_submit_input(bytes)
}
fn should_flush_terminal_input_after_contact(
forward_only: bool,
mode: &str,
pending_empty: bool,
) -> bool {
!forward_only && !pending_empty && mode != "view-only" && mode != "forward-only"
}
fn should_strip_stale_terminal_reports(
silent_for: Duration,
suppress_until: Option<Instant>,
) -> bool {
silent_for >= STALE_TERMINAL_INPUT_AFTER
|| suppress_until.is_some_and(|deadline| Instant::now() < deadline)
}
fn input_contains_focus_in(bytes: &[u8]) -> bool {
contains_bytes(bytes, b"\x1b[I") || contains_bytes(bytes, b"\x9bI")
}
fn strip_terminal_focus_reports(bytes: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(bytes.len());
let mut offset = 0;
while offset < bytes.len() {
match bytes.get(offset..) {
Some([0x1b, b'[', b'I' | b'O', ..]) => {
offset += 3;
}
Some([0x9b, b'I' | b'O', ..]) => {
offset += 2;
}
_ => {
out.push(bytes[offset]);
offset += 1;
}
}
}
out
}
fn should_repaint_idle_alternate_screen(
alternate_screen: bool,
last_terminal_frame_at: Instant,
last_attempt_at: Instant,
now: Instant,
) -> bool {
alternate_screen
&& now.duration_since(last_terminal_frame_at) >= ALT_SCREEN_IDLE_REPAINT_AFTER
&& now.duration_since(last_attempt_at) >= ALT_SCREEN_IDLE_REPAINT_AFTER
}
fn arm_stale_terminal_input_suppression(suppress_until: &mut Option<Instant>) {
*suppress_until = Some(Instant::now() + POST_RECONNECT_STALE_INPUT_GRACE);
}
async fn flush_pending_user_input(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
predictor: &mut Predictor,
pending: &mut VecDeque<PendingUserInput>,
pending_bytes: &mut usize,
) -> Result<()> {
while let Some(input) = pending.pop_front() {
let PendingUserInput {
bytes,
strip_mouse_reports,
} = input;
*pending_bytes = pending_bytes.saturating_sub(bytes.len());
let bytes = if strip_mouse_reports {
strip_stale_mouse_reports(&bytes)
} else {
bytes
};
if bytes.is_empty() {
continue;
}
if send_input(socket, addr, cred, send_seq, bytes.clone()).await? {
predictor.observe_input(&bytes)?;
} else {
requeue_pending_user_input_front(
pending,
pending_bytes,
PendingUserInput {
bytes,
strip_mouse_reports: false,
},
)?;
break;
}
}
Ok(())
}
fn strip_stale_mouse_reports(bytes: &[u8]) -> Vec<u8> {
let mut out = Vec::with_capacity(bytes.len());
let mut offset = 0;
while offset < bytes.len() {
if let Some(len) = stale_mouse_report_len(bytes, offset) {
offset += len;
continue;
}
if stale_mouse_report_maybe_incomplete(bytes, offset) {
break;
}
out.push(bytes[offset]);
offset += 1;
}
out
}
fn stale_mouse_report_len(bytes: &[u8], offset: usize) -> Option<usize> {
match bytes.get(offset).copied()? {
0x1b if bytes.get(offset + 1) == Some(&b'[') => match bytes.get(offset + 2).copied() {
Some(b'M') if offset + 6 <= bytes.len() => Some(6),
Some(b'I' | b'O') => Some(3),
Some(b'<') => mouse_report_params_len(bytes, offset + 3, 2).map(|len| len + 3),
Some(byte) if byte.is_ascii_digit() => {
mouse_report_params_len(bytes, offset + 2, 2).map(|len| len + 2)
}
_ => None,
},
0x9b => match bytes.get(offset + 1).copied() {
Some(b'M') if offset + 5 <= bytes.len() => Some(5),
Some(b'I' | b'O') => Some(2),
Some(b'<') => mouse_report_params_len(bytes, offset + 2, 2).map(|len| len + 2),
Some(byte) if byte.is_ascii_digit() => {
mouse_report_params_len(bytes, offset + 1, 2).map(|len| len + 1)
}
_ => None,
},
byte if byte.is_ascii_digit() => mouse_report_params_len(bytes, offset, 1)
.or_else(|| orphan_mouse_suffix_len(bytes, offset)),
b'M' | b'm' if offset == 0 => orphan_leading_mouse_terminator_len(bytes, offset),
_ => None,
}
}
fn mouse_report_params_len(bytes: &[u8], offset: usize, min_semicolons: usize) -> Option<usize> {
let mut semicolons = 0usize;
let mut cursor = offset;
while let Some(byte) = bytes.get(cursor).copied() {
match byte {
b'0'..=b'9' => cursor += 1,
b';' => {
semicolons += 1;
cursor += 1;
}
b'M' | b'm' if semicolons >= min_semicolons => return Some(cursor + 1 - offset),
_ => return None,
}
}
None
}
fn orphan_mouse_suffix_len(bytes: &[u8], offset: usize) -> Option<usize> {
if offset != 0 {
return None;
}
let mut cursor = offset;
while let Some(byte) = bytes.get(cursor).copied() {
match byte {
b'0'..=b'9' => cursor += 1,
b'M' | b'm' if cursor > offset => return Some(cursor + 1 - offset),
_ => return None,
}
}
None
}
fn orphan_leading_mouse_terminator_len(bytes: &[u8], offset: usize) -> Option<usize> {
if offset != 0 || !matches!(bytes.get(offset), Some(b'M' | b'm')) {
return None;
}
mouse_report_params_len(bytes, offset + 1, 1).map(|len| len + 1)
}
fn stale_mouse_report_maybe_incomplete(bytes: &[u8], offset: usize) -> bool {
let Some(rest) = bytes.get(offset..) else {
return false;
};
match rest {
[0x1b] | [0x1b, b'['] | [0x1b, b'[', b'<'] => true,
[0x1b, b'[', b'M', ..] if rest.len() < 6 => true,
[0x1b, b'[', b'<', params @ ..] | [0x1b, b'[', params @ ..]
if params_are_mouse_prefix(params, 2) =>
{
true
}
[0x9b] | [0x9b, b'<'] => true,
[0x9b, b'M', ..] if rest.len() < 5 => true,
[0x9b, b'<', params @ ..] | [0x9b, params @ ..] if params_are_mouse_prefix(params, 2) => {
true
}
params if params_are_mouse_prefix(params, 2) => true,
_ => false,
}
}
fn params_are_mouse_prefix(params: &[u8], min_semicolons: usize) -> bool {
let mut semicolons = 0usize;
let mut saw_digit = false;
for byte in params {
match byte {
b'0'..=b'9' => saw_digit = true,
b';' if saw_digit => semicolons += 1,
_ => return false,
}
}
saw_digit && semicolons >= min_semicolons
}
async fn flush_startup_input_if_ready(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
predictor: &mut Predictor,
pending: (&mut VecDeque<PendingUserInput>, &mut usize),
hold: (&mut Option<Instant>, &mut StartupGateMode),
) -> Result<()> {
let (startup_hold_until, startup_gate_mode) = hold;
let Some(deadline) = *startup_hold_until else {
return Ok(());
};
if !predictor.alternate_screen && Instant::now() < deadline {
return Ok(());
}
*startup_hold_until = None;
*startup_gate_mode = StartupGateMode::HoldControl;
let (pending, pending_bytes) = pending;
flush_pending_user_input(
socket,
addr,
cred,
send_seq,
predictor,
pending,
pending_bytes,
)
.await
}
async fn send_input(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
bytes: Vec<u8>,
) -> Result<bool> {
let body = protocol::to_body(&Input { bytes })?;
let packet = protocol::encode_encrypted(
PacketKind::Input,
cred.client_id,
*send_seq,
cred.last_rendered_seq,
&cred.session_key,
CLIENT_TO_SERVER,
&body,
)?;
*send_seq += 1;
send_terminal_udp(socket, &packet, addr).await
}
async fn send_stream_open(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
target_host: String,
target_port: u16,
) -> Result<()> {
let body = protocol::to_body(&StreamOpen {
stream_id,
target_host,
target_port,
})?;
send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamOpen, &body).await
}
async fn send_stream_open_ok(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
) -> Result<()> {
let body = protocol::to_body(&StreamOpenOk { stream_id })?;
send_stream_packet(
socket,
addr,
cred,
send_seq,
PacketKind::StreamOpenOk,
&body,
)
.await
}
async fn send_stream_open_reject(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
reason: String,
) -> Result<()> {
let body = protocol::to_body(&StreamOpenReject { stream_id, reason })?;
send_stream_packet(
socket,
addr,
cred,
send_seq,
PacketKind::StreamOpenReject,
&body,
)
.await
}
async fn send_stream_data(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
offset: u64,
bytes: Vec<u8>,
) -> Result<()> {
let body = protocol::to_body(&StreamData {
stream_id,
offset,
bytes,
})?;
send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamData, &body).await
}
#[allow(clippy::too_many_arguments)]
async fn queue_or_send_stream_data(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
bytes: Vec<u8>,
opened_streams: &HashSet<u64>,
stream_send_credit: &mut HashMap<u64, usize>,
stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
stream_next_send_offset: &mut HashMap<u64, u64>,
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
) -> Result<()> {
if !opened_streams.contains(&stream_id)
|| stream_send_credit.get(&stream_id).copied().unwrap_or(0) < bytes.len()
|| stream_pending_data
.get(&stream_id)
.is_some_and(|pending| !pending.is_empty())
{
stream_pending_data
.entry(stream_id)
.or_default()
.push_back(bytes);
return Ok(());
}
*stream_send_credit.entry(stream_id).or_default() -= bytes.len();
let offset = *stream_next_send_offset.entry(stream_id).or_default();
stream_next_send_offset.insert(stream_id, offset.saturating_add(bytes.len() as u64));
stream_sent_data.entry(stream_id).or_default().insert(
offset,
PendingStreamChunk {
offset,
bytes: bytes.clone(),
last_sent: Instant::now(),
attempts: 1,
},
);
send_stream_data(socket, addr, cred, send_seq, stream_id, offset, bytes).await
}
#[allow(clippy::too_many_arguments)]
async fn flush_stream_pending_data(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
stream_send_credit: &mut HashMap<u64, usize>,
stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
stream_next_send_offset: &mut HashMap<u64, u64>,
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
) -> Result<()> {
let Some(pending) = stream_pending_data.get_mut(&stream_id) else {
return Ok(());
};
while let Some(bytes) = pending.front() {
let credit = stream_send_credit.get(&stream_id).copied().unwrap_or(0);
if credit < bytes.len() {
break;
}
let bytes = pending.pop_front().expect("pending front exists");
*stream_send_credit.entry(stream_id).or_default() -= bytes.len();
let offset = *stream_next_send_offset.entry(stream_id).or_default();
stream_next_send_offset.insert(stream_id, offset.saturating_add(bytes.len() as u64));
stream_sent_data.entry(stream_id).or_default().insert(
offset,
PendingStreamChunk {
offset,
bytes: bytes.clone(),
last_sent: Instant::now(),
attempts: 1,
},
);
send_stream_data(socket, addr, cred, send_seq, stream_id, offset, bytes).await?;
}
if pending.is_empty() {
stream_pending_data.remove(&stream_id);
}
Ok(())
}
async fn retransmit_stream_data(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
) -> Result<()> {
let now = Instant::now();
let mut retransmit = Vec::new();
for (stream_id, chunks) in stream_sent_data.iter_mut() {
for chunk in chunks.values_mut() {
if now.duration_since(chunk.last_sent) < Duration::from_millis(200) {
continue;
}
chunk.last_sent = now;
chunk.attempts = chunk.attempts.saturating_add(1);
retransmit.push((*stream_id, chunk.offset, chunk.bytes.clone()));
}
}
for (stream_id, offset, bytes) in retransmit {
send_stream_data(socket, addr, cred, send_seq, stream_id, offset, bytes).await?;
}
Ok(())
}
async fn retransmit_stream_opens(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_pending_opens: &mut HashMap<u64, PendingStreamOpen>,
) -> Result<()> {
let now = Instant::now();
let mut retransmit = Vec::new();
for (stream_id, pending) in stream_pending_opens.iter_mut() {
if now.duration_since(pending.last_sent) < Duration::from_millis(200) {
continue;
}
pending.last_sent = now;
pending.attempts = pending.attempts.saturating_add(1);
retransmit.push((*stream_id, pending.target_host.clone(), pending.target_port));
}
for (stream_id, target_host, target_port) in retransmit {
send_stream_open(
socket,
addr,
cred,
send_seq,
stream_id,
target_host,
target_port,
)
.await?;
}
Ok(())
}
fn accept_stream_data(
stream_next_recv_offset: &mut HashMap<u64, u64>,
stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
data: StreamData,
) -> (Vec<Vec<u8>>, usize, u64) {
let stream_id = data.stream_id;
let expected = stream_next_recv_offset.entry(stream_id).or_insert(0);
if data.offset < *expected {
return (Vec::new(), 0, *expected);
}
if data.offset > *expected {
stream_recv_pending
.entry(stream_id)
.or_default()
.entry(data.offset)
.or_insert(data.bytes);
return (Vec::new(), 0, *expected);
}
let mut writes = vec![data.bytes];
let mut consumed = writes[0].len();
*expected = expected.saturating_add(consumed as u64);
while let Some(bytes) = stream_recv_pending
.entry(stream_id)
.or_default()
.remove(expected)
{
consumed = consumed.saturating_add(bytes.len());
*expected = expected.saturating_add(bytes.len() as u64);
writes.push(bytes);
}
if stream_recv_pending
.get(&stream_id)
.is_some_and(BTreeMap::is_empty)
{
stream_recv_pending.remove(&stream_id);
}
(writes, consumed, *expected)
}
#[allow(clippy::too_many_arguments)]
fn cleanup_stream_state(
stream_id: u64,
opened_streams: &mut HashSet<u64>,
stream_send_credit: &mut HashMap<u64, usize>,
stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
stream_pending_opens: &mut HashMap<u64, PendingStreamOpen>,
stream_next_send_offset: &mut HashMap<u64, u64>,
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
stream_next_recv_offset: &mut HashMap<u64, u64>,
stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
) {
opened_streams.remove(&stream_id);
stream_send_credit.remove(&stream_id);
stream_pending_data.remove(&stream_id);
stream_pending_opens.remove(&stream_id);
stream_next_send_offset.remove(&stream_id);
stream_sent_data.remove(&stream_id);
stream_next_recv_offset.remove(&stream_id);
stream_recv_pending.remove(&stream_id);
}
#[allow(clippy::too_many_arguments)]
fn stream_state_exists(
stream_id: u64,
opened_streams: &HashSet<u64>,
stream_send_credit: &HashMap<u64, usize>,
stream_pending_data: &HashMap<u64, VecDeque<Vec<u8>>>,
stream_pending_opens: &HashMap<u64, PendingStreamOpen>,
stream_next_send_offset: &HashMap<u64, u64>,
stream_sent_data: &HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
stream_next_recv_offset: &HashMap<u64, u64>,
stream_recv_pending: &HashMap<u64, BTreeMap<u64, Vec<u8>>>,
) -> bool {
opened_streams.contains(&stream_id)
|| stream_send_credit.contains_key(&stream_id)
|| stream_pending_data.contains_key(&stream_id)
|| stream_pending_opens.contains_key(&stream_id)
|| stream_next_send_offset.contains_key(&stream_id)
|| stream_sent_data.contains_key(&stream_id)
|| stream_next_recv_offset.contains_key(&stream_id)
|| stream_recv_pending.contains_key(&stream_id)
}
#[allow(clippy::too_many_arguments)]
fn retire_stream_state(
stream_id: u64,
opened_streams: &mut HashSet<u64>,
stream_send_credit: &mut HashMap<u64, usize>,
stream_pending_data: &mut HashMap<u64, VecDeque<Vec<u8>>>,
stream_pending_opens: &mut HashMap<u64, PendingStreamOpen>,
stream_next_send_offset: &mut HashMap<u64, u64>,
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
stream_next_recv_offset: &mut HashMap<u64, u64>,
stream_recv_pending: &mut HashMap<u64, BTreeMap<u64, Vec<u8>>>,
stream_retired: &mut HashSet<u64>,
stream_retired_order: &mut VecDeque<u64>,
) {
let had_state = stream_state_exists(
stream_id,
opened_streams,
stream_send_credit,
stream_pending_data,
stream_pending_opens,
stream_next_send_offset,
stream_sent_data,
stream_next_recv_offset,
stream_recv_pending,
);
cleanup_stream_state(
stream_id,
opened_streams,
stream_send_credit,
stream_pending_data,
stream_pending_opens,
stream_next_send_offset,
stream_sent_data,
stream_next_recv_offset,
stream_recv_pending,
);
if !had_state || STREAM_RETIRED_TOMBSTONES == 0 {
return;
}
if stream_retired.insert(stream_id) {
stream_retired_order.push_back(stream_id);
}
while stream_retired_order.len() > STREAM_RETIRED_TOMBSTONES {
if let Some(expired) = stream_retired_order.pop_front() {
stream_retired.remove(&expired);
}
}
}
fn ack_stream_data(
stream_sent_data: &mut HashMap<u64, BTreeMap<u64, PendingStreamChunk>>,
stream_send_credit: &mut HashMap<u64, usize>,
stream_id: u64,
received_offset: u64,
) {
let Some(sent) = stream_sent_data.get_mut(&stream_id) else {
return;
};
let acked_offsets: Vec<u64> = sent
.iter()
.filter_map(|(offset, chunk)| {
let end = chunk.offset.saturating_add(chunk.bytes.len() as u64);
(end <= received_offset).then_some(*offset)
})
.collect();
let mut acked_bytes = 0usize;
for offset in acked_offsets {
if let Some(chunk) = sent.remove(&offset) {
acked_bytes = acked_bytes.saturating_add(chunk.bytes.len());
}
}
if sent.is_empty() {
stream_sent_data.remove(&stream_id);
}
add_stream_credit(stream_send_credit, stream_id, acked_bytes);
}
fn add_stream_credit(stream_send_credit: &mut HashMap<u64, usize>, stream_id: u64, bytes: usize) {
let credit = stream_send_credit.entry(stream_id).or_default();
*credit = credit.saturating_add(bytes).min(STREAM_INITIAL_WINDOW);
}
async fn send_stream_window_adjust(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
bytes: usize,
received_offset: u64,
) -> Result<()> {
let body = protocol::to_body(&StreamWindowAdjust {
stream_id,
received_offset,
bytes: bytes.min(u32::MAX as usize) as u32,
})?;
send_stream_packet(
socket,
addr,
cred,
send_seq,
PacketKind::StreamWindowAdjust,
&body,
)
.await
}
async fn send_stream_close(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
stream_id: u64,
) -> Result<()> {
let body = protocol::to_body(&StreamClose { stream_id })?;
send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamClose, &body).await
}
async fn send_stream_packet(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
kind: PacketKind,
body: &[u8],
) -> Result<()> {
let packet = protocol::encode_encrypted(
kind,
cred.client_id,
*send_seq,
cred.last_rendered_seq,
&cred.session_key,
CLIENT_TO_SERVER,
body,
)?;
*send_seq += 1;
let _ = send_terminal_udp(socket, &packet, addr).await?;
Ok(())
}
// ---------------------------------------------------------------------------
// Predictive local echo (speculative echo).
//
// This is a clean-room Rust implementation written from the design described in
// the Mosh paper (Winstein & Balakrishnan, "Mosh: An Interactive Remote Shell
// for Mobile Clients", USENIX ATC 2012) and a conceptual reading of how a
// prediction overlay behaves. No Mosh source was copied or translated; the data
// structures, byte handling, confirmation strategy, and drawing here are
// original to dosh.
//
// Why dosh's approach differs from Mosh's: Mosh runs a full terminal emulator
// on the client and keeps a cell-accurate framebuffer, so it can confirm each
// predicted character by comparing the predicted cell to the actual cell the
// server produced. dosh's client has no client-side emulator -- it blits opaque
// server byte-frames straight to the real terminal. So dosh cannot do
// content-level confirmation. Instead it confirms predictions by *frame
// sequencing*: when the server emits a new authoritative frame (output_seq
// advances) that reflects the keystrokes we predicted, that frame is rendered
// and supersedes our overlay; we erase the speculative glyphs and let the
// server's bytes stand. To draw and erase safely without an emulator, dosh
// keeps a tiny model of just the current line near the cursor.
//
// Correctness over coverage: a visibly wrong/sticky prediction is worse than no
// prediction, so anything we cannot model with confidence (escape sequences,
// control chars other than backspace, multi-byte / wide chars, the right
// margin) ends the current epoch and stops further speculation until the next
// confirmation re-bases us against authoritative output.
/// Display policy for predictions, mirroring Mosh's off / adaptive / always
/// design (we call the adaptive mode "experimental" to match Mosh's naming for
/// the SRTT-gated mode that only shows predictions on a laggy link).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum PredictMode {
/// Never show predictions.
Off,
/// Show predictions only when the link is laggy (SRTT above a trigger).
Experimental,
/// Always show predictions whenever any are outstanding.
Always,
}
impl PredictMode {
fn parse(s: &str) -> Self {
match s.trim().to_ascii_lowercase().as_str() {
"off" | "never" | "false" | "no" => PredictMode::Off,
"always" => PredictMode::Always,
// "experimental" / "adaptive" / "auto" / anything else -> adaptive.
_ => PredictMode::Experimental,
}
}
}
/// SRTT (in milliseconds) above which the experimental mode begins showing
/// predictions, and below which it stops. Hysteresis avoids flapping on a link
/// hovering around the threshold. On a faster-than-this link, native local echo
/// already feels instant, so speculation only adds risk.
const SRTT_TRIGGER_HIGH_MS: f64 = 30.0;
const SRTT_TRIGGER_LOW_MS: f64 = 20.0;
/// SRTT (ms) above which we additionally underline ("flag") predicted cells so
/// the user can tell speculative glyphs from confirmed ones on a very slow link.
const FLAG_TRIGGER_HIGH_MS: f64 = 80.0;
const FLAG_TRIGGER_LOW_MS: f64 = 50.0;
/// If a prediction is still outstanding after this long, treat the link as
/// laggy enough to force display even if the SRTT estimate hasn't caught up yet
/// (covers a sudden latency spike on the very first slow keystroke).
const GLITCH_FORCE_MS: u128 = 250;
/// Enough recent bytes to join a CSI private-mode sequence split across frames
/// without retaining arbitrary terminal output.
const TERMINAL_OUTPUT_PARSE_TAIL: usize = 64;
fn alternate_screen_after_output(mut active: bool, bytes: &[u8]) -> bool {
let mut offset = 0usize;
while offset < bytes.len() {
let Some((next, transition)) = terminal_private_mode_transition(bytes, offset) else {
offset += 1;
continue;
};
if let Some(enable) = transition {
active = enable;
}
offset = next.max(offset + 1);
}
active
}
fn terminal_private_mode_transition(bytes: &[u8], offset: usize) -> Option<(usize, Option<bool>)> {
let mut cursor = offset;
match bytes.get(cursor).copied()? {
0x1b if bytes.get(cursor + 1) == Some(&b'[') => cursor += 2,
0x9b => cursor += 1,
_ => return None,
}
if bytes.get(cursor) != Some(&b'?') {
return None;
}
cursor += 1;
let params_start = cursor;
while let Some(byte) = bytes.get(cursor).copied() {
match byte {
b'0'..=b'9' | b';' | b':' => cursor += 1,
b'h' | b'l' => {
let params = &bytes[params_start..cursor];
let touches_alt = terminal_private_params_include_alt_screen(params);
return Some((cursor + 1, touches_alt.then_some(byte == b'h')));
}
0x40..=0x7e => return Some((cursor + 1, None)),
_ => return None,
}
}
None
}
fn terminal_private_params_include_alt_screen(params: &[u8]) -> bool {
params
.split(|byte| matches!(byte, b';' | b':'))
.any(|param| matches!(param, b"47" | b"1047" | b"1049"))
}
/// One speculatively-echoed character on the current line.
#[derive(Clone, Debug)]
struct PredictedCell {
/// The byte we drew (always a single printable ASCII byte, 0x20..=0x7e).
byte: u8,
/// Epoch this prediction belongs to (a keystroke burst).
epoch: u64,
}
/// Predictive local-echo engine. See the module comment above for the model.
struct Predictor {
mode: PredictMode,
enabled: bool,
/// True while the server is in the alternate screen (a full-screen TUI such
/// as vim/htop); we never speculate there because we cannot model arbitrary
/// cursor addressing safely.
alternate_screen: bool,
/// Tail of recent terminal output, retained so alternate-screen transitions
/// split across UDP frames are still detected.
output_parse_tail: Vec<u8>,
/// Predicted cells for the current line, left-to-right starting at the
/// column where the cursor sat when this run of predictions began.
cells: Vec<PredictedCell>,
/// Cursor offset, in cells, from the start of `cells`. Equals `cells.len()`
/// while typing at the end of the line; can be less after Left-arrow or
/// backspace within the predicted span.
cursor: usize,
/// Current keystroke-burst epoch. Bumped whenever we hit something we cannot
/// model (CR/LF, escape, control byte, wide char, right margin) so those
/// speculative cells become "tentative" and stop being drawn.
epoch: u64,
/// Highest epoch confirmed by authoritative server output. Cells whose epoch
/// is greater than this are tentative and are not displayed.
confirmed_epoch: u64,
/// How many columns the current overlay occupies on screen (0 = nothing
/// drawn). Used to erase exactly what was painted before re-rendering.
drawn_width: usize,
/// Whether the last draw underlined the cells (flagging on).
drawn_flagged: bool,
// --- SRTT estimation (display-only; never gates whether input is sent) ---
/// Smoothed round-trip time estimate in milliseconds, or None until we have
/// a sample.
srtt_ms: Option<f64>,
/// When the oldest still-outstanding prediction was made. Used both to
/// sample SRTT on the next frame and to force display on a latency spike.
oldest_pending_at: Option<Instant>,
/// Hysteresis latches.
srtt_trigger: bool,
flagging: bool,
}
impl Predictor {
/// Construct with the default adaptive (experimental) display policy.
#[cfg(test)]
fn new(enabled: bool) -> Self {
Self::with_mode(enabled, PredictMode::Experimental)
}
fn with_mode(enabled: bool, mode: PredictMode) -> Self {
Self {
mode,
enabled: enabled && mode != PredictMode::Off,
alternate_screen: false,
output_parse_tail: Vec::new(),
cells: Vec::new(),
cursor: 0,
epoch: 1,
confirmed_epoch: 0,
drawn_width: 0,
drawn_flagged: false,
srtt_ms: None,
oldest_pending_at: None,
srtt_trigger: false,
flagging: false,
}
}
/// Drop all speculative state. Called on reconnect/resume and screen resize,
/// where any cached line model would be stale.
fn reset(&mut self) {
let _ = self.erase_drawn();
self.cells.clear();
self.cursor = 0;
self.epoch += 1;
self.confirmed_epoch = self.epoch - 1;
self.alternate_screen = false;
self.output_parse_tail.clear();
self.oldest_pending_at = None;
}
/// Number of speculative cells currently outstanding (active in any epoch).
#[cfg(test)]
fn pending_len(&self) -> usize {
self.cells.len()
}
/// End the current keystroke burst (Mosh's "become tentative"). In dosh's
/// frame-confirmation model we cannot content-verify lingering cells, so the
/// safe choice is to drop the current epoch's speculative cells outright and
/// start a fresh epoch. This is how we bail out on anything we cannot model
/// (escape sequences, CR/LF, control bytes, wide chars, the right margin)
/// without risking a stale/wrong glyph.
fn become_tentative(&mut self) {
let _ = self.erase_drawn();
self.cells.clear();
self.cursor = 0;
self.epoch += 1;
self.oldest_pending_at = None;
}
/// Observe raw input bytes that are *also* being sent to the server. This is
/// display-only and never consumes input. The caller still forwards `bytes`
/// unconditionally.
fn observe_input(&mut self, bytes: &[u8]) -> Result<()> {
if !self.enabled || self.alternate_screen {
return Ok(());
}
// A large burst is almost certainly a paste, not interactive typing.
// Predicting it would risk wrapping a long dim run across many lines, so
// we bail: drop any overlay and let the authoritative frame draw it.
if bytes.len() > PREDICT_BURST_LIMIT {
self.become_tentative();
return self.redraw();
}
let mut i = 0;
while i < bytes.len() {
let b = bytes[i];
match b {
// Printable ASCII: append a predicted cell at the cursor.
0x20..=0x7e => {
self.predict_char(b);
i += 1;
}
// Backspace (^H) or DEL (^?): remove the previous predicted cell.
0x08 | 0x7f => {
self.predict_backspace();
i += 1;
}
// ESC: try to recognize a cursor-motion CSI/SS3 we can model
// (Left/Right within the line); otherwise become tentative.
0x1b => {
let consumed = self.predict_escape(&bytes[i..]);
if consumed == 0 {
// Unrecognized / incomplete escape: bail safely.
self.become_tentative();
i = bytes.len();
} else {
i += consumed;
}
}
// CR / LF and any other control byte: we cannot model where the
// cursor lands (shell prompt, scroll, etc.), so end the epoch.
_ => {
self.become_tentative();
i += 1;
}
}
}
if self.oldest_pending_at.is_none() && !self.cells.is_empty() {
self.oldest_pending_at = Some(Instant::now());
}
self.redraw()
}
/// Predict one printable character at the current cursor position.
fn predict_char(&mut self, byte: u8) {
// Right margin is tricky (shells, editors, and wrap behavior differ), so
// we refuse to predict past a conservative column budget and bail out
// instead of risking a wrong glyph at a wrap point.
if self.cells.len() >= PREDICT_MAX_LINE {
self.become_tentative();
return;
}
let cell = PredictedCell {
byte,
epoch: self.epoch,
};
if self.cursor < self.cells.len() {
// Typing over an existing predicted cell (after a Left arrow):
// overwrite in place, as a terminal in insert-off mode would.
self.cells[self.cursor] = cell;
} else {
self.cells.push(cell);
}
self.cursor += 1;
}
/// Predict a backspace: erase the previous predicted cell if we have one.
fn predict_backspace(&mut self) {
if self.cursor == 0 {
// We'd be backspacing past the start of our predicted span into
// territory we don't model -> bail safely.
self.become_tentative();
return;
}
self.cursor -= 1;
// Only safe to shrink the line if we're at its end; otherwise a
// mid-line backspace shifts everything left, which we don't model, so
// bail.
if self.cursor == self.cells.len() - 1 {
self.cells.pop();
} else {
self.become_tentative();
}
}
/// Recognize Left/Right cursor motion escape sequences we can model and
/// apply them to the local cursor. Returns the number of bytes consumed, or
/// 0 if `data` does not start with a sequence we handle.
///
/// Improvement over the previous dosh predictor (which bailed on any escape)
/// and parity with Mosh: predicting in-line arrow-key motion so cursor
/// repositioning feels instant too.
fn predict_escape(&mut self, data: &[u8]) -> usize {
// Both CSI (ESC [) and application-cursor SS3 (ESC O) are used for arrows.
if data.len() < 3 || data[0] != 0x1b {
return 0;
}
if data[1] != b'[' && data[1] != b'O' {
return 0;
}
match data[2] {
// Right arrow: only within already-predicted cells.
b'C' if self.cursor < self.cells.len() => {
self.cursor += 1;
3
}
// Left arrow.
b'D' if self.cursor > 0 => {
self.cursor -= 1;
3
}
// Recognized arrow but at the edge of our predicted span, or
// Up/Down/Home/End/etc. which cross lines or jump unpredictably:
// refuse so the caller becomes tentative rather than mispredicting.
_ => 0,
}
}
/// Observe authoritative server output. Detects alternate-screen transitions
/// and confirms outstanding predictions: a fresh frame means the server has
/// re-rendered the line, so our speculative glyphs are superseded and erased.
fn observe_output(&mut self, bytes: &[u8]) {
let mut parse = Vec::with_capacity(self.output_parse_tail.len() + bytes.len());
parse.extend_from_slice(&self.output_parse_tail);
parse.extend_from_slice(bytes);
let before = self.alternate_screen;
self.alternate_screen = alternate_screen_after_output(self.alternate_screen, &parse);
if self.alternate_screen != before {
let _ = self.discard_all();
}
self.output_parse_tail.clear();
let keep = parse.len().min(TERMINAL_OUTPUT_PARSE_TAIL);
self.output_parse_tail
.extend_from_slice(&parse[parse.len().saturating_sub(keep)..]);
}
/// Sample SRTT from a newly arrived frame and confirm/clear outstanding
/// predictions. Called once per accepted authoritative frame, *before* the
/// frame's bytes are written to the terminal so the server render lands on a
/// clean line. Returns nothing; drawing is handled by the caller's render
/// followed by `redraw` on the next input.
fn confirm_with_frame(&mut self) -> Result<()> {
if let Some(sent_at) = self.oldest_pending_at.take() {
let sample = sent_at.elapsed().as_secs_f64() * 1000.0;
self.update_srtt(sample);
}
// The frame is authoritative for the line; clear our overlay so the
// server's bytes are the only thing on screen. This is dosh's
// confirmation: frame arrival == the predicted keystrokes are now echoed
// by the server. Faster than waiting for per-cell content match.
self.discard_all()
}
fn update_srtt(&mut self, sample_ms: f64) {
// Standard exponentially weighted moving average (RFC 6298 style alpha).
const ALPHA: f64 = 0.125;
self.srtt_ms = Some(match self.srtt_ms {
None => sample_ms,
Some(prev) => (1.0 - ALPHA) * prev + ALPHA * sample_ms,
});
}
/// Test seam: pin the SRTT estimate and refresh the display triggers so the
/// experimental-mode threshold behavior can be exercised deterministically.
#[cfg(test)]
fn set_srtt_for_test(&mut self, ms: f64) {
self.srtt_ms = Some(ms);
self.update_triggers();
}
/// Test seam: backdate the oldest outstanding prediction so the glitch-force
/// display path (latency spike) can be exercised without real time passing.
#[cfg(test)]
fn backdate_pending_for_test(&mut self, ms: u64) {
self.oldest_pending_at = Some(Instant::now() - Duration::from_millis(ms));
}
/// Test seam: whether an overlay is currently painted on screen.
#[cfg(test)]
fn is_drawn(&self) -> bool {
self.drawn_width > 0
}
/// Whether we should *display* predictions right now under the active policy.
fn should_display(&self) -> bool {
match self.mode {
PredictMode::Off => false,
PredictMode::Always => true,
PredictMode::Experimental => {
if self.srtt_trigger {
return true;
}
// Force display on a latency spike even before SRTT catches up.
if let Some(at) = self.oldest_pending_at
&& at.elapsed().as_millis() >= GLITCH_FORCE_MS
{
return true;
}
false
}
}
}
/// Update the SRTT/flag hysteresis latches from the current estimate.
fn update_triggers(&mut self) {
let srtt = self.srtt_ms.unwrap_or(0.0);
if srtt > SRTT_TRIGGER_HIGH_MS {
self.srtt_trigger = true;
} else if srtt <= SRTT_TRIGGER_LOW_MS {
self.srtt_trigger = false;
}
if srtt > FLAG_TRIGGER_HIGH_MS {
self.flagging = true;
} else if srtt <= FLAG_TRIGGER_LOW_MS {
self.flagging = false;
}
}
/// The active (displayable) predicted bytes: cells in confirmed-or-current
/// epochs only, drawn from the start of the line up to the cursor budget.
fn visible_bytes(&self) -> Vec<u8> {
self.cells
.iter()
.filter(|c| c.epoch > self.confirmed_epoch || c.epoch == self.epoch)
.map(|c| c.byte)
.collect()
}
/// Repaint the prediction overlay: erase the old one, draw the new one if
/// policy says so. Uses save/restore cursor so it never moves the real
/// cursor and, when nothing should show, leaves the terminal untouched.
fn redraw(&mut self) -> Result<()> {
self.update_triggers();
// Always erase whatever we drew before so we never leave a stale glyph.
self.erase_drawn()?;
if !self.should_display() {
return Ok(());
}
let bytes = self.visible_bytes();
if bytes.is_empty() {
return Ok(());
}
let flag = self.flagging;
// Save cursor, dim (and optionally underline) predicted glyphs, draw,
// reset attributes, restore cursor.
let mut out = Vec::with_capacity(bytes.len() + 12);
out.extend_from_slice(b"\x1b7");
out.extend_from_slice(if flag { b"\x1b[2;4m" } else { b"\x1b[2m" });
out.extend_from_slice(&bytes);
out.extend_from_slice(b"\x1b[0m\x1b8");
emit_overlay(&out)?;
// Record exactly how many columns we painted so erasure overwrites the
// right amount even if the model changes before the next repaint.
self.drawn_width = bytes.len();
self.drawn_flagged = flag;
Ok(())
}
/// Re-evaluate the display policy on a select-loop timer tick (not just on a
/// keystroke). The experimental policy can decide to *start* showing purely
/// with the passage of time — `should_display` force-shows once a prediction
/// has been outstanding longer than `GLITCH_FORCE_MS` — so without this, a
/// sudden latency spike would only surface predictions on the *next* key,
/// one event late. Calling this from the status tick flips speculation on (or
/// off, or re-flags it) promptly. Repaints only when the on-screen result
/// would actually change, so an idle tick is a cheap no-op (no flicker).
fn refresh_policy(&mut self) -> Result<()> {
if !self.enabled || self.alternate_screen {
return Ok(());
}
self.update_triggers();
let want = self.should_display() && !self.visible_bytes().is_empty();
let have = self.drawn_width > 0;
let flag_changed = want && have && self.flagging != self.drawn_flagged;
if want != have || flag_changed {
self.redraw()?;
}
Ok(())
}
/// Erase the currently drawn overlay (if any) by overwriting the previously
/// painted columns with spaces, using save/restore cursor so the real cursor
/// and surrounding text are untouched.
fn erase_drawn(&mut self) -> Result<()> {
let width = self.drawn_width;
self.drawn_width = 0;
if width == 0 {
return Ok(());
}
let mut out = Vec::with_capacity(width + 4);
out.extend_from_slice(b"\x1b7");
out.resize(out.len() + width, b' ');
out.extend_from_slice(b"\x1b8");
emit_overlay(&out)
}
/// Erase the overlay and drop all speculative state, advancing the confirmed
/// epoch so nothing lingers.
fn discard_all(&mut self) -> Result<()> {
self.erase_drawn()?;
self.cells.clear();
self.cursor = 0;
self.epoch += 1;
self.confirmed_epoch = self.epoch - 1;
self.oldest_pending_at = None;
Ok(())
}
/// Erase the overlay without dropping the SRTT estimate or model; used right
/// before an authoritative frame is rendered so its bytes land cleanly.
/// (Kept as the public name the run loop calls.)
fn clear_pending(&mut self) -> Result<()> {
self.confirm_with_frame()
}
/// Whether a byte sequence is a plain printable run. Retained as a small
/// predicate exercised by the unit tests.
#[cfg(test)]
fn predictable(bytes: &[u8]) -> bool {
!bytes.is_empty() && bytes.iter().all(|byte| matches!(byte, 0x20..=0x7e))
}
}
/// Conservative cap on how many predicted cells we keep on one line before we
/// stop speculating (a backstop against runaway predictions; the right-margin /
/// wrap point is genuinely ambiguous without a client-side emulator).
const PREDICT_MAX_LINE: usize = 256;
/// Largest single input burst we treat as interactive typing. Anything larger
/// is assumed to be a paste and is not speculated (would risk a long dim run
/// wrapping across lines before the server frame supersedes it).
const PREDICT_BURST_LIMIT: usize = 32;
/// Write a prepared overlay byte sequence to the real terminal. Suppressed
/// under `cfg(test)` so unit tests exercise the prediction model without
/// emitting escape sequences to the test runner's terminal.
#[cfg(not(test))]
fn emit_overlay(bytes: &[u8]) -> Result<()> {
let mut stdout = std::io::stdout();
stdout.write_all(bytes)?;
stdout.flush()?;
Ok(())
}
#[cfg(test)]
fn emit_overlay(_bytes: &[u8]) -> Result<()> {
Ok(())
}
/// Resolve the prediction display policy. `DOSH_PREDICT_MODE` (off /
/// experimental / always) overrides for quick experiments; otherwise the
/// already-selected CLI/config policy is used.
fn selected_predict_mode(args: &Args, config: &dosh::config::ClientConfig) -> Result<PredictMode> {
if args.predict_always && args.predict_never {
return Err(anyhow!("--predict-always and --predict-never conflict"));
}
if args.predict_always {
return Ok(PredictMode::Always);
}
if args.predict_never {
return Ok(PredictMode::Off);
}
if let Some(mode) = args.predict_mode.as_deref() {
return Ok(PredictMode::parse(mode));
}
if let Ok(env) = std::env::var("DOSH_PREDICT_MODE") {
return Ok(PredictMode::parse(&env));
}
Ok(PredictMode::parse(&config.predict_mode))
}
fn selected_escape_key(
args: &Args,
config: &dosh::config::ClientConfig,
) -> Result<Option<Vec<u8>>> {
if let Some(raw) = args.escape_key.as_deref() {
return parse_escape_key(raw);
}
if let Ok(raw) = std::env::var("DOSH_ESCAPE_KEY") {
return parse_escape_key(&raw);
}
parse_escape_key(&config.escape_key)
}
fn parse_escape_key(raw: &str) -> Result<Option<Vec<u8>>> {
let value = raw.trim();
if value.eq_ignore_ascii_case("none")
|| value.eq_ignore_ascii_case("off")
|| value.eq_ignore_ascii_case("disabled")
{
return Ok(None);
}
if let Some(rest) = value.strip_prefix("^") {
anyhow::ensure!(
rest.chars().count() == 1,
"invalid control escape key {raw:?}"
);
let ch = rest.as_bytes()[0];
let byte = match ch {
b'?' => 0x7f,
b'@'..=b'_' => ch - b'@',
b'a'..=b'z' => ch - b'a' + 1,
_ => return Err(anyhow!("invalid control escape key {raw:?}")),
};
return Ok(Some(vec![byte]));
}
if let Some(hex) = value.strip_prefix("0x") {
let byte = u8::from_str_radix(hex, 16)
.with_context(|| format!("invalid hex escape key {raw:?}"))?;
return Ok(Some(vec![byte]));
}
if value.len() == 1 {
return Ok(Some(value.as_bytes().to_vec()));
}
Err(anyhow!(
"escape key must be none, ^X, 0xNN, or one ASCII character"
))
}
fn input_matches_escape(bytes: &[u8], escape_key: Option<&[u8]>) -> bool {
let Some(escape_key) = escape_key else {
return false;
};
!escape_key.is_empty() && bytes.windows(escape_key.len()).any(|w| w == escape_key)
}
/// Resolve whether the non-destructive disconnect status line is enabled.
/// `DOSH_DISCONNECT_STATUS` (0/false/off to disable) overrides for quick
/// experiments; otherwise the client config's `disconnect_status` applies,
/// defaulting to on.
fn resolve_disconnect_status() -> bool {
if let Ok(env) = std::env::var("DOSH_DISCONNECT_STATUS") {
let v = env.trim().to_ascii_lowercase();
return !matches!(v.as_str(), "0" | "false" | "off" | "no");
}
match load_client_config(None) {
Ok(cfg) => cfg.disconnect_status,
Err(_) => true,
}
}
fn contains_bytes(haystack: &[u8], needle: &[u8]) -> bool {
haystack
.windows(needle.len())
.any(|window| window == needle)
}
const FRAME_GAP_RESYNC_AFTER_MS: u64 = 250;
const FRAME_GAP_RESYNC_COOLDOWN_MS: u64 = 1_000;
const FRAME_GAP_PENDING_RESYNC_THRESHOLD: usize = 128;
#[derive(Default)]
struct FrameBuffer {
pending: BTreeMap<u64, Frame>,
gap_since: Option<Instant>,
last_resync_at: Option<Instant>,
}
impl FrameBuffer {
fn clear(&mut self) {
self.pending.clear();
self.gap_since = None;
self.last_resync_at = None;
}
fn accept(&mut self, frame: Frame, last_rendered_seq: &mut u64) -> Vec<Frame> {
if frame.snapshot {
if frame.output_seq < *last_rendered_seq {
return Vec::new();
}
self.pending.clear();
self.gap_since = None;
self.last_resync_at = None;
*last_rendered_seq = frame.output_seq;
return vec![frame];
}
if frame.output_seq <= *last_rendered_seq {
return Vec::new();
}
self.pending.entry(frame.output_seq).or_insert(frame);
let mut ready = Vec::new();
while let Some(frame) = self.pending.remove(&(*last_rendered_seq + 1)) {
*last_rendered_seq = frame.output_seq;
ready.push(frame);
}
self.refresh_gap(*last_rendered_seq);
ready
}
fn refresh_gap(&mut self, last_rendered_seq: u64) {
let waiting_on_missing = self
.pending
.first_key_value()
.is_some_and(|(&seq, _)| seq > last_rendered_seq.saturating_add(1));
if waiting_on_missing {
self.gap_since.get_or_insert_with(Instant::now);
} else {
self.gap_since = None;
self.last_resync_at = None;
}
}
fn resync_due(&mut self) -> bool {
let Some(gap_since) = self.gap_since else {
return false;
};
let now = Instant::now();
let waited = now.duration_since(gap_since);
if waited < Duration::from_millis(FRAME_GAP_RESYNC_AFTER_MS)
&& self.pending.len() < FRAME_GAP_PENDING_RESYNC_THRESHOLD
{
return false;
}
if self.last_resync_at.is_some_and(|last| {
now.duration_since(last) < Duration::from_millis(FRAME_GAP_RESYNC_COOLDOWN_MS)
}) {
return false;
}
self.last_resync_at = Some(now);
true
}
#[cfg(test)]
fn backdate_gap_for_test(&mut self, elapsed: Duration) {
self.gap_since = Some(Instant::now() - elapsed);
}
}
/// Read the current terminal size and, if it changed since `last_size`, send a
/// resize to the server. Shared by the SIGWINCH branch and the poll fallback.
async fn maybe_send_resize(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
last_size: &mut (u16, u16),
) -> Result<()> {
if let Ok((cols, rows)) = size()
&& cols > 0
&& rows > 0
&& (cols, rows) != *last_size
{
*last_size = (cols, rows);
if cred.mode != "view-only" && cred.mode != "forward-only" {
send_resize(socket, addr, cred, send_seq, cols, rows).await?;
}
}
Ok(())
}
async fn send_resize(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
cols: u16,
rows: u16,
) -> Result<()> {
let body = protocol::to_body(&Resize { cols, rows })?;
let packet = protocol::encode_encrypted(
PacketKind::Resize,
cred.client_id,
*send_seq,
cred.last_rendered_seq,
&cred.session_key,
CLIENT_TO_SERVER,
&body,
)?;
*send_seq += 1;
let _ = send_terminal_udp(socket, &packet, addr).await?;
Ok(())
}
async fn send_ack(
socket: &UdpSocket,
addr: SocketAddr,
cred: &CachedCredential,
send_seq: &mut u64,
) -> Result<()> {
let packet = protocol::encode_encrypted(
PacketKind::Ack,
cred.client_id,
*send_seq,
cred.last_rendered_seq,
&cred.session_key,
CLIENT_TO_SERVER,
b"",
)?;
*send_seq += 1;
let _ = send_terminal_udp(socket, &packet, addr).await?;
Ok(())
}
async fn detach_once(socket: &UdpSocket, cred: &CachedCredential, seq: u64) -> Result<()> {
let addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
let packet = protocol::encode_encrypted(
PacketKind::Detach,
cred.client_id,
seq,
cred.last_rendered_seq,
&cred.session_key,
CLIENT_TO_SERVER,
b"",
)?;
let _ = send_terminal_udp(socket, &packet, addr).await?;
Ok(())
}
fn render_frame(frame: &Frame) -> Result<()> {
let mut stdout = std::io::stdout();
stdout.write_all(&frame.bytes)?;
stdout.flush()?;
Ok(())
}
/// Seconds of silence from the server before the disconnect status line appears.
/// Short enough to give quick feedback on a lost link, long enough that a normal
/// idle period (the run loop only pings after 2s of quiet) never flashes it.
const DISCONNECT_STATUS_THRESHOLD_SECS: u64 = 2;
/// Non-destructive, mosh-style disconnect status bar.
///
/// When server packets stop arriving we paint one blue line on the terminal's
/// top row — e.g. `[dosh] reconnecting — last contact 3s ago` — using
/// save/restore cursor (ESC 7 / ESC 8) so the real application cursor never
/// moves. The run loop suppresses the bar while the server is in the alternate
/// screen because a saved cursor does not make writing row 1 non-destructive for
/// full-screen TUIs. The bar is refreshed on the status timer tick and cleared
/// the instant a packet resumes.
///
/// The decision logic and the rendered text are pure and unit-tested; only
/// [`DisconnectStatus::emit`] / [`DisconnectStatus::emit_clear`] touch stdout.
struct DisconnectStatus {
enabled: bool,
/// Whether a status line is currently painted on screen (so we know to clear
/// it and to avoid redundant repaints of identical text).
shown: bool,
/// The text currently painted, to skip no-op repaints (the elapsed seconds
/// only change once per second so most ticks are no-ops).
drawn_text: String,
}
/// What the run loop should do with the status line on a given tick or packet.
#[derive(Debug, Clone, PartialEq, Eq)]
enum StatusAction {
/// Leave the screen as-is (already correct).
None,
/// Paint/repaint the status bar with this exact text on the top row.
Show(String),
/// Erase the status line (link recovered or feature disabled).
Clear,
/// Drop tracked status state without touching the terminal.
Forget,
}
impl DisconnectStatus {
fn new(enabled: bool) -> Self {
Self {
enabled,
shown: false,
drawn_text: String::new(),
}
}
/// Pure: render the status text for a given elapsed silence. Kept separate
/// from any I/O so it can be asserted in tests.
fn render_text(elapsed: Duration) -> String {
let secs = elapsed.as_secs();
if secs >= 30 {
format!("[dosh] still reconnecting — last contact {secs}s ago")
} else {
format!("[dosh] reconnecting — last contact {secs}s ago")
}
}
/// Pure decision for a status timer tick: given how long the link has been
/// silent, decide whether to show (and with what text), clear, or do nothing.
/// Does not mutate screen state; the caller applies the returned action and
/// then calls [`Self::applied`].
fn on_tick(&self, stale: Duration) -> StatusAction {
if !self.enabled {
return if self.shown {
StatusAction::Clear
} else {
StatusAction::None
};
}
if stale >= Duration::from_secs(DISCONNECT_STATUS_THRESHOLD_SECS) {
let text = Self::render_text(stale);
if self.shown && text == self.drawn_text {
StatusAction::None
} else {
StatusAction::Show(text)
}
} else if self.shown {
StatusAction::Clear
} else {
StatusAction::None
}
}
/// Pure decision for "a packet just arrived": clear any visible status line
/// the instant contact resumes, otherwise do nothing.
fn on_contact(&self) -> StatusAction {
if self.shown {
StatusAction::Clear
} else {
StatusAction::None
}
}
/// Pure decision for periods where drawing the top-row status is unsafe
/// (alternate-screen TUIs). Forget any tracked bar without writing a clear
/// sequence, because clearing row 1 would corrupt the TUI just like showing
/// the bar would.
fn on_suppressed(&self) -> StatusAction {
if self.shown {
StatusAction::Forget
} else {
StatusAction::None
}
}
/// Record that an action was applied to the screen, updating internal state.
fn applied(&mut self, action: &StatusAction) {
match action {
StatusAction::Show(text) => {
self.shown = true;
self.drawn_text = text.clone();
}
StatusAction::Clear => {
self.shown = false;
self.drawn_text.clear();
}
StatusAction::Forget => {
self.shown = false;
self.drawn_text.clear();
}
StatusAction::None => {}
}
}
/// Apply an action to the real terminal (idempotent for `None`).
fn apply(&mut self, action: StatusAction) -> Result<()> {
match &action {
StatusAction::Show(text) => self.emit(text)?,
StatusAction::Clear => self.emit_clear()?,
StatusAction::Forget => {}
StatusAction::None => {}
}
self.applied(&action);
Ok(())
}
/// Paint `text` on the terminal's top row using save/restore cursor so the
/// app's cursor is untouched.
fn emit(&self, text: &str) -> Result<()> {
let (_, rows) = terminal_size();
let bytes = render_status_overlay(text, rows);
emit_status(&bytes)
}
/// Erase the top-row status bar, again with save/restore cursor.
fn emit_clear(&self) -> Result<()> {
let (_, rows) = terminal_size();
emit_status(&render_status_clear(rows))
}
}
/// Pure: build the escape sequence that paints `text` on a top blue status bar,
/// leaving the cursor where it was.
fn render_status_overlay(text: &str, _rows: u16) -> Vec<u8> {
let mut out = Vec::with_capacity(text.len() + 24);
out.extend_from_slice(b"\x1b7"); // save cursor
out.extend_from_slice(b"\x1b[1;1H"); // move to top row, column 1
out.extend_from_slice(b"\x1b[2K"); // clear entire row
out.extend_from_slice(b"\x1b[44;37m"); // blue background + white text
out.extend_from_slice(text.as_bytes());
out.extend_from_slice(b"\x1b[0m"); // reset attributes
out.extend_from_slice(b"\x1b8"); // restore cursor
out
}
/// Pure: build the escape sequence that clears the top status bar.
fn render_status_clear(_rows: u16) -> Vec<u8> {
let mut out = Vec::with_capacity(12);
out.extend_from_slice(b"\x1b7");
out.extend_from_slice(b"\x1b[1;1H");
out.extend_from_slice(b"\x1b[2K");
out.extend_from_slice(b"\x1b8");
out
}
#[cfg(not(test))]
fn emit_status(bytes: &[u8]) -> Result<()> {
let mut stdout = std::io::stdout();
stdout.write_all(bytes)?;
stdout.flush()?;
Ok(())
}
#[cfg(test)]
fn emit_status(_bytes: &[u8]) -> Result<()> {
Ok(())
}
fn cache_path(root: &str, server: &str, session: &str, mode: &str) -> std::path::PathBuf {
let safe = cache_key(server, session, mode);
expand_tilde(root).join(format!("{safe}.bin"))
}
fn cache_key(server: &str, session: &str, mode: &str) -> String {
format!("{server}_{session}_{mode}")
.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
.collect::<String>()
}
fn cache_server_prefix(server: &str) -> String {
format!("{server}_")
.chars()
.map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
.collect::<String>()
}
fn clear_cached_credentials(root: &str, server: &str) -> Result<usize> {
let dir = expand_tilde(root);
let Ok(entries) = fs::read_dir(&dir) else {
return Ok(0);
};
let prefix = cache_server_prefix(server);
let mut removed = 0usize;
for entry in entries {
let entry = entry?;
let path = entry.path();
if !path.is_file() {
continue;
}
let Some(name) = path.file_name().and_then(|value| value.to_str()) else {
continue;
};
if name.starts_with(&prefix) && name.ends_with(".bin") {
fs::remove_file(&path)?;
removed += 1;
}
}
Ok(removed)
}
fn load_cache(path: &std::path::Path) -> Result<CachedCredential> {
let raw = fs::read(path)?;
Ok(bincode::deserialize(&raw)?)
}
fn save_cache(path: &std::path::Path, cred: &CachedCredential) -> Result<()> {
if let Some(parent) = path.parent() {
fs::create_dir_all(parent)?;
}
fs::write(path, bincode::serialize(cred)?)?;
Ok(())
}
struct RawMode;
impl RawMode {
fn enter() -> Result<Self> {
flush_local_terminal_input();
enable_raw_mode()?;
flush_local_terminal_input();
drain_local_terminal_input();
Ok(Self)
}
}
#[cfg(unix)]
fn flush_local_terminal_input() {
unsafe {
libc::tcflush(libc::STDIN_FILENO, libc::TCIFLUSH);
}
}
#[cfg(not(unix))]
fn flush_local_terminal_input() {}
#[cfg(unix)]
fn drain_local_terminal_input() {
unsafe {
let fd = libc::STDIN_FILENO;
let flags = libc::fcntl(fd, libc::F_GETFL);
if flags < 0 {
return;
}
if libc::fcntl(fd, libc::F_SETFL, flags | libc::O_NONBLOCK) < 0 {
return;
}
let mut buf = [0u8; 1024];
loop {
let n = libc::read(fd, buf.as_mut_ptr().cast(), buf.len());
if n <= 0 {
break;
}
}
let _ = libc::fcntl(fd, libc::F_SETFL, flags);
}
}
#[cfg(not(unix))]
fn drain_local_terminal_input() {}
impl Drop for RawMode {
fn drop(&mut self) {
let mut stdout = std::io::stdout();
let _ = stdout.write_all(TERMINAL_CLEANUP);
let _ = stdout.flush();
let _ = disable_raw_mode();
}
}
const TERMINAL_CLEANUP: &[u8] = concat!(
"\x1b[?1l",
"\x1b[0m",
"\x1b[?25h",
"\x1b[?1003l",
"\x1b[?1002l",
"\x1b[?1001l",
"\x1b[?1000l",
"\x1b[?1004l",
"\x1b[?1015l",
"\x1b[?1006l",
"\x1b[?1005l",
"\x1b[?2004l",
"\x1b[?47l",
"\x1b[?1047l",
"\x1b[?1049l"
)
.as_bytes();
#[cfg(test)]
mod tests {
use super::{
ALT_SCREEN_IDLE_REPAINT_AFTER, CachedCredential, DisconnectStatus, DynamicForward,
FRAME_GAP_RESYNC_AFTER_MS, FrameBuffer, LocalForward, MAX_PENDING_USER_INPUT_BYTES,
POST_SUBMIT_ALL_INPUT_HOLD, PendingStreamOpen, PredictMode, Predictor,
RESTART_STATUS_SCRIPT, RemoteForward, STARTUP_INPUT_HOLD, SshConfig, StartupGateMode,
StatusAction, auth_allows, cache_key, cache_server_prefix, cleanup_stream_state,
clear_cached_credentials, ensure_tui_safe_status_overlay, input_contains_focus_in,
input_matches_escape, is_local_status_target, is_resume_response_for_client,
latest_release_download_url, load_first_native_identity_with_prompt, parse_dynamic_forward,
parse_escape_key, parse_local_forward, parse_remote_forward, parse_ssh_config,
post_submit_hold_duration, queue_pending_user_input, raw_contains_host_table,
recv_response_until, refresh_live_addr, release_tag_download_url,
release_tag_from_effective_url, release_version_from_tag, render_status_clear,
render_status_overlay, requested_env, resolved_startup_command, retire_stream_state,
retransmit_stream_opens, rewrite_forward_command, selected_predict_mode, selected_udp_host,
server_version_mismatch, should_flush_terminal_input_after_contact,
should_hold_during_startup_gate, should_hold_post_submit_input,
should_repaint_idle_alternate_screen, split_after_command_submit, ssh_destination_host,
ssh_username, ssh_with_user, startup_command, status_ssh_target, strip_stale_mouse_reports,
strip_terminal_focus_reports, terminal_private_mode_transition, toml_bare_key_or_quoted,
update_check_requested, update_version_status, upsert_managed_block, valid_forward_host,
vscode_safe_alias,
};
use dosh::config::{ClientConfig, CommandExtension, HostConfig};
use dosh::native::EnvVar;
use dosh::protocol::{self, Frame, PacketKind};
use dosh::udp::is_transient_udp_send_error;
use ed25519_dalek::{SigningKey, VerifyingKey};
use std::collections::{BTreeMap, HashMap, HashSet, VecDeque};
use std::fs;
use std::time::{Duration, Instant};
fn test_args(server: &str, command: &[&str]) -> super::Args {
super::Args {
server: Some(server.to_string()),
command: command.iter().map(|value| value.to_string()).collect(),
session: None,
new: false,
view_only: false,
local_auth: false,
auth: None,
no_cache: false,
remove: false,
replace: false,
attach_only: false,
local_forward: Vec::new(),
remote_forward: Vec::new(),
dynamic_forward: Vec::new(),
forward_agent: false,
forward_only: false,
background: false,
predict: false,
predict_mode: None,
predict_always: false,
predict_never: false,
escape_key: None,
ssh_port: None,
ssh_auth_command: None,
ssh_key: None,
ssh_known_hosts: None,
ssh_control_path: None,
dosh_port: None,
dosh_host: None,
verbose: 0,
}
}
#[test]
fn vscode_alias_sanitizes_host_names() {
assert_eq!(vscode_safe_alias("palav"), "palav");
assert_eq!(vscode_safe_alias("user@example.com"), "user-example-com");
assert_eq!(vscode_safe_alias(":///"), "host");
}
#[test]
fn cleanup_stream_state_removes_all_per_stream_state_only_for_target() {
let stream_id = 42;
let other_stream_id = 43;
let mut opened_streams = std::collections::HashSet::from([stream_id, other_stream_id]);
let mut stream_send_credit = HashMap::from([(stream_id, 12), (other_stream_id, 34)]);
let mut stream_pending_data = HashMap::from([
(stream_id, VecDeque::from([vec![1]])),
(other_stream_id, VecDeque::from([vec![2]])),
]);
let mut stream_pending_opens = HashMap::from([
(
stream_id,
PendingStreamOpen {
target_host: "127.0.0.1".to_string(),
target_port: 1234,
last_sent: Instant::now(),
attempts: 1,
},
),
(
other_stream_id,
PendingStreamOpen {
target_host: "127.0.0.1".to_string(),
target_port: 1235,
last_sent: Instant::now(),
attempts: 1,
},
),
]);
let mut stream_next_send_offset = HashMap::from([(stream_id, 55), (other_stream_id, 66)]);
let mut stream_sent_data = HashMap::from([
(
stream_id,
BTreeMap::from([(
0,
super::PendingStreamChunk {
offset: 0,
bytes: vec![3],
last_sent: Instant::now(),
attempts: 1,
},
)]),
),
(
other_stream_id,
BTreeMap::from([(
0,
super::PendingStreamChunk {
offset: 0,
bytes: vec![4],
last_sent: Instant::now(),
attempts: 1,
},
)]),
),
]);
let mut stream_next_recv_offset = HashMap::from([(stream_id, 77), (other_stream_id, 88)]);
let mut stream_recv_pending = HashMap::from([
(stream_id, BTreeMap::from([(77, vec![5])])),
(other_stream_id, BTreeMap::from([(88, vec![6])])),
]);
cleanup_stream_state(
stream_id,
&mut opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_pending_opens,
&mut stream_next_send_offset,
&mut stream_sent_data,
&mut stream_next_recv_offset,
&mut stream_recv_pending,
);
assert!(!opened_streams.contains(&stream_id));
assert!(!stream_send_credit.contains_key(&stream_id));
assert!(!stream_pending_data.contains_key(&stream_id));
assert!(!stream_pending_opens.contains_key(&stream_id));
assert!(!stream_next_send_offset.contains_key(&stream_id));
assert!(!stream_sent_data.contains_key(&stream_id));
assert!(!stream_next_recv_offset.contains_key(&stream_id));
assert!(!stream_recv_pending.contains_key(&stream_id));
assert!(opened_streams.contains(&other_stream_id));
assert!(stream_send_credit.contains_key(&other_stream_id));
assert!(stream_pending_data.contains_key(&other_stream_id));
assert!(stream_pending_opens.contains_key(&other_stream_id));
assert!(stream_next_send_offset.contains_key(&other_stream_id));
assert!(stream_sent_data.contains_key(&other_stream_id));
assert!(stream_next_recv_offset.contains_key(&other_stream_id));
assert!(stream_recv_pending.contains_key(&other_stream_id));
}
#[test]
fn retire_stream_state_tombstones_only_real_stream_state() {
let stream_id = 42;
let mut opened_streams = HashSet::from([stream_id]);
let mut stream_send_credit = HashMap::from([(stream_id, 12)]);
let mut stream_pending_data = HashMap::new();
let mut stream_pending_opens = HashMap::new();
let mut stream_next_send_offset = HashMap::new();
let mut stream_sent_data = HashMap::new();
let mut stream_next_recv_offset = HashMap::new();
let mut stream_recv_pending = HashMap::new();
let mut stream_retired = HashSet::new();
let mut stream_retired_order = VecDeque::new();
retire_stream_state(
stream_id,
&mut opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_pending_opens,
&mut stream_next_send_offset,
&mut stream_sent_data,
&mut stream_next_recv_offset,
&mut stream_recv_pending,
&mut stream_retired,
&mut stream_retired_order,
);
retire_stream_state(
99,
&mut opened_streams,
&mut stream_send_credit,
&mut stream_pending_data,
&mut stream_pending_opens,
&mut stream_next_send_offset,
&mut stream_sent_data,
&mut stream_next_recv_offset,
&mut stream_recv_pending,
&mut stream_retired,
&mut stream_retired_order,
);
assert!(!opened_streams.contains(&stream_id));
assert!(!stream_send_credit.contains_key(&stream_id));
assert!(stream_retired.contains(&stream_id));
assert!(!stream_retired.contains(&99));
assert_eq!(stream_retired_order, VecDeque::from([stream_id]));
}
#[test]
fn managed_ssh_block_is_replaced_not_duplicated() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("config.dosh");
upsert_managed_block(
&path,
"dosh-palav",
"# BEGIN DOSH dosh-palav\nHost dosh-palav\n HostName 127.0.0.1\n# END DOSH dosh-palav\n",
)
.unwrap();
upsert_managed_block(
&path,
"dosh-palav",
"# BEGIN DOSH dosh-palav\nHost dosh-palav\n HostName localhost\n# END DOSH dosh-palav\n",
)
.unwrap();
let raw = fs::read_to_string(path).unwrap();
assert_eq!(raw.matches("# BEGIN DOSH dosh-palav").count(), 1);
assert!(raw.contains("HostName localhost"));
assert!(!raw.contains("HostName 127.0.0.1"));
}
#[test]
fn parses_ssh_config_hostname() {
let raw = "user alice\nhostname server.example.com\nport 22\n";
assert_eq!(
parse_ssh_config(raw).hostname,
Some("server.example.com".to_string())
);
}
#[test]
fn parses_ssh_config_hostname_case_insensitively() {
let raw = "HostName 192.0.2.10\n";
assert_eq!(
parse_ssh_config(raw).hostname,
Some("192.0.2.10".to_string())
);
}
#[test]
fn parses_ssh_config_port() {
let parsed = parse_ssh_config("hostname example.com\nport 2222\n");
assert_eq!(parsed.hostname, Some("example.com".to_string()));
assert_eq!(parsed.port, Some(2222));
}
#[test]
fn parses_ssh_config_user_and_identity_files() {
let parsed = parse_ssh_config(
"user alice\nidentityfile ~/.ssh/id_ed25519\nidentityfile ~/.ssh/work\n",
);
assert_eq!(parsed.user, Some("alice".to_string()));
assert_eq!(
parsed.identity_files,
vec!["~/.ssh/id_ed25519".to_string(), "~/.ssh/work".to_string()]
);
}
#[test]
fn parses_ssh_config_native_parity_options() {
let parsed = parse_ssh_config(
"userknownhostsfile ~/.ssh/known_hosts ~/.ssh/work_hosts\n\
identitiesonly yes\n\
proxyjump bastion\n\
proxycommand ssh bastion -W %h:%p\n\
identityfile none\n\
sendenv LANG LC_*\n\
setenv DOSH_MODE=native DOSH_COLOR=true\n",
);
assert_eq!(
parsed.user_known_hosts_file,
Some("~/.ssh/known_hosts ~/.ssh/work_hosts".to_string())
);
assert!(parsed.identities_only);
assert_eq!(parsed.proxy_jump, Some("bastion".to_string()));
assert_eq!(
parsed.proxy_command,
Some("ssh bastion -W %h:%p".to_string())
);
assert_eq!(
parsed.send_env,
vec!["LANG".to_string(), "LC_*".to_string()]
);
assert_eq!(
parsed.set_env,
vec![
EnvVar {
name: "DOSH_MODE".to_string(),
value: "native".to_string()
},
EnvVar {
name: "DOSH_COLOR".to_string(),
value: "true".to_string()
}
]
);
assert!(parsed.identity_files.is_empty());
}
#[test]
fn requested_env_merges_config_host_and_ssh_setenv() {
let mut config = ClientConfig::default();
config.send_env.clear();
config.set_env.insert("DOSH_MODE".into(), "client".into());
config.set_env.insert("DOSH_KEEP".into(), "yes".into());
let mut host = HostConfig::default();
host.set_env.insert("DOSH_MODE".into(), "host".into());
let ssh_config = SshConfig {
set_env: vec![EnvVar {
name: "DOSH_SSH".to_string(),
value: "true".to_string(),
}],
..SshConfig::default()
};
assert_eq!(
requested_env(&config, &host, Some(&ssh_config)),
vec![
EnvVar {
name: "DOSH_KEEP".to_string(),
value: "yes".to_string()
},
EnvVar {
name: "DOSH_MODE".to_string(),
value: "host".to_string()
},
EnvVar {
name: "DOSH_SSH".to_string(),
value: "true".to_string()
}
]
);
}
#[test]
fn detects_existing_imported_host_tables() {
assert!(raw_contains_host_table(
"[homelab]\nssh = \"homelab\"\n",
"homelab"
));
assert!(raw_contains_host_table("[\"home box\"]\n", "home box"));
assert_eq!(toml_bare_key_or_quoted("homelab"), "homelab");
assert_eq!(toml_bare_key_or_quoted("home box"), "\"home box\"");
}
#[test]
fn strips_user_from_fallback_udp_host() {
assert_eq!(ssh_destination_host("alice@example.com"), "example.com");
}
#[test]
fn selected_udp_host_supports_ssh_auto_and_localhost() {
let ssh = SshConfig {
hostname: Some("server.lan".to_string()),
..SshConfig::default()
};
assert_eq!(
selected_udp_host(None, "alias", &ssh).unwrap(),
"server.lan"
);
assert_eq!(
selected_udp_host(Some("ssh"), "alias", &ssh).unwrap(),
"server.lan"
);
assert_eq!(
selected_udp_host(Some("auto"), "alias", &ssh).unwrap(),
"server.lan"
);
assert_eq!(
selected_udp_host(Some("localhost"), "alias", &ssh).unwrap(),
"127.0.0.1"
);
assert!(selected_udp_host(Some("any"), "alias", &ssh).is_err());
}
#[test]
fn parses_user_from_ssh_destination() {
assert_eq!(ssh_username("alice@example.com"), Some("alice".to_string()));
assert_eq!(ssh_username("example.com"), None);
}
#[test]
fn applies_host_config_user_to_plain_ssh_destination() {
assert_eq!(
ssh_with_user("example.com", Some("alice")),
"alice@example.com"
);
assert_eq!(
ssh_with_user("root@example.com", Some("alice")),
"root@example.com"
);
assert_eq!(
ssh_with_user("ssh://example.com", Some("alice")),
"ssh://example.com"
);
}
#[test]
fn auth_preference_accepts_auto_or_named_method() {
assert!(auth_allows("native,ssh", "native"));
assert!(auth_allows("native,ssh", "ssh"));
assert!(auth_allows("auto", "native"));
assert!(!auth_allows("native", "ssh"));
}
#[test]
fn encrypted_identity_file_prompts_and_loads() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("id_ed25519");
let signing = SigningKey::from_bytes(&[61u8; 32]);
write_encrypted_identity(&path, &signing, "let-me-in");
let mut config = ClientConfig::default();
config.identity_files.clear();
let loaded = load_first_native_identity_with_prompt(
&config,
Some(&path),
&SshConfig::default(),
|_| Ok("let-me-in".to_string()),
)
.unwrap();
assert_eq!(
loaded.public_key().key_data().ed25519().unwrap().as_ref(),
VerifyingKey::from(&signing).as_bytes()
);
}
#[test]
fn encrypted_identity_file_reports_wrong_passphrase() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("id_ed25519");
let signing = SigningKey::from_bytes(&[62u8; 32]);
write_encrypted_identity(&path, &signing, "correct");
let mut config = ClientConfig::default();
config.identity_files.clear();
let err = load_first_native_identity_with_prompt(
&config,
Some(&path),
&SshConfig::default(),
|_| Ok("wrong".to_string()),
)
.unwrap_err();
assert!(err.to_string().contains("no usable native identity"));
}
#[test]
fn identities_only_skips_dosh_config_identity_fallback() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("id_ed25519");
let signing = SigningKey::from_bytes(&[63u8; 32]);
write_identity(&path, &signing);
let config = ClientConfig {
identity_files: vec![path.display().to_string()],
..ClientConfig::default()
};
let ssh_config = SshConfig {
identities_only: true,
..SshConfig::default()
};
let err = load_first_native_identity_with_prompt(&config, None, &ssh_config, |_| {
unreachable!("unencrypted key should not prompt")
})
.unwrap_err();
assert!(err.to_string().contains("no usable native identity"));
}
#[test]
fn identities_only_still_allows_ssh_config_identity_files() {
let dir = tempfile::tempdir().unwrap();
let path = dir.path().join("id_ed25519");
let signing = SigningKey::from_bytes(&[64u8; 32]);
write_identity(&path, &signing);
let mut config = ClientConfig::default();
config.identity_files.clear();
let ssh_config = SshConfig {
identity_files: vec![path.display().to_string()],
identities_only: true,
..SshConfig::default()
};
let loaded = load_first_native_identity_with_prompt(&config, None, &ssh_config, |_| {
unreachable!("unencrypted key should not prompt")
})
.unwrap();
assert_eq!(
loaded.public_key().key_data().ed25519().unwrap().as_ref(),
VerifyingKey::from(&signing).as_bytes()
);
}
#[tokio::test]
async fn recv_response_until_ignores_unmatched_packets() {
let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let sender = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let addr = receiver.local_addr().unwrap();
let stale = protocol::encode_plain(PacketKind::Ping, [1u8; 16], 1, 0, b"stale").unwrap();
let valid = protocol::encode_plain(PacketKind::Pong, [2u8; 16], 2, 0, b"valid").unwrap();
sender.send_to(&stale, addr).await.unwrap();
sender.send_to(&valid, addr).await.unwrap();
let packet = recv_response_until(&receiver, std::time::Duration::from_secs(1), |packet| {
if packet.header.kind == PacketKind::Pong {
Ok(Some(packet))
} else {
Ok(None)
}
})
.await
.unwrap();
assert_eq!(packet.header.kind, PacketKind::Pong);
assert_eq!(packet.body, b"valid");
}
fn write_identity(path: &std::path::Path, signing: &SigningKey) {
let keypair = ssh_key::private::Ed25519Keypair::from(signing);
let private =
ssh_key::PrivateKey::new(ssh_key::private::KeypairData::from(keypair), "test").unwrap();
private
.write_openssh_file(path, ssh_key::LineEnding::LF)
.unwrap();
}
fn write_encrypted_identity(path: &std::path::Path, signing: &SigningKey, passphrase: &str) {
let keypair = ssh_key::private::Ed25519Keypair::from(signing);
let private =
ssh_key::PrivateKey::new(ssh_key::private::KeypairData::from(keypair), "test").unwrap();
let encrypted = private.encrypt(&mut rand::rngs::OsRng, passphrase).unwrap();
encrypted
.write_openssh_file(path, ssh_key::LineEnding::LF)
.unwrap();
}
fn test_frame(seq: u64, snapshot: bool) -> Frame {
Frame {
session: "test".to_string(),
output_seq: seq,
bytes: format!("frame-{seq}").into_bytes(),
snapshot,
closed: false,
}
}
#[test]
fn frame_buffer_renders_only_contiguous_frames() {
let mut buffer = FrameBuffer::default();
let mut last = 10;
assert!(buffer.accept(test_frame(12, false), &mut last).is_empty());
assert_eq!(last, 10);
let ready = buffer.accept(test_frame(11, false), &mut last);
assert_eq!(ready.len(), 2);
assert_eq!(ready[0].output_seq, 11);
assert_eq!(ready[1].output_seq, 12);
assert_eq!(last, 12);
}
#[test]
fn frame_buffer_ignores_duplicate_frames() {
let mut buffer = FrameBuffer::default();
let mut last = 4;
assert_eq!(buffer.accept(test_frame(5, false), &mut last).len(), 1);
assert!(buffer.accept(test_frame(5, false), &mut last).is_empty());
assert_eq!(last, 5);
}
#[test]
fn frame_buffer_accepts_snapshot_as_complete_state() {
let mut buffer = FrameBuffer::default();
let mut last = 10;
assert!(buffer.accept(test_frame(12, false), &mut last).is_empty());
let ready = buffer.accept(test_frame(20, true), &mut last);
assert_eq!(ready.len(), 1);
assert_eq!(ready[0].output_seq, 20);
assert_eq!(last, 20);
assert!(buffer.accept(test_frame(12, false), &mut last).is_empty());
}
#[test]
fn frame_buffer_keeps_gap_buffered_without_snapshot_resync() {
let mut buffer = FrameBuffer::default();
let mut last = 10;
assert!(buffer.accept(test_frame(13, false), &mut last).is_empty());
assert_eq!(last, 10);
assert!(buffer.accept(test_frame(14, false), &mut last).is_empty());
assert_eq!(last, 10);
assert!(!buffer.resync_due());
let ready = buffer.accept(test_frame(11, false), &mut last);
assert_eq!(ready.len(), 1);
assert_eq!(ready[0].output_seq, 11);
assert_eq!(last, 11);
let ready = buffer.accept(test_frame(12, false), &mut last);
assert_eq!(ready.len(), 3);
assert_eq!(
ready
.iter()
.map(|frame| frame.output_seq)
.collect::<Vec<_>>(),
vec![12, 13, 14]
);
assert_eq!(last, 14);
}
#[test]
fn frame_buffer_resyncs_old_gap_with_cooldown() {
let mut buffer = FrameBuffer::default();
let mut last = 10;
assert!(buffer.accept(test_frame(13, false), &mut last).is_empty());
buffer.backdate_gap_for_test(Duration::from_millis(FRAME_GAP_RESYNC_AFTER_MS + 10));
assert!(buffer.resync_due());
assert!(
!buffer.resync_due(),
"resync must be rate-limited so snapshots cannot churn the TUI"
);
let ready = buffer.accept(test_frame(20, true), &mut last);
assert_eq!(ready.len(), 1);
assert_eq!(ready[0].output_seq, 20);
assert_eq!(last, 20);
assert!(!buffer.resync_due());
}
#[test]
fn frame_buffer_ignores_stale_snapshot() {
let mut buffer = FrameBuffer::default();
let mut last = 20;
assert!(buffer.accept(test_frame(19, true), &mut last).is_empty());
assert_eq!(last, 20);
}
#[test]
fn predictor_only_accepts_plain_printable_input() {
assert!(Predictor::predictable(b"hello world"));
assert!(!Predictor::predictable(b""));
assert!(!Predictor::predictable(b"\n"));
assert!(!Predictor::predictable(b"\r"));
assert!(!Predictor::predictable(b"\x1b[A"));
assert!(!Predictor::predictable("snowman: \u{2603}".as_bytes()));
assert!(!Predictor::predictable("wide: 你".as_bytes()));
assert!(!Predictor::predictable("combine: e\u{301}".as_bytes()));
}
#[test]
fn predictor_tracks_alternate_screen_state() {
let mut predictor = Predictor::new(true);
predictor.observe_output(b"\x1b[?1049h");
assert!(predictor.alternate_screen);
predictor.observe_output(b"\x1b[?1049l");
assert!(!predictor.alternate_screen);
}
#[test]
fn predictor_tracks_legacy_and_combined_alternate_screen_modes() {
let mut predictor = Predictor::new(true);
predictor.observe_output(b"\x1b[?1000;1006;1049h");
assert!(predictor.alternate_screen);
predictor.observe_output(b"\x1b[?1006;1000l");
assert!(
predictor.alternate_screen,
"leaving mouse modes must not imply leaving alternate screen"
);
predictor.observe_output(b"\x1b[?1047l");
assert!(!predictor.alternate_screen);
predictor.observe_output(b"\x1b[?47h");
assert!(predictor.alternate_screen);
predictor.observe_output(b"\x1b[?47l");
assert!(!predictor.alternate_screen);
}
#[test]
fn predictor_tracks_alternate_screen_split_across_frames() {
let mut predictor = Predictor::new(true);
predictor.observe_output(b"\x1b");
assert!(!predictor.alternate_screen);
predictor.observe_output(b"[?");
assert!(!predictor.alternate_screen);
predictor.observe_output(b"10");
assert!(!predictor.alternate_screen);
predictor.observe_output(b"49h");
assert!(predictor.alternate_screen);
predictor.observe_output(b"\x1b");
assert!(predictor.alternate_screen);
predictor.observe_output(b"[?104");
assert!(predictor.alternate_screen);
predictor.observe_output(b"9l");
assert!(!predictor.alternate_screen);
}
#[test]
fn private_mode_parser_ignores_non_alt_modes() {
assert_eq!(
terminal_private_mode_transition(b"\x1b[?1000;1006h", 0),
Some((13, None))
);
assert_eq!(
terminal_private_mode_transition(b"\x1b[?1000;1049h", 0),
Some((13, Some(true)))
);
}
/// Typing printable characters builds an in-order predicted line and the
/// local cursor tracks the end of it.
#[test]
fn predictor_predicts_printable_run() {
let mut p = Predictor::with_mode(true, PredictMode::Always);
p.observe_input(b"hi").unwrap();
assert_eq!(p.visible_bytes(), b"hi");
assert_eq!(p.cursor, 2);
assert_eq!(p.pending_len(), 2);
}
/// An authoritative server frame confirms (and clears) outstanding
/// predictions: after a frame the overlay is empty.
#[test]
fn predictor_confirmation_clears_prediction() {
let mut p = Predictor::with_mode(true, PredictMode::Always);
p.observe_input(b"abc").unwrap();
assert_eq!(p.pending_len(), 3);
// Server echoes the keystrokes; a new frame arrives.
p.clear_pending().unwrap(); // run loop calls this before render_frame
assert_eq!(p.pending_len(), 0);
assert!(p.visible_bytes().is_empty());
}
/// A glitch -- server output that contradicts the prediction (here the
/// server enters the alternate screen mid-burst) -- discards the epoch's
/// predictions so nothing wrong is left on screen.
#[test]
fn predictor_glitch_discards_epoch() {
let mut p = Predictor::with_mode(true, PredictMode::Always);
p.observe_input(b"vim").unwrap();
assert_eq!(p.pending_len(), 3);
let epoch_before = p.epoch;
// Server diverges: enters a full-screen TUI. Our line prediction is now
// meaningless and must be dropped.
p.observe_output(b"\x1b[?1049h");
assert!(p.alternate_screen);
assert_eq!(p.pending_len(), 0);
// The contradicted epoch is fully retired (confirmed past).
assert!(p.confirmed_epoch >= epoch_before);
// While in the alternate screen we refuse to speculate at all.
p.observe_input(b"x").unwrap();
assert_eq!(p.pending_len(), 0);
}
/// Backspace prediction removes the last predicted cell from the end of the
/// line and moves the local cursor back.
#[test]
fn predictor_predicts_backspace() {
let mut p = Predictor::with_mode(true, PredictMode::Always);
p.observe_input(b"hello").unwrap();
assert_eq!(p.visible_bytes(), b"hello");
p.observe_input(b"\x7f").unwrap(); // DEL
assert_eq!(p.visible_bytes(), b"hell");
assert_eq!(p.cursor, 4);
p.observe_input(&[0x08]).unwrap(); // ^H
assert_eq!(p.visible_bytes(), b"hel");
assert_eq!(p.cursor, 3);
// Backspacing past the start of our predicted span bails safely (we
// don't model what's to the left), ending the epoch with no glyphs.
p.observe_input(b"\x7f\x7f\x7f\x7f").unwrap();
assert_eq!(p.pending_len(), 0);
}
/// Right-margin / line-wrap edge: we refuse to predict past a conservative
/// per-line budget instead of risking a wrong glyph at the wrap point.
/// Bytes are fed in small interactive bursts so the per-line budget (not the
/// paste guard) is what trips.
#[test]
fn predictor_line_wrap_edge_bails_at_budget() {
let mut p = Predictor::with_mode(true, PredictMode::Always);
let chunk = vec![b'x'; super::PREDICT_BURST_LIMIT];
let mut typed = 0;
while typed + chunk.len() <= super::PREDICT_MAX_LINE {
p.observe_input(&chunk).unwrap();
typed += chunk.len();
}
assert_eq!(p.pending_len(), super::PREDICT_MAX_LINE);
// One more character crosses the budget -> bail, drop the epoch.
p.observe_input(b"y").unwrap();
assert_eq!(p.pending_len(), 0);
}
/// A paste-sized burst is not speculated at all (no overlay), but the input
/// is still forwarded by the caller -- prediction is display-only.
#[test]
fn predictor_skips_paste_sized_burst() {
let mut p = Predictor::with_mode(true, PredictMode::Always);
let paste = vec![b'a'; super::PREDICT_BURST_LIMIT + 1];
p.observe_input(&paste).unwrap();
assert_eq!(p.pending_len(), 0);
}
/// Carriage return / newline cannot be modeled (prompt, scroll), so it ends
/// the epoch and clears the speculative line.
#[test]
fn predictor_newline_ends_epoch() {
let mut p = Predictor::with_mode(true, PredictMode::Always);
p.observe_input(b"ls").unwrap();
assert_eq!(p.pending_len(), 2);
p.observe_input(b"\r").unwrap();
assert_eq!(p.pending_len(), 0);
}
/// In-line arrow-key motion (improvement over the old printable-only
/// predictor): Left/Right move the local cursor within the predicted span
/// and a following keystroke overwrites in place.
#[test]
fn predictor_predicts_inline_arrows() {
let mut p = Predictor::with_mode(true, PredictMode::Always);
p.observe_input(b"abc").unwrap();
assert_eq!(p.cursor, 3);
p.observe_input(b"\x1b[D").unwrap(); // left
assert_eq!(p.cursor, 2);
p.observe_input(b"\x1bOD").unwrap(); // left (application/SS3)
assert_eq!(p.cursor, 1);
p.observe_input(b"\x1b[C").unwrap(); // right
assert_eq!(p.cursor, 2);
// Overwrite the cell at the cursor (index 2, the 'c') in place.
p.observe_input(b"Z").unwrap();
assert_eq!(p.visible_bytes(), b"abZ");
assert_eq!(p.cursor, 3);
}
/// Experimental (adaptive) mode shows no prediction while SRTT is below the
/// trigger, and shows it once SRTT rises above the trigger.
#[test]
fn predictor_experimental_respects_srtt_threshold() {
let mut p = Predictor::with_mode(true, PredictMode::Experimental);
// Fast link: no fresh prediction yet, no SRTT sample -> not displayed.
p.set_srtt_for_test(5.0);
p.observe_input(b"abc").unwrap();
assert!(p.pending_len() > 0); // model still tracks the prediction...
assert!(!p.should_display()); // ...but policy hides it on a fast link.
// Slow link: SRTT above the high trigger -> predictions are displayed.
p.set_srtt_for_test(60.0);
assert!(p.should_display());
}
/// Always mode displays regardless of SRTT; Off never does.
#[test]
fn predictor_mode_policies() {
let mut always = Predictor::with_mode(true, PredictMode::Always);
always.observe_input(b"x").unwrap();
assert!(always.should_display());
let off = Predictor::with_mode(true, PredictMode::Off);
assert!(!off.enabled);
assert!(!off.should_display());
assert_eq!(PredictMode::parse("off"), PredictMode::Off);
assert_eq!(PredictMode::parse("ALWAYS"), PredictMode::Always);
assert_eq!(PredictMode::parse("adaptive"), PredictMode::Experimental);
assert_eq!(PredictMode::parse("anything"), PredictMode::Experimental);
}
#[test]
fn cli_prediction_flags_override_config() {
let mut args = test_args("host", &[]);
let mut config = ClientConfig {
predict_mode: "experimental".to_string(),
..ClientConfig::default()
};
args.predict_always = true;
assert_eq!(
selected_predict_mode(&args, &config).unwrap(),
PredictMode::Always
);
args.predict_always = false;
args.predict_never = true;
assert_eq!(
selected_predict_mode(&args, &config).unwrap(),
PredictMode::Off
);
args.predict_never = false;
args.predict_mode = Some("always".to_string());
assert_eq!(
selected_predict_mode(&args, &config).unwrap(),
PredictMode::Always
);
args.predict_mode = None;
config.predict_mode = "off".to_string();
assert_eq!(
selected_predict_mode(&args, &config).unwrap(),
PredictMode::Off
);
}
#[test]
fn startup_command_joins_args_for_remote_shell() {
assert_eq!(startup_command(&[]), None);
assert_eq!(startup_command(&["tm".to_string()]), Some(b"tm\r".to_vec()));
assert_eq!(
startup_command(&["tm".to_string(), "work".to_string()]),
Some(b"tm work\r".to_vec())
);
}
#[test]
fn split_after_command_submit_holds_followup_input() {
assert_eq!(split_after_command_submit(b"tm"), None);
assert_eq!(
split_after_command_submit(b"tm\r"),
Some((b"tm\r".to_vec(), Vec::new()))
);
assert_eq!(
split_after_command_submit(b"tm\r\x1b[B"),
Some((b"tm\r".to_vec(), b"\x1b[B".to_vec()))
);
}
#[test]
fn post_submit_hold_only_catches_terminal_control_input() {
assert!(should_hold_post_submit_input(b"\x1b[B"));
assert!(should_hold_post_submit_input(b"\x9bB"));
assert!(!should_hold_post_submit_input(b"next command"));
assert!(!should_hold_post_submit_input(b"\r"));
assert!(!should_hold_post_submit_input(b""));
}
#[test]
fn startup_gate_keeps_later_input_behind_queued_control_input() {
assert!(should_hold_during_startup_gate(
StartupGateMode::HoldControl,
false,
b"\x1b[B"
));
assert!(should_hold_during_startup_gate(
StartupGateMode::HoldControl,
true,
b"x"
));
assert!(should_hold_during_startup_gate(
StartupGateMode::HoldAll,
false,
b"x"
));
assert!(!should_hold_during_startup_gate(
StartupGateMode::HoldControl,
false,
b"x"
));
}
#[test]
fn post_submit_hold_duration_has_short_printable_phase() {
assert_eq!(post_submit_hold_duration(b"\x1b[B"), STARTUP_INPUT_HOLD);
assert_eq!(post_submit_hold_duration(b"x"), POST_SUBMIT_ALL_INPUT_HOLD);
}
#[test]
fn contact_flushes_pending_terminal_input_for_interactive_sessions() {
assert!(should_flush_terminal_input_after_contact(
false, "normal", false
));
assert!(should_flush_terminal_input_after_contact(
false,
"read-write",
false
));
assert!(!should_flush_terminal_input_after_contact(
false, "normal", true
));
assert!(!should_flush_terminal_input_after_contact(
true, "normal", false
));
assert!(!should_flush_terminal_input_after_contact(
false,
"view-only",
false
));
assert!(!should_flush_terminal_input_after_contact(
false,
"forward-only",
false
));
}
#[test]
fn reconnect_refreshes_live_send_address_from_credentials() {
let mut addr = "127.0.0.1:50000".parse().unwrap();
let cred = CachedCredential {
server: "host".to_string(),
session: "term".to_string(),
mode: "normal".to_string(),
udp_host: "127.0.0.1".to_string(),
udp_port: 50001,
client_id: [1; 16],
session_key: [2; 32],
session_key_id: [3; 16],
attach_ticket: Vec::new(),
attach_ticket_psk: [4; 32],
last_rendered_seq: 0,
};
refresh_live_addr(&mut addr, &cred).unwrap();
assert_eq!(addr.port(), 50001);
}
#[test]
fn resume_accepts_unknown_client_reject_from_legacy_zero_id_server() {
let client_id = [7u8; 16];
let matching_reject = protocol::Packet {
header: protocol::Header {
kind: PacketKind::AttachReject,
flags: 0,
conn_id: client_id,
seq: 0,
ack: 0,
session_key_id: [0u8; 16],
body_len: 0,
},
body: Vec::new(),
};
let zero_id_reject = protocol::Packet {
header: protocol::Header {
conn_id: [0u8; 16],
..matching_reject.header.clone()
},
body: Vec::new(),
};
let unrelated_frame = protocol::Packet {
header: protocol::Header {
kind: PacketKind::Frame,
conn_id: [9u8; 16],
..matching_reject.header.clone()
},
body: Vec::new(),
};
assert!(is_resume_response_for_client(&matching_reject, client_id));
assert!(is_resume_response_for_client(&zero_id_reject, client_id));
assert!(!is_resume_response_for_client(&unrelated_frame, client_id));
}
#[tokio::test]
async fn pending_stream_open_is_retransmitted_until_answered() {
let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let sender = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let addr = receiver.local_addr().unwrap();
let cred = CachedCredential {
server: "host".to_string(),
session: "term".to_string(),
mode: "forward-only".to_string(),
udp_host: "127.0.0.1".to_string(),
udp_port: addr.port(),
client_id: [1; 16],
session_key: [2; 32],
session_key_id: protocol::session_key_id(&[2; 32]),
attach_ticket: Vec::new(),
attach_ticket_psk: [3; 32],
last_rendered_seq: 0,
};
let mut send_seq = 10;
let mut pending = HashMap::from([(
44,
PendingStreamOpen {
target_host: "127.0.0.1".to_string(),
target_port: 8080,
last_sent: std::time::Instant::now() - Duration::from_millis(250),
attempts: 1,
},
)]);
retransmit_stream_opens(&sender, addr, &cred, &mut send_seq, &mut pending)
.await
.unwrap();
let mut buf = [0u8; 2048];
let (n, _) = tokio::time::timeout(Duration::from_secs(1), receiver.recv_from(&mut buf))
.await
.unwrap()
.unwrap();
let packet = protocol::decode(&buf[..n]).unwrap();
assert_eq!(packet.header.kind, PacketKind::StreamOpen);
assert_eq!(packet.header.seq, 10);
let plain =
protocol::decrypt_body(&packet, &cred.session_key, protocol::CLIENT_TO_SERVER).unwrap();
let open: protocol::StreamOpen = protocol::from_body(&plain).unwrap();
assert_eq!(open.stream_id, 44);
assert_eq!(open.target_port, 8080);
assert_eq!(send_seq, 11);
assert_eq!(pending[&44].attempts, 2);
}
#[test]
fn update_check_accepts_only_check_flag() {
assert!(!update_check_requested(&[]).unwrap());
assert!(update_check_requested(&["--check".to_string()]).unwrap());
assert!(update_check_requested(&["-n".to_string()]).unwrap());
assert!(update_check_requested(&["--wat".to_string()]).is_err());
assert!(update_check_requested(&["--check".to_string(), "extra".to_string()]).is_err());
}
#[test]
fn status_uses_host_alias_and_user() {
let host = HostConfig {
ssh: Some("server.example.com".to_string()),
user: Some("palav".to_string()),
..HostConfig::default()
};
assert_eq!(
status_ssh_target("palav", &host),
"palav@server.example.com"
);
assert_eq!(
status_ssh_target("rawhost", &HostConfig::default()),
"rawhost"
);
}
#[test]
fn status_detects_local_targets() {
assert!(is_local_status_target("localhost"));
assert!(is_local_status_target("palav@127.0.0.1"));
assert!(is_local_status_target("::1"));
assert!(!is_local_status_target("server.example.com"));
}
#[test]
fn restart_script_has_systemd_and_nohup_fallback() {
assert!(RESTART_STATUS_SCRIPT.contains("systemctl --user restart dosh-server.service"));
assert!(RESTART_STATUS_SCRIPT.contains("nohup"));
assert!(RESTART_STATUS_SCRIPT.contains("dosh-server"));
}
#[test]
fn doctor_detects_server_version_mismatch() {
assert!(!server_version_mismatch(""));
assert!(!server_version_mismatch(env!("CARGO_PKG_VERSION")));
assert!(server_version_mismatch("0.0.0-old"));
}
#[test]
fn release_download_url_uses_latest_release_asset() {
assert_eq!(
latest_release_download_url(
"https://git.palav.dev/Palav/dosh.git",
"dosh-linux-x86_64.tar.gz"
)
.unwrap(),
"https://git.palav.dev/Palav/dosh/releases/latest/download/dosh-linux-x86_64.tar.gz"
);
assert!(
latest_release_download_url("git@git.palav.dev:Palav/dosh.git", "artifact").is_none()
);
}
#[test]
fn release_tag_parses_effective_latest_url() {
assert_eq!(
release_tag_from_effective_url(
"https://example.com/owner/dosh",
"https://example.com/owner/dosh/releases/tag/v1.0.0"
)
.as_deref(),
Some("v1.0.0")
);
assert_eq!(release_version_from_tag("v1.0.0"), "1.0.0");
assert_eq!(release_version_from_tag("2026.06.28"), "2026.06.28");
assert_eq!(
release_tag_download_url(
"https://example.com/owner/dosh.git",
"v1.0.0",
"dosh-linux-x86_64.tar.gz"
)
.as_deref(),
Some(
"https://example.com/owner/dosh/releases/download/v1.0.0/dosh-linux-x86_64.tar.gz"
)
);
assert!(release_tag_from_effective_url("https://example.com/owner/dosh", "").is_none());
}
#[test]
fn update_version_status_compares_numeric_parts() {
assert_eq!(update_version_status("0.1.5", "0.1.5"), "current");
assert_eq!(update_version_status("0.1.9", "0.1.10"), "update available");
assert_eq!(update_version_status("1.0.0", "0.9.9"), "different");
}
#[test]
fn resolved_startup_command_expands_global_extension() {
let mut config = ClientConfig::default();
config.extensions.insert(
"tm".to_string(),
CommandExtension {
command: Some("tm {args}".to_string()),
description: Some("server tmux dashboard".to_string()),
disabled: false,
},
);
let host = HostConfig::default();
assert_eq!(
resolved_startup_command(&["tm".into(), "work space".into()], &config, &host),
Some(b"tm 'work space'\r".to_vec())
);
}
#[test]
fn resolved_startup_command_uses_host_extension_override() {
let mut config = ClientConfig::default();
config.extensions.insert(
"tm".to_string(),
CommandExtension {
command: Some("tm {args}".to_string()),
description: None,
disabled: false,
},
);
let mut host = HostConfig::default();
host.extensions.insert(
"tm".to_string(),
CommandExtension {
command: Some("/opt/tm/bin/tm --remote {args}".to_string()),
description: None,
disabled: false,
},
);
assert_eq!(
resolved_startup_command(&["tm".into(), "dosh".into()], &config, &host),
Some(b"/opt/tm/bin/tm --remote 'dosh'\r".to_vec())
);
}
#[test]
fn resolved_startup_command_host_can_disable_global_extension() {
let mut config = ClientConfig::default();
config.extensions.insert(
"tm".to_string(),
CommandExtension {
command: Some("tm {args}".to_string()),
description: None,
disabled: false,
},
);
let mut host = HostConfig::default();
host.extensions.insert(
"tm".to_string(),
CommandExtension {
command: None,
description: None,
disabled: true,
},
);
assert_eq!(
resolved_startup_command(&["tm".into(), "dosh".into()], &config, &host),
None
);
}
#[test]
fn resolved_startup_command_keeps_raw_command_fallback_and_default_command() {
let config = ClientConfig::default();
let mut host = HostConfig {
default_command: Some("tm".to_string()),
..HostConfig::default()
};
assert_eq!(
resolved_startup_command(&["echo".into(), "ok".into()], &config, &host),
Some(b"echo ok\r".to_vec())
);
assert_eq!(
resolved_startup_command(&[], &config, &host),
Some(b"tm\r".to_vec())
);
host.default_command = None;
assert_eq!(resolved_startup_command(&[], &config, &host), None);
}
#[test]
fn parses_local_forward_with_default_bind() {
assert_eq!(
parse_local_forward("8080:127.0.0.1:80").unwrap(),
LocalForward {
bind_host: "127.0.0.1".to_string(),
listen_port: 8080,
target_host: "127.0.0.1".to_string(),
target_port: 80,
}
);
}
#[test]
fn parses_local_forward_with_explicit_bind() {
assert_eq!(
parse_local_forward("0.0.0.0:8080:localhost:8000").unwrap(),
LocalForward {
bind_host: "0.0.0.0".to_string(),
listen_port: 8080,
target_host: "localhost".to_string(),
target_port: 8000,
}
);
}
#[test]
fn parses_remote_forward_with_default_bind() {
assert_eq!(
parse_remote_forward("9090:127.0.0.1:5432").unwrap(),
RemoteForward {
bind_host: "127.0.0.1".to_string(),
listen_port: 9090,
target_host: "127.0.0.1".to_string(),
target_port: 5432,
}
);
}
#[test]
fn parses_dynamic_forward_with_default_bind() {
assert_eq!(
parse_dynamic_forward("1080").unwrap(),
DynamicForward {
bind_host: "127.0.0.1".to_string(),
listen_port: 1080,
}
);
}
#[test]
fn forward_command_rewrites_to_forward_only_connection() {
let args = test_args(
"forward",
&["homelab", "-L", "8080:127.0.0.1:80", "-D", "1080", "-f"],
);
let rewritten = rewrite_forward_command(args).unwrap();
assert_eq!(rewritten.server.as_deref(), Some("homelab"));
assert!(rewritten.forward_only);
assert!(rewritten.background);
assert_eq!(rewritten.command, Vec::<String>::new());
assert_eq!(rewritten.local_forward, vec!["8080:127.0.0.1:80"]);
assert_eq!(rewritten.dynamic_forward, vec!["1080"]);
}
#[test]
fn forward_command_requires_a_forward_request() {
let err = rewrite_forward_command(test_args("forward", &["homelab"]))
.unwrap_err()
.to_string();
assert!(err.contains("requires at least one"));
}
#[test]
fn forward_parser_rejects_zero_ports_and_control_chars() {
assert!(parse_local_forward("0:127.0.0.1:80").is_err());
assert!(parse_local_forward("8080:127.0.0.1:0").is_err());
assert!(parse_dynamic_forward("127.0.0.1:0").is_err());
assert!(!valid_forward_host("bad\nhost"));
assert!(!valid_forward_host("bad\thost"));
assert!(valid_forward_host("server.example.com"));
}
#[test]
fn recover_removes_only_matching_cached_credentials() {
let dir = tempfile::tempdir().unwrap();
let root = dir.path().display().to_string();
let matching = dir
.path()
.join(format!("{}.bin", cache_key("user@host", "default", "rw")));
let other = dir
.path()
.join(format!("{}.bin", cache_key("other", "default", "rw")));
let similar = dir
.path()
.join(format!("{}.bin", cache_key("user@host2", "default", "rw")));
fs::write(&matching, b"cached").unwrap();
fs::write(&other, b"cached").unwrap();
fs::write(&similar, b"cached").unwrap();
assert_eq!(cache_server_prefix("user@host"), "user_host_");
assert_eq!(clear_cached_credentials(&root, "user@host").unwrap(), 1);
assert!(!matching.exists());
assert!(other.exists());
assert!(similar.exists());
}
#[test]
fn pending_user_input_is_bounded_and_ordered() {
let mut pending = VecDeque::new();
let mut pending_bytes = 0usize;
queue_pending_user_input(&mut pending, &mut pending_bytes, b"abc".to_vec()).unwrap();
queue_pending_user_input(&mut pending, &mut pending_bytes, b"def".to_vec()).unwrap();
assert_eq!(pending_bytes, 6);
assert_eq!(pending.pop_front().unwrap().bytes, b"abc");
assert_eq!(pending.pop_front().unwrap().bytes, b"def");
let mut pending = VecDeque::new();
let mut pending_bytes = MAX_PENDING_USER_INPUT_BYTES;
assert!(queue_pending_user_input(&mut pending, &mut pending_bytes, b"x".to_vec()).is_err());
assert!(pending.is_empty());
assert_eq!(pending_bytes, MAX_PENDING_USER_INPUT_BYTES);
}
#[test]
fn transient_udp_send_errors_are_not_terminal_fatal() {
#[cfg(unix)]
{
assert!(is_transient_udp_send_error(
&std::io::Error::from_raw_os_error(libc::ENETUNREACH)
));
assert!(is_transient_udp_send_error(
&std::io::Error::from_raw_os_error(libc::EHOSTUNREACH)
));
assert!(is_transient_udp_send_error(
&std::io::Error::from_raw_os_error(libc::EADDRNOTAVAIL)
));
assert!(!is_transient_udp_send_error(
&std::io::Error::from_raw_os_error(libc::EINVAL)
));
}
#[cfg(windows)]
{
assert!(is_transient_udp_send_error(
&std::io::Error::from_raw_os_error(10051)
)); // WSAENETUNREACH
assert!(is_transient_udp_send_error(
&std::io::Error::from_raw_os_error(10065)
)); // WSAEHOSTUNREACH
assert!(is_transient_udp_send_error(
&std::io::Error::from_raw_os_error(10049)
)); // WSAEADDRNOTAVAIL
assert!(!is_transient_udp_send_error(
&std::io::Error::from_raw_os_error(10022)
)); // WSAEINVAL
}
}
#[test]
fn stale_mouse_reports_are_stripped_from_pending_input() {
let input = b"\x1b[<35;152;1Mhello\x1b[<0;107;10m";
assert_eq!(strip_stale_mouse_reports(input), b"hello");
}
#[test]
fn stale_mouse_tail_reports_are_stripped_from_pending_input() {
let input = b"35;152;1M35;149;1Mls\r";
assert_eq!(strip_stale_mouse_reports(input), b"ls\r");
assert_eq!(strip_stale_mouse_reports(b"152;1Mls\r"), b"ls\r");
}
#[test]
fn split_stale_mouse_suffixes_are_stripped_from_pending_input() {
assert_eq!(strip_stale_mouse_reports(b"1M35;149;1Mls\r"), b"ls\r");
assert_eq!(strip_stale_mouse_reports(b"10m"), b"");
}
#[test]
fn leading_mouse_terminator_from_split_report_is_stripped() {
assert_eq!(strip_stale_mouse_reports(b"M35;149;1Mls\r"), b"ls\r");
assert_eq!(strip_stale_mouse_reports(b"m35;149;1Mls\r"), b"ls\r");
}
#[test]
fn numeric_semicolon_input_survives_stale_filter() {
assert_eq!(strip_stale_mouse_reports(b"123;"), b"123;");
assert_eq!(strip_stale_mouse_reports(b"12;34"), b"12;34");
assert_eq!(strip_stale_mouse_reports(b"echo 1;"), b"echo 1;");
assert_eq!(strip_stale_mouse_reports(b"make test"), b"make test");
}
#[test]
fn incomplete_three_field_mouse_prefix_is_suppressed() {
assert_eq!(strip_stale_mouse_reports(b"35;152;"), b"");
assert_eq!(strip_stale_mouse_reports(b"\x1b[<35;152;"), b"");
assert_eq!(strip_stale_mouse_reports(b"\x9b35;152;"), b"");
}
#[test]
fn stale_focus_reports_are_stripped_from_pending_input() {
assert_eq!(strip_stale_mouse_reports(b"\x1b[Ihello\x1b[O"), b"hello");
assert_eq!(strip_stale_mouse_reports(b"\x9bIhello\x9bO"), b"hello");
}
#[test]
fn focus_in_reports_trigger_local_repaint_detection() {
assert!(input_contains_focus_in(b"\x1b[I"));
assert!(input_contains_focus_in(b"before\x9bIafter"));
assert!(!input_contains_focus_in(b"\x1b[O"));
assert!(!input_contains_focus_in(b"\x1b[A"));
}
#[test]
fn focus_reports_are_never_forwarded_as_user_input() {
assert_eq!(strip_terminal_focus_reports(b"\x1b[Ihello\x1b[O"), b"hello");
assert_eq!(strip_terminal_focus_reports(b"\x9bIhello\x9bO"), b"hello");
assert_eq!(strip_terminal_focus_reports(b"\x1b[A"), b"\x1b[A");
assert_eq!(
strip_terminal_focus_reports(b"before\x1b[Iafter"),
b"beforeafter"
);
}
#[test]
fn idle_repaint_only_runs_for_stale_alternate_screen() {
let now = Instant::now();
let stale = now - ALT_SCREEN_IDLE_REPAINT_AFTER - Duration::from_secs(1);
let recent = now - Duration::from_secs(1);
assert!(should_repaint_idle_alternate_screen(
true, stale, stale, now
));
assert!(!should_repaint_idle_alternate_screen(
false, stale, stale, now
));
assert!(!should_repaint_idle_alternate_screen(
true, recent, stale, now
));
assert!(!should_repaint_idle_alternate_screen(
true, stale, recent, now
));
}
#[test]
fn non_mouse_escape_input_survives_stale_filter() {
assert_eq!(strip_stale_mouse_reports(b"\x1b[A"), b"\x1b[A");
assert_eq!(
strip_stale_mouse_reports(b"\x1b[200~paste\x1b[201~"),
b"\x1b[200~paste\x1b[201~"
);
}
#[test]
fn escape_key_parser_accepts_mosh_style_controls() {
assert_eq!(parse_escape_key("^]").unwrap(), Some(vec![0x1d]));
assert_eq!(parse_escape_key("^C").unwrap(), Some(vec![0x03]));
assert_eq!(parse_escape_key("^?").unwrap(), Some(vec![0x7f]));
assert_eq!(parse_escape_key("0x1d").unwrap(), Some(vec![0x1d]));
assert_eq!(parse_escape_key("~").unwrap(), Some(vec![b'~']));
assert_eq!(parse_escape_key("none").unwrap(), None);
assert!(parse_escape_key("too-long").is_err());
}
#[test]
fn escape_key_matches_inside_stdin_chunk() {
assert!(input_matches_escape(b"abc\x1d", Some(&[0x1d])));
assert!(!input_matches_escape(b"abc", Some(&[0x1d])));
assert!(!input_matches_escape(b"abc\x1d", None));
}
// --- Item 1: disconnect status line state machine ---
/// The status line stays hidden while the link is fresh and only appears once
/// silence crosses the threshold; the rendered text reports elapsed seconds.
#[test]
fn disconnect_status_shows_after_threshold() {
let mut status = DisconnectStatus::new(true);
// Fresh link: nothing to do.
assert_eq!(status.on_tick(Duration::from_secs(0)), StatusAction::None);
assert_eq!(status.on_tick(Duration::from_secs(1)), StatusAction::None);
// Crossed the threshold: show with the elapsed seconds in the text.
let action = status.on_tick(Duration::from_secs(3));
assert_eq!(
action,
StatusAction::Show("[dosh] reconnecting — last contact 3s ago".to_string())
);
status.applied(&action);
assert!(status.shown);
}
/// Once shown, a same-second tick is a no-op but a new elapsed value repaints.
#[test]
fn disconnect_status_repaints_only_on_text_change() {
let mut status = DisconnectStatus::new(true);
let first = status.on_tick(Duration::from_secs(3));
status.applied(&first);
// Same elapsed second -> no redundant repaint.
assert_eq!(status.on_tick(Duration::from_secs(3)), StatusAction::None);
// Next second -> repaint with updated text.
assert_eq!(
status.on_tick(Duration::from_secs(4)),
StatusAction::Show("[dosh] reconnecting — last contact 4s ago".to_string())
);
}
/// Contact resuming clears a shown line immediately, and a sub-threshold tick
/// also clears it; clearing when nothing is shown is a no-op.
#[test]
fn disconnect_status_clears_on_contact_and_recovery() {
let mut status = DisconnectStatus::new(true);
// Nothing shown yet: contact is a no-op.
assert_eq!(status.on_contact(), StatusAction::None);
let show = status.on_tick(Duration::from_secs(5));
status.applied(&show);
assert!(status.shown);
// A packet arrives: clear immediately.
let clear = status.on_contact();
assert_eq!(clear, StatusAction::Clear);
status.applied(&clear);
assert!(!status.shown);
// Show again, then a sub-threshold tick (link recovered) clears it too.
let show = status.on_tick(Duration::from_secs(5));
status.applied(&show);
assert_eq!(status.on_tick(Duration::from_secs(0)), StatusAction::Clear);
}
/// When the feature is disabled the machine never shows, and clears anything
/// already on screen (e.g. config toggled off at runtime).
#[test]
fn disconnect_status_disabled_never_shows() {
let mut status = DisconnectStatus::new(false);
assert_eq!(status.on_tick(Duration::from_secs(10)), StatusAction::None);
// Simulate a stale "shown" flag and confirm it gets cleared.
status.shown = true;
assert_eq!(status.on_tick(Duration::from_secs(10)), StatusAction::Clear);
}
#[test]
fn disconnect_status_suppresses_without_touching_alternate_screen() {
let mut status = DisconnectStatus::new(true);
let show = status.on_tick(Duration::from_secs(3));
status.applied(&show);
assert!(status.shown);
let suppressed = status.on_suppressed();
assert_eq!(suppressed, StatusAction::Forget);
status.applied(&suppressed);
assert!(!status.shown);
assert_eq!(status.drawn_text, "");
}
/// The rendered overlay must preserve cursor position: save cursor (ESC 7),
/// park on the top row, paint the blue reconnect bar, then restore cursor
/// (ESC 8). The run loop suppresses this renderer entirely in full-screen
/// alternate-screen TUIs.
#[test]
fn disconnect_status_overlay_is_save_restore_top_blue_bar() {
let bytes = render_status_overlay("[dosh] reconnecting — last contact 3s ago", 24);
ensure_tui_safe_status_overlay(&bytes).unwrap();
assert!(bytes.starts_with(b"\x1b7"), "must save cursor first");
assert!(bytes.ends_with(b"\x1b8"), "must restore cursor last");
// Positions to the first row, column 1.
assert!(bytes.windows(6).any(|w| w == b"\x1b[1;1H"));
// Clears the row and uses a blue background with white text.
assert!(bytes.windows(4).any(|w| w == b"\x1b[2K"));
assert!(bytes.windows(8).any(|w| w == b"\x1b[44;37m"));
// Resets attributes before restoring the cursor.
assert!(bytes.windows(4).any(|w| w == b"\x1b[0m"));
let clear = render_status_clear(24);
assert!(clear.starts_with(b"\x1b7"));
assert!(clear.ends_with(b"\x1b8"));
assert!(clear.windows(4).any(|w| w == b"\x1b[2K"));
}
#[test]
fn disconnect_status_text_escalates_after_long_silence() {
assert_eq!(
DisconnectStatus::render_text(Duration::from_secs(29)),
"[dosh] reconnecting — last contact 29s ago"
);
assert_eq!(
DisconnectStatus::render_text(Duration::from_secs(30)),
"[dosh] still reconnecting — last contact 30s ago"
);
}
// --- Item 2: predictor policy re-evaluated on a timer tick ---
/// A latency spike must surface predictions on a *timer tick*, not only on
/// the next keystroke: `refresh_policy` force-shows once a prediction has been
/// outstanding past the glitch threshold, even though SRTT hasn't caught up.
#[test]
fn predictor_refresh_policy_force_shows_on_spike() {
let mut p = Predictor::with_mode(true, PredictMode::Experimental);
// Fast SRTT so the SRTT trigger stays off; only the spike timer can show.
p.set_srtt_for_test(5.0);
p.observe_input(b"abc").unwrap();
assert!(p.pending_len() > 0);
// Just typed: not outstanding long enough yet -> still hidden.
assert!(!p.is_drawn());
// Simulate the spike: the prediction has now been outstanding well past
// the glitch threshold. A keystroke-free timer tick must reveal it.
p.backdate_pending_for_test(super::GLITCH_FORCE_MS as u64 + 50);
p.refresh_policy().unwrap();
assert!(p.is_drawn(), "tick should reveal prediction on a spike");
}
/// `refresh_policy` also turns the overlay *off* on a tick once SRTT recovers
/// below the trigger (no keystroke required).
#[test]
fn predictor_refresh_policy_hides_on_recovery() {
let mut p = Predictor::with_mode(true, PredictMode::Experimental);
// Slow link: SRTT trigger latched on -> prediction is shown on input.
p.set_srtt_for_test(60.0);
p.observe_input(b"abc").unwrap();
assert!(p.is_drawn());
// Link recovers below the low trigger; a timer tick must hide it.
p.set_srtt_for_test(5.0);
p.refresh_policy().unwrap();
assert!(!p.is_drawn(), "tick should hide prediction after recovery");
}
/// An idle tick with nothing outstanding is a cheap no-op (no spurious draw).
#[test]
fn predictor_refresh_policy_idle_is_noop() {
let mut p = Predictor::with_mode(true, PredictMode::Experimental);
p.set_srtt_for_test(60.0);
p.refresh_policy().unwrap();
assert!(!p.is_drawn());
}
}
Reference in New Issue View Git Blame Copy Permalink