fix(client): handle close-flagged echo replies

crates/irtt-client/src/client.rs

@@ -483,6 +483,9 @@ impl Client {
                     packet_meta: meta.into(),
                 });
             }
+            if flags::has(reply.flags, flags::FLAG_CLOSE) {
+                self.close_from_peer(token, now, &mut events)?;
+            }
             Ok(events)
         } else if session.completed.contains(wire_seq) {
             update_highest_received(session, wire_seq);
@@ -501,7 +504,7 @@ impl Client {
             update_highest_received(session, wire_seq);
             session.completed.insert(wire_seq, timed_out.logical_seq);
 
-            Ok(vec![ClientEvent::LateReply {
+            let mut events = vec![ClientEvent::LateReply {
                 seq: wire_seq,
                 logical_seq: Some(timed_out.logical_seq),
                 highest_seen,
@@ -514,7 +517,11 @@ impl Client {
                 received_stats,
                 bytes: packet_len,
                 packet_meta: meta.into(),
-            }])
+            }];
+            if flags::has(reply.flags, flags::FLAG_CLOSE) {
+                self.close_from_peer(token, now, &mut events)?;
+            }
+            Ok(events)
         } else if session
             .highest_received_seq
             .is_some_and(|h| sequence_is_before(wire_seq, h))
@@ -543,6 +550,26 @@ impl Client {
         }
     }
 
+    fn close_from_peer(
+        &mut self,
+        token: u64,
+        now: ClientTimestamp,
+        events: &mut Vec<ClientEvent>,
+    ) -> Result<(), ClientError> {
+        clear_dscp_on_socket(&self.socket, self.remote)?;
+        self.phase = ClientPhase::Closed;
+        if let Some(session) = self.session.as_mut() {
+            session.timed_out.clear();
+        }
+        self.session = None;
+        events.push(ClientEvent::SessionClosed {
+            remote: self.remote,
+            token,
+            at: now,
+        });
+        Ok(())
+    }
+
     fn accept_open_reply(
         &mut self,
         reply: OpenReply,

crates/irtt-client/src/client/tests/close.rs

@@ -95,3 +95,81 @@ fn close_clears_timed_out_metadata() {
     assert!(client.session.is_none());
     server.join();
 }
+
+#[test]
+fn close_flagged_echo_reply_emits_reply_then_closes_without_sending_close() {
+    let params = default_params();
+    let server = start_fake_server({
+        let params = params.clone();
+        move |socket, tx| {
+            let (_, peer) = recv_request(&socket, &tx);
+            let reply = open_reply(FLAG_OPEN | FLAG_REPLY, TOKEN, &params, None);
+            socket.send_to(&reply, peer).unwrap();
+
+            let (request, _) = recv_request(&socket, &tx);
+            let seq = u32::from_le_bytes(request[12..16].try_into().unwrap());
+            let reply = echo_reply_packet_with_flags(
+                TOKEN,
+                seq,
+                &params,
+                &TimestampFields::default(),
+                None,
+                FLAG_REPLY | flags::FLAG_CLOSE,
+            );
+            socket.send_to(&reply, peer).unwrap();
+
+            socket
+                .set_read_timeout(Some(Duration::from_millis(250)))
+                .unwrap();
+            while recv_request_timeout(&socket, &tx).is_some() {}
+        }
+    });
+    let config = ClientConfig {
+        socket_config: crate::SocketConfig {
+            recv_timeout: Some(Duration::from_millis(200)),
+            ..Default::default()
+        },
+        ..default_test_config(server.addr)
+    };
+    let mut client = Client::connect(config).unwrap();
+    assert_open_started(client.open(ClientTimestamp::now()).unwrap());
+    client.send_probe().unwrap();
+
+    let events = client.recv_once().unwrap();
+    assert!(matches!(
+        events.first(),
+        Some(ClientEvent::EchoReply { .. })
+    ));
+    assert!(matches!(
+        events.get(1),
+        Some(ClientEvent::SessionClosed { token: TOKEN, .. })
+    ));
+    assert_eq!(events.len(), 2);
+    assert!(client.next_send_deadline().is_none());
+    assert!(matches!(
+        client.send_probe(),
+        Err(ClientError::AlreadyClosed)
+    ));
+
+    let first = server.rx.recv_timeout(Duration::from_millis(100)).unwrap();
+    let second = server.rx.recv_timeout(Duration::from_millis(100)).unwrap();
+    assert_eq!(first[3] & FLAG_OPEN, FLAG_OPEN);
+    assert_eq!(second[3] & flags::FLAG_CLOSE, 0);
+    assert!(server.rx.recv_timeout(Duration::from_millis(400)).is_err());
+    server.join();
+}
+
+#[test]
+fn normal_echo_reply_does_not_close_session() {
+    let params = default_params();
+    let (mut client, server) = open_client_with_echo_server(&params);
+    client.send_probe().unwrap();
+
+    let events = client.recv_once().unwrap();
+    assert!(matches!(events.as_slice(), [ClientEvent::EchoReply { .. }]));
+    assert!(client.next_send_deadline().is_some());
+    assert!(client.send_probe().is_ok());
+
+    client.close(ClientTimestamp::now()).unwrap();
+    server.join();
+}

crates/irtt-client/src/client/tests/support.rs

@@ -80,13 +80,24 @@ pub(super) fn echo_reply_packet(
     params: &Params,
     timestamps: &TimestampFields,
     hmac_key: Option<&[u8]>,
+) -> Vec<u8> {
+    echo_reply_packet_with_flags(token, seq, params, timestamps, hmac_key, FLAG_REPLY)
+}
+
+pub(super) fn echo_reply_packet_with_flags(
+    token: u64,
+    seq: u32,
+    params: &Params,
+    timestamps: &TimestampFields,
+    hmac_key: Option<&[u8]>,
+    flags: u8,
 ) -> Vec<u8> {
     let has_hmac = hmac_key.is_some();
     let layout = PacketLayout::echo(has_hmac, params);
     let packet_len = echo_packet_len(has_hmac, params);
     let mut packet = Vec::with_capacity(packet_len);
 
-    let mut flags = FLAG_REPLY;
+    let mut flags = flags;
     if has_hmac {
         flags |= FLAG_HMAC;
     }