Initial Dosh implementation
This commit is contained in:
Codex
@@ -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