Initial Dosh implementation
This commit is contained in:
Codex
@@ -0,0 +1,59 @@
|
||||
use anyhow::{Context, Result};
|
||||
use clap::Parser;
|
||||
use dosh::auth::{build_bootstrap, encode_bootstrap, load_or_create_server_secret};
|
||||
use dosh::config::load_server_config;
|
||||
|
||||
#[derive(Debug, Parser)]
|
||||
struct Args {
|
||||
#[arg(long, default_value_t = 1)]
|
||||
protocol: u8,
|
||||
#[arg(long)]
|
||||
nonce: String,
|
||||
#[arg(long, default_value = "default")]
|
||||
session: String,
|
||||
#[arg(long, default_value = "read-write")]
|
||||
mode: String,
|
||||
#[arg(long, default_value = "80x24")]
|
||||
size: String,
|
||||
#[arg(long, default_value = "dev")]
|
||||
client_version: String,
|
||||
#[arg(long)]
|
||||
udp_host: Option<String>,
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
let args = Args::parse();
|
||||
anyhow::ensure!(args.protocol == 1, "unsupported protocol {}", args.protocol);
|
||||
let config = load_server_config(None)?;
|
||||
let secret = load_or_create_server_secret(&config)?;
|
||||
let nonce = parse_nonce(&args.nonce)?;
|
||||
let size = parse_size(&args.size)?;
|
||||
let user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
|
||||
let udp_host = args.udp_host.unwrap_or_else(|| "127.0.0.1".to_string());
|
||||
let resp = build_bootstrap(
|
||||
&config,
|
||||
&secret,
|
||||
user,
|
||||
args.session,
|
||||
args.mode,
|
||||
size,
|
||||
nonce,
|
||||
udp_host,
|
||||
)?;
|
||||
println!("{}", encode_bootstrap(&resp)?);
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn parse_nonce(raw: &str) -> Result<[u8; 12]> {
|
||||
let bytes = base64::Engine::decode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, raw)
|
||||
.context("decode nonce")?;
|
||||
anyhow::ensure!(bytes.len() == 12, "nonce must decode to 12 bytes");
|
||||
let mut out = [0u8; 12];
|
||||
out.copy_from_slice(&bytes);
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn parse_size(raw: &str) -> Result<(u16, u16)> {
|
||||
let (cols, rows) = raw.split_once('x').context("size must be COLSxROWS")?;
|
||||
Ok((cols.parse()?, rows.parse()?))
|
||||
}
|
||||
@@ -0,0 +1,233 @@
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use clap::Parser;
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
use std::time::{Duration, Instant};
|
||||
|
||||
#[derive(Debug, Parser)]
|
||||
#[command(name = "dosh-bench")]
|
||||
struct Args {
|
||||
#[arg(long, default_value = "local")]
|
||||
server: String,
|
||||
#[arg(long, default_value = "default")]
|
||||
session: String,
|
||||
#[arg(long, default_value_t = 22)]
|
||||
ssh_port: u16,
|
||||
#[arg(long, default_value_t = 50000)]
|
||||
dosh_port: u16,
|
||||
#[arg(long)]
|
||||
dosh_host: Option<String>,
|
||||
#[arg(long, default_value_t = 3)]
|
||||
iterations: usize,
|
||||
#[arg(long)]
|
||||
local_auth: bool,
|
||||
#[arg(long)]
|
||||
client: Option<PathBuf>,
|
||||
#[arg(long, default_value = "dosh-auth")]
|
||||
ssh_auth_command: String,
|
||||
#[arg(long)]
|
||||
ssh_key: Option<PathBuf>,
|
||||
#[arg(long)]
|
||||
ssh_known_hosts: Option<PathBuf>,
|
||||
#[arg(long)]
|
||||
ssh_control_path: Option<PathBuf>,
|
||||
#[arg(long)]
|
||||
controlmaster: bool,
|
||||
#[arg(long)]
|
||||
no_cache: bool,
|
||||
#[arg(long)]
|
||||
assert_ssh_plus_ms: Option<f64>,
|
||||
}
|
||||
|
||||
fn main() -> Result<()> {
|
||||
let args = Args::parse();
|
||||
let client = args.client.clone().unwrap_or_else(default_client_path);
|
||||
let mut ssh_times = Vec::new();
|
||||
let mut dosh_times = Vec::new();
|
||||
let generated_control_path = if args.controlmaster {
|
||||
Some(std::env::temp_dir().join(format!("dosh-bench-control-{}", std::process::id())))
|
||||
} else {
|
||||
None
|
||||
};
|
||||
let control_path = generated_control_path
|
||||
.as_ref()
|
||||
.or(args.ssh_control_path.as_ref());
|
||||
let _controlmaster = if let Some(path) = generated_control_path.as_ref() {
|
||||
Some(ControlMaster::start(&args, path.clone())?)
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
for _ in 0..args.iterations.max(1) {
|
||||
if !args.local_auth {
|
||||
let mut ssh = Command::new("ssh");
|
||||
add_ssh_options(&mut ssh, &args, control_path);
|
||||
ssh.arg(&args.server).arg("true");
|
||||
ssh_times.push(time_command(&mut ssh)?);
|
||||
}
|
||||
|
||||
let mut cmd = Command::new(&client);
|
||||
cmd.arg("--attach-only")
|
||||
.arg("--session")
|
||||
.arg(&args.session)
|
||||
.arg("--dosh-port")
|
||||
.arg(args.dosh_port.to_string());
|
||||
if let Some(host) = &args.dosh_host {
|
||||
cmd.arg("--dosh-host").arg(host);
|
||||
}
|
||||
if args.local_auth {
|
||||
cmd.arg("--local-auth").arg(&args.server);
|
||||
} else {
|
||||
cmd.arg("--ssh-port")
|
||||
.arg(args.ssh_port.to_string())
|
||||
.arg("--ssh-auth-command")
|
||||
.arg(&args.ssh_auth_command);
|
||||
if args.no_cache {
|
||||
cmd.arg("--no-cache");
|
||||
}
|
||||
if let Some(key) = &args.ssh_key {
|
||||
cmd.arg("--ssh-key").arg(key);
|
||||
}
|
||||
if let Some(known_hosts) = &args.ssh_known_hosts {
|
||||
cmd.arg("--ssh-known-hosts").arg(known_hosts);
|
||||
}
|
||||
if let Some(control_path) = control_path {
|
||||
cmd.arg("--ssh-control-path").arg(control_path);
|
||||
}
|
||||
cmd.arg(&args.server);
|
||||
}
|
||||
dosh_times.push(time_command(&mut cmd)?);
|
||||
}
|
||||
|
||||
if !ssh_times.is_empty() {
|
||||
println!(
|
||||
"ssh_true_ms avg={:.2} samples={:?}",
|
||||
avg_ms(&ssh_times),
|
||||
ssh_times
|
||||
);
|
||||
}
|
||||
println!(
|
||||
"dosh_attach_ms avg={:.2} samples={:?}",
|
||||
avg_ms(&dosh_times),
|
||||
dosh_times
|
||||
);
|
||||
if let Some(margin) = args.assert_ssh_plus_ms {
|
||||
if ssh_times.is_empty() {
|
||||
return Err(anyhow!(
|
||||
"--assert-ssh-plus-ms requires non-local SSH benchmark"
|
||||
));
|
||||
}
|
||||
let ssh_avg = avg_ms(&ssh_times);
|
||||
let dosh_avg = avg_ms(&dosh_times);
|
||||
if dosh_avg > ssh_avg + margin {
|
||||
return Err(anyhow!(
|
||||
"dosh attach avg {dosh_avg:.2}ms exceeded ssh avg {ssh_avg:.2}ms + {margin:.2}ms"
|
||||
));
|
||||
}
|
||||
println!("gate ok: dosh avg {dosh_avg:.2}ms <= ssh avg {ssh_avg:.2}ms + {margin:.2}ms");
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn add_ssh_options(cmd: &mut Command, args: &Args, control_path: Option<&PathBuf>) {
|
||||
cmd.arg("-p").arg(args.ssh_port.to_string()).arg("-T");
|
||||
if let Some(key) = &args.ssh_key {
|
||||
cmd.arg("-i").arg(key);
|
||||
}
|
||||
if let Some(known_hosts) = &args.ssh_known_hosts {
|
||||
cmd.arg("-o")
|
||||
.arg(format!("UserKnownHostsFile={}", known_hosts.display()));
|
||||
}
|
||||
if let Some(control_path) = control_path {
|
||||
cmd.arg("-S").arg(control_path);
|
||||
}
|
||||
}
|
||||
|
||||
struct ControlMaster {
|
||||
server: String,
|
||||
ssh_port: u16,
|
||||
ssh_key: Option<PathBuf>,
|
||||
ssh_known_hosts: Option<PathBuf>,
|
||||
control_path: PathBuf,
|
||||
}
|
||||
|
||||
impl ControlMaster {
|
||||
fn start(args: &Args, control_path: PathBuf) -> Result<Self> {
|
||||
let mut cmd = Command::new("ssh");
|
||||
cmd.arg("-p")
|
||||
.arg(args.ssh_port.to_string())
|
||||
.arg("-T")
|
||||
.arg("-M")
|
||||
.arg("-S")
|
||||
.arg(&control_path)
|
||||
.arg("-f")
|
||||
.arg("-N")
|
||||
.arg("-o")
|
||||
.arg("ControlPersist=60")
|
||||
.arg("-o")
|
||||
.arg("ExitOnForwardFailure=yes");
|
||||
if let Some(key) = &args.ssh_key {
|
||||
cmd.arg("-i").arg(key);
|
||||
}
|
||||
if let Some(known_hosts) = &args.ssh_known_hosts {
|
||||
cmd.arg("-o")
|
||||
.arg(format!("UserKnownHostsFile={}", known_hosts.display()));
|
||||
}
|
||||
cmd.arg(&args.server);
|
||||
time_command(&mut cmd).context("start SSH ControlMaster")?;
|
||||
Ok(Self {
|
||||
server: args.server.clone(),
|
||||
ssh_port: args.ssh_port,
|
||||
ssh_key: args.ssh_key.clone(),
|
||||
ssh_known_hosts: args.ssh_known_hosts.clone(),
|
||||
control_path,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for ControlMaster {
|
||||
fn drop(&mut self) {
|
||||
let mut cmd = Command::new("ssh");
|
||||
cmd.arg("-p")
|
||||
.arg(self.ssh_port.to_string())
|
||||
.arg("-T")
|
||||
.arg("-S")
|
||||
.arg(&self.control_path)
|
||||
.arg("-O")
|
||||
.arg("exit");
|
||||
if let Some(key) = &self.ssh_key {
|
||||
cmd.arg("-i").arg(key);
|
||||
}
|
||||
if let Some(known_hosts) = &self.ssh_known_hosts {
|
||||
cmd.arg("-o")
|
||||
.arg(format!("UserKnownHostsFile={}", known_hosts.display()));
|
||||
}
|
||||
let _ = cmd.arg(&self.server).output();
|
||||
}
|
||||
}
|
||||
|
||||
fn time_command(cmd: &mut Command) -> Result<Duration> {
|
||||
let start = Instant::now();
|
||||
let output = cmd.output().with_context(|| format!("run {:?}", cmd))?;
|
||||
if !output.status.success() {
|
||||
return Err(anyhow!(
|
||||
"command failed {:?}\nstdout:\n{}\nstderr:\n{}",
|
||||
cmd,
|
||||
String::from_utf8_lossy(&output.stdout),
|
||||
String::from_utf8_lossy(&output.stderr)
|
||||
));
|
||||
}
|
||||
Ok(start.elapsed())
|
||||
}
|
||||
|
||||
fn avg_ms(samples: &[Duration]) -> f64 {
|
||||
let total: f64 = samples.iter().map(|d| d.as_secs_f64() * 1000.0).sum();
|
||||
total / samples.len() as f64
|
||||
}
|
||||
|
||||
fn default_client_path() -> PathBuf {
|
||||
std::env::current_exe()
|
||||
.ok()
|
||||
.and_then(|path| path.parent().map(|parent| parent.join("dosh-client")))
|
||||
.unwrap_or_else(|| PathBuf::from("dosh-client"))
|
||||
}
|
||||
@@ -0,0 +1,617 @@
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use base64::Engine;
|
||||
use base64::engine::general_purpose::URL_SAFE_NO_PAD;
|
||||
use clap::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_server_config};
|
||||
use dosh::crypto;
|
||||
use dosh::protocol::{
|
||||
self, AttachOk, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input, PacketKind,
|
||||
ResumeRequest, SERVER_TO_CLIENT, TicketAttachBody, TicketAttachEnvelope,
|
||||
TicketAttachOkEnvelope,
|
||||
};
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::fs;
|
||||
use std::io::{Read, Write};
|
||||
use std::net::{SocketAddr, ToSocketAddrs};
|
||||
use std::process::Command;
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::net::UdpSocket;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[derive(Debug, Parser)]
|
||||
#[command(name = "dosh-client")]
|
||||
struct Args {
|
||||
#[arg()]
|
||||
server: Option<String>,
|
||||
#[arg(long, default_value = "default")]
|
||||
session: String,
|
||||
#[arg(long)]
|
||||
view_only: bool,
|
||||
#[arg(long)]
|
||||
local_auth: bool,
|
||||
#[arg(long)]
|
||||
no_cache: bool,
|
||||
#[arg(long)]
|
||||
attach_only: bool,
|
||||
#[arg(long)]
|
||||
ssh_port: Option<u16>,
|
||||
#[arg(long, default_value = "dosh-auth")]
|
||||
ssh_auth_command: String,
|
||||
#[arg(long)]
|
||||
ssh_key: Option<std::path::PathBuf>,
|
||||
#[arg(long)]
|
||||
ssh_known_hosts: Option<std::path::PathBuf>,
|
||||
#[arg(long)]
|
||||
ssh_control_path: Option<std::path::PathBuf>,
|
||||
#[arg(long)]
|
||||
dosh_port: Option<u16>,
|
||||
#[arg(long)]
|
||||
dosh_host: Option<String>,
|
||||
}
|
||||
|
||||
#[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,
|
||||
}
|
||||
|
||||
#[tokio::main(flavor = "current_thread")]
|
||||
async fn main() -> Result<()> {
|
||||
let args = Args::parse();
|
||||
let config = load_client_config(None).unwrap_or_default();
|
||||
let server = args.server.unwrap_or(config.server);
|
||||
let session = args.session;
|
||||
let mode = if args.view_only || config.view_only {
|
||||
"view-only"
|
||||
} else {
|
||||
"read-write"
|
||||
}
|
||||
.to_string();
|
||||
let ssh_port = args.ssh_port.unwrap_or(config.ssh_port);
|
||||
let dosh_port = args.dosh_port.unwrap_or(config.dosh_port);
|
||||
let cache_path = cache_path(&config.credential_cache, &server, &session, &mode);
|
||||
let (cols, rows) = size().unwrap_or((80, 24));
|
||||
|
||||
let started = Instant::now();
|
||||
let target_udp_host = args
|
||||
.dosh_host
|
||||
.clone()
|
||||
.or_else(|| config.dosh_host.clone())
|
||||
.unwrap_or_else(|| {
|
||||
if args.local_auth {
|
||||
"127.0.0.1".to_string()
|
||||
} else {
|
||||
ssh_destination_host(&server)
|
||||
}
|
||||
});
|
||||
|
||||
let credential = if !args.no_cache {
|
||||
load_cache(&cache_path).ok()
|
||||
} else {
|
||||
None
|
||||
};
|
||||
|
||||
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;
|
||||
eprintln!(
|
||||
"dosh timing credential_lookup_end={}ms",
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
match try_resume(&socket, &cached, cols, rows).await {
|
||||
Ok((frame, cred)) => {
|
||||
eprintln!(
|
||||
"dosh timing udp_resume_ready={}ms",
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
if !args.no_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)).await;
|
||||
}
|
||||
Err(err) => {
|
||||
eprintln!("dosh resume failed, trying ticket attach before SSH: {err:#}");
|
||||
if config.cache_attach_tickets && !args.no_cache {
|
||||
match try_ticket_attach(&socket, &cached, cols, rows).await {
|
||||
Ok((frame, cred)) => {
|
||||
eprintln!(
|
||||
"dosh timing udp_ticket_attach_ready={}ms",
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
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)).await;
|
||||
}
|
||||
Err(err) => {
|
||||
eprintln!(
|
||||
"dosh ticket attach failed, falling back to SSH bootstrap: {err:#}"
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
let bootstrap_start = Instant::now();
|
||||
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,
|
||||
&args.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
|
||||
};
|
||||
eprintln!(
|
||||
"dosh timing ssh_bootstrap={}ms",
|
||||
bootstrap_start.elapsed().as_millis()
|
||||
);
|
||||
|
||||
let (ok, mut cred) = bootstrap_attach(&socket, &server, &bootstrap, cols, rows).await?;
|
||||
cred.last_rendered_seq = ok.initial_seq;
|
||||
if !args.no_cache {
|
||||
save_cache(&cache_path, &cred)?;
|
||||
}
|
||||
eprintln!(
|
||||
"dosh timing terminal_ready={}ms",
|
||||
started.elapsed().as_millis()
|
||||
);
|
||||
let first = Frame {
|
||||
session: ok.session,
|
||||
output_seq: ok.initial_seq,
|
||||
bytes: ok.snapshot,
|
||||
snapshot: true,
|
||||
};
|
||||
if args.attach_only {
|
||||
render_frame(&first)?;
|
||||
detach_once(&socket, &cred, 2).await?;
|
||||
return Ok(());
|
||||
}
|
||||
run_terminal(socket, cred, Some(first)).await
|
||||
}
|
||||
|
||||
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 ssh_bootstrap(
|
||||
server: &str,
|
||||
ssh_port: u16,
|
||||
ssh_auth_command: &str,
|
||||
ssh_key: Option<&std::path::Path>,
|
||||
ssh_known_hosts: Option<&std::path::Path>,
|
||||
ssh_control_path: Option<&std::path::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");
|
||||
command.arg("-p").arg(ssh_port.to_string()).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("--nonce")
|
||||
.arg(nonce_b64)
|
||||
.arg("--session")
|
||||
.arg(session)
|
||||
.arg("--mode")
|
||||
.arg(mode)
|
||||
.arg("--size")
|
||||
.arg(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 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 bootstrap_attach(
|
||||
socket: &UdpSocket,
|
||||
server_name: &str,
|
||||
bootstrap: &BootstrapResponse,
|
||||
cols: u16,
|
||||
rows: u16,
|
||||
) -> Result<(AttachOk, CachedCredential)> {
|
||||
let addr = resolve_addr(&bootstrap.udp_host, bootstrap.udp_port)?;
|
||||
let req = BootstrapAttachRequest {
|
||||
bootstrap: bootstrap.clone(),
|
||||
cols,
|
||||
rows,
|
||||
};
|
||||
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 mut buf = vec![0u8; 65535];
|
||||
let (n, _) = tokio::time::timeout(Duration::from_secs(5), socket.recv_from(&mut buf)).await??;
|
||||
let packet = protocol::decode(&buf[..n])?;
|
||||
if packet.header.kind != PacketKind::AttachOk {
|
||||
return Err(anyhow!("attach rejected or unexpected response"));
|
||||
}
|
||||
let plain = protocol::decrypt_body(&packet, &bootstrap.session_key, SERVER_TO_CLIENT)?;
|
||||
let ok: AttachOk = protocol::from_body(&plain)?;
|
||||
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,
|
||||
) -> 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,
|
||||
})?;
|
||||
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 mut buf = vec![0u8; 65535];
|
||||
let (n, _) =
|
||||
tokio::time::timeout(Duration::from_millis(700), socket.recv_from(&mut buf)).await??;
|
||||
let packet = protocol::decode(&buf[..n])?;
|
||||
if packet.header.kind != PacketKind::AttachOk {
|
||||
return Err(anyhow!("ticket attach rejected"));
|
||||
}
|
||||
let envelope: TicketAttachOkEnvelope = protocol::from_body(&packet.body)?;
|
||||
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 plain = crypto::open(
|
||||
&response_key,
|
||||
&envelope.server_nonce,
|
||||
b"dosh-ticket-attach-ok-v1",
|
||||
&envelope.ciphertext,
|
||||
)?;
|
||||
let ok: AttachOk = protocol::from_body(&plain)?;
|
||||
let frame = Frame {
|
||||
session: ok.session.clone(),
|
||||
output_seq: ok.initial_seq,
|
||||
bytes: ok.snapshot.clone(),
|
||||
snapshot: true,
|
||||
};
|
||||
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(
|
||||
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 mut buf = vec![0u8; 65535];
|
||||
let (n, _) =
|
||||
tokio::time::timeout(Duration::from_millis(700), socket.recv_from(&mut buf)).await??;
|
||||
let packet = protocol::decode(&buf[..n])?;
|
||||
if packet.header.kind != PacketKind::ResumeOk {
|
||||
return Err(anyhow!("resume rejected"));
|
||||
}
|
||||
let plain = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT)?;
|
||||
let frame: Frame = protocol::from_body(&plain)?;
|
||||
let mut next = cached.clone();
|
||||
next.last_rendered_seq = frame.output_seq;
|
||||
Ok((frame, next))
|
||||
}
|
||||
|
||||
async fn run_terminal(
|
||||
socket: UdpSocket,
|
||||
mut cred: CachedCredential,
|
||||
first_frame: Option<Frame>,
|
||||
) -> Result<()> {
|
||||
let _raw = RawMode::enter()?;
|
||||
let addr = resolve_addr(&cred.udp_host, cred.udp_port)?;
|
||||
let mut send_seq = 2u64;
|
||||
if let Some(frame) = first_frame {
|
||||
render_frame(&frame)?;
|
||||
cred.last_rendered_seq = frame.output_seq;
|
||||
send_ack(&socket, addr, &cred, &mut send_seq).await?;
|
||||
}
|
||||
|
||||
let (stdin_tx, mut stdin_rx) = mpsc::unbounded_channel::<Vec<u8>>();
|
||||
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,
|
||||
}
|
||||
}
|
||||
})?;
|
||||
|
||||
let mut recv_buf = vec![0u8; 65535];
|
||||
loop {
|
||||
tokio::select! {
|
||||
Some(bytes) = stdin_rx.recv() => {
|
||||
if bytes == [0x1d] {
|
||||
break;
|
||||
}
|
||||
if cred.mode != "view-only" {
|
||||
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?;
|
||||
}
|
||||
}
|
||||
recv = socket.recv_from(&mut recv_buf) => {
|
||||
let (n, _) = recv?;
|
||||
let packet = protocol::decode(&recv_buf[..n])?;
|
||||
match packet.header.kind {
|
||||
PacketKind::Frame | PacketKind::ResumeOk => {
|
||||
let plain = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT)?;
|
||||
let frame: Frame = protocol::from_body(&plain)?;
|
||||
render_frame(&frame)?;
|
||||
cred.last_rendered_seq = frame.output_seq;
|
||||
send_ack(&socket, addr, &cred, &mut send_seq).await?;
|
||||
}
|
||||
PacketKind::Pong => {}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
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(())
|
||||
}
|
||||
|
||||
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 _ = disable_raw_mode();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,743 @@
|
||||
use anyhow::{Context, Result, anyhow};
|
||||
use clap::{Parser, Subcommand};
|
||||
use dosh::auth::{
|
||||
build_bootstrap, encode_bootstrap, load_or_create_server_secret, open_attach_ticket,
|
||||
verify_bootstrap,
|
||||
};
|
||||
use dosh::config::{ServerConfig, load_server_config};
|
||||
use dosh::crypto;
|
||||
use dosh::protocol::{
|
||||
self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
|
||||
PacketKind, ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT, TicketAttachBody,
|
||||
TicketAttachEnvelope, TicketAttachOkEnvelope,
|
||||
};
|
||||
use dosh::pty::{PtyHandle, PtyOutput, spawn_pty_session};
|
||||
use std::collections::{HashMap, VecDeque};
|
||||
use std::net::SocketAddr;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::{Duration, Instant};
|
||||
use tokio::net::UdpSocket;
|
||||
use tokio::sync::mpsc;
|
||||
|
||||
#[derive(Debug, Parser)]
|
||||
#[command(name = "dosh-server")]
|
||||
struct Args {
|
||||
#[command(subcommand)]
|
||||
command: Command,
|
||||
}
|
||||
|
||||
#[derive(Debug, Subcommand)]
|
||||
enum Command {
|
||||
Serve {
|
||||
#[arg(long)]
|
||||
config: Option<std::path::PathBuf>,
|
||||
},
|
||||
Auth {
|
||||
#[arg(long, default_value_t = 1)]
|
||||
protocol: u8,
|
||||
#[arg(long)]
|
||||
nonce: String,
|
||||
#[arg(long, default_value = "default")]
|
||||
session: String,
|
||||
#[arg(long, default_value = "read-write")]
|
||||
mode: String,
|
||||
#[arg(long, default_value = "80x24")]
|
||||
size: String,
|
||||
#[arg(long, default_value = "dev")]
|
||||
client_version: String,
|
||||
#[arg(long)]
|
||||
udp_host: Option<String>,
|
||||
},
|
||||
}
|
||||
|
||||
#[tokio::main(flavor = "current_thread")]
|
||||
async fn main() -> Result<()> {
|
||||
let args = Args::parse();
|
||||
match args.command {
|
||||
Command::Serve { config } => serve(config).await,
|
||||
Command::Auth {
|
||||
protocol,
|
||||
nonce,
|
||||
session,
|
||||
mode,
|
||||
size,
|
||||
client_version: _,
|
||||
udp_host,
|
||||
} => {
|
||||
anyhow::ensure!(protocol == 1, "unsupported protocol {protocol}");
|
||||
let config = load_server_config(None)?;
|
||||
let secret = load_or_create_server_secret(&config)?;
|
||||
let nonce = parse_nonce(&nonce)?;
|
||||
let size = parse_size(&size)?;
|
||||
let user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
|
||||
let udp_host = udp_host.unwrap_or_else(|| "127.0.0.1".to_string());
|
||||
let resp =
|
||||
build_bootstrap(&config, &secret, user, session, mode, size, nonce, udp_host)?;
|
||||
println!("{}", encode_bootstrap(&resp)?);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn serve(config_path: Option<std::path::PathBuf>) -> Result<()> {
|
||||
let config = load_server_config(config_path)?;
|
||||
let secret = load_or_create_server_secret(&config)?;
|
||||
let bind = format!("{}:{}", config.bind, config.port);
|
||||
let socket = Arc::new(
|
||||
UdpSocket::bind(&bind)
|
||||
.await
|
||||
.with_context(|| format!("bind {bind}"))?,
|
||||
);
|
||||
eprintln!("dosh-server listening on {bind}");
|
||||
|
||||
let (pty_tx, mut pty_rx) = mpsc::unbounded_channel();
|
||||
let state = Arc::new(Mutex::new(ServerState::new(
|
||||
config.clone(),
|
||||
secret,
|
||||
pty_tx.clone(),
|
||||
)));
|
||||
{
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
for session in config.prewarm_sessions.clone() {
|
||||
locked.ensure_session(&session, 80, 24)?;
|
||||
}
|
||||
}
|
||||
|
||||
let output_state = Arc::clone(&state);
|
||||
let output_socket = Arc::clone(&socket);
|
||||
tokio::spawn(async move {
|
||||
while let Some(output) = pty_rx.recv().await {
|
||||
if let Err(err) = broadcast_output(&output_state, &output_socket, output).await {
|
||||
eprintln!("broadcast error: {err:#}");
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let retransmit_state = Arc::clone(&state);
|
||||
let retransmit_socket = Arc::clone(&socket);
|
||||
tokio::spawn(async move {
|
||||
let mut interval = tokio::time::interval(Duration::from_millis(100));
|
||||
loop {
|
||||
interval.tick().await;
|
||||
if let Err(err) = retransmit_pending(&retransmit_state, &retransmit_socket).await {
|
||||
eprintln!("retransmit error: {err:#}");
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let mut buf = vec![0u8; 65535];
|
||||
loop {
|
||||
let (n, peer) = socket.recv_from(&mut buf).await?;
|
||||
if let Err(err) = handle_packet(&state, &socket, peer, &buf[..n]).await {
|
||||
eprintln!("packet from {peer}: {err:#}");
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
struct ServerState {
|
||||
config: ServerConfig,
|
||||
secret: [u8; 32],
|
||||
pty_tx: mpsc::UnboundedSender<PtyOutput>,
|
||||
sessions: HashMap<String, Session>,
|
||||
}
|
||||
|
||||
struct Session {
|
||||
pty: PtyHandle,
|
||||
parser: vt100::Parser,
|
||||
clients: HashMap<[u8; 16], ClientState>,
|
||||
output_seq: u64,
|
||||
recent: VecDeque<Vec<u8>>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ClientState {
|
||||
endpoint: SocketAddr,
|
||||
mode: String,
|
||||
session_key: [u8; 32],
|
||||
last_acked: u64,
|
||||
replay: ReplayWindow,
|
||||
send_seq: u64,
|
||||
cols: u16,
|
||||
rows: u16,
|
||||
last_seen: Instant,
|
||||
pending: VecDeque<PendingFrame>,
|
||||
last_screen: Option<vt100::Screen>,
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct PendingFrame {
|
||||
output_seq: u64,
|
||||
packet: Vec<u8>,
|
||||
last_sent: Instant,
|
||||
attempts: u8,
|
||||
}
|
||||
|
||||
impl ServerState {
|
||||
fn new(
|
||||
config: ServerConfig,
|
||||
secret: [u8; 32],
|
||||
pty_tx: mpsc::UnboundedSender<PtyOutput>,
|
||||
) -> Self {
|
||||
Self {
|
||||
config,
|
||||
secret,
|
||||
pty_tx,
|
||||
sessions: HashMap::new(),
|
||||
}
|
||||
}
|
||||
|
||||
fn ensure_session(&mut self, name: &str, cols: u16, rows: u16) -> Result<()> {
|
||||
if self.sessions.contains_key(name) {
|
||||
return Ok(());
|
||||
}
|
||||
let pty = spawn_pty_session(
|
||||
name.to_string(),
|
||||
&self.config.shell,
|
||||
cols.max(1),
|
||||
rows.max(1),
|
||||
self.pty_tx.clone(),
|
||||
)?;
|
||||
self.sessions.insert(
|
||||
name.to_string(),
|
||||
Session {
|
||||
pty,
|
||||
parser: vt100::Parser::new(rows.max(1), cols.max(1), self.config.scrollback),
|
||||
clients: HashMap::new(),
|
||||
output_seq: 0,
|
||||
recent: VecDeque::with_capacity(self.config.scrollback),
|
||||
},
|
||||
);
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_packet(
|
||||
state: &Arc<Mutex<ServerState>>,
|
||||
socket: &Arc<UdpSocket>,
|
||||
peer: SocketAddr,
|
||||
raw: &[u8],
|
||||
) -> Result<()> {
|
||||
let packet = protocol::decode(raw)?;
|
||||
match packet.header.kind {
|
||||
PacketKind::BootstrapAttachRequest => {
|
||||
handle_bootstrap_attach(state, socket, peer, packet.body).await
|
||||
}
|
||||
PacketKind::TicketAttachRequest => {
|
||||
handle_ticket_attach(state, socket, peer, packet.body).await
|
||||
}
|
||||
PacketKind::ResumeRequest => handle_resume(state, socket, peer, &packet).await,
|
||||
PacketKind::Input => handle_input(state, peer, &packet).await,
|
||||
PacketKind::Resize => handle_resize(state, peer, &packet).await,
|
||||
PacketKind::Ping => handle_ping(state, socket, peer, &packet).await,
|
||||
PacketKind::Ack => handle_ack(state, &packet).await,
|
||||
PacketKind::Detach => handle_detach(state, &packet).await,
|
||||
_ => Ok(()),
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_bootstrap_attach(
|
||||
state: &Arc<Mutex<ServerState>>,
|
||||
socket: &Arc<UdpSocket>,
|
||||
peer: SocketAddr,
|
||||
body: Vec<u8>,
|
||||
) -> Result<()> {
|
||||
let req: BootstrapAttachRequest = protocol::from_body(&body)?;
|
||||
let (client_id, key, key_id, session_name, mode, output_seq, snapshot) = {
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
if !verify_bootstrap(&req.bootstrap, &locked.secret)? {
|
||||
return send_reject(socket, peer, "invalid or expired bootstrap").await;
|
||||
}
|
||||
if !locked.sessions.contains_key(&req.bootstrap.session) {
|
||||
if locked.config.create_on_attach {
|
||||
locked.ensure_session(&req.bootstrap.session, req.cols, req.rows)?;
|
||||
} else {
|
||||
return send_reject(socket, peer, "session does not exist").await;
|
||||
}
|
||||
}
|
||||
let session = locked
|
||||
.sessions
|
||||
.get_mut(&req.bootstrap.session)
|
||||
.expect("session exists");
|
||||
let client_id = crypto::random_16();
|
||||
let snapshot = session.parser.screen().state_formatted();
|
||||
let screen = session.parser.screen().clone();
|
||||
let output_seq = session.output_seq;
|
||||
session.clients.insert(
|
||||
client_id,
|
||||
ClientState {
|
||||
endpoint: peer,
|
||||
mode: req.bootstrap.mode.clone(),
|
||||
session_key: req.bootstrap.session_key,
|
||||
last_acked: output_seq,
|
||||
replay: ReplayWindow::default(),
|
||||
send_seq: 1,
|
||||
cols: req.cols,
|
||||
rows: req.rows,
|
||||
last_seen: Instant::now(),
|
||||
pending: VecDeque::new(),
|
||||
last_screen: Some(screen),
|
||||
},
|
||||
);
|
||||
(
|
||||
client_id,
|
||||
req.bootstrap.session_key,
|
||||
req.bootstrap.session_key_id,
|
||||
req.bootstrap.session.clone(),
|
||||
req.bootstrap.mode.clone(),
|
||||
output_seq,
|
||||
snapshot,
|
||||
)
|
||||
};
|
||||
let ok = AttachOk {
|
||||
client_id,
|
||||
session: session_name,
|
||||
mode,
|
||||
session_key: key,
|
||||
session_key_id: key_id,
|
||||
initial_seq: output_seq,
|
||||
snapshot,
|
||||
};
|
||||
let body = protocol::to_body(&ok)?;
|
||||
let out = protocol::encode_encrypted(
|
||||
PacketKind::AttachOk,
|
||||
client_id,
|
||||
1,
|
||||
0,
|
||||
&key,
|
||||
SERVER_TO_CLIENT,
|
||||
&body,
|
||||
)?;
|
||||
socket.send_to(&out, peer).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_ticket_attach(
|
||||
state: &Arc<Mutex<ServerState>>,
|
||||
socket: &Arc<UdpSocket>,
|
||||
peer: SocketAddr,
|
||||
body: Vec<u8>,
|
||||
) -> Result<()> {
|
||||
let env: TicketAttachEnvelope = protocol::from_body(&body)?;
|
||||
let (ticket, request_plain) = {
|
||||
let locked = state.lock().expect("server state poisoned");
|
||||
if !locked.config.allow_attach_tickets {
|
||||
return send_reject(socket, peer, "attach tickets disabled").await;
|
||||
}
|
||||
let ticket = open_attach_ticket(&locked.secret, &env.ticket)?;
|
||||
let request_key = crypto::hkdf32(
|
||||
&ticket.psk,
|
||||
&env.client_nonce,
|
||||
b"dosh/ticket-attach-request/v1",
|
||||
)?;
|
||||
let request_plain = crypto::open(
|
||||
&request_key,
|
||||
&env.client_nonce,
|
||||
b"dosh-ticket-attach-request-v1",
|
||||
&env.ciphertext,
|
||||
)?;
|
||||
(ticket, request_plain)
|
||||
};
|
||||
let req: TicketAttachBody = protocol::from_body(&request_plain)?;
|
||||
if req.session != ticket.session || req.mode != ticket.mode {
|
||||
return send_reject(socket, peer, "ticket scope mismatch").await;
|
||||
}
|
||||
|
||||
let session_key = crypto::random_32();
|
||||
let session_key_id = {
|
||||
let digest = crypto::sha256(&session_key);
|
||||
let mut out = [0u8; 16];
|
||||
out.copy_from_slice(&digest[..16]);
|
||||
out
|
||||
};
|
||||
let (client_id, output_seq, snapshot) = {
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
if !locked.sessions.contains_key(&req.session) {
|
||||
if locked.config.create_on_attach {
|
||||
locked.ensure_session(&req.session, req.cols, req.rows)?;
|
||||
} else {
|
||||
return send_reject(socket, peer, "session does not exist").await;
|
||||
}
|
||||
}
|
||||
let session = locked
|
||||
.sessions
|
||||
.get_mut(&req.session)
|
||||
.expect("session exists");
|
||||
let client_id = crypto::random_16();
|
||||
let snapshot = session.parser.screen().state_formatted();
|
||||
let screen = session.parser.screen().clone();
|
||||
let output_seq = session.output_seq;
|
||||
session.clients.insert(
|
||||
client_id,
|
||||
ClientState {
|
||||
endpoint: peer,
|
||||
mode: req.mode.clone(),
|
||||
session_key,
|
||||
last_acked: output_seq,
|
||||
replay: ReplayWindow::default(),
|
||||
send_seq: 1,
|
||||
cols: req.cols,
|
||||
rows: req.rows,
|
||||
last_seen: Instant::now(),
|
||||
pending: VecDeque::new(),
|
||||
last_screen: Some(screen),
|
||||
},
|
||||
);
|
||||
(client_id, output_seq, snapshot)
|
||||
};
|
||||
|
||||
let ok = AttachOk {
|
||||
client_id,
|
||||
session: req.session,
|
||||
mode: req.mode,
|
||||
session_key,
|
||||
session_key_id,
|
||||
initial_seq: output_seq,
|
||||
snapshot,
|
||||
};
|
||||
let ok_plain = protocol::to_body(&ok)?;
|
||||
let server_nonce = crypto::random_12();
|
||||
let mut salt = Vec::with_capacity(24);
|
||||
salt.extend_from_slice(&env.client_nonce);
|
||||
salt.extend_from_slice(&server_nonce);
|
||||
let response_key = crypto::hkdf32(&ticket.psk, &salt, b"dosh/ticket-attach-ok/v1")?;
|
||||
let ciphertext = crypto::seal(
|
||||
&response_key,
|
||||
&server_nonce,
|
||||
b"dosh-ticket-attach-ok-v1",
|
||||
&ok_plain,
|
||||
)?;
|
||||
let envelope = TicketAttachOkEnvelope {
|
||||
server_nonce,
|
||||
ciphertext,
|
||||
};
|
||||
let body = protocol::to_body(&envelope)?;
|
||||
let out = protocol::encode_plain(PacketKind::AttachOk, client_id, 1, 0, &body)?;
|
||||
socket.send_to(&out, peer).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_reject(socket: &UdpSocket, peer: SocketAddr, reason: &str) -> Result<()> {
|
||||
let body = protocol::to_body(&AttachReject {
|
||||
reason: reason.to_string(),
|
||||
})?;
|
||||
let out = protocol::encode_plain(PacketKind::AttachReject, [0u8; 16], 0, 0, &body)?;
|
||||
socket.send_to(&out, peer).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_resume(
|
||||
state: &Arc<Mutex<ServerState>>,
|
||||
socket: &Arc<UdpSocket>,
|
||||
peer: SocketAddr,
|
||||
packet: &protocol::Packet,
|
||||
) -> Result<()> {
|
||||
let (key, session_name) = match find_client_key(state, &packet.header.conn_id) {
|
||||
Ok(found) => found,
|
||||
Err(_) => return send_reject(socket, peer, "unknown client").await,
|
||||
};
|
||||
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
|
||||
let req: ResumeRequest = protocol::from_body(&body)?;
|
||||
let (send_seq, output_seq, snapshot) = {
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
let session = locked
|
||||
.sessions
|
||||
.get_mut(&req.session)
|
||||
.ok_or_else(|| anyhow!("unknown session"))?;
|
||||
let client = session
|
||||
.clients
|
||||
.get_mut(&packet.header.conn_id)
|
||||
.ok_or_else(|| anyhow!("unknown client"))?;
|
||||
if !client.replay.accept(packet.header.seq) {
|
||||
return Ok(());
|
||||
}
|
||||
client.endpoint = peer;
|
||||
client.last_acked = req.last_rendered_seq;
|
||||
client.cols = req.cols;
|
||||
client.rows = req.rows;
|
||||
client.last_seen = Instant::now();
|
||||
client.send_seq += 1;
|
||||
let snapshot = session.parser.screen().state_formatted();
|
||||
client.last_screen = Some(session.parser.screen().clone());
|
||||
(client.send_seq, session.output_seq, snapshot)
|
||||
};
|
||||
let frame = Frame {
|
||||
session: session_name,
|
||||
output_seq,
|
||||
bytes: snapshot,
|
||||
snapshot: true,
|
||||
};
|
||||
let body = protocol::to_body(&frame)?;
|
||||
let out = protocol::encode_encrypted(
|
||||
PacketKind::ResumeOk,
|
||||
packet.header.conn_id,
|
||||
send_seq,
|
||||
packet.header.seq,
|
||||
&key,
|
||||
SERVER_TO_CLIENT,
|
||||
&body,
|
||||
)?;
|
||||
socket.send_to(&out, peer).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_input(
|
||||
state: &Arc<Mutex<ServerState>>,
|
||||
peer: SocketAddr,
|
||||
packet: &protocol::Packet,
|
||||
) -> Result<()> {
|
||||
let (key, session_name) = find_client_key(state, &packet.header.conn_id)?;
|
||||
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
|
||||
let input: Input = protocol::from_body(&body)?;
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
let session = locked
|
||||
.sessions
|
||||
.get_mut(&session_name)
|
||||
.ok_or_else(|| anyhow!("unknown session"))?;
|
||||
let client = session
|
||||
.clients
|
||||
.get_mut(&packet.header.conn_id)
|
||||
.ok_or_else(|| anyhow!("unknown client"))?;
|
||||
if !client.replay.accept(packet.header.seq) {
|
||||
return Ok(());
|
||||
}
|
||||
if client.endpoint != peer {
|
||||
client.endpoint = peer;
|
||||
}
|
||||
client.last_seen = Instant::now();
|
||||
if client.mode == "view-only" {
|
||||
return Ok(());
|
||||
}
|
||||
session.pty.write_all(&input.bytes)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_resize(
|
||||
state: &Arc<Mutex<ServerState>>,
|
||||
peer: SocketAddr,
|
||||
packet: &protocol::Packet,
|
||||
) -> Result<()> {
|
||||
let (key, session_name) = find_client_key(state, &packet.header.conn_id)?;
|
||||
let body = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
|
||||
let resize: Resize = protocol::from_body(&body)?;
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
let session = locked
|
||||
.sessions
|
||||
.get_mut(&session_name)
|
||||
.ok_or_else(|| anyhow!("unknown session"))?;
|
||||
let client = session
|
||||
.clients
|
||||
.get_mut(&packet.header.conn_id)
|
||||
.ok_or_else(|| anyhow!("unknown client"))?;
|
||||
if !client.replay.accept(packet.header.seq) {
|
||||
return Ok(());
|
||||
}
|
||||
if client.mode != "view-only" {
|
||||
client.endpoint = peer;
|
||||
client.cols = resize.cols;
|
||||
client.rows = resize.rows;
|
||||
session.pty.resize(resize.cols, resize.rows)?;
|
||||
session.parser.set_size(resize.rows, resize.cols);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_ping(
|
||||
state: &Arc<Mutex<ServerState>>,
|
||||
socket: &Arc<UdpSocket>,
|
||||
peer: SocketAddr,
|
||||
packet: &protocol::Packet,
|
||||
) -> Result<()> {
|
||||
let (key, _) = find_client_key(state, &packet.header.conn_id)?;
|
||||
let seq = {
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
let mut found = None;
|
||||
for session in locked.sessions.values_mut() {
|
||||
if let Some(client) = session.clients.get_mut(&packet.header.conn_id) {
|
||||
if !client.replay.accept(packet.header.seq) {
|
||||
return Ok(());
|
||||
}
|
||||
client.last_seen = Instant::now();
|
||||
client.send_seq += 1;
|
||||
found = Some(client.send_seq);
|
||||
break;
|
||||
}
|
||||
}
|
||||
found.ok_or_else(|| anyhow!("unknown client"))?
|
||||
};
|
||||
let out = protocol::encode_encrypted(
|
||||
PacketKind::Pong,
|
||||
packet.header.conn_id,
|
||||
seq,
|
||||
packet.header.seq,
|
||||
&key,
|
||||
SERVER_TO_CLIENT,
|
||||
b"",
|
||||
)?;
|
||||
socket.send_to(&out, peer).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_detach(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet) -> Result<()> {
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
for session in locked.sessions.values_mut() {
|
||||
session.clients.remove(&packet.header.conn_id);
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_ack(state: &Arc<Mutex<ServerState>>, packet: &protocol::Packet) -> Result<()> {
|
||||
let (key, _) = find_client_key(state, &packet.header.conn_id)?;
|
||||
let _ = protocol::decrypt_body(packet, &key, CLIENT_TO_SERVER)?;
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
for session in locked.sessions.values_mut() {
|
||||
if let Some(client) = session.clients.get_mut(&packet.header.conn_id) {
|
||||
if !client.replay.accept(packet.header.seq) {
|
||||
return Ok(());
|
||||
}
|
||||
client.last_seen = Instant::now();
|
||||
client.last_acked = packet.header.ack;
|
||||
while client
|
||||
.pending
|
||||
.front()
|
||||
.is_some_and(|pending| pending.output_seq <= packet.header.ack)
|
||||
{
|
||||
client.pending.pop_front();
|
||||
}
|
||||
return Ok(());
|
||||
}
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn broadcast_output(
|
||||
state: &Arc<Mutex<ServerState>>,
|
||||
socket: &Arc<UdpSocket>,
|
||||
output: PtyOutput,
|
||||
) -> Result<()> {
|
||||
let sends = {
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
let scrollback = locked.config.scrollback;
|
||||
let retransmit_window = locked.config.retransmit_window;
|
||||
let session = locked
|
||||
.sessions
|
||||
.get_mut(&output.session)
|
||||
.ok_or_else(|| anyhow!("unknown session"))?;
|
||||
session.parser.process(&output.bytes);
|
||||
session.output_seq += 1;
|
||||
let output_seq = session.output_seq;
|
||||
session.recent.push_back(output.bytes.clone());
|
||||
while session.recent.len() > scrollback {
|
||||
session.recent.pop_front();
|
||||
}
|
||||
let mut sends = Vec::new();
|
||||
for (client_id, client) in session.clients.iter_mut() {
|
||||
client.send_seq += 1;
|
||||
let current_screen = session.parser.screen().clone();
|
||||
let mut snapshot = false;
|
||||
let mut bytes = if client.pending.len() >= retransmit_window {
|
||||
client.pending.clear();
|
||||
snapshot = true;
|
||||
current_screen.state_formatted()
|
||||
} else if let Some(prev) = &client.last_screen {
|
||||
current_screen.state_diff(prev)
|
||||
} else {
|
||||
snapshot = true;
|
||||
current_screen.state_formatted()
|
||||
};
|
||||
if bytes.is_empty() {
|
||||
bytes = output.bytes.clone();
|
||||
}
|
||||
let frame = Frame {
|
||||
session: output.session.clone(),
|
||||
output_seq,
|
||||
bytes,
|
||||
snapshot,
|
||||
};
|
||||
let body = protocol::to_body(&frame)?;
|
||||
let packet = protocol::encode_encrypted(
|
||||
PacketKind::Frame,
|
||||
*client_id,
|
||||
client.send_seq,
|
||||
client.last_acked,
|
||||
&client.session_key,
|
||||
SERVER_TO_CLIENT,
|
||||
&body,
|
||||
)?;
|
||||
while client.pending.len() >= retransmit_window {
|
||||
client.pending.pop_front();
|
||||
}
|
||||
client.last_screen = Some(current_screen);
|
||||
client.pending.push_back(PendingFrame {
|
||||
output_seq,
|
||||
packet: packet.clone(),
|
||||
last_sent: Instant::now(),
|
||||
attempts: 0,
|
||||
});
|
||||
sends.push((client.endpoint, packet));
|
||||
}
|
||||
sends
|
||||
};
|
||||
for (endpoint, packet) in sends {
|
||||
socket.send_to(&packet, endpoint).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn retransmit_pending(
|
||||
state: &Arc<Mutex<ServerState>>,
|
||||
socket: &Arc<UdpSocket>,
|
||||
) -> Result<()> {
|
||||
let sends = {
|
||||
let mut locked = state.lock().expect("server state poisoned");
|
||||
let now = Instant::now();
|
||||
let mut sends = Vec::new();
|
||||
for session in locked.sessions.values_mut() {
|
||||
for client in session.clients.values_mut() {
|
||||
for pending in client.pending.iter_mut() {
|
||||
if pending.output_seq <= client.last_acked {
|
||||
continue;
|
||||
}
|
||||
if now.duration_since(pending.last_sent) >= Duration::from_millis(200)
|
||||
&& pending.attempts < 8
|
||||
{
|
||||
pending.last_sent = now;
|
||||
pending.attempts += 1;
|
||||
sends.push((client.endpoint, pending.packet.clone()));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
sends
|
||||
};
|
||||
for (endpoint, packet) in sends {
|
||||
socket.send_to(&packet, endpoint).await?;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn find_client_key(
|
||||
state: &Arc<Mutex<ServerState>>,
|
||||
client_id: &[u8; 16],
|
||||
) -> Result<([u8; 32], String)> {
|
||||
let locked = state.lock().expect("server state poisoned");
|
||||
for (name, session) in &locked.sessions {
|
||||
if let Some(client) = session.clients.get(client_id) {
|
||||
return Ok((client.session_key, name.clone()));
|
||||
}
|
||||
}
|
||||
Err(anyhow!("unknown client"))
|
||||
}
|
||||
|
||||
fn parse_nonce(raw: &str) -> Result<[u8; 12]> {
|
||||
let bytes = base64::Engine::decode(&base64::engine::general_purpose::URL_SAFE_NO_PAD, raw)
|
||||
.context("decode nonce")?;
|
||||
anyhow::ensure!(bytes.len() == 12, "nonce must decode to 12 bytes");
|
||||
let mut out = [0u8; 12];
|
||||
out.copy_from_slice(&bytes);
|
||||
Ok(out)
|
||||
}
|
||||
|
||||
fn parse_size(raw: &str) -> Result<(u16, u16)> {
|
||||
let (cols, rows) = raw.split_once('x').context("size must be COLSxROWS")?;
|
||||
Ok((cols.parse()?, rows.parse()?))
|
||||
}
|
||||
Reference in New Issue
Block a user