From c8cf39db127335ec226093b25d4f1b456382f865 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Fabr=C3=ADcio=20Bracht?= Date: Mon, 21 Sep 2026 10:56:06 -0300 Subject: [PATCH] fix quic partial-frame desync via per-peer writer task --- Cargo.lock | 175 +++++++++++++++- crates/mqdb-cluster/Cargo.toml | 1 + .../src/cluster/quic_transport.rs | 111 +++++----- crates/mqdb-cluster/src/cluster/transport.rs | 2 + .../tests/quic_transport_framing.rs | 191 ++++++++++++++++++ 5 files changed, 424 insertions(+), 56 deletions(-) create mode 100644 crates/mqdb-cluster/tests/quic_transport_framing.rs diff --git a/Cargo.lock b/Cargo.lock index ae97329..ecc86e1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -94,6 +94,45 @@ dependencies = [ "password-hash", ] +[[package]] +name = "asn1-rs" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f43a50ac4fdca5df8e885c21b835997f0a1cdee65494a6847694a98652d9d8" +dependencies = [ + "asn1-rs-derive", + "asn1-rs-impl", + "displaydoc", + "nom", + "num-traits", + "rusticata-macros", + "thiserror 2.0.18", + "time", +] + +[[package]] +name = "asn1-rs-derive" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3109e49b1e4909e9db6515a30c633684d68cdeaa252f215214cb4fa1a5bfee2c" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "asn1-rs-impl" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b18050c2cd6fe86c3a76584ef5e0baf286d038cda203eb6223df2cc413565f7" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "async-trait" version = "0.1.89" @@ -146,6 +185,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64ct" version = "1.8.3" @@ -173,6 +218,15 @@ dependencies = [ "syn", ] +[[package]] +name = "bit-vec" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b71798fca2c1fe1086445a7258a4bc81e6e49dcd24c8d0dd9a1e57395b603f51" +dependencies = [ + "serde", +] + [[package]] name = "bitflags" version = "2.11.0" @@ -502,6 +556,20 @@ version = "2.10.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea" +[[package]] +name = "der-parser" +version = "10.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07da5016415d5a3c4dd39b11ed26f915f52fc4e0dc197d87908bc916e51bc1a6" +dependencies = [ + "asn1-rs", + "displaydoc", + "nom", + "num-bigint", + "num-traits", + "rusticata-macros", +] + [[package]] name = "deranged" version = "0.5.8" @@ -1004,7 +1072,7 @@ version = "0.1.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "futures-channel", "futures-util", @@ -1285,10 +1353,10 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0529410abe238729a60b108898784df8984c87f6054c9c4fcacc47e4803c1ce1" dependencies = [ "aws-lc-rs", - "base64", + "base64 0.22.1", "getrandom 0.2.17", "js-sys", - "pem", + "pem 3.0.6", "serde", "serde_json", "signature", @@ -1407,6 +1475,12 @@ version = "0.3.17" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + [[package]] name = "mio" version = "1.2.0" @@ -1424,7 +1498,7 @@ version = "0.8.29" dependencies = [ "arc-swap", "argon2", - "base64", + "base64 0.22.1", "comfy-table", "fjall", "flume", @@ -1457,7 +1531,7 @@ dependencies = [ name = "mqdb-cli" version = "0.8.39" dependencies = [ - "base64", + "base64 0.22.1", "bebytes", "clap", "comfy-table", @@ -1497,6 +1571,7 @@ dependencies = [ "mqdb-vault", "mqtt5", "quinn", + "rcgen", "regex", "ring", "rustls", @@ -1539,7 +1614,7 @@ dependencies = [ name = "mqdb-vault" version = "0.1.5" dependencies = [ - "base64", + "base64 0.22.1", "mqdb-agent", "mqdb-core", "mqtt5", @@ -1560,7 +1635,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "c61d9de73200ed2112ed25ebcdd5e74b22aa5c89065fd3554e111e3353e9d5f0" dependencies = [ "argon2", - "base64", + "base64 0.22.1", "bebytes", "bytes", "flume", @@ -1618,6 +1693,16 @@ dependencies = [ "web-time", ] +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + [[package]] name = "nu-ansi-term" version = "0.50.3" @@ -1661,6 +1746,15 @@ dependencies = [ "autocfg", ] +[[package]] +name = "oid-registry" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12f40cff3dde1b6087cc5d5f5d4d65712f34016a03ed60e9c08dcc392736b5b7" +dependencies = [ + "asn1-rs", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -1796,7 +1890,17 @@ version = "3.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" dependencies = [ - "base64", + "base64 0.22.1", + "serde_core", +] + +[[package]] +name = "pem" +version = "4.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d354a98a3d1251555de99e8fdd8afda05573c31b82f59063a7b0a29b5527f120" +dependencies = [ + "base64 0.23.1", "serde_core", ] @@ -2065,6 +2169,20 @@ dependencies = [ "rand_core 0.10.1", ] +[[package]] +name = "rcgen" +version = "0.14.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8774e05a7d0de114588e6a28fe7e71694b82614ed569d86d8b389dfbc98b8ad8" +dependencies = [ + "pem 4.0.0", + "ring", + "rustls-pki-types", + "time", + "x509-parser", + "yasna", +] + [[package]] name = "redox_syscall" version = "0.5.18" @@ -2109,7 +2227,7 @@ version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ab3f43e3283ab1488b624b44b0e988d0acea0b3214e694730a055cb6b2efa801" dependencies = [ - "base64", + "base64 0.22.1", "bytes", "encoding_rs", "futures-channel", @@ -2175,6 +2293,15 @@ dependencies = [ "semver", ] +[[package]] +name = "rusticata-macros" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "faf0c4a6ece9950b9abdb62b1cfcf2a68b3b67a10ba445b3bb85be2a293d0632" +dependencies = [ + "nom", +] + [[package]] name = "rustix" version = "1.1.4" @@ -2858,7 +2985,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec" dependencies = [ "async-trait", - "base64", + "base64 0.22.1", "bytes", "http", "http-body", @@ -3577,12 +3704,40 @@ version = "0.6.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" +[[package]] +name = "x509-parser" +version = "0.18.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d43b0f71ce057da06bc0851b23ee24f3f86190b07203dd8f567d0b706a185202" +dependencies = [ + "asn1-rs", + "data-encoding", + "der-parser", + "lazy_static", + "nom", + "oid-registry", + "ring", + "rusticata-macros", + "thiserror 2.0.18", + "time", +] + [[package]] name = "xxhash-rust" version = "0.8.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "fdd20c5420375476fbd4394763288da7eb0cc0b8c11deed431a91562af7335d3" +[[package]] +name = "yasna" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5f6765e852b9b4dc8e2a76843e4d64d1cea8e79bcde0b6901aea8e7c7f08282" +dependencies = [ + "bit-vec", + "time", +] + [[package]] name = "yoke" version = "0.8.2" diff --git a/crates/mqdb-cluster/Cargo.toml b/crates/mqdb-cluster/Cargo.toml index fb3a9a3..7fd1a6e 100644 --- a/crates/mqdb-cluster/Cargo.toml +++ b/crates/mqdb-cluster/Cargo.toml @@ -46,5 +46,6 @@ dev-insecure = ["mqdb-agent/dev-insecure"] http-api = ["mqdb-agent/http-api"] [dev-dependencies] +rcgen = "0.14.10" tempfile.workspace = true tracing-subscriber.workspace = true diff --git a/crates/mqdb-cluster/src/cluster/quic_transport.rs b/crates/mqdb-cluster/src/cluster/quic_transport.rs index 713cd1a..ae3078a 100644 --- a/crates/mqdb-cluster/src/cluster/quic_transport.rs +++ b/crates/mqdb-cluster/src/cluster/quic_transport.rs @@ -11,16 +11,15 @@ use std::net::SocketAddr; use std::path::Path; use std::sync::atomic::{AtomicBool, Ordering}; use std::sync::{Arc, Mutex}; -use std::time::Duration; use tokio::sync::{Notify, RwLock}; use tracing::{debug, error, info, trace, warn}; -const SEND_TIMEOUT_MS: u64 = 5000; const INBOX_CHANNEL_CAPACITY: usize = 16384; +const PEER_SEND_QUEUE_CAPACITY: usize = 1024; struct PeerConnection { _connection: Connection, - send_stream: tokio::sync::Mutex, + writer_tx: flume::Sender>, } const MAX_MESSAGE_SIZE: usize = 10 * 1024 * 1024; @@ -236,27 +235,25 @@ impl QuicDirectTransport { "connected to peer via QUIC" ); - let (send_stream, recv_stream) = connection + let (mut send_stream, recv_stream) = connection .open_bi() .await .map_err(|e| TransportError::SendFailed(format!("failed to open stream: {e}")))?; let header = self.node_id.get().to_be_bytes(); - { - let mut stream = tokio::sync::Mutex::new(send_stream); - stream - .get_mut() - .write_all(&header) - .await - .map_err(|e| TransportError::SendFailed(format!("failed to send header: {e}")))?; - - let peer_conn = PeerConnection { - _connection: connection.clone(), - send_stream: stream, - }; - - self.peers.write().await.insert(peer_id, peer_conn); - } + send_stream + .write_all(&header) + .await + .map_err(|e| TransportError::SendFailed(format!("failed to send header: {e}")))?; + + let (writer_tx, writer_rx) = flume::bounded(PEER_SEND_QUEUE_CAPACITY); + let peer_conn = PeerConnection { + _connection: connection.clone(), + writer_tx, + }; + self.peers.write().await.insert(peer_id, peer_conn); + + tokio::spawn(peer_writer_task(send_stream, writer_rx, peer_id)); let inbox_tx = self.inbox_tx.clone(); let notify = self.message_notify.clone(); @@ -286,43 +283,45 @@ impl QuicDirectTransport { buf } + fn frame_message(&self, message: &ClusterMessage) -> Result, TransportError> { + let payload = self.serialize_message(message); + let len = u32::try_from(payload.len()) + .map_err(|_| TransportError::SendFailed("message too large to frame".to_string()))?; + + let mut frame = Vec::with_capacity(4 + payload.len()); + frame.extend_from_slice(&len.to_be_bytes()); + frame.extend_from_slice(&payload); + Ok(frame) + } + async fn send_to_peer( &self, peer_id: NodeId, message: &ClusterMessage, ) -> Result<(), TransportError> { - let payload = self.serialize_message(message); - - #[allow(clippy::cast_possible_truncation)] - let len_prefix = (payload.len() as u32).to_be_bytes(); + let frame = self.frame_message(message)?; let peers = self.peers.read().await; let peer = peers .get(&peer_id) .ok_or(TransportError::NodeNotFound(peer_id))?; - let mut stream = peer.send_stream.lock().await; - - let timeout = Duration::from_millis(SEND_TIMEOUT_MS); - - tokio::time::timeout(timeout, stream.write_all(&len_prefix)) - .await - .map_err(|_| TransportError::SendFailed("send timeout (length prefix)".to_string()))? - .map_err(|e| TransportError::SendFailed(format!("failed to write length: {e}")))?; - - tokio::time::timeout(timeout, stream.write_all(&payload)) - .await - .map_err(|_| TransportError::SendFailed("send timeout (payload)".to_string()))? - .map_err(|e| TransportError::SendFailed(format!("failed to write payload: {e}")))?; - - trace!( - from = self.node_id.get(), - to = peer_id.get(), - msg_type = message.type_name(), - "sent QUIC message" - ); - - Ok(()) + match peer.writer_tx.try_send(frame) { + Ok(()) => { + trace!( + from = self.node_id.get(), + to = peer_id.get(), + msg_type = message.type_name(), + "queued QUIC message" + ); + Ok(()) + } + Err(flume::TrySendError::Full(_)) => Err(TransportError::SendQueueFull(peer_id)), + Err(flume::TrySendError::Disconnected(_)) => Err(TransportError::SendFailed(format!( + "peer {} writer task ended", + peer_id.get() + ))), + } } } @@ -500,18 +499,38 @@ async fn handle_incoming_connection( info!(peer = peer_node.get(), "accepted incoming QUIC connection"); + let (writer_tx, writer_rx) = flume::bounded(PEER_SEND_QUEUE_CAPACITY); { let peer_conn = PeerConnection { _connection: connection.clone(), - send_stream: tokio::sync::Mutex::new(send_stream), + writer_tx, }; peers.write().await.insert(peer_node, peer_conn); } + tokio::spawn(peer_writer_task(send_stream, writer_rx, peer_node)); + receiver_task(recv_stream, peer_node, inbox_tx, notify, local_node).await; Ok(()) } +async fn peer_writer_task( + mut send_stream: SendStream, + writer_rx: flume::Receiver>, + peer_node: NodeId, +) { + trace!(peer = peer_node.get(), "peer writer task started"); + + while let Ok(frame) = writer_rx.recv_async().await { + if let Err(e) = send_stream.write_all(&frame).await { + warn!(peer = peer_node.get(), error = %e, "peer writer failed, tearing down stream"); + break; + } + } + + debug!(peer = peer_node.get(), "peer writer task ended"); +} + async fn receiver_task( mut recv_stream: RecvStream, peer_node: NodeId, diff --git a/crates/mqdb-cluster/src/cluster/transport.rs b/crates/mqdb-cluster/src/cluster/transport.rs index bf7a4f5..6192fca 100644 --- a/crates/mqdb-cluster/src/cluster/transport.rs +++ b/crates/mqdb-cluster/src/cluster/transport.rs @@ -530,6 +530,7 @@ pub enum TransportError { PartitionNotFound(PartitionId), NetworkPartitioned, SendFailed(String), + SendQueueFull(NodeId), NotConnected, } @@ -540,6 +541,7 @@ impl std::fmt::Display for TransportError { Self::PartitionNotFound(id) => write!(f, "partition {} not found", id.get()), Self::NetworkPartitioned => write!(f, "network partitioned"), Self::SendFailed(msg) => write!(f, "send failed: {msg}"), + Self::SendQueueFull(id) => write!(f, "send queue full for node {}", id.get()), Self::NotConnected => write!(f, "not connected"), } } diff --git a/crates/mqdb-cluster/tests/quic_transport_framing.rs b/crates/mqdb-cluster/tests/quic_transport_framing.rs new file mode 100644 index 0000000..157c7b9 --- /dev/null +++ b/crates/mqdb-cluster/tests/quic_transport_framing.rs @@ -0,0 +1,191 @@ +use std::net::SocketAddr; +use std::sync::Arc; +use std::time::Duration; + +use mqdb_cluster::{ + ClusterMessage, ClusterTransport, Heartbeat, NodeId, QuicDirectTransport, TransportError, +}; +use quinn::{Connection, Endpoint, RecvStream, SendStream, ServerConfig, TransportConfig, VarInt}; +use rcgen::{BasicConstraints, CertificateParams, IsCa, Issuer, KeyPair}; +use rustls::pki_types::{CertificateDer, PrivateKeyDer}; +use tokio::sync::oneshot; + +const STALL_WINDOW: u32 = 16 * 1024; +const SEND_ATTEMPTS: u32 = 8192; + +struct Certs { + ca_pem: String, + leaf_pem: String, + leaf_key_pem: String, + leaf_der: CertificateDer<'static>, + leaf_key_der: Vec, +} + +fn generate_certs() -> Certs { + let mut ca_params = CertificateParams::new(Vec::new()).unwrap(); + ca_params.is_ca = IsCa::Ca(BasicConstraints::Unconstrained); + let ca_key = KeyPair::generate().unwrap(); + let ca_cert = ca_params.self_signed(&ca_key).unwrap(); + let ca_pem = ca_cert.pem(); + + let leaf_params = CertificateParams::new(vec!["localhost".to_string()]).unwrap(); + let leaf_key = KeyPair::generate().unwrap(); + let issuer = Issuer::new(ca_params, ca_key); + let leaf_cert = leaf_params.signed_by(&leaf_key, &issuer).unwrap(); + + Certs { + ca_pem, + leaf_pem: leaf_cert.pem(), + leaf_key_pem: leaf_key.serialize_pem(), + leaf_der: leaf_cert.der().clone(), + leaf_key_der: leaf_key.serialize_der(), + } +} + +fn stalling_server_endpoint(certs: &Certs) -> Endpoint { + let key = PrivateKeyDer::try_from(certs.leaf_key_der.clone()).unwrap(); + let server_crypto = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(vec![certs.leaf_der.clone()], key) + .unwrap(); + + let mut server_config = ServerConfig::with_crypto(Arc::new( + quinn::crypto::rustls::QuicServerConfig::try_from(server_crypto).unwrap(), + )); + + let mut transport = TransportConfig::default(); + transport.stream_receive_window(VarInt::from_u32(STALL_WINDOW)); + server_config.transport_config(Arc::new(transport)); + + Endpoint::server(server_config, "127.0.0.1:0".parse().unwrap()).unwrap() +} + +async fn far_end( + endpoint: Endpoint, + release: oneshot::Receiver<()>, + frame_tx: flume::Sender>, +) { + let connection: Connection = endpoint.accept().await.unwrap().await.unwrap(); + let (send_stream, mut recv): (SendStream, RecvStream) = connection.accept_bi().await.unwrap(); + + let mut header = [0u8; 2]; + recv.read_exact(&mut header).await.unwrap(); + + release.await.ok(); + + loop { + let mut len_buf = [0u8; 4]; + if recv.read_exact(&mut len_buf).await.is_err() { + break; + } + let len = u32::from_be_bytes(len_buf) as usize; + let mut payload = vec![0u8; len]; + if recv.read_exact(&mut payload).await.is_err() { + break; + } + if frame_tx.send(payload).is_err() { + break; + } + } + + drop(send_stream); + drop(connection); +} + +async fn collect_frame(frame_rx: &flume::Receiver>) -> Vec { + tokio::time::timeout(Duration::from_secs(10), frame_rx.recv_async()) + .await + .expect("timed out waiting for a framed message") + .expect("far-end frame channel closed") +} + +#[tokio::test] +async fn backpressure_drops_messages_but_keeps_the_stream_framed() { + let _ = rustls::crypto::ring::default_provider().install_default(); + + let certs = generate_certs(); + let dir = tempfile::tempdir().unwrap(); + let ca_path = dir.path().join("ca.pem"); + let leaf_path = dir.path().join("leaf.pem"); + let leaf_key_path = dir.path().join("leaf.key"); + std::fs::write(&ca_path, &certs.ca_pem).unwrap(); + std::fs::write(&leaf_path, &certs.leaf_pem).unwrap(); + std::fs::write(&leaf_key_path, &certs.leaf_key_pem).unwrap(); + + let endpoint = stalling_server_endpoint(&certs); + let far_addr: SocketAddr = endpoint.local_addr().unwrap(); + + let (release_tx, release_rx) = oneshot::channel(); + let (frame_tx, frame_rx) = flume::unbounded(); + let far_handle = tokio::spawn(far_end(endpoint, release_rx, frame_tx)); + + let local = NodeId::validated(1).unwrap(); + let peer = NodeId::validated(2).unwrap(); + let transport = QuicDirectTransport::new(local); + transport.set_ca_file(ca_path.clone()); + transport + .bind("127.0.0.1:0".parse().unwrap(), &leaf_path, &leaf_key_path) + .await + .unwrap(); + transport.connect_to_peer(peer, far_addr).await.unwrap(); + + let mut ok_count = 0u32; + let mut terminated_by_drop = false; + let mut blocked = false; + for tick in 0..SEND_ATTEMPTS { + let message = ClusterMessage::Heartbeat(Heartbeat::create(local, u64::from(tick))); + match tokio::time::timeout(Duration::from_millis(500), transport.send(peer, message)).await + { + Err(_) => { + blocked = true; + break; + } + Ok(Ok(())) => ok_count += 1, + Ok(Err(TransportError::SendQueueFull(_))) => { + terminated_by_drop = true; + break; + } + Ok(Err(other)) => panic!("unexpected send error: {other}"), + } + } + + assert!( + !blocked, + "a send blocked under backpressure instead of returning immediately" + ); + assert!( + terminated_by_drop, + "expected SendQueueFull once the per-peer queue filled" + ); + assert!( + ok_count > 0, + "expected some sends to be accepted before the queue filled" + ); + + release_tx.send(()).ok(); + + for _ in 0..ok_count { + let frame = collect_frame(&frame_rx).await; + assert!( + frame.len() >= 3, + "frame shorter than a cluster message header" + ); + assert_eq!( + &frame[0..2], + &local.get().to_be_bytes(), + "frame did not begin with the sender node id (stream mis-framed)" + ); + } + + let follow_up = ClusterMessage::Heartbeat(Heartbeat::create(local, u64::MAX)); + transport.send(peer, follow_up).await.unwrap(); + let frame = collect_frame(&frame_rx).await; + assert!(frame.len() >= 3); + assert_eq!( + &frame[0..2], + &local.get().to_be_bytes(), + "post-backpressure frame mis-framed" + ); + + far_handle.abort(); +}