From 8f549f4968228a221d30b45f08e4a19ce44b0ccf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabr=C3=ADcio=20Bracht?= Date: Mon, 12 Oct 2026 14:14:12 -0700 Subject: [PATCH] serve quic client data flows on bidirectional streams --- CHANGELOG.md | 10 +- QUIC_IMPLEMENTATION_SPEC.md | 10 +- .../src/broker/client_handler/bidi_flow.rs | 417 ++++++++++++++++++ crates/mqtt5/src/broker/client_handler/mod.rs | 47 +- .../src/broker/client_handler/publish.rs | 8 + .../src/broker/client_handler/subscribe.rs | 30 +- crates/mqtt5/src/broker/quic_acceptor.rs | 193 ++++++-- crates/mqtt5/tests/quic_flow_streams.rs | 297 ++++++++++++- specs/tla/bidi-client-flow/BidiClientFlow.cfg | 19 + specs/tla/bidi-client-flow/BidiClientFlow.tla | 270 ++++++++++++ .../BidiClientFlow_NEG_recovery.cfg | 12 + .../BidiClientFlow_NEG_takeover.cfg | 12 + .../BidiClientFlow_V_close_blind.cfg | 12 + .../BidiClientFlow_V_detached_drop.cfg | 12 + ...BidiClientFlow_V_detached_serverstream.cfg | 12 + .../BidiClientFlow_V_route_control.cfg | 12 + .../BidiClientFlow_V_route_current.cfg | 12 + .../BidiClientFlow_V_takeover_current.cfg | 12 + .../BidiClientFlow_V_takeover_latest.cfg | 12 + specs/tla/bidi-client-flow/README.md | 100 +++++ 20 files changed, 1430 insertions(+), 79 deletions(-) create mode 100644 crates/mqtt5/src/broker/client_handler/bidi_flow.rs create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow.cfg create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow.tla create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow_NEG_recovery.cfg create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow_NEG_takeover.cfg create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow_V_close_blind.cfg create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow_V_detached_drop.cfg create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow_V_detached_serverstream.cfg create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow_V_route_control.cfg create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow_V_route_current.cfg create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow_V_takeover_current.cfg create mode 100644 specs/tla/bidi-client-flow/BidiClientFlow_V_takeover_latest.cfg create mode 100644 specs/tla/bidi-client-flow/README.md diff --git a/CHANGELOG.md b/CHANGELOG.md index 1750521b..a9961b81 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,13 +10,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed - **The QUIC broker refuses a client stream it does not support instead of closing the connection** (#207). Any client-opened bidirectional stream that was not a flow discard closed the connection with `ERROR_PROTOCOL_L0` (0xB4), including user-defined flows (0x14) and client data flows, which MQoQ §9.20 allows on bidirectional streams. A user-defined flow on a unidirectional stream was read as MQTT, so depending on its bytes it closed the connection as a malformed packet or was left open. As §9.10 describes, only the stream is now refused, and the connection stays up: - - a user-defined flow, or a client data flow on a bidirectional stream, with `ERROR_FLOW_REFUSED` (0xBE); + - a user-defined flow with `ERROR_FLOW_REFUSED` (0xBE); - a server data flow header, or a client data flow whose flow ID has the server bit set, with `ERROR_BAD_FLOW_ID` (0xB6); - a control flow header on any stream but the control stream, or a bidirectional stream without a flow header, with `ERROR_STREAM_TYPE` (0xB5). A refused bidirectional stream gets STOP_SENDING and RESET_STREAM with the code, a unidirectional one STOP_SENDING. Unidirectional streams without a flow header are still read as MQTT. - **Discarding a flow now removes its subscriptions** (#207). A discard removed the flow's state but left its flow-bound subscriptions, so the broker kept delivering to a flow the client had discarded. §9.21.2 counts subscriptions as part of the flow state that §9.16 says to discard. The broker now removes them before it answers the discard with FIN. -- **Only a flow header followed directly by FIN is a discard** (§9.16). A bidirectional stream whose header carried the discard flags was treated as a discard even when MQTT data followed, and the data was dropped. Such a stream is now refused with `ERROR_FLOW_REFUSED`. +- **Only a flow header followed directly by FIN is a discard** (§9.16). A bidirectional stream whose header carried the discard flags was treated as a discard even when MQTT data followed, and the data was dropped. Such a stream is now served as a client data flow. +- **The QUIC broker serves client data flows on bidirectional streams** (#207). MQoQ §9.20 lets a client open a data flow on a bidirectional stream, and §9.1.2 keeps the flow's traffic on that flow. The broker now answers each packet on the stream it arrived on (PUBACK, PUBREC, PUBREL, PUBCOMP, SUBACK, UNSUBACK and PINGRESP) and delivers the flow's subscriptions there, retained messages included. DISCONNECT and AUTH stay on the control stream. + - A newer stream for the same flow takes it over and the older one is reset (§9.14). A stream older than the newest one the flow has had is reset when it arrives, even if the newer one has already ended. + - A FIN from the client ends the flow: the broker answers with FIN and removes the flow's subscriptions (§9.13). If the client resets the stream instead, the flow keeps its subscriptions, and its QoS 1 and 2 messages wait for the client's next stream for it, up to the client's Receive Maximum; QoS 0 messages are dropped meanwhile (§9.21.2). If more messages arrive than that, the broker discards the flow. + - A flow that ends inside a packet is reset with `ERROR_INCOMPLETE_PACKET` (0xBA) and kept (§9.13.1). + + The design is checked by a TLA+ model in `specs/tla/bidi-client-flow`. ## [mqtt5 0.48.0] - 2026-10-11 diff --git a/QUIC_IMPLEMENTATION_SPEC.md b/QUIC_IMPLEMENTATION_SPEC.md index 21182c00..26c6e544 100644 --- a/QUIC_IMPLEMENTATION_SPEC.md +++ b/QUIC_IMPLEMENTATION_SPEC.md @@ -283,9 +283,11 @@ graph TD #### Broker Connection Handling -The broker spawns three concurrent tasks per QUIC connection: a **client handler** that processes packets from the control stream, a **datagram reader** that decodes unreliable datagrams and forwards them via mpsc channel, and a **data stream acceptor** that accepts unidirectional streams and spawns per-stream readers with flow header parsing. +The broker spawns these concurrent tasks per QUIC connection: a **client handler** that processes packets from the control stream, a **datagram reader** that decodes unreliable datagrams and forwards them via mpsc channel, a **unidirectional stream acceptor** that spawns per-stream readers with flow header parsing, and a **bidirectional stream acceptor** for client data flows and flow discards. -Data stream readers detect flow headers by inspecting the first byte (0x11-0x13) and register flows in a per-connection `FlowRegistry`. +Data stream readers detect flow headers by inspecting the first byte (0x11-0x14) and register flows in a per-connection `FlowRegistry`. + +A client data flow on a bidirectional stream is served on that stream: the broker's answers to the flow's packets and the deliveries for its subscriptions go back on it. A newer stream for the same flow takes it over and the older one is reset. While the flow has no stream, its QoS 1/2 deliveries wait for the next one, up to the client's Receive Maximum. A flow header followed directly by FIN discards the flow. The design is modelled in `specs/tla/bidi-client-flow`. ### Internal Types @@ -393,3 +395,7 @@ The project uses Quinn 0.11.x. Both client and broker set ALPN to `mqtt`. **Connection migration (`quic_migration_tests.rs`):** - test_quic_migration_detected_by_server, test_quic_migration_qos1_survives - test_quic_multiple_migrations, test_migrate_non_quic_returns_error, test_migrate_not_connected_returns_error + +**Broker flow streams (`quic_flow_streams.rs`):** +- refused streams and their error codes, flow discards +- bidirectional client data flows: answers and deliveries on the flow's stream, takeover, late older streams, messages waiting for a new stream, FIN, incomplete packets, retained messages diff --git a/crates/mqtt5/src/broker/client_handler/bidi_flow.rs b/crates/mqtt5/src/broker/client_handler/bidi_flow.rs new file mode 100644 index 00000000..a6c8eded --- /dev/null +++ b/crates/mqtt5/src/broker/client_handler/bidi_flow.rs @@ -0,0 +1,417 @@ +#[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] +pub(crate) use quic::{BidiFlow, BidiFlowEvent, OpenedStream, StreamEnding}; + +#[cfg(any(target_arch = "wasm32", not(feature = "transport-quic")))] +pub(crate) enum BidiFlowEvent {} + +#[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] +mod quic { + use super::super::subscribe::FlowActivity; + use super::super::ClientHandler; + use crate::broker::router::RoutableMessage; + use crate::broker::storage::QueuedMessage; + use crate::error::Result; + use crate::packet::Packet; + use crate::session::quic_flow::FlowState; + use crate::QoS; + use mqtt5_protocol::QuicStreamCode; + use quinn::{SendStream, VarInt}; + use std::collections::VecDeque; + use tokio::sync::oneshot; + use tracing::{debug, warn}; + + const MAX_BIDI_FLOWS: usize = 256; + + pub(crate) enum StreamEnding { + Finished, + Incomplete, + Aborted, + } + + pub(crate) struct OpenedStream { + pub(crate) flow: FlowState, + pub(crate) stream_id: u64, + pub(crate) send: SendStream, + pub(crate) abort: oneshot::Sender, + } + + pub(crate) enum BidiFlowEvent { + Opened(Box), + Packet { + flow_id: u64, + stream_id: u64, + packet: Packet, + }, + Closed { + flow_id: u64, + stream_id: u64, + ending: StreamEnding, + }, + } + + struct FlowStream { + id: u64, + send: SendStream, + abort: oneshot::Sender, + } + + impl FlowStream { + fn close(mut self, code: QuicStreamCode) { + if let Err(e) = self.send.reset(VarInt::from_u32(code.code())) { + debug!( + stream_id = self.id, + "flow stream already closed before reset: {e}" + ); + } + if self.abort.send(code).is_err() { + debug!(stream_id = self.id, "flow stream reader already ended"); + } + } + + fn finish(mut self) { + if let Err(e) = self.send.finish() { + debug!( + stream_id = self.id, + "flow stream already closed before finish: {e}" + ); + } + } + } + + pub(crate) struct LiveFlow { + current: Option, + highest: u64, + detached: VecDeque, + } + + pub(crate) enum BidiFlow { + Live(LiveFlow), + Ended { highest: u64 }, + } + + impl BidiFlow { + fn highest(&self) -> u64 { + match self { + Self::Live(live) => live.highest, + Self::Ended { highest } => *highest, + } + } + } + + impl ClientHandler { + pub(in super::super) async fn handle_bidi_event( + &mut self, + event: BidiFlowEvent, + ) -> Result { + match event { + BidiFlowEvent::Opened(opened) => { + self.open_bidi_stream(opened).await; + Ok(FlowActivity::Lifecycle) + } + BidiFlowEvent::Packet { + flow_id, + stream_id, + packet, + } => { + self.handle_bidi_packet(flow_id, stream_id, packet).await?; + Ok(FlowActivity::Packet) + } + BidiFlowEvent::Closed { + flow_id, + stream_id, + ending, + } => { + self.close_bidi_stream(flow_id, stream_id, ending).await; + Ok(FlowActivity::Lifecycle) + } + } + } + + async fn open_bidi_stream(&mut self, opened: Box) { + let OpenedStream { + flow, + stream_id, + send, + abort, + } = *opened; + let flow_id = flow.id.raw(); + let stream = FlowStream { + id: stream_id, + send, + abort, + }; + if self + .bidi_flows + .get(&flow_id) + .is_some_and(|known| stream_id <= known.highest()) + { + debug!( + flow_id, + stream_id, "aborting a stream older than the flow's newest" + ); + stream.close(QuicStreamCode::NoError); + return; + } + if let Some(BidiFlow::Live(live)) = self.bidi_flows.get_mut(&flow_id) { + if let Some(previous) = live.current.replace(stream) { + debug!( + flow_id, + stream_id, + previous = previous.id, + "flow taken over" + ); + previous.close(QuicStreamCode::NoError); + } + live.highest = stream_id; + let waiting = std::mem::take(&mut live.detached); + self.touch_flow(flow.id).await; + self.requeue_for_flow(flow_id, waiting); + return; + } + if !self.admit_bidi_flow(flow_id) || !self.register_flow(flow).await { + warn!(flow_id, "refusing bidirectional flow: flow limit reached"); + stream.close(QuicStreamCode::FlowRefused); + return; + } + debug!(flow_id, stream_id, "bidirectional client data flow opened"); + self.bidi_flows.insert( + flow_id, + BidiFlow::Live(LiveFlow { + current: Some(stream), + highest: stream_id, + detached: VecDeque::new(), + }), + ); + } + + fn admit_bidi_flow(&mut self, flow_id: u64) -> bool { + if self.bidi_flows.contains_key(&flow_id) || self.bidi_flows.len() < MAX_BIDI_FLOWS { + return true; + } + let ended = self + .bidi_flows + .iter() + .find_map(|(id, flow)| matches!(flow, BidiFlow::Ended { .. }).then_some(*id)); + ended.is_some_and(|id| self.bidi_flows.remove(&id).is_some()) + } + + async fn register_flow(&self, flow: FlowState) -> bool { + let Some(registry) = self.flow_registry.as_ref() else { + return true; + }; + let mut registry = registry.lock().await; + if registry.contains(flow.id) { + registry.touch(flow.id); + return true; + } + registry.register_flow(flow) + } + + async fn touch_flow(&self, flow_id: crate::transport::flow::FlowId) { + if let Some(registry) = self.flow_registry.as_ref() { + registry.lock().await.touch(flow_id); + } + } + + fn requeue_for_flow(&self, flow_id: u64, waiting: VecDeque) { + if waiting.is_empty() { + return; + } + let (Some(queue), Some(client_id)) = (self.queue.as_ref(), self.client_id.as_ref()) + else { + warn!( + flow_id, + count = waiting.len(), + "no session queue for the flow's waiting messages" + ); + return; + }; + debug!( + flow_id, + count = waiting.len(), + "delivering messages that waited for the flow" + ); + let messages = waiting + .into_iter() + .map(|routable| { + let qos = routable.publish.qos; + QueuedMessage::new(routable.publish, client_id.clone(), qos, None) + .with_target_flow(Some(flow_id)) + }) + .collect(); + queue.requeue_front_in_epoch(messages, self.queue_epoch); + queue.notify(); + } + + async fn handle_bidi_packet( + &mut self, + flow_id: u64, + stream_id: u64, + packet: Packet, + ) -> Result<()> { + if !self.is_current_stream(flow_id, stream_id) { + debug!( + flow_id, + stream_id, "dropping packet from a stream that is not the flow's current one" + ); + return Ok(()); + } + self.pending_external_flow_id = Some(flow_id); + self.reply_flow = Some(flow_id); + let result = self.handle_packet(packet).await; + self.reply_flow = None; + self.pending_external_flow_id = None; + result + } + + fn is_current_stream(&self, flow_id: u64, stream_id: u64) -> bool { + matches!( + self.bidi_flows.get(&flow_id), + Some(BidiFlow::Live(LiveFlow { current: Some(stream), .. })) if stream.id == stream_id + ) + } + + async fn close_bidi_stream(&mut self, flow_id: u64, stream_id: u64, ending: StreamEnding) { + if !self.is_current_stream(flow_id, stream_id) { + return; + } + let Some(stream) = self.detach_stream(flow_id) else { + return; + }; + match ending { + StreamEnding::Finished => { + debug!(flow_id, stream_id, "flow finished by the client"); + stream.finish(); + self.end_bidi_flow(flow_id).await; + } + StreamEnding::Incomplete => { + warn!(flow_id, stream_id, "flow ended inside a packet"); + stream.close(QuicStreamCode::IncompletePacket); + } + StreamEnding::Aborted => { + debug!(flow_id, stream_id, "flow aborted by the client"); + stream.close(QuicStreamCode::NoError); + } + } + } + + fn detach_stream(&mut self, flow_id: u64) -> Option { + match self.bidi_flows.get_mut(&flow_id) { + Some(BidiFlow::Live(live)) => live.current.take(), + _ => None, + } + } + + async fn end_bidi_flow(&mut self, flow_id: u64) { + let Some(flow) = self.bidi_flows.get_mut(&flow_id) else { + return; + }; + let highest = flow.highest(); + if let BidiFlow::Live(live) = std::mem::replace(flow, BidiFlow::Ended { highest }) { + if let Some(stream) = live.current { + stream.close(QuicStreamCode::NoError); + } + } + if let Some(registry) = self.flow_registry.as_ref() { + registry + .lock() + .await + .remove(crate::transport::flow::FlowId::from(flow_id)); + } + self.handle_flow_closed(flow_id).await; + } + + pub(in super::super) fn discard_bidi_flow(&mut self, flow_id: u64) { + let Some(flow) = self.bidi_flows.get_mut(&flow_id) else { + return; + }; + let highest = flow.highest(); + if let BidiFlow::Live(live) = std::mem::replace(flow, BidiFlow::Ended { highest }) { + if let Some(stream) = live.current { + stream.close(QuicStreamCode::DiscardState); + } + } + } + + pub(in super::super) async fn hold_for_flow( + &mut self, + routable: RoutableMessage, + ) -> Option { + let Some(flow_id) = routable.target_flow else { + return Some(routable); + }; + let capacity = usize::from(self.client_receive_maximum); + match self.bidi_flows.get_mut(&flow_id) { + None => Some(routable), + Some(BidiFlow::Ended { .. }) => { + debug!(flow_id, "dropping a message for a flow that has ended"); + None + } + Some(BidiFlow::Live(live)) => { + if live.current.is_some() { + return Some(routable); + } + if routable.publish.qos == QoS::AtMostOnce { + debug!( + flow_id, + "dropping a QoS 0 message while the flow has no stream" + ); + return None; + } + if live.detached.len() < capacity { + live.detached.push_back(routable); + return None; + } + warn!( + flow_id, + capacity, "flow has no stream and its waiting messages are at the limit; discarding the flow" + ); + self.end_bidi_flow(flow_id).await; + None + } + } + } + + pub(in super::super) fn is_bidi_flow(&self, flow_id: u64) -> bool { + self.bidi_flows.contains_key(&flow_id) + } + + pub(in super::super) async fn write_buffer_to_flow(&mut self, flow_id: u64) -> Result<()> { + let Some(BidiFlow::Live(LiveFlow { + current: Some(stream), + .. + })) = self.bidi_flows.get_mut(&flow_id) + else { + debug!(flow_id, "dropping a write for a flow without a stream"); + return Ok(()); + }; + if let Err(e) = stream.send.write_all(&self.write_buffer).await { + debug!(flow_id, "flow stream write failed, detaching the flow: {e}"); + if let Some(stream) = self.detach_stream(flow_id) { + stream.close(QuicStreamCode::NoError); + } + } + Ok(()) + } + + pub(in super::super) fn take_detached_messages( + &mut self, + client_id: &str, + ) -> Vec { + self.bidi_flows + .iter_mut() + .filter_map(|(flow_id, flow)| match flow { + BidiFlow::Live(live) => Some((*flow_id, std::mem::take(&mut live.detached))), + BidiFlow::Ended { .. } => None, + }) + .flat_map(|(flow_id, waiting)| { + waiting.into_iter().map(move |routable| { + let qos = routable.publish.qos; + QueuedMessage::new(routable.publish, client_id.to_string(), qos, None) + .with_target_flow(Some(flow_id)) + }) + }) + .collect() + } + } +} diff --git a/crates/mqtt5/src/broker/client_handler/mod.rs b/crates/mqtt5/src/broker/client_handler/mod.rs index ec9b6d68..22f568ef 100644 --- a/crates/mqtt5/src/broker/client_handler/mod.rs +++ b/crates/mqtt5/src/broker/client_handler/mod.rs @@ -1,9 +1,13 @@ mod auth; +mod bidi_flow; mod connect; mod lifecycle; mod publish; mod subscribe; +#[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] +pub(crate) use bidi_flow::{BidiFlowEvent, OpenedStream, StreamEnding}; + use crate::broker::auth::AuthProvider; use crate::broker::config::BrokerConfig; use crate::broker::events::{ClientConnectEvent, ClientDisconnectEvent}; @@ -161,6 +165,10 @@ pub struct ClientHandler { pub(super) pending_target_flow: Option, pub(super) flow_closed_rx: Option>, pub(super) flow_discard_rx: Option)>>, + pub(super) bidi_flow_rx: Option>, + pub(super) reply_flow: Option, + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] + pub(super) bidi_flows: HashMap, #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] pub(super) flow_registry: Option>>, @@ -283,6 +291,10 @@ impl ClientHandler { pending_target_flow: None, flow_closed_rx: None, flow_discard_rx: None, + bidi_flow_rx: None, + reply_flow: None, + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] + bidi_flows: HashMap::new(), #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] flow_registry: None, #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] @@ -342,6 +354,7 @@ impl ClientHandler { self } + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] #[must_use] pub(crate) fn with_flow_discard_rx( mut self, @@ -351,6 +364,13 @@ impl ClientHandler { self } + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] + #[must_use] + pub(crate) fn with_bidi_flow_rx(mut self, rx: mpsc::Receiver) -> Self { + self.bidi_flow_rx = Some(rx); + self + } + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] #[must_use] pub fn with_flow_registry( @@ -803,6 +823,25 @@ impl ClientHandler { } let write_timeout = self.write_timeout(); let packet_name = packet.packet_type_name(); + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] + if let Some(flow_id) = self + .reply_flow + .filter(|_| !matches!(packet, Packet::Disconnect(_) | Packet::Auth(_))) + { + self.write_buffer.clear(); + encode_packet_to_buffer(&packet, &mut self.write_buffer)?; + return timeout(write_timeout, self.write_buffer_to_flow(flow_id)) + .await + .map_err(|_| { + warn!( + packet = packet_name, + timeout = ?write_timeout, + flow_id, + "Flow stream write stalled; treating the client as gone" + ); + MqttError::Timeout + })?; + } timeout(write_timeout, self.transport.write_packet(packet)) .await .map_err(|_| { @@ -971,6 +1010,8 @@ impl ClientHandler { self.outbound_inflight.clear(); self.awaiting_pubcomp.clear(); unsent.append(&mut self.held); + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] + unsent.extend(self.take_detached_messages(&client_id)); while let Some(routable) = self.qos1_rx.recv().await { let qos = routable.publish.qos; unsent.push( @@ -1083,8 +1124,10 @@ impl ClientHandler { } } - flow_event = Self::next_flow_event(self.flow_closed_rx.as_mut(), self.flow_discard_rx.as_mut()) => { - self.handle_flow_event(flow_event).await; + flow_event = Self::next_flow_event(self.flow_closed_rx.as_mut(), self.flow_discard_rx.as_mut(), self.bidi_flow_rx.as_mut()) => { + if self.handle_flow_event(flow_event).await? == subscribe::FlowActivity::Packet { + last_packet_time = Instant::now(); + } } routable = async { diff --git a/crates/mqtt5/src/broker/client_handler/publish.rs b/crates/mqtt5/src/broker/client_handler/publish.rs index f14ffeef..527d4ff2 100644 --- a/crates/mqtt5/src/broker/client_handler/publish.rs +++ b/crates/mqtt5/src/broker/client_handler/publish.rs @@ -619,6 +619,11 @@ impl ClientHandler { } pub(super) async fn send_routable(&mut self, routable: RoutableMessage) -> Result<()> { + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] + let Some(routable) = self.hold_for_flow(routable).await + else { + return Ok(()); + }; self.pending_target_flow = routable.target_flow; let result = self.send_publish(routable.publish).await; self.pending_target_flow = None; @@ -776,6 +781,9 @@ impl ClientHandler { if self.quic_connection.is_some() { if let Some(flow_id) = self.pending_target_flow { + if self.is_bidi_flow(flow_id) { + return self.write_buffer_to_flow(flow_id).await; + } if self.server_stream_manager.is_none() { self.server_stream_manager = Some(self.build_server_stream_manager()); } diff --git a/crates/mqtt5/src/broker/client_handler/subscribe.rs b/crates/mqtt5/src/broker/client_handler/subscribe.rs index 2656d0bc..52f1301e 100644 --- a/crates/mqtt5/src/broker/client_handler/subscribe.rs +++ b/crates/mqtt5/src/broker/client_handler/subscribe.rs @@ -18,11 +18,19 @@ use tracing::{debug, warn}; use crate::broker::router::{MessageRouter, RoutableMessage, Subscribed, Unsubscribed}; use crate::broker::session_slot::SessionSlotGuard; +use super::bidi_flow::BidiFlowEvent; use super::ClientHandler; pub(super) enum FlowEvent { Closed(u64), Discarded(u64, oneshot::Sender<()>), + Bidi(BidiFlowEvent), +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) enum FlowActivity { + Packet, + Lifecycle, } struct AcceptedFilter<'a> { @@ -348,7 +356,7 @@ impl ClientHandler { } let routable = RoutableMessage { publish: msg, - target_flow: None, + target_flow: self.reply_flow, }; self.offer_retained(routable); } @@ -439,6 +447,7 @@ impl ClientHandler { pub(super) async fn next_flow_event( closed_rx: Option<&mut mpsc::Receiver>, discard_rx: Option<&mut mpsc::Receiver<(u64, oneshot::Sender<()>)>>, + bidi_rx: Option<&mut mpsc::Receiver>, ) -> Option { let closed = async { match closed_rx { @@ -455,16 +464,28 @@ impl ClientHandler { None => std::future::pending().await, } }; + let bidi = async { + match bidi_rx { + Some(rx) => rx.recv().await.map(FlowEvent::Bidi), + None => std::future::pending().await, + } + }; tokio::select! { event = closed => event, event = discarded => event, + event = bidi => event, } } - pub(super) async fn handle_flow_event(&mut self, event: Option) { + pub(super) async fn handle_flow_event( + &mut self, + event: Option, + ) -> Result { match event { Some(FlowEvent::Closed(flow_id)) => self.handle_flow_closed(flow_id).await, Some(FlowEvent::Discarded(flow_id, done)) => { + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] + self.discard_bidi_flow(flow_id); self.handle_flow_closed(flow_id).await; if done.send(()).is_err() { debug!( @@ -473,8 +494,13 @@ impl ClientHandler { ); } } + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] + Some(FlowEvent::Bidi(event)) => return self.handle_bidi_event(event).await, + #[cfg(any(target_arch = "wasm32", not(feature = "transport-quic")))] + Some(FlowEvent::Bidi(event)) => match event {}, None => {} } + Ok(FlowActivity::Lifecycle) } pub(super) async fn handle_flow_closed(&mut self, flow_id: u64) { diff --git a/crates/mqtt5/src/broker/quic_acceptor.rs b/crates/mqtt5/src/broker/quic_acceptor.rs index 3f385f3f..c27da0cf 100644 --- a/crates/mqtt5/src/broker/quic_acceptor.rs +++ b/crates/mqtt5/src/broker/quic_acceptor.rs @@ -1,5 +1,5 @@ use crate::broker::auth::AuthProvider; -use crate::broker::client_handler::ClientHandler; +use crate::broker::client_handler::{BidiFlowEvent, ClientHandler, OpenedStream, StreamEnding}; use crate::broker::config::BrokerConfig; use crate::broker::resource_monitor::ResourceMonitor; use crate::broker::router::MessageRouter; @@ -567,6 +567,7 @@ async fn run_quic_handler_inner( let (packet_tx, packet_rx) = mpsc::channel::<(Packet, Option)>(100); let (flow_closed_tx, flow_closed_rx) = mpsc::channel::(32); let (flow_discard_tx, flow_discard_rx) = mpsc::channel(32); + let (bidi_flow_tx, bidi_flow_rx) = mpsc::channel(100); let flow_registry = Arc::new(Mutex::new(FlowRegistry::new(256))); let (protocol_version_tx, protocol_version_rx) = watch::channel(None); @@ -608,6 +609,7 @@ async fn run_quic_handler_inner( .with_skip_bridge_forwarding(skip_bridge_forwarding) .with_flow_closed_rx(flow_closed_rx) .with_flow_discard_rx(flow_discard_rx) + .with_bidi_flow_rx(bidi_flow_rx) .with_flow_registry(flow_registry.clone()); let handler_label = label; @@ -633,10 +635,14 @@ async fn run_quic_handler_inner( label, ); spawn_bi_accept_loop( - connection.clone(), - flow_registry.clone(), - flow_discard_tx, - peer_addr, + BidiStreamContext { + connection: connection.clone(), + peer_addr, + flow_registry: flow_registry.clone(), + flow_discard_tx, + bidi_flow_tx, + protocol_version_rx: protocol_version_rx.clone(), + }, label, ); spawn_quic_stats_sampler(connection.clone(), peer_addr); @@ -764,25 +770,24 @@ fn spawn_datagram_reader( }); } -fn spawn_bi_accept_loop( +#[derive(Clone)] +struct BidiStreamContext { connection: Arc, + peer_addr: SocketAddr, flow_registry: Arc>, flow_discard_tx: mpsc::Sender<(u64, oneshot::Sender<()>)>, - peer_addr: SocketAddr, - label: &'static str, -) { + bidi_flow_tx: mpsc::Sender, + protocol_version_rx: watch::Receiver>, +} + +fn spawn_bi_accept_loop(context: BidiStreamContext, label: &'static str) { tokio::spawn(async move { + let peer_addr = context.peer_addr; loop { - match connection.accept_bi().await { + match context.connection.accept_bi().await { Ok((send, recv)) => { debug!("{} bidirectional stream accepted from {}", label, peer_addr); - spawn_bi_stream_handler( - send, - recv, - peer_addr, - flow_registry.clone(), - flow_discard_tx.clone(), - ); + spawn_bi_stream_handler(send, recv, context.clone()); } Err(e) => { let reason = crate::transport::quic_error::parse_connection_error(&e); @@ -864,18 +869,19 @@ fn refuse_stream( } } -async fn ends_after_header(recv: &mut RecvStream) -> bool { - matches!(recv.read_chunk(1, true).await, Ok(None)) +async fn chunk_after_header(recv: &mut RecvStream) -> Option> { + match recv.read_chunk(usize::MAX, true).await { + Ok(chunk) => Some(chunk.map(|chunk| chunk.bytes)), + Err(e) => { + debug!("stream ended before its first packet: {e}"); + None + } + } } -fn spawn_bi_stream_handler( - mut send: SendStream, - mut recv: RecvStream, - peer_addr: SocketAddr, - flow_registry: Arc>, - flow_discard_tx: mpsc::Sender<(u64, oneshot::Sender<()>)>, -) { +fn spawn_bi_stream_handler(mut send: SendStream, mut recv: RecvStream, context: BidiStreamContext) { tokio::spawn(async move { + let peer_addr = context.peer_addr; let result = match try_read_flow_header(&mut recv).await { Ok(StreamOpening::Flow(result)) => result, Ok(StreamOpening::Refused(code)) => { @@ -893,34 +899,123 @@ fn spawn_bi_stream_handler( return; }; - if !flags.is_discard_signal() - || !result.leftover.is_empty() - || !ends_after_header(&mut recv).await - { - refuse_stream(&mut send, &mut recv, peer_addr, QuicStreamCode::FlowRefused); - return; - } + let buffer = if flags.is_discard_signal() && result.leftover.is_empty() { + match chunk_after_header(&mut recv).await { + Some(Some(bytes)) => BytesMut::from(bytes.as_ref()), + Some(None) => { + discard_flow(send, flow_id, &context).await; + return; + } + None => return, + } + } else { + result.leftover + }; - if let Some(state) = flow_registry.lock().await.remove(flow_id) { - debug!(flow_id = ?flow_id, "discarded flow state for {peer_addr}: {state:?}"); - } + let flow = FlowState::new_client_data(flow_id, flags, result.expire); + serve_bidi_flow(send, recv, flow, buffer, &context).await; + }); +} - let (discarded_tx, discarded_rx) = oneshot::channel(); - if flow_discard_tx - .send((flow_id.raw(), discarded_tx)) - .await - .is_err() - || discarded_rx.await.is_err() - { - debug!(flow_id = ?flow_id, "client handler ended before the flow was discarded"); - } +async fn discard_flow(mut send: SendStream, flow_id: FlowId, context: &BidiStreamContext) { + let peer_addr = context.peer_addr; + if let Some(state) = context.flow_registry.lock().await.remove(flow_id) { + debug!(flow_id = ?flow_id, "discarded flow state for {peer_addr}: {state:?}"); + } - if let Err(e) = send.finish() { - debug!(flow_id = ?flow_id, "discard stream from {peer_addr} already closed: {e}"); + let (discarded_tx, discarded_rx) = oneshot::channel(); + if context + .flow_discard_tx + .send((flow_id.raw(), discarded_tx)) + .await + .is_err() + || discarded_rx.await.is_err() + { + debug!(flow_id = ?flow_id, "client handler ended before the flow was discarded"); + } + + if let Err(e) = send.finish() { + debug!(flow_id = ?flow_id, "discard stream from {peer_addr} already closed: {e}"); + } + + debug!(flow_id = ?flow_id, "completed discard handshake for {peer_addr}"); +} + +async fn serve_bidi_flow( + send: SendStream, + mut recv: RecvStream, + flow: FlowState, + mut buffer: BytesMut, + context: &BidiStreamContext, +) { + let peer_addr = context.peer_addr; + let Some(protocol_version) = + negotiated_protocol_version(context.protocol_version_rx.clone()).await + else { + debug!("client handler ended before CONNECT from {peer_addr}"); + return; + }; + let flow_id = flow.id.raw(); + let stream_id = u64::from(recv.id()); + let (abort_tx, mut abort_rx) = oneshot::channel(); + let opened = OpenedStream { + flow, + stream_id, + send, + abort: abort_tx, + }; + if context + .bidi_flow_tx + .send(BidiFlowEvent::Opened(Box::new(opened))) + .await + .is_err() + { + debug!(flow_id, "client handler ended before the flow opened"); + return; + } + + let ending = loop { + tokio::select! { + code = &mut abort_rx => { + let code = code.unwrap_or(QuicStreamCode::NoError); + if let Err(e) = recv.stop(quinn::VarInt::from_u32(code.code())) { + debug!(flow_id, stream_id, "flow stream already closed before stop: {e}"); + } + return; + } + read = read_packet_with_buffer(&mut recv, &mut buffer, protocol_version) => match read { + Ok(packet) => { + let event = BidiFlowEvent::Packet { flow_id, stream_id, packet }; + if context.bidi_flow_tx.send(event).await.is_err() { + debug!(flow_id, "client handler ended, stopping flow reader"); + return; + } + } + Err(MqttError::ClientClosed) if buffer.is_empty() => break StreamEnding::Finished, + Err(MqttError::ClientClosed) => break StreamEnding::Incomplete, + Err(MqttError::ConnectionError(reason)) => { + debug!(flow_id, stream_id, "flow stream from {peer_addr} ended: {reason}"); + break StreamEnding::Aborted; + } + Err(e) => { + close_on_malformed_packet(&context.connection, peer_addr, "data flow", &e); + return; + } + } } + }; - debug!(flow_id = ?flow_id, "completed discard handshake for {peer_addr}"); - }); + let closed = BidiFlowEvent::Closed { + flow_id, + stream_id, + ending, + }; + if context.bidi_flow_tx.send(closed).await.is_err() { + debug!( + flow_id, + "client handler ended before the flow stream closed" + ); + } } const CONTROL_FLOW_CLOSE_GRACE: Duration = Duration::from_secs(1); diff --git a/crates/mqtt5/tests/quic_flow_streams.rs b/crates/mqtt5/tests/quic_flow_streams.rs index c5e3eb8f..25fdddcd 100644 --- a/crates/mqtt5/tests/quic_flow_streams.rs +++ b/crates/mqtt5/tests/quic_flow_streams.rs @@ -296,17 +296,6 @@ async fn a_unidirectional_client_data_flow_is_served() { assert_eq!(on_uni(&PINGREQ).await, (StreamEnd::Open, STILL_SERVED)); } -#[tokio::test] -async fn a_bidirectional_client_data_flow_is_refused() { - assert_eq!( - on_bidi(&with_pingreq(data_flow(1, 0, FlowFlags::default())), false).await, - ( - StreamEnd::Refused(code(QuicStreamCode::FlowRefused)), - STILL_SERVED - ) - ); -} - #[tokio::test] async fn a_discard_header_ended_by_fin_is_answered_with_fin() { assert_eq!( @@ -315,17 +304,6 @@ async fn a_discard_header_ended_by_fin_is_answered_with_fin() { ); } -#[tokio::test] -async fn a_discard_header_followed_by_data_is_not_a_discard() { - assert_eq!( - on_bidi(&with_pingreq(data_flow(1, 0, FlowFlags::discard())), false).await, - ( - StreamEnd::Refused(code(QuicStreamCode::FlowRefused)), - STILL_SERVED - ) - ); -} - #[tokio::test] async fn a_unidirectional_flow_ending_inside_a_packet_keeps_the_session() { let mut flow = data_flow(1, 0, FlowFlags::default()); @@ -422,3 +400,278 @@ async fn discarding_a_flow_removes_its_subscriptions() { assert!(delivered_to_aborted_flow(false).await); assert!(!delivered_to_aborted_flow(true).await); } + +const PINGRESP: [u8; 2] = [0xD0, 0x00]; + +async fn read_frame(recv: &mut quinn::RecvStream) -> Option> { + let mut frame = vec![0u8; 2]; + tokio::time::timeout(Duration::from_secs(2), recv.read_exact(&mut frame)) + .await + .ok()? + .ok()?; + let mut body = vec![0u8; usize::from(frame[1])]; + recv.read_exact(&mut body).await.ok()?; + frame.extend_from_slice(&body); + Some(frame) +} + +async fn reset_code(recv: &mut quinn::RecvStream) -> Option { + match tokio::time::timeout(Duration::from_secs(2), recv.read_to_end(1024)).await { + Ok(Err(quinn::ReadToEndError::Read(quinn::ReadError::Reset(code)))) => { + Some(code.into_inner()) + } + _ => None, + } +} + +fn qos1_subscribe_frame(topic: &str) -> Vec { + let mut frame = subscribe_frame(topic); + let last = frame.len() - 1; + frame[last] = 0x01; + frame +} + +fn qos1_publish_frame(topic: &str, packet_id: u16, payload: &[u8]) -> Vec { + let mut body = Vec::new(); + body.extend_from_slice( + &u16::try_from(topic.len()) + .expect("short topic") + .to_be_bytes(), + ); + body.extend_from_slice(topic.as_bytes()); + body.extend_from_slice(&packet_id.to_be_bytes()); + body.push(0x00); + body.extend_from_slice(payload); + let mut frame = vec![0x32, u8::try_from(body.len()).expect("short PUBLISH")]; + frame.extend_from_slice(&body); + frame +} + +fn payload_of(frame: &[u8]) -> &[u8] { + let topic_len = usize::from(u16::from_be_bytes([frame[2], frame[3]])); + let packet_id_len = if frame[0] & 0x06 == 0 { 0 } else { 2 }; + let properties_at = 4 + topic_len + packet_id_len; + &frame[properties_at + 1 + usize::from(frame[properties_at])..] +} + +async fn open_flow( + split: &QuicSplitResult, + sequence: u64, + first: &[u8], +) -> (quinn::SendStream, quinn::RecvStream) { + let (mut send, recv) = split.connection.open_bi().await.expect("open bidi flow"); + let mut bytes = data_flow(sequence, 300, FlowFlags::default()); + bytes.extend_from_slice(first); + send.write_all(&bytes).await.expect("write flow"); + (send, recv) +} + +async fn publish_from_control( + split: &mut QuicSplitResult, + topic: &str, + packet_id: u16, + payload: &[u8], +) { + split + .send + .write_all(&qos1_publish_frame(topic, packet_id, payload)) + .await + .expect("publish on control"); + let puback = read_frame(&mut split.recv) + .await + .expect("PUBACK on control"); + assert_eq!(puback[0], 0x40); +} + +async fn nothing_off_the_flow(split: &mut QuicSplitResult) -> bool { + let server_stream = + tokio::time::timeout(Duration::from_millis(500), split.connection.accept_bi()) + .await + .is_err(); + server_stream && outcome(split).await == STILL_SERVED +} + +#[tokio::test] +async fn a_bidirectional_client_data_flow_is_answered_on_its_own_stream() { + for flags in [FlowFlags::default(), FlowFlags::discard()] { + let (broker, mut split) = connected_session("bidi-served").await; + let (mut send, mut recv) = split.connection.open_bi().await.expect("open bidi"); + send.write_all(&with_pingreq(data_flow(1, 0, flags))) + .await + .expect("write flow"); + assert_eq!(read_frame(&mut recv).await, Some(PINGRESP.to_vec())); + assert_eq!(outcome(&mut split).await, STILL_SERVED); + broker.stop().await; + } +} + +#[tokio::test] +async fn subscriptions_on_a_bidirectional_flow_are_delivered_on_it() { + let (broker, mut split) = connected_session("bidi-subscribe").await; + let (_send, mut recv) = open_flow(&split, 1, &qos1_subscribe_frame("bidi/a")).await; + assert_eq!( + read_frame(&mut recv).await, + Some(vec![0x90, 0x04, 0x00, 0x01, 0x00, 0x01]) + ); + publish_from_control(&mut split, "bidi/a", 7, b"one").await; + let delivered = read_frame(&mut recv).await.expect("PUBLISH on the flow"); + assert_eq!(delivered[0] & 0xF0, 0x30); + assert_eq!(payload_of(&delivered), b"one"); + assert!(nothing_off_the_flow(&mut split).await); + broker.stop().await; +} + +#[tokio::test] +async fn a_newer_stream_takes_the_flow_over() { + let (broker, mut split) = connected_session("bidi-takeover").await; + let (_old_send, mut old_recv) = open_flow(&split, 1, &qos1_subscribe_frame("bidi/t")).await; + assert!(read_frame(&mut old_recv).await.is_some()); + let (_new_send, mut new_recv) = open_flow(&split, 1, &PINGREQ).await; + assert_eq!(read_frame(&mut new_recv).await, Some(PINGRESP.to_vec())); + assert_eq!( + reset_code(&mut old_recv).await, + Some(code(QuicStreamCode::NoError)) + ); + publish_from_control(&mut split, "bidi/t", 8, b"after").await; + let delivered = read_frame(&mut new_recv) + .await + .expect("PUBLISH on the new stream"); + assert_eq!(payload_of(&delivered), b"after"); + assert!(nothing_off_the_flow(&mut split).await); + broker.stop().await; +} + +async fn late_older_stream_is_refused(newer_aborted: bool) { + let (broker, mut split) = connected_session("bidi-late").await; + let (mut older_send, mut older_recv) = split.connection.open_bi().await.expect("open older"); + let (mut newer_send, mut newer_recv) = + open_flow(&split, 1, &qos1_subscribe_frame("bidi/l")).await; + assert!(read_frame(&mut newer_recv).await.is_some()); + if newer_aborted { + newer_send + .reset(quinn::VarInt::from_u32( + QuicStreamCode::FlowCancelled.code(), + )) + .expect("abort the newer stream"); + assert_eq!( + reset_code(&mut newer_recv).await, + Some(code(QuicStreamCode::NoError)) + ); + } + older_send + .write_all(&with_pingreq(data_flow(1, 300, FlowFlags::default()))) + .await + .expect("write older"); + assert_eq!( + reset_code(&mut older_recv).await, + Some(code(QuicStreamCode::NoError)) + ); + if !newer_aborted { + publish_from_control(&mut split, "bidi/l", 9, b"newer").await; + let delivered = read_frame(&mut newer_recv) + .await + .expect("PUBLISH on the newer stream"); + assert_eq!(payload_of(&delivered), b"newer"); + } + assert!(nothing_off_the_flow(&mut split).await); + broker.stop().await; +} + +#[tokio::test] +async fn an_older_stream_arriving_late_does_not_take_the_flow_back() { + late_older_stream_is_refused(false).await; + late_older_stream_is_refused(true).await; +} + +#[tokio::test] +async fn messages_wait_for_a_flow_that_lost_its_stream() { + let (broker, mut split) = connected_session("bidi-detached").await; + let (mut first_send, mut first_recv) = + open_flow(&split, 1, &qos1_subscribe_frame("bidi/d")).await; + assert!(read_frame(&mut first_recv).await.is_some()); + first_send + .reset(quinn::VarInt::from_u32( + QuicStreamCode::FlowCancelled.code(), + )) + .expect("abort the flow's stream"); + assert_eq!( + reset_code(&mut first_recv).await, + Some(code(QuicStreamCode::NoError)) + ); + split + .send + .write_all(&publish_frame("bidi/d", b"qos0")) + .await + .expect("QoS 0 publish"); + publish_from_control(&mut split, "bidi/d", 10, b"waited").await; + assert!(nothing_off_the_flow(&mut split).await); + let (mut second_send, mut second_recv) = open_flow(&split, 1, &[]).await; + let delivered = read_frame(&mut second_recv).await.expect("waiting PUBLISH"); + assert_eq!(delivered[0] & 0x06, 0x02); + assert_eq!(payload_of(&delivered), b"waited"); + let packet_id = [delivered[10], delivered[11]]; + second_send + .write_all(&[0x40, 0x02, packet_id[0], packet_id[1], 0xC0, 0x00]) + .await + .expect("PUBACK and PINGREQ on the flow"); + assert_eq!(read_frame(&mut second_recv).await, Some(PINGRESP.to_vec())); + broker.stop().await; +} + +#[tokio::test] +async fn finishing_a_bidirectional_flow_ends_it() { + let (broker, mut split) = connected_session("bidi-finish").await; + let (mut send, mut recv) = open_flow(&split, 1, &subscribe_frame("bidi/f")).await; + assert!(read_frame(&mut recv).await.is_some()); + send.finish().expect("finish the flow"); + let rest = tokio::time::timeout(Duration::from_secs(2), recv.read_to_end(64)) + .await + .expect("flow finished by the broker") + .expect("flow ends with FIN"); + assert!(rest.is_empty()); + split + .send + .write_all(&publish_frame("bidi/f", b"gone")) + .await + .expect("publish"); + assert!(nothing_off_the_flow(&mut split).await); + broker.stop().await; +} + +#[tokio::test] +async fn a_bidirectional_flow_ending_inside_a_packet_is_reset() { + let (broker, mut split) = connected_session("bidi-incomplete").await; + let (mut send, mut recv) = open_flow(&split, 1, &PINGREQ[..1]).await; + send.finish().expect("finish inside a packet"); + assert_eq!( + reset_code(&mut recv).await, + Some(code(QuicStreamCode::IncompletePacket)) + ); + assert_eq!(outcome(&mut split).await, STILL_SERVED); + broker.stop().await; +} + +#[tokio::test] +async fn retained_messages_for_a_bidirectional_flow_arrive_on_it() { + let (broker, mut split) = connected_session("bidi-retained").await; + let mut retained = publish_frame("bidi/r", b"kept"); + retained[0] = 0x31; + split + .send + .write_all(&retained) + .await + .expect("retained publish"); + assert!(outcome(&mut split).await.session_alive); + let (_send, mut recv) = open_flow(&split, 1, &qos1_subscribe_frame("bidi/r")).await; + assert_eq!( + read_frame(&mut recv).await, + Some(vec![0x90, 0x04, 0x00, 0x01, 0x00, 0x01]) + ); + let delivered = read_frame(&mut recv) + .await + .expect("retained PUBLISH on the flow"); + assert_eq!(delivered[0] & 0x01, 0x01); + assert_eq!(payload_of(&delivered), b"kept"); + assert!(nothing_off_the_flow(&mut split).await); + broker.stop().await; +} diff --git a/specs/tla/bidi-client-flow/BidiClientFlow.cfg b/specs/tla/bidi-client-flow/BidiClientFlow.cfg new file mode 100644 index 00000000..19074e4c --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow.cfg @@ -0,0 +1,19 @@ +SPECIFICATION Spec +CONSTANTS + MaxStreams = 3 + MaxReq = 2 + MaxMsg = 2 + TAKEOVER = "highest" + ROUTE = "stream" + CLOSE = "check" + DETACHED = "queue" +INVARIANTS + TypeOK + ReplyOnOwnStream + NothingOffFlow + CurrentIsNewest + NewestLiveIsCurrentWhenQuiet + AnsweredWhenQuiet + DeliveredWhenQuiet + NoSilentDrop +CHECK_DEADLOCK FALSE diff --git a/specs/tla/bidi-client-flow/BidiClientFlow.tla b/specs/tla/bidi-client-flow/BidiClientFlow.tla new file mode 100644 index 00000000..8852bfc5 --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow.tla @@ -0,0 +1,270 @@ +--------------------------- MODULE BidiClientFlow --------------------------- +EXTENDS Naturals, Sequences, FiniteSets, TLC + +CONSTANTS MaxStreams, MaxReq, MaxMsg, TAKEOVER, ROUTE, CLOSE, DETACHED + +ASSUME TAKEOVER \in {"highest", "current", "latest"} +ASSUME ROUTE \in {"stream", "current", "control"} +ASSUME CLOSE \in {"check", "blind"} +ASSUME DETACHED \in {"queue", "serverstream", "drop"} + +Streams == 1..MaxStreams +Control == 0 + +VARIABLES + cstate, + inq, + started, + ended, + babort, + chan, + cur, + handled, + out, + readFrom, + sent, + sub, + pending, + published, + got, + dropped, + discardReq, + discarded, + detachedMsgs + +vars == <> + +Reply(r) == [k |-> "reply", r |-> r] +Msg(m) == [k |-> "msg", m |-> m] + +TypeOK == + /\ cstate \in [Streams -> {"none", "open", "aborted"}] + /\ started \in [Streams -> BOOLEAN] + /\ ended \in [Streams -> BOOLEAN] + /\ babort \in [Streams -> BOOLEAN] + /\ cur \in Streams \cup {0} + /\ handled \subseteq Streams + /\ sub \in BOOLEAN + /\ discardReq \in BOOLEAN + /\ discarded \in BOOLEAN + +Init == + /\ cstate = [s \in Streams |-> "none"] + /\ inq = [s \in Streams |-> <<>>] + /\ started = [s \in Streams |-> FALSE] + /\ ended = [s \in Streams |-> FALSE] + /\ babort = [s \in Streams |-> FALSE] + /\ chan = <<>> + /\ cur = 0 + /\ handled = {} + /\ out = [s \in Streams \cup {Control} |-> <<>>] + /\ readFrom = {} + /\ sent = 0 + /\ sub = TRUE + /\ pending = <<>> + /\ published = 0 + /\ got = {} + /\ dropped = {} + /\ discardReq = FALSE + /\ discarded = FALSE + /\ detachedMsgs = {} + +ClientOpen(s) == + /\ ~discarded + /\ cstate[s] = "none" + /\ \A t \in Streams : t < s => cstate[t] # "none" + /\ cstate' = [cstate EXCEPT ![s] = "open"] + /\ UNCHANGED <> + +ClientSend(s) == + /\ cstate[s] = "open" + /\ sent < MaxReq + /\ sent' = sent + 1 + /\ inq' = [inq EXCEPT ![s] = Append(@, sent + 1)] + /\ UNCHANGED <> + +ClientAbort(s) == + /\ cstate[s] = "open" + /\ cstate' = [cstate EXCEPT ![s] = "aborted"] + /\ inq' = [inq EXCEPT ![s] = <<>>] + /\ UNCHANGED <> + +ClientDiscard == + /\ ~discarded + /\ \E s \in Streams : cstate[s] # "none" + /\ \A s \in Streams : cstate[s] # "open" + /\ discarded' = TRUE + /\ discardReq' = TRUE + /\ UNCHANGED <> + +ReaderOpen(s) == + /\ cstate[s] # "none" + /\ ~started[s] + /\ started' = [started EXCEPT ![s] = TRUE] + /\ chan' = Append(chan, [t |-> "open", s |-> s]) + /\ UNCHANGED <> + +ReaderRead(s) == + /\ started[s] + /\ ~ended[s] + /\ ~babort[s] + /\ inq[s] # <<>> + /\ chan' = Append(chan, [t |-> "req", s |-> s, r |-> Head(inq[s])]) + /\ readFrom' = readFrom \cup {<>} + /\ inq' = [inq EXCEPT ![s] = Tail(@)] + /\ UNCHANGED <> + +ReaderEnd(s) == + /\ started[s] + /\ ~ended[s] + /\ cstate[s] = "aborted" \/ babort[s] + /\ ended' = [ended EXCEPT ![s] = TRUE] + /\ chan' = Append(chan, [t |-> "closed", s |-> s]) + /\ UNCHANGED <> + +AbortStream(b, s) == IF s = 0 THEN b ELSE [b EXCEPT ![s] = TRUE] + +Highest(S) == IF S = {} THEN 0 ELSE CHOOSE x \in S : \A y \in S : y <= x + +Supersedes(s) == + CASE TAKEOVER = "highest" -> s > Highest(handled) + [] TAKEOVER = "current" -> s > cur + [] TAKEOVER = "latest" -> TRUE + +HandleOpen(e) == + /\ handled' = handled \cup {e.s} + /\ IF Supersedes(e.s) + THEN /\ cur' = e.s + /\ babort' = AbortStream(babort, cur) + ELSE /\ babort' = AbortStream(babort, e.s) + /\ UNCHANGED cur + /\ UNCHANGED <> + +ReplyTarget(e) == + CASE ROUTE = "stream" -> IF cur = e.s THEN e.s ELSE -1 + [] ROUTE = "current" -> IF cur # 0 THEN cur ELSE -1 + [] ROUTE = "control" -> Control + +HandleReq(e) == + /\ LET target == ReplyTarget(e) + IN out' = IF target = -1 THEN out + ELSE [out EXCEPT ![target] = Append(@, Reply(e.r))] + /\ UNCHANGED <> + +HandleClosed(e) == + /\ cur' = IF CLOSE = "blind" \/ cur = e.s THEN 0 ELSE cur + /\ UNCHANGED <> + +Handle == + /\ chan # <<>> + /\ chan' = Tail(chan) + /\ LET e == Head(chan) + IN CASE e.t = "open" -> HandleOpen(e) + [] e.t = "req" -> HandleReq(e) + [] e.t = "closed" -> HandleClosed(e) + /\ UNCHANGED <> + +HandleDiscard == + /\ discardReq + /\ discardReq' = FALSE + /\ sub' = FALSE + /\ pending' = <<>> + /\ cur' = 0 + /\ babort' = AbortStream(babort, cur) + /\ UNCHANGED <> + +Publish == + /\ sub + /\ published < MaxMsg + /\ published' = published + 1 + /\ pending' = Append(pending, published + 1) + /\ detachedMsgs' = IF cur = 0 /\ handled # {} THEN detachedMsgs \cup {published + 1} ELSE detachedMsgs + /\ UNCHANGED <> + +DeliverOnFlow == + /\ pending # <<>> + /\ cur # 0 + /\ ~babort[cur] + /\ cstate[cur] = "open" + /\ out' = [out EXCEPT ![cur] = Append(@, Msg(Head(pending)))] + /\ got' = got \cup {Head(pending)} + /\ pending' = Tail(pending) + /\ UNCHANGED <> + +DeliverDetached == + /\ pending # <<>> + /\ cur = 0 + /\ DETACHED # "queue" + /\ pending' = Tail(pending) + /\ IF DETACHED = "serverstream" + THEN /\ out' = [out EXCEPT ![Control] = Append(@, Msg(Head(pending)))] + /\ got' = got \cup {Head(pending)} + /\ UNCHANGED dropped + ELSE /\ dropped' = dropped \cup {Head(pending)} + /\ UNCHANGED <> + /\ UNCHANGED <> + +BrokerNext == + \/ \E s \in Streams : ReaderOpen(s) \/ ReaderRead(s) \/ ReaderEnd(s) + \/ Handle + \/ HandleDiscard + \/ DeliverOnFlow + \/ DeliverDetached + +ClientNext == + \/ \E s \in Streams : ClientOpen(s) \/ ClientSend(s) \/ ClientAbort(s) + \/ ClientDiscard + \/ Publish + +Next == BrokerNext \/ ClientNext + +Spec == Init /\ [][Next]_vars + +Quiescent == ~ENABLED BrokerNext + +NewestLive(s) == + /\ cstate[s] = "open" + /\ ~babort[s] + /\ \A t \in Streams : t > s => cstate[t] = "none" + +ReplyOnOwnStream == + \A s \in Streams : \A i \in 1..Len(out[s]) : + out[s][i].k = "reply" => <> \in readFrom + +NothingOffFlow == out[Control] = <<>> + +CurrentIsNewest == cur # 0 => cur = Highest(handled) + +NewestLiveIsCurrentWhenQuiet == + Quiescent => \A s \in Streams : NewestLive(s) => cur = s + +AnsweredWhenQuiet == + Quiescent => \A s \in Streams : NewestLive(s) => + \A p \in readFrom : p[1] = s => + \E i \in 1..Len(out[s]) : out[s][i] = Reply(p[2]) + +DeliveredWhenQuiet == + Quiescent => ((\E s \in Streams : NewestLive(s)) => pending = <<>>) + +NoSilentDrop == dropped = {} + +NEG_TakeoverDelivers == + ~(\E s \in Streams : s > 1 /\ babort[s - 1] /\ \E i \in 1..Len(out[s]) : out[s][i].k = "msg") + +NEG_RecoveryDelivers == ~(\E s \in Streams : s > 1 /\ \E i \in 1..Len(out[s]) : out[s][i].k = "msg" /\ out[s][i].m \in detachedMsgs /\ cstate[s - 1] = "aborted") + +============================================================================= diff --git a/specs/tla/bidi-client-flow/BidiClientFlow_NEG_recovery.cfg b/specs/tla/bidi-client-flow/BidiClientFlow_NEG_recovery.cfg new file mode 100644 index 00000000..265ad398 --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow_NEG_recovery.cfg @@ -0,0 +1,12 @@ +SPECIFICATION Spec +CONSTANTS + MaxStreams = 3 + MaxReq = 2 + MaxMsg = 2 + TAKEOVER = "highest" + ROUTE = "stream" + CLOSE = "check" + DETACHED = "queue" +INVARIANTS + NEG_RecoveryDelivers +CHECK_DEADLOCK FALSE diff --git a/specs/tla/bidi-client-flow/BidiClientFlow_NEG_takeover.cfg b/specs/tla/bidi-client-flow/BidiClientFlow_NEG_takeover.cfg new file mode 100644 index 00000000..80beaba8 --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow_NEG_takeover.cfg @@ -0,0 +1,12 @@ +SPECIFICATION Spec +CONSTANTS + MaxStreams = 3 + MaxReq = 2 + MaxMsg = 2 + TAKEOVER = "highest" + ROUTE = "stream" + CLOSE = "check" + DETACHED = "queue" +INVARIANTS + NEG_TakeoverDelivers +CHECK_DEADLOCK FALSE diff --git a/specs/tla/bidi-client-flow/BidiClientFlow_V_close_blind.cfg b/specs/tla/bidi-client-flow/BidiClientFlow_V_close_blind.cfg new file mode 100644 index 00000000..4125bc38 --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow_V_close_blind.cfg @@ -0,0 +1,12 @@ +SPECIFICATION Spec +CONSTANTS + MaxStreams = 3 + MaxReq = 2 + MaxMsg = 2 + TAKEOVER = "highest" + ROUTE = "stream" + CLOSE = "blind" + DETACHED = "queue" +INVARIANTS + NewestLiveIsCurrentWhenQuiet +CHECK_DEADLOCK FALSE diff --git a/specs/tla/bidi-client-flow/BidiClientFlow_V_detached_drop.cfg b/specs/tla/bidi-client-flow/BidiClientFlow_V_detached_drop.cfg new file mode 100644 index 00000000..6436466d --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow_V_detached_drop.cfg @@ -0,0 +1,12 @@ +SPECIFICATION Spec +CONSTANTS + MaxStreams = 3 + MaxReq = 2 + MaxMsg = 2 + TAKEOVER = "highest" + ROUTE = "stream" + CLOSE = "check" + DETACHED = "drop" +INVARIANTS + NoSilentDrop +CHECK_DEADLOCK FALSE diff --git a/specs/tla/bidi-client-flow/BidiClientFlow_V_detached_serverstream.cfg b/specs/tla/bidi-client-flow/BidiClientFlow_V_detached_serverstream.cfg new file mode 100644 index 00000000..ea7c48cd --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow_V_detached_serverstream.cfg @@ -0,0 +1,12 @@ +SPECIFICATION Spec +CONSTANTS + MaxStreams = 3 + MaxReq = 2 + MaxMsg = 2 + TAKEOVER = "highest" + ROUTE = "stream" + CLOSE = "check" + DETACHED = "serverstream" +INVARIANTS + NothingOffFlow +CHECK_DEADLOCK FALSE diff --git a/specs/tla/bidi-client-flow/BidiClientFlow_V_route_control.cfg b/specs/tla/bidi-client-flow/BidiClientFlow_V_route_control.cfg new file mode 100644 index 00000000..be3d3872 --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow_V_route_control.cfg @@ -0,0 +1,12 @@ +SPECIFICATION Spec +CONSTANTS + MaxStreams = 3 + MaxReq = 2 + MaxMsg = 2 + TAKEOVER = "highest" + ROUTE = "control" + CLOSE = "check" + DETACHED = "queue" +INVARIANTS + NothingOffFlow +CHECK_DEADLOCK FALSE diff --git a/specs/tla/bidi-client-flow/BidiClientFlow_V_route_current.cfg b/specs/tla/bidi-client-flow/BidiClientFlow_V_route_current.cfg new file mode 100644 index 00000000..9d27bfd7 --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow_V_route_current.cfg @@ -0,0 +1,12 @@ +SPECIFICATION Spec +CONSTANTS + MaxStreams = 3 + MaxReq = 2 + MaxMsg = 2 + TAKEOVER = "highest" + ROUTE = "current" + CLOSE = "check" + DETACHED = "queue" +INVARIANTS + ReplyOnOwnStream +CHECK_DEADLOCK FALSE diff --git a/specs/tla/bidi-client-flow/BidiClientFlow_V_takeover_current.cfg b/specs/tla/bidi-client-flow/BidiClientFlow_V_takeover_current.cfg new file mode 100644 index 00000000..a24469a3 --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow_V_takeover_current.cfg @@ -0,0 +1,12 @@ +SPECIFICATION Spec +CONSTANTS + MaxStreams = 3 + MaxReq = 2 + MaxMsg = 2 + TAKEOVER = "current" + ROUTE = "stream" + CLOSE = "check" + DETACHED = "queue" +INVARIANTS + CurrentIsNewest +CHECK_DEADLOCK FALSE diff --git a/specs/tla/bidi-client-flow/BidiClientFlow_V_takeover_latest.cfg b/specs/tla/bidi-client-flow/BidiClientFlow_V_takeover_latest.cfg new file mode 100644 index 00000000..16be24e9 --- /dev/null +++ b/specs/tla/bidi-client-flow/BidiClientFlow_V_takeover_latest.cfg @@ -0,0 +1,12 @@ +SPECIFICATION Spec +CONSTANTS + MaxStreams = 3 + MaxReq = 2 + MaxMsg = 2 + TAKEOVER = "latest" + ROUTE = "stream" + CLOSE = "check" + DETACHED = "queue" +INVARIANTS + CurrentIsNewest +CHECK_DEADLOCK FALSE diff --git a/specs/tla/bidi-client-flow/README.md b/specs/tla/bidi-client-flow/README.md new file mode 100644 index 00000000..56862747 --- /dev/null +++ b/specs/tla/bidi-client-flow/README.md @@ -0,0 +1,100 @@ +# Client data flows on bidirectional QUIC streams + +Model for #207: how the broker serves a client data flow (0x12) that the client opens on a +bidirectional stream. MQoQ §9.1.2 requires the broker to answer on the same flow: PUBACK, +PUBREC, PUBREL and PUBCOMP "must be exchanged in the same data flow", PINGRESP goes "through +the same data flow where it received the PINGREQ", and "the subscriber must expect messages for +subscriptions from the same flow it subscribed to". §9.14 lets the client take a flow over with +a newer stream, and "greater QUIC Stream ID of the same QUIC stream type always takes +precedence". + +## What is modelled + +One client data flow with a subscription, up to `MaxStreams` bidirectional streams for it +(numbered in the order the client opens them), `MaxReq` requests and `MaxMsg` published +messages. + +- The client opens streams in order, sends requests on an open stream, aborts streams, and + discards the flow once none of its streams is open. +- Each stream has its own broker reader. A reader announces the stream, forwards each request, + and announces the end. All readers feed one channel to the client handler, so events from + different streams interleave in any order while each stream's own events stay in order. +- The client handler keeps the flow's current stream (`cur`), answers requests, and delivers the + subscription's messages. A discard arrives on its own path, as it does in the broker. +- Writes are recorded per stream in `out`; `out[0]` stands for any stream that is not one of the + flow's own, such as the control stream or a broker-opened server stream. + +## Design rules and their alternatives + +| Constant | Proposed | Alternatives | +|---|---|---| +| `TAKEOVER` | `highest`: a stream becomes current only if its ID is above every stream the flow has had | `current`: above the current stream; `latest`: the last one announced wins | +| `ROUTE` | `stream`: answer on the stream the request came from, only while it is current, else drop | `current`: answer on whatever stream is current; `control`: answer on the control stream | +| `CLOSE` | `check`: a stream's end clears `cur` only if it is the current stream | `blind`: any stream's end clears `cur` | +| `DETACHED` | `queue`: while the flow has no current stream, its messages wait | `serverstream`: send them on a stream the broker opens; `drop`: drop them | + +## Properties + +- `ReplyOnOwnStream`: every reply on a stream answers a request read from that stream. +- `NothingOffFlow`: nothing for the flow is written outside its own streams. +- `CurrentIsNewest`: the current stream is the highest stream the handler has seen. +- `NewestLiveIsCurrentWhenQuiet`: once the broker has nothing left to do, the client's newest + open stream that the broker has not aborted is the current one. +- `AnsweredWhenQuiet`: once the broker is quiet, every request read from that stream is answered. +- `DeliveredWhenQuiet`: once the broker is quiet and the flow has a live stream, no message is + left waiting. +- `NoSilentDrop`: no message is dropped. + +`NEG_TakeoverDelivers` and `NEG_RecoveryDelivers` are reachability controls and must be +violated: a message delivered on a stream that took over from an older one, and a message +published while the flow had lost its stream that is delivered once a new stream recovers it. + +## Results + +All runs use `MaxStreams = 3`, `MaxReq = 2`, `MaxMsg = 2`, checked with tla-mcp to completion or +to the first violation. + +| Config | Result | +|---|---| +| `BidiClientFlow.cfg` (proposed) | all invariants hold, 3,478,866 states, depth 31, state space exhausted | +| `BidiClientFlow_V_takeover_current.cfg` | `CurrentIsNewest` violated at depth 10 | +| `BidiClientFlow_V_takeover_latest.cfg` | `CurrentIsNewest` violated at depth 7 | +| `BidiClientFlow_V_route_current.cfg` | `ReplyOnOwnStream` violated at depth 10 | +| `BidiClientFlow_V_route_control.cfg` | `NothingOffFlow` violated at depth 7 | +| `BidiClientFlow_V_close_blind.cfg` | `NewestLiveIsCurrentWhenQuiet` violated at depth 9 | +| `BidiClientFlow_V_detached_serverstream.cfg` | `NothingOffFlow` violated at depth 3 | +| `BidiClientFlow_V_detached_drop.cfg` | `NoSilentDrop` violated at depth 3 | +| `BidiClientFlow_NEG_takeover.cfg` | violated at depth 9, as required | +| `BidiClientFlow_NEG_recovery.cfg` | violated at depth 12, as required | + +The `current` takeover rule was the first proposal. Its counterexample: the client opens streams +1 and 2, stream 2 is announced and becomes current, stream 2 ends so `cur` is cleared, and then +stream 1's late announcement makes the older stream current. Comparing against the highest +stream the flow has ever had, rather than the current one, closes it. + +The recovery control's trace: stream 1 is current, the client aborts it and the broker clears +`cur`, a message is published and waits, the client opens stream 2, and the waiting message is +delivered on stream 2. + +## Implementation + +The broker follows the proposed configuration. + +- Readers: `serve_bidi_flow` in `crates/mqtt5/src/broker/quic_acceptor.rs`, one per stream, sends `Opened`, `Packet` and `Closed` events on one channel to the client handler (`chan`). A reader the handler aborts stops without a `Closed` event. +- Handler: `crates/mqtt5/src/broker/client_handler/bidi_flow.rs` keeps each flow's current stream and the highest stream ID it has had (`cur`, `Highest(handled)`). + - `TAKEOVER = highest` is `open_bidi_stream`. + - `ROUTE = stream` is `handle_bidi_packet`, which also drops a packet from a stream that is not current instead of only its reply. + - `CLOSE = check` is `close_bidi_stream`. + - `DETACHED = queue` is `hold_for_flow`, with the waiting messages requeued at the front of the session queue when a new stream takes the flow over. +- Beyond the model: QoS 0 messages are dropped while the flow has no stream, as §9.21.2 allows, and the waiting QoS 1/2 messages are capped at the client's Receive Maximum. A FIN from the client ends the flow and removes its subscriptions (§9.13). + +`crates/mqtt5/tests/quic_flow_streams.rs` exercises the rules against a running broker. Against main, all eight bidirectional flow tests fail and the eight others pass. Each rule was also broken in the implementation, one at a time, to check that a test catches it: + +| Mutation | Failing test | +|---|---| +| takeover compared against the current stream | `an_older_stream_arriving_late_does_not_take_the_flow_back` | +| messages dropped while the flow has no stream | `messages_wait_for_a_flow_that_lost_its_stream` | +| answers written to the control stream | six tests, including `a_bidirectional_client_data_flow_is_answered_on_its_own_stream` | +| retained messages not bound to the flow | `retained_messages_for_a_bidirectional_flow_arrive_on_it` | + +`CLOSE = check` has no test that fails without it. The handler aborts a superseded stream's reader, which then stops without sending `Closed`. A stale `Closed` therefore needs the client's own reset to reach the reader before the handler's abort does, and no test can force that ordering.