diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 9cced944..14e7000a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -349,6 +349,8 @@ jobs: lcov2.info lcov3.info target/llvm-cov-target/cargo-llvm-cov.profdata + target/llvm-cov-target/cargo-llvm-cov2.profdata + target/integration-llvm-cov2.profdata target/llvm-cov-target/*.profraw retention-days: 7 diff --git a/Cargo.toml b/Cargo.toml index 8e13078e..e037db4f 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -103,6 +103,10 @@ clap = { version = "4.5", features = ["derive"] } name = "no_io_quic_async_client_example" required-features = ["quic"] +[[example]] +name = "no_io_quic_multi_stream_example" +required-features = ["quic"] + [[example]] name = "no_io_quic_client_example" required-features = ["quic"] diff --git a/docs/TODO.md b/docs/TODO.md index f892a7cf..37b9c117 100644 --- a/docs/TODO.md +++ b/docs/TODO.md @@ -23,6 +23,7 @@ - [ ] Client support MQTT v3 - [ ] UnixDomain socket - [ ] QUIC support (multi stream) +- [ ] FFI/language wrappers for QUIC source bind/rebind and local-address-change notification - [ ] Service discovery for distributed deployments - [ ] WebSocket transport support - [ ] Packet inspection utilities diff --git a/examples/no_io_quic_multi_stream_example.rs b/examples/no_io_quic_multi_stream_example.rs new file mode 100644 index 00000000..5f1bf190 --- /dev/null +++ b/examples/no_io_quic_multi_stream_example.rs @@ -0,0 +1,192 @@ +// SPDX-License-Identifier: MPL-2.0 + +//! Demonstrates MQTT-over-QUIC multi-stream support with the sans-I/O +//! `QuicMqttEngine`. +//! +//! The engine opens a primary *control* stream that carries only session-level +//! packets (CONNECT/CONNACK, PING, DISCONNECT, AUTH). This example then opens two +//! dedicated data streams once connected: +//! +//! - a **sub stream**: SUBSCRIBE is sent via `subscribe_on`; the SUBACK and any +//! delivered messages — plus the PUBACK we send for QoS 1 deliveries — flow on +//! that same stream. +//! - a **pub stream**: a PUBLISH is sent via `publish_on`; the QoS 1 PUBACK comes +//! back on that same stream. +//! +//! Responses never cross-fire between streams or onto the control stream. + +use flowsdk::mqtt_client::commands::{PublishCommand, SubscribeCommand}; +use flowsdk::mqtt_client::engine::{MqttEvent, QuicMqttEngine}; +use flowsdk::mqtt_client::opts::MqttClientOptions; +use std::future::Future; +use std::net::{ToSocketAddrs, UdpSocket}; +use std::pin::Pin; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::task::{Context, Poll}; +use std::time::{Duration, Instant}; + +struct ProtocolDriver { + engine: QuicMqttEngine, + socket: UdpSocket, + server_addr: std::net::SocketAddr, + last_tick: Instant, + /// Handles of the dedicated streams, opened once connected. + sub_stream: Option, + pub_stream: Option, + setup_done: bool, + running: Arc, + exit_when_rcvd: bool, +} + +impl Future for ProtocolDriver { + type Output = Result<(), Box>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + if !self.running.load(Ordering::SeqCst) { + return Poll::Ready(Ok(())); + } + + let now = Instant::now(); + let mut buf = [0u8; 2048]; + while let Ok((len, remote)) = self.socket.recv_from(&mut buf) { + if remote == self.server_addr { + self.engine + .handle_datagram(buf[..len].to_vec(), remote, now); + } + } + + if now.duration_since(self.last_tick) >= Duration::from_millis(10) { + let events = self.engine.handle_tick(now); + for event in events { + match event { + MqttEvent::Connected(_) => println!("Multi-stream: MQTT Connected!"), + MqttEvent::Published(res) => { + println!("Multi-stream: Message Published: ID={:?}", res.packet_id) + } + MqttEvent::Disconnected(reason) => { + println!("Multi-stream: MQTT Disconnected: {:?}", reason); + return Poll::Ready(Ok(())); + } + _ => {} + } + } + self.last_tick = now; + } + + let mut dags = self.engine.take_outgoing_datagrams(); + while let Some((dest, data)) = dags.pop_front() { + let _ = self.socket.send_to(&data, dest); + } + + // Once connected, open a dedicated sub stream and pub stream, then + // subscribe and publish on them respectively. + if self.engine.is_connected() && !self.setup_done { + let sub_stream = self.engine.open_data_stream()?; + let pub_stream = self.engine.open_data_stream()?; + self.sub_stream = Some(sub_stream); + self.pub_stream = Some(pub_stream); + println!( + "Multi-stream: opened sub stream {} and pub stream {} ({} data streams total)", + sub_stream, + pub_stream, + self.engine.data_stream_count() + ); + + // SUBSCRIBE on the sub stream — SUBACK arrives on the same stream. + if let Ok(sub_cmd) = SubscribeCommand::builder() + .add_topic("test/topic/multistream", 1) + .build() + { + let _ = self.engine.subscribe_on(sub_stream, sub_cmd); + } + + // PUBLISH on the pub stream — the QoS 1 PUBACK arrives on that stream. + if let Ok(pub_cmd) = PublishCommand::builder() + .topic("test/topic/multistream") + .payload("Hello from a dedicated QUIC pub stream!".to_string()) + .qos(1) + .build() + { + let _ = self.engine.publish_on(pub_stream, pub_cmd); + } + + self.setup_done = true; + if self.exit_when_rcvd { + self.running.store(false, Ordering::SeqCst); + } + } + + cx.waker().wake_by_ref(); + Poll::Pending + } +} + +/// Demonstrates the multi-stream `QuicMqttEngine` API against a broker. +pub fn run_example(exit_when_rcvd: bool) -> Result<(), Box> { + // Install the process-default rustls crypto provider. The OpenSSL provider is + // only available on the size-optimised `quic-openssl` builds; the mainstream + // `quic` build uses ring. + #[cfg(feature = "quic-proto-openssl")] + let _ = rustls_openssl::default_provider().install_default(); + #[cfg(not(feature = "quic-proto-openssl"))] + let _ = rustls::crypto::ring::default_provider().install_default(); + + let mqtt_opts = MqttClientOptions::builder() + .client_id("no-io-quic-multi-stream-client") + .peer("broker.emqx.io:14567") + .build(); + + let mut root_store = rustls::RootCertStore::empty(); + for cert in rustls_native_certs::load_native_certs()? { + root_store.add(cert).ok(); + } + let crypto = rustls::ClientConfig::builder() + .with_root_certificates(root_store) + .with_no_client_auth(); + + let socket = UdpSocket::bind("0.0.0.0:0")?; + socket.set_nonblocking(true)?; + let server_addr = ("broker.emqx.io", 14567u16) + .to_socket_addrs()? + .next() + .ok_or("DNS Failure")?; + + let mut engine = QuicMqttEngine::new(mqtt_opts)?; + engine.connect(server_addr, "broker.emqx.io", crypto, Instant::now())?; + + let running = Arc::new(AtomicBool::new(true)); + let r = running.clone(); + ctrlc::set_handler(move || { + r.store(false, Ordering::SeqCst); + })?; + + let driver = ProtocolDriver { + engine, + socket, + server_addr, + last_tick: Instant::now(), + sub_stream: None, + pub_stream: None, + setup_done: false, + running, + exit_when_rcvd, + }; + + futures::executor::block_on(driver) +} + +fn main() -> Result<(), Box> { + run_example(false) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + #[ignore = "requires network access to broker.emqx.io:14567"] + fn test_example() { + run_example(true).unwrap(); + } +} diff --git a/flowsdk_ffi/src/engine.rs b/flowsdk_ffi/src/engine.rs index 849b5d59..2073f741 100644 --- a/flowsdk_ffi/src/engine.rs +++ b/flowsdk_ffi/src/engine.rs @@ -81,7 +81,7 @@ impl MqttEngineFFI { pub fn handle_incoming(&self, data: Vec) -> Vec { let mut engine = self.engine.lock().unwrap(); let events = engine.handle_incoming(&data); - let mapped: Vec<_> = events.into_iter().map(map_event).collect(); + let mapped: Vec<_> = events.into_iter().filter_map(map_event).collect(); self.events.lock().unwrap().extend(mapped.iter().cloned()); mapped } @@ -90,7 +90,7 @@ impl MqttEngineFFI { let now = self.start_time + Duration::from_millis(now_ms); let mut engine = self.engine.lock().unwrap(); let events = engine.handle_tick(now); - let mapped: Vec<_> = events.into_iter().map(map_event).collect(); + let mapped: Vec<_> = events.into_iter().filter_map(map_event).collect(); self.events.lock().unwrap().extend(mapped.iter().cloned()); mapped } @@ -116,7 +116,7 @@ impl MqttEngineFFI { pub fn take_events(&self) -> Vec { let mut events = std::mem::take(&mut *self.events.lock().unwrap()); let engine_events = self.engine.lock().unwrap().take_events(); - events.extend(engine_events.into_iter().map(map_event)); + events.extend(engine_events.into_iter().filter_map(map_event)); events } @@ -183,43 +183,148 @@ impl MqttEngineFFI { } } -fn map_event(event: MqttEvent) -> MqttEventFFI { +fn map_event(event: MqttEvent) -> Option { match event { - MqttEvent::Connected(res) => MqttEventFFI::Connected(ConnectionResultFFI { + MqttEvent::Connected(res) => Some(MqttEventFFI::Connected(ConnectionResultFFI { reason_code: res.reason_code, session_present: res.session_present, - }), - MqttEvent::Disconnected(code) => MqttEventFFI::Disconnected { reason_code: code }, - MqttEvent::MessageReceived(msg) => MqttEventFFI::MessageReceived(MqttMessageFFI { + })), + MqttEvent::Disconnected(code) => Some(MqttEventFFI::Disconnected { reason_code: code }), + MqttEvent::PublishReceived { .. } | MqttEvent::PubRelReceived { .. } => None, + MqttEvent::MessageReceived(msg) => Some(MqttEventFFI::MessageReceived(MqttMessageFFI { topic: msg.topic_name, payload: msg.payload, qos: msg.qos, retain: msg.retain, - }), - MqttEvent::Published(res) => MqttEventFFI::Published(PublishResultFFI { + })), + MqttEvent::Published(res) => Some(MqttEventFFI::Published(PublishResultFFI { packet_id: res.packet_id, reason_code: res.reason_code, qos: res.qos, - }), - MqttEvent::Subscribed(res) => MqttEventFFI::Subscribed(SubscribeResultFFI { + })), + MqttEvent::Subscribed(res) => Some(MqttEventFFI::Subscribed(SubscribeResultFFI { packet_id: res.packet_id, reason_codes: res.reason_codes, - }), - MqttEvent::Unsubscribed(res) => MqttEventFFI::Unsubscribed(UnsubscribeResultFFI { + })), + MqttEvent::Unsubscribed(res) => Some(MqttEventFFI::Unsubscribed(UnsubscribeResultFFI { packet_id: res.packet_id, reason_codes: res.reason_codes, - }), - MqttEvent::PingResponse(res) => MqttEventFFI::PingResponse { + })), + MqttEvent::PingResponse(res) => Some(MqttEventFFI::PingResponse { success: res.success, - }, - MqttEvent::Error(err) => MqttEventFFI::Error { + }), + MqttEvent::Error(err) => Some(MqttEventFFI::Error { message: format!("{:?}", err), - }, - MqttEvent::ReconnectNeeded => MqttEventFFI::ReconnectNeeded, - MqttEvent::ReconnectScheduled { attempt, delay } => MqttEventFFI::ReconnectScheduled { - attempt, - delay_ms: delay.as_millis() as u64, - }, + }), + MqttEvent::TransportClosed { .. } => None, + MqttEvent::StreamClosed { + stream_id, + reason, + by_peer, + } => Some(MqttEventFFI::StreamClosed { + stream_id, + reason, + by_peer, + }), + MqttEvent::StreamReset { + stream_id, + error_code, + } => Some(MqttEventFFI::StreamReset { + stream_id, + error_code, + }), + MqttEvent::StreamStopped { + stream_id, + error_code, + } => Some(MqttEventFFI::StreamStopped { + stream_id, + error_code, + }), + MqttEvent::ZeroRttStatusChanged { .. } => None, + MqttEvent::ReconnectNeeded => Some(MqttEventFFI::ReconnectNeeded), + MqttEvent::ReconnectScheduled { attempt, delay } => { + Some(MqttEventFFI::ReconnectScheduled { + attempt, + delay_ms: delay.as_millis() as u64, + }) + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use flowsdk::mqtt_client::engine::QuicZeroRttStatus; + use std::time::Duration; + + #[test] + fn zero_rtt_status_event_is_not_reported_as_ffi_error() { + let event = MqttEvent::ZeroRttStatusChanged { + status: QuicZeroRttStatus::Attempted, + }; + + assert!(map_event(event).is_none()); + } + + #[test] + fn transport_closed_event_is_not_reported_as_ffi_error() { + let event = MqttEvent::TransportClosed { + reason: "connection closed".to_string(), + by_peer: true, + error_code: Some(0), + }; + + assert!(map_event(event).is_none()); + } + + #[test] + fn maps_stream_and_reconnect_events() { + assert!(matches!( + map_event(MqttEvent::StreamClosed { + stream_id: 7, + reason: "recv_finished".to_string(), + by_peer: true, + }), + Some(MqttEventFFI::StreamClosed { + stream_id: 7, + by_peer: true, + .. + }) + )); + assert!(matches!( + map_event(MqttEvent::StreamReset { + stream_id: 8, + error_code: 42, + }), + Some(MqttEventFFI::StreamReset { + stream_id: 8, + error_code: 42, + }) + )); + assert!(matches!( + map_event(MqttEvent::StreamStopped { + stream_id: 9, + error_code: 43, + }), + Some(MqttEventFFI::StreamStopped { + stream_id: 9, + error_code: 43, + }) + )); + assert!(matches!( + map_event(MqttEvent::ReconnectNeeded), + Some(MqttEventFFI::ReconnectNeeded) + )); + assert!(matches!( + map_event(MqttEvent::ReconnectScheduled { + attempt: 3, + delay: Duration::from_millis(250), + }), + Some(MqttEventFFI::ReconnectScheduled { + attempt: 3, + delay_ms: 250, + }) + )); } } @@ -338,7 +443,7 @@ impl TlsMqttEngineFFI { pub fn handle_tick(&self, now_ms: u64) -> Vec { let now = self.start_time + Duration::from_millis(now_ms); let events = self.engine.lock().unwrap().handle_tick(now); - let mapped: Vec<_> = events.into_iter().map(map_event).collect(); + let mapped: Vec<_> = events.into_iter().filter_map(map_event).collect(); self.events.lock().unwrap().extend(mapped.iter().cloned()); mapped } @@ -346,7 +451,7 @@ impl TlsMqttEngineFFI { pub fn take_events(&self) -> Vec { let mut events = std::mem::take(&mut *self.events.lock().unwrap()); let engine_events = self.engine.lock().unwrap().take_events(); - events.extend(engine_events.into_iter().map(map_event)); + events.extend(engine_events.into_iter().filter_map(map_event)); events } @@ -608,7 +713,7 @@ impl QuicMqttEngineFFI { let now = self.start_time + Duration::from_millis(now_ms); let mut engine = self.engine.lock().unwrap(); let events = engine.handle_tick(now); - let mapped: Vec<_> = events.into_iter().map(map_event).collect(); + let mapped: Vec<_> = events.into_iter().filter_map(map_event).collect(); self.events.lock().unwrap().extend(mapped.iter().cloned()); mapped } @@ -616,7 +721,7 @@ impl QuicMqttEngineFFI { pub fn take_events(&self) -> Vec { let mut events = std::mem::take(&mut *self.events.lock().unwrap()); let engine_events = self.engine.lock().unwrap().take_events(); - events.extend(engine_events.into_iter().map(map_event)); + events.extend(engine_events.into_iter().filter_map(map_event)); events } @@ -1624,6 +1729,9 @@ pub unsafe extern "C" fn mqtt_event_list_get_tag(ptr: *const MqttEventListFFI, i MqttEventFFI::Error { .. } => 8, MqttEventFFI::ReconnectNeeded => 9, MqttEventFFI::ReconnectScheduled { .. } => 10, + MqttEventFFI::StreamClosed { .. } => 11, + MqttEventFFI::StreamReset { .. } => 12, + MqttEventFFI::StreamStopped { .. } => 13, } } else { 0 @@ -1742,6 +1850,78 @@ pub unsafe extern "C" fn mqtt_event_list_get_error_message( std::ptr::null_mut() } +/// # Safety +/// +/// This function is unsafe because it dereferences a raw pointer to `MqttEventListFFI`. +#[no_mangle] +pub unsafe extern "C" fn mqtt_event_list_get_stream_id( + ptr: *const MqttEventListFFI, + index: usize, +) -> u64 { + if let Some(list) = ptr.as_ref() { + match list.events.get(index) { + Some(MqttEventFFI::StreamClosed { stream_id, .. }) + | Some(MqttEventFFI::StreamReset { stream_id, .. }) + | Some(MqttEventFFI::StreamStopped { stream_id, .. }) => *stream_id, + _ => 0, + } + } else { + 0 + } +} + +/// # Safety +/// +/// This function is unsafe because it dereferences a raw pointer to `MqttEventListFFI`. +#[no_mangle] +pub unsafe extern "C" fn mqtt_event_list_get_stream_error_code( + ptr: *const MqttEventListFFI, + index: usize, +) -> u64 { + if let Some(list) = ptr.as_ref() { + match list.events.get(index) { + Some(MqttEventFFI::StreamReset { error_code, .. }) + | Some(MqttEventFFI::StreamStopped { error_code, .. }) => *error_code, + _ => 0, + } + } else { + 0 + } +} + +/// # Safety +/// +/// This function is unsafe because it dereferences a raw pointer to `MqttEventListFFI` +/// and returns an allocated `c_char` pointer that must be freed using `mqtt_engine_free_string`. +#[no_mangle] +pub unsafe extern "C" fn mqtt_event_list_get_stream_close_reason( + ptr: *const MqttEventListFFI, + index: usize, +) -> *mut c_char { + if let Some(list) = ptr.as_ref() { + if let Some(MqttEventFFI::StreamClosed { reason, .. }) = list.events.get(index) { + return CString::new(reason.clone()).unwrap().into_raw(); + } + } + std::ptr::null_mut() +} + +/// # Safety +/// +/// This function is unsafe because it dereferences a raw pointer to `MqttEventListFFI`. +#[no_mangle] +pub unsafe extern "C" fn mqtt_event_list_get_stream_closed_by_peer( + ptr: *const MqttEventListFFI, + index: usize, +) -> i32 { + if let Some(list) = ptr.as_ref() { + if let Some(MqttEventFFI::StreamClosed { by_peer, .. }) = list.events.get(index) { + return i32::from(*by_peer); + } + } + -1 +} + /// # Safety /// /// This function is unsafe because it dereferences a raw pointer to `QuicMqttEngineFFI` diff --git a/flowsdk_ffi/src/engine/ffi_types.rs b/flowsdk_ffi/src/engine/ffi_types.rs index 6e1147e1..49d9b1cb 100644 --- a/flowsdk_ffi/src/engine/ffi_types.rs +++ b/flowsdk_ffi/src/engine/ffi_types.rs @@ -42,15 +42,37 @@ pub struct UnsubscribeResultFFI { #[derive(Clone)] pub enum MqttEventFFI { Connected(ConnectionResultFFI), - Disconnected { reason_code: Option }, + Disconnected { + reason_code: Option, + }, MessageReceived(MqttMessageFFI), Published(PublishResultFFI), Subscribed(SubscribeResultFFI), Unsubscribed(UnsubscribeResultFFI), - PingResponse { success: bool }, - Error { message: String }, + PingResponse { + success: bool, + }, + Error { + message: String, + }, + StreamClosed { + stream_id: u64, + reason: String, + by_peer: bool, + }, + StreamReset { + stream_id: u64, + error_code: u64, + }, + StreamStopped { + stream_id: u64, + error_code: u64, + }, ReconnectNeeded, - ReconnectScheduled { attempt: u32, delay_ms: u64 }, + ReconnectScheduled { + attempt: u32, + delay_ms: u64, + }, } #[cfg_attr(feature = "uniffi-bindings", derive(uniffi::Record))] diff --git a/src/mqtt_client/engine.rs b/src/mqtt_client/engine.rs index e6a1fbbf..c50218d3 100644 --- a/src/mqtt_client/engine.rs +++ b/src/mqtt_client/engine.rs @@ -1,11 +1,28 @@ // SPDX-License-Identifier: MPL-2.0 #[cfg(feature = "quic-proto")] -use quinn_proto::{ClientConfig, Connection, ConnectionHandle, Endpoint, EndpointConfig, StreamId}; -use std::collections::VecDeque; +use quinn_proto::{ + ClientConfig, Connection, ConnectionError, ConnectionHandle, Dir, Endpoint, EndpointConfig, + StreamId, VarInt, +}; +#[cfg(feature = "quic-proto")] +use rustls::{ + client::{ + ClientSessionMemoryCache, ClientSessionStore, Resumption, Tls12ClientSessionValue, + Tls13ClientSessionValue, + }, + pki_types::ServerName, + NamedGroup, +}; #[cfg(feature = "quic-proto")] -use std::sync::Arc; +use std::collections::HashMap; +use std::collections::VecDeque; use std::time::{Duration, Instant}; +#[cfg(feature = "quic-proto")] +use std::{ + fmt, + sync::{Arc, Mutex}, +}; use crate::mqtt_serde::control_packet::MqttPacket; use crate::mqtt_serde::mqttv3::{ @@ -35,6 +52,39 @@ use super::opts::MqttClientOptions; /// For MQTT v3.1.1 messages, the v5-specific fields (Properties) will be empty. pub type MqttMessage = MqttPublish; +/// Configuration for QUIC 0-RTT/session-ticket support. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct QuicZeroRttConfig { + /// Maximum number of TLS session entries retained in memory. + pub session_cache_size: usize, + /// Replay exact early MQTT bytes as 1-RTT if the peer rejects 0-RTT. + pub replay_on_reject: bool, +} + +impl Default for QuicZeroRttConfig { + fn default() -> Self { + Self { + session_cache_size: 256, + replay_on_reject: true, + } + } +} + +/// Current QUIC 0-RTT state for [`QuicMqttEngine`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)] +pub enum QuicZeroRttStatus { + /// 0-RTT is not enabled for this engine/connection. + Disabled, + /// 0-RTT was enabled but no resumable session ticket was available. + Unavailable, + /// 0-RTT keys were available and early data is being attempted. + Attempted, + /// The peer accepted early data. + Accepted, + /// The peer rejected early data. + Rejected, +} + /// Events emitted by the MqttEngine to be handled by the application (I/O layer) #[derive(Debug, serde::Serialize)] pub enum MqttEvent { @@ -43,9 +93,55 @@ pub enum MqttEvent { Published(PublishResult), Subscribed(SubscribeResult), Unsubscribed(UnsubscribeResult), + /// Incoming PUBLISH metadata. + /// + /// Emitted immediately before `MessageReceived`. `stream` identifies the + /// logical channel when the packet arrived on a QUIC data stream. + PublishReceived { + packet_id: Option, + stream: Option, + }, MessageReceived(MqttMessage), + /// Incoming QoS 2 PUBREL received. + /// + /// Exposed so manual-ack clients can wait for PUBREL before sending PUBCOMP. + /// `stream` identifies the logical channel when the packet arrived on a QUIC + /// data stream. + PubRelReceived { + packet_id: u16, + stream: Option, + }, PingResponse(PingResult), Error(MqttClientError), + /// The underlying QUIC transport connection was lost or closed, with details + /// from quinn-proto (peer close, idle timeout, local close, reset, etc.). + /// + /// `by_peer` is true if the peer initiated the close; `error_code` carries the + /// QUIC application/transport error code when one was signalled. + TransportClosed { + reason: String, + by_peer: bool, + error_code: Option, + }, + /// A QUIC stream was closed gracefully. + /// + /// `by_peer` is true when the peer finished its send side. It is false when + /// our finished send side was acknowledged by the peer. + StreamClosed { + stream_id: u64, + reason: String, + by_peer: bool, + }, + /// The peer aborted its send side with RESET_STREAM. + StreamReset { + stream_id: u64, + error_code: u64, + }, + /// The peer aborted our send side with STOP_SENDING. + StreamStopped { + stream_id: u64, + error_code: u64, + }, /// Signal that a reconnection is needed (e.g. after keep-alive timeout) ReconnectNeeded, /// Reconnection scheduled with exponential backoff @@ -53,6 +149,10 @@ pub enum MqttEvent { attempt: u32, delay: Duration, }, + /// QUIC 0-RTT status changed for the Sans-I/O QUIC engine. + ZeroRttStatusChanged { + status: QuicZeroRttStatus, + }, } /// A "Sans-I/O" MQTTv3.1.1/v5.0 protocol engine. @@ -99,6 +199,13 @@ pub struct MqttEngine { parser: MqttParser, outgoing_buffer: VecDeque>, + // Retransmissions (MQTT v3) for packets originally sent on a specific logical + // channel (e.g. a QUIC data stream). Drained by the transport via + // `take_stream_retransmissions` and re-sent on that same channel so the QoS + // 1/2 handshake never crosses streams. Untagged retransmissions still use + // `outgoing_buffer`. + stream_retransmissions: VecDeque<(u64, Vec)>, + // Pending operations tracking (state only) inflight_queue: InflightQueue, @@ -143,6 +250,7 @@ impl MqttEngine { last_packet_received: Instant::now(), parser, outgoing_buffer: VecDeque::new(), + stream_retransmissions: VecDeque::new(), events: Vec::new(), reconnect_attempts: 0, next_reconnect_at: None, @@ -173,6 +281,28 @@ impl MqttEngine { self.is_connected = false; } + /// Reset all transport-tied state for a fresh underlying transport (e.g. a + /// QUIC reconnect), discarding any bytes queued for the *old* connection so + /// they cannot be replayed before the new CONNECT. + /// + /// Clears the outgoing byte buffer, pending per-stream retransmissions, and + /// the inbound parser (any partial packet from the old transport). The MQTT + /// session and the inflight queue are **retained** so QoS 1/2 messages can be + /// resumed if the server reports `session_present`. + pub fn reset_for_new_transport(&mut self) { + self.is_connected = false; + self.outgoing_buffer.clear(); + self.stream_retransmissions.clear(); + self.parser = MqttParser::new(self.options.parser_buffer_size, self.options.mqtt_version); + // Cancel any pending reconnect deadline so it cannot fire another + // ReconnectNeeded while a new transport handshake is in progress. The + // attempt counter is left intact until CONNACK resets it. + self.next_reconnect_at = None; + let now = Instant::now(); + self.last_packet_sent = now; + self.last_packet_received = now; + } + /// Schedule the next reconnection attempt using exponential backoff. /// /// Logic: `delay = min(base * 2^attempts, max)`. @@ -244,8 +374,13 @@ impl MqttEngine { match self.parser.next_packet() { Ok(Some(packet)) => { self.last_packet_received = Instant::now(); - let packet_events = self.handle_packet(packet); + let (packet_events, responses) = self.handle_packet(packet, None); self.events.extend(packet_events); + // Single-stream transports (TCP/TLS, QUIC control stream) send + // responses back over the same shared outgoing buffer. + for response in responses { + let _ = self.enqueue_packet(response); + } } Ok(None) => break, Err(e) => { @@ -285,8 +420,11 @@ impl MqttEngine { let keep_alive = Duration::from_secs(self.options.keep_alive as u64); - // 1. Keep-alive: Send PING if needed - if keep_alive > Duration::ZERO && now.duration_since(self.last_packet_sent) >= keep_alive { + // 1. Keep-alive: Send PING if needed (unless automatic keep-alive disabled) + if self.options.auto_keepalive + && keep_alive > Duration::ZERO + && now.duration_since(self.last_packet_sent) >= keep_alive + { self.send_ping(); self.last_packet_sent = now; } @@ -313,11 +451,16 @@ impl MqttEngine { fn handle_retransmissions(&mut self, now: Instant) -> Vec { let events = Vec::new(); - let expired = self.inflight_queue.get_expired(now); - for mut packet in expired { + let expired = self.inflight_queue.get_expired_with_stream(now); + for (mut packet, stream) in expired { packet.set_dup(true); if let Ok(bytes) = packet.to_bytes() { - self.outgoing_buffer.push_back(bytes); + match stream { + // Retransmit on the originating channel (e.g. QUIC data stream) + // to preserve the same-stream QoS handshake. + Some(stream_id) => self.stream_retransmissions.push_back((stream_id, bytes)), + None => self.outgoing_buffer.push_back(bytes), + } self.last_packet_sent = now; } } @@ -325,6 +468,14 @@ impl MqttEngine { events } + /// Drain retransmissions that must be re-sent on a specific logical channel. + /// + /// Each entry is `(stream_handle, encoded_bytes)`. The transport routes each + /// onto the named channel. Empty for single-stream transports. + pub fn take_stream_retransmissions(&mut self) -> VecDeque<(u64, Vec)> { + std::mem::take(&mut self.stream_retransmissions) + } + /// Returns the exact timestamp of the next required wake-up. /// /// The runtime loop should sleep until this timestamp to avoid busy-waiting. @@ -345,8 +496,8 @@ impl MqttEngine { let mut next = None; let keep_alive = Duration::from_secs(self.options.keep_alive as u64); - // 2. Keep-alive timer (send PING) - if keep_alive > Duration::ZERO { + // 2. Keep-alive timer (send PING) — only when automatic keep-alive is on + if self.options.auto_keepalive && keep_alive > Duration::ZERO { let ping_deadline = self.last_packet_sent + keep_alive; next = Some(ping_deadline); } @@ -459,6 +610,105 @@ impl MqttEngine { Ok(pid) } + /// Encode a PUBLISH packet immediately and return its bytes, bypassing the + /// internal priority queue / outgoing buffer. + /// + /// This is used by transports that need to route individual packets to a + /// specific destination (e.g. a dedicated QUIC stream for MQTT-over-QUIC + /// multi-stream support). Protocol state shared across streams — packet-id + /// allocation and the inflight queue for QoS 1/2 — is still managed here, so + /// acknowledgements continue to flow through the normal engine machinery. + /// + /// Returns the assigned packet id (for QoS > 0) together with the encoded + /// bytes. The caller is responsible for actually transmitting the bytes. + /// + /// `stream` records the logical channel the packet is sent on so that QoS 1/2 + /// retransmissions are routed back onto the same channel (see + /// [`take_stream_retransmissions`](Self::take_stream_retransmissions)). + pub fn publish_encoded( + &mut self, + mut command: PublishCommand, + stream: Option, + ) -> Result<(Option, Vec), MqttClientError> { + let pid = if command.qos > 0 { + if let Some(pid) = command.packet_id { + Some(pid) + } else { + let pid = self.next_packet_id()?; + command.packet_id = Some(pid); + Some(pid) + } + } else { + None + }; + + let packet = if self.options.mqtt_version == 5 { + MqttPacket::Publish5(command.to_mqtt_publish()) + } else { + MqttPacket::Publish3(command.to_mqttv3_publish()) + }; + + // Encode before registering inflight so a serialization failure leaves + // no dangling inflight entry. + let bytes = packet.to_bytes().map_err(MqttClientError::from)?; + + // QoS > 0 messages must be tracked for retransmission/acknowledgement. + // `push_with_stream` enforces the receive-maximum limit and errors if exceeded. + if command.qos > 0 { + self.inflight_queue.push_with_stream( + pid.expect("QoS > 0 always assigns a packet id"), + packet, + command.qos, + stream, + )?; + } + + self.last_packet_sent = Instant::now(); + Ok((pid, bytes)) + } + + /// Process a single already-parsed MQTT packet that arrived on a specific + /// logical channel (e.g. one QUIC stream). + /// + /// This mirrors the per-packet handling done inside [`handle_incoming`], but + /// takes a decoded [`MqttPacket`] and, instead of enqueueing any protocol + /// responses onto the shared outgoing buffer, **returns the encoded response + /// bytes** so the caller can write them back on the *same* channel the packet + /// arrived on. This is what guarantees no cross-stream firing for MQTT over + /// QUIC: a PUBACK/PUBREC/PUBREL/PUBCOMP is always sent on the stream that + /// carried the packet it answers. + /// + /// `stream` identifies the channel the packet arrived on, used to tag any + /// inflight entry (QoS 2 PUBREL) created while handling it. + /// + /// Returns `(events, response_bytes)`. + pub fn ingest_stream_packet( + &mut self, + packet: MqttPacket, + stream: u64, + ) -> (Vec, Vec) { + self.last_packet_received = Instant::now(); + let (packet_events, responses) = self.handle_packet(packet, Some(stream)); + self.events.extend(packet_events); + + let mut response_bytes = Vec::new(); + for response in responses { + match response.to_bytes() { + Ok(bytes) => response_bytes.extend(bytes), + Err(e) => self.events.push(MqttEvent::Error(MqttClientError::from(e))), + } + } + if !response_bytes.is_empty() { + self.last_packet_sent = Instant::now(); + } + + // Acknowledgements may have freed inflight capacity; flush any control + // packets waiting in the priority queue (these target the control path). + self.process_queue(); + + (self.take_events(), response_bytes) + } + /// Queue a SUBSCRIBE packet. /// /// Be aware that this might fail immediately with `MqttClientError::BufferFull` @@ -519,22 +769,117 @@ impl MqttEngine { Ok(pid) } - /// Queue a DISCONNECT packet and update state to disconnected. - pub fn disconnect(&mut self) { - if !self.is_connected { - return; - } + /// Encode a SUBSCRIBE packet immediately and return its bytes, bypassing the + /// shared outgoing buffer. + /// + /// Like [`publish_encoded`](Self::publish_encoded), this is used by the QUIC + /// transport to route a SUBSCRIBE onto a dedicated data stream so the SUBACK + /// (and subsequently delivered messages) flow on that same stream. Packet-id + /// allocation and inflight tracking remain centralized. + pub fn subscribe_encoded( + &mut self, + mut command: SubscribeCommand, + stream: Option, + ) -> Result<(u16, Vec), MqttClientError> { + let pid = if let Some(pid) = command.packet_id { + pid + } else { + let pid = self.next_packet_id()?; + command.packet_id = Some(pid); + pid + }; + + let packet = if self.options.mqtt_version == 5 { + MqttPacket::Subscribe5(subscribev5::MqttSubscribe::new( + pid, + command.subscriptions, + command.properties, + )) + } else { + let v3_subs = command + .subscriptions + .into_iter() + .map(|s| subscribev3::SubscriptionTopic { + topic_filter: s.topic_filter, + qos: s.qos, + }) + .collect(); + MqttPacket::Subscribe3(subscribev3::MqttSubscribe::new(pid, v3_subs)) + }; + + let bytes = packet.to_bytes().map_err(MqttClientError::from)?; + self.inflight_queue + .push_with_stream(pid, packet, 1, stream)?; + self.last_packet_sent = Instant::now(); + Ok((pid, bytes)) + } + + /// Encode an UNSUBSCRIBE packet immediately and return its bytes, bypassing + /// the shared outgoing buffer. See [`subscribe_encoded`](Self::subscribe_encoded). + pub fn unsubscribe_encoded( + &mut self, + mut command: UnsubscribeCommand, + stream: Option, + ) -> Result<(u16, Vec), MqttClientError> { + let pid = if let Some(pid) = command.packet_id { + pid + } else { + let pid = self.next_packet_id()?; + command.packet_id = Some(pid); + pid + }; let packet = if self.options.mqtt_version == 5 { + MqttPacket::Unsubscribe5(unsubscribev5::MqttUnsubscribe::new( + pid, + command.topics.clone(), + command.properties, + )) + } else { + MqttPacket::Unsubscribe3(unsubscribev3::MqttUnsubscribe::new(pid, command.topics)) + }; + + let bytes = packet.to_bytes().map_err(MqttClientError::from)?; + self.inflight_queue + .push_with_stream(pid, packet, 1, stream)?; + self.last_packet_sent = Instant::now(); + Ok((pid, bytes)) + } + + fn disconnect_packet(&self) -> MqttPacket { + if self.options.mqtt_version == 5 { MqttPacket::Disconnect5(disconnectv5::MqttDisconnect::new(0, Vec::new())) } else { MqttPacket::Disconnect3(disconnectv3::MqttDisconnect::new()) - }; + } + } + /// Queue a DISCONNECT packet and update state to disconnected, ignoring a + /// full outgoing buffer (best-effort). Prefer [`try_disconnect`](Self::try_disconnect) + /// when the caller needs to know the DISCONNECT was actually queued. + pub fn disconnect(&mut self) { + if !self.is_connected { + return; + } + let packet = self.disconnect_packet(); let _ = self.enqueue_packet(packet); self.is_connected = false; } + /// Queue a DISCONNECT packet, propagating [`MqttClientError::BufferFull`] if + /// the outgoing buffer is full. The session is marked disconnected only once + /// the packet has actually been queued. A no-op (returns `Ok`) if already + /// disconnected. + pub fn try_disconnect(&mut self) -> Result<(), MqttClientError> { + if !self.is_connected { + return Ok(()); + } + let packet = self.disconnect_packet(); + self.enqueue_packet(packet)?; + self.is_connected = false; + Ok(()) + } + pub fn auth(&mut self, reason_code: u8, properties: Vec) { if self.options.mqtt_version == 5 { let auth = authv5::MqttAuth::new(reason_code, properties); @@ -544,8 +889,29 @@ impl MqttEngine { // --- Internal Helpers --- - fn handle_packet(&mut self, packet: MqttPacket) -> Vec { + /// Handle a single decoded incoming packet. + /// + /// Returns the application-facing events together with any **direct response + /// packets** that the protocol requires us to send in reply to *this* packet + /// (e.g. PUBACK for an incoming QoS 1 PUBLISH, PUBREC/PUBCOMP for QoS 2, + /// PUBREL after a PUBREC). These responses are returned rather than enqueued + /// so the caller can route them to the correct destination — for MQTT over + /// QUIC this means replying on the *same stream* the packet arrived on, + /// avoiding cross-stream firing. + /// + /// Session-level resends (session resumption after CONNACK) are not direct + /// responses and are still pushed onto the shared outgoing buffer. + /// + /// `stream` is the logical channel the packet arrived on; it is recorded on + /// any inflight entry created while handling it (the QoS 2 PUBREL), so that + /// retransmission stays on the same channel. + fn handle_packet( + &mut self, + packet: MqttPacket, + stream: Option, + ) -> (Vec, Vec) { let mut events = Vec::new(); + let mut responses: Vec = Vec::new(); match packet { MqttPacket::ConnAck5(ack) => { self.is_connected = ack.reason_code == 0; @@ -626,17 +992,19 @@ impl MqttEngine { MqttPacket::Publish5(p) => { let qos = p.qos; let pid = p.packet_id; + events.push(MqttEvent::PublishReceived { + packet_id: pid, + stream, + }); events.push(MqttEvent::MessageReceived(p)); - if qos == 1 { + if self.options.auto_ack && qos == 1 { if let Some(pid) = pid { - let ack = MqttPacket::PubAck5(MqttPubAck::new(pid, 0, Vec::new())); - let _ = self.enqueue_packet(ack); + responses.push(MqttPacket::PubAck5(MqttPubAck::new(pid, 0, Vec::new()))); } - } else if qos == 2 { + } else if self.options.auto_ack && qos == 2 { if let Some(pid) = pid { - let rec = MqttPacket::PubRec5(MqttPubRec::new(pid, 0, Vec::new())); - let _ = self.enqueue_packet(rec); + responses.push(MqttPacket::PubRec5(MqttPubRec::new(pid, 0, Vec::new()))); } } } @@ -653,21 +1021,23 @@ impl MqttEngine { p.dup, Vec::new(), ); + events.push(MqttEvent::PublishReceived { + packet_id: pid, + stream, + }); events.push(MqttEvent::MessageReceived(p5)); - if qos == 1 { + if self.options.auto_ack && qos == 1 { if let Some(pid) = pid { - let ack = MqttPacket::PubAck3( + responses.push(MqttPacket::PubAck3( crate::mqtt_serde::mqttv3::puback::MqttPubAck::new(pid), - ); - let _ = self.enqueue_packet(ack); + )); } - } else if qos == 2 { + } else if self.options.auto_ack && qos == 2 { if let Some(pid) = pid { - let rec = MqttPacket::PubRec3( + responses.push(MqttPacket::PubRec3( crate::mqtt_serde::mqttv3::pubrec::MqttPubRec::new(pid), - ); - let _ = self.enqueue_packet(rec); + )); } } } @@ -712,8 +1082,13 @@ impl MqttEngine { // Update entry to store PUBREL for retransmission entry.packet = rel.clone(); entry.sent_at = Instant::now(); - let _ = self.inflight_queue.push(entry.packet_id, entry.packet, 2); - let _ = self.enqueue_packet(rel); + let _ = self.inflight_queue.push_with_stream( + entry.packet_id, + entry.packet, + 2, + stream, + ); + responses.push(rel); } } MqttPacket::PubRec3(rec) => { @@ -722,19 +1097,38 @@ impl MqttEngine { // Update entry to store PUBREL for retransmission entry.packet = rel.clone(); entry.sent_at = Instant::now(); - let _ = self.inflight_queue.push(entry.packet_id, entry.packet, 2); - let _ = self.enqueue_packet(rel); + let _ = self.inflight_queue.push_with_stream( + entry.packet_id, + entry.packet, + 2, + stream, + ); + responses.push(rel); } } MqttPacket::PubRel5(rel) => { - let comp = MqttPacket::PubComp5(MqttPubComp::new(rel.packet_id, 0, Vec::new())); - let _ = self.enqueue_packet(comp); + events.push(MqttEvent::PubRelReceived { + packet_id: rel.packet_id, + stream, + }); + if self.options.auto_ack { + responses.push(MqttPacket::PubComp5(MqttPubComp::new( + rel.packet_id, + 0, + Vec::new(), + ))); + } } MqttPacket::PubRel3(rel) => { - let comp = MqttPacket::PubComp3( - crate::mqtt_serde::mqttv3::pubcomp::MqttPubComp::new(rel.message_id), - ); - let _ = self.enqueue_packet(comp); + events.push(MqttEvent::PubRelReceived { + packet_id: rel.message_id, + stream, + }); + if self.options.auto_ack { + responses.push(MqttPacket::PubComp3( + crate::mqtt_serde::mqttv3::pubcomp::MqttPubComp::new(rel.message_id), + )); + } } MqttPacket::PubComp5(comp) => { if let Some(entry) = self.inflight_queue.acknowledge(comp.packet_id) { @@ -762,18 +1156,31 @@ impl MqttEngine { } _ => {} } - events + (events, responses) } - pub fn send_ping(&mut self) { - let packet = if self.options.mqtt_version == 5 { + fn pingreq_packet(&self) -> MqttPacket { + if self.options.mqtt_version == 5 { MqttPacket::PingReq5(pingreqv5::MqttPingReq::new()) } else { MqttPacket::PingReq3(pingreqv3::MqttPingReq::new()) - }; + } + } + + /// Queue a PINGREQ, ignoring a full outgoing buffer (best-effort, used by the + /// automatic keep-alive timer). + pub fn send_ping(&mut self) { + let packet = self.pingreq_packet(); let _ = self.enqueue_packet(packet); } + /// Queue a PINGREQ, propagating [`MqttClientError::BufferFull`] if the + /// outgoing buffer is full, so callers can confirm it was actually queued. + pub fn try_send_ping(&mut self) -> Result<(), MqttClientError> { + let packet = self.pingreq_packet(); + self.enqueue_packet(packet) + } + pub fn enqueue_packet(&mut self, packet: MqttPacket) -> Result<(), MqttClientError> { if self.outgoing_buffer.len() >= self.options.max_outgoing_packet_count { return Err(MqttClientError::BufferFull { @@ -849,10 +1256,217 @@ impl MqttEngine { } } +/// State for a single QUIC data stream used by [`QuicMqttEngine`]. +/// +/// Every QUIC stream is an independent, ordered byte stream that carries its own +/// sequence of MQTT packets. Because partial packets from different streams must +/// never be concatenated, each stream owns a dedicated [`MqttParser`] for framing +/// inbound bytes and a buffer for outbound bytes awaiting transmission. +#[cfg(feature = "quic-proto")] +struct QuicStream { + parser: MqttParser, + outgoing: Vec, + /// Number of logical packets currently buffered in `outgoing` (reset to 0 once + /// the buffer fully drains). Used to bound the buffer like the engine's + /// `max_outgoing_packet_count`, so a stalled stream cannot grow without limit. + pending_packets: usize, + /// Maximum number of buffered packets before [`MqttClientError::BufferFull`]. + max_packets: usize, + /// A FIN has been requested; the send side is finished once `outgoing` drains. + finishing: bool, + /// The send side has been finished or reset; further writes are rejected. + finished: bool, + /// The receive side has seen FIN or RESET; further reads are skipped. + recv_closed: bool, +} + +#[cfg(feature = "quic-proto")] +impl QuicStream { + fn new(parser_buffer_size: usize, mqtt_version: u8, max_packets: usize) -> Self { + Self { + parser: MqttParser::new(parser_buffer_size, mqtt_version), + outgoing: Vec::new(), + pending_packets: 0, + max_packets, + finishing: false, + finished: false, + recv_closed: false, + } + } + + fn with_outgoing( + parser_buffer_size: usize, + mqtt_version: u8, + max_packets: usize, + outgoing: Vec, + pending_packets: usize, + ) -> Self { + Self { + parser: MqttParser::new(parser_buffer_size, mqtt_version), + outgoing, + pending_packets, + max_packets, + finishing: false, + finished: false, + recv_closed: false, + } + } +} + +#[cfg(feature = "quic-proto")] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum EarlyStreamRole { + Control, + DefaultPub, + DefaultSub, + ExplicitData, +} + +#[cfg(feature = "quic-proto")] +#[derive(Debug, Clone)] +struct EarlyStreamJournal { + original_stream_id: StreamId, + role: EarlyStreamRole, + bytes: Vec, + packet_count: usize, +} + +#[cfg(feature = "quic-proto")] +struct ClearableQuicSessionCache { + size: usize, + inner: Mutex, +} + +#[cfg(feature = "quic-proto")] +impl ClearableQuicSessionCache { + fn new(size: usize) -> Self { + Self { + size, + inner: Mutex::new(ClientSessionMemoryCache::new(size)), + } + } + + fn size(&self) -> usize { + self.size + } + + fn clear(&self) { + let mut guard = self.inner.lock().unwrap(); + let _ = std::mem::replace(&mut *guard, ClientSessionMemoryCache::new(self.size)); + } +} + +#[cfg(feature = "quic-proto")] +impl fmt::Debug for ClearableQuicSessionCache { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("ClearableQuicSessionCache") + .field("size", &self.size) + .field("sessions", &"[redacted]") + .finish() + } +} + +#[cfg(feature = "quic-proto")] +impl ClientSessionStore for ClearableQuicSessionCache { + fn set_kx_hint(&self, server_name: ServerName<'static>, group: NamedGroup) { + self.inner.lock().unwrap().set_kx_hint(server_name, group); + } + + fn kx_hint(&self, server_name: &ServerName<'_>) -> Option { + self.inner.lock().unwrap().kx_hint(server_name) + } + + fn set_tls12_session(&self, server_name: ServerName<'static>, value: Tls12ClientSessionValue) { + self.inner + .lock() + .unwrap() + .set_tls12_session(server_name, value); + } + + fn tls12_session(&self, server_name: &ServerName<'_>) -> Option { + self.inner.lock().unwrap().tls12_session(server_name) + } + + fn remove_tls12_session(&self, server_name: &ServerName<'static>) { + self.inner.lock().unwrap().remove_tls12_session(server_name); + } + + fn insert_tls13_ticket( + &self, + server_name: ServerName<'static>, + value: Tls13ClientSessionValue, + ) { + self.inner + .lock() + .unwrap() + .insert_tls13_ticket(server_name, value); + } + + fn take_tls13_ticket( + &self, + server_name: &ServerName<'static>, + ) -> Option { + self.inner.lock().unwrap().take_tls13_ticket(server_name) + } +} + +/// Error returned when writing to a stream whose send side has been finished or +/// reset. +#[cfg(feature = "quic-proto")] +fn stream_finished_error() -> MqttClientError { + MqttClientError::InvalidState { + expected: "an open (non-finished) stream".to_string(), + actual: "stream finished or reset".to_string(), + } +} + +/// Convert an application error code into a QUIC [`VarInt`], rejecting values +/// that exceed the 62-bit QUIC varint range. +#[cfg(feature = "quic-proto")] +fn quic_error_code(error_code: u64) -> Result { + VarInt::from_u64(error_code).map_err(|_| MqttClientError::InvalidConfiguration { + field: "error_code".to_string(), + reason: "QUIC error code exceeds 2^62-1".to_string(), + }) +} + +/// Map a quinn-proto [`ConnectionError`] to `(by_peer, error_code)` for the +/// [`MqttEvent::TransportClosed`] event. +#[cfg(feature = "quic-proto")] +fn transport_close_details(reason: &ConnectionError) -> (bool, Option) { + match reason { + ConnectionError::ApplicationClosed(c) => (true, Some(u64::from(c.error_code))), + // Peer's QUIC stack aborted with a transport error code — preserve it. + ConnectionError::ConnectionClosed(c) => (true, Some(u64::from(c.error_code))), + ConnectionError::Reset => (true, None), + ConnectionError::VersionMismatch + | ConnectionError::TransportError(_) + | ConnectionError::TimedOut + | ConnectionError::LocallyClosed + | ConnectionError::CidsExhausted => (false, None), + } +} + /// A "Sans-I/O" MQTT over QUIC protocol engine. /// /// This engine combines the `MqttEngine` (MQTT state machine) with `quinn_proto` (QUIC state machine) /// to provide a complete MQTT-over-QUIC implementation that does not perform any direct I/O. +/// +/// # Multi-stream support +/// +/// MQTT-over-QUIC can spread traffic across multiple QUIC streams to avoid +/// head-of-line blocking. This engine opens a primary **control stream** for the +/// MQTT handshake and control packets, and supports additional **data streams**: +/// +/// - [`open_data_stream`](Self::open_data_stream) opens a new client-initiated +/// bidirectional stream and returns an opaque handle. +/// - [`publish_on`](Self::publish_on) routes a PUBLISH onto a specific data stream. +/// - Server-initiated bidirectional streams are accepted automatically and their +/// inbound packets are fed into the shared protocol state machine. +/// +/// MQTT protocol state (packet-id allocation, the inflight queue for QoS 1/2, +/// session, keep-alive) is shared across all streams; only the wire framing is +/// per-stream. #[cfg(feature = "quic-proto")] pub struct QuicMqttEngine { mqtt_engine: MqttEngine, @@ -860,13 +1474,50 @@ pub struct QuicMqttEngine { connection: Option, connection_handle: Option, - // The bidirectional stream used for MQTT control packets - mqtt_stream: Option, + // The primary bidirectional stream used for the MQTT handshake and control packets. + control_stream: Option, + + // Additional data streams (client- or server-initiated), keyed by raw QUIC stream id. + // Each carries its own framing parser and outgoing buffer. + data_streams: HashMap, + + // Lazily-opened default data streams backing the convenience `publish` / + // `subscribe` / `unsubscribe` methods, so that PUBLISH and SUBSCRIBE traffic + // never lands on the session control stream. + default_pub_stream: Option, + default_sub_stream: Option, + + // Raw bytes injected onto the control stream via the low-level escape hatch + // (`send_raw_on` / `send_packet_on`), flushed alongside the engine's session + // packets. Used for negative testing. + control_outgoing: Vec, + + // Connection parameters retained from the last `connect` so `reconnect` can + // re-establish on the same endpoint without the caller rebuilding the config. + client_config: Option, + server_addr: Option, + server_name: Option, + + // Optional QUIC 0-RTT/session-ticket policy and cache. The cache is retained + // across connections until explicitly cleared or replaced with another size. + zero_rtt_config: Option, + zero_rtt_status: QuicZeroRttStatus, + zero_rtt_cache: Option>, + pending_transport_events: VecDeque, + early_stream_journal: Vec, + + // A deferred graceful QUIC close, applied after the queued MQTT DISCONNECT has + // been flushed (see `disconnect_and_close`). + pending_close: Option<(VarInt, bytes::Bytes)>, + // A FIN requested for the control stream's send side, applied after its + // buffered bytes are flushed. + control_finishing: bool, + // The control stream's send side has been finished or reset; further writes + // (including raw injection) are rejected. + control_finished: bool, // Outgoing UDP datagrams to be sent by the application outgoing_datagrams: VecDeque<(std::net::SocketAddr, Vec)>, - // Internal buffer for reading from QUIC stream and feeding to MQTT engine - // stream_read_buffer removed } #[cfg(feature = "quic-proto")] @@ -889,7 +1540,22 @@ impl QuicMqttEngine { endpoint, connection: None, connection_handle: None, - mqtt_stream: None, + control_stream: None, + data_streams: HashMap::new(), + default_pub_stream: None, + default_sub_stream: None, + control_outgoing: Vec::new(), + client_config: None, + server_addr: None, + server_name: None, + zero_rtt_config: None, + zero_rtt_status: QuicZeroRttStatus::Disabled, + zero_rtt_cache: None, + pending_transport_events: VecDeque::new(), + early_stream_journal: Vec::new(), + pending_close: None, + control_finishing: false, + control_finished: false, outgoing_datagrams: VecDeque::new(), }) } @@ -906,17 +1572,113 @@ impl QuicMqttEngine { crypto_config.alpn_protocols = vec![b"mqtt".to_vec()]; } - // Wrap in quinn config - let mut client_config = ClientConfig::new(Arc::new( - quinn_proto::crypto::rustls::QuicClientConfig::try_from(crypto_config).map_err( - |e| MqttClientError::InternalError { - message: format!("Failed to create QUIC client config: {}", e), - }, - )?, - )); + let client_config = Self::client_config_from_crypto(crypto_config)?; - // Disable unreliable datagrams (buffer size 0 / None) - let mut transport = quinn_proto::TransportConfig::default(); + // Retain parameters so `reconnect` can re-establish without rebuilding. + self.client_config = Some(client_config); + self.server_addr = Some(server_addr); + self.server_name = Some(server_name.to_string()); + self.zero_rtt_config = None; + self.zero_rtt_status = QuicZeroRttStatus::Disabled; + self.pending_transport_events.clear(); + self.early_stream_journal.clear(); + + self.establish(now) + } + + /// Connect with QUIC 0-RTT/session-ticket support enabled. + /// + /// Early MQTT commands are only accepted while the returned status is + /// [`QuicZeroRttStatus::Attempted`]. If the peer rejects 0-RTT and + /// `replay_on_reject` is true, the engine replays the exact early bytes as + /// 1-RTT data on replacement streams. + pub fn connect_with_zero_rtt( + &mut self, + server_addr: std::net::SocketAddr, + server_name: &str, + mut crypto_config: rustls::ClientConfig, + zero_rtt_config: QuicZeroRttConfig, + now: Instant, + ) -> Result<(), MqttClientError> { + if crypto_config.alpn_protocols.is_empty() { + crypto_config.alpn_protocols = vec![b"mqtt".to_vec()]; + } + + crypto_config.enable_early_data = true; + let cache = self.zero_rtt_cache_for_size(zero_rtt_config.session_cache_size); + crypto_config.resumption = Resumption::store(cache); + + let client_config = Self::client_config_from_crypto(crypto_config)?; + + self.client_config = Some(client_config); + self.server_addr = Some(server_addr); + self.server_name = Some(server_name.to_string()); + self.zero_rtt_config = Some(zero_rtt_config); + self.zero_rtt_status = QuicZeroRttStatus::Unavailable; + self.pending_transport_events.clear(); + self.early_stream_journal.clear(); + + self.establish(now) + } + + /// Current QUIC 0-RTT state. + pub fn zero_rtt_status(&self) -> QuicZeroRttStatus { + self.zero_rtt_status + } + + /// Clear any retained QUIC TLS resumption/session tickets. + pub fn clear_quic_session_cache(&mut self) { + if let Some(cache) = &self.zero_rtt_cache { + cache.clear(); + } + } + + /// Re-establish the QUIC connection on the existing endpoint, reusing the + /// configuration captured by the previous [`connect`](Self::connect). + /// + /// Resets all stream state (the new connection starts fresh streams) and the + /// MQTT connected flag; the MQTT handshake is re-driven once the new QUIC + /// handshake completes. Returns an error if [`connect`](Self::connect) has not + /// been called yet. + pub fn reconnect(&mut self, now: Instant) -> Result<(), MqttClientError> { + if self.client_config.is_none() { + return Err(MqttClientError::InvalidState { + expected: "a prior successful connect()".to_string(), + actual: "reconnect called before connect".to_string(), + }); + } + self.reset_connection_state(); + self.establish(now) + } + + /// Notify quinn-proto that the embedding runtime has switched to a new local UDP address. + /// + /// The sans-I/O engine does not own sockets, so the caller is responsible for + /// rebinding/replacing the UDP socket used to send [`take_outgoing_datagrams`](Self::take_outgoing_datagrams). + /// This hook mirrors Quinn's async `Endpoint::rebind` behavior for active + /// connections: it rotates to a fresh remote connection ID when available and + /// queues a PING so the peer can validate the new path. + pub fn notify_local_address_changed(&mut self) -> Result<(), MqttClientError> { + if let Some(conn) = &mut self.connection { + conn.local_address_changed(); + } + Ok(()) + } + + fn client_config_from_crypto( + crypto_config: rustls::ClientConfig, + ) -> Result { + // Wrap in quinn config + let mut client_config = ClientConfig::new(Arc::new( + quinn_proto::crypto::rustls::QuicClientConfig::try_from(crypto_config).map_err( + |e| MqttClientError::InternalError { + message: format!("Failed to create QUIC client config: {}", e), + }, + )?, + )); + + // Disable unreliable datagrams (buffer size 0 / None) + let mut transport = quinn_proto::TransportConfig::default(); transport.datagram_receive_buffer_size(None); // Set max_idle_timeout to prevent QUIC from timing out before MQTT keepalive mechanism // Use 120 seconds to accommodate MQTT keepalive (typically 30-60s) with 2x multiplier for safety @@ -927,14 +1689,469 @@ impl QuicMqttEngine { })?; transport.max_idle_timeout(Some(idle_timeout)); client_config.transport_config(Arc::new(transport)); + Ok(client_config) + } + + fn zero_rtt_cache_for_size(&mut self, size: usize) -> Arc { + if let Some(cache) = &self.zero_rtt_cache { + if cache.size() == size { + return Arc::clone(cache); + } + } + + let cache = Arc::new(ClearableQuicSessionCache::new(size)); + self.zero_rtt_cache = Some(Arc::clone(&cache)); + cache + } + + fn journal_stream_open( + early_stream_journal: &mut Vec, + stream_id: StreamId, + role: EarlyStreamRole, + ) { + early_stream_journal.push(EarlyStreamJournal { + original_stream_id: stream_id, + role, + bytes: Vec::new(), + packet_count: 0, + }); + } + + fn journal_stream_bytes( + early_stream_journal: &mut [EarlyStreamJournal], + stream_id: StreamId, + bytes: &[u8], + packet_count_delta: usize, + ) { + if bytes.is_empty() && packet_count_delta == 0 { + return; + } + if let Some(entry) = early_stream_journal + .iter_mut() + .find(|entry| entry.original_stream_id == stream_id) + { + entry.bytes.extend_from_slice(bytes); + entry.packet_count += packet_count_delta; + } + } + + fn append_control_outgoing_bytes( + control_outgoing: &mut Vec, + early_stream_journal: &mut [EarlyStreamJournal], + zero_rtt_status: QuicZeroRttStatus, + control_stream: Option, + bytes: &[u8], + packet_count_delta: usize, + ) { + if bytes.is_empty() { + return; + } + control_outgoing.extend_from_slice(bytes); + if zero_rtt_status == QuicZeroRttStatus::Attempted { + if let Some(stream_id) = control_stream { + Self::journal_stream_bytes( + early_stream_journal, + stream_id, + bytes, + packet_count_delta, + ); + } + } + } + + fn mark_zero_rtt_unavailable(&mut self) { + self.zero_rtt_status = QuicZeroRttStatus::Unavailable; + self.pending_transport_events + .push_back(MqttEvent::ZeroRttStatusChanged { + status: QuicZeroRttStatus::Unavailable, + }); + } + + fn begin_zero_rtt_attempt(&mut self, control_stream: Option) { + let Some(stream_id) = control_stream else { + self.mark_zero_rtt_unavailable(); + return; + }; + + self.zero_rtt_status = QuicZeroRttStatus::Attempted; + self.pending_transport_events + .push_back(MqttEvent::ZeroRttStatusChanged { + status: QuicZeroRttStatus::Attempted, + }); + self.control_stream = Some(stream_id); + Self::journal_stream_open( + &mut self.early_stream_journal, + stream_id, + EarlyStreamRole::Control, + ); + self.mqtt_engine.connect(); + let connect_bytes = self.mqtt_engine.take_outgoing(); + Self::append_control_outgoing_bytes( + &mut self.control_outgoing, + &mut self.early_stream_journal, + self.zero_rtt_status, + self.control_stream, + &connect_bytes, + usize::from(!connect_bytes.is_empty()), + ); + } + + fn ensure_command_allowed_before_mqtt_connected(&self) -> Result<(), MqttClientError> { + if self.mqtt_engine.is_connected() || self.zero_rtt_status == QuicZeroRttStatus::Attempted { + return Ok(()); + } + + Err(MqttClientError::InvalidState { + expected: "a connected MQTT session or attempted QUIC 0-RTT".to_string(), + actual: format!( + "MQTT not connected; 0-RTT status {:?}", + self.zero_rtt_status + ), + }) + } + + #[allow(clippy::too_many_arguments)] + fn replay_early_stream_journals( + conn: &mut Connection, + journals: Vec, + parser_buffer_size: usize, + mqtt_version: u8, + max_packets: usize, + control_stream: &mut Option, + data_streams: &mut HashMap, + default_pub_stream: &mut Option, + default_sub_stream: &mut Option, + control_outgoing: &mut Vec, + events: &mut Vec, + ) { + for journal in journals { + let Some(stream_id) = conn.streams().open(Dir::Bi) else { + events.push(MqttEvent::Error(MqttClientError::InternalError { + message: "QUIC stream limit reached while replaying rejected 0-RTT data" + .to_string(), + })); + break; + }; + + if stream_id != journal.original_stream_id { + events.push(MqttEvent::Error(MqttClientError::InternalError { + message: format!( + "0-RTT replay stream id changed from {} to {}", + u64::from(journal.original_stream_id), + u64::from(stream_id) + ), + })); + } + + match journal.role { + EarlyStreamRole::Control => { + *control_stream = Some(stream_id); + control_outgoing.extend_from_slice(&journal.bytes); + } + EarlyStreamRole::DefaultPub + | EarlyStreamRole::DefaultSub + | EarlyStreamRole::ExplicitData => { + data_streams.insert( + stream_id, + QuicStream::with_outgoing( + parser_buffer_size, + mqtt_version, + max_packets, + journal.bytes, + journal.packet_count, + ), + ); + match journal.role { + EarlyStreamRole::DefaultPub => *default_pub_stream = Some(stream_id), + EarlyStreamRole::DefaultSub => *default_sub_stream = Some(stream_id), + EarlyStreamRole::ExplicitData | EarlyStreamRole::Control => {} + } + } + } + } + } + + /// Clear all per-connection state (streams, buffers, MQTT connected flag). + fn reset_connection_state(&mut self) { + self.connection = None; + self.connection_handle = None; + self.control_stream = None; + self.data_streams.clear(); + self.default_pub_stream = None; + self.default_sub_stream = None; + self.control_outgoing.clear(); + self.outgoing_datagrams.clear(); + self.pending_close = None; + self.control_finishing = false; + self.control_finished = false; + self.pending_transport_events.clear(); + self.early_stream_journal.clear(); + // Discard any MQTT bytes queued for the old transport so they are not + // replayed before the next CONNECT. + self.mqtt_engine.reset_for_new_transport(); + } + + fn clear_data_stream_refs_fields( + data_streams: &mut HashMap, + default_pub_stream: &mut Option, + default_sub_stream: &mut Option, + stream_id: StreamId, + ) { + data_streams.remove(&stream_id); + if *default_pub_stream == Some(stream_id) { + *default_pub_stream = None; + } + if *default_sub_stream == Some(stream_id) { + *default_sub_stream = None; + } + } + + fn handle_control_stream_terminal_fields( + control_stream: &mut Option, + control_outgoing: &mut Vec, + control_finishing: &mut bool, + control_finished: &mut bool, + mqtt_engine: &mut MqttEngine, + mqtt_events: &mut Vec, + ) { + *control_stream = None; + control_outgoing.clear(); + *control_finishing = false; + *control_finished = true; + mqtt_engine.handle_connection_lost(); + mqtt_events.push(MqttEvent::Disconnected(None)); + } + + #[cfg(test)] + fn handle_stream_stopped( + &mut self, + stream_id: StreamId, + error_code: VarInt, + mqtt_events: &mut Vec, + ) { + Self::handle_stream_stopped_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.default_pub_stream, + &mut self.default_sub_stream, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + stream_id, + error_code, + mqtt_events, + ); + } + + #[allow(clippy::too_many_arguments)] + fn handle_stream_stopped_fields( + control_stream: &mut Option, + data_streams: &mut HashMap, + default_pub_stream: &mut Option, + default_sub_stream: &mut Option, + control_outgoing: &mut Vec, + control_finishing: &mut bool, + control_finished: &mut bool, + mqtt_engine: &mut MqttEngine, + stream_id: StreamId, + error_code: VarInt, + mqtt_events: &mut Vec, + ) { + if Some(stream_id) == *control_stream { + mqtt_events.push(MqttEvent::StreamStopped { + stream_id: u64::from(stream_id), + error_code: u64::from(error_code), + }); + Self::handle_control_stream_terminal_fields( + control_stream, + control_outgoing, + control_finishing, + control_finished, + mqtt_engine, + mqtt_events, + ); + } else if data_streams.contains_key(&stream_id) { + mqtt_events.push(MqttEvent::StreamStopped { + stream_id: u64::from(stream_id), + error_code: u64::from(error_code), + }); + Self::clear_data_stream_refs_fields( + data_streams, + default_pub_stream, + default_sub_stream, + stream_id, + ); + } + } + + #[cfg(test)] + fn handle_stream_reset( + &mut self, + stream_id: StreamId, + error_code: VarInt, + mqtt_events: &mut Vec, + ) { + Self::handle_stream_reset_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.default_pub_stream, + &mut self.default_sub_stream, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + stream_id, + error_code, + mqtt_events, + ); + } + + #[allow(clippy::too_many_arguments)] + fn handle_stream_reset_fields( + control_stream: &mut Option, + data_streams: &mut HashMap, + default_pub_stream: &mut Option, + default_sub_stream: &mut Option, + control_outgoing: &mut Vec, + control_finishing: &mut bool, + control_finished: &mut bool, + mqtt_engine: &mut MqttEngine, + stream_id: StreamId, + error_code: VarInt, + mqtt_events: &mut Vec, + ) { + if Some(stream_id) == *control_stream { + mqtt_events.push(MqttEvent::StreamReset { + stream_id: u64::from(stream_id), + error_code: u64::from(error_code), + }); + Self::handle_control_stream_terminal_fields( + control_stream, + control_outgoing, + control_finishing, + control_finished, + mqtt_engine, + mqtt_events, + ); + } else if data_streams.contains_key(&stream_id) { + mqtt_events.push(MqttEvent::StreamReset { + stream_id: u64::from(stream_id), + error_code: u64::from(error_code), + }); + Self::clear_data_stream_refs_fields( + data_streams, + default_pub_stream, + default_sub_stream, + stream_id, + ); + } + } + + #[cfg(test)] + fn handle_stream_closed( + &mut self, + stream_id: StreamId, + reason: &'static str, + by_peer: bool, + mqtt_events: &mut Vec, + ) { + Self::handle_stream_closed_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + stream_id, + reason, + by_peer, + mqtt_events, + ); + } + + #[allow(clippy::too_many_arguments)] + fn handle_stream_closed_fields( + control_stream: &mut Option, + data_streams: &mut HashMap, + control_outgoing: &mut Vec, + control_finishing: &mut bool, + control_finished: &mut bool, + mqtt_engine: &mut MqttEngine, + stream_id: StreamId, + reason: &'static str, + by_peer: bool, + mqtt_events: &mut Vec, + ) { + if Some(stream_id) == *control_stream { + mqtt_events.push(MqttEvent::StreamClosed { + stream_id: u64::from(stream_id), + reason: reason.to_string(), + by_peer, + }); + if by_peer { + Self::handle_control_stream_terminal_fields( + control_stream, + control_outgoing, + control_finishing, + control_finished, + mqtt_engine, + mqtt_events, + ); + } + } else if let Some(ds) = data_streams.get_mut(&stream_id) { + mqtt_events.push(MqttEvent::StreamClosed { + stream_id: u64::from(stream_id), + reason: reason.to_string(), + by_peer, + }); + if by_peer { + ds.recv_closed = true; + } else { + ds.finished = true; + } + } + } - let (ch, conn) = self + /// Open a fresh QUIC connection using the retained configuration. + fn establish(&mut self, now: Instant) -> Result<(), MqttClientError> { + let client_config = + self.client_config + .clone() + .ok_or_else(|| MqttClientError::InvalidState { + expected: "stored client config".to_string(), + actual: "none".to_string(), + })?; + let server_addr = self + .server_addr + .ok_or_else(|| MqttClientError::InvalidState { + expected: "stored server address".to_string(), + actual: "none".to_string(), + })?; + let server_name = + self.server_name + .clone() + .ok_or_else(|| MqttClientError::InvalidState { + expected: "stored server name".to_string(), + actual: "none".to_string(), + })?; + + let (ch, mut conn) = self .endpoint - .connect(now, client_config, server_addr, server_name) + .connect(now, client_config, server_addr, &server_name) .map_err(|e| MqttClientError::InternalError { message: format!("Failed to create QUIC connection: {}", e), })?; + if self.zero_rtt_config.is_some() { + if conn.has_0rtt() { + let control_stream = conn.streams().open(Dir::Bi); + self.begin_zero_rtt_attempt(control_stream); + } else { + self.mark_zero_rtt_unavailable(); + } + } + self.connection = Some(conn); self.connection_handle = Some(ch); @@ -987,58 +2204,415 @@ impl QuicMqttEngine { /// Drive time-dependent logic for both QUIC and MQTT state machines. pub fn handle_tick(&mut self, now: Instant) -> Vec { - let mut mqtt_events = Vec::new(); + let mut mqtt_events: Vec = self.pending_transport_events.drain(..).collect(); - // 1. Drive QUIC Endpoint - handle_call removed in 0.11? - // self.endpoint.handle_call(now); + let parser_buffer_size = self.mqtt_engine.options().parser_buffer_size; + let mqtt_version = self.mqtt_engine.mqtt_version(); + let max_packets = self.mqtt_engine.options().max_outgoing_packet_count; - // 2. Drive QUIC Connection + // Drive QUIC Connection if let Some(conn) = &mut self.connection { conn.handle_timeout(now); - // Check for new streams or data + // 1. Drain connection events. while let Some(event) = conn.poll() { match event { - quinn_proto::Event::Stream(_stream_id) => { - // Stream event (readable/writable etc) - } - quinn_proto::Event::Connected - // QUIC Handshake done. Open a bidirectional stream for MQTT. - if self.mqtt_stream.is_none() => - { - if let Some(stream_id) = conn.streams().open(quinn_proto::Dir::Bi) { - self.mqtt_stream = Some(stream_id); - self.mqtt_engine.connect(); + quinn_proto::Event::Stream(stream_event) => match stream_event { + quinn_proto::StreamEvent::Finished { id } => { + Self::handle_stream_closed_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + id, + "send_finished", + false, + &mut mqtt_events, + ); + } + quinn_proto::StreamEvent::Stopped { id, error_code } => { + Self::handle_stream_stopped_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.default_pub_stream, + &mut self.default_sub_stream, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + id, + error_code, + &mut mqtt_events, + ); + } + quinn_proto::StreamEvent::Opened { .. } + | quinn_proto::StreamEvent::Readable { .. } + | quinn_proto::StreamEvent::Writable { .. } + | quinn_proto::StreamEvent::Available { .. } => { + // Readable/writable/open events are handled by polling + // the streams directly below. + } + }, + quinn_proto::Event::Connected => { + let zero_rtt_was_attempted = + self.zero_rtt_status == QuicZeroRttStatus::Attempted; + + if zero_rtt_was_attempted { + if conn.accepted_0rtt() { + self.zero_rtt_status = QuicZeroRttStatus::Accepted; + self.early_stream_journal.clear(); + mqtt_events.push(MqttEvent::ZeroRttStatusChanged { + status: QuicZeroRttStatus::Accepted, + }); + } else { + self.zero_rtt_status = QuicZeroRttStatus::Rejected; + mqtt_events.push(MqttEvent::ZeroRttStatusChanged { + status: QuicZeroRttStatus::Rejected, + }); + + let replay_on_reject = self + .zero_rtt_config + .map(|config| config.replay_on_reject) + .unwrap_or(false); + let journals = std::mem::take(&mut self.early_stream_journal); + self.control_stream = None; + self.data_streams.clear(); + self.default_pub_stream = None; + self.default_sub_stream = None; + self.control_outgoing.clear(); + self.pending_close = None; + self.control_finishing = false; + self.control_finished = false; + + if replay_on_reject { + Self::replay_early_stream_journals( + conn, + journals, + parser_buffer_size, + mqtt_version, + max_packets, + &mut self.control_stream, + &mut self.data_streams, + &mut self.default_pub_stream, + &mut self.default_sub_stream, + &mut self.control_outgoing, + &mut mqtt_events, + ); + } + } + } + + // QUIC handshake done. Open the control stream for MQTT + // unless it already exists from 0-RTT or a rejected + // attempt intentionally left retry to the caller. + if self.control_stream.is_none() && !zero_rtt_was_attempted { + if let Some(stream_id) = conn.streams().open(Dir::Bi) { + self.control_stream = Some(stream_id); + self.mqtt_engine.connect(); + } } } - quinn_proto::Event::ConnectionLost { .. } => { - self.mqtt_engine.handle_connection_lost(); + quinn_proto::Event::ConnectionLost { reason } => { + // Drop transport-tied MQTT output so a stale packet is not + // replayed before CONNECT on the next transport. + self.mqtt_engine.reset_for_new_transport(); + self.control_stream = None; + self.data_streams.clear(); + self.default_pub_stream = None; + self.default_sub_stream = None; + self.control_outgoing.clear(); + self.early_stream_journal.clear(); + self.pending_close = None; + self.control_finishing = false; + self.control_finished = false; + // Surface the QUIC-level detail in addition to the + // MQTT-level Disconnected signal. + let (by_peer, error_code) = transport_close_details(&reason); + mqtt_events.push(MqttEvent::TransportClosed { + reason: reason.to_string(), + by_peer, + error_code, + }); mqtt_events.push(MqttEvent::Disconnected(None)); } _ => {} } } - // 3. Transfer data: QUIC Stream -> MqttEngine - if let Some(stream_id) = self.mqtt_stream { - let mut stream = conn.recv_stream(stream_id); - // 0.11 read(ordered) -> Chunks - // Fix lifetime by ensuring the Result temporary is dropped - let read_result = stream.read(true); - if let Ok(mut chunks) = read_result { - while let Ok(Some(chunk)) = chunks.next(16384) { - mqtt_events.extend(self.mqtt_engine.handle_incoming(&chunk.bytes)); + // 2. Accept server-initiated bidirectional streams as data streams. + while let Some(stream_id) = conn.streams().accept(Dir::Bi) { + self.data_streams.entry(stream_id).or_insert_with(|| { + QuicStream::new(parser_buffer_size, mqtt_version, max_packets) + }); + } + + // 3a. Read the control stream into the shared internal parser. + if let Some(stream_id) = self.control_stream { + let reset_code = { + let mut stream = conn.recv_stream(stream_id); + stream.received_reset().ok().flatten() + }; + if let Some(error_code) = reset_code { + Self::handle_stream_reset_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.default_pub_stream, + &mut self.default_sub_stream, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + stream_id, + error_code, + &mut mqtt_events, + ); + } + } + if let Some(stream_id) = self.control_stream { + let mut recv_finished = false; + let mut read_reset = None; + { + let mut stream = conn.recv_stream(stream_id); + // Bind the Result to a local so its temporary (which borrows + // `stream`) is dropped before `stream` itself. + let read_result = stream.read(true); + if let Ok(mut chunks) = read_result { + loop { + match chunks.next(16384) { + Ok(Some(chunk)) => { + mqtt_events + .extend(self.mqtt_engine.handle_incoming(&chunk.bytes)); + } + Ok(None) => { + recv_finished = true; + break; + } + Err(quinn_proto::ReadError::Reset(error_code)) => { + read_reset = Some(error_code); + break; + } + Err(_) => break, + } + } + } + } + if let Some(error_code) = read_reset { + Self::handle_stream_reset_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.default_pub_stream, + &mut self.default_sub_stream, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + stream_id, + error_code, + &mut mqtt_events, + ); + } else if recv_finished { + Self::handle_stream_closed_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + stream_id, + "recv_finished", + true, + &mut mqtt_events, + ); + } + } + + // 3b. Read each data stream through its own parser, feeding complete + // packets into the shared protocol state machine. + let data_ids: Vec = self.data_streams.keys().copied().collect(); + let max_event_count = self.mqtt_engine.options().max_event_count; + for stream_id in data_ids { + if !self.data_streams.contains_key(&stream_id) { + continue; + } + + // Back-pressure: stop reading inbound data once the application + // event buffer is full (mirrors `handle_incoming`). Undecoded bytes + // stay in quinn-proto's receive buffer, flow-controlling the peer, + // and are processed on a later tick after events are drained. + if mqtt_events.len() >= max_event_count { + break; + } + + let reset_code = { + let mut stream = conn.recv_stream(stream_id); + stream.received_reset().ok().flatten() + }; + if let Some(error_code) = reset_code { + Self::handle_stream_reset_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.default_pub_stream, + &mut self.default_sub_stream, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + stream_id, + error_code, + &mut mqtt_events, + ); + continue; + } + + // Back-pressure: if this stream's bounded send buffer is already + // full, don't pull more inbound data this tick. quinn-proto then + // stops granting the peer flow-control credit, so the buffer cannot + // grow without bound while we are unable to emit the matching + // QoS 1/2 acknowledgements. + let at_capacity = self + .data_streams + .get(&stream_id) + .map(|ds| ds.pending_packets >= ds.max_packets) + .unwrap_or(true); + if at_capacity { + continue; + } + + let recv_closed = self + .data_streams + .get(&stream_id) + .map(|ds| ds.recv_closed) + .unwrap_or(true); + if !recv_closed { + let mut recv_finished = false; + let mut read_reset = None; + { + let mut stream = conn.recv_stream(stream_id); + let read_result = stream.read(true); + if let Ok(mut chunks) = read_result { + loop { + match chunks.next(16384) { + Ok(Some(chunk)) => { + if let Some(ds) = self.data_streams.get_mut(&stream_id) { + ds.parser.feed(&chunk.bytes); + } + } + Ok(None) => { + recv_finished = true; + break; + } + Err(quinn_proto::ReadError::Reset(error_code)) => { + read_reset = Some(error_code); + break; + } + Err(_) => break, + } + } + } + } + if let Some(error_code) = read_reset { + Self::handle_stream_reset_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.default_pub_stream, + &mut self.default_sub_stream, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + stream_id, + error_code, + &mut mqtt_events, + ); + continue; + } else if recv_finished { + Self::handle_stream_closed_fields( + &mut self.control_stream, + &mut self.data_streams, + &mut self.control_outgoing, + &mut self.control_finishing, + &mut self.control_finished, + &mut self.mqtt_engine, + stream_id, + "recv_finished", + true, + &mut mqtt_events, + ); } } + // `chunks`/`stream` are dropped here, releasing the connection + // borrow needed to drive the MQTT engine. Disjoint field borrows + // (data_streams + mqtt_engine) keep this valid while `conn` lives. + Self::drain_data_stream( + &mut self.data_streams, + &mut self.mqtt_engine, + stream_id, + &mut mqtt_events, + ); } - // 4. Transfer data: MqttEngine -> QUIC Stream - let outgoing_bytes = self.mqtt_engine.take_outgoing(); - if !outgoing_bytes.is_empty() { - if let Some(stream_id) = self.mqtt_stream { + // 4a. Flush the control stream. `control_outgoing` is a persistent + // buffer: the engine's session packets are appended to it, then it + // is written with partial-write retention (like the data streams), + // so a flow-controlled or partial write never silently drops bytes + // such as a queued MQTT DISCONNECT. + // Stop pulling new session output once a FIN has been requested or + // applied — those bytes could never be written on the finished send + // side and would block the deferred close below forever. + if !self.control_finishing && !self.control_finished { + let session_bytes = self.mqtt_engine.take_outgoing(); + Self::append_control_outgoing_bytes( + &mut self.control_outgoing, + &mut self.early_stream_journal, + self.zero_rtt_status, + self.control_stream, + &session_bytes, + 0, + ); + } + if let Some(stream_id) = self.control_stream { + // Flush buffered bytes (partial-write retention) until the send side + // is finished; once finished we no longer write. + if !self.control_finished && !self.control_outgoing.is_empty() { let mut stream = conn.send_stream(stream_id); - // @TODO: handle error - let _ = stream.write(&outgoing_bytes); + if let Ok(written) = stream.write(&self.control_outgoing) { + self.control_outgoing.drain(..written); + } + } + // Deferred FIN: finish only once the control buffer is fully + // flushed, so no queued bytes are lost. + if self.control_finishing + && !self.control_finished + && self.control_outgoing.is_empty() + { + let _ = conn.send_stream(stream_id).finish(); + self.control_finishing = false; + self.control_finished = true; + } + } + + // 4b. Flush each data stream's outgoing buffer, retaining any bytes + // that could not be written yet (stream-level back-pressure). + for (stream_id, ds) in self.data_streams.iter_mut() { + if !ds.outgoing.is_empty() { + let mut stream = conn.send_stream(*stream_id); + if let Ok(written) = stream.write(&ds.outgoing) { + ds.outgoing.drain(..written); + } + // Once fully drained the buffered-packet count is cleared, + // freeing capacity for new publishes/subscribes on this stream. + if ds.outgoing.is_empty() { + ds.pending_packets = 0; + } + } + // Deferred FIN: finish the send side only once all buffered bytes + // have actually been written, so queued publishes are not lost. + if ds.finishing && !ds.finished && ds.outgoing.is_empty() { + let _ = conn.send_stream(*stream_id).finish(); + ds.finishing = false; + ds.finished = true; } } @@ -1050,12 +2624,46 @@ impl QuicMqttEngine { // @TODO: Reuse the same buffer; poll_transmit writes the datagram payload into `buf` each time. buf.clear(); } + + // 6. Deferred graceful close: apply only once the control buffer + // (including the queued MQTT DISCONNECT) has been fully written and + // transmitted above. If the control stream is still draining under + // flow control, defer to a later tick so the DISCONNECT is not + // dropped. Then emit the CONNECTION_CLOSE after the DISCONNECT. + if self.pending_close.is_some() && self.control_outgoing.is_empty() { + if let Some((code, reason)) = self.pending_close.take() { + conn.close(now, code, reason); + self.mqtt_engine.handle_connection_lost(); + let mut close_buf = Vec::new(); + while let Some(transmit) = conn.poll_transmit(now, 1, &mut close_buf) { + self.outgoing_datagrams + .push_back((transmit.destination, close_buf.clone())); + close_buf.clear(); + } + } + } } - // 6. Drive MqttEngine tick + // 7. Drive MqttEngine tick let tick_events = self.mqtt_engine.handle_tick(now); mqtt_events.extend(tick_events); + // 8. Route per-stream (MQTT v3) retransmissions back onto their + // originating data stream; they are flushed on the next tick. + // Best-effort: skipped if the stream is gone or its buffer is full — + // the inflight timer will retry on a later tick. + let retransmissions = self.mqtt_engine.take_stream_retransmissions(); + for (stream, bytes) in retransmissions { + if let Some(stream_id) = self + .data_streams + .keys() + .copied() + .find(|id| u64::from(*id) == stream) + { + let _ = self.enqueue_on_stream(stream_id, &bytes); + } + } + mqtt_events } @@ -1064,44 +2672,649 @@ impl QuicMqttEngine { } pub fn take_events(&mut self) -> Vec { - self.mqtt_engine.take_events() + let mut events: Vec = self.pending_transport_events.drain(..).collect(); + events.extend(self.mqtt_engine.take_events()); + events } - /// Delegate: Queue a PUBLISH packet. - pub fn publish(&mut self, command: PublishCommand) -> Result, MqttClientError> { - self.mqtt_engine.publish(command) + /// Open a new client-initiated bidirectional QUIC data stream. + /// + /// Returns an opaque handle (the raw QUIC stream id) that can be passed to + /// [`publish_on`](Self::publish_on), [`subscribe_on`](Self::subscribe_on) and + /// [`unsubscribe_on`](Self::unsubscribe_on). A stream becomes a "pub stream" + /// or "sub stream" simply by virtue of the traffic routed onto it; the broker + /// replies (SUBACK, and QoS 1/2 acknowledgements) on the same stream. + /// + /// Fails if the connection has not been established yet, or if the peer's + /// stream limit (`initial_max_streams_bidi`) has been reached. + pub fn open_data_stream(&mut self) -> Result { + Ok(self + .open_bidi_stream_with_role(EarlyStreamRole::ExplicitData)? + .into()) } - /// Delegate: Queue a SUBSCRIBE packet. - pub fn subscribe(&mut self, command: SubscribeCommand) -> Result { - self.mqtt_engine.subscribe(command) - } + /// Internal: open a bidirectional stream and register per-stream state. + fn open_bidi_stream_with_role( + &mut self, + role: EarlyStreamRole, + ) -> Result { + self.ensure_command_allowed_before_mqtt_connected()?; - /// Delegate: Queue an UNSUBSCRIBE packet. - pub fn unsubscribe(&mut self, command: UnsubscribeCommand) -> Result { - self.mqtt_engine.unsubscribe(command) + let parser_buffer_size = self.mqtt_engine.options().parser_buffer_size; + let mqtt_version = self.mqtt_engine.mqtt_version(); + let max_packets = self.mqtt_engine.options().max_outgoing_packet_count; + + let conn = self + .connection + .as_mut() + .ok_or_else(|| MqttClientError::InvalidState { + expected: "an established QUIC connection".to_string(), + actual: "no connection".to_string(), + })?; + + let stream_id = + conn.streams() + .open(Dir::Bi) + .ok_or_else(|| MqttClientError::InternalError { + message: "QUIC stream limit reached: cannot open new data stream".to_string(), + })?; + + self.data_streams.insert( + stream_id, + QuicStream::new(parser_buffer_size, mqtt_version, max_packets), + ); + if self.zero_rtt_status == QuicZeroRttStatus::Attempted { + Self::journal_stream_open(&mut self.early_stream_journal, stream_id, role); + } + Ok(stream_id) } - /// Delegate: Queue a DISCONNECT packet. - pub fn disconnect(&mut self) { - self.mqtt_engine.disconnect(); + /// Resolve a public stream handle to a live data-stream id, rejecting unknown + /// handles and the control stream. + fn resolve_stream(&self, handle: u64) -> Result { + self.data_streams + .keys() + .copied() + .find(|id| u64::from(*id) == handle) + .ok_or_else(|| MqttClientError::InvalidState { + expected: "a handle from open_data_stream".to_string(), + actual: format!("unknown data stream {}", handle), + }) } - /// Check if the MQTT session is connected. - pub fn is_connected(&self) -> bool { - self.mqtt_engine.is_connected() + /// Reject the operation if the target stream's outgoing buffer is already at + /// its packet-count limit, so callers do not allocate packet ids / inflight + /// slots for a packet that cannot be buffered. + fn ensure_stream_capacity(&self, stream_id: StreamId) -> Result<(), MqttClientError> { + if let Some(ds) = self.data_streams.get(&stream_id) { + if ds.finishing || ds.finished { + return Err(stream_finished_error()); + } + if ds.pending_packets >= ds.max_packets { + return Err(MqttClientError::BufferFull { + buffer_type: "quic_stream_outgoing".to_string(), + capacity: ds.max_packets, + }); + } + } + Ok(()) } -} -#[cfg(test)] -mod tests { - use super::*; - use crate::mqtt_client::opts::MqttClientOptions; - use crate::mqtt_serde::mqttv5::pingreqv5; + /// Append one already-encoded packet to a data stream's bounded outgoing + /// buffer, returning [`MqttClientError::BufferFull`] if the stream is at its + /// packet-count limit, or an error if the stream has been finished/reset. + fn enqueue_on_stream( + &mut self, + stream_id: StreamId, + bytes: &[u8], + ) -> Result<(), MqttClientError> { + if let Some(ds) = self.data_streams.get_mut(&stream_id) { + if ds.finishing || ds.finished { + return Err(stream_finished_error()); + } + if ds.pending_packets >= ds.max_packets { + return Err(MqttClientError::BufferFull { + buffer_type: "quic_stream_outgoing".to_string(), + capacity: ds.max_packets, + }); + } + ds.outgoing.extend_from_slice(bytes); + ds.pending_packets += 1; + if self.zero_rtt_status == QuicZeroRttStatus::Attempted { + Self::journal_stream_bytes(&mut self.early_stream_journal, stream_id, bytes, 1); + } + } + Ok(()) + } - #[test] - fn test_outgoing_buffer_limit() { - let options = MqttClientOptions::builder() + /// Decode and process complete packets buffered on a data stream's parser, + /// routing any protocol response (PUBACK/PUBREC/PUBREL/PUBCOMP) back onto the + /// *same* stream and counting it against the stream's bounded send buffer. + /// + /// Stops once the send buffer reaches `max_packets`, leaving the undecoded + /// bytes in the parser for a later tick (back-pressure). This is an associated + /// function taking the two fields directly so it can run while the QUIC + /// `connection` field is mutably borrowed elsewhere in `handle_tick`. + fn drain_data_stream( + data_streams: &mut HashMap, + mqtt_engine: &mut MqttEngine, + stream_id: StreamId, + events: &mut Vec, + ) { + let max_event_count = mqtt_engine.options().max_event_count; + loop { + // Back-pressure: stop once the event buffer is full (mirrors + // `handle_incoming`). This bounds the number of events produced in a + // single tick even for packets that emit events but no response bytes + // (e.g. a stream full of QoS 0 PUBLISHes). The undecoded bytes remain + // in this stream's parser for a later tick. + if events.len() >= max_event_count { + break; + } + + // Fetch the next packet only while there is buffer capacity; drops the + // `data_streams` borrow before calling into the engine. + let packet = match data_streams.get_mut(&stream_id) { + Some(ds) if ds.pending_packets < ds.max_packets => match ds.parser.next_packet() { + Ok(Some(packet)) => packet, + Ok(None) => break, + Err(e) => { + events.push(MqttEvent::Error(MqttClientError::from(e))); + break; + } + }, + // Unknown stream, or send buffer full → stop (back-pressure). + _ => break, + }; + + let (evs, resp) = mqtt_engine.ingest_stream_packet(packet, u64::from(stream_id)); + events.extend(evs); + if !resp.is_empty() { + if let Some(ds) = data_streams.get_mut(&stream_id) { + // If the send side has been finished/reset we can no longer + // transmit on this stream; surface the message event but drop + // the unsendable response rather than letting it accumulate. + if !ds.finished { + ds.outgoing.extend_from_slice(&resp); + ds.pending_packets += 1; + } + } + } + } + } + + /// Lazily open (once) the default data stream used by [`publish`](Self::publish). + fn ensure_default_pub_stream(&mut self) -> Result { + if let Some(id) = self.default_pub_stream { + return Ok(id); + } + let id = self.open_bidi_stream_with_role(EarlyStreamRole::DefaultPub)?; + self.default_pub_stream = Some(id); + Ok(id) + } + + /// Lazily open (once) the default data stream used by + /// [`subscribe`](Self::subscribe) / [`unsubscribe`](Self::unsubscribe). + fn ensure_default_sub_stream(&mut self) -> Result { + if let Some(id) = self.default_sub_stream { + return Ok(id); + } + let id = self.open_bidi_stream_with_role(EarlyStreamRole::DefaultSub)?; + self.default_sub_stream = Some(id); + Ok(id) + } + + /// Publish on the default pub data stream, opening it on first use. + /// + /// PUBLISH traffic is never placed on the session control stream. The QoS 1/2 + /// acknowledgement handshake completes on this same stream. + pub fn publish(&mut self, command: PublishCommand) -> Result, MqttClientError> { + self.ensure_command_allowed_before_mqtt_connected()?; + let stream_id = self.ensure_default_pub_stream()?; + self.ensure_stream_capacity(stream_id)?; + let (pid, bytes) = self + .mqtt_engine + .publish_encoded(command, Some(u64::from(stream_id)))?; + self.enqueue_on_stream(stream_id, &bytes)?; + Ok(pid) + } + + /// Publish onto a specific data stream previously returned by + /// [`open_data_stream`](Self::open_data_stream). + /// + /// Packet-id allocation and QoS 1/2 inflight tracking are still handled by the + /// shared MQTT state machine, so acknowledgements are processed normally. The + /// entire QoS 1/2 handshake (PUBACK, or PUBREC/PUBREL/PUBCOMP) completes on + /// this same stream — there is no cross-stream firing. + pub fn publish_on( + &mut self, + stream: u64, + command: PublishCommand, + ) -> Result, MqttClientError> { + self.ensure_command_allowed_before_mqtt_connected()?; + let stream_id = self.resolve_stream(stream)?; + self.ensure_stream_capacity(stream_id)?; + let (pid, bytes) = self.mqtt_engine.publish_encoded(command, Some(stream))?; + self.enqueue_on_stream(stream_id, &bytes)?; + Ok(pid) + } + + /// Subscribe on the default sub data stream, opening it on first use. + /// + /// SUBSCRIBE traffic is never placed on the session control stream; the SUBACK + /// and any messages delivered for the subscription flow on this same stream. + pub fn subscribe(&mut self, command: SubscribeCommand) -> Result { + self.ensure_command_allowed_before_mqtt_connected()?; + let stream_id = self.ensure_default_sub_stream()?; + self.ensure_stream_capacity(stream_id)?; + let (pid, bytes) = self + .mqtt_engine + .subscribe_encoded(command, Some(u64::from(stream_id)))?; + self.enqueue_on_stream(stream_id, &bytes)?; + Ok(pid) + } + + /// Subscribe onto a specific data stream. The SUBACK and delivered messages + /// for this subscription arrive on the same stream. + pub fn subscribe_on( + &mut self, + stream: u64, + command: SubscribeCommand, + ) -> Result { + self.ensure_command_allowed_before_mqtt_connected()?; + let stream_id = self.resolve_stream(stream)?; + self.ensure_stream_capacity(stream_id)?; + let (pid, bytes) = self.mqtt_engine.subscribe_encoded(command, Some(stream))?; + self.enqueue_on_stream(stream_id, &bytes)?; + Ok(pid) + } + + /// Unsubscribe on the default sub data stream. + pub fn unsubscribe(&mut self, command: UnsubscribeCommand) -> Result { + self.ensure_command_allowed_before_mqtt_connected()?; + let stream_id = self.ensure_default_sub_stream()?; + self.ensure_stream_capacity(stream_id)?; + let (pid, bytes) = self + .mqtt_engine + .unsubscribe_encoded(command, Some(u64::from(stream_id)))?; + self.enqueue_on_stream(stream_id, &bytes)?; + Ok(pid) + } + + /// Unsubscribe onto a specific data stream. + pub fn unsubscribe_on( + &mut self, + stream: u64, + command: UnsubscribeCommand, + ) -> Result { + self.ensure_command_allowed_before_mqtt_connected()?; + let stream_id = self.resolve_stream(stream)?; + self.ensure_stream_capacity(stream_id)?; + let (pid, bytes) = self + .mqtt_engine + .unsubscribe_encoded(command, Some(stream))?; + self.enqueue_on_stream(stream_id, &bytes)?; + Ok(pid) + } + + /// Number of active data streams, excluding the control stream. + pub fn data_stream_count(&self) -> usize { + self.data_streams.len() + } + + // --- Connection close controls --- + + /// Gracefully close the QUIC connection immediately, sending a + /// CONNECTION_CLOSE frame with the given application error code and reason. + /// + /// Use `error_code = 0` for a normal close. This is an **immediate** QUIC + /// close: quinn-proto abandons any unacknowledged stream data, so it does not + /// guarantee delivery of buffered bytes — use [`disconnect_and_close`](Self::disconnect_and_close) + /// when an MQTT DISCONNECT must reach the peer first. The connection enters its + /// closing state and the CONNECTION_CLOSE frame is emitted on the next + /// [`handle_tick`](Self::handle_tick); the MQTT session is marked disconnected. + pub fn close(&mut self, error_code: u64, reason: &[u8]) -> Result<(), MqttClientError> { + let code = quic_error_code(error_code)?; + // Supersede any deferred close so we don't close twice. + self.pending_close = None; + let conn = self.require_connection()?; + conn.close(Instant::now(), code, bytes::Bytes::copy_from_slice(reason)); + self.mqtt_engine.handle_connection_lost(); + Ok(()) + } + + /// Send an MQTT DISCONNECT on the control stream, then gracefully close the + /// QUIC connection **after** the DISCONNECT has been transmitted. + /// + /// The QUIC close is deferred to [`handle_tick`](Self::handle_tick): the queued + /// DISCONNECT is written and put on the wire first, then CONNECTION_CLOSE is + /// emitted, so the peer has a chance to observe the MQTT-level disconnect. + /// (Delivery is still best-effort — QUIC does not wait for the DISCONNECT to be + /// acknowledged.) The MQTT session is marked disconnected immediately. + pub fn disconnect_and_close( + &mut self, + error_code: u64, + reason: &[u8], + ) -> Result<(), MqttClientError> { + let code = quic_error_code(error_code)?; + // Validate a connection exists before queuing anything. + let _ = self.require_connection()?; + // Queue the DISCONNECT first; if the outgoing buffer is full, surface the + // error and do NOT arm the close, so we never close the QUIC connection + // having silently dropped the MQTT DISCONNECT. (No-op if not connected.) + self.mqtt_engine.try_disconnect()?; + self.pending_close = Some((code, bytes::Bytes::copy_from_slice(reason))); + Ok(()) + } + + /// Silently/abruptly drop the QUIC connection without sending CONNECTION_CLOSE. + /// + /// The peer is left to detect the loss via its idle timeout. Useful for + /// simulating dead-peer / abrupt-disconnect scenarios. All local stream state + /// is cleared. Connection parameters are retained so [`reconnect`](Self::reconnect) + /// still works afterwards. + pub fn close_silent(&mut self) { + self.reset_connection_state(); + } + + // --- Per-stream controls --- + + /// Cleanly finish (FIN) the send side of a stream — the control stream or a + /// data stream. + /// + /// The FIN is **deferred** to [`handle_tick`](Self::handle_tick) and applied + /// only once the stream's buffered outbound bytes have been written, so bytes + /// queued by `publish_on`/`send_raw_on` before the next tick are not lost. + /// After finishing, further writes to the stream are rejected. + /// + /// Finishing the **control stream** ends the MQTT session: the layer is marked + /// disconnected so it produces no further keep-alive/control packets that could + /// never be sent on the finished send side. + pub fn finish_stream(&mut self, stream: u64) -> Result<(), MqttClientError> { + let stream_id = self.resolve_any_stream(stream)?; + // Validate a connection exists; the FIN itself is applied on the next tick. + let _ = self.require_connection()?; + if Some(stream_id) == self.control_stream { + // Capture anything the engine has already queued so it is flushed + // before the FIN, then stop the MQTT layer from producing more control + // traffic (auto-PINGREQ, etc.). + let pending = self.mqtt_engine.take_outgoing(); + Self::append_control_outgoing_bytes( + &mut self.control_outgoing, + &mut self.early_stream_journal, + self.zero_rtt_status, + self.control_stream, + &pending, + 0, + ); + self.mqtt_engine.handle_connection_lost(); + self.control_finishing = true; + } else if let Some(ds) = self.data_streams.get_mut(&stream_id) { + ds.finishing = true; + } + Ok(()) + } + + /// Reset (RESET_STREAM) the send side of a stream with the given error code, + /// discarding any buffered outbound bytes. After resetting, further writes to + /// the stream are rejected. + pub fn reset_stream(&mut self, stream: u64, error_code: u64) -> Result<(), MqttClientError> { + let stream_id = self.resolve_any_stream(stream)?; + let code = quic_error_code(error_code)?; + if Some(stream_id) == self.control_stream { + self.control_outgoing.clear(); + self.control_finishing = false; + self.control_finished = true; + // Resetting the control send side ends the MQTT session. + self.mqtt_engine.handle_connection_lost(); + } else if let Some(ds) = self.data_streams.get_mut(&stream_id) { + ds.outgoing.clear(); + ds.pending_packets = 0; + ds.finishing = false; + ds.finished = true; + } + let conn = self.require_connection()?; + conn.send_stream(stream_id) + .reset(code) + .map_err(|e| MqttClientError::InternalError { + message: format!("QUIC stream reset failed: {:?}", e), + }) + } + + /// Ask the peer to stop sending on a stream (STOP_SENDING) with the given + /// error code. + pub fn stop_stream(&mut self, stream: u64, error_code: u64) -> Result<(), MqttClientError> { + let stream_id = self.resolve_any_stream(stream)?; + let code = quic_error_code(error_code)?; + let conn = self.require_connection()?; + conn.recv_stream(stream_id) + .stop(code) + .map_err(|e| MqttClientError::InternalError { + message: format!("QUIC stream stop failed: {:?}", e), + }) + } + + // --- Keep-alive / ping --- + + /// Queue an MQTT PINGREQ on the control stream (manual keep-alive). + /// + /// Rejected until the MQTT session is connected (CONNACK received), so a + /// PINGREQ can never be queued ahead of CONNECT on a fresh transport, and + /// rejected once the control stream's send side has been finished or reset. + pub fn ping(&mut self) -> Result<(), MqttClientError> { + if self.control_finishing || self.control_finished { + return Err(stream_finished_error()); + } + if !self.mqtt_engine.is_connected() { + return Err(MqttClientError::InvalidState { + expected: "a connected MQTT session".to_string(), + actual: "not connected".to_string(), + }); + } + // Propagate a full outgoing buffer instead of silently dropping the ping. + self.mqtt_engine.try_send_ping() + } + + /// Send a QUIC-level PING frame (transport keep-alive), independent of MQTT. + pub fn quic_ping(&mut self) -> Result<(), MqttClientError> { + self.require_connection()?.ping(); + Ok(()) + } + + // --- Internal helpers shared by the controls above --- + + fn require_connection(&mut self) -> Result<&mut Connection, MqttClientError> { + self.connection + .as_mut() + .ok_or_else(|| MqttClientError::InvalidState { + expected: "an established QUIC connection".to_string(), + actual: "no connection".to_string(), + }) + } + + /// Resolve a handle to either the control stream or a known data stream. + fn resolve_any_stream(&self, handle: u64) -> Result { + if Some(handle) == self.control_stream.map(u64::from) { + return Ok(self.control_stream.unwrap()); + } + self.resolve_stream(handle) + } + + /// The control stream handle, available once the QUIC handshake has completed. + /// + /// Can be passed to [`send_raw_on`](Self::send_raw_on) / + /// [`send_packet_on`](Self::send_packet_on) to target the control stream + /// directly (including with packets that do not normally belong there). + pub fn control_stream_id(&self) -> Option { + self.control_stream.map(u64::from) + } + + // --- Low-level escape hatch (negative / conformance testing) --- + // + // The routing helpers above keep traffic spec-correct by default. flowSDK is + // also used to exercise *bad* behaviour — e.g. sending a CONNECT on a data + // stream, or PUBLISH on the control stream — so the following middle-layer + // API deliberately bypasses all routing rules and protocol-state bookkeeping. + + /// Write raw bytes onto an arbitrary stream (the control stream or any data + /// stream), bypassing all MQTT routing rules and protocol state. + /// + /// Intended for negative testing: malformed frames, packets sent on the + /// "wrong" stream, etc. The bytes are buffered on the target stream and + /// flushed on the next [`handle_tick`](Self::handle_tick). + /// + /// Errors if `stream` is neither the control stream nor a known data stream, + /// or if the stream's send side has already been finished or reset. + pub fn send_raw_on(&mut self, stream: u64, bytes: &[u8]) -> Result<(), MqttClientError> { + if Some(stream) == self.control_stream.map(u64::from) { + if self.control_finishing || self.control_finished { + return Err(stream_finished_error()); + } + Self::append_control_outgoing_bytes( + &mut self.control_outgoing, + &mut self.early_stream_journal, + self.zero_rtt_status, + self.control_stream, + bytes, + 0, + ); + return Ok(()); + } + let stream_id = self.resolve_stream(stream)?; + if let Some(ds) = self.data_streams.get_mut(&stream_id) { + if ds.finishing || ds.finished { + return Err(stream_finished_error()); + } + ds.outgoing.extend_from_slice(bytes); + if self.zero_rtt_status == QuicZeroRttStatus::Attempted { + Self::journal_stream_bytes(&mut self.early_stream_journal, stream_id, bytes, 0); + } + } + Ok(()) + } + + /// Encode an arbitrary MQTT packet and write it onto an arbitrary stream, + /// bypassing routing rules. Unlike the high-level methods this does **not** + /// allocate packet ids or track inflight state, so the caller has full control + /// (e.g. to send `MqttPacket::Connect5(..)` on a data stream). + pub fn send_packet_on( + &mut self, + stream: u64, + packet: MqttPacket, + ) -> Result<(), MqttClientError> { + let bytes = packet.to_bytes().map_err(MqttClientError::from)?; + self.send_raw_on(stream, &bytes) + } + + /// Send a success PUBACK on a specific QUIC stream. + /// + /// This is a typed convenience wrapper around [`send_packet_on`](Self::send_packet_on) + /// for manual-ack tests. It does not update protocol/inflight state. + pub fn puback_on(&mut self, stream: u64, packet_id: u16) -> Result<(), MqttClientError> { + self.send_packet_on(stream, self.puback_packet(packet_id)) + } + + /// Send a success PUBREC on a specific QUIC stream. + /// + /// This is a typed convenience wrapper around [`send_packet_on`](Self::send_packet_on) + /// for manual-ack tests. It does not update protocol/inflight state. + pub fn pubrec_on(&mut self, stream: u64, packet_id: u16) -> Result<(), MqttClientError> { + self.send_packet_on(stream, self.pubrec_packet(packet_id)) + } + + /// Send a success PUBREL on a specific QUIC stream. + /// + /// This is a typed convenience wrapper around [`send_packet_on`](Self::send_packet_on) + /// for manual-ack tests. It does not update protocol/inflight state. + pub fn pubrel_on(&mut self, stream: u64, packet_id: u16) -> Result<(), MqttClientError> { + self.send_packet_on(stream, self.pubrel_packet(packet_id)) + } + + /// Send a success PUBCOMP on a specific QUIC stream. + /// + /// This is a typed convenience wrapper around [`send_packet_on`](Self::send_packet_on) + /// for manual-ack tests. It does not update protocol/inflight state. + pub fn pubcomp_on(&mut self, stream: u64, packet_id: u16) -> Result<(), MqttClientError> { + self.send_packet_on(stream, self.pubcomp_packet(packet_id)) + } + + fn puback_packet(&self, packet_id: u16) -> MqttPacket { + if self.mqtt_engine.mqtt_version() == 5 { + MqttPacket::PubAck5(MqttPubAck::new(packet_id, 0, Vec::new())) + } else { + MqttPacket::PubAck3(crate::mqtt_serde::mqttv3::puback::MqttPubAck::new( + packet_id, + )) + } + } + + fn pubrec_packet(&self, packet_id: u16) -> MqttPacket { + if self.mqtt_engine.mqtt_version() == 5 { + MqttPacket::PubRec5(MqttPubRec::new(packet_id, 0, Vec::new())) + } else { + MqttPacket::PubRec3(crate::mqtt_serde::mqttv3::pubrec::MqttPubRec::new( + packet_id, + )) + } + } + + fn pubrel_packet(&self, packet_id: u16) -> MqttPacket { + if self.mqtt_engine.mqtt_version() == 5 { + MqttPacket::PubRel5(MqttPubRel::new(packet_id, 0, Vec::new())) + } else { + MqttPacket::PubRel3(crate::mqtt_serde::mqttv3::pubrel::MqttPubRel::new( + packet_id, + )) + } + } + + fn pubcomp_packet(&self, packet_id: u16) -> MqttPacket { + if self.mqtt_engine.mqtt_version() == 5 { + MqttPacket::PubComp5(MqttPubComp::new(packet_id, 0, Vec::new())) + } else { + MqttPacket::PubComp3(crate::mqtt_serde::mqttv3::pubcomp::MqttPubComp::new( + packet_id, + )) + } + } + + /// Borrow the underlying sans-I/O MQTT protocol engine (the "middle layer"). + /// + /// Exposed so tests can drive the protocol state machine directly — inspect + /// state, or enqueue packets that the routing helpers would place elsewhere. + pub fn engine(&self) -> &MqttEngine { + &self.mqtt_engine + } + + /// Mutably borrow the underlying sans-I/O MQTT protocol engine. + /// + /// Note: packets enqueued via [`MqttEngine::enqueue_packet`] flow on the + /// control stream. To place arbitrary packets on a data stream, encode them + /// and use [`send_packet_on`](Self::send_packet_on) / [`send_raw_on`](Self::send_raw_on). + pub fn engine_mut(&mut self) -> &mut MqttEngine { + &mut self.mqtt_engine + } + + /// Delegate: Queue a DISCONNECT packet. + pub fn disconnect(&mut self) { + self.mqtt_engine.disconnect(); + } + + /// Check if the MQTT session is connected. + pub fn is_connected(&self) -> bool { + self.mqtt_engine.is_connected() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::mqtt_client::opts::MqttClientOptions; + use crate::mqtt_serde::mqttv5::pingreqv5; + + #[test] + fn test_outgoing_buffer_limit() { + let options = MqttClientOptions::builder() .max_outgoing_packet_count(2) .build(); let mut engine = MqttEngine::new(options); @@ -1165,4 +3378,1366 @@ mod tests { _ => panic!("Expected second PingResponse"), } } + + use crate::mqtt_client::commands::PublishCommand; + + #[test] + fn test_publish_encoded_qos0_returns_bytes_without_inflight() { + let mut engine = MqttEngine::new(MqttClientOptions::builder().build()); + engine.connect(); // initialize session for packet-id allocation + let _ = engine.take_outgoing(); // discard the CONNECT bytes + + let cmd = PublishCommand::builder() + .topic("a/b") + .payload("hi".to_string()) + .qos(0) + .build() + .unwrap(); + + let (pid, bytes) = engine.publish_encoded(cmd, None).unwrap(); + assert!(pid.is_none(), "QoS 0 must not allocate a packet id"); + assert!(!bytes.is_empty()); + // QoS 0 is never tracked for retransmission. + assert!(engine.inflight_queue.is_empty()); + // publish_encoded bypasses the outgoing buffer entirely. + assert!(engine.outgoing_buffer.is_empty()); + } + + #[test] + fn test_publish_encoded_qos1_tracks_inflight_and_acks() { + let mut engine = MqttEngine::new(MqttClientOptions::builder().build()); + engine.connect(); + engine.is_connected = true; + let _ = engine.take_outgoing(); + + let cmd = PublishCommand::builder() + .topic("a/b") + .payload("hi".to_string()) + .qos(1) + .build() + .unwrap(); + + let (pid, bytes) = engine.publish_encoded(cmd, Some(7)).unwrap(); + let pid = pid.expect("QoS 1 must allocate a packet id"); + assert!(!bytes.is_empty()); + assert!(!engine.inflight_queue.is_empty()); + + // Acknowledge it through ingest_stream_packet (as if it arrived on a data stream). + let ack = MqttPacket::PubAck5(MqttPubAck::new(pid, 0, Vec::new())); + let (events, resp) = engine.ingest_stream_packet(ack, 7); + assert!( + events.iter().any(|e| matches!(e, MqttEvent::Published(_))), + "expected a Published event after PUBACK" + ); + // A terminal PUBACK requires no response on the stream. + assert!(resp.is_empty()); + assert!( + engine.inflight_queue.is_empty(), + "inflight entry must be cleared after PUBACK" + ); + } + + #[test] + fn test_ingest_stream_packet_emits_event() { + let mut engine = MqttEngine::new(MqttClientOptions::builder().build()); + let (events, resp) = engine.ingest_stream_packet( + MqttPacket::PingResp5(crate::mqtt_serde::mqttv5::pingrespv5::MqttPingResp::new()), + 0, + ); + assert_eq!(events.len(), 1); + assert!(matches!(events[0], MqttEvent::PingResponse(_))); + assert!(resp.is_empty()); + } + + #[test] + fn test_v3_retransmission_routes_to_origin_stream() { + // MQTT v3 QoS 1 publish sent on data stream 9 must, on retransmission, + // be routed back to stream 9 — never the shared/control outgoing buffer. + let options = MqttClientOptions::builder() + .mqtt_version(3) + .keep_alive(0) // disable keep-alive PING so the buffer stays clean + .retransmission_timeout_ms(10) + .build(); + let mut engine = MqttEngine::new(options); + engine.connect(); + engine.is_connected = true; + let _ = engine.take_outgoing(); // discard CONNECT + + let cmd = PublishCommand::builder() + .topic("t") + .payload("x".to_string()) + .qos(1) + .build() + .unwrap(); + engine.publish_encoded(cmd, Some(9)).unwrap(); + + // Drive the tick well past the retransmission timeout. + let future = Instant::now() + Duration::from_secs(3600); + let _ = engine.handle_tick(future); + + let retrans = engine.take_stream_retransmissions(); + assert_eq!( + retrans.len(), + 1, + "expected one stream-routed retransmission" + ); + assert_eq!( + retrans[0].0, 9, + "retransmission must target the origin stream" + ); + assert!( + engine.outgoing_buffer.is_empty(), + "v3 stream retransmission must not land on the shared/control buffer" + ); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_send_raw_on_routes_to_target_stream() { + use quinn_proto::{Dir, Side, StreamId}; + + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + + // Simulate an established control stream and one data stream. + let control = StreamId::new(Side::Client, Dir::Bi, 0); + let data = StreamId::new(Side::Client, Dir::Bi, 1); + engine.control_stream = Some(control); + engine + .data_streams + .insert(data, QuicStream::new(1024, 5, 1000)); + + // Raw bytes targeting the data stream land in that stream's buffer. + engine.send_raw_on(u64::from(data), &[1, 2, 3]).unwrap(); + assert_eq!( + engine.data_streams.get(&data).unwrap().outgoing, + vec![1, 2, 3] + ); + assert!(engine.control_outgoing.is_empty()); + + // Raw bytes targeting the control stream land in the control buffer only — + // no cross-fire to the data stream. + engine.send_raw_on(u64::from(control), &[9]).unwrap(); + assert_eq!(engine.control_outgoing, vec![9]); + assert_eq!( + engine.data_streams.get(&data).unwrap().outgoing, + vec![1, 2, 3] + ); + + // Sending a CONNECT (a control packet) on the data stream is allowed — + // this is exactly the kind of bad behaviour flowSDK must be able to drive. + let connect = connectv5::MqttConnect::new( + "bad-client".to_string(), + None, + None, + None, + 60, + true, + Vec::new(), + ); + engine + .send_packet_on(u64::from(data), MqttPacket::Connect5(connect)) + .unwrap(); + assert!(engine.data_streams.get(&data).unwrap().outgoing.len() > 3); + + // An unknown stream handle is rejected. + assert!(engine.send_raw_on(9999, &[0]).is_err()); + + // After the send side is finished/reset, raw writes are rejected rather + // than silently buffering unsendable bytes. + engine.data_streams.get_mut(&data).unwrap().finished = true; + assert!(engine.send_raw_on(u64::from(data), &[0]).is_err()); + engine.control_finished = true; + assert!(engine.send_raw_on(u64::from(control), &[0]).is_err()); + } + + #[cfg(feature = "quic-proto")] + fn parse_packets(bytes: &[u8], mqtt_version: u8) -> Vec { + use crate::mqtt_serde::parser::ParseOk; + + let mut packets = Vec::new(); + let mut offset = 0; + while offset < bytes.len() { + match MqttPacket::from_bytes_with_version(&bytes[offset..], mqtt_version).unwrap() { + ParseOk::Packet(packet, consumed) => { + packets.push(packet); + offset += consumed; + } + ParseOk::Continue(_, _) => panic!("unexpected partial packet"), + ParseOk::TopicName(_, _) => panic!("unexpected topic-name parse result"), + } + } + packets + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_manual_ack_helpers_route_v5_packets_to_target_stream() { + use quinn_proto::{Dir, Side, StreamId}; + + let mut engine = QuicMqttEngine::new( + MqttClientOptions::builder() + .mqtt_version(5) + .auto_ack(false) + .build(), + ) + .unwrap(); + let data = StreamId::new(Side::Client, Dir::Bi, 1); + engine + .data_streams + .insert(data, QuicStream::new(1024, 5, 1000)); + + engine.puback_on(u64::from(data), 11).unwrap(); + engine.pubrec_on(u64::from(data), 12).unwrap(); + engine.pubrel_on(u64::from(data), 13).unwrap(); + engine.pubcomp_on(u64::from(data), 14).unwrap(); + + let bytes = &engine.data_streams.get(&data).unwrap().outgoing; + let packets = parse_packets(bytes, 5); + assert_eq!(packets.len(), 4); + assert!( + matches!(&packets[0], MqttPacket::PubAck5(p) if p.packet_id == 11 && p.reason_code == 0) + ); + assert!( + matches!(&packets[1], MqttPacket::PubRec5(p) if p.packet_id == 12 && p.reason_code == 0) + ); + assert!( + matches!(&packets[2], MqttPacket::PubRel5(p) if p.packet_id == 13 && p.reason_code == 0) + ); + assert!( + matches!(&packets[3], MqttPacket::PubComp5(p) if p.packet_id == 14 && p.reason_code == 0) + ); + assert!(engine.control_outgoing.is_empty()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_manual_ack_helpers_route_v3_packets_to_target_stream() { + use quinn_proto::{Dir, Side, StreamId}; + + let mut engine = QuicMqttEngine::new( + MqttClientOptions::builder() + .mqtt_version(3) + .auto_ack(false) + .build(), + ) + .unwrap(); + let data = StreamId::new(Side::Client, Dir::Bi, 1); + engine + .data_streams + .insert(data, QuicStream::new(1024, 3, 1000)); + + engine.puback_on(u64::from(data), 21).unwrap(); + engine.pubrec_on(u64::from(data), 22).unwrap(); + engine.pubrel_on(u64::from(data), 23).unwrap(); + engine.pubcomp_on(u64::from(data), 24).unwrap(); + + let bytes = &engine.data_streams.get(&data).unwrap().outgoing; + let packets = parse_packets(bytes, 3); + assert_eq!(packets.len(), 4); + assert!(matches!(&packets[0], MqttPacket::PubAck3(p) if p.message_id == 21)); + assert!(matches!(&packets[1], MqttPacket::PubRec3(p) if p.message_id == 22)); + assert!(matches!(&packets[2], MqttPacket::PubRel3(p) if p.message_id == 23)); + assert!(matches!(&packets[3], MqttPacket::PubComp3(p) if p.message_id == 24)); + assert!(engine.control_outgoing.is_empty()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_manual_ack_helpers_preserve_send_packet_on_errors() { + use quinn_proto::{Dir, Side, StreamId}; + + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + let data = StreamId::new(Side::Client, Dir::Bi, 1); + engine + .data_streams + .insert(data, QuicStream::new(1024, 5, 1000)); + + assert!(engine.puback_on(9999, 1).is_err()); + + engine.data_streams.get_mut(&data).unwrap().finished = true; + assert!(engine.pubrec_on(u64::from(data), 2).is_err()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_drain_data_stream_bounds_response_buffer() { + use quinn_proto::{Dir, Side, StreamId}; + + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + engine.mqtt_engine.is_connected = true; + + let max_packets = 3; + let data = StreamId::new(Side::Client, Dir::Bi, 1); + engine + .data_streams + .insert(data, QuicStream::new(16384, 5, max_packets)); + + // Feed many incoming QoS 1 PUBLISH packets into this stream's parser; each + // would normally generate a PUBACK response buffered on the same stream. + let mut bytes = Vec::new(); + for pid in 1..=10u16 { + let publish = MqttPublish::new_with_prop( + 1, + "t/x".to_string(), + Some(pid), + b"p".to_vec(), + false, + false, + Vec::new(), + ); + bytes.extend(MqttPacket::Publish5(publish).to_bytes().unwrap()); + } + engine + .data_streams + .get_mut(&data) + .unwrap() + .parser + .feed(&bytes); + + let mut events = Vec::new(); + QuicMqttEngine::drain_data_stream( + &mut engine.data_streams, + &mut engine.mqtt_engine, + data, + &mut events, + ); + + let ds = engine.data_streams.get(&data).unwrap(); + // Back-pressure: ingestion stops at the buffer cap rather than growing + // unbounded, and the packet count reflects the buffered responses. + assert_eq!(ds.pending_packets, max_packets); + assert!(!ds.outgoing.is_empty()); + // Only `max_packets` publishes were processed; the rest stay in the parser. + let received = events + .iter() + .filter(|e| matches!(e, MqttEvent::MessageReceived(_))) + .count(); + assert_eq!(received, max_packets); + + // Draining again without freeing the buffer makes no further progress. + let mut more = Vec::new(); + QuicMqttEngine::drain_data_stream( + &mut engine.data_streams, + &mut engine.mqtt_engine, + data, + &mut more, + ); + assert_eq!( + engine.data_streams.get(&data).unwrap().pending_packets, + max_packets + ); + assert!(more.is_empty()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_drain_data_stream_bounds_events_for_qos0() { + use quinn_proto::{Dir, Side, StreamId}; + + // QoS 0 PUBLISHes produce events but no response bytes, so the send-buffer + // cap alone would not bound them — the event-count cap must. + let max_event_count = 4; + let opts = MqttClientOptions::builder() + .max_event_count(max_event_count) + .build(); + let mut engine = QuicMqttEngine::new(opts).unwrap(); + engine.mqtt_engine.is_connected = true; + + let data = StreamId::new(Side::Client, Dir::Bi, 1); + engine + .data_streams + .insert(data, QuicStream::new(16384, 5, 1000)); + + let mut bytes = Vec::new(); + for _ in 0..20 { + let publish = MqttPublish::new_with_prop( + 0, + "t/x".to_string(), + None, + b"p".to_vec(), + false, + false, + Vec::new(), + ); + bytes.extend(MqttPacket::Publish5(publish).to_bytes().unwrap()); + } + engine + .data_streams + .get_mut(&data) + .unwrap() + .parser + .feed(&bytes); + + let mut events = Vec::new(); + QuicMqttEngine::drain_data_stream( + &mut engine.data_streams, + &mut engine.mqtt_engine, + data, + &mut events, + ); + + // Event production is bounded by max_event_count even though no response + // bytes are generated (pending_packets stays 0). + assert_eq!(events.len(), max_event_count); + assert_eq!(engine.data_streams.get(&data).unwrap().pending_packets, 0); + } + + #[test] + fn test_ingest_stream_packet_qos1_publish_acks_on_same_stream() { + // An incoming QoS 1 PUBLISH must produce PUBACK bytes for the caller to + // write back on the originating stream — not enqueued to the shared buffer. + let mut engine = MqttEngine::new(MqttClientOptions::builder().build()); + engine.is_connected = true; + + let publish = MqttPublish::new_with_prop( + 1, + "t/1".to_string(), + Some(42), + b"payload".to_vec(), + false, + false, + Vec::new(), + ); + let (events, resp) = engine.ingest_stream_packet(MqttPacket::Publish5(publish), 5); + + assert!(events + .iter() + .any(|e| matches!(e, MqttEvent::MessageReceived(_)))); + // The response is a PUBACK (v5 fixed header 0x40) and must be returned to + // the caller, never placed on the shared/control outgoing buffer. + assert!( + !resp.is_empty(), + "QoS 1 PUBLISH must yield a PUBACK response" + ); + assert_eq!(resp[0] & 0xF0, 0x40, "response must be a PUBACK"); + assert!( + engine.outgoing_buffer.is_empty(), + "ack must not cross-fire onto the shared/control buffer" + ); + } + + #[test] + fn test_auto_ack_false_suppresses_shared_publish_acks() { + let mut engine = MqttEngine::new( + MqttClientOptions::builder() + .auto_ack(false) + .mqtt_version(5) + .build(), + ); + engine.is_connected = true; + + let qos1 = MqttPublish::new_with_prop( + 1, + "t/qos1".to_string(), + Some(42), + b"payload".to_vec(), + false, + false, + Vec::new(), + ); + let events = engine.handle_incoming(&MqttPacket::Publish5(qos1).to_bytes().unwrap()); + assert!(events + .iter() + .any(|e| matches!(e, MqttEvent::MessageReceived(_)))); + assert!( + engine.take_outgoing().is_empty(), + "auto_ack(false) must suppress PUBACK" + ); + + let qos2 = MqttPublish::new_with_prop( + 2, + "t/qos2".to_string(), + Some(43), + b"payload".to_vec(), + false, + false, + Vec::new(), + ); + let events = engine.handle_incoming(&MqttPacket::Publish5(qos2).to_bytes().unwrap()); + assert!(events + .iter() + .any(|e| matches!(e, MqttEvent::MessageReceived(_)))); + assert!( + engine.take_outgoing().is_empty(), + "auto_ack(false) must suppress PUBREC" + ); + } + + #[test] + fn test_auto_ack_false_suppresses_shared_pubrel_ack() { + let mut engine = MqttEngine::new( + MqttClientOptions::builder() + .auto_ack(false) + .mqtt_version(5) + .build(), + ); + engine.is_connected = true; + + let pubrel = MqttPacket::PubRel5(MqttPubRel::new(44, 0, Vec::new())); + let events = engine.handle_incoming(&pubrel.to_bytes().unwrap()); + assert!(matches!( + events.as_slice(), + [MqttEvent::PubRelReceived { + packet_id: 44, + stream: None + }] + )); + assert!( + engine.take_outgoing().is_empty(), + "auto_ack(false) must suppress PUBCOMP" + ); + } + + #[test] + fn test_auto_ack_false_suppresses_stream_publish_and_pubrel_acks() { + let mut engine = MqttEngine::new( + MqttClientOptions::builder() + .auto_ack(false) + .mqtt_version(5) + .build(), + ); + engine.is_connected = true; + + let publish = MqttPublish::new_with_prop( + 1, + "t/1".to_string(), + Some(45), + b"payload".to_vec(), + false, + false, + Vec::new(), + ); + let (events, resp) = engine.ingest_stream_packet(MqttPacket::Publish5(publish), 5); + assert!(events + .iter() + .any(|e| matches!(e, MqttEvent::MessageReceived(_)))); + assert!(resp.is_empty(), "auto_ack(false) must suppress PUBACK"); + + let pubrel = MqttPacket::PubRel5(MqttPubRel::new(46, 0, Vec::new())); + let (events, resp) = engine.ingest_stream_packet(pubrel, 5); + assert!(matches!( + events.as_slice(), + [MqttEvent::PubRelReceived { + packet_id: 46, + stream: Some(5) + }] + )); + assert!(resp.is_empty(), "auto_ack(false) must suppress PUBCOMP"); + } + + #[test] + fn test_auto_ack_false_suppresses_v3_publish_ack() { + let mut engine = MqttEngine::new( + MqttClientOptions::builder() + .auto_ack(false) + .mqtt_version(3) + .build(), + ); + engine.is_connected = true; + + let publish = crate::mqtt_serde::mqttv3::publish::MqttPublish::new( + "t/v3".to_string(), + 1, + b"payload".to_vec(), + Some(47), + false, + false, + ); + let events = engine.handle_incoming(&MqttPacket::Publish3(publish).to_bytes().unwrap()); + assert!(events + .iter() + .any(|e| matches!(e, MqttEvent::MessageReceived(_)))); + assert!( + engine.take_outgoing().is_empty(), + "auto_ack(false) must suppress v3 PUBACK" + ); + } + + #[test] + fn test_try_send_ping_and_disconnect_report_buffer_full() { + let mut engine = MqttEngine::new( + MqttClientOptions::builder() + .max_outgoing_packet_count(1) + .build(), + ); + engine.connect(); // queues CONNECT, filling the single-slot buffer + engine.is_connected = true; + assert_eq!(engine.outgoing_buffer.len(), 1); + + // Fallible paths must report BufferFull rather than silently dropping. + assert!(matches!( + engine.try_send_ping(), + Err(MqttClientError::BufferFull { .. }) + )); + assert!(matches!( + engine.try_disconnect(), + Err(MqttClientError::BufferFull { .. }) + )); + // A failed try_disconnect must not mark the session disconnected. + assert!(engine.is_connected()); + + // After draining, both succeed and disconnect updates state. + let _ = engine.take_outgoing(); + assert!(engine.try_send_ping().is_ok()); + let _ = engine.take_outgoing(); + assert!(engine.try_disconnect().is_ok()); + assert!(!engine.is_connected()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_quic_ping_propagates_buffer_full() { + let mut engine = QuicMqttEngine::new( + MqttClientOptions::builder() + .max_outgoing_packet_count(1) + .build(), + ) + .unwrap(); + engine.mqtt_engine.is_connected = true; + // Fill the single outgoing slot, then ping() must report BufferFull + // instead of returning Ok without queuing a PINGREQ. + engine.mqtt_engine.send_ping(); + assert!(matches!( + engine.ping(), + Err(MqttClientError::BufferFull { .. }) + )); + } + + #[test] + fn test_reset_for_new_transport_clears_outbound_keeps_inflight() { + let mut engine = MqttEngine::new(MqttClientOptions::builder().build()); + engine.connect(); // queues CONNECT into outgoing_buffer + engine.is_connected = true; + + // A QoS 1 publish (tracked inflight) and a PINGREQ (queued outbound). + let cmd = PublishCommand::builder() + .topic("t") + .payload("x".to_string()) + .qos(1) + .build() + .unwrap(); + engine.publish_encoded(cmd, Some(3)).unwrap(); + engine.send_ping(); + // A pending reconnect deadline from a prior timeout. + engine.next_reconnect_at = Some(Instant::now() + Duration::from_secs(30)); + assert!(!engine.outgoing_buffer.is_empty()); + assert!(!engine.inflight_queue.is_empty()); + + engine.reset_for_new_transport(); + + // The stale reconnect deadline must be cancelled so it cannot fire another + // ReconnectNeeded during the new handshake. + assert!(engine.next_reconnect_at.is_none()); + + // Transport-tied output is discarded so it cannot be replayed before the + // next CONNECT, but the inflight queue is retained for session resumption. + assert!( + engine.outgoing_buffer.is_empty(), + "stale outbound bytes must be cleared on reset" + ); + assert!(engine.stream_retransmissions.is_empty()); + assert!(!engine.is_connected()); + assert!( + !engine.inflight_queue.is_empty(), + "inflight must be retained for session resumption" + ); + } + + #[test] + fn test_auto_keepalive_toggle() { + let tick_at = Instant::now() + Duration::from_millis(1100); + + // Disabled: no PINGREQ emitted even though keep_alive has elapsed. + let mut off = MqttEngine::new( + MqttClientOptions::builder() + .keep_alive(1) + .auto_keepalive(false) + .build(), + ); + off.connect(); + off.is_connected = true; + let _ = off.take_outgoing(); + let _ = off.handle_tick(tick_at); + assert!( + off.outgoing_buffer.is_empty(), + "no PINGREQ when auto_keepalive is disabled" + ); + + // Enabled: PINGREQ is emitted on the keep-alive timer. + let mut on = MqttEngine::new( + MqttClientOptions::builder() + .keep_alive(1) + .auto_keepalive(true) + .build(), + ); + on.connect(); + on.is_connected = true; + let _ = on.take_outgoing(); + let _ = on.handle_tick(tick_at); + assert!( + !on.outgoing_buffer.is_empty(), + "PINGREQ expected when auto_keepalive is enabled" + ); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_transport_close_details_mapping() { + use quinn_proto::{ApplicationClose, ConnectionClose, TransportErrorCode, VarInt}; + + let (by_peer, code) = transport_close_details(&ConnectionError::TimedOut); + assert!(!by_peer); + assert_eq!(code, None); + + let (by_peer, _) = transport_close_details(&ConnectionError::Reset); + assert!(by_peer, "Reset is peer-initiated"); + + let (by_peer, _) = transport_close_details(&ConnectionError::LocallyClosed); + assert!(!by_peer, "LocallyClosed is local"); + + // Peer application close preserves the application error code. + let app = ConnectionError::ApplicationClosed(ApplicationClose { + error_code: VarInt::from_u32(5), + reason: bytes::Bytes::new(), + }); + assert_eq!(transport_close_details(&app), (true, Some(5))); + + // Peer transport close preserves the transport error code (P3 regression). + let transport = ConnectionError::ConnectionClosed(ConnectionClose { + error_code: TransportErrorCode::APPLICATION_ERROR, + frame_type: None, + reason: bytes::Bytes::new(), + }); + let (by_peer, code) = transport_close_details(&transport); + assert!(by_peer); + assert_eq!(code, Some(u64::from(TransportErrorCode::APPLICATION_ERROR))); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_quic_error_code_rejects_overflow() { + assert!(quic_error_code(0).is_ok()); + assert!(quic_error_code((1u64 << 62) - 1).is_ok()); + assert!(quic_error_code(u64::MAX).is_err()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_stream_abort_events_clear_data_stream_refs() { + use quinn_proto::{Dir, Side, StreamId}; + + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + let reset_stream = StreamId::new(Side::Client, Dir::Bi, 1); + engine + .data_streams + .insert(reset_stream, QuicStream::new(1024, 5, 1000)); + engine.default_pub_stream = Some(reset_stream); + + let mut events = Vec::new(); + engine.handle_stream_reset(reset_stream, VarInt::from_u32(42), &mut events); + + assert!(matches!( + events.as_slice(), + [MqttEvent::StreamReset { + stream_id, + error_code: 42 + }] if *stream_id == u64::from(reset_stream) + )); + assert!(!engine.data_streams.contains_key(&reset_stream)); + assert_eq!(engine.default_pub_stream, None); + + let stopped_stream = StreamId::new(Side::Client, Dir::Bi, 2); + engine + .data_streams + .insert(stopped_stream, QuicStream::new(1024, 5, 1000)); + engine.default_sub_stream = Some(stopped_stream); + + events.clear(); + engine.handle_stream_stopped(stopped_stream, VarInt::from_u32(7), &mut events); + + assert!(matches!( + events.as_slice(), + [MqttEvent::StreamStopped { + stream_id, + error_code: 7 + }] if *stream_id == u64::from(stopped_stream) + )); + assert!(!engine.data_streams.contains_key(&stopped_stream)); + assert_eq!(engine.default_sub_stream, None); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_stream_closed_events_mark_half_close_state() { + use quinn_proto::{Dir, Side, StreamId}; + + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + let data = StreamId::new(Side::Client, Dir::Bi, 1); + engine + .data_streams + .insert(data, QuicStream::new(1024, 5, 1000)); + + let mut events = Vec::new(); + engine.handle_stream_closed(data, "recv_finished", true, &mut events); + + assert!(matches!( + events.as_slice(), + [MqttEvent::StreamClosed { + stream_id, + reason, + by_peer: true + }] if *stream_id == u64::from(data) && reason == "recv_finished" + )); + assert!(engine.data_streams.get(&data).unwrap().recv_closed); + assert!(!engine.data_streams.get(&data).unwrap().finished); + + events.clear(); + engine.handle_stream_closed(data, "send_finished", false, &mut events); + + assert!(matches!( + events.as_slice(), + [MqttEvent::StreamClosed { + stream_id, + reason, + by_peer: false + }] if *stream_id == u64::from(data) && reason == "send_finished" + )); + assert!(engine.data_streams.get(&data).unwrap().finished); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_control_stream_abort_emits_disconnect() { + use quinn_proto::{Dir, Side, StreamId}; + + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + let control = StreamId::new(Side::Client, Dir::Bi, 0); + engine.control_stream = Some(control); + engine.control_outgoing.extend_from_slice(&[1, 2, 3]); + engine.mqtt_engine.is_connected = true; + + let mut events = Vec::new(); + engine.handle_stream_reset(control, VarInt::from_u32(9), &mut events); + + assert!(matches!( + events.as_slice(), + [ + MqttEvent::StreamReset { + stream_id, + error_code: 9 + }, + MqttEvent::Disconnected(None) + ] if *stream_id == u64::from(control) + )); + assert_eq!(engine.control_stream, None); + assert!(engine.control_outgoing.is_empty()); + assert!(!engine.mqtt_engine.is_connected()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_control_stream_stop_and_peer_close_emit_disconnect() { + use quinn_proto::{Dir, Side, StreamId}; + + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + let control = StreamId::new(Side::Client, Dir::Bi, 0); + engine.control_stream = Some(control); + engine.control_outgoing.extend_from_slice(&[1, 2, 3]); + engine.control_finishing = true; + engine.mqtt_engine.is_connected = true; + + let mut events = Vec::new(); + engine.handle_stream_stopped(control, VarInt::from_u32(11), &mut events); + + assert!(matches!( + events.as_slice(), + [ + MqttEvent::StreamStopped { + stream_id, + error_code: 11 + }, + MqttEvent::Disconnected(None) + ] if *stream_id == u64::from(control) + )); + assert_eq!(engine.control_stream, None); + assert!(engine.control_outgoing.is_empty()); + assert!(!engine.control_finishing); + assert!(engine.control_finished); + assert!(!engine.mqtt_engine.is_connected()); + + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + engine.control_stream = Some(control); + engine.mqtt_engine.is_connected = true; + + events.clear(); + engine.handle_stream_closed(control, "recv_finished", true, &mut events); + + assert!(matches!( + events.as_slice(), + [ + MqttEvent::StreamClosed { + stream_id, + reason, + by_peer: true + }, + MqttEvent::Disconnected(None) + ] if *stream_id == u64::from(control) && reason == "recv_finished" + )); + assert_eq!(engine.control_stream, None); + assert!(engine.control_finished); + assert!(!engine.mqtt_engine.is_connected()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_controls_require_connection() { + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + + // reconnect before any connect() has stored config. + assert!(engine.reconnect(Instant::now()).is_err()); + + // Connection-level / stream-level controls error without a connection. + assert!(engine.close(0, b"bye").is_err()); + assert!(engine.quic_ping().is_err()); + assert!(engine.finish_stream(0).is_err()); + assert!(engine.reset_stream(0, 1).is_err()); + assert!(engine.stop_stream(0, 1).is_err()); + + // close_silent is always safe and clears state. + engine.close_silent(); + + // ping() is rejected until the MQTT session is connected, so a PINGREQ + // can never be queued ahead of CONNECT on a fresh transport. + assert!(engine.ping().is_err()); + + // Once connected, ping() queues a PINGREQ. + engine.mqtt_engine.is_connected = true; + engine.ping().unwrap(); + assert!( + !engine.mqtt_engine.outgoing_buffer.is_empty(), + "ping() must queue a PINGREQ once connected" + ); + + // Rejected again once the control stream is finished. + engine.control_finished = true; + assert!(engine.ping().is_err()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_notify_local_address_changed_noop_and_connected() { + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + + assert!(engine.notify_local_address_changed().is_ok()); + assert!(engine.connection.is_none()); + + engine + .connect( + "127.0.0.1:4433".parse().unwrap(), + "localhost", + quic_test_crypto_config(), + Instant::now(), + ) + .unwrap(); + assert!(engine.connection.is_some()); + assert!(engine.notify_local_address_changed().is_ok()); + assert!(engine.connection.is_some()); + } + + #[cfg(feature = "quic-proto")] + fn quic_test_crypto_config() -> rustls::ClientConfig { + #[cfg(feature = "quic-proto-openssl")] + let _ = rustls_openssl::default_provider().install_default(); + #[cfg(not(feature = "quic-proto-openssl"))] + let _ = rustls::crypto::ring::default_provider().install_default(); + + rustls::ClientConfig::builder() + .with_root_certificates(rustls::RootCertStore::empty()) + .with_no_client_auth() + } + + #[cfg(feature = "quic-proto")] + fn connected_quic_engine_without_peer_limits() -> QuicMqttEngine { + let mut engine = QuicMqttEngine::new( + MqttClientOptions::builder() + .max_outgoing_packet_count(1) + .build(), + ) + .unwrap(); + engine + .connect( + "127.0.0.1:4433".parse().unwrap(), + "localhost", + quic_test_crypto_config(), + Instant::now(), + ) + .unwrap(); + engine.mqtt_engine.is_connected = true; + engine + } + + #[cfg(feature = "quic-proto")] + fn mqtt_connected_quic_engine() -> QuicMqttEngine { + let mut engine = QuicMqttEngine::new( + MqttClientOptions::builder() + .max_outgoing_packet_count(1) + .build(), + ) + .unwrap(); + engine.mqtt_engine.connect(); + let _ = engine.mqtt_engine.take_outgoing(); + let events = engine + .mqtt_engine + .handle_incoming(&[0x20, 0x03, 0x00, 0x00, 0x00]); + assert!(matches!(events.as_slice(), [MqttEvent::Connected(_)])); + engine + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_quic_data_stream_api_routes_commands_and_reuses_defaults() { + use quinn_proto::{Dir, Side, StreamId}; + + // GIVEN: A QUIC connection exists, but no peer transport parameters have + // been received to grant client-initiated bidirectional stream capacity. + let mut engine = connected_quic_engine_without_peer_limits(); + + // WHEN: The caller asks the public API to open a new data stream. + // THEN: The request fails instead of creating untracked stream state. + assert!(matches!( + engine.open_data_stream(), + Err(MqttClientError::InternalError { .. }) + )); + + // GIVEN: A connected MQTT session with one known QUIC data stream. + let mut engine = mqtt_connected_quic_engine(); + let stream_id = StreamId::new(Side::Client, Dir::Bi, 1); + engine + .data_streams + .insert(stream_id, QuicStream::new(1024, 5, 1)); + let stream = u64::from(stream_id); + assert_eq!(engine.data_stream_count(), 1); + assert!(engine.control_stream_id().is_none()); + + // WHEN: A publish is targeted at that explicit stream. + // THEN: The MQTT packet is encoded, tracked, and buffered on the stream. + let publish_pid = engine + .publish_on( + stream, + PublishCommand::simple("topic/a", b"payload".to_vec(), 1, false), + ) + .unwrap(); + assert_eq!(publish_pid, Some(1)); + assert!(engine + .data_streams + .values() + .any(|ds| ds.pending_packets == 1 && !ds.outgoing.is_empty())); + + let bad_stream = stream + 4; + // WHEN: A publish targets an unknown stream handle. + // THEN: The API rejects the command before allocating MQTT state. + assert!(matches!( + engine.publish_on( + bad_stream, + PublishCommand::simple("topic/b", b"payload".to_vec(), 0, false) + ), + Err(MqttClientError::InvalidState { .. }) + )); + + // GIVEN: A fresh engine that has not registered the previous stream id. + let mut engine = mqtt_connected_quic_engine(); + // WHEN: A subscribe uses a handle from another engine instance. + // THEN: The handle is treated as unknown and rejected. + let sub_pid = engine + .subscribe_on(stream, SubscribeCommand::single("topic/sub", 0)) + .expect_err("stream handle from another engine must be rejected"); + assert!(matches!(sub_pid, MqttClientError::InvalidState { .. })); + + // GIVEN: The stream handle is registered on this engine. + engine + .data_streams + .insert(stream_id, QuicStream::new(1024, 5, 1)); + // WHEN: A subscribe is targeted at the explicit stream. + // THEN: The SUBSCRIBE packet is allocated and buffered on that stream. + let sub_pid = engine + .subscribe_on(stream, SubscribeCommand::single("topic/sub", 0)) + .unwrap(); + assert_eq!(sub_pid, 1); + + // GIVEN: A connected MQTT session with one known QUIC data stream. + let mut engine = mqtt_connected_quic_engine(); + engine + .data_streams + .insert(stream_id, QuicStream::new(1024, 5, 1)); + // WHEN: An unsubscribe is targeted at the explicit stream. + // THEN: The UNSUBSCRIBE packet is allocated and buffered on that stream. + let unsub_pid = engine + .unsubscribe_on( + stream, + UnsubscribeCommand::from_topics(vec!["topic/sub".to_string()]), + ) + .unwrap(); + assert_eq!(unsub_pid, 1); + + // GIVEN: The default publish stream is already known. + let mut engine = mqtt_connected_quic_engine(); + engine + .data_streams + .insert(stream_id, QuicStream::new(1024, 5, 1)); + engine.default_pub_stream = Some(stream_id); + // WHEN: The high-level publish API is used. + // THEN: It reuses the default publish stream instead of opening another. + let publish_pid = engine + .publish(PublishCommand::simple( + "topic/default", + b"x".to_vec(), + 1, + false, + )) + .unwrap(); + assert_eq!(publish_pid, Some(1)); + let default_pub = engine.default_pub_stream.unwrap(); + assert_eq!(engine.ensure_default_pub_stream().unwrap(), default_pub); + + // GIVEN: The default subscribe stream is already known. + let mut engine = mqtt_connected_quic_engine(); + engine + .data_streams + .insert(stream_id, QuicStream::new(1024, 5, 2)); + engine.default_sub_stream = Some(stream_id); + // WHEN: The high-level subscribe and unsubscribe APIs are used. + // THEN: Both commands reuse the default subscribe stream. + let sub_pid = engine + .subscribe(SubscribeCommand::single("topic/default", 0)) + .unwrap(); + assert_eq!(sub_pid, 1); + let default_sub = engine.default_sub_stream.unwrap(); + let unsub_pid = engine + .unsubscribe(UnsubscribeCommand::from_topics(vec![ + "topic/default".to_string() + ])) + .unwrap(); + assert_eq!(unsub_pid, 2); + assert_eq!(engine.ensure_default_sub_stream().unwrap(), default_sub); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_quic_data_stream_capacity_and_finished_errors() { + use quinn_proto::{Dir, Side, StreamId}; + + let mut engine = mqtt_connected_quic_engine(); + let stream_id = StreamId::new(Side::Client, Dir::Bi, 1); + engine + .data_streams + .insert(stream_id, QuicStream::new(1024, 5, 1)); + let stream = u64::from(stream_id); + + engine + .data_streams + .get_mut(&stream_id) + .unwrap() + .pending_packets = 1; + assert!(matches!( + engine.publish_on( + stream, + PublishCommand::simple("topic/full", b"payload".to_vec(), 1, false) + ), + Err(MqttClientError::BufferFull { + buffer_type, + capacity: 1, + }) if buffer_type == "quic_stream_outgoing" + )); + + let ds = engine.data_streams.get_mut(&stream_id).unwrap(); + ds.pending_packets = 0; + ds.finished = true; + assert!(matches!( + engine.subscribe_on(stream, SubscribeCommand::single("topic/finished", 0)), + Err(MqttClientError::InvalidState { .. }) + )); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_zero_rtt_config_default_and_connect_disables() { + let default = QuicZeroRttConfig::default(); + assert_eq!(default.session_cache_size, 256); + assert!(default.replay_on_reject); + + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + engine.zero_rtt_config = Some(default); + engine.zero_rtt_status = QuicZeroRttStatus::Attempted; + engine + .pending_transport_events + .push_back(MqttEvent::ZeroRttStatusChanged { + status: QuicZeroRttStatus::Attempted, + }); + QuicMqttEngine::journal_stream_open( + &mut engine.early_stream_journal, + StreamId::new(quinn_proto::Side::Client, Dir::Bi, 0), + EarlyStreamRole::Control, + ); + + engine + .connect( + "127.0.0.1:4433".parse().unwrap(), + "localhost", + quic_test_crypto_config(), + Instant::now(), + ) + .unwrap(); + + assert_eq!(engine.zero_rtt_config, None); + assert_eq!(engine.zero_rtt_status(), QuicZeroRttStatus::Disabled); + assert!(engine.pending_transport_events.is_empty()); + assert!(engine.early_stream_journal.is_empty()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_connect_with_zero_rtt_emits_unavailable_without_ticket() { + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + let config = QuicZeroRttConfig { + session_cache_size: 8, + replay_on_reject: false, + }; + + engine + .connect_with_zero_rtt( + "127.0.0.1:4433".parse().unwrap(), + "localhost", + quic_test_crypto_config(), + config, + Instant::now(), + ) + .unwrap(); + + assert_eq!(engine.zero_rtt_config, Some(config)); + assert_eq!(engine.zero_rtt_status(), QuicZeroRttStatus::Unavailable); + assert!(engine.zero_rtt_cache.is_some()); + + let events = engine.take_events(); + assert_eq!(events.len(), 1); + assert!(matches!( + events[0], + MqttEvent::ZeroRttStatusChanged { + status: QuicZeroRttStatus::Unavailable + } + )); + assert!(engine.take_events().is_empty()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_zero_rtt_early_control_open_failure_falls_back_to_unavailable() { + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + engine.zero_rtt_config = Some(QuicZeroRttConfig::default()); + + engine.begin_zero_rtt_attempt(None); + + assert_eq!(engine.zero_rtt_status(), QuicZeroRttStatus::Unavailable); + assert_eq!(engine.control_stream, None); + assert!(engine.control_outgoing.is_empty()); + assert!(engine.early_stream_journal.is_empty()); + + let events = engine.take_events(); + assert_eq!(events.len(), 1); + assert!(matches!( + events[0], + MqttEvent::ZeroRttStatusChanged { + status: QuicZeroRttStatus::Unavailable + } + )); + assert!(engine + .publish(PublishCommand::simple( + "topic", + b"payload".to_vec(), + 0, + false + )) + .is_err()); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_clearable_quic_session_cache_delegates_and_redacts() { + let cache = ClearableQuicSessionCache::new(16); + let name = ServerName::try_from("example.com").unwrap().to_owned(); + + cache.set_kx_hint(name.clone(), NamedGroup::X25519); + assert_eq!(cache.kx_hint(&name), Some(NamedGroup::X25519)); + + let debug = format!("{:?}", cache); + assert!(debug.contains("[redacted]")); + assert!(!debug.contains("example.com")); + + cache.clear(); + assert_eq!(cache.kx_hint(&name), None); + } + + #[cfg(feature = "quic-proto")] + #[test] + fn test_early_stream_journal_records_enqueue_not_write() { + let mut engine = QuicMqttEngine::new(MqttClientOptions::builder().build()).unwrap(); + engine.zero_rtt_status = QuicZeroRttStatus::Attempted; + + let control = StreamId::new(quinn_proto::Side::Client, Dir::Bi, 0); + let data = StreamId::new(quinn_proto::Side::Client, Dir::Bi, 1); + engine.control_stream = Some(control); + QuicMqttEngine::journal_stream_open( + &mut engine.early_stream_journal, + control, + EarlyStreamRole::Control, + ); + QuicMqttEngine::journal_stream_open( + &mut engine.early_stream_journal, + data, + EarlyStreamRole::DefaultPub, + ); + engine + .data_streams + .insert(data, QuicStream::new(1024, 5, 10)); + + QuicMqttEngine::append_control_outgoing_bytes( + &mut engine.control_outgoing, + &mut engine.early_stream_journal, + engine.zero_rtt_status, + engine.control_stream, + &[0x10, 0x00], + 1, + ); + engine.enqueue_on_stream(data, &[0x30, 0x01, b'x']).unwrap(); + + let control_entry = engine + .early_stream_journal + .iter() + .find(|entry| entry.original_stream_id == control) + .unwrap(); + assert_eq!(control_entry.bytes, vec![0x10, 0x00]); + assert_eq!(control_entry.packet_count, 1); + + let data_entry = engine + .early_stream_journal + .iter() + .find(|entry| entry.original_stream_id == data) + .unwrap(); + assert_eq!(data_entry.bytes, vec![0x30, 0x01, b'x']); + assert_eq!(data_entry.packet_count, 1); + + // Simulate a partial QUIC write draining the stream buffer. The replay + // journal must keep the full original enqueue bytes. + engine + .data_streams + .get_mut(&data) + .unwrap() + .outgoing + .drain(..1); + let data_entry = engine + .early_stream_journal + .iter() + .find(|entry| entry.original_stream_id == data) + .unwrap(); + assert_eq!(data_entry.bytes, vec![0x30, 0x01, b'x']); + } } diff --git a/src/mqtt_client/inflight.rs b/src/mqtt_client/inflight.rs index a373288b..c708c459 100644 --- a/src/mqtt_client/inflight.rs +++ b/src/mqtt_client/inflight.rs @@ -19,6 +19,11 @@ pub struct InflightEntry { pub retry_count: u32, /// QoS level (1 or 2) pub qos: u8, + /// Logical channel the packet was originally sent on (e.g. a QUIC data + /// stream handle). `None` for single-stream transports (TCP/TLS, QUIC control + /// stream). Retransmissions are routed back onto this same channel so the QoS + /// 1/2 handshake never crosses streams. + pub stream: Option, } /// A queue for managing inflight QoS 1 and QoS 2 messages. @@ -71,12 +76,25 @@ impl InflightQueue { } } - /// Push a message into the inflight queue + /// Push a message into the inflight queue (single-stream / unrouted). pub fn push( &mut self, packet_id: u16, packet: MqttPacket, qos: u8, + ) -> Result<(), MqttClientError> { + self.push_with_stream(packet_id, packet, qos, None) + } + + /// Push a message into the inflight queue, recording the logical channel + /// (e.g. QUIC data stream) it was sent on so retransmissions can be routed + /// back onto the same channel. + pub fn push_with_stream( + &mut self, + packet_id: u16, + packet: MqttPacket, + qos: u8, + stream: Option, ) -> Result<(), MqttClientError> { // Validation for PUBLISH should happen in MqttEngine before calling push, // but we keep a check here as a safety measure. @@ -100,6 +118,7 @@ impl InflightQueue { sent_at: now, retry_count: 0, qos, + stream, }; self.entries.insert(packet_id, entry); @@ -126,8 +145,19 @@ impl InflightQueue { } } - /// Get all expired messages that need retransmission (MQTT 3.1.1 only) + /// Get all expired messages that need retransmission (MQTT 3.1.1 only). + /// This is a BAD design in MQTT 3.1.1 pub fn get_expired(&mut self, now: Instant) -> Vec { + self.get_expired_with_stream(now) + .into_iter() + .map(|(packet, _stream)| packet) + .collect() + } + + /// Like [`get_expired`](Self::get_expired) but also reports the logical + /// channel each packet was originally sent on, so the caller can retransmit + /// it on the same channel (e.g. the originating QUIC data stream). + pub fn get_expired_with_stream(&mut self, now: Instant) -> Vec<(MqttPacket, Option)> { let mut expired = Vec::new(); // MQTT v5.0: MUST NOT retransmit while connection is active @@ -147,7 +177,7 @@ impl InflightQueue { if entry.sent_at == sent_at { entry.retry_count += 1; entry.sent_at = now; - expired.push(entry.packet.clone()); + expired.push((entry.packet.clone(), entry.stream)); // Re-queue for next timeout self.deadline_queue.push_back((pid, now)); } diff --git a/src/mqtt_client/mod.rs b/src/mqtt_client/mod.rs index cce44529..9895cdc1 100644 --- a/src/mqtt_client/mod.rs +++ b/src/mqtt_client/mod.rs @@ -29,6 +29,8 @@ pub use commands::{ SubscribeCommand, SubscribeCommandBuilder, UnsubscribeCommand, }; pub use engine::{MqttEngine, MqttEvent, MqttMessage}; +#[cfg(feature = "quic-proto")] +pub use engine::{QuicMqttEngine, QuicZeroRttConfig, QuicZeroRttStatus}; pub use error::{MqttClientError, MqttClientResult}; pub use no_io_client::NoIoMqttClient; pub use opts::{MqttClientOptions, MqttClientOptionsBuilder}; diff --git a/src/mqtt_client/opts.rs b/src/mqtt_client/opts.rs index 02435164..806f0fe6 100644 --- a/src/mqtt_client/opts.rs +++ b/src/mqtt_client/opts.rs @@ -126,6 +126,14 @@ pub struct MqttClientOptions { /// - For high-connection-count scenarios with small packets, use 1500 (1 MTU) /// to reduce per-connection memory overhead. pub parser_buffer_size: usize, + + /// Whether the engine automatically sends MQTT PINGREQ keep-alive packets. + /// - Default: true + /// - Set to `false` to suppress automatic PINGREQ even when `keep_alive > 0`. + /// The negotiated keep-alive is still advertised in CONNECT; this only stops + /// the client-side timer from emitting PINGREQ, which is useful for testing + /// QUIC/data-stream keep-alive behaviour or sending PINGREQ manually. + pub auto_keepalive: bool, } impl Default for MqttClientOptions { @@ -162,6 +170,7 @@ impl Default for MqttClientOptions { max_event_count: 1000, receive_maximum: 65535, parser_buffer_size: 16384, + auto_keepalive: true, } } } @@ -203,6 +212,16 @@ impl MqttClientOptions { self } + /// Enable or disable automatic MQTT PINGREQ keep-alive (default: enabled). + /// + /// When disabled, the client will not emit PINGREQ on the keep-alive timer + /// even if `keep_alive > 0`; send PINGREQ manually instead. Useful for + /// keep-alive / idle-timeout testing. + pub fn auto_keepalive(mut self, enabled: bool) -> Self { + self.auto_keepalive = enabled; + self + } + /// Set the username for authentication pub fn username(mut self, username: impl Into) -> Self { self.username = Some(username.into()); diff --git a/src/mqtt_client/tokio_async_client.rs b/src/mqtt_client/tokio_async_client.rs index 38e8a8bb..1f3ed546 100644 --- a/src/mqtt_client/tokio_async_client.rs +++ b/src/mqtt_client/tokio_async_client.rs @@ -2,6 +2,8 @@ use std::collections::HashMap; use std::io; +#[cfg(feature = "quic")] +use std::net::SocketAddr; use std::pin::Pin; use std::time::{Duration, Instant}; @@ -256,6 +258,10 @@ pub struct TokioAsyncClientConfig { /// Only used when connecting via quic:// URLs (requires `quic` feature) #[cfg(feature = "quic")] pub quic_datagram_receive_buffer_size: usize, + /// QUIC transport: Local UDP address to bind before connecting + /// Only used when connecting via quic:// URLs (requires `quic` feature) + #[cfg(feature = "quic")] + pub quic_local_bind_addr: Option, /// QUIC transport: Enable TLS key logging (writes to SSLKEYLOGFILE) /// Only used when connecting via quic:// URLs (requires `quic` feature) #[cfg(feature = "quic")] @@ -301,6 +307,8 @@ impl Default for TokioAsyncClientConfig { #[cfg(feature = "quic")] quic_datagram_receive_buffer_size: 0, // Disable datagram #[cfg(feature = "quic")] + quic_local_bind_addr: None, // Use OS-selected source address by default + #[cfg(feature = "quic")] quic_enable_key_log: false, // Disable key logging by default #[cfg(feature = "rustls-tls")] tls_enable_key_log: false, // Disable key logging by default @@ -775,6 +783,16 @@ impl ConfigBuilder { self } + /// Bind QUIC connections to a local UDP address before connecting. + /// + /// Use port `0` to let the OS assign an ephemeral source port. + /// Only applies when connecting via quic:// URLs (requires `quic` feature) + #[cfg(feature = "quic")] + pub fn quic_local_bind_addr(mut self, addr: SocketAddr) -> Self { + self.config.quic_local_bind_addr = Some(addr); + self + } + /// Enable TLS key logging for QUIC connections /// /// When enabled, TLS session keys are written to the file specified by @@ -1766,9 +1784,17 @@ impl TokioClientWorker { } self.event_handler.on_unsubscribed(&res).await; } + MqttEvent::PublishReceived { .. } => { + // Low-level delivery metadata; MessageReceived carries the + // application payload for this client. + } MqttEvent::MessageReceived(publish) => { self.event_handler.on_message_received(&publish).await; } + MqttEvent::PubRelReceived { .. } => { + // Low-level manual QoS2 acknowledgement detail; the tokio + // client continues to use the engine's normal auto-ack path. + } MqttEvent::PingResponse(res) => { if let Some(tx) = self.pending_ping.take() { let _ = tx.send(res.clone()); @@ -1778,6 +1804,20 @@ impl TokioClientWorker { MqttEvent::Error(err) => { self.event_handler.on_error(&err).await; } + MqttEvent::TransportClosed { .. } => { + // QUIC-transport-specific detail; the tokio (TCP/TLS) client + // surfaces loss via Disconnected/ReconnectNeeded instead. + } + MqttEvent::StreamClosed { .. } + | MqttEvent::StreamReset { .. } + | MqttEvent::StreamStopped { .. } => { + // QUIC-stream-specific detail; this client does not expose + // the Sans-I/O QUIC multi-stream API. + } + MqttEvent::ZeroRttStatusChanged { .. } => { + // QUIC-transport-specific detail; this client does not expose + // the Sans-I/O QUIC 0-RTT API. + } MqttEvent::ReconnectNeeded => { // Engine requesting action. // If we have a stream, it means we timed out -> handle_connection_lost (to clean up) @@ -2125,6 +2165,10 @@ impl TokioClientWorker { if self.config.quic_enable_key_log { builder = builder.enable_key_log(true); } + + if let Some(local_bind_addr) = self.config.quic_local_bind_addr { + builder = builder.local_bind_addr(local_bind_addr); + } let cfg = builder.build(); let transport = QuicTransport::connect_with_config(addr, cfg) diff --git a/src/mqtt_client/tokio_quic_client.rs b/src/mqtt_client/tokio_quic_client.rs index d97bcb7b..e0d054b6 100644 --- a/src/mqtt_client/tokio_quic_client.rs +++ b/src/mqtt_client/tokio_quic_client.rs @@ -4,7 +4,7 @@ use super::commands::{PublishCommand, SubscribeCommand, UnsubscribeCommand}; use super::engine::{MqttEvent, QuicMqttEngine}; use super::opts::MqttClientOptions; use std::collections::VecDeque; -use std::net::SocketAddr; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::time::{Duration, Instant}; use tokio::net::UdpSocket; use tokio::sync::{mpsc, oneshot}; @@ -24,8 +24,13 @@ pub enum QuicCommand { server_addr: SocketAddr, server_name: String, crypto: Box, + local_bind_addr: Option, resp: oneshot::Sender>, }, + Rebind { + local_addr: SocketAddr, + resp: oneshot::Sender>, + }, Publish { cmd: PublishCommand, resp: oneshot::Sender>>, @@ -68,6 +73,17 @@ impl TokioQuicMqttClient { server_addr: SocketAddr, server_name: String, crypto: rustls::ClientConfig, + ) -> CommandResult<()> { + self.connect_with_bind(server_addr, server_name, crypto, None) + .await + } + + pub async fn connect_with_bind( + &self, + server_addr: SocketAddr, + server_name: String, + crypto: rustls::ClientConfig, + local_bind_addr: Option, ) -> CommandResult<()> { let (tx, rx) = oneshot::channel(); self.command_tx @@ -75,6 +91,7 @@ impl TokioQuicMqttClient { server_addr, server_name, crypto: Box::new(crypto), + local_bind_addr, resp: tx, }) .await @@ -84,6 +101,20 @@ impl TokioQuicMqttClient { .and_then(|r| r) } + pub async fn rebind(&self, local_addr: SocketAddr) -> CommandResult { + let (tx, rx) = oneshot::channel(); + self.command_tx + .send(QuicCommand::Rebind { + local_addr, + resp: tx, + }) + .await + .map_err(|_| "Failed to send rebind command")?; + rx.await + .map_err(|_| "Command response channel dropped".into()) + .and_then(|r| r) + } + pub async fn publish(&self, cmd: PublishCommand) -> CommandResult> { let (tx, rx) = oneshot::channel(); self.command_tx @@ -148,10 +179,10 @@ async fn run_engine_loop( tokio::select! { Some(cmd) = command_rx.recv() => { match cmd { - QuicCommand::Connect { server_addr, server_name, crypto, resp } => { + QuicCommand::Connect { server_addr, server_name, crypto, local_bind_addr, resp } => { let res = (|| -> CommandResult<()> { - //@TODO: bind to a specific address - let std_socket = std::net::UdpSocket::bind("0.0.0.0:0")?; + let bind_addr = bind_addr_for_peer(server_addr, local_bind_addr)?; + let std_socket = std::net::UdpSocket::bind(bind_addr)?; std_socket.set_nonblocking(true)?; socket = Some(UdpSocket::from_std(std_socket)?); engine.connect(server_addr, &server_name, *crypto, Instant::now())?; @@ -159,6 +190,17 @@ async fn run_engine_loop( })(); let _ = resp.send(res); } + QuicCommand::Rebind { local_addr, resp } => { + let res = (|| -> CommandResult { + let std_socket = std::net::UdpSocket::bind(local_addr)?; + std_socket.set_nonblocking(true)?; + let bound_addr = std_socket.local_addr()?; + socket = Some(UdpSocket::from_std(std_socket)?); + engine.notify_local_address_changed()?; + Ok(bound_addr) + })(); + let _ = resp.send(res); + } QuicCommand::Publish { cmd, resp } => { let _ = resp.send(engine.publish(cmd).map_err(|e| e.into())); } @@ -220,3 +262,26 @@ async fn run_engine_loop( } } } + +fn bind_addr_for_peer( + peer: SocketAddr, + requested: Option, +) -> CommandResult { + if let Some(addr) = requested { + if addr.is_ipv4() != peer.is_ipv4() { + return Err(format!( + "QUIC local bind address family ({}) does not match peer address family ({})", + addr, peer + ) + .into()); + } + return Ok(addr); + } + + let ip = if peer.is_ipv4() { + IpAddr::V4(Ipv4Addr::UNSPECIFIED) + } else { + IpAddr::V6(Ipv6Addr::UNSPECIFIED) + }; + Ok(SocketAddr::new(ip, 0)) +} diff --git a/src/mqtt_client/transport/quic.rs b/src/mqtt_client/transport/quic.rs index 2866d98e..22a08188 100644 --- a/src/mqtt_client/transport/quic.rs +++ b/src/mqtt_client/transport/quic.rs @@ -6,6 +6,7 @@ mod imp { use super::super::{Transport, TransportError}; use async_trait::async_trait; + use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr}; use std::pin::Pin; use std::task::{Context, Poll}; use tokio::io::{AsyncRead, AsyncWrite, ReadBuf}; @@ -95,6 +96,8 @@ mod imp { pub datagram_receive_buffer_size: usize, /// Enable TLS key logging (writes to the file specified by SSLKEYLOGFILE env var) pub enable_key_log: bool, + /// Optional local UDP address to bind before connecting. + pub local_bind_addr: Option, } /// Builder for `QuicConfig` to simplify ergonomic construction. @@ -107,6 +110,7 @@ mod imp { insecure_skip_verify: bool, datagram_receive_buffer_size: usize, enable_key_log: bool, + local_bind_addr: Option, } impl QuicConfigBuilder { @@ -286,6 +290,20 @@ mod imp { self } + /// Bind the QUIC UDP socket to a specific local address before connecting. + /// + /// Use port `0` to let the OS assign an ephemeral source port. + pub fn local_bind_addr(mut self, addr: SocketAddr) -> Self { + self.local_bind_addr = Some(addr); + self + } + + /// Bind the QUIC UDP socket to a specific local IP with an ephemeral port. + pub fn local_bind_ip(mut self, ip: IpAddr) -> Self { + self.local_bind_addr = Some(SocketAddr::new(ip, 0)); + self + } + /// Finalize the builder into a `QuicConfig`. pub fn build(self) -> QuicConfig { QuicConfig { @@ -297,6 +315,7 @@ mod imp { insecure_skip_verify: self.insecure_skip_verify, datagram_receive_buffer_size: self.datagram_receive_buffer_size, enable_key_log: self.enable_key_log, + local_bind_addr: self.local_bind_addr, } } } @@ -313,6 +332,7 @@ mod imp { insecure_skip_verify: false, datagram_receive_buffer_size: 0, enable_key_log: false, + local_bind_addr: None, } } } @@ -454,8 +474,10 @@ mod imp { trpt_cfg.datagram_receive_buffer_size(Some(cfg.datagram_receive_buffer_size)); quinn_crypto.transport_config(Arc::new(trpt_cfg)); - // Create an endpoint bound to an ephemeral UDP port - let socket = std::net::UdpSocket::bind("0.0.0.0:0").map_err(|e| { + let bind_addr = bind_addr_for_peer(peer, cfg.local_bind_addr)?; + + // Create an endpoint bound to an ephemeral UDP port, or the caller's requested source address. + let socket = std::net::UdpSocket::bind(bind_addr).map_err(|e| { TransportError::ConnectionFailed(format!("QUIC endpoint create failed: {}", e)) })?; let runtime = quinn::default_runtime() @@ -494,6 +516,21 @@ mod imp { recv, }) } + + /// Rebind the underlying QUIC endpoint to a new UDP socket. + /// + /// This changes the source address used for future packets on all active + /// connections managed by the endpoint. Use port `0` to let the OS assign + /// a new ephemeral source port. + pub fn rebind(&self, local_addr: SocketAddr) -> Result { + let socket = std::net::UdpSocket::bind(local_addr).map_err(|e| { + TransportError::ConnectionFailed(format!("QUIC endpoint rebind failed: {}", e)) + })?; + self.endpoint.rebind(socket).map_err(|e| { + TransportError::ConnectionFailed(format!("QUIC endpoint rebind failed: {}", e)) + })?; + self.endpoint.local_addr().map_err(TransportError::Io) + } } #[async_trait] @@ -530,6 +567,28 @@ mod imp { endpoint.local_addr().ok().map(|s| s.to_string()) } + fn bind_addr_for_peer( + peer: SocketAddr, + requested: Option, + ) -> Result { + if let Some(addr) = requested { + if addr.is_ipv4() != peer.is_ipv4() { + return Err(TransportError::InvalidAddress(format!( + "QUIC local bind address family ({}) does not match peer address family ({})", + addr, peer + ))); + } + return Ok(addr); + } + + let ip = if peer.is_ipv4() { + IpAddr::V4(Ipv4Addr::UNSPECIFIED) + } else { + IpAddr::V6(Ipv6Addr::UNSPECIFIED) + }; + Ok(SocketAddr::new(ip, 0)) + } + impl AsyncRead for QuicTransport { fn poll_read( mut self: Pin<&mut Self>, diff --git a/swift/Examples/QuicClientExample/main.swift b/swift/Examples/QuicClientExample/main.swift index 8fc83f7f..527cd0fc 100644 --- a/swift/Examples/QuicClientExample/main.swift +++ b/swift/Examples/QuicClientExample/main.swift @@ -199,6 +199,12 @@ while nowMs(since: engineStartMs) < runDurationMs { print("ℹ Reconnect attempt \(attempt) in \(delayMs)ms") case .pingResponse(let success): print("ℹ Ping response: success=\(success)") + case .streamClosed(let streamId, let reason, let byPeer): + print("ℹ Stream \(streamId) closed: \(reason), byPeer=\(byPeer)") + case .streamReset(let streamId, let errorCode): + print("ℹ Stream \(streamId) reset: errorCode=\(errorCode)") + case .streamStopped(let streamId, let errorCode): + print("ℹ Stream \(streamId) stopped: errorCode=\(errorCode)") case .unsubscribed(_): break } diff --git a/swift/Examples/TcpClientExample/main.swift b/swift/Examples/TcpClientExample/main.swift index 0e5556c4..7abc6e87 100644 --- a/swift/Examples/TcpClientExample/main.swift +++ b/swift/Examples/TcpClientExample/main.swift @@ -122,6 +122,12 @@ func handleEvent(_ event: MqttEventFfi, engine: MqttEngineFfi, fd: Int32, print("ℹ Reconnect attempt \(attempt) in \(delayMs)ms") case .pingResponse(let success): print("ℹ Ping response: success=\(success)") + case .streamClosed(_, _, _): + break + case .streamReset(_, _): + break + case .streamStopped(_, _): + break case .unsubscribed(_): break } diff --git a/tests/engine_logic_tests.rs b/tests/engine_logic_tests.rs index fee9c56a..cca4db7f 100644 --- a/tests/engine_logic_tests.rs +++ b/tests/engine_logic_tests.rs @@ -537,8 +537,15 @@ fn test_incoming_publish_qos1() { )); let events = engine.handle_incoming(&publish.to_bytes().unwrap()); - assert_eq!(events.len(), 1); - match &events[0] { + assert_eq!(events.len(), 2); + assert!(matches!( + events[0], + MqttEvent::PublishReceived { + packet_id: Some(100), + stream: None + } + )); + match &events[1] { MqttEvent::MessageReceived(p) => { assert_eq!(p.topic_name, "t"); assert_eq!(p.payload, b"data"); @@ -673,8 +680,15 @@ fn test_incoming_publish_qos2() { )); let events = engine.handle_incoming(&publish.to_bytes().unwrap()); - assert_eq!(events.len(), 1); - match &events[0] { + assert_eq!(events.len(), 2); + match events[0] { + MqttEvent::PublishReceived { packet_id, stream } => { + assert_eq!(packet_id, Some(pid)); + assert_eq!(stream, None); + } + _ => panic!("Expected PublishReceived event"), + } + match &events[1] { MqttEvent::MessageReceived(p) => { assert_eq!(p.qos, 2); assert_eq!(p.packet_id, Some(pid)); diff --git a/tests/key_log_tests.rs b/tests/key_log_tests.rs index 93da624c..d47cb024 100644 --- a/tests/key_log_tests.rs +++ b/tests/key_log_tests.rs @@ -15,6 +15,26 @@ fn quic_config_enable_key_log() { assert!(!cfg_default.enable_key_log); } +/// Verify that QuicConfig builder wires local_bind_addr into the built config. +#[cfg(feature = "quic")] +#[test] +fn quic_config_local_bind_addr() { + use flowsdk::mqtt_client::transport::quic::QuicConfig; + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + + let bind_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0); + let cfg = QuicConfig::builder().local_bind_addr(bind_addr).build(); + assert_eq!(cfg.local_bind_addr, Some(bind_addr)); + + let cfg_from_ip = QuicConfig::builder() + .local_bind_ip(IpAddr::V4(Ipv4Addr::LOCALHOST)) + .build(); + assert_eq!(cfg_from_ip.local_bind_addr, Some(bind_addr)); + + let cfg_default = QuicConfig::builder().build(); + assert_eq!(cfg_default.local_bind_addr, None); +} + /// Verify that RustlsTlsConfig builder wires enable_key_log into the built config. #[cfg(feature = "rustls-tls")] #[test] @@ -66,6 +86,23 @@ fn tokio_async_config_quic_key_log() { assert!(!config_default.quic_enable_key_log); } +/// Verify TokioAsyncClientConfig builder exposes quic_local_bind_addr. +#[cfg(feature = "quic")] +#[test] +fn tokio_async_config_quic_local_bind_addr() { + use flowsdk::mqtt_client::TokioAsyncClientConfig; + use std::net::{IpAddr, Ipv4Addr, SocketAddr}; + + let bind_addr = SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0); + let config = TokioAsyncClientConfig::builder() + .quic_local_bind_addr(bind_addr) + .build(); + assert_eq!(config.quic_local_bind_addr, Some(bind_addr)); + + let config_default = TokioAsyncClientConfig::builder().build(); + assert_eq!(config_default.quic_local_bind_addr, None); +} + /// Verify TokioAsyncClientConfig builder exposes tls_enable_key_log. #[cfg(feature = "rustls-tls")] #[test]