diff --git a/Cargo.lock b/Cargo.lock index a17abb0..2d71b13 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1245,10 +1245,12 @@ dependencies = [ "rand 0.10.2", "rand_core 0.10.1", "rcgen", + "rustls", "serde", "sha2 0.11.0", "thiserror 1.0.69", "time", + "tokio", "zeroize", ] @@ -1273,6 +1275,7 @@ dependencies = [ "rand 0.8.7", "tokio", "tracing", + "wtransport", ] [[package]] @@ -1282,6 +1285,7 @@ dependencies = [ "async-trait", "mtp-codec", "mtp-common", + "mtp-crypto", "rcgen", "rustls", "rustls-native-certs", @@ -1334,9 +1338,11 @@ dependencies = [ "http", "mtp-codec", "mtp-common", + "mtp-crypto", "mtp-host", "mtp-transport", "quinn", + "rand 0.10.2", "rcgen", "rustls", "thiserror 2.0.18", diff --git a/client/src/crypto.rs b/client/src/crypto.rs index 01c9d50..c6dcea7 100644 --- a/client/src/crypto.rs +++ b/client/src/crypto.rs @@ -14,14 +14,14 @@ pub(crate) fn unexpected_response_type_error( )) } -pub(crate) fn verify_host_challenge( +pub(crate) async fn verify_host_challenge( challenge: &CommunicationValue, host_pk: &mtp_crypto::PublicKeyBundle, id: u64, server_challenge: u128, require_pq: bool, ) -> Result<(), CommunicationError> { - use mtp_crypto::{auth, verify_ed25519, verify_ml_dsa}; + use mtp_crypto::{auth, verify_ed25519}; let sig = match challenge.get_data(DataType::Signature) { DataValue::Bytes(b) => b.clone(), @@ -49,18 +49,27 @@ pub(crate) fn verify_host_challenge( } let payload = auth::challenge_payload(id, server_challenge); - verify_ed25519(&host_pk.sig_cl_public_key, &payload, &sig).map_err(|_| { - CommunicationError::AuthenticationFailed("Host challenge signature invalid".into()) - })?; - if !pq_sig.is_empty() && verify_ml_dsa(&host_pk.sig_pq_public_key, &payload, &pq_sig).is_err() { - return Err(CommunicationError::AuthenticationFailed( - "Host challenge PQ signature invalid".into(), - )); + if pq_sig.is_empty() { + verify_ed25519(&host_pk.sig_cl_public_key, &payload, &sig).map_err(|_| { + CommunicationError::AuthenticationFailed("Host challenge signature invalid".into()) + })?; + } else { + mtp_crypto::sign_parallel::verify_dual_parallel( + host_pk.sig_cl_public_key.clone(), + host_pk.sig_pq_public_key.clone(), + payload, + sig, + pq_sig, + ) + .await + .map_err(|_| { + CommunicationError::AuthenticationFailed("Host challenge signature invalid".into()) + })?; } Ok(()) } -pub(crate) fn verify_host_final( +pub(crate) async fn verify_host_final( response: &CommunicationValue, host_pk: &mtp_crypto::PublicKeyBundle, id: u64, @@ -68,7 +77,7 @@ pub(crate) fn verify_host_final( server_challenge: u128, require_pq: bool, ) -> Result<(), CommunicationError> { - use mtp_crypto::{auth, verify_ed25519, verify_ml_dsa}; + use mtp_crypto::{auth, verify_ed25519}; match response.get_data(DataType::ClientNonce) { DataValue::UnsignedNumber(n) if *n == client_nonce => {} @@ -98,12 +107,20 @@ pub(crate) fn verify_host_final( } let payload = auth::host_final_payload(id, client_nonce, server_challenge); - verify_ed25519(&host_pk.sig_cl_public_key, &payload, &sig) + if pq_sig.is_empty() { + verify_ed25519(&host_pk.sig_cl_public_key, &payload, &sig).map_err(|_| { + CommunicationError::AuthenticationFailed("Host signature invalid".into()) + })?; + } else { + mtp_crypto::sign_parallel::verify_dual_parallel( + host_pk.sig_cl_public_key.clone(), + host_pk.sig_pq_public_key.clone(), + payload, + sig, + pq_sig, + ) + .await .map_err(|_| CommunicationError::AuthenticationFailed("Host signature invalid".into()))?; - if !pq_sig.is_empty() && verify_ml_dsa(&host_pk.sig_pq_public_key, &payload, &pq_sig).is_err() { - return Err(CommunicationError::AuthenticationFailed( - "Host PQ signature invalid".into(), - )); } Ok(()) } @@ -141,33 +158,36 @@ pub(crate) fn negotiated_version( } } -pub(crate) fn signed_challenge_response( +pub(crate) async fn signed_challenge_response( keys: &mtp_crypto::Keyring, - proof_payload: &[u8], + proof_payload: Vec, client_nonce: u128, ) -> Result { use mtp_crypto::{Ed25519Signer, MlDsaSigner, SignatureScheme}; let signer = Ed25519Signer::new(&keys.sig_cl_secret_key) .map_err(|e| CommunicationError::Other(e.to_string()))?; - let signature = signer - .sign(proof_payload) - .map_err(|e| CommunicationError::Other(e.to_string()))?; - let mut proof = CommunicationValue::new(CommunicationType::ChallengeResponse) .add_typed_default( DataType::ClientNonce, DataValue::UnsignedNumber(client_nonce), - ) - .add_typed_default(DataType::Signature, DataValue::Bytes(signature)); + ); - if !keys.sig_pq_secret_key.as_bytes().is_empty() { + if keys.sig_pq_secret_key.as_bytes().is_empty() { + let signature = signer + .sign(&proof_payload) + .map_err(|e| CommunicationError::Other(e.to_string()))?; + proof = proof.add_typed_default(DataType::Signature, DataValue::Bytes(signature)); + } else { let pq_signer = MlDsaSigner::new(&keys.sig_pq_secret_key, &keys.sig_pq_public_key) .map_err(|e| CommunicationError::Other(e.to_string()))?; - let pq_signature = pq_signer - .sign(proof_payload) - .map_err(|e| CommunicationError::Other(e.to_string()))?; - proof = proof.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_signature)); + let (signature, pq_signature) = + mtp_crypto::sign_parallel::sign_dual_parallel(signer, pq_signer, proof_payload) + .await + .map_err(|e| CommunicationError::Other(e.to_string()))?; + proof = proof + .add_typed_default(DataType::Signature, DataValue::Bytes(signature)) + .add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_signature)); } Ok(proof) @@ -212,7 +232,8 @@ pub(crate) async fn receive_verified_challenge( bound_id, server_challenge, require_pq, - )?; + ) + .await?; Ok(server_challenge) } diff --git a/client/src/lib.rs b/client/src/lib.rs index 04162f8..7003686 100644 --- a/client/src/lib.rs +++ b/client/src/lib.rs @@ -127,7 +127,7 @@ impl MTPClient { CommunicationError::Other("host returned an invalid negotiated version".into()) })?, mtp_common::HandshakeOutcome::Rejected { reason } => { - sender.close(); + sender.close().await; return Err(CommunicationError::Other(reason.to_string())); } }; @@ -192,7 +192,7 @@ impl MTPClient { ident = ident.add_typed_default(DataType::Description, DataValue::Str(desc.clone())); } if let Err(e) = sender.send(&ident).await { - sender.close(); + sender.close().await; return Err(e); } @@ -209,7 +209,7 @@ impl MTPClient { { Ok(c) => c, Err(e) => { - sender.close(); + sender.close().await; return Err(e); } }; @@ -222,22 +222,23 @@ impl MTPClient { client_nonce, ); - let proof = match crypto::signed_challenge_response(keys, &proof_payload, client_nonce) { + let proof = match crypto::signed_challenge_response(keys, proof_payload, client_nonce).await + { Ok(p) => p, Err(e) => { - sender.close(); + sender.close().await; return Err(e); } }; if let Err(e) = sender.send(&proof).await { - sender.close(); + sender.close().await; return Err(e); } let response = match receiver.receive().await { Ok(r) => r, Err(e) => { - sender.close(); + sender.close().await; return Err(e); } }; @@ -249,7 +250,7 @@ impl MTPClient { ) })?; if response.get_type() != expected_type { - sender.close(); + sender.close().await; return Err(crypto::unexpected_response_type_error( "auth_connect", expected_type, @@ -257,7 +258,7 @@ impl MTPClient { )); } if let Err(e) = crypto::check_connected(&response, "Server rejected authentication") { - sender.close(); + sender.close().await; return Err(e); } if let Err(e) = crypto::verify_host_final( @@ -267,8 +268,10 @@ impl MTPClient { client_nonce, server_challenge, config.require_pq, - ) { - sender.close(); + ) + .await + { + sender.close().await; return Err(e); } @@ -340,7 +343,7 @@ impl MTPClient { register.add_typed_default(DataType::Description, DataValue::Str(desc.clone())); } if let Err(e) = sender.send(®ister).await { - sender.close(); + sender.close().await; return Err(e); } @@ -357,7 +360,7 @@ impl MTPClient { { Ok(c) => c, Err(e) => { - sender.close(); + sender.close().await; return Err(e); } }; @@ -370,22 +373,23 @@ impl MTPClient { client_nonce, ); - let proof = match crypto::signed_challenge_response(keys, &proof_payload, client_nonce) { + let proof = match crypto::signed_challenge_response(keys, proof_payload, client_nonce).await + { Ok(p) => p, Err(e) => { - sender.close(); + sender.close().await; return Err(e); } }; if let Err(e) = sender.send(&proof).await { - sender.close(); + sender.close().await; return Err(e); } let response = match receiver.receive().await { Ok(r) => r, Err(e) => { - sender.close(); + sender.close().await; return Err(e); } }; @@ -395,7 +399,7 @@ impl MTPClient { CommunicationError::Other("RegisterResponse is absent from the type map".into()) })?; if response.get_type() != expected_type { - sender.close(); + sender.close().await; return Err(crypto::unexpected_response_type_error( "auth_register", expected_type, @@ -403,13 +407,13 @@ impl MTPClient { )); } if let Err(e) = crypto::check_connected(&response, "Server rejected registration") { - sender.close(); + sender.close().await; return Err(e); } let assigned_id = match response.get_data(DataType::Id) { DataValue::UnsignedNumber(n) => *n as u64, _ => { - sender.close(); + sender.close().await; return Err(CommunicationError::AuthenticationFailed( "Missing assigned ID".into(), )); @@ -422,8 +426,10 @@ impl MTPClient { client_nonce, server_challenge, config.require_pq, - ) { - sender.close(); + ) + .await + { + sender.close().await; return Err(e); } diff --git a/client/src/ping.rs b/client/src/ping.rs index 167fb19..9309cbc 100644 --- a/client/src/ping.rs +++ b/client/src/ping.rs @@ -57,7 +57,7 @@ pub(crate) fn start_ping_session( } _ = ticker.tick() => { if max_missed_pings > 0 && !pending.is_empty() && pending.len() >= max_missed_pings { - sender.close(); + sender.close().await; break; } @@ -80,7 +80,7 @@ pub(crate) fn start_ping_session( } let id = ping.get_id(); if sender.send(&ping).await.is_err() { - sender.close(); + sender.close().await; break; } pending.insert(id, Instant::now()); diff --git a/crypto/Cargo.toml b/crypto/Cargo.toml index e75e2f9..c55176e 100644 --- a/crypto/Cargo.toml +++ b/crypto/Cargo.toml @@ -26,9 +26,11 @@ ml-dsa = { version = "0.1.1", optional = true } serde = { version = "1", optional = true, features = ["derive"] } rcgen = { version = "0.14", optional = true } time = { version = "0.3", optional = true } +tokio = { version = "1", features = ["macros", "rt"], optional = true } +rustls = "0.23.41" [features] -default = ["chacha20poly1305", "ed25519-dalek", "hkdf", "sha2", "ml-dsa"] +default = ["chacha20poly1305", "ed25519-dalek", "hkdf", "sha2", "ml-dsa", "parallel"] # Enabling ml-dsa by default ensures dual-signature support in the handshake # without requiring a separate PQC feature flag in host/client crates. full = ["default", "aes-gcm"] @@ -38,3 +40,4 @@ wasm = ["getrandom/wasm_js"] hkdf = ["dep:hkdf", "dep:sha2"] sha2 = ["dep:sha2"] tls = ["dep:rcgen", "dep:time"] +parallel = ["dep:tokio"] diff --git a/crypto/src/keypair.rs b/crypto/src/keypair.rs index 0407edc..61c5b44 100644 --- a/crypto/src/keypair.rs +++ b/crypto/src/keypair.rs @@ -200,6 +200,35 @@ impl Keyring { } } + /// Generates the independent KEM, classical-signature, and PQ-signature keys concurrently. + #[cfg(all( + feature = "mlkem-tls", + feature = "ml-dsa", + feature = "ed25519-dalek", + feature = "parallel" + ))] + pub async fn generate_parallel() -> Self { + let kem_handle = tokio::task::spawn_blocking(crate::kem::HybridKem::generate_keypair); + let ed_handle = tokio::task::spawn_blocking(crate::sign::Ed25519Signer::generate); + let pq_handle = tokio::task::spawn_blocking(crate::sign::MlDsaSigner::generate); + + let (kem_result, ed_result, pq_result) = tokio::join!(kem_handle, ed_handle, pq_handle); + let (kem_sk, kem_pk) = kem_result.expect("key generation task must not panic"); + let (_ed_signer, sig_cl_sk, sig_cl_pk) = + ed_result.expect("key generation task must not panic"); + let (_pq_signer, sig_pq_sk, sig_pq_pk) = + pq_result.expect("key generation task must not panic"); + + Self { + kem_public_key: kem_pk, + kem_secret_key: kem_sk, + sig_pq_public_key: sig_pq_pk, + sig_pq_secret_key: sig_pq_sk, + sig_cl_public_key: sig_cl_pk, + sig_cl_secret_key: sig_cl_sk, + } + } + pub fn public_key_bundle(&self) -> PublicKeyBundle { PublicKeyBundle { kem_public_key: self.kem_public_key.clone(), diff --git a/crypto/src/lib.rs b/crypto/src/lib.rs index 9fb1579..43ab5a4 100644 --- a/crypto/src/lib.rs +++ b/crypto/src/lib.rs @@ -3,6 +3,10 @@ pub mod auth; pub mod error; pub mod keypair; +use std::sync::Once; + +static CRYPTO_INIT: Once = Once::new(); + #[cfg(feature = "sha2")] pub mod hash; @@ -12,6 +16,9 @@ pub mod kdf; #[cfg(any(feature = "ed25519-dalek", feature = "ml-dsa"))] pub mod sign; +#[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa", feature = "parallel"))] +pub mod sign_parallel; + #[cfg(any(feature = "ed25519-dalek", feature = "ml-dsa"))] pub use sign::SigAlgorithm; @@ -59,6 +66,16 @@ pub use kem::{Encapsulated, HybridKem}; pub use enc::EncryptionType; +/// Install Rustls' AWS-LC provider once for the entire process. +/// +/// Rustls only accepts one process-wide default provider. Calling this helper +/// from every TLS entry point makes that initialization idempotent. +pub fn ensure_crypto_provider() { + CRYPTO_INIT.call_once(|| { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + }); +} + #[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] pub use enc::{decrypt_with, encrypt_for}; diff --git a/crypto/src/sign_parallel.rs b/crypto/src/sign_parallel.rs new file mode 100644 index 0000000..1eba5c2 --- /dev/null +++ b/crypto/src/sign_parallel.rs @@ -0,0 +1,73 @@ +//! Parallel helpers for dual (classical + post-quantum) signatures. + +use std::sync::Arc; + +use tokio::task; + +use crate::{ + CryptoError, SignaturePqPublicKey, SignaturePublicKey, SignatureScheme, verify_ed25519, + verify_ml_dsa, +}; + +/// Signs a message with the classical and PQ schemes concurrently on Tokio's blocking pool. +pub async fn sign_dual_parallel( + ed_signer: S1, + pq_signer: S2, + message: Vec, +) -> Result<(Vec, Vec), CryptoError> +where + S1: SignatureScheme + Send + 'static, + S2: SignatureScheme + Send + 'static, +{ + let ed_message = message.clone(); + let ed_handle = task::spawn_blocking(move || ed_signer.sign(&ed_message)); + let pq_handle = task::spawn_blocking(move || pq_signer.sign(&message)); + + let (ed_result, pq_result) = tokio::join!(ed_handle, pq_handle); + let ed_signature = ed_result.map_err(|_| CryptoError::SigningFailed)??; + let pq_signature = pq_result.map_err(|_| CryptoError::SigningFailed)??; + Ok((ed_signature, pq_signature)) +} + +/// Signs with an owned classical signer and a shared PQ signer. +/// +/// This avoids reconstructing the ML-DSA signing key when a host signs both +/// the challenge and the final response in one authentication handshake. +pub async fn sign_dual_parallel_shared_pq( + ed_signer: S1, + pq_signer: Arc, + message: Vec, +) -> Result<(Vec, Vec), CryptoError> +where + S1: SignatureScheme + Send + 'static, + S2: SignatureScheme + Send + Sync + 'static, +{ + let ed_message = message.clone(); + let ed_handle = task::spawn_blocking(move || ed_signer.sign(&ed_message)); + let pq_handle = task::spawn_blocking(move || pq_signer.sign(&message)); + + let (ed_result, pq_result) = tokio::join!(ed_handle, pq_handle); + let ed_signature = ed_result.map_err(|_| CryptoError::SigningFailed)??; + let pq_signature = pq_result.map_err(|_| CryptoError::SigningFailed)??; + Ok((ed_signature, pq_signature)) +} + +/// Verifies the classical and PQ signatures concurrently on Tokio's blocking pool. +pub async fn verify_dual_parallel( + ed_public_key: SignaturePublicKey, + pq_public_key: SignaturePqPublicKey, + message: Vec, + ed_signature: Vec, + pq_signature: Vec, +) -> Result<(), CryptoError> { + let ed_message = message.clone(); + let ed_handle = + task::spawn_blocking(move || verify_ed25519(&ed_public_key, &ed_message, &ed_signature)); + let pq_handle = + task::spawn_blocking(move || verify_ml_dsa(&pq_public_key, &message, &pq_signature)); + + let (ed_result, pq_result) = tokio::join!(ed_handle, pq_handle); + ed_result.map_err(|_| CryptoError::VerificationFailed)??; + pq_result.map_err(|_| CryptoError::VerificationFailed)??; + Ok(()) +} diff --git a/docs/NATIVE-HOST-WEB-SERVER.md b/docs/NATIVE-HOST-WEB-SERVER.md index 8579d50..61a9e69 100644 --- a/docs/NATIVE-HOST-WEB-SERVER.md +++ b/docs/NATIVE-HOST-WEB-SERVER.md @@ -4,6 +4,10 @@ WebTransport sessions are returned by `accept()` for application messages. `MTPWebServer` and `MTPHost` cannot bind the same IP and port. +The repository's combined server example registers `/` on `MTPWebServer` and +returns `OK` while the process is running. The route is served over HTTP/3 at +`https://localhost:8080/` on the same QUIC endpoint as WebTransport MTP sessions. + ## WebServerConfig | Builder | Default | Purpose | @@ -73,7 +77,7 @@ let mut server = MTPWebServer::new(host, web).await?; while let Some(connection) = server.accept().await? { // connection: WebMTPConnection - while let Ok(message) = connection.receiver.receive().await { + while let Ok(message) = connection.receive().await { println!("received MTP message {}", message.get_id()); } } @@ -81,7 +85,27 @@ while let Some(connection) = server.accept().await? { `server.accept()` returns `Option` for each WebTransport session. HTTP/3 routes do not surface through `accept()` because the server dispatches them internally. `WebMTPConnection` retains the negotiated version, codec, request path, description, sender, and receiver used by native MTP connections. -WebTransport sessions are unauthenticated. With the `crypto` feature enabled, construction rejects any `AuthenticationPolicy` other than `Unauthenticated`. The connection has `AuthState::Unauthenticated` and a random 48-bit client ID when crypto fields are compiled in; `guest_id_generator` is not used by this adapter. +### Authentication + +`MTPWebServer` does not impose its own authentication policy. It respects the `AuthenticationPolicy` set on the supplied `HostConfig`: + +| Policy | Behavior | +|--------|----------| +| `Unauthenticated` (default) | No authentication handshake is performed. The connection has `AuthState::Unauthenticated` and a random 48-bit client ID. `guest_id_generator` is not used by this adapter. | +| `AllowAuthentication` | The server accepts the first message. If it is an `Identification` or `Register` message, a full challenge-response handshake is performed. If it is an ordinary opening message, the connection remains unauthenticated. | +| `ForceAuthentication` | The server requires a valid `Identification` or `Register` message as the first frame and performs the challenge-response handshake. Any other opening message is rejected. | + +When authentication is required or allowed and the client presents credentials, the server performs the same Ed25519/ML-DSA challenge-response handshake used by native MTP host connections: + +1. The client sends `Identification` (with a client ID) or `Register` (with a public-key bundle). +2. The server looks up or accepts the client's public keys, generates a random 128-bit server nonce, and signs a challenge payload with its host keyring. +3. The client responds with a proof signed by its own keys. +4. The server verifies the proof, assigns the client ID, and sends a final signed response. + +On success, the connection has `AuthState::Authenticated`, the assigned `client_id`, and `client_public_key` populated. On failure, `accept()` returns `AcceptError::AuthenticationFailed` (or `AcceptError::AuthenticationTimedOut` if the handshake exceeds `host_config.auth_timeout`). + +`MTPWebServer::new` returns `CommunicationError` for certificate parsing, certificate loading, and bind failures. It does **not** reject `HostConfig` based on `AuthenticationPolicy`; any policy is accepted at construction time. + ## Errors diff --git a/example/Cargo.lock b/example/Cargo.lock index 79f5bb2..23258d9 100644 --- a/example/Cargo.lock +++ b/example/Cargo.lock @@ -232,6 +232,7 @@ dependencies = [ "mtp", "rand 0.10.2", "tokio", + "tracing-subscriber", ] [[package]] @@ -1185,10 +1186,12 @@ dependencies = [ "rand 0.10.2", "rand_core 0.10.1", "rcgen", + "rustls", "serde", "sha2 0.11.0", "thiserror 1.0.69", "time", + "tokio", "zeroize", ] @@ -1213,6 +1216,7 @@ dependencies = [ "rand 0.8.7", "tokio", "tracing", + "wtransport", ] [[package]] @@ -1222,6 +1226,7 @@ dependencies = [ "async-trait", "mtp-codec", "mtp-common", + "mtp-crypto", "rcgen", "rustls", "rustls-native-certs", @@ -1251,24 +1256,17 @@ dependencies = [ "http", "mtp-codec", "mtp-common", + "mtp-crypto", "mtp-host", "mtp-transport", "quinn", + "rand 0.10.2", "rustls", "thiserror 2.0.18", "tokio", "tracing", ] -[[package]] -name = "mtp-webserver-example" -version = "0.2.0" -dependencies = [ - "http", - "mtp", - "tokio", -] - [[package]] name = "nom" version = "7.1.3" @@ -1279,6 +1277,15 @@ dependencies = [ "minimal-lexical", ] +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + [[package]] name = "num-bigint" version = "0.4.8" @@ -1874,8 +1881,11 @@ dependencies = [ "base64", "hex", "mtp", + "rustls", "serde_json", "tokio", + "tokio-rustls", + "tracing-subscriber", ] [[package]] @@ -1921,6 +1931,15 @@ dependencies = [ "sponge-cursor", ] +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + [[package]] name = "shlex" version = "2.0.1" @@ -2100,6 +2119,15 @@ dependencies = [ "syn", ] +[[package]] +name = "thread_local" +version = "1.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ad99c4c6d32803332c548b1af0540b357b3f5fc0be8f6c6bfe8b2e6ae784070" +dependencies = [ + "cfg-if", +] + [[package]] name = "time" version = "0.3.53" @@ -2183,6 +2211,16 @@ dependencies = [ "syn", ] +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + [[package]] name = "tokio-util" version = "0.7.18" @@ -2226,6 +2264,32 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" dependencies = [ "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "nu-ansi-term", + "sharded-slab", + "smallvec", + "thread_local", + "tracing-core", + "tracing-log", ] [[package]] @@ -2286,6 +2350,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + [[package]] name = "version_check" version = "0.9.5" diff --git a/example/Cargo.toml b/example/Cargo.toml index e14bf36..48506d1 100644 --- a/example/Cargo.toml +++ b/example/Cargo.toml @@ -1,6 +1,6 @@ [workspace] members = [ "server", - "client", "keygen", "webserver", + "client", "keygen", ] resolver = "3" diff --git a/example/client/Cargo.toml b/example/client/Cargo.toml index 460880b..e56862d 100644 --- a/example/client/Cargo.toml +++ b/example/client/Cargo.toml @@ -11,3 +11,4 @@ path = "src/main.rs" mtp = { version = "0.2.0", path = "../../", features = ["client", "crypto", "files", "pipes"] } tokio = { version = "1", features = ["full"] } rand = "0.10.1" +tracing-subscriber = "0.3.23" diff --git a/example/client/src/auth.rs b/example/client/src/auth.rs index d51640d..56e1304 100644 --- a/example/client/src/auth.rs +++ b/example/client/src/auth.rs @@ -1,4 +1,6 @@ -use std::fs; +use std::time::Instant; + +use tokio::fs; use mtp::client::{ClientConfig, MTPClient, MTPConnection}; use mtp::crypto::{ @@ -14,16 +16,25 @@ pub async fn connect_or_register( let keyring_path = format!("{key_prefix}.mk"); let id_path = format!("{key_prefix}.id"); + let file_load_started = Instant::now(); if let (Ok(keyring), Ok(id)) = ( load_keyring_raw(&keyring_path), - fs::read_to_string(&id_path), + fs::read_to_string(&id_path).await, ) { let client_id: u64 = id.trim().parse()?; - println!("Loaded client keys (ID: {client_id})"); + println!( + "Loaded client keys (ID: {client_id}) in {:?}", + file_load_started.elapsed() + ); config.client_id = client_id; + let auth_started = Instant::now(); let conn = MTPClient::auth_connect(config, &keyring, &host_public_key).await?; - println!("Authenticated (version {})", conn.version); + println!( + "Authenticated (version {}) in {:?}", + conn.version, + auth_started.elapsed() + ); return Ok((conn, keyring)); } @@ -45,7 +56,7 @@ pub async fn connect_or_register( println!("Registered with ID: {}", conn.client_id); save_keyring_raw(&keyring, &keyring_path)?; - fs::write(&id_path, conn.client_id.to_string())?; + fs::write(&id_path, conn.client_id.to_string()).await?; println!("Saved client keys -> {keyring_path}"); Ok((conn, keyring)) diff --git a/example/client/src/main.rs b/example/client/src/main.rs index 037449e..4e4b7eb 100644 --- a/example/client/src/main.rs +++ b/example/client/src/main.rs @@ -20,6 +20,7 @@ fn dev_cert_path() -> String { #[tokio::main] async fn main() -> Result<(), Box> { + tracing_subscriber::fmt::init(); let cert_path = dev_cert_path(); let cert_pem = fs::read(&cert_path).unwrap_or_else(|e| { panic!( @@ -29,10 +30,9 @@ async fn main() -> Result<(), Box> { let host_public_key = match load_public_key_bundle("host.mpkb") { Ok(bundle) => bundle, Err(e) => { - return Err(format!( - "Missing host.mpkb: run the server first to export it ({e})" - ) - .into()); + return Err( + format!("Missing host.mpkb: run the server first to export it ({e})").into(), + ); } }; @@ -49,7 +49,7 @@ async fn main() -> Result<(), Box> { println!("\n--- Pipe demo ---"); pipes::run_pipe_demo(&conn, 1).await?; - conn.sender.close(); + conn.sender.close().await; println!("\nDone"); Ok(()) } diff --git a/example/client/src/pipes.rs b/example/client/src/pipes.rs index f09ee70..5ae6edd 100644 --- a/example/client/src/pipes.rs +++ b/example/client/src/pipes.rs @@ -50,7 +50,6 @@ pub async fn run_pipe_demo( .finish() .await .map_err(|e| mtp::common::PipeError::IoError(e.to_string()))?; - println!(" [pipe {i}.{run}] writer: data sent and finished"); Ok::<(), mtp::common::PipeError>(()) } Ok(None) => { diff --git a/example/server/Cargo.toml b/example/server/Cargo.toml index 0f85b13..890fa91 100644 --- a/example/server/Cargo.toml +++ b/example/server/Cargo.toml @@ -8,8 +8,11 @@ name = "server" path = "src/main.rs" [dependencies] -mtp = { version = "0.2.0", path = "../../", features = ["crypto", "tls", "host", "files", "pipes"] } +mtp = { version = "0.2.0", path = "../../", features = ["crypto", "tls", "web-server", "files", "pipes"] } tokio = { version = "1", features = ["full"] } +tokio-rustls = "0.26" +rustls = "0.23" serde_json = { version = "1" } hex = "0.4" base64 = "0.22" +tracing-subscriber = "0.3.23" diff --git a/example/server/src/clients.rs b/example/server/src/clients.rs index 4dfed94..241de6f 100644 --- a/example/server/src/clients.rs +++ b/example/server/src/clients.rs @@ -1,14 +1,14 @@ use std::collections::HashMap; -use std::fs; use std::sync::{Arc, Mutex}; +use tokio::fs; use mtp::crypto::PublicKeyBundle; -pub fn load_client_db( +pub async fn load_client_db( path: &str, ) -> Result<(Arc>>, Arc>), Box> { - let clients_map = match fs::read_to_string(path) { + let clients_map = match fs::read_to_string(path).await { Ok(data) => match serde_json::from_str(&data) { Ok(clients) => clients, Err(e) => { diff --git a/example/server/src/keys.rs b/example/server/src/keys.rs index b315db0..53f83af 100644 --- a/example/server/src/keys.rs +++ b/example/server/src/keys.rs @@ -1,4 +1,4 @@ -use std::fs; +use tokio::fs; use mtp::crypto::Keyring; use mtp::files::{load_keyring_raw, save_keyring_raw, save_public_key_bundle}; @@ -20,14 +20,16 @@ pub fn load_or_generate_host_keys( Ok((HOST_ID, keyring)) } -pub fn export_host_public_keys(host_keyring: &Keyring) -> Result<(), Box> { +pub async fn export_host_public_keys( + host_keyring: &Keyring, +) -> Result<(), Box> { let bundle = host_keyring.public_key_bundle(); save_public_key_bundle(&bundle, "host.mpkb")?; /* The web client fetches the bundle as hex over HTTP. */ let bundle_hex = hex::encode(bundle.as_bytes()); - fs::write("host_public_key_bundle.hex", &bundle_hex)?; - fs::create_dir_all("web-client/public")?; - fs::write("web-client/public/host_public_key_bundle.hex", &bundle_hex)?; + fs::write("host_public_key_bundle.hex", &bundle_hex).await?; + fs::create_dir_all("web-client/public").await?; + fs::write("web-client/public/host_public_key_bundle.hex", &bundle_hex).await?; Ok(()) } diff --git a/example/server/src/main.rs b/example/server/src/main.rs index 20c7071..1affc0a 100644 --- a/example/server/src/main.rs +++ b/example/server/src/main.rs @@ -2,14 +2,19 @@ mod clients; mod handlers; mod keys; mod tls; +#[path = "web-server.rs"] +mod web_server; -use mtp::host::{AuthenticationPolicy, HostConfig, MTPHost}; +use mtp::host::HostConfig; use mtp::type_map::TypeMap; use std::future::Future; use std::path::Path; use std::pin::Pin; use std::sync::Arc; +const CONNECTION_IDLE_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(300); +const MAX_MESSAGES_PER_CONNECTION: u64 = 10_000; + fn dev_cert_paths() -> (String, String) { let cert = std::env::var("MTP_DEV_CERT").unwrap_or_else(|_| { if Path::new("example/dev-cert/cert.pem").exists() { @@ -29,62 +34,81 @@ fn dev_cert_paths() -> (String, String) { } async fn handle_pipe_loopback( - conn: &mtp::host::MTPConnection, - req: mtp::host::PipeRequest, + conn: &mtp::webserver::WebMTPConnection, + request: mtp::host::PipeRequest< + mtp::webserver::WebMtpSender, + mtp::webserver::H3TransportReceiver, + >, ) -> Result<(), Box> { - let pipe_id = req.id(); + let pipe_id = request.id(); + println!(" [loopback] Accepting pipe {pipe_id} ..."); + let mut reader = request.accept().await?; + + let return_pipe = conn.create_pipe("loopback").await?; println!( - " [loopback] Pipe request: id={pipe_id} description={:?}", - req.description() + " [loopback] Requested return pipe {}; waiting for client acceptance ...", + return_pipe.pipe_id() ); + let Some(mut writer) = return_pipe.wait().await? else { + return Err("client denied the return pipe".into()); + }; - println!(" [loopback] Calling accept() for pipe {pipe_id} ..."); - let mut reader = req.accept().await?; - println!(" [loopback] Pipe {pipe_id} accepted, reading data ..."); - - let handle = conn.create_pipe("loopback").await?; - println!( - " [loopback] Return pipe created (id={}), waiting for client ...", - handle.pipe_id() - ); - - match handle.wait().await? { - Some(mut writer) => { - println!(" [loopback] Client accepted return pipe, echoing incoming bytes ..."); - let mut total = 0usize; - let mut buf = [0u8; 16 * 1024]; - loop { - let n = tokio::io::AsyncReadExt::read(&mut reader, &mut buf).await?; - if n == 0 { - break; - } - total += n; - tokio::io::AsyncWriteExt::write_all(&mut writer, &buf[..n]).await?; - } - writer.finish().await?; - println!( - " [loopback] Pipe {pipe_id} loopback complete ({} bytes)", - total - ); - } - None => { - eprintln!(" [loopback] Return pipe denied by client for pipe {pipe_id}"); - } - } - + let copied = tokio::io::copy(&mut reader, &mut writer).await?; + writer.finish_async().await?; + println!(" [loopback] Pipe {pipe_id} complete ({copied} bytes)"); Ok(()) } #[tokio::main] async fn main() -> Result<(), Box> { + tracing_subscriber::fmt::init(); let (cert_path, key_path) = dev_cert_paths(); - let (cert_pem, key_pem) = tls::load_or_generate_tls(&cert_path, &key_path)?; - let cert_hash = tls::certificate_sha256_hex(&cert_pem)?; - tls::export_webtransport_cert_hash(&cert_hash)?; + let (cert_pem, key_pem) = tls::load_or_generate_tls(&cert_path, &key_path).await?; + let cert_hash = tls::certificate_sha256_hex(&cert_pem).await?; + tls::export_webtransport_cert_hash(&cert_hash).await?; println!("WebTransport certificate sha256: {cert_hash}"); let (_host_id, host_keyring) = keys::load_or_generate_host_keys("host.mk")?; - keys::export_host_public_keys(&host_keyring)?; + keys::export_host_public_keys(&host_keyring).await?; + + let (clients, next_id) = clients::load_client_db("clients.json").await?; + let clients_for_get = clients.clone(); + let get_existing_client = move |id: u64, _description: Option| { + let clients = clients_for_get.clone(); + Box::pin(async move { clients.lock().ok()?.get(&id).cloned() }) + as Pin> + Send>> + }; + let clients_for_register = clients.clone(); + let next_id_for_register = next_id.clone(); + let complete_register = move |bundle: mtp::crypto::PublicKeyBundle, + _description: Option| { + let clients = clients_for_register.clone(); + let next_id = next_id_for_register.clone(); + + Box::pin(async move { + let id = { + let mut next = next_id.lock().expect("client id mutex poisoned"); + let id = *next; + *next += 1; + id + }; + + let json = { + let mut db = clients.lock().expect("client database mutex poisoned"); + db.insert(id, bundle); + serde_json::to_string_pretty(&*db).ok() + }; + + if let Some(json) = json { + if let Err(error) = tokio::fs::write("clients.json", json).await { + eprintln!("Failed to persist clients.json: {error}"); + } + } + + println!("Registered new client with ID: {id}"); + id + }) as Pin + Send>> + }; let decrypt_keyring = Arc::new( match mtp::crypto::Keyring::from_bytes(&host_keyring.to_bytes()) { @@ -95,51 +119,7 @@ async fn main() -> Result<(), Box> { }, ); - let (clients, next_id) = clients::load_client_db("clients.json")?; - - let clients_for_get = clients.clone(); - let get_existing_user = move |id: u64, _description: Option| { - let clients = clients_for_get.clone(); - Box::pin(async move { - let result = clients.lock().unwrap().get(&id).cloned(); - if result.is_some() { - println!("Auth lookup: client ID {id} found"); - } else { - eprintln!("Auth lookup: unknown client ID {id}"); - } - result - }) as Pin> + Send>> - }; - - let clients_for_register = clients.clone(); - let next_id_for_register = next_id.clone(); - let clients_path = "clients.json".to_string(); - let complete_register = move |bundle: mtp::crypto::PublicKeyBundle, - _description: Option| { - let db_arc = clients_for_register.clone(); - let nid_arc = next_id_for_register.clone(); - let path = clients_path.clone(); - Box::pin(async move { - let mut db = db_arc.lock().unwrap(); - let mut nid = nid_arc.lock().unwrap(); - let id = *nid; - *nid += 1; - db.insert(id, bundle); - match serde_json::to_string_pretty(&*db) { - Ok(json) => match std::fs::write(&path, json) { - Ok(()) => {} - Err(e) => eprintln!("Failed to persist client database to {path}: {e}"), - }, - Err(e) => { - eprintln!("Failed to serialize client database after registering {id}: {e}") - } - } - println!("Registered new client with ID: {}", id); - id - }) as Pin + Send>> - }; - - println!("Starting MTP server on port 8080 ..."); + println!("Starting integrated MTP web server on port 8080 ..."); let config = HostConfig::new( std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED), @@ -149,13 +129,21 @@ async fn main() -> Result<(), Box> { ) .with_authentication( host_keyring, - Box::new(get_existing_user), + Box::new(get_existing_client), Box::new(complete_register), - ) - .with_authentication_policy(AuthenticationPolicy::ForceAuthentication); + ); - let mut host = MTPHost::new(config).await?; - println!("Server listening on {}", host.local_addr()); + let _https = web_server::spawn_https( + std::net::SocketAddr::new(config.ip, config.port), + &config.tls_fullchain, + &config.tls_key, + ) + .await?; + let mut host = mtp::webserver::MTPWebServer::new(config, web_server::config()?).await?; + println!( + "Server listening on https://{} (TCP HTTPS + UDP WebTransport)", + host.local_addr() + ); while let Some(conn) = host.accept().await? { let decrypt_keyring = Arc::clone(&decrypt_keyring); @@ -170,56 +158,81 @@ async fn main() -> Result<(), Box> { let tm: &TypeMap = conn.codec.registry().get(&conn.version).unwrap(); println!("Waiting for messages / pipe requests ..."); - loop { - tokio::select! { - biased; + let mut pipe_open = true; + let mut message_open = true; + let mut messages_received = 0_u64; - pipe_req = conn.receive_pipe() => { - match pipe_req { - Ok(req) => { - println!(" Pipe request: id={} desc={:?}", req.id(), req.description()); - if let Err(e) = handle_pipe_loopback(&conn, req).await { - eprintln!(" Pipe loopback error: {e}"); - } - } - Err(e) => { - println!("Pipe channel closed: {e}"); - break; - } - } - } - msg = conn.receive() => { - match msg { - Ok(msg) => { - println!("Received: {msg}"); - let response = match handlers::process_and_respond( - &msg, - tm, - conn.client_public_key.as_ref(), - &decrypt_keyring, - ) { - Ok(response) => response, - Err(e) => { - eprintln!("Failed to build response: {e}"); - continue; + while pipe_open || message_open { + let activity = tokio::time::timeout(CONNECTION_IDLE_TIMEOUT, async { + tokio::select! { + biased; + pipe_request = conn.receive_pipe(), if pipe_open => { + match pipe_request { + Ok(request) => { + if let Err(error) = handle_pipe_loopback(&conn, request).await { + eprintln!(" [loopback] Pipe error: {error}"); } - }; - println!("Sending: {response}"); - if let Err(e) = conn.sender.send(&response).await { - eprintln!("Send error: {e}"); - break; + } + Err(mtp::common::CommunicationError::StreamClosed) + | Err(mtp::common::CommunicationError::ClosedByPeer) => { + println!("Pipe channel closed normally"); + pipe_open = false; + } + Err(error) => { + println!("Pipe channel closed: {error}"); + pipe_open = false; } } - Err(e) => { - println!("Connection ended: {e}"); - break; + } + message = conn.receive(), if message_open => { + match message { + Ok(message) => { + messages_received += 1; + println!("Received: {message}"); + match handlers::process_and_respond( + &message, + tm, + conn.client_public_key.as_ref(), + &decrypt_keyring, + ) { + Ok(response) => { + println!("Sending: {response}"); + if let Err(error) = conn.sender.send(&response).await { + eprintln!("Send error: {error}"); + pipe_open = false; + message_open = false; + } + } + Err(error) => { + eprintln!("Failed to build response: {error}"); + } + } + } + Err(mtp::common::CommunicationError::StreamClosed) + | Err(mtp::common::CommunicationError::ClosedByPeer) => { + println!("Message channel closed normally"); + message_open = false; + } + Err(error) => { + println!("Message channel closed: {error}"); + message_open = false; + } } } } + }) + .await; + + if activity.is_err() { + println!("Connection idle timeout reached"); + break; + } + if messages_received >= MAX_MESSAGES_PER_CONNECTION { + println!("Connection message limit reached"); + break; } } - conn.sender.close(); println!("Connection closed\n"); }); } diff --git a/example/server/src/tls.rs b/example/server/src/tls.rs index 3ae25b1..4b53f00 100644 --- a/example/server/src/tls.rs +++ b/example/server/src/tls.rs @@ -1,34 +1,34 @@ use base64::Engine; -use std::fs; use std::path::Path; +use tokio::fs; -pub fn load_or_generate_tls( +pub async fn load_or_generate_tls( cert_path: &str, key_path: &str, ) -> Result<(Vec, Vec), Box> { - if let (Ok(c), Ok(k)) = (fs::read(cert_path), fs::read(key_path)) { + if let (Ok(c), Ok(k)) = (fs::read(cert_path).await, fs::read(key_path).await) { println!("Using existing TLS cert from {cert_path}"); return Ok((c, k)); } println!("Generating self-signed TLS certificate ..."); if let Some(parent) = Path::new(cert_path).parent() { - fs::create_dir_all(parent)?; + fs::create_dir_all(parent).await?; } if let Some(parent) = Path::new(key_path).parent() { - fs::create_dir_all(parent)?; + fs::create_dir_all(parent).await?; } let (cert_pem, key_pem) = mtp::crypto::tls::generate_self_signed_cert("localhost")?; - fs::write(cert_path, &cert_pem)?; - fs::write(key_path, &key_pem)?; + fs::write(cert_path, &cert_pem).await?; + fs::write(key_path, &key_pem).await?; println!("Wrote {cert_path} and {key_path}"); Ok((cert_pem, key_pem)) } -pub fn certificate_sha256_hex(cert: &[u8]) -> Result> { +pub async fn certificate_sha256_hex(cert: &[u8]) -> Result> { let der = if cert.starts_with(b"-----BEGIN CERTIFICATE-----") { let pem = std::str::from_utf8(cert)?; let base64 = pem @@ -43,14 +43,14 @@ pub fn certificate_sha256_hex(cert: &[u8]) -> Result Result<(), Box> { +pub async fn export_webtransport_cert_hash(hash: &str) -> Result<(), Box> { let public_dir = if Path::new("web-client").exists() { Path::new("web-client/public") } else { Path::new("example/web-client/public") }; - fs::create_dir_all(public_dir)?; - fs::write(public_dir.join("mtp_dev_cert_hash.txt"), hash)?; + fs::create_dir_all(public_dir).await?; + fs::write(public_dir.join("mtp_dev_cert_hash.txt"), hash).await?; let dev_cert_dir = if Path::new("dev-cert").exists() { Path::new("dev-cert") @@ -58,7 +58,7 @@ pub fn export_webtransport_cert_hash(hash: &str) -> Result<(), Box Http3Response { + response + .header("content-type", "text/plain; charset=utf-8") + .body("OK") +} + +pub fn config() -> Result { + WebServerConfig::new().route("/", ok) +} + +/// Starts the conventional HTTPS side of the example host. WebTransport uses +/// UDP/QUIC on the same port; browsers still need TCP/TLS to navigate to a URL. +pub async fn spawn_https( + address: SocketAddr, + certificate_pem: &[u8], + key_pem: &[u8], +) -> io::Result> { + // The TCP listener is created before the QUIC endpoint, so it must select + // rustls' process-wide provider itself. + mtp::crypto::ensure_crypto_provider(); + let certificates = rustls::pki_types::CertificateDer::pem_slice_iter(certificate_pem) + .collect::, _>>() + .map_err(io::Error::other)?; + let key = PrivateKeyDer::from_pem_slice(key_pem).map_err(io::Error::other)?; + let tls = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(certificates, key) + .map_err(io::Error::other)?; + let listener = TcpListener::bind(address).await?; + let acceptor = TlsAcceptor::from(Arc::new(tls)); + let asset_root = web_client_dist(); + match &asset_root { + Some(_) => println!( + "HTTPS web client available at https://localhost:{}", + address.port() + ), + None => eprintln!( + "Web client build not found; HTTPS will show setup instructions. Run `pnpm --dir example/web-client build`." + ), + } + Ok(tokio::spawn(async move { + loop { + let Ok((stream, _)) = listener.accept().await else { + break; + }; + let acceptor = acceptor.clone(); + let asset_root = asset_root.clone(); + tokio::spawn(async move { + let Ok(mut stream) = acceptor.accept(stream).await else { + return; + }; + let _ = serve_https_request(&mut stream, &asset_root).await; + }); + } + })) +} + +fn web_client_dist() -> Option { + [ + PathBuf::from("web-client/dist"), + PathBuf::from("example/web-client/dist"), + ] + .into_iter() + .find(|path| path.join("index.html").is_file()) +} + +async fn serve_https_request(stream: &mut S, asset_root: &Option) -> io::Result<()> +where + S: AsyncRead + AsyncWrite + Unpin, +{ + let mut request = [0; 16 * 1024]; + let size = stream.read(&mut request).await?; + let request = std::str::from_utf8(&request[..size]).unwrap_or_default(); + let path = request + .lines() + .next() + .and_then(|line| line.split_whitespace().nth(1)) + .unwrap_or("/") + .split('?') + .next() + .unwrap_or("/"); + let (status, content_type, body) = match asset_root { + Some(asset_root) => { + let relative = path.trim_start_matches('/'); + let candidate = asset_root.join(relative); + let file = if relative.is_empty() || !candidate.is_file() || relative.contains("..") { + asset_root.join("index.html") + } else { + candidate + }; + let content_type = content_type(&file); + match tokio::fs::read(&file).await { + Ok(body) => ("200 OK", content_type, body), + Err(_) => ("404 Not Found", "text/plain; charset=utf-8", b"Not found".to_vec()), + } + } + None => ( + "503 Service Unavailable", + "text/html; charset=utf-8", + b"MTP web client not built

Run pnpm --dir example/web-client build.

".to_vec(), + ), + }; + + let response = format!( + "HTTP/1.1 {status}\r\nContent-Type: {content_type}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n", + body.len() + ); + stream.write_all(response.as_bytes()).await?; + stream.write_all(&body).await?; + stream.shutdown().await +} + +fn content_type(file: &Path) -> &'static str { + match file.extension().and_then(|extension| extension.to_str()) { + Some("html") => "text/html; charset=utf-8", + Some("js") => "text/javascript; charset=utf-8", + Some("css") => "text/css; charset=utf-8", + Some("wasm") => "application/wasm", + Some("svg") => "image/svg+xml", + Some("json") => "application/json", + _ => "application/octet-stream", + } +} diff --git a/example/web-client/src/main.ts b/example/web-client/src/main.ts index e1979d0..d9316a6 100644 --- a/example/web-client/src/main.ts +++ b/example/web-client/src/main.ts @@ -44,10 +44,10 @@ let activeClient: ReturnType extends Promise let micStream: MediaStream | null = null; let mediaRecorder: MediaRecorder | null = null; let activePipeWriter: MTPPipeWriter | null = null; +let loopbackAudioContext: AudioContext | null = null; +let micStreamGeneration = 0; let pipeSendCount = 0; let pendingPipeReaders: MTPPipeReader[] = []; -let audioContext: AudioContext | null = null; -let micMonitorAudio: HTMLAudioElement | null = null; let currentPipePingMs: number | null = null; let lastPipeSendStartedAt = 0; let currentPipeId: number | null = null; @@ -56,6 +56,13 @@ let currentPipeState = "idle"; let loopbackPlaybackCount = 0; let hasPipeRequestHandler = false; +// ===== AUDIO LOOPBACK STATE ===== +// We accumulate all chunks into a single Blob, then decode and play it +// when the pipe closes. decodeAudioData needs a complete file, not fragments. +let loopbackBlobParts: BlobPart[] = []; +let loopbackMimeType = ""; +let loopbackAudioElement: HTMLAudioElement | null = null; + const credentialStorage: MTPCredentialStorage = { getItem: (key) => localStorage.getItem(key), setItem: (key, value) => localStorage.setItem(key, value), @@ -126,11 +133,18 @@ function setPipeState( function getPipeId(handle: unknown): number | null { if (handle && typeof handle === "object") { const candidate = handle as Record; - const value = + let value = candidate.pipeId ?? candidate.pipe_id ?? candidate["pipe-id"] ?? candidate.id; + if (typeof value === "function") { + try { + value = value.call(handle); + } catch { + return null; + } + } if (typeof value === "number" && Number.isFinite(value)) { return value; } @@ -138,63 +152,64 @@ function getPipeId(handle: unknown): number | null { return null; } -function ensureAudioContext() { - if (!audioContext) { - audioContext = new AudioContext(); - } - return audioContext; +// ===== FIXED AUDIO LOOPBACK: accumulate chunks, play as single file ===== + +function startLoopbackAccumulation(mimeType: string) { + loopbackBlobParts = []; + loopbackMimeType = mimeType; + pipeLog("Loopback: accumulating audio chunks..."); } -async function playLoopbackAudio(chunks: BlobPart[], mimeType: string) { - if (chunks.length === 0) return; - const blob = new Blob(chunks, { type: mimeType }); - const url = URL.createObjectURL(blob); - const audio = new Audio(url); - audio.autoplay = true; - audio.onended = () => URL.revokeObjectURL(url); - audio.onerror = () => URL.revokeObjectURL(url); +function queueLoopbackChunk(data: Uint8Array) { + loopbackBlobParts.push(data.slice()); +} + +async function finishLoopbackPlayback() { + if (loopbackBlobParts.length === 0) { + pipeLog("Loopback: no chunks received.", "error"); + return; + } + + // Stop any previous playback + if (loopbackAudioElement) { + loopbackAudioElement.pause(); + const src = loopbackAudioElement.src; + loopbackAudioElement.src = ""; + if (src.startsWith("blob:")) { + URL.revokeObjectURL(src); + } + loopbackAudioElement = null; + } + + // Concatenate all chunks into one Blob + const blob = new Blob(loopbackBlobParts, { type: loopbackMimeType }); + loopbackBlobParts = []; + + pipeLog(`Loopback: assembled ${blob.size} bytes, decoding...`); + try { - await ensureAudioContext().resume(); - await audio.play(); + const arrayBuffer = await blob.arrayBuffer(); + const audioContext = new AudioContext(); + const audioBuffer = await audioContext.decodeAudioData(arrayBuffer); + + const source = audioContext.createBufferSource(); + source.buffer = audioBuffer; + source.connect(audioContext.destination); + source.start(); + loopbackPlaybackCount += 1; updateMetrics(); - pipeLog(`Loopback playback started (${chunks.length} chunks).`); + pipeLog(`Loopback playback started (${audioBuffer.duration.toFixed(2)}s).`); + + // Clean up audio context when done + source.onended = () => { + audioContext.close().catch(() => {}); + }; } catch (e) { - URL.revokeObjectURL(url); - pipeLog(`Loopback playback failed: ${e}`, "error"); + pipeLog(`Loopback decode/playback failed: ${e}`, "error"); } } -async function startMicPlayback(stream: MediaStream) { - if (micMonitorAudio) { - micMonitorAudio.pause(); - micMonitorAudio.srcObject = null; - micMonitorAudio = null; - } - - const audio = new Audio(); - audio.autoplay = true; - audio.controls = false; - audio.muted = false; - audio.srcObject = stream; - micMonitorAudio = audio; - - try { - await ensureAudioContext().resume(); - await audio.play(); - pipeLog("Microphone monitoring playback started."); - } catch (e) { - pipeLog(`Microphone monitoring playback failed: ${e}`, "error"); - } -} - -function stopMicPlayback() { - if (!micMonitorAudio) return; - micMonitorAudio.pause(); - micMonitorAudio.srcObject = null; - micMonitorAudio = null; -} - function renderStructured(value: unknown): string { return JSON.stringify(value, (_key, item) => { if (typeof item === "bigint") { @@ -375,7 +390,6 @@ async function connect() { log("WebTransport is not supported in this browser.", "error"); return; } - const hostPk = hexToBytes(HOST_PUBLIC_KEY.value); saveHostPublicKey(); await loadDevCertHash(); @@ -401,7 +415,7 @@ async function connect() { const activeClientId = await client.auth(); clientId = activeClientId; loadKeys(); - log(`Connected as client ${activeClientId}`); + log(`Connected as authenticated client ${activeClientId}`); log("\nSending typed Ping..."); await client.send( @@ -451,7 +465,9 @@ async function startMicStreaming() { STOP_MIC.disabled = false; setPipeState("creating", { pipeId: null, description: "mic-audio" }); pipeLog("Microphone acquired. Creating pipe ..."); - await startMicPlayback(micStream); + pipeLog( + "Microphone monitoring is off; playback will use the server loopback.", + ); if (!hasPipeRequestHandler) { activeClient.setOnPipeRequest(async (request) => { @@ -487,6 +503,7 @@ async function startMicStreaming() { return; } activePipeWriter = writer; + const streamGeneration = ++micStreamGeneration; setPipeState("streaming", { pipeId: getPipeId(writer) ?? pipeId, @@ -500,12 +517,22 @@ async function startMicStreaming() { const mimeType = MediaRecorder.isTypeSupported("audio/webm;codecs=opus") ? "audio/webm;codecs=opus" : "audio/webm"; - mediaRecorder = new MediaRecorder(micStream, { mimeType }); + const recorder = new MediaRecorder(micStream, { mimeType }); + mediaRecorder = recorder; - mediaRecorder.ondataavailable = async (event) => { - if (event.data.size === 0 || !activeClient) return; + recorder.ondataavailable = async (event) => { + // A final chunk can be queued before recorder.stop(). Do not use the + // captured writer unless this is still the current active stream. + if ( + event.data.size === 0 || + !activeClient || + micStreamGeneration !== streamGeneration || + mediaRecorder !== recorder || + activePipeWriter !== writer + ) { + return; + } - const startTime = performance.now(); pipeSendCount++; const chunkNum = pipeSendCount; @@ -513,6 +540,14 @@ async function startMicStreaming() { lastPipeSendStartedAt = performance.now(); const buffer = await event.data.arrayBuffer(); const data = new Uint8Array(buffer); + // arrayBuffer() yields, so shutdown may have happened meanwhile. + if ( + micStreamGeneration !== streamGeneration || + mediaRecorder !== recorder || + activePipeWriter !== writer + ) { + return; + } await writer.write(data); currentPipePingMs = performance.now() - lastPipeSendStartedAt; updateMetrics(); @@ -521,27 +556,29 @@ async function startMicStreaming() { } }; - mediaRecorder.start(200); // emit data every 200ms + recorder.start(200); // emit data every 200ms updateMetrics(); pipeLog("Streaming started (200ms chunks)."); } - async function readLoopbackPipe(reader: MTPPipeReader) { const startTime = performance.now(); let totalBytes = 0; let chunkCount = 0; - const chunks: BlobPart[] = []; const mimeType = MediaRecorder.isTypeSupported("audio/webm;codecs=opus") ? "audio/webm;codecs=opus" : "audio/webm"; try { + loopbackBlobParts = []; + loopbackMimeType = mimeType; + pipeLog("Loopback: accumulating chunks..."); + while (true) { const data = await reader.read(); if (data == null) break; // EOF totalBytes += data.length; chunkCount++; - chunks.push(data.slice().buffer); + loopbackBlobParts.push(data.slice()); } } catch (e) { pipeLog(` Return pipe read error: ${e}`, "error"); @@ -554,26 +591,90 @@ async function readLoopbackPipe(reader: MTPPipeReader) { `delay=${elapsed.toFixed(1)}ms`, ); setPipeState("loopback-ready", { pingMs: elapsed }); - await playLoopbackAudio(chunks, mimeType); + + // Decode and play the complete recording + if (loopbackBlobParts.length > 0) { + try { + const blob = new Blob(loopbackBlobParts, { type: loopbackMimeType }); + const arrayBuffer = await blob.arrayBuffer(); + + if (!loopbackAudioContext) { + loopbackAudioContext = new AudioContext(); + } + const audioBuffer = + await loopbackAudioContext.decodeAudioData(arrayBuffer); + + const source = loopbackAudioContext.createBufferSource(); + source.buffer = audioBuffer; + source.connect(loopbackAudioContext.destination); + source.start(); + + loopbackPlaybackCount += 1; + updateMetrics(); + pipeLog( + `Loopback playback started (${audioBuffer.duration.toFixed(2)}s).`, + ); + } catch (e) { + pipeLog(`Loopback decode failed: ${e}`, "error"); + } + } // Clean up the reader from the pending list const idx = pendingPipeReaders.indexOf(reader); if (idx >= 0) pendingPipeReaders.splice(idx, 1); } +// ===== CRITICAL FIX: stopMicStreaming must capture the final chunk ===== async function stopMicStreaming() { - if (mediaRecorder && mediaRecorder.state !== "inactive") { - mediaRecorder.stop(); - mediaRecorder = null; + micStreamGeneration++; + const recorder = mediaRecorder; + mediaRecorder = null; + const writer = activePipeWriter; + activePipeWriter = null; + + // STOPPING STRATEGY: + // 1. Request a final dataavailable event by calling requestData() if needed, + // then stop(). The final event contains the WebM trailer. + // 2. Wait for that final event to be processed (it writes through the pipe). + // 3. Only THEN close the pipe writer. + + if (recorder) { + // Create a promise that resolves when the final dataavailable fires + const finalChunkPromise = new Promise((resolve) => { + const originalHandler = recorder.ondataavailable; + recorder.ondataavailable = async (event) => { + // Call the original handler first so the chunk gets written to the pipe + if (originalHandler) { + await originalHandler.call(recorder, event); + } + // The final chunk from stop() has a 'type' but no special marker. + // MediaRecorder state will be 'inactive' after the final event. + if (recorder.state === "inactive") { + resolve(); + } + }; + }); + + if (recorder.state !== "inactive") { + recorder.stop(); + } + + // Wait up to 1 second for the final chunk to be captured and written + await Promise.race([ + finalChunkPromise, + new Promise((_, reject) => + setTimeout(() => reject(new Error("final chunk timeout")), 1000), + ), + ]).catch(() => { + pipeLog("Warning: final chunk may not have been captured", "error"); + }); } - if (activePipeWriter) { + if (writer) { try { - await activePipeWriter.close(); + await writer.close(); } catch (e) { pipeLog(`Pipe close error: ${e}`, "error"); - } finally { - activePipeWriter = null; } } @@ -581,9 +682,6 @@ async function stopMicStreaming() { micStream.getTracks().forEach((track) => track.stop()); micStream = null; } - stopMicPlayback(); - - // Close pending pipe readers pendingPipeReaders = []; setPipeState("stopped", { pipeId: currentPipeId, diff --git a/example/webserver/Cargo.toml b/example/webserver/Cargo.toml deleted file mode 100644 index 6f5ad02..0000000 --- a/example/webserver/Cargo.toml +++ /dev/null @@ -1,10 +0,0 @@ -[package] -name = "mtp-webserver-example" -version = "0.2.0" -edition = "2024" -publish = false - -[dependencies] -http = "1" -mtp = { path = "../..", features = ["web-server"] } -tokio = { version = "1", features = ["full"] } diff --git a/example/webserver/src/main.rs b/example/webserver/src/main.rs deleted file mode 100644 index ee2d89a..0000000 --- a/example/webserver/src/main.rs +++ /dev/null @@ -1,38 +0,0 @@ -use http::StatusCode; -use mtp::{ - host::HostConfig, - webserver::{Http3Request, Http3Response, MTPWebServer, WebServerConfig}, -}; - -async fn health(_request: Http3Request, response: Http3Response) -> Http3Response { - response - .status(StatusCode::OK) - .header("content-type", "application/json") - .body(r#"{"status":"ok"}"#) -} - -#[tokio::main] -async fn main() -> Result<(), Box> { - let config = HostConfig::new( - "0.0.0.0".parse()?, - 4433, - std::fs::read("cert.pem")?, - std::fs::read("key.pem")?, - ); - let web = WebServerConfig::new().route("/health", health)?.mtp_path("/mtp"); - let mut server = MTPWebServer::new(config, web).await?; - println!("listening on {}", server.local_addr()); - - while let Some(connection) = server.accept().await? { - println!( - "MTP client connected: path={}, version={}, description={:?}", - connection.path, connection.version, connection.description - ); - tokio::spawn(async move { - while let Ok(message) = connection.receiver.receive().await { - println!("received MTP message {}", message.get_id()); - } - }); - } - Ok(()) -} diff --git a/host/Cargo.toml b/host/Cargo.toml index d75c2be..3d2b287 100644 --- a/host/Cargo.toml +++ b/host/Cargo.toml @@ -11,6 +11,7 @@ mtp-crypto = { version = "0.2.0", path = "../crypto", optional = true } rand = "0.8" tokio = { version = "1", features = ["macros", "rt", "time", "sync"] } tracing = "0.1" +wtransport = "0.7" [features] diff --git a/host/src/connection.rs b/host/src/connection.rs index 76f2e5c..da092ed 100644 --- a/host/src/connection.rs +++ b/host/src/connection.rs @@ -1,3 +1,5 @@ +#[cfg(feature = "pipes")] +use mtp_codec::{CommunicationType, DataType, DataValue}; use mtp_codec::{CommunicationValue, Version, registry::VersionedCodec}; use mtp_common::CommunicationError; #[cfg(feature = "pipes")] @@ -8,7 +10,9 @@ use tokio::sync::{Mutex, mpsc}; #[cfg(feature = "crypto")] use crate::error::random_client_id; #[cfg(feature = "pipes")] -use crate::pipe::{PipeDispatcher, PipeRequest}; +use crate::pipe::{PipeDispatcher, PipeReceiver, PipeRequest, PipeSender, run_dispatcher}; +#[cfg(feature = "pipes")] +use mtp_transport::Policy; mod connection_capability { pub trait Sealed {} @@ -44,7 +48,11 @@ impl MtpReceiverLike for mtp_transport::G } } -pub struct MTPConnection { +pub struct MTPConnection< + S = mtp_transport::Sender, + R = mtp_transport::Receiver, + P = wtransport::RecvStream, +> { pub version: Version, pub codec: VersionedCodec, pub sender: S, @@ -58,9 +66,11 @@ pub struct MTPConnection #[cfg(feature = "pipes")] pub(crate) app_rx: Mutex>>, #[cfg(feature = "pipes")] - pub(crate) pipe_req_rx: Mutex>, + pub(crate) pipe_req_rx: Mutex>>, #[cfg(feature = "pipes")] - pub(crate) pipe_dispatcher: Arc, + pub(crate) pipe_dispatcher: Arc>, + #[cfg(not(feature = "pipes"))] + pub(crate) _pipe_stream: std::marker::PhantomData

, pub description: Option, pub(crate) _dispatcher_task: tokio::task::JoinHandle<()>, #[cfg(feature = "crypto")] @@ -71,7 +81,13 @@ pub struct MTPConnection pub client_public_key: Option, } -impl MTPConnection { +#[cfg(feature = "pipes")] +impl MTPConnection +where + S: PipeSender, + R: PipeReceiver

, + P: tokio::io::AsyncRead + Send + Unpin + 'static, +{ /// Construct an MTP connection from an alternative transport backend. /// /// Native `MTPHost` users continue to receive the default @@ -86,25 +102,60 @@ impl MTPConnection { path: String, description: Option, ) -> Self { - #[cfg(feature = "pipes")] - let (_, app_rx) = mpsc::channel::>(1); - #[cfg(feature = "pipes")] - let (_, pipe_req_rx) = mpsc::channel::(1); - #[cfg(feature = "pipes")] - let dispatcher = Arc::new(PipeDispatcher::default_for_external()); + let policy = Arc::new(Policy::default()); + let (app_tx, app_rx) = mpsc::channel(policy.receiver_queue_capacity); + let (pipe_req_tx, pipe_req_rx) = mpsc::channel(policy.receiver_queue_capacity); + let dispatcher = Arc::new(PipeDispatcher { + pending_creations: Mutex::new(std::collections::HashMap::new()), + pending_pipes: Mutex::new(std::collections::HashMap::new()), + policy, + }); + let task = tokio::spawn(run_dispatcher( + receiver.clone(), + sender.clone(), + app_tx, + pipe_req_tx, + dispatcher.clone(), + )); Self { version, codec, sender, receiver, path, - #[cfg(feature = "pipes")] app_rx: Mutex::new(app_rx), - #[cfg(feature = "pipes")] pipe_req_rx: Mutex::new(pipe_req_rx), - #[cfg(feature = "pipes")] pipe_dispatcher: dispatcher, description, + _dispatcher_task: task, + #[cfg(feature = "crypto")] + auth_state: crate::error::AuthState::Unauthenticated, + #[cfg(feature = "crypto")] + client_id: random_client_id(), + #[cfg(feature = "crypto")] + client_public_key: None, + } + } +} + +#[cfg(not(feature = "pipes"))] +impl MTPConnection { + pub fn from_transport_parts( + version: Version, + codec: VersionedCodec, + sender: S, + receiver: R, + path: String, + description: Option, + ) -> Self { + Self { + version, + codec, + sender, + receiver, + path, + description, + _pipe_stream: std::marker::PhantomData, _dispatcher_task: tokio::spawn(async {}), #[cfg(feature = "crypto")] auth_state: crate::error::AuthState::Unauthenticated, @@ -124,3 +175,59 @@ impl MTPConnection { Ok(message) } } + +#[cfg(feature = "pipes")] +impl MTPConnection +where + S: PipeSender, + P: tokio::io::AsyncRead + Send + Unpin + 'static, +{ + pub async fn receive(&self) -> Result { + let mut rx = self.app_rx.lock().await; + match rx.recv().await { + Some(Ok(mut message)) => { + message.set_type_map(self.codec.type_map()); + Ok(message) + } + Some(Err(error)) => Err(error), + None => Err(CommunicationError::StreamClosed), + } + } + + pub async fn create_pipe( + &self, + description: &str, + ) -> Result, mtp_common::PipeError> { + let pipe_id = rand::random::(); + let (response_tx, response_rx) = tokio::sync::oneshot::channel(); + self.pipe_dispatcher + .pending_creations + .lock() + .await + .insert(pipe_id, response_tx); + + let request = CommunicationValue::new(CommunicationType::PipeRequest) + .with_id(pipe_id) + .add_typed_default(DataType::Description, DataValue::Str(description.into())); + self.sender + .send_pipe_message(&request) + .await + .map_err(mtp_common::PipeError::from)?; + + Ok(crate::pipe::PipeHandle { + pipe_id, + description: description.to_owned(), + sender: self.sender.clone(), + response_rx, + }) + } + + pub async fn receive_pipe(&self) -> Result, CommunicationError> { + self.pipe_req_rx + .lock() + .await + .recv() + .await + .ok_or(CommunicationError::StreamClosed) + } +} diff --git a/host/src/handshake.rs b/host/src/handshake.rs index 278d1d8..3715298 100644 --- a/host/src/handshake.rs +++ b/host/src/handshake.rs @@ -1,4 +1,4 @@ -#[cfg(any(feature = "crypto", feature = "pipes"))] +#[cfg(feature = "crypto")] use mtp_codec::{CommunicationType, CommunicationValue}; use mtp_codec::{ DataType, DataValue, Version, @@ -7,6 +7,7 @@ use mtp_codec::{ use mtp_common::RejectionReason; use mtp_transport::{Receiver, Sender}; use std::sync::Arc; +use std::time::Instant; #[cfg(feature = "pipes")] use tokio::sync::mpsc; @@ -18,9 +19,9 @@ use crate::connection::MTPConnection; use crate::error::AuthState; use crate::error::{AcceptError, extract_version, send_accepted, send_rejection}; #[cfg(feature = "pipes")] -use crate::pipe::run_dispatcher; +use crate::pipe::PipeDispatcher; #[cfg(feature = "pipes")] -use crate::pipe::{PipeDispatcher, PipeRequest}; +use crate::pipe::run_dispatcher; pub struct MTPHost { pub(crate) transport: mtp_transport::Host, @@ -71,11 +72,16 @@ impl MTPHost { } if self.handshakes.is_empty() { + let incoming_started = Instant::now(); match self.transport.next().await { Some((sender, receiver)) => { + tracing::debug!(elapsed = ?incoming_started.elapsed(), "host accept loop: dispatch authentication handshake"); let context = self.context.clone(); self.handshakes.spawn(async move { - context.accept_pair_timed(sender, receiver).await + let handshake_started = Instant::now(); + let result = context.accept_pair_timed(sender, receiver).await; + tracing::debug!(elapsed = ?handshake_started.elapsed(), success = result.is_ok(), "host accept loop: authentication handshake finished"); + result }); continue; } @@ -99,9 +105,15 @@ impl MTPHost { incoming = self.transport.next() => { match incoming { Some((sender, receiver)) => { + tracing::debug!("host accept loop: dispatch authentication handshake"); let context = self.context.clone(); self.handshakes - .spawn(async move { context.accept_pair_timed(sender, receiver).await }); + .spawn(async move { + let handshake_started = Instant::now(); + let result = context.accept_pair_timed(sender, receiver).await; + tracing::debug!(elapsed = ?handshake_started.elapsed(), success = result.is_ok(), "host accept loop: authentication handshake finished"); + result + }); } None => self.transport_closed = true, } @@ -172,7 +184,7 @@ impl HandshakeContext { }, ) .await; - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "authentication not allowed on this host".into(), )); @@ -187,7 +199,7 @@ impl HandshakeContext { }, ) .await; - sender.close(); + sender.close().await; return Err(AcceptError::MissingVersion); } }; @@ -208,7 +220,7 @@ impl HandshakeContext { }, ) .await; - sender.close(); + sender.close().await; return Err(AcceptError::UnsupportedVersion(client_version)); } }; @@ -256,7 +268,7 @@ impl HandshakeContext { }, ) .await; - sender.close(); + sender.close().await; return Err(AcceptError::MissingVersion); } }; @@ -273,7 +285,7 @@ impl HandshakeContext { }, ) .await; - sender.close(); + sender.close().await; return Err(AcceptError::UnsupportedVersion(client_version)); } }; @@ -359,14 +371,6 @@ impl HandshakeContext { receiver.respond_to_pings(sender.clone()); } - #[cfg(feature = "pipes")] - let (_, app_rx) = mpsc::channel::< - Result, - >(1); - #[cfg(feature = "pipes")] - let (_, pipe_req_rx) = mpsc::channel::(1); - #[cfg(feature = "pipes")] - let dispatcher = Arc::new(PipeDispatcher); let task = tokio::spawn(async {}); MTPConnection { @@ -375,12 +379,7 @@ impl HandshakeContext { sender, receiver, path: "/".to_string(), - #[cfg(feature = "pipes")] - app_rx: tokio::sync::Mutex::new(app_rx), - #[cfg(feature = "pipes")] - pipe_req_rx: tokio::sync::Mutex::new(pipe_req_rx), - #[cfg(feature = "pipes")] - pipe_dispatcher: dispatcher, + _pipe_stream: std::marker::PhantomData, description, _dispatcher_task: task, } @@ -451,11 +450,6 @@ impl HandshakeContext { receiver.respond_to_pings(sender.clone()); } - let (_, app_rx) = mpsc::channel::< - Result, - >(1); - let (_, pipe_req_rx) = mpsc::channel::(1); - let dispatcher = Arc::new(PipeDispatcher); let task = tokio::spawn(async {}); MTPConnection { @@ -464,9 +458,7 @@ impl HandshakeContext { sender, receiver, path: "/".to_string(), - app_rx: tokio::sync::Mutex::new(app_rx), - pipe_req_rx: tokio::sync::Mutex::new(pipe_req_rx), - pipe_dispatcher: dispatcher, + _pipe_stream: std::marker::PhantomData, description, _dispatcher_task: task, auth_state, @@ -477,60 +469,6 @@ impl HandshakeContext { } } -#[cfg(feature = "pipes")] -impl MTPConnection { - pub async fn receive( - &self, - ) -> Result { - let mut rx = self.app_rx.lock().await; - match rx.recv().await { - Some(Ok(mut message)) => { - message.set_type_map(self.codec.type_map()); - Ok(message) - } - Some(Err(error)) => Err(error), - None => Err(mtp_common::CommunicationError::StreamClosed), - } - } - - pub async fn create_pipe( - &self, - description: &str, - ) -> Result { - let pipe_id = rand::random::(); - let (tx, rx) = tokio::sync::oneshot::channel(); - - { - let mut pending = self.pipe_dispatcher.pending_creations.lock().await; - pending.insert(pipe_id, tx); - } - - let request = CommunicationValue::new(CommunicationType::PipeRequest) - .with_id(pipe_id) - .add_typed_default(DataType::Description, DataValue::Str(description.into())); - - self.sender - .send(&request) - .await - .map_err(mtp_common::PipeError::from)?; - - Ok(crate::pipe::PipeHandle { - pipe_id, - description: description.to_string(), - sender: self.sender.clone(), - response_rx: rx, - }) - } - - pub async fn receive_pipe(&self) -> Result { - let mut rx = self.pipe_req_rx.lock().await; - match rx.recv().await { - Some(req) => Ok(req), - None => Err(mtp_common::CommunicationError::StreamClosed), - } - } -} - #[cfg(feature = "crypto")] enum Flow { Login { @@ -581,21 +519,21 @@ impl HandshakeContext { let hello = match receiver.receive().await { Ok(m) => m, Err(e) => { - sender.close(); + sender.close().await; return Err(AcceptError::Receive(e)); } }; let version_str = match hello.get_data(DataType::Version) { DataValue::Str(s) => s.clone(), _ => { - sender.close(); + sender.close().await; return Err(AcceptError::MissingVersion); } }; let client_version = match Version::parse(&version_str) { Some(v) => v, None => { - sender.close(); + sender.close().await; return Err(AcceptError::MissingVersion); } }; @@ -611,7 +549,7 @@ impl HandshakeContext { let cid = match hello.get_data(DataType::Id) { DataValue::UnsignedNumber(n) => *n as u64, _ => { - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "missing client id".into(), )); @@ -628,7 +566,7 @@ impl HandshakeContext { DataValue::Str("unknown client id".into()), ); let _ = sender.send(&rejection).await; - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "unknown client id".into(), )); @@ -644,7 +582,7 @@ impl HandshakeContext { AcceptError::AuthenticationFailed("invalid public key bundle".into()) })?, _ => { - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "missing public keys".into(), )); @@ -656,7 +594,7 @@ impl HandshakeContext { CommunicationType::RegisterResponse, ) } else { - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "unexpected authentication message".into(), )); @@ -685,11 +623,11 @@ impl HandshakeContext { client_version: Version, description: Option, ) -> Result, AcceptError> { - use mtp_crypto::{ - Ed25519Signer, MlDsaSigner, SignatureScheme, auth, verify_ed25519, verify_ml_dsa, - }; + use mtp_crypto::{Ed25519Signer, MlDsaSigner, SignatureScheme, auth, verify_ed25519}; + let handshake_started = Instant::now(); let tm = mtp_codec::TypeMap::latest(); + let negotiate_started = Instant::now(); let negotiated = match self .registry .negotiate(std::slice::from_ref(&client_version)) @@ -703,10 +641,11 @@ impl HandshakeContext { }, ) .await; - sender.close(); + sender.close().await; return Err(AcceptError::UnsupportedVersion(client_version)); } }; + tracing::debug!(elapsed = ?negotiate_started.elapsed(), "authentication handshake: version negotiation"); let pq_enabled = !self .config .host_keyring @@ -729,30 +668,43 @@ impl HandshakeContext { }, ) .await; - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "PQ authentication is required but the host PQ key is absent".into(), )); } - let host_sign = |payload: &[u8]| -> Result<(Vec, Vec), AcceptError> { - let signer = Ed25519Signer::new(&self.config.host_keyring.sig_cl_secret_key) - .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?; - let sig = signer - .sign(payload) - .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?; - let pq_sig = if pq_enabled { - let pq = MlDsaSigner::new( + let signer_init_started = Instant::now(); + let host_pq_signer = if pq_enabled { + Some(Arc::new( + MlDsaSigner::new( &self.config.host_keyring.sig_pq_secret_key, &self.config.host_keyring.sig_pq_public_key, ) + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?, + )) + } else { + None + }; + tracing::debug!(elapsed = ?signer_init_started.elapsed(), "authentication handshake: signer initialization"); + + let host_sign = |payload: Vec| async { + let signer = Ed25519Signer::new(&self.config.host_keyring.sig_cl_secret_key) .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?; - pq.sign(payload) - .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))? + if let Some(pq_signer) = host_pq_signer.as_ref() { + mtp_crypto::sign_parallel::sign_dual_parallel_shared_pq( + signer, + Arc::clone(pq_signer), + payload, + ) + .await + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string())) } else { - Vec::new() - }; - Ok((sig, pq_sig)) + let sig = signer + .sign(&payload) + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?; + Ok((sig, Vec::new())) + } }; let challenge_id = match &flow { @@ -761,8 +713,10 @@ impl HandshakeContext { }; let server_challenge: u128 = rand::random(); + let sign_challenge_started = Instant::now(); let (chal_sig, chal_pq_sig) = - host_sign(&auth::challenge_payload(challenge_id, server_challenge))?; + host_sign(auth::challenge_payload(challenge_id, server_challenge)).await?; + tracing::debug!(elapsed = ?sign_challenge_started.elapsed(), "authentication handshake: sign challenge"); let mut challenge_msg = CommunicationValue::new(CommunicationType::Challenge) .add_typed_default( @@ -782,20 +736,24 @@ impl HandshakeContext { challenge_msg = challenge_msg .add_typed_default(DataType::PqSignature, DataValue::Bytes(chal_pq_sig)); } + let send_challenge_started = Instant::now(); if let Err(e) = sender.send(&challenge_msg).await { - sender.close(); + sender.close().await; return Err(AcceptError::Send(e)); } + tracing::debug!(elapsed = ?send_challenge_started.elapsed(), "authentication handshake: send challenge"); + let receive_proof_started = Instant::now(); let proof = match receiver.receive().await { Ok(m) => m, Err(e) => { - sender.close(); + sender.close().await; return Err(AcceptError::Receive(e)); } }; + tracing::debug!(elapsed = ?receive_proof_started.elapsed(), "authentication handshake: receive client proof"); if Some(proof.get_type()) != CommunicationType::ChallengeResponse.try_to_id(&tm) { - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "missing challenge response".into(), )); @@ -803,7 +761,7 @@ impl HandshakeContext { let client_nonce = match proof.get_data(DataType::ClientNonce) { DataValue::UnsignedNumber(n) => *n, _ => { - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "missing client nonce".into(), )); @@ -812,7 +770,7 @@ impl HandshakeContext { let sig_bytes = match proof.get_data(DataType::Signature) { DataValue::Bytes(b) => b.clone(), _ => { - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "missing challenge signature".into(), )); @@ -837,18 +795,24 @@ impl HandshakeContext { }; let has_client_pq_key = !bundle.sig_pq_public_key.as_bytes().is_empty(); - let pq_ok = if self.config.require_pq { - has_client_pq_key - && !pq_sig_bytes.is_empty() - && verify_ml_dsa(&bundle.sig_pq_public_key, &proof_payload, &pq_sig_bytes).is_ok() + let verify_proof_started = Instant::now(); + let proof_ok = if pq_sig_bytes.is_empty() { + !self.config.require_pq + && verify_ed25519(&bundle.sig_cl_public_key, &proof_payload, &sig_bytes).is_ok() + } else if has_client_pq_key { + mtp_crypto::sign_parallel::verify_dual_parallel( + bundle.sig_cl_public_key.clone(), + bundle.sig_pq_public_key.clone(), + proof_payload, + sig_bytes, + pq_sig_bytes, + ) + .await + .is_ok() } else { - pq_sig_bytes.is_empty() - || (has_client_pq_key - && verify_ml_dsa(&bundle.sig_pq_public_key, &proof_payload, &pq_sig_bytes) - .is_ok()) + false }; - let proof_ok = - verify_ed25519(&bundle.sig_cl_public_key, &proof_payload, &sig_bytes).is_ok() && pq_ok; + tracing::debug!(elapsed = ?verify_proof_started.elapsed(), "authentication handshake: verify client proof"); if !proof_ok { send_rejection( @@ -858,12 +822,13 @@ impl HandshakeContext { }, ) .await; - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "client proof signature invalid".into(), )); } + let register_started = Instant::now(); let (assigned_id, client_bundle) = match flow { Flow::Login { id, bundle } => (id, bundle), Flow::Register { bundle, .. } => { @@ -872,12 +837,16 @@ impl HandshakeContext { (new_id, bundle) } }; + tracing::debug!(elapsed = ?register_started.elapsed(), "authentication handshake: registration callback"); - let (host_sig, host_pq_sig) = host_sign(&auth::host_final_payload( + let sign_final_started = Instant::now(); + let (host_sig, host_pq_sig) = host_sign(auth::host_final_payload( assigned_id, client_nonce, server_challenge, - ))?; + )) + .await?; + tracing::debug!(elapsed = ?sign_final_started.elapsed(), "authentication handshake: sign final response"); let mut response = CommunicationValue::new(response_type) .add_typed_default(DataType::Connected, DataValue::BoolTrue) @@ -894,14 +863,17 @@ impl HandshakeContext { response.add_typed_default(DataType::PqSignature, DataValue::Bytes(host_pq_sig)); } + let send_final_started = Instant::now(); if let Err(e) = sender.send(&response).await { - sender.close(); + sender.close().await; return Err(AcceptError::Send(e)); } if let Err(e) = sender.finish_stream().await { - sender.close(); + sender.close().await; return Err(AcceptError::Send(e)); } + tracing::debug!(elapsed = ?send_final_started.elapsed(), "authentication handshake: send final response"); + tracing::debug!(elapsed = ?handshake_started.elapsed(), "authentication handshake: complete"); let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) { Some(codec) => codec, @@ -939,14 +911,14 @@ impl HandshakeContext { let version_str = match hello.get_data(DataType::Version) { DataValue::Str(s) => s.clone(), _ => { - sender.close(); + sender.close().await; return Err(AcceptError::MissingVersion); } }; let client_version = match Version::parse(&version_str) { Some(v) => v, None => { - sender.close(); + sender.close().await; return Err(AcceptError::MissingVersion); } }; @@ -959,16 +931,16 @@ impl HandshakeContext { if Some(hello.get_type()) == CommunicationType::Register.try_to_id(&tm) { let bundle = match hello.get_data(DataType::PublicKeys) { DataValue::Bytes(b) => PublicKeyBundle::from_bytes(b).map_err(|_| { - sender.close(); AcceptError::AuthenticationFailed("invalid public key bundle".into()) })?, _ => { - sender.close(); + sender.close().await; return Err(AcceptError::AuthenticationFailed( "missing public keys".into(), )); } }; + sender.close().await; let pk_bytes = bundle.as_bytes(); return self .complete_auth_handshake( @@ -1019,7 +991,7 @@ impl HandshakeContext { }, ) .await; - sender.close(); + sender.close().await; return Err(AcceptError::UnsupportedVersion(client_version)); } }; @@ -1046,7 +1018,7 @@ impl HandshakeContext { ))); } - sender.close(); + sender.close().await; Err(AcceptError::AuthenticationFailed( "unexpected message type".into(), )) diff --git a/host/src/lib.rs b/host/src/lib.rs index 960462e..df64b68 100644 --- a/host/src/lib.rs +++ b/host/src/lib.rs @@ -20,6 +20,8 @@ pub use mtp_transport::Sender; pub use mtp_common::PipeError; #[cfg(feature = "pipes")] pub use mtp_transport::PipeWriter; +#[cfg(feature = "pipes")] +pub use pipe::PipeRequest; pub use mtp_codec::registry::Registry; @@ -31,11 +33,15 @@ pub use error::AuthState; #[cfg(test)] mod tests { use super::*; - use mtp_codec::{CommunicationType, DataType, DataValue, registry::VersionedCodec}; + #[cfg(not(feature = "pipes"))] + use mtp_codec::registry::VersionedCodec; + use mtp_codec::{CommunicationType, DataType, DataValue}; + #[cfg(not(feature = "pipes"))] #[derive(Clone, Debug, PartialEq, Eq)] struct AlternateSender; + #[cfg(not(feature = "pipes"))] #[derive(Clone, Debug, PartialEq, Eq)] struct AlternateReceiver; @@ -80,19 +86,21 @@ mod tests { Ok(()) } + #[cfg(not(feature = "pipes"))] #[tokio::test] async fn alternative_transports_use_the_shared_connection_type() { let registry = Registry::builtin(); let version = mtp_codec::Version(1, 0); let codec = VersionedCodec::for_version(registry, version.clone()).unwrap(); - let connection = MTPConnection::from_transport_parts( - version.clone(), - codec, - AlternateSender, - AlternateReceiver, - "/mtp".into(), - Some("browser".into()), - ); + let connection: MTPConnection = + MTPConnection::from_transport_parts( + version.clone(), + codec, + AlternateSender, + AlternateReceiver, + "/mtp".into(), + Some("browser".into()), + ); assert_eq!(connection.version, version); assert_eq!(connection.path, "/mtp"); assert_eq!(connection.description.as_deref(), Some("browser")); diff --git a/host/src/pipe.rs b/host/src/pipe.rs index 6463e04..ed3c623 100644 --- a/host/src/pipe.rs +++ b/host/src/pipe.rs @@ -1,25 +1,106 @@ use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue}; - -use mtp_common::PipeError; - -use mtp_transport::{PipeReader, Policy, Receiver, Sender}; - +use mtp_common::{CommunicationError, PipeError}; +use mtp_transport::{PipeReader, PipeWriter, Policy, TransportEvent}; use std::collections::HashMap; - use std::sync::Arc; - use tokio::sync::{Mutex, mpsc}; -use tracing::debug; +/// The sender operations needed by the transport-independent pipe protocol. +pub trait PipeSender: Clone + Send + Sync + 'static { + type Writer: tokio::io::AsyncWrite + Send + Unpin + 'static; -pub struct PipeHandle { + fn send_pipe_message( + &self, + message: &CommunicationValue, + ) -> impl std::future::Future> + Send; + + fn open_pipe_stream( + &self, + pipe_id: u32, + description: &str, + ) -> impl std::future::Future, CommunicationError>> + Send; +} + +/// The receiver operations needed by the transport-independent pipe protocol. +pub trait PipeReceiver

: Clone + Send + Sync + 'static +where + P: tokio::io::AsyncRead + Send + Unpin + 'static, +{ + fn receive_pipe_event( + &self, + ) -> impl std::future::Future, CommunicationError>> + Send; +} + +impl PipeSender for mtp_transport::Sender { + type Writer = wtransport::SendStream; + + async fn send_pipe_message( + &self, + message: &CommunicationValue, + ) -> Result<(), CommunicationError> { + self.send(message).await + } + + async fn open_pipe_stream( + &self, + pipe_id: u32, + description: &str, + ) -> Result, CommunicationError> { + self.open_pipe(pipe_id, description).await + } +} + +impl PipeReceiver for mtp_transport::Receiver { + async fn receive_pipe_event( + &self, + ) -> Result, CommunicationError> { + self.receive_event().await + } +} + +impl PipeSender for mtp_transport::GenericSender +where + C: mtp_transport::TransportConnection, + C::SendStream: tokio::io::AsyncWrite + Send + Unpin + 'static, +{ + type Writer = C::SendStream; + + async fn send_pipe_message( + &self, + message: &CommunicationValue, + ) -> Result<(), CommunicationError> { + self.send(message).await + } + + async fn open_pipe_stream( + &self, + pipe_id: u32, + description: &str, + ) -> Result, CommunicationError> { + self.open_pipe(pipe_id, description).await + } +} + +impl PipeReceiver for mtp_transport::GenericReceiver +where + C: mtp_transport::TransportConnection, + C::RecvStream: tokio::io::AsyncRead + Send + Unpin + 'static, +{ + async fn receive_pipe_event( + &self, + ) -> Result, CommunicationError> { + self.receive_event().await + } +} + +pub struct PipeHandle { pub(crate) pipe_id: u32, pub(crate) description: String, - pub(crate) sender: Sender, + pub(crate) sender: S, pub(crate) response_rx: tokio::sync::oneshot::Receiver>, } -impl PipeHandle { +impl PipeHandle { pub fn pipe_id(&self) -> u32 { self.pipe_id } @@ -28,31 +109,33 @@ impl PipeHandle { &self.description } - pub async fn wait(self) -> Result, PipeError> { + pub async fn wait(self) -> Result>, PipeError> { match self.response_rx.await { - Ok(Ok(true)) => { - let writer = self - .sender - .open_pipe(self.pipe_id, &self.description) - .await - .map_err(PipeError::from)?; - Ok(Some(writer)) - } + Ok(Ok(true)) => self + .sender + .open_pipe_stream(self.pipe_id, &self.description) + .await + .map(Some) + .map_err(PipeError::from), Ok(Ok(false)) => Ok(None), - Ok(Err(e)) => Err(e), + Ok(Err(error)) => Err(error), Err(_) => Err(PipeError::StreamClosed), } } } -pub struct PipeRequest { +pub struct PipeRequest { pub(crate) pipe_id: u32, pub(crate) description: String, - pub(crate) sender: Sender, - pub(crate) dispatcher: Arc, + pub(crate) sender: S, + pub(crate) dispatcher: Arc>, } -impl PipeRequest { +impl PipeRequest +where + S: PipeSender, + P: tokio::io::AsyncRead + Send + Unpin + 'static, +{ pub fn id(&self) -> u32 { self.pipe_id } @@ -61,133 +144,106 @@ impl PipeRequest { &self.description } - pub async fn accept(self) -> Result { + pub async fn accept(self) -> Result, PipeError> { let (pipe_tx, pipe_rx) = tokio::sync::oneshot::channel(); - { - let mut pending = self.dispatcher.pending_pipes.lock().await; - pending.insert(self.pipe_id, pipe_tx); - } + self.dispatcher + .pending_pipes + .lock() + .await + .insert(self.pipe_id, pipe_tx); - let resp = CommunicationValue::new(CommunicationType::PipeResponse) + let response = CommunicationValue::new(CommunicationType::PipeResponse) .with_id(self.pipe_id) .add_typed_default(DataType::Accepted, DataValue::BoolTrue); - self.sender.send(&resp).await.map_err(PipeError::from)?; + self.sender + .send_pipe_message(&response) + .await + .map_err(PipeError::from)?; - let timeout = self.dispatcher.policy.read_timeout; - tokio::time::timeout(timeout, pipe_rx) + tokio::time::timeout(self.dispatcher.policy.read_timeout, pipe_rx) .await .map_err(|_| PipeError::HandshakeTimeout)? .map_err(|_| PipeError::StreamClosed) } pub async fn deny(self) -> Result<(), PipeError> { - let resp = CommunicationValue::new(CommunicationType::PipeResponse) + let response = CommunicationValue::new(CommunicationType::PipeResponse) .with_id(self.pipe_id) .add_typed_default(DataType::Accepted, DataValue::BoolFalse); - self.sender.send(&resp).await.map_err(PipeError::from)?; - Ok(()) + self.sender + .send_pipe_message(&response) + .await + .map_err(PipeError::from) } } -pub(crate) struct PipeDispatcher { +pub(crate) struct PipeDispatcher

{ pub(crate) pending_creations: Mutex>>>, - pub(crate) pending_pipes: Mutex>>, + pub(crate) pending_pipes: Mutex>>>, pub(crate) policy: Arc, } -impl PipeDispatcher { - pub(crate) fn default_for_external() -> Self { - Self { - pending_creations: Mutex::new(HashMap::new()), - pending_pipes: Mutex::new(HashMap::new()), - policy: Arc::new(Policy::default()), - } - } -} - -pub(crate) async fn run_dispatcher( - receiver: Receiver, - sender: Sender, - app_tx: mpsc::Sender>, - pipe_req_tx: mpsc::Sender, - dispatcher: Arc, -) { +pub(crate) async fn run_dispatcher( + receiver: R, + sender: S, + app_tx: mpsc::Sender>, + pipe_req_tx: mpsc::Sender>, + dispatcher: Arc>, +) where + S: PipeSender, + R: PipeReceiver

, + P: tokio::io::AsyncRead + Send + Unpin + 'static, +{ let pipe_req_type = CommunicationType::PipeRequest.try_to_id(&mtp_codec::TypeMap::latest()); let pipe_resp_type = CommunicationType::PipeResponse.try_to_id(&mtp_codec::TypeMap::latest()); loop { - match receiver.receive_event().await { - Ok(mtp_transport::TransportEvent::Message(msg)) => { - debug!( - target = "mtp.host", - message_type = ?msg.get_type(), - message_id = msg.get_id(), - "dispatcher received message" - ); - if Some(msg.get_type()) == pipe_req_type { - let pipe_id = msg.get_id(); - let description = msg.get_str(DataType::Description).unwrap_or("").to_string(); - debug!( - target = "mtp.host", - pipe_id, description, "dispatcher classified pipe request" - ); - let req = PipeRequest { - pipe_id, - description, + match receiver.receive_pipe_event().await { + Ok(TransportEvent::Message(message)) => { + if Some(message.get_type()) == pipe_req_type { + let request = PipeRequest { + pipe_id: message.get_id(), + description: message + .get_str(DataType::Description) + .unwrap_or("") + .to_owned(), sender: sender.clone(), dispatcher: dispatcher.clone(), }; - let _ = pipe_req_tx.send(req).await; + let _ = pipe_req_tx.send(request).await; continue; } - - if Some(msg.get_type()) == pipe_resp_type { - let pipe_id = msg.get_id(); - let accepted = msg.get_bool(DataType::Accepted).unwrap_or(false); - debug!( - target = "mtp.host", - pipe_id, accepted, "dispatcher classified pipe response" - ); + if Some(message.get_type()) == pipe_resp_type { let mut pending = dispatcher.pending_creations.lock().await; - if let Some(tx) = pending.remove(&pipe_id) { - let _ = tx.send(Ok(accepted)); + if let Some(reply) = pending.remove(&message.get_id()) { + let _ = + reply.send(Ok(message.get_bool(DataType::Accepted).unwrap_or(false))); } continue; } - - if app_tx.send(Ok(msg)).await.is_err() { + if app_tx.send(Ok(message)).await.is_err() { break; } } - Ok(mtp_transport::TransportEvent::Pipe(reader)) => { + Ok(TransportEvent::Pipe(reader)) => { let pipe_id = reader.pipe_id(); - debug!( - target = "mtp.host", - pipe_id, - description = reader.description(), - "dispatcher received pipe stream" - ); let mut pending = dispatcher.pending_pipes.lock().await; - if let Some(tx) = pending.remove(&pipe_id) { - let _ = tx.send(reader); + if let Some(reply) = pending.remove(&pipe_id) { + let _ = reply.send(reader); continue; } - - debug!( - target = "mtp.host", - pipe_id, "dispatcher treating pipe stream as pipe request" - ); - let req = PipeRequest { + drop(pending); + let request = PipeRequest { pipe_id, - description: reader.description().to_string(), + description: reader.description().to_owned(), sender: sender.clone(), dispatcher: dispatcher.clone(), }; - let _ = pipe_req_tx.send(req).await; + let _ = pipe_req_tx.send(request).await; } - Err(e) => { - if app_tx.send(Err(e)).await.is_err() { + Err(error) => { + if app_tx.send(Err(error)).await.is_err() { break; } } diff --git a/mtp-webserver/Cargo.toml b/mtp-webserver/Cargo.toml index b2b6cc0..d7d9777 100644 --- a/mtp-webserver/Cargo.toml +++ b/mtp-webserver/Cargo.toml @@ -8,6 +8,7 @@ mtp-common = { version = "0.2.0", path = "../common" } mtp-codec = { version = "0.2.0", path = "../codec", features = ["registry"] } mtp-host = { version = "0.2.0", path = "../host" } mtp-transport = { version = "0.2.0", path = "../transport" } +mtp-crypto = { version = "0.2.0", path = "../crypto" } bytes = "1" http = "1" tokio = { version = "1", features = ["io-util", "macros", "net", "rt", "sync", "time"] } @@ -19,11 +20,12 @@ rustls = "0.23" tracing = "0.1" thiserror = "2" async-trait = "0.1" +rand = { version = "0.10.1", optional = true } [dev-dependencies] rcgen = "0.14" [features] default = [] -crypto = ["mtp-host/crypto"] +crypto = ["mtp-host/crypto", "dep:rand"] pipes = ["mtp-host/pipes", "mtp-transport/pipes"] diff --git a/mtp-webserver/src/lib.rs b/mtp-webserver/src/lib.rs index 13e45a8..9612b45 100644 --- a/mtp-webserver/src/lib.rs +++ b/mtp-webserver/src/lib.rs @@ -11,6 +11,8 @@ mod stream; mod transport; pub use error::WebServerError; +#[cfg(feature = "pipes")] +pub use mtp_transport::TransportEvent; pub use router::{HttpHandler, Router, RouterError}; pub use server::{MTPWebServer, WebServerConfig, WebServerMetrics}; pub use stream::{Http3Request, Http3Response}; diff --git a/mtp-webserver/src/server.rs b/mtp-webserver/src/server.rs index dd4ce20..c97fb62 100644 --- a/mtp-webserver/src/server.rs +++ b/mtp-webserver/src/server.rs @@ -53,7 +53,7 @@ impl WebServerConfig { max_request_body: 4 * 1024 * 1024, max_connections: 256, request_timeout: Duration::from_secs(30), - drain_timeout: Duration::from_secs(10), + drain_timeout: Duration::from_secs(5), metrics: None, } } @@ -138,18 +138,12 @@ impl MTPWebServer { host_config: HostConfig, web_config: WebServerConfig, ) -> Result { - #[cfg(feature = "crypto")] - if !matches!( - host_config.authentication_policy, - mtp_host::AuthenticationPolicy::Unauthenticated - ) { - return Err(CommunicationError::Other( - "web authentication is not supported yet; use Unauthenticated".into(), - )); - } + let host_config = Arc::new(host_config); let endpoint = build_endpoint(&host_config)?; let driver_endpoint = endpoint.clone(); - let (mtp_tx, mtp_incoming) = tokio::sync::mpsc::channel(16); + // A completed MTP handshake must never block the endpoint driver just + // because the application is briefly slow to call `accept()`. + let (mtp_tx, mtp_incoming) = tokio::sync::mpsc::channel(web_config.max_connections.max(1)); let (shutdown_tx, shutdown_rx) = watch::channel(()); let connection_semaphore = Arc::new(Semaphore::new(web_config.max_connections)); let driver_config = DriverConfig { @@ -160,6 +154,7 @@ impl MTPWebServer { drain_timeout: web_config.drain_timeout, send_pongs: host_config.send_pongs, policy: host_config.policy, + host_config, metrics: web_config.metrics, }; let driver = tokio::spawn(run_driver( @@ -230,11 +225,12 @@ struct DriverConfig { drain_timeout: Duration, send_pongs: bool, policy: mtp_transport::Policy, + host_config: Arc, metrics: Option>, } fn build_endpoint(config: &HostConfig) -> Result { - let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + mtp_crypto::ensure_crypto_provider(); let certificates = rustls::pki_types::CertificateDer::pem_slice_iter(&config.tls_fullchain) .collect::, _>>() .map_err(|_| CommunicationError::CertificateLoadFailed)?; @@ -270,6 +266,7 @@ async fn run_driver( drain_timeout, send_pongs, policy, + host_config, metrics, } = config; let mut connection_tasks = tokio::task::JoinSet::new(); @@ -283,14 +280,20 @@ async fn run_driver( let Some(incoming) = incoming else { break; }; - let permit = match connection_semaphore.clone().acquire_owned().await { + // Do not await capacity here: doing so would prevent this loop + // from observing shutdown while all connection slots are in use. + let permit = match connection_semaphore.clone().try_acquire_owned() { Ok(permit) => permit, - Err(_) => break, + Err(_) => { + tracing::debug!("rejecting QUIC connection at configured connection limit"); + continue; + } }; let router = router.clone(); let mtp_path = mtp_path.clone(); let mtp_tx = mtp_tx.clone(); let metrics = metrics.clone(); + let host_config = host_config.clone(); connection_tasks.spawn(async move { let _permit = permit; let connect_start = std::time::Instant::now(); @@ -364,7 +367,11 @@ async fn run_driver( return; } }; - tasks.spawn(run_session_requests( + // The WebTransport session request driver must outlive this + // endpoint request task. Keep it detached so handing the MTP + // connection to the application does not wait for the session + // (which is intentionally an open-ended accept loop). + tokio::spawn(run_session_requests( session.clone(), router.clone(), max_request_body, @@ -372,9 +379,25 @@ async fn run_driver( metrics.clone(), )); let result = - accept_web_connection(session, mtp_path, connection, send_pongs, policy) + accept_web_connection(session, mtp_path, connection, send_pongs, policy, host_config.clone()) .await; - let _ = mtp_tx.send(result).await; + match mtp_tx.try_send(result) { + Ok(()) => { + // The detached session driver remains active while the + // delivered MTP connection keeps the session alive. + } + Err(tokio::sync::mpsc::error::TrySendError::Full(result)) => { + tracing::warn!("MTP connection backlog is full; dropping connection"); + if let Ok(connection) = result { + connection.sender.close(); + } + } + Err(tokio::sync::mpsc::error::TrySendError::Closed(result)) => { + if let Ok(connection) = result { + connection.sender.close(); + } + } + } return; } let router = router.clone(); diff --git a/mtp-webserver/src/transport.rs b/mtp-webserver/src/transport.rs index 8bc6f48..c320295 100644 --- a/mtp-webserver/src/transport.rs +++ b/mtp-webserver/src/transport.rs @@ -5,12 +5,16 @@ use mtp_codec::{ }; use mtp_common::CommunicationError; use mtp_host::AcceptError; +use mtp_host::HostConfig; use mtp_transport::{ GenericReceiver, GenericSender, Policy, TransportConnection, TransportRecvStream, TransportSendStream, }; use std::sync::Arc; +#[cfg(feature = "crypto")] +use std::time::Instant; use tokio::io::{AsyncReadExt, AsyncWriteExt}; +use tracing::error; type Session = h3_webtransport::server::WebTransportSession; type H3SendStream = h3_webtransport::stream::SendStream, Bytes>; @@ -47,6 +51,13 @@ impl TransportSendStream for H3TransportSender { self.stream .write_all(buf) .await + .map_err(|_| CommunicationError::StreamError)?; + // Control/authentication frames use a persistent stream. h3 keeps + // those writes buffered until flushed; without this the peer can wait + // for the challenge while the server waits for its proof. + self.stream + .flush() + .await .map_err(|_| CommunicationError::StreamError) } @@ -65,7 +76,16 @@ impl TransportRecvStream for H3TransportReceiver { .read_exact(buf) .await .map(|_| ()) - .map_err(|_| CommunicationError::StreamError) + .map_err(|error| { + if error.kind() == std::io::ErrorKind::UnexpectedEof { + // Browser control frames are sent on one-frame uni streams. + // Reaching FIN while looking for another frame is normal. + return CommunicationError::StreamClosed; + } + error!("[mtp-webserver] receive stream read_exact failed ({} bytes): {error}", buf.len()); + tracing::warn!(len = buf.len(), %error, "WebTransport receive stream read_exact failed"); + CommunicationError::StreamError + }) } async fn read_chunk(&mut self, max: usize) -> Result>, CommunicationError> { @@ -76,7 +96,14 @@ impl TransportRecvStream for H3TransportReceiver { buf.truncate(size); Ok(Some(buf)) } - Err(_) => Err(CommunicationError::StreamError), + Err(error) => { + error!( + "[mtp-webserver] receive stream read failed (max {} bytes): {error}", + max + ); + tracing::warn!(max, %error, "WebTransport receive stream read failed"); + Err(CommunicationError::StreamError) + } } } } @@ -129,14 +156,50 @@ impl TransportConnection for H3TransportConnection { } async fn accept_uni(&self) -> Result { + const MAX_CONSECUTIVE_ERRORS: u32 = 10; + const INITIAL_RETRY_DELAY: std::time::Duration = std::time::Duration::from_millis(20); + + let mut consecutive_errors = 0_u32; loop { match self.session.accept_uni().await { Ok(Some((id, stream))) if id == self.session.session_id() => { return Ok(H3TransportReceiver { stream }); } - Ok(Some(_)) => continue, + Ok(Some(_)) => { + consecutive_errors = 0; + continue; + } Ok(None) => return Err(CommunicationError::StreamClosed), - Err(_) => return Err(CommunicationError::StreamError), + Err(error) => { + // A browser can reset an individual pipe stream while it + // is stopping MediaRecorder. h3-webtransport reports that + // through accept_uni even though the QUIC connection is + // still healthy. Do not turn that stream-local failure + // into a connection-wide MTP failure. + if self.quinn.close_reason().is_some() { + return Err(CommunicationError::StreamClosed); + } + + consecutive_errors += 1; + if consecutive_errors > MAX_CONSECUTIVE_ERRORS { + tracing::warn!( + %error, + consecutive_errors, + "WebTransport receive-stream accept repeatedly failed" + ); + return Err(CommunicationError::StreamError); + } + + let multiplier = 1_u32 << consecutive_errors.saturating_sub(1).min(5); + let retry_delay = INITIAL_RETRY_DELAY * multiplier; + tracing::debug!( + %error, + consecutive_errors, + ?retry_delay, + "retrying transient WebTransport receive-stream error" + ); + tokio::time::sleep(retry_delay).await; + } } } } @@ -155,7 +218,8 @@ impl TransportConnection for H3TransportConnection { /// Shared host MTP connection instantiated with HTTP/3 stream adapters. pub type WebMtpSender = GenericSender; pub type WebMtpReceiver = GenericReceiver; -pub type WebMTPConnection = mtp_host::MTPConnection; +pub type WebMTPConnection = + mtp_host::MTPConnection; pub(crate) async fn accept_web_connection( session: Arc, @@ -163,7 +227,33 @@ pub(crate) async fn accept_web_connection( quinn: quinn::Connection, send_pongs: bool, policy: Policy, + host_config: Arc, ) -> Result { + #[cfg(feature = "crypto")] + { + tokio::time::timeout( + host_config.auth_timeout, + accept_web_connection_inner(session, path, quinn, send_pongs, policy, host_config), + ) + .await + .unwrap_or(Err(AcceptError::AuthenticationTimedOut)) + } + + #[cfg(not(feature = "crypto"))] + accept_web_connection_inner(session, path, quinn, send_pongs, policy, host_config).await +} + +async fn accept_web_connection_inner( + session: Arc, + path: String, + quinn: quinn::Connection, + send_pongs: bool, + policy: Policy, + _host_config: Arc, +) -> Result { + #[cfg(feature = "crypto")] + let auth_handshake_started = Instant::now(); + let max_message_size = policy.max_message_size; let transport = H3TransportConnection::new(session, quinn); let policy = Arc::new(policy); let receiver = WebMtpReceiver::new(transport.clone(), policy.clone()); @@ -187,12 +277,261 @@ pub(crate) async fn accept_web_connection( if send_pongs { receiver.respond_to_pings(sender.clone()).await; } - Ok(mtp_host::MTPConnection::from_transport_parts( + let connection = mtp_host::MTPConnection::from_transport_parts( negotiated, codec, sender, receiver, path, - description, - )) + description.clone(), + ); + #[cfg(feature = "crypto")] + let mut connection = connection; + #[cfg(feature = "crypto")] + if !matches!( + _host_config.authentication_policy, + mtp_host::AuthenticationPolicy::Unauthenticated + ) { + use mtp_crypto::{ + Ed25519Signer, MlDsaSigner, PublicKeyBundle, SignatureScheme, auth, verify_ed25519, + verify_ml_dsa, + }; + let tm = mtp_codec::TypeMap::latest(); + let client_lookup_started = Instant::now(); + let (client_id, client_bundle, response_type) = if Some(first.get_type()) + == mtp_codec::CommunicationType::Identification.try_to_id(&tm) + { + let id = match first.get_data(DataType::Id) { + DataValue::UnsignedNumber(value) => *value as u64, + _ => { + return Err(AcceptError::AuthenticationFailed( + "missing client id".into(), + )); + } + }; + let bundle = (_host_config.get_existing_client)(id, description.clone()) + .await + .ok_or_else(|| AcceptError::AuthenticationFailed("unknown client id".into()))?; + ( + id, + bundle, + mtp_codec::CommunicationType::IdentificationResponse, + ) + } else if Some(first.get_type()) == mtp_codec::CommunicationType::Register.try_to_id(&tm) { + let bundle = match first.get_data(DataType::PublicKeys) { + DataValue::Bytes(bytes) => PublicKeyBundle::from_bytes(bytes).map_err(|_| { + AcceptError::AuthenticationFailed("invalid public key bundle".into()) + })?, + _ => { + return Err(AcceptError::AuthenticationFailed( + "missing public keys".into(), + )); + } + }; + (0, bundle, mtp_codec::CommunicationType::RegisterResponse) + } else { + return Err(AcceptError::AuthenticationFailed( + "unexpected authentication message".into(), + )); + }; + tracing::debug!(elapsed = ?client_lookup_started.elapsed(), "web authentication handshake: identify client"); + + let signer_init_started = Instant::now(); + let host_pq_signer = if !_host_config + .host_keyring + .sig_pq_secret_key + .as_bytes() + .is_empty() + { + Some( + MlDsaSigner::new( + &_host_config.host_keyring.sig_pq_secret_key, + &_host_config.host_keyring.sig_pq_public_key, + ) + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?, + ) + } else { + None + }; + tracing::debug!(elapsed = ?signer_init_started.elapsed(), "web authentication handshake: signer initialization"); + + let server_challenge: u128 = rand::random(); + let host_sign = |payload: &[u8]| -> Result<(Vec, Vec), AcceptError> { + let signer = Ed25519Signer::new(&_host_config.host_keyring.sig_cl_secret_key) + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?; + let sig = signer + .sign(payload) + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?; + let pq = if let Some(pq_signer) = host_pq_signer.as_ref() { + pq_signer + .sign(payload) + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))? + } else { + Vec::new() + }; + Ok((sig, pq)) + }; + let sign_challenge_started = Instant::now(); + let (sig, pq_sig) = host_sign(&auth::challenge_payload(client_id, server_challenge))?; + tracing::debug!(elapsed = ?sign_challenge_started.elapsed(), "web authentication handshake: sign challenge"); + let mut challenge = + mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Challenge) + .add_typed_default( + DataType::ServerNonce, + DataValue::UnsignedNumber(server_challenge), + ) + .add_typed_default(DataType::Signature, DataValue::Bytes(sig)) + .add_typed_default( + DataType::RequirePq, + if _host_config.require_pq { + DataValue::BoolTrue + } else { + DataValue::BoolFalse + }, + ); + if !pq_sig.is_empty() { + challenge = + challenge.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_sig)); + } + let send_challenge_started = Instant::now(); + connection + .sender + .send(&challenge) + .await + .map_err(AcceptError::Send)?; + tracing::debug!(elapsed = ?send_challenge_started.elapsed(), "web authentication handshake: send challenge"); + let receive_proof_started = Instant::now(); + let proof = connection.receive().await.map_err(AcceptError::Receive)?; + tracing::debug!(elapsed = ?receive_proof_started.elapsed(), "web authentication handshake: receive client proof"); + if Some(proof.get_type()) != mtp_codec::CommunicationType::ChallengeResponse.try_to_id(&tm) + { + return Err(AcceptError::AuthenticationFailed( + "missing challenge response".into(), + )); + } + let nonce = match proof.get_data(DataType::ClientNonce) { + DataValue::UnsignedNumber(n) => *n, + _ => { + return Err(AcceptError::AuthenticationFailed( + "missing client nonce".into(), + )); + } + }; + let signature = match proof.get_data(DataType::Signature) { + DataValue::Bytes(bytes) => bytes, + _ => { + return Err(AcceptError::AuthenticationFailed( + "missing challenge signature".into(), + )); + } + }; + let pq_signature = match proof.get_data(DataType::PqSignature) { + DataValue::Bytes(bytes) => bytes.as_slice(), + _ => &[], + }; + let payload = if first.get_type() + == mtp_codec::CommunicationType::Register + .try_to_id(&tm) + .unwrap() + { + auth::register_proof_payload( + &version.to_string(), + &client_bundle.as_bytes(), + server_challenge, + nonce, + ) + } else { + auth::login_proof_payload(&version.to_string(), client_id, server_challenge, nonce) + }; + let verify_proof_started = Instant::now(); + if verify_ed25519(&client_bundle.sig_cl_public_key, &payload, signature).is_err() { + return Err(AcceptError::AuthenticationFailed( + "client proof signature invalid".into(), + )); + } + let client_has_pq = !client_bundle.sig_pq_public_key.as_bytes().is_empty(); + if _host_config.require_pq + && (!client_has_pq + || pq_signature.is_empty() + || verify_ml_dsa(&client_bundle.sig_pq_public_key, &payload, pq_signature).is_err()) + { + return Err(AcceptError::AuthenticationFailed( + "client PQ proof signature invalid".into(), + )); + } + tracing::debug!(elapsed = ?verify_proof_started.elapsed(), "web authentication handshake: verify client proof"); + let register_started = Instant::now(); + let assigned_id = if response_type == mtp_codec::CommunicationType::RegisterResponse { + (_host_config.complete_register)(client_bundle.clone(), description.clone()).await + } else { + client_id + }; + tracing::debug!(elapsed = ?register_started.elapsed(), "web authentication handshake: registration callback"); + let sign_final_started = Instant::now(); + let (final_sig, final_pq) = host_sign(&auth::host_final_payload( + assigned_id, + nonce, + server_challenge, + ))?; + tracing::debug!(elapsed = ?sign_final_started.elapsed(), "web authentication handshake: sign final response"); + let mut response = mtp_codec::CommunicationValue::new(response_type) + .add_typed_default(DataType::Connected, DataValue::BoolTrue) + .add_typed_default(DataType::Id, DataValue::UnsignedNumber(assigned_id as u128)) + .add_typed_default(DataType::ClientNonce, DataValue::UnsignedNumber(nonce)) + .add_typed_default(DataType::Signature, DataValue::Bytes(final_sig)) + .add_typed_default(DataType::Version, DataValue::Str(version.to_string())); + if !final_pq.is_empty() { + response = + response.add_typed_default(DataType::PqSignature, DataValue::Bytes(final_pq)); + } + let send_final_started = Instant::now(); + connection + .sender + .send(&response) + .await + .map_err(AcceptError::Send)?; + connection + .sender + .finish_stream() + .await + .map_err(AcceptError::Send)?; + tracing::debug!(elapsed = ?send_final_started.elapsed(), "web authentication handshake: send final response"); + tracing::debug!(elapsed = ?auth_handshake_started.elapsed(), "web authentication handshake: complete"); + connection.receiver.set_max_message_size(max_message_size); + connection.auth_state = mtp_host::AuthState::Authenticated; + connection.client_id = assigned_id; + connection.client_public_key = Some(client_bundle); + return Ok(connection); + } + + // Complete the opening handshake for unauthenticated connections. Native clients + // wait for this response before sending application messages. + let response = + mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::IdentificationResponse) + .add_typed_default( + mtp_codec::DataType::Connected, + mtp_codec::DataValue::BoolTrue, + ) + .add_typed_default( + mtp_codec::DataType::Version, + mtp_codec::DataValue::Str(connection.version.to_string()), + ) + .add_typed_default( + mtp_codec::DataType::Id, + // WebTransport connections currently do not expose the host's guest + // ID through MTPConnection; unauthenticated clients do not need it. + mtp_codec::DataValue::UnsignedNumber(0), + ); + connection + .sender + .send(&response) + .await + .map_err(AcceptError::Send)?; + connection + .sender + .finish_stream() + .await + .map_err(AcceptError::Send)?; + connection.receiver.set_max_message_size(max_message_size); + Ok(connection) } diff --git a/src/vite/index.ts b/src/vite/index.ts index 9ecbf62..0653c8e 100644 --- a/src/vite/index.ts +++ b/src/vite/index.ts @@ -371,6 +371,13 @@ export function mtp(options: MTPVitePluginOptions): VitePlugin { await buildIfNeeded(state); return { + // The generated wasm-bindgen JavaScript imports its sibling `.wasm` + // by a relative URL. Prebundling it independently lets Vite retain an + // older wrapper while the plugin has rebuilt the wasm binary, which + // produces missing closure-export errors at runtime. + optimizeDeps: { + exclude: ["mtp", "mtp/raw", "mtp/type-map"], + }, resolve: { preserveSymlinks: true, alias: { diff --git a/transport/Cargo.toml b/transport/Cargo.toml index 335b0ae..bc227ff 100644 --- a/transport/Cargo.toml +++ b/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", diff --git a/transport/src/client.rs b/transport/src/client.rs index 42d906c..6136927 100644 --- a/transport/src/client.rs +++ b/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)) } diff --git a/transport/src/connection.rs b/transport/src/connection.rs index b135564..726ba1a 100644 --- a/transport/src/connection.rs +++ b/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; diff --git a/transport/src/generic_connection.rs b/transport/src/generic_connection.rs index 4835be7..e389e30 100644 --- a/transport/src/generic_connection.rs +++ b/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 GenericSender { 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 { pipes: Arc>>>, connection: C, ping_sender: Arc>>>, + max_message_size: Arc, } impl Clone for GenericReceiver { @@ -159,6 +176,7 @@ impl Clone for GenericReceiver { 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 GenericReceiver { #[cfg(feature = "pipes")] let (pipe_tx, pipe_rx) = mpsc::channel(policy.receiver_queue_capacity); let ping_sender: Arc>>> = 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 GenericReceiver { #[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 GenericReceiver { .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 GenericReceiver { 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 GenericReceiver { 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 GenericReceiver { pipes: Arc::new(Mutex::new(pipe_rx)), connection, ping_sender, + max_message_size, } } pub async fn respond_to_pings(&self, sender: GenericSender) { *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 { self.incoming .lock() diff --git a/transport/src/host.rs b/transport/src/host.rs index 8e40a7b..65b4fe8 100644 --- a/transport/src/host.rs +++ b/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 { - 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, ) { + 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( diff --git a/transport/tests/integration.rs b/transport/tests/integration.rs index 477b0ce..16d42de 100644 --- a/transport/tests/integration.rs +++ b/transport/tests/integration.rs @@ -105,7 +105,7 @@ async fn test_explicit_development_tls() -> Result<(), Box Result<(), Box> 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> { 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> { // 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> { .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 Result<(), Box Result<(), Box