diff --git a/src/bin/dosh-client.rs b/src/bin/dosh-client.rs index c973c52..5989cfa 100644 --- a/src/bin/dosh-client.rs +++ b/src/bin/dosh-client.rs @@ -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")); - } - let plain = protocol::decrypt_body(&packet, &bootstrap.session_key, SERVER_TO_CLIENT)?; - let ok: AttachOk = protocol::from_body(&plain)?; + if packet.header.kind != PacketKind::AttachOk + || packet.header.session_key_id != expected_key_id + { + return Ok(None); + } + let Ok(plain) = protocol::decrypt_body(&packet, &bootstrap.session_key, SERVER_TO_CLIENT) + else { + return Ok(None); + }; + Ok(protocol::from_body::(&plain).ok()) + }) + .await?; let cred = CachedCredential { server: server_name.to_string(), session: ok.session.clone(), @@ -1729,33 +1735,36 @@ 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")); - } - 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)?; + if packet.header.kind != PacketKind::AttachOk { + return Ok(None); + } + let Ok(envelope) = protocol::from_body::(&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); + let response_key = crypto::hkdf32( + &cached.attach_ticket_psk, + &salt, + b"dosh/ticket-attach-ok/v1", + )?; + let Ok(plain) = crypto::open( + &response_key, + &envelope.server_nonce, + b"dosh-ticket-attach-ok-v1", + &envelope.ciphertext, + ) else { + return Ok(None); + }; + Ok(protocol::from_body::(&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")); - } - let plain = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT)?; - let frame: Frame = protocol::from_body(&plain)?; + if packet.header.kind != PacketKind::ResumeOk { + return Ok(None); + } + let Ok(plain) = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT) + else { + return Ok(None); + }; + Ok(protocol::from_body::(&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])?; - if packet.header.conn_id != cached.client_id { - return Err(anyhow!("resume received stale client packet")); - } - 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")); - } - let plain = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT)?; - let frame: Frame = protocol::from_body(&plain)?; + if packet.header.kind != PacketKind::ResumeOk { + return Ok(None); + } + let Ok(plain) = protocol::decrypt_body(&packet, &cached.session_key, SERVER_TO_CLIENT) + else { + return Ok(None); + }; + Ok(protocol::from_body::(&plain).ok()) + }) + .await?; let mut next = cached.clone(); next.last_rendered_seq = frame.output_seq; Ok((frame, next)) } +async fn recv_response_until(socket: &UdpSocket, wait: Duration, mut accept: F) -> Result +where + F: FnMut(protocol::Packet) -> Result>, +{ + 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 =