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

Keep Dosh alive across stale UDP reconnects
ci / test (push) Has been cancelled
Details
ci / remote-bench (push) Has been cancelled
Details

Browse Source
This commit is contained in:
DuProcess
2026-06-13 10:20:03 -04:00
parent 93374162ff
commit 80166eae15
5 changed files with 161 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/bin/dosh-client.rs
+152 -64
View File
@@ -164,63 +164,64 @@ async fn main() -> Result<()> {
cached.udp_host = target_udp_host.clone();
cached.udp_port = dosh_port;
log_timing(args.verbose, "credential_lookup_end", started.elapsed());
match try_resume(&socket, &cached, cols, rows).await {
Ok((frame, cred)) => {
log_timing(args.verbose, "udp_resume_ready", started.elapsed());
if !args.no_cache {
if config.cache_attach_tickets && !args.no_cache {
match try_ticket_attach(&socket, &cached, cols, rows).await {
Ok((frame, cred)) => {
log_timing(args.verbose, "udp_ticket_attach_ready", started.elapsed());
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),
predict,
startup_command,
config.reconnect_timeout_secs,
)
.await;
}
Err(err) => {
log_debug(
args.verbose,
2,
&format!("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)) => {
log_timing(args.verbose, "udp_ticket_attach_ready", started.elapsed());
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),
predict,
startup_command,
config.reconnect_timeout_secs,
)
.await;
}
Err(err) => {
log_debug(
args.verbose,
2,
&format!(
"dosh ticket attach failed, falling back to SSH bootstrap: {err:#}"
),
);
}
if args.attach_only {
render_frame(&frame)?;
detach_once(&socket, &cred, 2).await?;
return Ok(());
}
return run_terminal(
socket,
cred,
Some(frame),
predict,
startup_command,
config.reconnect_timeout_secs,
)
.await;
}
Err(err) => {
log_debug(
args.verbose,
2,
&format!(
"dosh ticket attach failed, falling back to SSH bootstrap: {err:#}"
),
);
}
}
} else {
match try_resume_fresh_process(&socket, &cached, cols, rows).await {
Ok((frame, cred)) => {
log_timing(args.verbose, "udp_resume_ready", started.elapsed());
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),
predict,
startup_command,
config.reconnect_timeout_secs,
)
.await;
}
Err(err) => {
log_debug(
args.verbose,
2,
&format!("dosh resume failed, falling back to SSH bootstrap: {err:#}"),
);
}
}
}
@@ -745,7 +746,7 @@ async fn try_ticket_attach(
Ok((frame, next))
}
async fn try_resume(
async fn try_resume_fresh_process(
socket: &UdpSocket,
cached: &CachedCredential,
cols: u16,
@@ -787,6 +788,53 @@ async fn try_resume(
Ok((frame, next))
}
async fn try_live_resume(
socket: &UdpSocket,
cached: &CachedCredential,
send_seq: &mut u64,
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,
*send_seq,
cached.last_rendered_seq,
&cached.session_key,
CLIENT_TO_SERVER,
&body,
)?;
*send_seq += 1;
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.conn_id != cached.client_id {
return Err(anyhow!("resume received stale client packet"));
}
if packet.header.kind != PacketKind::ResumeOk {
if packet.header.kind == PacketKind::AttachReject {
let reject: AttachReject = protocol::from_body(&packet.body)?;
return Err(anyhow!("resume rejected: {}", reject.reason));
}
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,
@@ -865,12 +913,35 @@ async fn run_terminal(
}
recv = socket.recv_from(&mut recv_buf) => {
let (n, _) = recv?;
last_packet_at = Instant::now();
let packet = protocol::decode(&recv_buf[..n])?;
let Ok(packet) = protocol::decode(&recv_buf[..n]) else {
continue;
};
if packet.header.conn_id != [0u8; 16] && packet.header.conn_id != cred.client_id {
continue;
}
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)?;
let Ok(plain) = protocol::decrypt_body(&packet, &cred.session_key, SERVER_TO_CLIENT) else {
if let Some(frame) = reconnect(
&socket,
&mut cred,
&mut send_seq,
last_size,
&mut frame_buffer,
&mut predictor,
)
.await?
{
render_frame(&frame)?;
predictor.observe_output(&frame.bytes);
last_packet_at = Instant::now();
}
continue;
};
let Ok(frame) = protocol::from_body::<Frame>(&plain) else {
continue;
};
last_packet_at = Instant::now();
let frames = frame_buffer.accept(frame, &mut cred.last_rendered_seq);
for frame in frames {
predictor.clear_pending()?;
@@ -883,11 +954,13 @@ async fn run_terminal(
}
send_ack(&socket, addr, &cred, &mut send_seq).await?;
}
PacketKind::Pong => {}
PacketKind::Pong => {
last_packet_at = Instant::now();
}
PacketKind::AttachReject => {
let reject: AttachReject = protocol::from_body(&packet.body)?;
if reject.reason == "unknown client" {
if let Some(frame) = reconnect_with_ticket(
if let Some(frame) = reconnect(
&socket,
&mut cred,
&mut send_seq,
@@ -911,7 +984,7 @@ async fn run_terminal(
_ = status_tick.tick() => {
let stale = last_packet_at.elapsed();
if stale >= Duration::from_secs(reconnect_timeout_secs.max(1)) {
if let Some(frame) = reconnect_with_ticket(
if let Some(frame) = reconnect(
&socket,
&mut cred,
&mut send_seq,
@@ -945,7 +1018,7 @@ async fn run_terminal(
Ok(())
}
async fn reconnect_with_ticket(
async fn reconnect(
socket: &UdpSocket,
cred: &mut CachedCredential,
send_seq: &mut u64,
@@ -953,6 +1026,21 @@ async fn reconnect_with_ticket(
frame_buffer: &mut FrameBuffer,
predictor: &mut Predictor,
) -> Result<Option<Frame>> {
if let Ok((frame, next)) = try_live_resume(socket, cred, send_seq, size.0, size.1).await {
*cred = next;
frame_buffer.clear();
predictor.reset();
cred.last_rendered_seq = frame.output_seq;
send_ack(
socket,
resolve_addr(&cred.udp_host, cred.udp_port)?,
cred,
send_seq,
)
.await?;
return Ok(Some(frame));
}
let Ok((frame, next)) = try_ticket_attach(socket, cred, size.0, size.1).await else {
return Ok(None);
};