Brought Example up to spec

transport/Cargo.toml

@@ -6,6 +6,7 @@ edition = "2024"
 [dependencies]
 mtp-codec = { version = "0.2.0", path = "../codec" }
 mtp-common = { version = "0.2.0", path = "../common" }
+mtp-crypto = { version = "0.2.0", path = "../crypto" }
 wtransport = { version = "0.7.1", default-features = false, features = [
     "aws-lc-rs",
     "quinn",

transport/src/client.rs

@@ -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))
 }
 

transport/src/connection.rs

@@ -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;

transport/src/generic_connection.rs

@@ -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()

transport/src/host.rs

@@ -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(

transport/tests/integration.rs

@@ -105,7 +105,7 @@ async fn test_explicit_development_tls() -> Result<(), Box<dyn std::error::Error
     let (client_tx, _client_rx) = connect_with_config(&url, client_config).await?;
     let (_host_tx, _host_rx) = h.next().await.ok_or("host did not accept connection")?;
 
-    client_tx.close();
+    client_tx.close().await;
     h.shutdown();
     Ok(())
 }
@@ -133,8 +133,8 @@ async fn test_send_receive_roundtrip() -> Result<(), Box<dyn std::error::Error>>
     assert_numbered_message(&client_received, CommunicationType::Pong, 99, &tm);
 
     // Close both sides
-    client_tx.close();
-    host_tx.close();
+    client_tx.close().await;
+    host_tx.close().await;
     Ok(())
 }
 
@@ -167,7 +167,7 @@ async fn test_concurrent_messages() -> Result<(), Box<dyn std::error::Error>> {
         assert_numbered_message(&received, CommunicationType::Pong, i * 10, &tm);
     }
 
-    client_tx.close();
+    client_tx.close().await;
     Ok(())
 }
 
@@ -178,7 +178,6 @@ async fn test_close_detection() -> Result<(), Box<dyn std::error::Error>> {
     // Send a message then close
     let msg = CommunicationValue::new(CommunicationType::Ping);
     client_tx.send(&msg).await?;
-    client_tx.close();
 
     // Host should still receive the message
     let tm = TypeMap::latest();
@@ -190,6 +189,8 @@ async fn test_close_detection() -> Result<(), Box<dyn std::error::Error>> {
             .expect("test type must be mapped")
     );
 
+    client_tx.close().await;
+
     // Host should get an error or closed signal on next receive
     let result = host_rx.receive().await;
     assert!(result.is_err());
@@ -244,8 +245,8 @@ async fn test_drop_receiver_keeps_sender_alive() -> Result<(), Box<dyn std::erro
     let got = client_rx.receive().await?;
     assert_numbered_message(&got, CommunicationType::Pong, 7, &tm);
 
-    client_tx.close();
-    host_tx.close();
+    client_tx.close().await;
+    host_tx.close().await;
     Ok(())
 }
 
@@ -267,8 +268,8 @@ async fn test_persistent_stream_reopens_after_local_finish()
     let received2 = host_rx.receive().await?;
     assert_numbered_message(&received2, CommunicationType::Pong, 22, &tm);
 
-    client_tx.close();
-    host_tx.close();
+    client_tx.close().await;
+    host_tx.close().await;
     drop(client_rx);
     Ok(())
 }
@@ -295,7 +296,7 @@ async fn test_receiver_backpressure_with_small_queue() -> Result<(), Box<dyn std
         assert_numbered_message(&received, CommunicationType::Ping, i, &tm);
     }
 
-    client_tx.close();
+    client_tx.close().await;
     drop(client_rx);
     h.shutdown();
     Ok(())
@@ -330,7 +331,7 @@ async fn test_max_frames_per_stream_enforced() -> Result<(), Box<dyn std::error:
     let second = host_rx.receive().await;
     assert!(second.is_err(), "stream should be closed after frame limit");
 
-    client_tx.close();
+    client_tx.close().await;
     h.shutdown();
     Ok(())
 }
@@ -375,7 +376,7 @@ async fn test_semaphore_saturation_with_concurrent_streams()
         assert_numbered_message(&received, CommunicationType::Ping, i, &tm);
     }
 
-    client_tx.close();
+    client_tx.close().await;
     h.shutdown();
     Ok(())
 }