From 4f8ea9a135a56a08a4603a4d8c3d4f0882b5d5b9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabr=C3=ADcio=20Bracht?= Date: Sun, 11 Oct 2026 19:47:22 -0700 Subject: [PATCH] refuse unsupported quic client streams without closing the connection --- CHANGELOG.md | 13 + crates/mqtt5/src/broker/client_handler/mod.rs | 23 +- .../src/broker/client_handler/subscribe.rs | 47 ++ crates/mqtt5/src/broker/quic_acceptor.rs | 183 ++++---- crates/mqtt5/tests/quic_flow_streams.rs | 424 ++++++++++++++++++ 5 files changed, 581 insertions(+), 109 deletions(-) create mode 100644 crates/mqtt5/tests/quic_flow_streams.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index d124c1aa..1750521b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,6 +5,19 @@ All notable changes to this project will be documented in this file. The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [Unreleased] + +### 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 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`. + ## [mqtt5 0.48.0] - 2026-10-11 ### Breaking diff --git a/crates/mqtt5/src/broker/client_handler/mod.rs b/crates/mqtt5/src/broker/client_handler/mod.rs index 1e18d44d..ec9b6d68 100644 --- a/crates/mqtt5/src/broker/client_handler/mod.rs +++ b/crates/mqtt5/src/broker/client_handler/mod.rs @@ -160,6 +160,7 @@ pub struct ClientHandler { pub(super) skip_bridge_forwarding: bool, pub(super) pending_target_flow: Option, pub(super) flow_closed_rx: Option>, + pub(super) flow_discard_rx: Option)>>, #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] pub(super) flow_registry: Option>>, @@ -281,6 +282,7 @@ impl ClientHandler { skip_bridge_forwarding: false, pending_target_flow: None, flow_closed_rx: None, + flow_discard_rx: None, #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] flow_registry: None, #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] @@ -340,6 +342,15 @@ impl ClientHandler { self } + #[must_use] + pub(crate) fn with_flow_discard_rx( + mut self, + rx: mpsc::Receiver<(u64, oneshot::Sender<()>)>, + ) -> Self { + self.flow_discard_rx = Some(rx); + self + } + #[cfg(all(not(target_arch = "wasm32"), feature = "transport-quic"))] #[must_use] pub fn with_flow_registry( @@ -1072,16 +1083,8 @@ impl ClientHandler { } } - closed_flow_id = async { - if let Some(ref mut rx) = self.flow_closed_rx { - rx.recv().await - } else { - std::future::pending::>().await - } - } => { - if let Some(flow_id) = closed_flow_id { - self.handle_flow_closed(flow_id).await; - } + 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; } routable = async { diff --git a/crates/mqtt5/src/broker/client_handler/subscribe.rs b/crates/mqtt5/src/broker/client_handler/subscribe.rs index b9fc9086..2656d0bc 100644 --- a/crates/mqtt5/src/broker/client_handler/subscribe.rs +++ b/crates/mqtt5/src/broker/client_handler/subscribe.rs @@ -12,6 +12,7 @@ use crate::packet::Packet; use crate::protocol::v5::reason_codes::ReasonCode; use crate::validation::{parse_shared_subscription, topic_matches_filter, validate_topic_filter}; use crate::QoS; +use tokio::sync::{mpsc, oneshot}; use tracing::{debug, warn}; use crate::broker::router::{MessageRouter, RoutableMessage, Subscribed, Unsubscribed}; @@ -19,6 +20,11 @@ use crate::broker::session_slot::SessionSlotGuard; use super::ClientHandler; +pub(super) enum FlowEvent { + Closed(u64), + Discarded(u64, oneshot::Sender<()>), +} + struct AcceptedFilter<'a> { filter: &'a crate::packet::subscribe::TopicFilter, stored: StoredSubscription, @@ -430,6 +436,47 @@ impl ClientHandler { result } + pub(super) async fn next_flow_event( + closed_rx: Option<&mut mpsc::Receiver>, + discard_rx: Option<&mut mpsc::Receiver<(u64, oneshot::Sender<()>)>>, + ) -> Option { + let closed = async { + match closed_rx { + Some(rx) => rx.recv().await.map(FlowEvent::Closed), + None => std::future::pending().await, + } + }; + let discarded = async { + match discard_rx { + Some(rx) => rx + .recv() + .await + .map(|(flow_id, done)| FlowEvent::Discarded(flow_id, done)), + None => std::future::pending().await, + } + }; + tokio::select! { + event = closed => event, + event = discarded => event, + } + } + + pub(super) async fn handle_flow_event(&mut self, event: Option) { + match event { + Some(FlowEvent::Closed(flow_id)) => self.handle_flow_closed(flow_id).await, + Some(FlowEvent::Discarded(flow_id, done)) => { + self.handle_flow_closed(flow_id).await; + if done.send(()).is_err() { + debug!( + flow_id, + "discard stream ended before its flow was discarded" + ); + } + } + None => {} + } + } + pub(super) async fn handle_flow_closed(&mut self, flow_id: u64) { let Some(ref client_id) = self.client_id else { return; diff --git a/crates/mqtt5/src/broker/quic_acceptor.rs b/crates/mqtt5/src/broker/quic_acceptor.rs index f62cf5b1..3f385f3f 100644 --- a/crates/mqtt5/src/broker/quic_acceptor.rs +++ b/crates/mqtt5/src/broker/quic_acceptor.rs @@ -11,8 +11,10 @@ use crate::packet::{FixedHeader, Packet}; use crate::session::quic_flow::{FlowRegistry, FlowState}; use crate::transport::flow::{ FlowFlags, FlowHeader, FlowId, FLOW_TYPE_CLIENT_DATA, FLOW_TYPE_CONTROL, FLOW_TYPE_SERVER_DATA, + FLOW_TYPE_USER_DEFINED, }; use bytes::{Buf, Bytes, BytesMut}; +use mqtt5_protocol::QuicStreamCode; use quinn::{Connection, Endpoint, RecvStream, SendStream, ServerConfig}; use rustls::pki_types::{CertificateDer, PrivateKeyDer}; use rustls::server::WebPkiClientVerifier; @@ -23,7 +25,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::time::Duration; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; -use tokio::sync::{broadcast, mpsc, watch, Mutex}; +use tokio::sync::{broadcast, mpsc, oneshot, watch, Mutex}; use tracing::{debug, error, instrument, trace, warn}; use super::tls_acceptor::TlsAcceptorConfig; @@ -332,10 +334,15 @@ pub async fn accept_quic_stream( fn is_flow_header_byte(b: u8) -> bool { matches!( b, - FLOW_TYPE_CONTROL | FLOW_TYPE_CLIENT_DATA | FLOW_TYPE_SERVER_DATA + FLOW_TYPE_CONTROL | FLOW_TYPE_CLIENT_DATA | FLOW_TYPE_SERVER_DATA | FLOW_TYPE_USER_DEFINED ) } +enum StreamOpening { + Flow(FlowHeaderResult), + Refused(QuicStreamCode), +} + struct FlowHeaderResult { flow_id: Option, flags: Option, @@ -344,28 +351,28 @@ struct FlowHeaderResult { } #[instrument(skip(recv), level = "debug")] -async fn try_read_flow_header(recv: &mut RecvStream) -> Result { +async fn try_read_flow_header(recv: &mut RecvStream) -> Result { let chunk = recv .read_chunk(1, true) .await .map_err(|e| MqttError::ConnectionError(format!("Failed to read stream: {e}")))?; let Some(chunk) = chunk else { - return Ok(FlowHeaderResult { + return Ok(StreamOpening::Flow(FlowHeaderResult { flow_id: None, flags: None, expire: None, leftover: BytesMut::new(), - }); + })); }; if chunk.bytes.is_empty() { - return Ok(FlowHeaderResult { + return Ok(StreamOpening::Flow(FlowHeaderResult { flow_id: None, flags: None, expire: None, leftover: BytesMut::new(), - }); + })); } let first_byte = chunk.bytes[0]; @@ -373,12 +380,12 @@ async fn try_read_flow_header(recv: &mut RecvStream) -> Result if !is_flow_header_byte(first_byte) { let mut leftover = BytesMut::with_capacity(chunk.bytes.len()); leftover.extend_from_slice(&chunk.bytes); - return Ok(FlowHeaderResult { + return Ok(StreamOpening::Flow(FlowHeaderResult { flow_id: None, flags: None, expire: None, leftover, - }); + })); } let mut header_buf = Vec::with_capacity(32); @@ -417,40 +424,27 @@ async fn try_read_flow_header(recv: &mut RecvStream) -> Result Ok(build_flow_header_result(flow_header, leftover)) } -fn build_flow_header_result(flow_header: FlowHeader, leftover: BytesMut) -> FlowHeaderResult { +fn build_flow_header_result(flow_header: FlowHeader, leftover: BytesMut) -> StreamOpening { match flow_header { - FlowHeader::Control(h) => { - trace!(flow_id = ?h.flow_id, "Parsed control flow header"); - FlowHeaderResult { - flow_id: Some(h.flow_id), - flags: Some(h.flags), - expire: None, - leftover, - } - } - FlowHeader::ClientData(h) | FlowHeader::ServerData(h) => { + FlowHeader::ClientData(h) if h.flow_id.is_client_initiated() => { let expire = if h.expire_interval > 0 { Some(Duration::from_secs(h.expire_interval)) } else { None }; trace!(flow_id = ?h.flow_id, expire = ?expire, "Parsed data flow header"); - FlowHeaderResult { + StreamOpening::Flow(FlowHeaderResult { flow_id: Some(h.flow_id), flags: Some(h.flags), expire, leftover, - } + }) } - FlowHeader::UserDefined(_) => { - trace!("Ignoring user-defined flow header"); - FlowHeaderResult { - flow_id: None, - flags: None, - expire: None, - leftover, - } + FlowHeader::ClientData(_) | FlowHeader::ServerData(_) => { + StreamOpening::Refused(QuicStreamCode::BadFlowId) } + FlowHeader::Control(_) => StreamOpening::Refused(QuicStreamCode::StreamType), + FlowHeader::UserDefined(_) => StreamOpening::Refused(QuicStreamCode::FlowRefused), } } @@ -572,6 +566,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 flow_registry = Arc::new(Mutex::new(FlowRegistry::new(256))); let (protocol_version_tx, protocol_version_rx) = watch::channel(None); @@ -612,6 +607,7 @@ async fn run_quic_handler_inner( .with_quic_protocol_version_tx(protocol_version_tx) .with_skip_bridge_forwarding(skip_bridge_forwarding) .with_flow_closed_rx(flow_closed_rx) + .with_flow_discard_rx(flow_discard_rx) .with_flow_registry(flow_registry.clone()); let handler_label = label; @@ -636,7 +632,13 @@ async fn run_quic_handler_inner( peer_addr, label, ); - spawn_bi_accept_loop(connection.clone(), flow_registry.clone(), peer_addr, label); + spawn_bi_accept_loop( + connection.clone(), + flow_registry.clone(), + flow_discard_tx, + peer_addr, + label, + ); spawn_quic_stats_sampler(connection.clone(), peer_addr); spawn_uni_accept_loop( connection, @@ -765,6 +767,7 @@ fn spawn_datagram_reader( fn spawn_bi_accept_loop( connection: Arc, flow_registry: Arc>, + flow_discard_tx: mpsc::Sender<(u64, oneshot::Sender<()>)>, peer_addr: SocketAddr, label: &'static str, ) { @@ -773,12 +776,12 @@ fn spawn_bi_accept_loop( match connection.accept_bi().await { Ok((send, recv)) => { debug!("{} bidirectional stream accepted from {}", label, peer_addr); - spawn_discard_handler( + spawn_bi_stream_handler( send, recv, peer_addr, flow_registry.clone(), - connection.clone(), + flow_discard_tx.clone(), ); } Err(e) => { @@ -846,58 +849,39 @@ fn decode_datagram_packet(data: &Bytes, protocol_version: u8) -> Option= error_level { - warn!( - ?code, - error_level, err_tolerance, "Stream error from {peer_addr}, resetting stream" - ); - if let Err(e) = send.reset(quinn::VarInt::from_u32(code.code())) { - debug!("stream from {peer_addr} already closed before reset: {e}"); - } - return; + if let Err(e) = send.reset(quinn::VarInt::from_u32(code.code())) { + debug!("stream from {peer_addr} already closed before reset: {e}"); } +} - warn!( - ?code, - error_level, err_tolerance, "Error exceeds tolerance for {peer_addr}, closing connection" - ); - connection.close( - quinn::VarInt::from_u32(mqtt5_protocol::QuicConnectionCode::ProtocolLevel0.code()), - code.to_string().as_bytes(), - ); +async fn ends_after_header(recv: &mut RecvStream) -> bool { + matches!(recv.read_chunk(1, true).await, Ok(None)) } -fn spawn_discard_handler( +fn spawn_bi_stream_handler( mut send: SendStream, mut recv: RecvStream, peer_addr: SocketAddr, flow_registry: Arc>, - connection: Arc, + flow_discard_tx: mpsc::Sender<(u64, oneshot::Sender<()>)>, ) { tokio::spawn(async move { let result = match try_read_flow_header(&mut recv).await { - Ok(result) => result, + Ok(StreamOpening::Flow(result)) => result, + Ok(StreamOpening::Refused(code)) => { + refuse_stream(&mut send, &mut recv, peer_addr, code); + return; + } Err(e) => { warn!("failed to read flow header on bi stream from {peer_addr}: {e}"); return; @@ -905,40 +889,30 @@ fn spawn_discard_handler( }; let (Some(flow_id), Some(flags)) = (result.flow_id, result.flags) else { - handle_stream_error( - &mut send, - &connection, - peer_addr, - mqtt5_protocol::QuicStreamCode::NoFlowState, - 0, - ); + refuse_stream(&mut send, &mut recv, peer_addr, QuicStreamCode::StreamType); return; }; - if !flags.is_discard_signal() { - handle_stream_error( - &mut send, - &connection, - peer_addr, - mqtt5_protocol::QuicStreamCode::NotFlowOwner, - flags.err_tolerance, - ); + 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; } + if let Some(state) = flow_registry.lock().await.remove(flow_id) { + debug!(flow_id = ?flow_id, "discarded flow state for {peer_addr}: {state:?}"); + } + + 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() { - let mut registry = flow_registry.lock().await; - if let Some(state) = registry.remove(flow_id) { - debug!( - flow_id = ?flow_id, - "discarded flow state for {peer_addr}: {state:?}" - ); - } else { - debug!( - flow_id = ?flow_id, - "no flow state to discard for {peer_addr} (already expired or unknown)" - ); - } + debug!(flow_id = ?flow_id, "client handler ended before the flow was discarded"); } if let Err(e) = send.finish() { @@ -1022,7 +996,14 @@ fn spawn_data_stream_reader( return; }; let (flow_id, mut buffer) = match try_read_flow_header(&mut recv).await { - Ok(result) => { + Ok(StreamOpening::Refused(code)) => { + debug!(%code, "refusing data stream from {peer_addr}"); + if let Err(e) = recv.stop(quinn::VarInt::from_u32(code.code())) { + debug!("data stream from {peer_addr} already closed before stop: {e}"); + } + return; + } + Ok(StreamOpening::Flow(result)) => { let flow_id = if let (Some(id), Some(flags)) = (result.flow_id, result.flags) { let state = FlowState::new_client_data(id, flags, result.expire); let mut registry = flow_registry.lock().await; @@ -1058,10 +1039,14 @@ fn spawn_data_stream_reader( break; } } - Err(MqttError::ClientClosed) => { + Err(MqttError::ClientClosed) if buffer.is_empty() => { debug!(flow_id = ?flow_id, "QUIC data stream closed from {}", peer_addr); break; } + Err(MqttError::ClientClosed) => { + warn!(flow_id = ?flow_id, "QUIC data stream from {peer_addr} ended inside a packet"); + break; + } Err(MqttError::ConnectionError(reason)) => { debug!(flow_id = ?flow_id, "QUIC data stream from {peer_addr} ended: {reason}"); break; diff --git a/crates/mqtt5/tests/quic_flow_streams.rs b/crates/mqtt5/tests/quic_flow_streams.rs new file mode 100644 index 00000000..c5e3eb8f --- /dev/null +++ b/crates/mqtt5/tests/quic_flow_streams.rs @@ -0,0 +1,424 @@ +#![cfg(feature = "broker")] +#![cfg(feature = "transport-quic")] + +use bytes::BytesMut; +use mqtt5::broker::config::{ + BrokerConfig, QuicConfig as BrokerQuicConfig, StorageBackend, StorageConfig, +}; +use mqtt5::broker::{BrokerShutdownHandle, MqttBroker}; +use mqtt5::error::Result; +use mqtt5::transport::flow::{ControlFlowHeader, DataFlowHeader, FlowFlags, FlowId}; +use mqtt5::transport::{QuicConfig, QuicSplitResult, QuicTransport}; +use mqtt5::Transport; +use mqtt5_protocol::QuicStreamCode; +use std::net::SocketAddr; +use std::path::PathBuf; +use std::sync::Once; +use std::time::Duration; +use tokio::task::JoinHandle; + +const CERT: &str = "../../test_certs/server.pem"; +const KEY: &str = "../../test_certs/server.key"; +const PINGREQ: [u8; 2] = [0xC0, 0x00]; +const USER_DEFINED_FLOWS: [&[u8]; 2] = [&[0x14, 0x02, b'h', b'i'], &[0x14, 0x61, 0x70, 0x70]]; + +#[derive(Debug, PartialEq)] +struct Outcome { + session_alive: bool, + close_code: Option, +} + +const STILL_SERVED: Outcome = Outcome { + session_alive: true, + close_code: None, +}; + +static CRYPTO_PROVIDER: Once = Once::new(); + +struct Broker { + shutdown: BrokerShutdownHandle, + run: JoinHandle>, +} + +impl Broker { + async fn start() -> (Self, SocketAddr) { + CRYPTO_PROVIDER.call_once(|| { + rustls::crypto::ring::default_provider() + .install_default() + .expect("install crypto provider"); + }); + let config = BrokerConfig::default() + .with_storage(StorageConfig::default().with_backend(StorageBackend::Memory)) + .with_bind_address(([127, 0, 0, 1], 0)) + .with_quic( + BrokerQuicConfig::new(PathBuf::from(CERT), PathBuf::from(KEY)) + .with_bind_address("127.0.0.1:0".parse::().expect("bind address")), + ); + let mut broker = MqttBroker::with_config(config) + .await + .expect("broker starts"); + let addr = broker.quic_local_addr().expect("QUIC endpoint bound"); + let shutdown = broker.shutdown_handle(); + let run = tokio::spawn(async move { broker.run().await }); + (Self { shutdown, run }, addr) + } + + async fn stop(self) { + self.shutdown.shutdown(); + tokio::time::timeout(Duration::from_secs(5), self.run) + .await + .expect("broker stops") + .expect("broker task joins") + .expect("broker run succeeds"); + } +} + +fn connect_frame(client_id: &str) -> Vec { + let mut body = vec![ + 0x00, 0x04, b'M', b'Q', b'T', b'T', 0x05, 0x02, 0x00, 0x3C, 0x00, + ]; + body.extend_from_slice( + &u16::try_from(client_id.len()) + .expect("client id fits u16") + .to_be_bytes(), + ); + body.extend_from_slice(client_id.as_bytes()); + let mut frame = vec![0x10, u8::try_from(body.len()).expect("short CONNECT")]; + frame.extend_from_slice(&body); + frame +} + +async fn connected_session(client_id: &str) -> (Broker, QuicSplitResult) { + let (broker, addr) = Broker::start().await; + let mut transport = + QuicTransport::new(QuicConfig::new(addr, "localhost").with_verify_server_cert(false)); + transport.connect().await.expect("QUIC handshake"); + let mut split = transport.into_split().expect("split transport"); + split + .send + .write_all(&connect_frame(client_id)) + .await + .expect("send CONNECT"); + let mut header = [0u8; 2]; + split + .recv + .read_exact(&mut header) + .await + .expect("read CONNACK header"); + assert_eq!(header[0], 0x20); + let mut body = vec![0u8; usize::from(header[1])]; + split + .recv + .read_exact(&mut body) + .await + .expect("read CONNACK body"); + (broker, split) +} + +async fn outcome(split: &mut QuicSplitResult) -> Outcome { + let mut response = [0u8; 2]; + let session_alive = split.send.write_all(&PINGREQ).await.is_ok() + && matches!( + tokio::time::timeout(Duration::from_secs(2), split.recv.read_exact(&mut response)) + .await, + Ok(Ok(())) + ) + && response == [0xD0, 0x00]; + let close_code = tokio::time::timeout(Duration::from_secs(3), split.connection.closed()) + .await + .ok() + .and_then(|error| match error { + quinn::ConnectionError::ApplicationClosed(close) => Some(close.error_code.into_inner()), + _ => None, + }); + Outcome { + session_alive, + close_code, + } +} + +#[derive(Debug, PartialEq)] +enum StreamEnd { + Stopped(u64), + Refused(u64), + Finished, + Open, +} + +fn code(code: QuicStreamCode) -> u64 { + u64::from(code.code()) +} + +fn data_flow(sequence: u64, expire: u64, flags: FlowFlags) -> Vec { + let mut buf = BytesMut::new(); + DataFlowHeader::client(FlowId::client(sequence), expire, flags).encode(&mut buf); + buf.to_vec() +} + +fn with_pingreq(mut header: Vec) -> Vec { + header.extend_from_slice(&PINGREQ); + header +} + +fn server_flow() -> Vec { + let mut buf = BytesMut::new(); + DataFlowHeader::server(FlowId::server(1), 0, FlowFlags::default()).encode(&mut buf); + with_pingreq(buf.to_vec()) +} + +fn client_flow_with_server_id() -> Vec { + let mut buf = BytesMut::new(); + DataFlowHeader::client(FlowId::server(1), 0, FlowFlags::default()).encode(&mut buf); + with_pingreq(buf.to_vec()) +} + +fn second_control_flow() -> Vec { + let mut buf = BytesMut::new(); + ControlFlowHeader::new(FlowFlags::default()).encode(&mut buf); + with_pingreq(buf.to_vec()) +} + +async fn on_uni(bytes: &[u8]) -> (StreamEnd, Outcome) { + let (broker, mut split) = connected_session("uni-flow").await; + let mut stream = split.connection.open_uni().await.expect("open uni"); + stream.write_all(bytes).await.expect("write uni"); + let end = match tokio::time::timeout(Duration::from_secs(1), stream.stopped()).await { + Ok(Ok(Some(code))) => StreamEnd::Stopped(code.into_inner()), + Ok(_) => StreamEnd::Finished, + Err(_) => StreamEnd::Open, + }; + let result = outcome(&mut split).await; + broker.stop().await; + (end, result) +} + +async fn on_bidi(bytes: &[u8], finish: bool) -> (StreamEnd, Outcome) { + let (broker, mut split) = connected_session("bidi-flow").await; + let (mut send, mut recv) = split.connection.open_bi().await.expect("open bidi"); + send.write_all(bytes).await.expect("write bidi"); + if finish { + send.finish().expect("finish bidi"); + } + let end = match tokio::time::timeout(Duration::from_secs(1), recv.read_to_end(64)).await { + Ok(Ok(data)) if data.is_empty() => StreamEnd::Finished, + Ok(Err(quinn::ReadToEndError::Read(quinn::ReadError::Reset(code)))) => { + let stopped = tokio::time::timeout(Duration::from_secs(1), send.stopped()).await; + assert!( + matches!(stopped, Ok(Ok(Some(stop))) if stop == code), + "send side was not stopped with the reset code: {stopped:?}" + ); + StreamEnd::Refused(code.into_inner()) + } + Ok(other) => panic!("unexpected bidi stream end: {other:?}"), + Err(_) => StreamEnd::Open, + }; + let result = outcome(&mut split).await; + broker.stop().await; + (end, result) +} + +#[tokio::test] +async fn user_defined_flows_are_refused_without_closing_the_connection() { + for flow in USER_DEFINED_FLOWS { + assert_eq!( + on_uni(flow).await, + ( + StreamEnd::Stopped(code(QuicStreamCode::FlowRefused)), + STILL_SERVED + ) + ); + assert_eq!( + on_bidi(flow, false).await, + ( + StreamEnd::Refused(code(QuicStreamCode::FlowRefused)), + STILL_SERVED + ) + ); + } +} + +#[tokio::test] +async fn flows_only_the_server_may_open_are_refused_as_a_bad_flow_id() { + for flow in [server_flow(), client_flow_with_server_id()] { + assert_eq!( + on_uni(&flow).await, + ( + StreamEnd::Stopped(code(QuicStreamCode::BadFlowId)), + STILL_SERVED + ) + ); + assert_eq!( + on_bidi(&flow, false).await, + ( + StreamEnd::Refused(code(QuicStreamCode::BadFlowId)), + STILL_SERVED + ) + ); + } +} + +#[tokio::test] +async fn a_second_control_flow_is_refused_as_a_stream_type_error() { + let flow = second_control_flow(); + assert_eq!( + on_uni(&flow).await, + ( + StreamEnd::Stopped(code(QuicStreamCode::StreamType)), + STILL_SERVED + ) + ); + assert_eq!( + on_bidi(&flow, false).await, + ( + StreamEnd::Refused(code(QuicStreamCode::StreamType)), + STILL_SERVED + ) + ); +} + +#[tokio::test] +async fn a_bidirectional_stream_without_a_flow_header_is_refused() { + assert_eq!( + on_bidi(&PINGREQ, false).await, + ( + StreamEnd::Refused(code(QuicStreamCode::StreamType)), + STILL_SERVED + ) + ); +} + +#[tokio::test] +async fn a_unidirectional_client_data_flow_is_served() { + assert_eq!( + on_uni(&with_pingreq(data_flow(1, 0, FlowFlags::default()))).await, + (StreamEnd::Open, STILL_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!( + on_bidi(&data_flow(1, 0, FlowFlags::discard()), true).await, + (StreamEnd::Finished, STILL_SERVED) + ); +} + +#[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()); + flow.extend_from_slice(&PINGREQ[..1]); + let (broker, mut split) = connected_session("incomplete-flow").await; + let mut stream = split.connection.open_uni().await.expect("open uni"); + stream.write_all(&flow).await.expect("write uni"); + stream.finish().expect("finish uni"); + tokio::time::sleep(Duration::from_millis(200)).await; + assert_eq!(outcome(&mut split).await, STILL_SERVED); + broker.stop().await; +} + +fn subscribe_frame(topic: &str) -> Vec { + let mut body = vec![0x00, 0x01, 0x00]; + body.extend_from_slice( + &u16::try_from(topic.len()) + .expect("short topic") + .to_be_bytes(), + ); + body.extend_from_slice(topic.as_bytes()); + body.push(0x00); + let mut frame = vec![0x82, u8::try_from(body.len()).expect("short SUBSCRIBE")]; + frame.extend_from_slice(&body); + frame +} + +fn publish_frame(topic: &str, 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.push(0x00); + body.extend_from_slice(payload); + let mut frame = vec![0x30, u8::try_from(body.len()).expect("short PUBLISH")]; + frame.extend_from_slice(&body); + frame +} + +async fn delivered_to_aborted_flow(discard: bool) -> bool { + let (broker, mut split) = connected_session("discard-subscriptions").await; + let preserved = FlowFlags { + err_tolerance: 2, + persistent_subscriptions: 1, + ..FlowFlags::default() + }; + let mut flow = split.connection.open_uni().await.expect("open flow"); + let mut bytes = data_flow(5, 300, preserved); + bytes.extend_from_slice(&subscribe_frame("flows/discard")); + flow.write_all(&bytes).await.expect("subscribe on flow"); + let mut suback = [0u8; 5]; + split + .recv + .read_exact(&mut suback) + .await + .expect("read SUBACK"); + assert_eq!(suback, [0x90, 0x04, 0x00, 0x01, 0x00]); + flow.reset(quinn::VarInt::from_u32( + QuicStreamCode::FlowCancelled.code(), + )) + .expect("abort flow"); + tokio::time::sleep(Duration::from_millis(200)).await; + + if discard { + let (mut send, mut recv) = split.connection.open_bi().await.expect("open discard"); + send.write_all(&data_flow(5, 0, FlowFlags::discard())) + .await + .expect("send discard"); + send.finish().expect("finish discard"); + let response = tokio::time::timeout(Duration::from_secs(2), recv.read_to_end(16)) + .await + .expect("discard answered") + .expect("discard ends with FIN"); + assert!(response.is_empty()); + } + + split + .send + .write_all(&publish_frame("flows/discard", b"after")) + .await + .expect("publish"); + let delivered = tokio::time::timeout(Duration::from_secs(2), split.connection.accept_bi()) + .await + .is_ok(); + broker.stop().await; + delivered +} + +#[tokio::test] +async fn discarding_a_flow_removes_its_subscriptions() { + assert!(delivered_to_aborted_flow(false).await); + assert!(!delivered_to_aborted_flow(true).await); +}