diff --git a/install.sh b/install.sh index 8d8ffc8..5393b32 100755 --- a/install.sh +++ b/install.sh @@ -11,6 +11,7 @@ from_current=0 start_server=1 force_config=0 update_cache="${DOSH_UPDATE_CACHE:-$HOME/.cache/dosh/source}" +quiet="${DOSH_UPDATE_QUIET:-0}" usage() { cat <<'EOF' @@ -136,10 +137,12 @@ else need git mkdir -p "$(dirname "$update_cache")" if [ -d "$update_cache/.git" ]; then + [ "$quiet" = "1" ] && echo "Updating Dosh source" git -C "$update_cache" remote set-url origin "$repo" git -C "$update_cache" fetch --depth 1 origin main git -C "$update_cache" checkout -q -B main FETCH_HEAD else + [ "$quiet" = "1" ] && echo "Downloading Dosh source" rm -rf "$update_cache" git clone --depth 1 --branch main "$repo" "$update_cache" >/dev/null fi @@ -147,10 +150,19 @@ else fi cd "$src_dir" +[ "$quiet" = "1" ] && echo "Building Dosh $role" if [ "$role" = "client" ]; then - cargo build --release --bin dosh-client + if [ "$quiet" = "1" ]; then + cargo build -q --release --bin dosh-client + else + cargo build --release --bin dosh-client + fi else - cargo build --release + if [ "$quiet" = "1" ]; then + cargo build -q --release + else + cargo build --release + fi fi bindir="$prefix/bin" @@ -245,8 +257,34 @@ dosh_port = $port default_session = "new" reconnect_timeout_secs = 5 view_only = false +predict = false cache_attach_tickets = true credential_cache = "~/.local/share/dosh/credentials" +EOF + fi + + hosts_config="$config_dir/hosts.toml" + if [ "$force_config" -eq 1 ] || [ ! -f "$hosts_config" ]; then + default_server="${server:-user@example.com}" + if [ -n "$dosh_host" ]; then + host_udp="$dosh_host" + else + host_udp="$default_server" + fi + cat >"$hosts_config" <, + #[arg()] + command: Vec, #[arg(long)] session: Option, #[arg(long)] @@ -82,7 +84,20 @@ async fn main() -> Result<()> { if args.server.as_deref() == Some("update") { return run_update(&config); } - let server = args.server.unwrap_or(config.server); + let requested_server = args.server.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 server = host.ssh.clone().unwrap_or_else(|| requested_server.clone()); + let startup_command = startup_command(&args.command).or_else(|| { + host.default_command + .as_deref() + .map(startup_command_from_string) + }); + let predict = args.predict || host.predict.unwrap_or(config.predict); let session = select_session(args.session.as_deref(), args.new); let mode = if args.view_only || config.view_only { "view-only" @@ -96,26 +111,42 @@ async fn main() -> Result<()> { .clone() .or_else(|| config.ssh_auth_command.clone()) .unwrap_or_else(|| "~/.local/bin/dosh-auth".to_string()); - let dosh_port = args.dosh_port.unwrap_or(config.dosh_port); - let cache_path = cache_path(&config.credential_cache, &server, &session, &mode); + 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) = size().unwrap_or((80, 24)); let started = Instant::now(); - let target_udp_host = - if let Some(host) = args.dosh_host.clone().or_else(|| config.dosh_host.clone()) { - host - } else if args.local_auth { - "127.0.0.1".to_string() - } else { - ssh_config_hostname(&server, ssh_port).unwrap_or_else(|err| { - log_debug( - args.verbose, - 2, - &format!("dosh could not resolve SSH config hostname, using {server}: {err:#}"), - ); - ssh_destination_host(&server) - }) - }; + 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 { + ssh_config_hostname(&server, ssh_port).unwrap_or_else(|err| { + log_debug( + args.verbose, + 2, + &format!("dosh could not resolve SSH config hostname, using {server}: {err:#}"), + ); + ssh_destination_host(&server) + }) + }; let credential = if !args.no_cache { load_cache(&cache_path).ok() @@ -140,7 +171,15 @@ async fn main() -> Result<()> { detach_once(&socket, &cred, 2).await?; return Ok(()); } - return run_terminal(socket, cred, Some(frame), args.predict).await; + return run_terminal( + socket, + cred, + Some(frame), + predict, + startup_command, + config.reconnect_timeout_secs, + ) + .await; } Err(err) => { log_debug( @@ -158,7 +197,15 @@ async fn main() -> Result<()> { detach_once(&socket, &cred, 2).await?; return Ok(()); } - return run_terminal(socket, cred, Some(frame), args.predict).await; + return run_terminal( + socket, + cred, + Some(frame), + predict, + startup_command, + config.reconnect_timeout_secs, + ) + .await; } Err(err) => { log_debug( @@ -215,7 +262,63 @@ async fn main() -> Result<()> { detach_once(&socket, &cred, 2).await?; return Ok(()); } - run_terminal(socket, cred, Some(first), args.predict).await + run_terminal( + socket, + cred, + Some(first), + predict, + startup_command, + config.reconnect_timeout_secs, + ) + .await +} + +fn startup_command(args: &[String]) -> Option> { + if args.is_empty() { + return None; + } + let mut command = args.join(" "); + command.push('\n'); + Some(command.into_bytes()) +} + +fn startup_command_from_string(command: &str) -> Vec { + let mut command = command.to_string(); + command.push('\n'); + command.into_bytes() +} + +fn run_sessions_command( + server: &str, + ssh_port: Option, + 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 status = command + .arg(server) + .arg("sh") + .arg("-lc") + .arg(script) + .status() + .context("run remote sessions command")?; + if !status.success() { + return Err(anyhow!("sessions command failed with status {status}")); + } + Ok(()) } fn run_update(config: &dosh::config::ClientConfig) -> Result<()> { @@ -230,7 +333,7 @@ fn run_update(config: &dosh::config::ClientConfig) -> Result<()> { .join("dosh") .join("source"); let mut script = format!( - "curl -fsSL {} | DOSH_REPO={} DOSH_PORT={} DOSH_UPDATE_CACHE={} sh -s -- client", + "curl -fsSL {} | DOSH_REPO={} DOSH_PORT={} DOSH_UPDATE_CACHE={} DOSH_UPDATE_QUIET=1 sh -s -- client", shell_word(&installer), shell_word(&repo), config.update_port.unwrap_or(config.dosh_port), @@ -600,6 +703,8 @@ async fn run_terminal( mut cred: CachedCredential, first_frame: Option, predict: bool, + startup_command: Option>, + reconnect_timeout_secs: u64, ) -> Result<()> { let _raw = RawMode::enter()?; let addr = resolve_addr(&cred.udp_host, cred.udp_port)?; @@ -619,6 +724,11 @@ async fn run_terminal( 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" + { + send_input(&socket, addr, &cred, &mut send_seq, bytes).await?; + } let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::>(); std::thread::Builder::new() @@ -648,18 +758,7 @@ async fn run_terminal( } if cred.mode != "view-only" { predictor.observe_input(&bytes)?; - 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?; + send_input(&socket, addr, &cred, &mut send_seq, bytes).await?; } } None => break, @@ -698,14 +797,46 @@ async fn run_terminal( PacketKind::Pong => {} PacketKind::AttachReject => { let reject: AttachReject = protocol::from_body(&packet.body)?; - return Err(anyhow!("server rejected session: {}", reject.reason)); + if reject.reason == "unknown client" { + if let Some(frame) = reconnect_with_ticket( + &socket, + &mut cred, + &mut send_seq, + last_size, + &mut frame_buffer, + &mut predictor, + ) + .await? + { + render_frame(&frame)?; + predictor.observe_output(&frame.bytes); + last_packet_at = Instant::now(); + } + } else { + return Err(anyhow!("server rejected session: {}", reject.reason)); + } } _ => {} } } _ = status_tick.tick() => { let stale = last_packet_at.elapsed(); - if stale >= Duration::from_secs(2) { + if stale >= Duration::from_secs(reconnect_timeout_secs.max(1)) { + if let Some(frame) = reconnect_with_ticket( + &socket, + &mut cred, + &mut send_seq, + last_size, + &mut frame_buffer, + &mut predictor, + ) + .await? + { + 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, @@ -725,6 +856,54 @@ async fn run_terminal( Ok(()) } +async fn reconnect_with_ticket( + socket: &UdpSocket, + cred: &mut CachedCredential, + send_seq: &mut u64, + size: (u16, u16), + frame_buffer: &mut FrameBuffer, + predictor: &mut Predictor, +) -> Result> { + let Ok((frame, next)) = try_ticket_attach(socket, cred, size.0, size.1).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, +) -> 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(()) +} + struct Predictor { enabled: bool, alternate_screen: bool, @@ -740,6 +919,11 @@ impl Predictor { } } + fn reset(&mut self) { + self.pending.clear(); + self.alternate_screen = false; + } + fn observe_input(&mut self, bytes: &[u8]) -> Result<()> { if !self.enabled || self.alternate_screen { return Ok(()); @@ -813,6 +997,10 @@ struct FrameBuffer { } impl FrameBuffer { + fn clear(&mut self) { + self.pending.clear(); + } + fn accept(&mut self, frame: Frame, last_rendered_seq: &mut u64) -> Vec { if frame.snapshot { if frame.output_seq < *last_rendered_seq { @@ -958,7 +1146,9 @@ const TERMINAL_CLEANUP: &[u8] = concat!( #[cfg(test)] mod tests { - use super::{FrameBuffer, Predictor, parse_ssh_config_hostname, ssh_destination_host}; + use super::{ + FrameBuffer, Predictor, parse_ssh_config_hostname, ssh_destination_host, startup_command, + }; use dosh::protocol::Frame; #[test] @@ -1061,4 +1251,14 @@ mod tests { predictor.observe_output(b"\x1b[?1049l"); assert!(!predictor.alternate_screen); } + + #[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()) + ); + } } diff --git a/src/config.rs b/src/config.rs index 1382559..cec0aa2 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,5 +1,6 @@ use anyhow::{Context, Result}; use serde::{Deserialize, Serialize}; +use std::collections::HashMap; use std::fs; use std::path::PathBuf; @@ -54,10 +55,27 @@ pub struct ClientConfig { pub default_session: String, pub reconnect_timeout_secs: u64, pub view_only: bool, + #[serde(default)] + pub predict: bool, pub cache_attach_tickets: bool, pub credential_cache: String, } +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct HostConfig { + pub ssh: Option, + pub dosh_host: Option, + pub port: Option, + pub default_command: Option, + pub predict: Option, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct HostsConfig { + #[serde(flatten)] + pub hosts: HashMap, +} + impl Default for ClientConfig { fn default() -> Self { Self { @@ -71,6 +89,7 @@ impl Default for ClientConfig { default_session: "new".to_string(), reconnect_timeout_secs: 5, view_only: false, + predict: false, cache_attach_tickets: true, credential_cache: "~/.local/share/dosh/credentials".to_string(), } @@ -95,6 +114,15 @@ pub fn load_client_config(path: Option) -> Result { toml::from_str(&raw).with_context(|| format!("parse {}", path.display())) } +pub fn load_hosts_config(path: Option) -> Result { + let path = path.unwrap_or_else(|| expand_tilde("~/.config/dosh/hosts.toml")); + if !path.exists() { + return Ok(HostsConfig::default()); + } + let raw = fs::read_to_string(&path).with_context(|| format!("read {}", path.display()))?; + toml::from_str(&raw).with_context(|| format!("parse {}", path.display())) +} + pub fn expand_tilde(path: &str) -> PathBuf { if let Some(rest) = path.strip_prefix("~/") { if let Some(home) = dirs::home_dir() {