Palav/dosh
1
0
Fork 0
Code Issues Pull Requests Actions 46 Packages Projects Releases 51 Wiki Activity

Add Dosh host workflows and reconnect QoL
ci / test (push) Has been cancelled
Details
ci / remote-bench (push) Has been cancelled
Details

Browse Source
This commit is contained in:
DuProcess
2026-06-12 13:48:44 -04:00
parent c56e8f9fe1
commit e3f61a4efb
3 changed files with 306 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/bin/dosh-client.rs
+238 -38
View File
@@ -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
+28
View File
@@ -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() {