Ignore stale packets during attach resume
ci / test (push) Has been cancelled
ci / remote-bench (push) Has been cancelled

This commit is contained in:
DuProcess
2026-06-13 22:21:20 -04:00
parent 76e2e980f1
commit 38e2fa091a
+103 -38
View File
@@ -1662,18 +1662,24 @@ async fn bootstrap_attach(
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 {
if packet.header.kind == PacketKind::AttachReject {
let expected_key_id = protocol::session_key_id(&bootstrap.session_key);
let ok = recv_response_until(socket, Duration::from_secs(5), |packet| {
if packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16] {
let reject: AttachReject = protocol::from_body(&packet.body)?;
return Err(anyhow!("attach rejected: {}", reject.reason));
}
return Err(anyhow!("attach rejected or unexpected response"));
if packet.header.kind != PacketKind::AttachOk
|| packet.header.session_key_id != expected_key_id
{
return Ok(None);
}
let plain = protocol::decrypt_body(&packet, &bootstrap.session_key, SERVER_TO_CLIENT)?;
let ok: AttachOk = protocol::from_body(&plain)?;
let Ok(plain) = protocol::decrypt_body(&packet, &bootstrap.session_key, SERVER_TO_CLIENT)
else {
return Ok(None);
};
Ok(protocol::from_body::<AttachOk>(&plain).ok())
})
.await?;
let cred = CachedCredential {
server: server_name.to_string(),
session: ok.session.clone(),
@@ -1729,18 +1735,17 @@ async fn try_ticket_attach(
)?;
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 {
if packet.header.kind == PacketKind::AttachReject {
let ok = recv_response_until(socket, Duration::from_millis(700), |packet| {
if packet.header.kind == PacketKind::AttachReject && packet.header.conn_id == [0u8; 16] {
let reject: AttachReject = protocol::from_body(&packet.body)?;
return Err(anyhow!("ticket attach rejected: {}", reject.reason));
}
return Err(anyhow!("ticket attach rejected"));
if packet.header.kind != PacketKind::AttachOk {
return Ok(None);
}
let envelope: TicketAttachOkEnvelope = protocol::from_body(&packet.body)?;
let Ok(envelope) = protocol::from_body::<TicketAttachOkEnvelope>(&packet.body) else {
return Ok(None);
};
let mut salt = Vec::with_capacity(24);
salt.extend_from_slice(&client_nonce);
salt.extend_from_slice(&envelope.server_nonce);
@@ -1749,13 +1754,17 @@ async fn try_ticket_attach(
&salt,
b"dosh/ticket-attach-ok/v1",
)?;
let plain = crypto::open(
let Ok(plain) = crypto::open(
&response_key,
&envelope.server_nonce,
b"dosh-ticket-attach-ok-v1",
&envelope.ciphertext,
)?;
let ok: AttachOk = protocol::from_body(&plain)?;
) else {
return Ok(None);
};
Ok(protocol::from_body::<AttachOk>(&plain).ok())
})
.await?;
let frame = Frame {
session: ok.session.clone(),
output_seq: ok.initial_seq,
@@ -1795,19 +1804,24 @@ async fn try_resume_fresh_process(
&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 {
let frame = recv_response_until(socket, Duration::from_millis(700), |packet| {
if packet.header.conn_id != cached.client_id {
return Ok(None);
}
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"));
if packet.header.kind != PacketKind::ResumeOk {
return Ok(None);
}
let plain = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT)?;
let frame: Frame = protocol::from_body(&plain)?;
let Ok(plain) = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT)
else {
return Ok(None);
};
Ok(protocol::from_body::<Frame>(&plain).ok())
})
.await?;
let mut next = cached.clone();
next.last_rendered_seq = frame.output_seq;
Ok((frame, next))
@@ -1839,27 +1853,53 @@ async fn try_live_resume(
)?;
*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])?;
let frame = recv_response_until(socket, Duration::from_millis(700), |packet| {
if packet.header.conn_id != cached.client_id {
return Err(anyhow!("resume received stale client packet"));
return Ok(None);
}
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"));
if packet.header.kind != PacketKind::ResumeOk {
return Ok(None);
}
let plain = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT)?;
let frame: Frame = protocol::from_body(&plain)?;
let Ok(plain) = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT)
else {
return Ok(None);
};
Ok(protocol::from_body::<Frame>(&plain).ok())
})
.await?;
let mut next = cached.clone();
next.last_rendered_seq = frame.output_seq;
Ok((frame, next))
}
async fn recv_response_until<T, F>(socket: &UdpSocket, wait: Duration, mut accept: F) -> Result<T>
where
F: FnMut(protocol::Packet) -> Result<Option<T>>,
{
let started = Instant::now();
let mut buf = vec![0u8; 65535];
loop {
let elapsed = started.elapsed();
if elapsed >= wait {
return Err(anyhow!("timed out waiting for server response"));
}
let remaining = wait - elapsed;
let (n, _) = tokio::time::timeout(remaining, socket.recv_from(&mut buf))
.await
.context("timed out waiting for server response")??;
let Ok(packet) = protocol::decode(&buf[..n]) else {
continue;
};
if let Some(response) = accept(packet)? {
return Ok(response);
}
}
}
async fn run_terminal(
socket: UdpSocket,
mut cred: CachedCredential,
@@ -2800,10 +2840,11 @@ mod tests {
DynamicForward, FrameBuffer, LocalForward, Predictor, RemoteForward, SshConfig,
auth_allows, load_first_native_identity_with_prompt, parse_dynamic_forward,
parse_local_forward, parse_remote_forward, parse_ssh_config, raw_contains_host_table,
ssh_destination_host, ssh_username, startup_command, toml_bare_key_or_quoted,
recv_response_until, ssh_destination_host, ssh_username, startup_command,
toml_bare_key_or_quoted,
};
use dosh::config::ClientConfig;
use dosh::protocol::Frame;
use dosh::protocol::{self, Frame, PacketKind};
use ed25519_dalek::{SigningKey, VerifyingKey};
#[test]
@@ -2981,6 +3022,30 @@ mod tests {
);
}
#[tokio::test]
async fn recv_response_until_ignores_unmatched_packets() {
let receiver = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let sender = tokio::net::UdpSocket::bind("127.0.0.1:0").await.unwrap();
let addr = receiver.local_addr().unwrap();
let stale = protocol::encode_plain(PacketKind::Ping, [1u8; 16], 1, 0, b"stale").unwrap();
let valid = protocol::encode_plain(PacketKind::Pong, [2u8; 16], 2, 0, b"valid").unwrap();
sender.send_to(&stale, addr).await.unwrap();
sender.send_to(&valid, addr).await.unwrap();
let packet = recv_response_until(&receiver, std::time::Duration::from_secs(1), |packet| {
if packet.header.kind == PacketKind::Pong {
Ok(Some(packet))
} else {
Ok(None)
}
})
.await
.unwrap();
assert_eq!(packet.header.kind, PacketKind::Pong);
assert_eq!(packet.body, b"valid");
}
fn write_identity(path: &std::path::Path, signing: &SigningKey) {
let keypair = ssh_key::private::Ed25519Keypair::from(signing);
let private =