Files
dosh/src/bin/dosh-server.rs
T
DuProcess 6a0f12c59e
ci / test (push) Has been cancelled
ci / remote-bench (push) Has been cancelled
Add native remote port forwarding
2026-06-13 15:41:48 -04:00

1619 lines
54 KiB
Rust

use anyhow::{Context, Result, anyhow};
use clap::{Parser, Subcommand};
use dosh::auth::{
build_attach_ticket, build_bootstrap, encode_bootstrap, load_or_create_server_secret, now_secs,
open_attach_ticket, verify_bootstrap,
};
use dosh::config::{ServerConfig, load_server_config};
use dosh::crypto;
use dosh::native::{
ForwardingKind, ForwardingRequest, NativeAuthOk, NativeServerHello, derive_native_session_key,
generate_native_ephemeral, host_public_key, host_public_key_line, load_or_create_host_key,
sign_server_hello, verify_native_user_auth_from_config,
};
use dosh::protocol::{
self, AttachOk, AttachReject, BootstrapAttachRequest, CLIENT_TO_SERVER, Frame, Input,
NativeAuthOkBody, NativeClientHelloBody, NativeServerHelloBody, NativeUserAuthBody, PacketKind,
ReplayWindow, Resize, ResumeRequest, SERVER_TO_CLIENT, StreamClose, StreamData, StreamOpen,
StreamOpenOk, StreamOpenReject, 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::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::{TcpListener, TcpStream, 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: Option<String>,
#[arg(long)]
host_key: bool,
#[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,
host_key,
session,
mode,
size,
client_version: _,
udp_host,
} => {
anyhow::ensure!(protocol == 1, "unsupported protocol {protocol}");
let config = load_server_config(None)?;
if host_key {
let signing_key = load_or_create_host_key(&config)?;
println!("{}", host_public_key_line(&host_public_key(&signing_key)));
return Ok(());
}
let secret = load_or_create_server_secret(&config)?;
let nonce = parse_nonce(nonce.as_deref().context("--nonce is required")?)?;
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 cleanup_state = Arc::clone(&state);
tokio::spawn(async move {
let mut interval = tokio::time::interval(Duration::from_secs(5));
loop {
interval.tick().await;
cleanup_disconnected_clients(&cleanup_state);
}
});
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>,
pending_native: HashMap<[u8; 16], PendingNativeAuth>,
next_server_stream_id: u64,
}
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>,
allowed_forwardings: Vec<ForwardingRequest>,
stream_writers: HashMap<u64, mpsc::UnboundedSender<Vec<u8>>>,
}
#[derive(Clone)]
struct PendingFrame {
output_seq: u64,
packet: Vec<u8>,
last_sent: Instant,
attempts: u8,
}
struct PendingNativeAuth {
client: dosh::native::NativeClientHello,
server: NativeServerHello,
session_key: [u8; 32],
peer: SocketAddr,
created_at: Instant,
}
impl ServerState {
fn new(
config: ServerConfig,
secret: [u8; 32],
pty_tx: mpsc::UnboundedSender<PtyOutput>,
) -> Self {
Self {
config,
secret,
pty_tx,
sessions: HashMap::new(),
pending_native: HashMap::new(),
next_server_stream_id: 1u64 << 63,
}
}
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::NativeClientHello => {
handle_native_client_hello(state, socket, peer, packet.body).await
}
PacketKind::NativeUserAuth => handle_native_user_auth(state, socket, peer, &packet).await,
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
| PacketKind::Resize
| PacketKind::Ping
| PacketKind::Ack
| PacketKind::StreamOpen
| PacketKind::StreamOpenOk
| PacketKind::StreamOpenReject
| PacketKind::StreamData
| PacketKind::StreamClose
| PacketKind::StreamEof
| PacketKind::StreamWindowAdjust
if find_client_key(state, &packet.header.conn_id).is_err() =>
{
send_reject_to_client(socket, peer, packet.header.conn_id, "unknown client").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,
PacketKind::StreamOpen => handle_stream_open(state, socket, peer, &packet).await,
PacketKind::StreamOpenOk => handle_stream_open_ok(state, peer, &packet).await,
PacketKind::StreamOpenReject => handle_stream_open_reject(state, peer, &packet).await,
PacketKind::StreamData => handle_stream_data(state, peer, &packet).await,
PacketKind::StreamClose => handle_stream_close(state, &packet).await,
_ => Ok(()),
}
}
async fn handle_native_client_hello(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
body: Vec<u8>,
) -> Result<()> {
let req: NativeClientHelloBody = protocol::from_body(&body)?;
let built: Result<([u8; 16], NativeServerHello)> = {
let mut locked = state.lock().expect("server state poisoned");
if !locked.config.native_auth {
Err(anyhow!("native auth disabled"))
} else if !req
.hello
.supported_aead
.iter()
.any(|algorithm| algorithm == "chacha20poly1305")
{
Err(anyhow!("native auth requires chacha20poly1305"))
} else if !req
.hello
.supported_user_key_algorithms
.iter()
.any(|algorithm| algorithm == "ssh-ed25519")
{
Err(anyhow!("native auth requires ssh-ed25519"))
} else {
let current_user = std::env::var("USER").unwrap_or_else(|_| "unknown".to_string());
if req.hello.requested_user != current_user {
Err(anyhow!("native auth user mismatch"))
} else {
let host_signing = load_or_create_host_key(&locked.config)?;
let (server_secret, server_public) = generate_native_ephemeral();
let mut hello = NativeServerHello {
protocol_version: dosh::native::NATIVE_PROTOCOL_VERSION,
server_random: crypto::random_32(),
server_ephemeral_public: server_public,
host_key: host_public_key(&host_signing),
chosen_aead: "chacha20poly1305".to_string(),
server_key_epoch: 1,
auth_challenge: crypto::random_32(),
rate_limit_remaining: Some(locked.config.native_auth_rate_limit_per_minute),
host_signature: Vec::new(),
};
sign_server_hello(&host_signing, &req.hello, &mut hello)?;
let session_key = derive_native_session_key(
&server_secret,
req.hello.client_ephemeral_public,
&req.hello,
&hello,
)?;
let mut pending_id = [0u8; 16];
pending_id.copy_from_slice(&hello.auth_challenge[..16]);
locked.pending_native.insert(
pending_id,
PendingNativeAuth {
client: req.hello,
server: hello.clone(),
session_key,
peer,
created_at: Instant::now(),
},
);
Ok((pending_id, hello))
}
}
};
let (pending_id, hello) = match built {
Ok(built) => built,
Err(err) => return send_reject(socket, peer, &err.to_string()).await,
};
let body = protocol::to_body(&NativeServerHelloBody { hello })?;
let out = protocol::encode_plain(PacketKind::NativeServerHello, pending_id, 1, 0, &body)?;
socket.send_to(&out, peer).await?;
Ok(())
}
async fn handle_native_user_auth(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
peer: SocketAddr,
packet: &protocol::Packet,
) -> Result<()> {
let pending = {
let mut locked = state.lock().expect("server state poisoned");
locked.pending_native.remove(&packet.header.conn_id)
};
let Some(pending) = pending else {
return send_reject_to_client(socket, peer, packet.header.conn_id, "unknown native auth")
.await;
};
if pending.peer != peer {
return send_reject_to_client(
socket,
peer,
packet.header.conn_id,
"native auth peer changed",
)
.await;
}
if pending.created_at.elapsed() > Duration::from_secs(30) {
return send_reject_to_client(socket, peer, packet.header.conn_id, "native auth expired")
.await;
}
let body = protocol::decrypt_body(packet, &pending.session_key, CLIENT_TO_SERVER)?;
let req: NativeUserAuthBody = protocol::from_body(&body)?;
let attached: Result<(
[u8; 16],
[u8; 16],
Vec<u8>,
[u8; 32],
String,
String,
u64,
Vec<u8>,
Vec<ForwardingRequest>,
)> = {
let mut locked = state.lock().expect("server state poisoned");
verify_native_user_auth_from_config(
&locked.config,
&pending.client,
&pending.server,
&req.auth,
)
.context("verify native user auth")
.and_then(|_| {
let session_name = pending.client.requested_session.clone();
let mode = pending.client.requested_mode.clone();
let cols = pending.client.terminal_size.0;
let rows = pending.client.terminal_size.1;
if !locked.sessions.contains_key(&session_name) {
if locked.config.create_on_attach {
locked.ensure_session(&session_name, cols, rows)?;
} else {
return Err(anyhow!("session does not exist"));
}
}
let session_key = pending.session_key;
let key_id = {
let digest = crypto::sha256(&session_key);
let mut out = [0u8; 16];
out.copy_from_slice(&digest[..16]);
out
};
let attach_ticket_psk = crypto::random_32();
let issued_at = now_secs()?;
let attach_ticket = build_attach_ticket(
&locked.secret,
crypto::sha256(&locked.secret),
1,
pending.client.requested_user.clone(),
session_name.clone(),
mode.clone(),
issued_at,
issued_at + locked.config.attach_ticket_ttl_secs,
&attach_ticket_psk,
)?;
let session = locked
.sessions
.get_mut(&session_name)
.expect("session exists");
if mode != "view-only" {
session.pty.resize(cols, rows)?;
session.parser.set_size(rows, cols);
}
let client_id = crypto::random_16();
let snapshot = screen_snapshot(session.parser.screen());
let screen = session.parser.screen().clone();
let output_seq = session.output_seq;
session.clients.insert(
client_id,
ClientState {
endpoint: peer,
mode: mode.clone(),
session_key,
last_acked: output_seq,
replay: ReplayWindow::default(),
send_seq: 1,
cols,
rows,
last_seen: Instant::now(),
pending: VecDeque::new(),
last_screen: Some(screen),
allowed_forwardings: req.auth.requested_forwardings.clone(),
stream_writers: HashMap::new(),
},
);
Ok((
client_id,
key_id,
attach_ticket,
attach_ticket_psk,
session_name,
mode,
output_seq,
snapshot,
req.auth.requested_forwardings.clone(),
))
})
};
let (
client_id,
key_id,
attach_ticket,
attach_ticket_psk,
session_name,
mode,
output_seq,
snapshot,
requested_forwardings,
) = match attached {
Ok(attached) => attached,
Err(err) => {
return send_reject_to_client(socket, peer, packet.header.conn_id, &err.to_string())
.await;
}
};
if let Err(err) = start_remote_forwards(
state,
socket,
client_id,
requested_forwardings.clone(),
pending.client.requested_session.clone(),
)
.await
{
remove_client(state, client_id);
return send_reject_to_client(socket, peer, packet.header.conn_id, &err.to_string()).await;
}
let ok = NativeAuthOk {
client_id,
session: session_name,
mode,
session_key: pending.session_key,
session_key_id: key_id,
attach_ticket,
attach_ticket_psk,
initial_seq: output_seq,
snapshot,
policy_flags: Vec::new(),
};
let body = protocol::to_body(&NativeAuthOkBody { ok })?;
let out = protocol::encode_encrypted(
PacketKind::NativeAuthOk,
client_id,
1,
packet.header.seq,
&pending.session_key,
SERVER_TO_CLIENT,
&body,
)?;
socket.send_to(&out, peer).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");
if req.bootstrap.mode != "view-only" {
session.pty.resize(req.cols, req.rows)?;
session.parser.set_size(req.rows, req.cols);
}
let client_id = crypto::random_16();
let snapshot = screen_snapshot(session.parser.screen());
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),
allowed_forwardings: Vec::new(),
stream_writers: HashMap::new(),
},
);
(
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");
if req.mode != "view-only" {
session.pty.resize(req.cols, req.rows)?;
session.parser.set_size(req.rows, req.cols);
}
let client_id = crypto::random_16();
let snapshot = screen_snapshot(session.parser.screen());
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),
allowed_forwardings: Vec::new(),
stream_writers: HashMap::new(),
},
);
(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<()> {
send_reject_to_client(socket, peer, [0u8; 16], reason).await
}
async fn send_reject_to_client(
socket: &UdpSocket,
peer: SocketAddr,
client_id: [u8; 16],
reason: &str,
) -> Result<()> {
let body = protocol::to_body(&AttachReject {
reason: reason.to_string(),
})?;
let out = protocol::encode_plain(PacketKind::AttachReject, client_id, 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;
if client.mode != "view-only" {
session.pty.resize(req.cols, req.rows)?;
session.parser.set_size(req.rows, req.cols);
}
let snapshot = screen_snapshot(session.parser.screen());
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,
closed: false,
};
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(())
}
fn remove_client(state: &Arc<Mutex<ServerState>>, client_id: [u8; 16]) {
let mut locked = state.lock().expect("server state poisoned");
for session in locked.sessions.values_mut() {
session.clients.remove(&client_id);
}
}
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 handle_stream_open(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
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 open: StreamOpen = protocol::from_body(&body)?;
let allowed = {
let mut locked = state.lock().expect("server state poisoned");
let config = locked.config.clone();
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(());
}
client.endpoint = peer;
client.last_seen = Instant::now();
stream_open_allowed(&config, client, &open)
};
if let Err(err) = allowed {
return send_stream_open_reject(
state,
socket,
packet.header.conn_id,
open.stream_id,
err.to_string(),
)
.await;
}
let target = format!("{}:{}", open.target_host, open.target_port);
let stream = match TcpStream::connect((open.target_host.as_str(), open.target_port)).await {
Ok(stream) => stream,
Err(err) => {
return send_stream_open_reject(
state,
socket,
packet.header.conn_id,
open.stream_id,
format!("connect {target}: {err}"),
)
.await;
}
};
let (mut reader, mut writer) = stream.into_split();
let (writer_tx, mut writer_rx) = mpsc::unbounded_channel::<Vec<u8>>();
{
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"))?;
client.stream_writers.insert(open.stream_id, writer_tx);
}
tokio::spawn(async move {
while let Some(bytes) = writer_rx.recv().await {
if writer.write_all(&bytes).await.is_err() {
break;
}
}
});
send_stream_open_ok(state, socket, packet.header.conn_id, open.stream_id).await?;
let state = Arc::clone(state);
let socket = Arc::clone(socket);
let client_id = packet.header.conn_id;
let stream_id = open.stream_id;
tokio::spawn(async move {
let mut buf = [0u8; 16 * 1024];
loop {
match reader.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if let Err(err) = send_stream_data_to_client(
&state,
&socket,
client_id,
stream_id,
buf[..n].to_vec(),
)
.await
{
eprintln!("stream send error: {err:#}");
break;
}
}
Err(_) => break,
}
}
let _ = send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
});
Ok(())
}
async fn handle_stream_data(
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 data: StreamData = protocol::from_body(&body)?;
let writer = {
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(());
}
client.endpoint = peer;
client.last_seen = Instant::now();
client.stream_writers.get(&data.stream_id).cloned()
};
if let Some(writer) = writer {
let _ = writer.send(data.bytes);
}
Ok(())
}
async fn handle_stream_open_ok(
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 _ok: StreamOpenOk = 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(());
}
client.endpoint = peer;
client.last_seen = Instant::now();
Ok(())
}
async fn handle_stream_open_reject(
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 reject: StreamOpenReject = 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(());
}
client.endpoint = peer;
client.last_seen = Instant::now();
client.stream_writers.remove(&reject.stream_id);
Ok(())
}
async fn handle_stream_close(
state: &Arc<Mutex<ServerState>>,
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 close: StreamClose = protocol::from_body(&body)?;
let mut locked = state.lock().expect("server state poisoned");
let Some(session) = locked.sessions.get_mut(&session_name) else {
return Ok(());
};
let Some(client) = session.clients.get_mut(&packet.header.conn_id) else {
return Ok(());
};
if !client.replay.accept(packet.header.seq) {
return Ok(());
}
client.last_seen = Instant::now();
client.stream_writers.remove(&close.stream_id);
Ok(())
}
async fn start_remote_forwards(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
requested_forwardings: Vec<ForwardingRequest>,
session_name: String,
) -> Result<()> {
let remote_forwards = requested_forwardings
.into_iter()
.filter(|forward| forward.kind == ForwardingKind::Remote)
.collect::<Vec<_>>();
if remote_forwards.is_empty() {
return Ok(());
}
let config = {
let locked = state.lock().expect("server state poisoned");
locked.config.clone()
};
anyhow::ensure!(
config.allow_remote_forwarding,
"remote forwarding disabled on server"
);
for forward in remote_forwards {
let bind_host = forward
.bind_host
.clone()
.unwrap_or_else(|| "127.0.0.1".to_string());
anyhow::ensure!(
is_loopback_bind(&bind_host),
"remote forwarding may only bind loopback in this server build"
);
let listen_port = forward.listen_port;
let target_host = forward
.target_host
.clone()
.ok_or_else(|| anyhow!("remote forward target host is required"))?;
let target_port = forward
.target_port
.ok_or_else(|| anyhow!("remote forward target port is required"))?;
let bind = format!("{bind_host}:{listen_port}");
let listener = TcpListener::bind(&bind)
.await
.with_context(|| format!("bind remote forward {bind}"))?;
let state = Arc::clone(state);
let socket = Arc::clone(socket);
let session_name = session_name.clone();
tokio::spawn(async move {
loop {
let Ok((stream, _)) = listener.accept().await else {
break;
};
let stream_id = allocate_server_stream_id(&state);
let (mut reader, mut writer) = stream.into_split();
let (writer_tx, mut writer_rx) = mpsc::unbounded_channel::<Vec<u8>>();
{
let mut locked = state.lock().expect("server state poisoned");
if let Some(session) = locked.sessions.get_mut(&session_name)
&& let Some(client) = session.clients.get_mut(&client_id)
{
client.stream_writers.insert(stream_id, writer_tx);
} else {
break;
}
}
tokio::spawn(async move {
while let Some(bytes) = writer_rx.recv().await {
if writer.write_all(&bytes).await.is_err() {
break;
}
}
});
if let Err(err) = send_stream_open_to_client(
&state,
&socket,
client_id,
stream_id,
target_host.clone(),
target_port,
)
.await
{
eprintln!("remote forward open error: {err:#}");
break;
}
let state = Arc::clone(&state);
let socket = Arc::clone(&socket);
tokio::spawn(async move {
let mut buf = [0u8; 16 * 1024];
loop {
match reader.read(&mut buf).await {
Ok(0) => break,
Ok(n) => {
if let Err(err) = send_stream_data_to_client(
&state,
&socket,
client_id,
stream_id,
buf[..n].to_vec(),
)
.await
{
eprintln!("remote forward data error: {err:#}");
break;
}
}
Err(_) => break,
}
}
let _ =
send_stream_close_to_client(&state, &socket, client_id, stream_id).await;
});
}
});
}
Ok(())
}
fn allocate_server_stream_id(state: &Arc<Mutex<ServerState>>) -> u64 {
let mut locked = state.lock().expect("server state poisoned");
let stream_id = locked.next_server_stream_id;
locked.next_server_stream_id = locked.next_server_stream_id.wrapping_add(1).max(1u64 << 63);
stream_id
}
fn is_loopback_bind(host: &str) -> bool {
matches!(host, "127.0.0.1" | "localhost" | "::1")
}
fn stream_open_allowed(
config: &ServerConfig,
client: &ClientState,
open: &StreamOpen,
) -> Result<()> {
anyhow::ensure!(config.allow_tcp_forwarding, "TCP forwarding disabled");
let allowed = client.allowed_forwardings.iter().any(|forward| {
forward.kind == ForwardingKind::Local
&& forward.target_host.as_deref() == Some(open.target_host.as_str())
&& forward.target_port == Some(open.target_port)
});
anyhow::ensure!(
allowed,
"stream target {}:{} was not requested during auth",
open.target_host,
open.target_port
);
Ok(())
}
async fn send_stream_open_ok(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
) -> Result<()> {
let body = protocol::to_body(&StreamOpenOk { stream_id })?;
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenOk, body).await
}
async fn send_stream_open_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
target_host: String,
target_port: u16,
) -> Result<()> {
let body = protocol::to_body(&StreamOpen {
stream_id,
target_host,
target_port,
})?;
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpen, body).await
}
async fn send_stream_open_reject(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
reason: String,
) -> Result<()> {
let body = protocol::to_body(&StreamOpenReject { stream_id, reason })?;
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamOpenReject, body).await
}
async fn send_stream_data_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
bytes: Vec<u8>,
) -> Result<()> {
let body = protocol::to_body(&StreamData { stream_id, bytes })?;
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamData, body).await
}
async fn send_stream_close_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
stream_id: u64,
) -> Result<()> {
{
let mut locked = state.lock().expect("server state poisoned");
for session in locked.sessions.values_mut() {
if let Some(client) = session.clients.get_mut(&client_id) {
client.stream_writers.remove(&stream_id);
break;
}
}
}
let body = protocol::to_body(&StreamClose { stream_id })?;
send_stream_packet_to_client(state, socket, client_id, PacketKind::StreamClose, body).await
}
async fn send_stream_packet_to_client(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
client_id: [u8; 16],
kind: PacketKind,
body: Vec<u8>,
) -> Result<()> {
let send = {
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(&client_id) {
client.send_seq += 1;
let packet = protocol::encode_encrypted(
kind,
client_id,
client.send_seq,
client.last_acked,
&client.session_key,
SERVER_TO_CLIENT,
&body,
)?;
found = Some((client.endpoint, packet));
break;
}
}
found
};
if let Some((endpoint, packet)) = send {
socket.send_to(&packet, endpoint).await?;
}
Ok(())
}
async fn broadcast_output(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
output: PtyOutput,
) -> Result<()> {
if output.exited {
return broadcast_session_exit(state, socket, output.session).await;
}
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 (bytes, snapshot) = if client.pending.len() >= retransmit_window {
client.pending.clear();
(screen_snapshot(&current_screen), true)
} else {
(output.bytes.clone(), false)
};
let frame = Frame {
session: output.session.clone(),
output_seq,
bytes,
snapshot,
closed: false,
};
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 broadcast_session_exit(
state: &Arc<Mutex<ServerState>>,
socket: &Arc<UdpSocket>,
session_name: String,
) -> Result<()> {
let sends = {
let mut locked = state.lock().expect("server state poisoned");
let Some(mut session) = locked.sessions.remove(&session_name) else {
return Ok(());
};
session.output_seq += 1;
let output_seq = session.output_seq;
let mut sends = Vec::new();
for (client_id, mut client) in session.clients.drain() {
client.send_seq += 1;
let frame = Frame {
session: session_name.clone(),
output_seq,
bytes: b"\r\n[dosh session exited]\r\n".to_vec(),
snapshot: false,
closed: true,
};
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,
)?;
sends.push((client.endpoint, packet));
}
sends
};
for (endpoint, packet) in sends {
socket.send_to(&packet, endpoint).await?;
}
Ok(())
}
fn screen_snapshot(screen: &vt100::Screen) -> Vec<u8> {
let mut bytes = Vec::new();
if screen.alternate_screen() {
bytes.extend_from_slice(b"\x1b[?1049h");
} else {
bytes.extend_from_slice(b"\x1b[?1049l");
}
bytes.extend(screen.state_formatted());
bytes
}
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 cleanup_disconnected_clients(state: &Arc<Mutex<ServerState>>) {
let mut locked = state.lock().expect("server state poisoned");
let timeout = Duration::from_secs(locked.config.client_timeout_secs.max(1));
let now = Instant::now();
for session in locked.sessions.values_mut() {
session
.clients
.retain(|_, client| now.duration_since(client.last_seen) <= timeout);
}
}
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()?))
}