Add Dosh host workflows and reconnect QoL

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" <<EOF
+# Example:
+# [palav]
+# ssh = "palav"
+# dosh_host = "palav.dev"
+# port = 50000
+# default_command = "tm"
+# predict = false
+
+[default]
+ssh = "$default_server"
+dosh_host = "$host_udp"
+port = $port
+predict = false
 EOF
   fi
 fi

src/bin/dosh-client.rs

@@ -6,7 +6,7 @@ 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_server_config};
+use dosh::config::{expand_tilde, load_client_config, load_hosts_config, load_server_config};
 use dosh::crypto;
 use dosh::protocol::{
     self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
@@ -28,6 +28,8 @@ use tokio::sync::mpsc;
 struct Args {
     #[arg()]
     server: Option<String>,
+    #[arg()]
+    command: Vec<String>,
     #[arg(long)]
     session: Option<String>,
     #[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<Vec<u8>> {
+    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<u8> {
+    let mut command = command.to_string();
+    command.push('\n');
+    command.into_bytes()
+}
+
+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 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<Frame>,
     predict: bool,
+    startup_command: Option<Vec<u8>>,
+    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::<Vec<u8>>();
     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<Option<Frame>> {
+    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<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(())
+}
+
 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<Frame> {
         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())
+        );
+    }
 }

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<String>,
+    pub dosh_host: Option<String>,
+    pub port: Option<u16>,
+    pub default_command: Option<String>,
+    pub predict: Option<bool>,
+}
+
+#[derive(Debug, Clone, Default, Serialize, Deserialize)]
+pub struct HostsConfig {
+    #[serde(flatten)]
+    pub hosts: HashMap<String, HostConfig>,
+}
+
 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<PathBuf>) -> Result<ClientConfig> {
     toml::from_str(&raw).with_context(|| format!("parse {}", path.display()))
 }
 
+pub fn load_hosts_config(path: Option<PathBuf>) -> Result<HostsConfig> {
+    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() {