Brought Example up to spec
Some checks failed
CI / checks (push) Failing after 3m29s

This commit is contained in:
Alex Emmet 2026-07-19 00:29:24 +02:00
commit 1b796d0ce7
46 changed files with 1755 additions and 691 deletions

View file

@ -1,4 +1,5 @@
use std::sync::Arc;
use std::time::Instant;
use mtp_common::CommunicationError;
use rustls::{ClientConfig as RustlsClientConfig, RootCertStore, pki_types::pem::PemObject};
@ -139,8 +140,10 @@ pub async fn connect_with_config(
url: &str,
config: ClientConfig,
) -> Result<(Sender, Receiver), CommunicationError> {
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
let connect_started = Instant::now();
mtp_crypto::ensure_crypto_provider();
let config_started = Instant::now();
let client_config = if config.insecure_certificate_verification {
#[cfg(feature = "insecure-tls")]
{
@ -168,14 +171,19 @@ pub async fn connect_with_config(
} else {
configure_client_system_roots(&config.policy)?
};
tracing::debug!(elapsed = ?config_started.elapsed(), "client connect: configure TLS");
let endpoint_started = Instant::now();
let endpoint = Endpoint::client(client_config)
.map_err(|e| CommunicationError::Other(format!("Endpoint creation failed: {}", e)))?;
tracing::debug!(elapsed = ?endpoint_started.elapsed(), "client connect: create endpoint");
let transport_connect_started = Instant::now();
let connection = endpoint
.connect(url)
.await
.map_err(|e| CommunicationError::ConnectingError(e.to_string()))?;
tracing::debug!(elapsed = ?transport_connect_started.elapsed(), "client connect: establish WebTransport session");
let handle = Arc::new(ConnectionHandle::new());
let policy = Arc::new(config.policy);
@ -183,6 +191,7 @@ pub async fn connect_with_config(
let sender = Sender::new(connection.clone(), handle.clone(), policy.clone());
let receiver = Receiver::new(connection, handle, policy);
tracing::debug!(elapsed = ?connect_started.elapsed(), "client connect: complete");
Ok((sender, receiver))
}

View file

@ -471,7 +471,7 @@ impl Sender {
}
#[instrument(skip(self), level = "trace")]
pub fn close(&self) {
pub fn close_immediate(&self) {
info!(target = "mtp.transport", "fire-and-forget close requested");
let connection = self.connection.clone();
let handle = self.handle.clone();
@ -514,7 +514,7 @@ impl Sender {
#[instrument(skip(self), level = "trace")]
/// Initiate a best-effort graceful close and wait for the configured force-close delay.
pub async fn close_and_wait(&self) {
pub async fn close(&self) {
info!(target = "mtp.transport", "graceful close initiated");
let connection = self.connection.clone();
let handle = self.handle.clone();
@ -535,11 +535,11 @@ impl Sender {
match timeout(policy.write_timeout, close_write).await {
Ok(Ok(())) => {}
Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => warn!(
"[Sender] close_and_wait failed: peer sent STOP_SENDING (error code {code})"
),
Ok(Err(e)) => warn!("[Sender] close_and_wait failed: {e}"),
Err(_) => warn!("[Sender] close_and_wait timed out"),
Ok(Err(wtransport::error::StreamWriteError::Stopped(code))) => {
warn!("[Sender] close failed: peer sent STOP_SENDING (error code {code})")
}
Ok(Err(e)) => warn!("[Sender] close failed: {e}"),
Err(_) => warn!("[Sender] close timed out"),
}
} else {
let _ = Self::send_close_frame(&connection, &policy).await;

View file

@ -10,6 +10,7 @@ use crate::{
use mtp_codec::CommunicationValue;
use mtp_common::CommunicationError;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use tokio::sync::{Mutex, RwLock, Semaphore, mpsc};
use tokio::time::timeout;
@ -132,6 +133,21 @@ impl<C: TransportConnection> GenericSender<C> {
self.connection
.close(self.policy.application_close_code, b"mtp-close");
}
/// Finish the current persistent stream.
///
/// This is used by hosts that put the opening/authentication exchange on
/// a persistent stream and then transition to application streams.
pub async fn finish_stream(&self) -> Result<(), CommunicationError> {
let _lock = self.send_lock.lock().await;
let mut stream = self.persistent.lock().await;
let Some(mut stream) = stream.take() else {
return Ok(());
};
timeout(self.policy.write_timeout, stream.finish())
.await
.map_err(|_| CommunicationError::StreamError)?
}
pub fn is_closed(&self) -> bool {
self.connection.close_reason().is_some()
}
@ -149,6 +165,7 @@ pub struct GenericReceiver<C: TransportConnection> {
pipes: Arc<Mutex<mpsc::Receiver<PipeReader<C::RecvStream>>>>,
connection: C,
ping_sender: Arc<RwLock<Option<GenericSender<C>>>>,
max_message_size: Arc<AtomicU64>,
}
impl<C: TransportConnection> Clone for GenericReceiver<C> {
@ -159,6 +176,7 @@ impl<C: TransportConnection> Clone for GenericReceiver<C> {
pipes: self.pipes.clone(),
connection: self.connection.clone(),
ping_sender: self.ping_sender.clone(),
max_message_size: self.max_message_size.clone(),
}
}
}
@ -169,9 +187,11 @@ impl<C: TransportConnection> GenericReceiver<C> {
#[cfg(feature = "pipes")]
let (pipe_tx, pipe_rx) = mpsc::channel(policy.receiver_queue_capacity);
let ping_sender: Arc<RwLock<Option<GenericSender<C>>>> = Arc::new(RwLock::new(None));
let max_message_size = Arc::new(AtomicU64::new(policy.handshake_max_message_size));
let task_ping_sender = ping_sender.clone();
let task_connection = connection.clone();
let task_policy = policy.clone();
let task_max_message_size = max_message_size.clone();
tokio::spawn(async move {
let limit = Arc::new(Semaphore::new(
task_policy.max_concurrent_stream_tasks.max(1),
@ -201,6 +221,7 @@ impl<C: TransportConnection> GenericReceiver<C> {
#[cfg(feature = "pipes")]
let pipe_tx = pipe_tx.clone();
let policy = task_policy.clone();
let max_message_size = task_max_message_size.clone();
let permit = limit.clone();
let ping_sender = task_ping_sender.clone();
tokio::spawn(async move {
@ -214,42 +235,65 @@ impl<C: TransportConnection> GenericReceiver<C> {
.max_frames_per_stream
.is_some_and(|max| frames >= max)
{
let _ = tx.send(Err(CommunicationError::StreamError)).await;
break;
}
let mut len = [0; 4];
match timeout(policy.read_timeout, stream.read_exact(&mut len)).await {
Ok(Ok(())) => {}
Ok(Err(_)) | Err(_) => break,
Ok(Err(CommunicationError::StreamClosed)) => break,
Ok(Err(error)) => {
tracing::error!(
"[mtp-transport] frame header read failed: {error}"
);
tracing::warn!(%error, "MTP receive stream failed while reading frame header");
break;
}
Err(error) => {
tracing::error!(
"[mtp-transport] frame header read timed out: {error}"
);
tracing::warn!(%error, "MTP receive stream timed out while reading frame header");
break;
}
}
let len = u32::from_be_bytes(len);
if len == policy.close_frame_len {
let _ = tx.send(Err(CommunicationError::StreamClosed)).await;
break;
}
if len as u64 > policy.max_message_size {
let _ = tx.send(Err(CommunicationError::MessageTooLarge)).await;
if len as u64 > max_message_size.load(Ordering::Relaxed) {
tracing::warn!(len, "MTP receive stream frame is too large");
break;
}
let target_len = len as usize;
let mut body = Vec::new();
if body.try_reserve(target_len.min(16 * 1024)).is_err() {
let _ = tx.send(Err(CommunicationError::MessageTooLarge)).await;
tracing::warn!(
target_len,
"MTP receive stream could not reserve frame body"
);
break;
}
while body.len() < target_len {
let chunk_len = (target_len - body.len()).min(16 * 1024);
let mut chunk = [0u8; 16 * 1024];
if !matches!(
timeout(
policy.read_timeout,
stream.read_exact(&mut chunk[..chunk_len]),
)
.await,
Ok(Ok(()))
) || body.try_reserve(chunk_len).is_err()
let body_read = timeout(
policy.read_timeout,
stream.read_exact(&mut chunk[..chunk_len]),
)
.await;
if !matches!(&body_read, Ok(Ok(())))
|| body.try_reserve(chunk_len).is_err()
{
let _ = tx.send(Err(CommunicationError::StreamError)).await;
tracing::error!(
"[mtp-transport] frame body read failed ({} bytes): {:?}",
chunk_len,
body_read
);
tracing::warn!(
pipe_chunk_len = chunk_len,
?body_read,
"MTP receive stream failed while reading frame body"
);
break;
}
body.extend_from_slice(&chunk[..chunk_len]);
@ -261,9 +305,7 @@ impl<C: TransportConnection> GenericReceiver<C> {
let message = match CommunicationValue::from_bytes(&body) {
Ok(message) => message,
Err(_) => {
let _ = tx
.send(Err(CommunicationError::ParseCommunicationValue))
.await;
tracing::warn!("MTP receive stream contained an invalid frame");
break;
}
};
@ -285,6 +327,8 @@ impl<C: TransportConnection> GenericReceiver<C> {
pipe_id,
};
tracing::debug!(pipe_id, description = %pipe_reader.description, "classified incoming pipe stream");
if pipe_tx.send(pipe_reader).await.is_err() {
break;
}
@ -322,11 +366,18 @@ impl<C: TransportConnection> GenericReceiver<C> {
pipes: Arc::new(Mutex::new(pipe_rx)),
connection,
ping_sender,
max_message_size,
}
}
pub async fn respond_to_pings(&self, sender: GenericSender<C>) {
*self.ping_sender.write().await = Some(sender);
}
/// Switch from the handshake frame limit to the application frame limit.
pub fn set_max_message_size(&self, max_message_size: u64) {
self.max_message_size
.store(max_message_size, Ordering::Relaxed);
}
pub async fn receive(&self) -> Result<CommunicationValue, CommunicationError> {
self.incoming
.lock()

View file

@ -3,6 +3,7 @@ use mtp_common::CommunicationError;
use rustls::pki_types::{PrivateKeyDer, pem::PemObject};
use std::net::{IpAddr, SocketAddr};
use std::sync::Arc;
use std::time::Instant;
use tracing::debug;
use wtransport::{Connection as WTConnection, Endpoint, ServerConfig};
@ -100,7 +101,7 @@ pub async fn host_with_config(
port: u16,
config: HostConfig,
) -> Result<Host, CommunicationError> {
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
mtp_crypto::ensure_crypto_provider();
let (cert_pem, key_pem) = match config.credentials {
HostCredentials::Pem { cert_pem, key_pem } => (cert_pem, key_pem),
@ -121,8 +122,11 @@ pub async fn host_with_config(
let task = tokio::spawn(async move {
loop {
let accept_started = Instant::now();
let incoming_session = endpoint.accept().await;
tracing::debug!(elapsed = ?accept_started.elapsed(), "host accept loop: received QUIC connection");
let session_started = Instant::now();
let request = match incoming_session.await {
Ok(req) => req,
Err(e) => {
@ -130,7 +134,9 @@ pub async fn host_with_config(
continue;
}
};
tracing::debug!(elapsed = ?session_started.elapsed(), "host accept loop: complete WebTransport handshake");
let request_accept_started = Instant::now();
let connection = match request
.accept_with_headers([("sec-webtransport-http3-draft02", "1")])
.await
@ -141,6 +147,7 @@ pub async fn host_with_config(
continue;
}
};
tracing::debug!(elapsed = ?request_accept_started.elapsed(), "host accept loop: accept WebTransport request");
let incoming_tx = incoming_tx.clone();
tokio::spawn(handle_connection(connection, incoming_tx, policy.clone()));
@ -159,11 +166,14 @@ async fn handle_connection(
tx: tokio::sync::mpsc::Sender<(Sender, Receiver)>,
policy: Arc<Policy>,
) {
let setup_started = Instant::now();
let handle = Arc::new(ConnectionHandle::new());
let sender = Sender::new(connection.clone(), handle.clone(), policy.clone());
let receiver = Receiver::new_for_handshake(connection, handle, policy);
let _ = tx.send((sender, receiver)).await;
if tx.send((sender, receiver)).await.is_ok() {
tracing::debug!(elapsed = ?setup_started.elapsed(), "host accept loop: hand connection to authentication");
}
}
async fn configure_server(