Ignore stale packets during attach resume

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::<AttachOk>(&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::<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);
+        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::<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"));
-    }
-    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::<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])?;
-    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::<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 =