DuProcess
f6ead86db4
ci / test (push) Has been cancelled
ci / fuzz-smoke (push) Has been cancelled
ci / package-release (linux-x86_64, ubuntu-latest) (push) Has been cancelled
ci / package-release (macos-aarch64, macos-14) (push) Has been cancelled
ci / package-release (macos-x86_64, macos-13) (push) Has been cancelled
ci / package-release (windows-x86_64, windows-latest) (push) Has been cancelled
ci / remote-bench (push) Has been cancelled
5612 lines
200 KiB
Rust
5612 lines
200 KiB
Rust
use anyhow::{Context, Result, anyhow};
|
|
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::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 serde::{Deserialize, Serialize};
|
|
use std::collections::BTreeMap;
|
|
use std::collections::HashMap;
|
|
use std::collections::HashSet;
|
|
use std::collections::VecDeque;
|
|
use std::fs;
|
|
use std::io::{Read, Write};
|
|
use std::net::{
|
|
Ipv4Addr, Ipv6Addr, SocketAddr, TcpListener as StdTcpListener, TcpStream as StdTcpStream,
|
|
ToSocketAddrs,
|
|
};
|
|
use std::path::{Path, PathBuf};
|
|
use std::process::{Command, Stdio};
|
|
use std::sync::Arc;
|
|
use std::sync::atomic::{AtomicU64, Ordering};
|
|
use std::time::{Duration, Instant};
|
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
|
use tokio::net::{TcpListener, TcpStream, UdpSocket, UnixStream};
|
|
use tokio::signal::unix::{SignalKind, signal};
|
|
use tokio::sync::mpsc;
|
|
|
|
const STREAM_INITIAL_WINDOW: usize = 1024 * 1024;
|
|
|
|
/// 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)]
|
|
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)]
|
|
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,
|
|
}
|
|
|
|
#[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("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("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 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 target_udp_host = if let Some(host) = args
|
|
.dosh_host
|
|
.clone()
|
|
.or_else(|| host.dosh_host.clone())
|
|
.or_else(|| config.dosh_host.clone())
|
|
{
|
|
host
|
|
} else if args.local_auth {
|
|
"127.0.0.1".to_string()
|
|
} else {
|
|
match ssh_config(&server, ssh_port) {
|
|
Ok(parsed) => {
|
|
let hostname = parsed
|
|
.hostname
|
|
.clone()
|
|
.unwrap_or_else(|| ssh_destination_host(&server));
|
|
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());
|
|
ssh_destination_host(&server)
|
|
}
|
|
}
|
|
};
|
|
|
|
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,
|
|
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,
|
|
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,
|
|
)
|
|
.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,
|
|
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,
|
|
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 bottom-row overlay");
|
|
|
|
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_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 udp_host = args
|
|
.dosh_host
|
|
.clone()
|
|
.or_else(|| host_config.dosh_host.clone())
|
|
.or_else(|| config.dosh_host.clone())
|
|
.unwrap_or_else(|| {
|
|
parsed_ssh_config
|
|
.hostname
|
|
.clone()
|
|
.unwrap_or_else(|| 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
|
|
}
|
|
);
|
|
println!("[ok] known host: trusted");
|
|
println!("[ok] native auth: authorized as {}", check.requested_user);
|
|
println!(
|
|
"[ok] forwarding policy: tcp={} remote={} agent={}",
|
|
check.allow_tcp_forwarding,
|
|
check.allow_remote_forwarding,
|
|
check.allow_agent_forwarding
|
|
);
|
|
Ok(())
|
|
}
|
|
Err(err) => {
|
|
println!("[fail] native/auth check: {err:#}");
|
|
Err(err.context("dosh doctor failed"))
|
|
}
|
|
}
|
|
}
|
|
|
|
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('\n');
|
|
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}", "ed_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('\n');
|
|
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);
|
|
}
|
|
}
|
|
|
|
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,
|
|
"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<()> {
|
|
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 script = 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'"#;
|
|
let remote_command = format!("sh -c {}", shell_word(script));
|
|
let status = command
|
|
.arg(server)
|
|
.arg(remote_command)
|
|
.status()
|
|
.context("run remote sessions command")?;
|
|
if !status.success() {
|
|
return Err(anyhow!("sessions command 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 latest_release_tag_download_url(repo: &str, artifact: &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);
|
|
let prefix = format!("{web}/releases/tag/");
|
|
let Some(tag) = effective.strip_prefix(&prefix) else {
|
|
return Ok(None);
|
|
};
|
|
if tag.is_empty() {
|
|
return Ok(None);
|
|
}
|
|
Ok(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 {
|
|
println!("dosh {}", env!("CARGO_PKG_VERSION"));
|
|
println!("repo: {repo}");
|
|
println!("installer: {installer}");
|
|
if let Some(latest_url) = latest_release_download_url(&repo, artifact) {
|
|
let mut status = "missing";
|
|
let mut display_url = latest_url.clone();
|
|
if url_reachable(&latest_url)? {
|
|
status = "available";
|
|
} else if let Some(tag_url) = latest_release_tag_download_url(&repo, artifact)?
|
|
&& url_reachable(&tag_url)?
|
|
{
|
|
status = "available";
|
|
display_url = tag_url;
|
|
}
|
|
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 mut script = format!(
|
|
"curl -fsSL {} | DOSH_REPO={} DOSH_PORT={} DOSH_UPDATE_CACHE={} DOSH_UPDATE_QUIET=1 DOSH_USE_PREBUILT=1 sh -s -- client",
|
|
shell_word(&installer),
|
|
shell_word(&repo),
|
|
config.update_port.unwrap_or(config.dosh_port),
|
|
shell_word(&update_cache.display().to_string())
|
|
);
|
|
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 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, ephemeral session: the server recognizes this name shape
|
|
// (`protocol::is_implicit_session_name`) and does NOT persist it across
|
|
// restarts, since you can never reattach to a random name.
|
|
protocol::generate_implicit_session_name()
|
|
}
|
|
|
|
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('[') {
|
|
if let Some((host, _)) = stripped.split_once(']') {
|
|
return host.to_string();
|
|
}
|
|
}
|
|
without_path
|
|
.split_once(':')
|
|
.map_or(without_path, |(host, _)| host)
|
|
.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 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>,
|
|
) -> 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,
|
|
)?;
|
|
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))
|
|
}
|
|
|
|
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(),
|
|
)?;
|
|
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>,
|
|
) -> 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());
|
|
}
|
|
|
|
match load_first_native_identity(config, cli_identity, ssh_config) {
|
|
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_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?;
|
|
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?;
|
|
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 packet.header.conn_id != 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?;
|
|
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 packet.header.conn_id != 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?;
|
|
let mut next = cached.clone();
|
|
next.last_rendered_seq = frame.output_seq;
|
|
Ok((frame, next))
|
|
}
|
|
|
|
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);
|
|
}
|
|
}
|
|
}
|
|
|
|
async fn run_terminal(
|
|
socket: UdpSocket,
|
|
mut cred: CachedCredential,
|
|
first_frame: Option<Frame>,
|
|
predict: bool,
|
|
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<()> {
|
|
let _raw = if forward_only {
|
|
None
|
|
} else {
|
|
Some(RawMode::enter()?)
|
|
};
|
|
let addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
|
|
let mut send_seq = 2u64;
|
|
let mut last_packet_at = Instant::now();
|
|
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 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.
|
|
let mut winch = signal(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 predict_mode = resolve_predict_mode();
|
|
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.
|
|
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_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();
|
|
if let Some(frame) = first_frame {
|
|
if !forward_only {
|
|
render_frame(&frame)?;
|
|
predictor.observe_output(&frame.bytes);
|
|
}
|
|
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
|
|
{
|
|
send_input(&socket, addr, &cred, &mut send_seq, bytes).await?;
|
|
}
|
|
|
|
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(bytes) => {
|
|
if bytes == [0x1d] {
|
|
break;
|
|
}
|
|
if cred.mode != "view-only" && cred.mode != "forward-only" {
|
|
predictor.observe_input(&bytes)?;
|
|
send_input(&socket, addr, &cred, &mut send_seq, bytes).await?;
|
|
}
|
|
}
|
|
None => break,
|
|
}
|
|
}
|
|
_ = async {
|
|
match winch.as_mut() {
|
|
Some(w) => { w.recv().await; }
|
|
None => 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?;
|
|
}
|
|
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?
|
|
{
|
|
if !forward_only {
|
|
render_frame(&frame)?;
|
|
predictor.observe_output(&frame.bytes);
|
|
}
|
|
last_packet_at = Instant::now();
|
|
}
|
|
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);
|
|
}
|
|
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();
|
|
}
|
|
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"",
|
|
)?;
|
|
socket.send_to(&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?
|
|
{
|
|
if !forward_only {
|
|
render_frame(&frame)?;
|
|
predictor.observe_output(&frame.bytes);
|
|
}
|
|
last_packet_at = Instant::now();
|
|
}
|
|
} else {
|
|
return Err(anyhow!("server rejected session: {}", reject.reason));
|
|
}
|
|
}
|
|
PacketKind::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();
|
|
// 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 {
|
|
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();
|
|
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();
|
|
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;
|
|
}
|
|
opened_streams.remove(&reject.stream_id);
|
|
stream_send_credit.remove(&reject.stream_id);
|
|
stream_pending_data.remove(&reject.stream_id);
|
|
stream_next_send_offset.remove(&reject.stream_id);
|
|
stream_sent_data.remove(&reject.stream_id);
|
|
stream_next_recv_offset.remove(&reject.stream_id);
|
|
stream_recv_pending.remove(&reject.stream_id);
|
|
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;
|
|
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 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();
|
|
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);
|
|
opened_streams.remove(&close.stream_id);
|
|
stream_send_credit.remove(&close.stream_id);
|
|
stream_pending_data.remove(&close.stream_id);
|
|
stream_next_send_offset.remove(&close.stream_id);
|
|
stream_sent_data.remove(&close.stream_id);
|
|
stream_next_recv_offset.remove(&close.stream_id);
|
|
stream_recv_pending.remove(&close.stream_id);
|
|
stream_writers.remove(&close.stream_id);
|
|
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);
|
|
}
|
|
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);
|
|
opened_streams.remove(&stream_id);
|
|
stream_send_credit.remove(&stream_id);
|
|
stream_pending_data.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);
|
|
stream_writers.remove(&stream_id);
|
|
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?
|
|
{
|
|
if !forward_only {
|
|
render_frame(&frame)?;
|
|
predictor.observe_output(&frame.bytes);
|
|
}
|
|
last_packet_at = Instant::now();
|
|
}
|
|
} 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;
|
|
socket.send_to(&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()?;
|
|
}
|
|
// 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 = 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?;
|
|
}
|
|
}
|
|
}
|
|
// 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))
|
|
}
|
|
|
|
async fn send_input(
|
|
socket: &UdpSocket,
|
|
addr: SocketAddr,
|
|
cred: &CachedCredential,
|
|
send_seq: &mut u64,
|
|
bytes: Vec<u8>,
|
|
) -> Result<()> {
|
|
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;
|
|
socket.send_to(&packet, addr).await?;
|
|
Ok(())
|
|
}
|
|
|
|
async fn send_stream_open(
|
|
socket: &UdpSocket,
|
|
addr: SocketAddr,
|
|
cred: &CachedCredential,
|
|
send_seq: &mut u64,
|
|
stream_id: u64,
|
|
target_host: String,
|
|
target_port: u16,
|
|
) -> Result<()> {
|
|
let body = protocol::to_body(&StreamOpen {
|
|
stream_id,
|
|
target_host,
|
|
target_port,
|
|
})?;
|
|
send_stream_packet(socket, addr, cred, send_seq, PacketKind::StreamOpen, &body).await
|
|
}
|
|
|
|
async fn send_stream_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
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
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(());
|
|
};
|
|
loop {
|
|
let Some(bytes) = pending.front() else {
|
|
break;
|
|
};
|
|
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(())
|
|
}
|
|
|
|
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)
|
|
}
|
|
|
|
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;
|
|
socket.send_to(&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;
|
|
|
|
/// 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,
|
|
|
|
/// 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,
|
|
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.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]) {
|
|
if contains_bytes(bytes, b"\x1b[?1049h")
|
|
|| contains_bytes(bytes, b"\x1b[?1047h")
|
|
|| contains_bytes(bytes, b"\x1b[?47h")
|
|
{
|
|
self.alternate_screen = true;
|
|
let _ = self.discard_all();
|
|
}
|
|
if contains_bytes(bytes, b"\x1b[?1049l")
|
|
|| contains_bytes(bytes, b"\x1b[?1047l")
|
|
|| contains_bytes(bytes, b"\x1b[?47l")
|
|
{
|
|
self.alternate_screen = false;
|
|
let _ = self.discard_all();
|
|
}
|
|
}
|
|
|
|
/// 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 client config's
|
|
/// `predict_mode` is used, defaulting to the adaptive (experimental) policy.
|
|
fn resolve_predict_mode() -> PredictMode {
|
|
if let Ok(env) = std::env::var("DOSH_PREDICT_MODE") {
|
|
return PredictMode::parse(&env);
|
|
}
|
|
match load_client_config(None) {
|
|
Ok(cfg) => PredictMode::parse(&cfg.predict_mode),
|
|
Err(_) => PredictMode::Experimental,
|
|
}
|
|
}
|
|
|
|
/// 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)
|
|
}
|
|
|
|
#[derive(Default)]
|
|
struct FrameBuffer {
|
|
pending: BTreeMap<u64, Frame>,
|
|
}
|
|
|
|
impl FrameBuffer {
|
|
fn clear(&mut self) {
|
|
self.pending.clear();
|
|
}
|
|
|
|
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();
|
|
*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);
|
|
}
|
|
ready
|
|
}
|
|
}
|
|
|
|
/// 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;
|
|
socket.send_to(&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;
|
|
socket.send_to(&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"",
|
|
)?;
|
|
socket.send_to(&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 line.
|
|
///
|
|
/// When server packets stop arriving we paint a single, unobtrusive line on the
|
|
/// terminal's *bottom row* — e.g. reverse-video cyan `[dosh] last contact 3s ago
|
|
/// — reconnecting…` — using save/restore cursor (ESC 7 / ESC 8) so the real
|
|
/// application cursor never moves and full-screen TUIs are not corrupted. (The
|
|
/// old in-band overlay wrote over the active line and corrupted TUIs; this draws
|
|
/// only the last row, parks the cursor there, then restores, touching nothing
|
|
/// the app owns.) The line 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 line with this exact text on the bottom row.
|
|
Show(String),
|
|
/// Erase the status line (link recovered or feature disabled).
|
|
Clear,
|
|
}
|
|
|
|
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
|
|
}
|
|
}
|
|
|
|
/// 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::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::None => {}
|
|
}
|
|
self.applied(&action);
|
|
Ok(())
|
|
}
|
|
|
|
/// Paint `text` on the terminal's bottom row using save/restore cursor so the
|
|
/// app's cursor is untouched. Sequence: save cursor (ESC 7), move to the last
|
|
/// row col 1, clear that row, set reverse-video cyan, write text, reset
|
|
/// attributes, restore cursor (ESC 8).
|
|
fn emit(&self, text: &str) -> Result<()> {
|
|
let (_, rows) = terminal_size();
|
|
let bytes = render_status_overlay(text, rows);
|
|
emit_status(&bytes)
|
|
}
|
|
|
|
/// Erase the bottom-row status line, 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 `rows`-tall terminal's
|
|
/// bottom row with reverse-video cyan, 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
|
|
// Move to bottom row, column 1.
|
|
out.extend_from_slice(format!("\x1b[{};1H", rows.max(1)).as_bytes());
|
|
out.extend_from_slice(b"\x1b[2K"); // clear entire row
|
|
out.extend_from_slice(b"\x1b[7;36m"); // reverse video + cyan
|
|
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 bottom-row status line.
|
|
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(format!("\x1b[{};1H", rows.max(1)).as_bytes());
|
|
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 = format!("{server}_{session}_{mode}")
|
|
.chars()
|
|
.map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
|
|
.collect::<String>();
|
|
expand_tilde(root).join(format!("{safe}.bin"))
|
|
}
|
|
|
|
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> {
|
|
enable_raw_mode()?;
|
|
Ok(Self)
|
|
}
|
|
}
|
|
|
|
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[?1015l",
|
|
"\x1b[?1006l",
|
|
"\x1b[?1005l",
|
|
"\x1b[?2004l",
|
|
"\x1b[?1049l"
|
|
)
|
|
.as_bytes();
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::{
|
|
DisconnectStatus, DynamicForward, FrameBuffer, LocalForward, PredictMode, Predictor,
|
|
RemoteForward, SshConfig, StatusAction, auth_allows, ensure_tui_safe_status_overlay,
|
|
latest_release_download_url, load_first_native_identity_with_prompt, parse_dynamic_forward,
|
|
parse_local_forward, parse_remote_forward, parse_ssh_config, raw_contains_host_table,
|
|
recv_response_until, render_status_clear, render_status_overlay, requested_env,
|
|
resolved_startup_command, rewrite_forward_command, ssh_destination_host, ssh_username,
|
|
ssh_with_user, startup_command, toml_bare_key_or_quoted, update_check_requested,
|
|
valid_forward_host,
|
|
};
|
|
use dosh::config::{ClientConfig, CommandExtension, HostConfig};
|
|
use dosh::native::EnvVar;
|
|
use dosh::protocol::{self, Frame, PacketKind};
|
|
use ed25519_dalek::{SigningKey, VerifyingKey};
|
|
use std::time::Duration;
|
|
|
|
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,
|
|
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 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 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 mut config = ClientConfig::default();
|
|
config.identity_files = vec![path.display().to_string()];
|
|
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_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"\x1b[A"));
|
|
assert!(!Predictor::predictable("snowman: \u{2603}".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);
|
|
}
|
|
|
|
/// 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 startup_command_joins_args_for_remote_shell() {
|
|
assert_eq!(startup_command(&[]), None);
|
|
assert_eq!(startup_command(&["tm".to_string()]), Some(b"tm\n".to_vec()));
|
|
assert_eq!(
|
|
startup_command(&["tm".to_string(), "work".to_string()]),
|
|
Some(b"tm work\n".to_vec())
|
|
);
|
|
}
|
|
|
|
#[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 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 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'\n".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'\n".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\n".to_vec())
|
|
);
|
|
assert_eq!(
|
|
resolved_startup_command(&[], &config, &host),
|
|
Some(b"tm\n".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"));
|
|
}
|
|
|
|
// --- 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);
|
|
}
|
|
|
|
/// The rendered overlay must be non-destructive: save cursor (ESC 7), park on
|
|
/// the bottom row, paint reverse-video cyan, then restore cursor (ESC 8) so
|
|
/// the application's cursor and full-screen TUIs are never disturbed.
|
|
#[test]
|
|
fn disconnect_status_overlay_is_save_restore_bottom_row() {
|
|
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 last row (row 24), column 1.
|
|
assert!(bytes.windows(7).any(|w| w == b"\x1b[24;1H"));
|
|
// Clears the row and uses reverse-video (7) + cyan (36).
|
|
assert!(bytes.windows(4).any(|w| w == b"\x1b[2K"));
|
|
assert!(bytes.windows(7).any(|w| w == b"\x1b[7;36m"));
|
|
// 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());
|
|
}
|
|
}
|