diff --git a/crates/irtt-client/src/client.rs b/crates/irtt-client/src/client.rs index b9b2c4f..0d76d48 100644 --- a/crates/irtt-client/src/client.rs +++ b/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, + ) -> 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, diff --git a/crates/irtt-client/src/client/tests/close.rs b/crates/irtt-client/src/client/tests/close.rs index 804686b..63b8962 100644 --- a/crates/irtt-client/src/client/tests/close.rs +++ b/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, ¶ms, 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, + ¶ms, + &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(¶ms); + 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(); +} diff --git a/crates/irtt-client/src/client/tests/support.rs b/crates/irtt-client/src/client/tests/support.rs index 9f847c2..59c7bcd 100644 --- a/crates/irtt-client/src/client/tests/support.rs +++ b/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 { + 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 { 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; }