From cab2cd7a52137eee2f25a08b9b3e7830c981e9ce Mon Sep 17 00:00:00 2001 From: Alex Date: Tue, 28 Jul 2026 18:49:40 +0200 Subject: [PATCH] [Fix] Syncronized Webserver & Host behaviour, Fixed the 10 sec default wait on auth --- example/.gitignore | 1 + example/Cargo.lock | 4 + example/client/Cargo.toml | 2 + example/client/src/auth.rs | 16 +- example/client/src/main.rs | 42 +- example/client/src/messages.rs | 15 +- example/client/src/metrics.rs | 554 +++++++++++++++++ example/client/src/pipes.rs | 15 +- example/server/Cargo.toml | 2 + example/server/src/main.rs | 69 ++- example/server/src/metrics.rs | 883 +++++++++++++++++++++++++++ host/src/connection.rs | 77 ++- host/src/engine.rs | 825 +++++++++++++++++++++++++ host/src/error.rs | 41 +- host/src/handshake.rs | 910 +++------------------------- host/src/lib.rs | 2 + mtp-webserver/src/h3.rs | 46 +- mtp-webserver/src/server.rs | 17 +- mtp-webserver/src/transport.rs | 345 ++++++++--- transport/src/connection.rs | 24 +- transport/src/generic_connection.rs | 114 +++- transport/src/host.rs | 51 +- 22 files changed, 2978 insertions(+), 1077 deletions(-) create mode 100644 example/client/src/metrics.rs create mode 100644 example/server/src/metrics.rs create mode 100644 host/src/engine.rs diff --git a/example/.gitignore b/example/.gitignore index c8a5b0c..7f5a6d9 100644 --- a/example/.gitignore +++ b/example/.gitignore @@ -14,3 +14,4 @@ web-client/dist/ client.id *.mk *.mpkb +metrics/ diff --git a/example/Cargo.lock b/example/Cargo.lock index c84365a..7b0f6cc 100644 --- a/example/Cargo.lock +++ b/example/Cargo.lock @@ -237,6 +237,8 @@ version = "0.2.0" dependencies = [ "mtp", "rand 0.10.2", + "serde", + "serde_json", "tokio", "tracing-subscriber", ] @@ -1988,6 +1990,8 @@ dependencies = [ "hex", "http", "mtp", + "rand 0.10.2", + "serde", "serde_json", "tokio", "tracing-subscriber", diff --git a/example/client/Cargo.toml b/example/client/Cargo.toml index e56862d..38fe2ad 100644 --- a/example/client/Cargo.toml +++ b/example/client/Cargo.toml @@ -12,3 +12,5 @@ mtp = { version = "0.2.0", path = "../../", features = ["client", "crypto", "fil tokio = { version = "1", features = ["full"] } rand = "0.10.1" tracing-subscriber = "0.3.23" +serde = { version = "1", features = ["derive"] } +serde_json = "1" diff --git a/example/client/src/auth.rs b/example/client/src/auth.rs index 56e1304..cd947d7 100644 --- a/example/client/src/auth.rs +++ b/example/client/src/auth.rs @@ -1,4 +1,4 @@ -use std::time::Instant; +use std::time::{Duration, Instant}; use tokio::fs; @@ -12,7 +12,7 @@ pub async fn connect_or_register( mut config: ClientConfig, host_public_key: PublicKeyBundle, key_prefix: &str, -) -> Result<(MTPConnection, Keyring), Box> { +) -> Result<(MTPConnection, Keyring, String, Duration), Box> { let keyring_path = format!("{key_prefix}.mk"); let id_path = format!("{key_prefix}.id"); @@ -30,12 +30,12 @@ pub async fn connect_or_register( config.client_id = client_id; let auth_started = Instant::now(); let conn = MTPClient::auth_connect(config, &keyring, &host_public_key).await?; + let auth_duration = auth_started.elapsed(); println!( "Authenticated (version {}) in {:?}", - conn.version, - auth_started.elapsed() + conn.version, auth_duration ); - return Ok((conn, keyring)); + return Ok((conn, keyring, "connect".into(), auth_duration)); } println!("No existing keys found: registering new client"); @@ -52,12 +52,14 @@ pub async fn connect_or_register( sig_sk, ); + let reg_started = Instant::now(); let conn = MTPClient::auth_register(config, &keyring, &host_public_key).await?; - println!("Registered with ID: {}", conn.client_id); + let reg_duration = reg_started.elapsed(); + println!("Registered with ID: {} in {:?}", conn.client_id, reg_duration); save_keyring_raw(&keyring, &keyring_path)?; fs::write(&id_path, conn.client_id.to_string()).await?; println!("Saved client keys -> {keyring_path}"); - Ok((conn, keyring)) + Ok((conn, keyring, "register".into(), reg_duration)) } diff --git a/example/client/src/main.rs b/example/client/src/main.rs index 4e4b7eb..c7358c6 100644 --- a/example/client/src/main.rs +++ b/example/client/src/main.rs @@ -1,9 +1,11 @@ mod auth; +mod metrics; mod messages; mod pipes; use std::fs; use std::path::Path; +use std::time::Duration; use mtp::client::ClientConfig; use mtp::files::load_public_key_bundle; @@ -36,6 +38,8 @@ async fn main() -> Result<(), Box> { } }; + let mut client_metrics = metrics::ClientMetrics::load("metrics/client_sessions.json"); + println!("Connecting to 127.0.0.1:8080 ..."); let config = ClientConfig::new("https://127.0.0.1:8080") @@ -43,11 +47,43 @@ async fn main() -> Result<(), Box> { .with_description("MTP example client"); let server_bundle = host_public_key.clone(); - let (conn, keyring) = auth::connect_or_register(config, host_public_key, "client").await?; - messages::send_and_receive(&conn, &keyring, &server_bundle).await?; + let (conn, keyring, auth_method, auth_duration) = + match auth::connect_or_register(config, host_public_key, "client").await { + Ok(result) => result, + Err(e) => { + let mut builder = metrics::SessionBuilder::new("failed", Duration::from_secs(0)); + builder.set_error(e.to_string()); + client_metrics.record_session(builder.build()); + client_metrics.save("metrics/client_sessions.json"); + client_metrics.build_overview("metrics/client_overview.json"); + return Err(e); + } + }; + + let mut builder = metrics::SessionBuilder::new(&auth_method, auth_duration); + + let roundtrip = messages::send_and_receive(&conn, &keyring, &server_bundle).await?; + builder.set_message_roundtrip(roundtrip); println!("\n--- Pipe demo ---"); - pipes::run_pipe_demo(&conn, 1).await?; + let pipe_results = pipes::run_pipe_demo(&conn, 1).await?; + for result in &pipe_results { + builder.add_pipe_result(result.clone()); + } + + let session_record = builder.build(); + println!( + "\nSession {} complete: auth={}ms, msg_roundtrip={}ms, pipes={} results, pipe_bytes={}", + session_record.session_id, + session_record.auth_duration_ms, + session_record.message_roundtrip_ms, + session_record.pipe_results.len(), + session_record.total_pipe_bytes, + ); + + client_metrics.record_session(session_record); + client_metrics.save("metrics/client_sessions.json"); + client_metrics.build_overview("metrics/client_overview.json"); conn.sender.close().await; println!("\nDone"); diff --git a/example/client/src/messages.rs b/example/client/src/messages.rs index 8959c92..90da68f 100644 --- a/example/client/src/messages.rs +++ b/example/client/src/messages.rs @@ -1,3 +1,5 @@ +use std::time::{Duration, Instant}; + use mtp::client::MTPConnection; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; use mtp::crypto::{Ed25519Signer, EncryptionType, Keyring, PublicKeyBundle, SigAlgorithm}; @@ -89,17 +91,22 @@ pub async fn send_and_receive( conn: &MTPConnection, keyring: &Keyring, server_bundle: &PublicKeyBundle, -) -> Result<(), Box> { +) -> Result> { let msg = build_demo_message(conn.client_id, keyring, server_bundle)?; println!("Sending: {msg}"); + let start = Instant::now(); conn.sender.send(&msg).await?; match conn.receive().await { Ok(resp) => { + let roundtrip = start.elapsed(); println!("Received: {resp}"); + println!("Message round-trip: {:.3}ms", roundtrip.as_secs_f64() * 1000.0); + Ok(roundtrip) + } + Err(e) => { + eprintln!("Receive error: {e}"); + Err(e.into()) } - Err(e) => eprintln!("Receive error: {e}"), } - - Ok(()) } diff --git a/example/client/src/metrics.rs b/example/client/src/metrics.rs new file mode 100644 index 0000000..59c7e98 --- /dev/null +++ b/example/client/src/metrics.rs @@ -0,0 +1,554 @@ +use serde::{Deserialize, Serialize}; +use std::path::Path; +use std::time::{Duration, SystemTime, UNIX_EPOCH}; + +fn now_epoch_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +fn generate_session_id() -> String { + let ts = now_epoch_secs(); + let rand_part: u32 = rand::random(); + format!("{ts}-{rand_part:08x}") +} + +// --------------------------------------------------------------------------- +// Persisted data types +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct PipeResult { + pub size: usize, + pub iteration: usize, + pub total_ms: f64, + pub data_only_ms: f64, + pub bytes_matched: bool, +} + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct ClientSessionRecord { + pub session_id: String, + pub timestamp: u64, + pub auth_method: String, + pub auth_duration_ms: f64, + pub error: Option, + pub message_roundtrip_ms: f64, + pub pipe_results: Vec, + pub total_pipe_bytes: u64, + pub overall_pipe_avg_total_ms: f64, + pub overall_pipe_avg_data_ms: f64, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct ClientAggregateStats { + pub total_sessions: u64, + pub auth_failures: u64, + pub avg_auth_duration_ms: f64, + pub avg_message_roundtrip_ms: f64, + pub avg_pipe_total_ms: f64, + pub avg_pipe_data_ms: f64, + pub total_pipe_bytes: u64, + pub avg_pipe_throughput_mbps: f64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ClientOverview { + pub total_sessions: u64, + pub aggregate: ClientAggregateStats, + pub sessions: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct ClientMetricsFile { + pub sessions: Vec, +} + +// --------------------------------------------------------------------------- +// Live metrics state +// --------------------------------------------------------------------------- + +pub struct ClientMetrics { + sessions: Vec, +} + +impl ClientMetrics { + pub fn new() -> Self { + Self { + sessions: Vec::new(), + } + } + + pub fn load(path: &str) -> Self { + let file = std::fs::read_to_string(path) + .ok() + .and_then(|s| serde_json::from_str::(&s).ok()); + + Self { + sessions: file.map(|f| f.sessions).unwrap_or_default(), + } + } + + pub fn save(&self, path: &str) { + let data = ClientMetricsFile { + sessions: self.sessions.clone(), + }; + if let Some(parent) = Path::new(path).parent() { + let _ = std::fs::create_dir_all(parent); + } + let json = serde_json::to_string_pretty(&data).unwrap_or_default(); + let _ = std::fs::write(path, json); + } + + pub fn record_session(&mut self, record: ClientSessionRecord) { + self.sessions.push(record); + } + + pub fn build_overview(&self, overview_path: &str) { + let total = self.sessions.len() as u64; + + if total == 0 { + let overview = ClientOverview { + total_sessions: 0, + aggregate: ClientAggregateStats::default(), + sessions: Vec::new(), + }; + if let Some(parent) = Path::new(overview_path).parent() { + let _ = std::fs::create_dir_all(parent); + } + let json = serde_json::to_string_pretty(&overview).unwrap_or_default(); + let _ = std::fs::write(overview_path, json); + return; + } + + let mut auth_sum: f64 = 0.0; + let mut msg_sum: f64 = 0.0; + let mut pipe_total_sum: f64 = 0.0; + let mut pipe_data_sum: f64 = 0.0; + let mut total_pipe_bytes: u64 = 0; + let mut total_pipe_duration_secs: f64 = 0.0; + let mut auth_failures: u64 = 0; + let mut success_count: u64 = 0; + + for s in &self.sessions { + if s.error.is_some() { + auth_failures += 1; + } else { + success_count += 1; + auth_sum += s.auth_duration_ms; + msg_sum += s.message_roundtrip_ms; + pipe_total_sum += s.overall_pipe_avg_total_ms; + pipe_data_sum += s.overall_pipe_avg_data_ms; + total_pipe_bytes += s.total_pipe_bytes; + for pr in &s.pipe_results { + total_pipe_duration_secs += pr.total_ms / 1000.0; + } + } + } + + let divisor = if success_count > 0 { success_count } else { 1 }; + + let aggregate = ClientAggregateStats { + total_sessions: total, + auth_failures, + avg_auth_duration_ms: auth_sum / divisor as f64, + avg_message_roundtrip_ms: msg_sum / divisor as f64, + avg_pipe_total_ms: pipe_total_sum / divisor as f64, + avg_pipe_data_ms: pipe_data_sum / divisor as f64, + total_pipe_bytes, + avg_pipe_throughput_mbps: if total_pipe_duration_secs > 0.0 { + (total_pipe_bytes as f64 / 1_048_576.0) / total_pipe_duration_secs + } else { + 0.0 + }, + }; + + let overview = ClientOverview { + total_sessions: total, + aggregate, + sessions: self.sessions.clone(), + }; + + if let Some(parent) = Path::new(overview_path).parent() { + let _ = std::fs::create_dir_all(parent); + } + let json = serde_json::to_string_pretty(&overview).unwrap_or_default(); + let _ = std::fs::write(overview_path, json); + } +} + +// --------------------------------------------------------------------------- +// Builder for constructing a session record piece by piece +// --------------------------------------------------------------------------- + +pub struct SessionBuilder { + session_id: String, + timestamp: u64, + auth_method: String, + auth_duration_ms: f64, + error: Option, + message_roundtrip_ms: f64, + pipe_results: Vec, +} + +impl SessionBuilder { + pub fn new(auth_method: &str, auth_duration: Duration) -> Self { + Self { + session_id: generate_session_id(), + timestamp: now_epoch_secs(), + auth_method: auth_method.to_string(), + auth_duration_ms: auth_duration.as_secs_f64() * 1000.0, + error: None, + message_roundtrip_ms: 0.0, + pipe_results: Vec::new(), + } + } + + pub fn set_error(&mut self, error: String) { + self.error = Some(error); + } + + pub fn set_message_roundtrip(&mut self, duration: Duration) { + self.message_roundtrip_ms = duration.as_secs_f64() * 1000.0; + } + + pub fn add_pipe_result(&mut self, result: PipeResult) { + self.pipe_results.push(result); + } + + pub fn build(self) -> ClientSessionRecord { + let total_pipe_bytes: u64 = self + .pipe_results + .iter() + .map(|r| r.size as u64) + .sum(); + + let overall_pipe_avg_total_ms = if self.pipe_results.is_empty() { + 0.0 + } else { + self.pipe_results.iter().map(|r| r.total_ms).sum::() + / self.pipe_results.len() as f64 + }; + + let overall_pipe_avg_data_ms = if self.pipe_results.is_empty() { + 0.0 + } else { + self.pipe_results.iter().map(|r| r.data_only_ms).sum::() + / self.pipe_results.len() as f64 + }; + + ClientSessionRecord { + session_id: self.session_id, + timestamp: self.timestamp, + auth_method: self.auth_method, + auth_duration_ms: self.auth_duration_ms, + error: self.error, + message_roundtrip_ms: self.message_roundtrip_ms, + pipe_results: self.pipe_results, + total_pipe_bytes, + overall_pipe_avg_total_ms, + overall_pipe_avg_data_ms, + } + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + fn tmp_path(name: &str) -> String { + let dir = std::env::temp_dir().join("mtp_client_metrics_test"); + let _ = std::fs::create_dir_all(&dir); + dir.join(name).to_str().unwrap().to_string() + } + + #[test] + fn test_pipe_result_roundtrip() { + let pr = PipeResult { + size: 1024, + iteration: 0, + total_ms: 5.5, + data_only_ms: 3.2, + bytes_matched: true, + }; + let json = serde_json::to_string(&pr).unwrap(); + let decoded: PipeResult = serde_json::from_str(&json).unwrap(); + assert_eq!(pr, decoded); + } + + #[test] + fn test_client_session_roundtrip() { + let record = ClientSessionRecord { + session_id: "test-session".into(), + timestamp: 12345, + auth_method: "connect".into(), + auth_duration_ms: 42.5, + error: None, + message_roundtrip_ms: 10.3, + pipe_results: vec![ + PipeResult { + size: 64, + iteration: 0, + total_ms: 1.0, + data_only_ms: 0.5, + bytes_matched: true, + }, + PipeResult { + size: 256, + iteration: 0, + total_ms: 2.0, + data_only_ms: 1.0, + bytes_matched: true, + }, + ], + total_pipe_bytes: 320, + overall_pipe_avg_total_ms: 1.5, + overall_pipe_avg_data_ms: 0.75, + }; + + let json = serde_json::to_string(&record).unwrap(); + let decoded: ClientSessionRecord = serde_json::from_str(&json).unwrap(); + assert_eq!(record, decoded); + } + + #[test] + fn test_client_metrics_load_missing() { + let metrics = ClientMetrics::load("/nonexistent/path.json"); + assert!(metrics.sessions.is_empty()); + } + + #[test] + fn test_multiple_client_sessions() { + let path = tmp_path("multi_session.json"); + let mut metrics = ClientMetrics::load(&path); + + for i in 0..3 { + let mut builder = SessionBuilder::new("connect", Duration::from_millis(10 + i)); + builder.set_message_roundtrip(Duration::from_millis(5 + i)); + builder.add_pipe_result(PipeResult { + size: 64, + iteration: 0, + total_ms: 1.0 + i as f64, + data_only_ms: 0.5 + i as f64 * 0.5, + bytes_matched: true, + }); + metrics.record_session(builder.build()); + } + + metrics.save(&path); + + let metrics2 = ClientMetrics::load(&path); + assert_eq!(metrics2.sessions.len(), 3); + assert_eq!(metrics2.sessions[0].auth_method, "connect"); + assert_eq!(metrics2.sessions[1].pipe_results[0].size, 64); + + let _ = std::fs::remove_file(&path); + } + + #[test] + fn test_client_overview_stats() { + let mut metrics = ClientMetrics::new(); + + for i in 0..4 { + let mut builder = SessionBuilder::new("connect", Duration::from_millis(20)); + builder.set_message_roundtrip(Duration::from_millis(10 + i as u64)); + builder.add_pipe_result(PipeResult { + size: 256, + iteration: 0, + total_ms: 2.0, + data_only_ms: 1.0, + bytes_matched: true, + }); + metrics.record_session(builder.build()); + } + + let overview_path = tmp_path("client_overview.json"); + metrics.build_overview(&overview_path); + + let json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: ClientOverview = serde_json::from_str(&json).unwrap(); + + assert_eq!(overview.total_sessions, 4); + assert_eq!(overview.aggregate.avg_auth_duration_ms, 20.0); + assert_eq!(overview.aggregate.avg_message_roundtrip_ms, 11.5); + assert_eq!(overview.aggregate.avg_pipe_total_ms, 2.0); + assert_eq!(overview.aggregate.avg_pipe_data_ms, 1.0); + assert_eq!(overview.aggregate.total_pipe_bytes, 1024); + assert!(overview.aggregate.avg_pipe_throughput_mbps > 0.0); + assert_eq!(overview.sessions.len(), 4); + + let _ = std::fs::remove_file(&overview_path); + } + + #[test] + fn test_client_overview_empty() { + let metrics = ClientMetrics::new(); + let overview_path = tmp_path("client_empty_overview.json"); + metrics.build_overview(&overview_path); + + let json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: ClientOverview = serde_json::from_str(&json).unwrap(); + assert_eq!(overview.total_sessions, 0); + assert!(overview.sessions.is_empty()); + + let _ = std::fs::remove_file(&overview_path); + } + + #[test] + fn test_auth_failure_recording() { + let path = tmp_path("auth_failure.json"); + let overview_path = tmp_path("auth_failure_overview.json"); + + let mut metrics = ClientMetrics::load(&path); + + // Successful session + let mut b1 = SessionBuilder::new("connect", Duration::from_millis(42)); + b1.set_message_roundtrip(Duration::from_millis(10)); + metrics.record_session(b1.build()); + + // Failed auth session + let mut b2 = SessionBuilder::new("connect", Duration::from_millis(5000)); + b2.set_error("authentication timed out".into()); + metrics.record_session(b2.build()); + + // Another successful session + let mut b3 = SessionBuilder::new("register", Duration::from_millis(100)); + b3.set_message_roundtrip(Duration::from_millis(8)); + metrics.record_session(b3.build()); + + metrics.save(&path); + let metrics2 = ClientMetrics::load(&path); + assert_eq!(metrics2.sessions.len(), 3); + assert!(metrics2.sessions[0].error.is_none()); + assert_eq!( + metrics2.sessions[1].error.as_deref(), + Some("authentication timed out") + ); + assert!(metrics2.sessions[2].error.is_none()); + + metrics2.build_overview(&overview_path); + let overview_json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: ClientOverview = serde_json::from_str(&overview_json).unwrap(); + + assert_eq!(overview.total_sessions, 3); + assert_eq!(overview.aggregate.auth_failures, 1); + // Averages should only count successful sessions + assert!((overview.aggregate.avg_auth_duration_ms - 71.0).abs() < 0.01); // (42+100)/2 + assert!((overview.aggregate.avg_message_roundtrip_ms - 9.0).abs() < 0.01); // (10+8)/2 + + let _ = std::fs::remove_file(&path); + let _ = std::fs::remove_file(&overview_path); + } + + // ----------------------------------------------------------------------- + // Integration-style tests + // ----------------------------------------------------------------------- + + fn make_pr(size: usize, iteration: usize, total_ms: f64, data_only_ms: f64) -> PipeResult { + PipeResult { + size, + iteration, + total_ms, + data_only_ms, + bytes_matched: true, + } + } + + #[test] + fn test_full_client_lifecycle() { + let path = tmp_path("client_lifecycle.json"); + let overview_path = tmp_path("client_lifecycle_overview.json"); + + let mut metrics = ClientMetrics::load(&path); + + let mut b1 = SessionBuilder::new("connect", Duration::from_millis(42)); + b1.set_message_roundtrip(Duration::from_millis(10)); + b1.add_pipe_result(make_pr(64, 0, 1.5, 0.8)); + b1.add_pipe_result(make_pr(256, 0, 2.5, 1.2)); + metrics.record_session(b1.build()); + + let mut b2 = SessionBuilder::new("register", Duration::from_millis(150)); + b2.set_message_roundtrip(Duration::from_millis(15)); + b2.add_pipe_result(make_pr(64, 0, 2.0, 1.0)); + b2.add_pipe_result(make_pr(1024, 0, 5.0, 3.0)); + metrics.record_session(b2.build()); + + metrics.save(&path); + let metrics2 = ClientMetrics::load(&path); + assert_eq!(metrics2.sessions.len(), 2); + + let s1 = &metrics2.sessions[0]; + assert_eq!(s1.auth_method, "connect"); + assert!((s1.auth_duration_ms - 42.0).abs() < 0.01); + assert!((s1.message_roundtrip_ms - 10.0).abs() < 0.01); + assert_eq!(s1.pipe_results.len(), 2); + assert_eq!(s1.total_pipe_bytes, 320); + assert!((s1.overall_pipe_avg_total_ms - 2.0).abs() < 0.01); + assert!((s1.overall_pipe_avg_data_ms - 1.0).abs() < 0.01); + + let s2 = &metrics2.sessions[1]; + assert_eq!(s2.auth_method, "register"); + assert_eq!(s2.pipe_results.len(), 2); + assert_eq!(s2.total_pipe_bytes, 1088); + + metrics2.build_overview(&overview_path); + let overview_json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: ClientOverview = serde_json::from_str(&overview_json).unwrap(); + assert_eq!(overview.total_sessions, 2); + assert!((overview.aggregate.avg_auth_duration_ms - 96.0).abs() < 0.01); + assert!((overview.aggregate.avg_message_roundtrip_ms - 12.5).abs() < 0.01); + assert_eq!(overview.aggregate.total_pipe_bytes, 1408); + assert!(overview.aggregate.avg_pipe_throughput_mbps > 0.0); + + let _ = std::fs::remove_file(&path); + let _ = std::fs::remove_file(&overview_path); + } + + #[test] + fn test_cross_session_accumulation() { + let path = tmp_path("client_accumulate.json"); + let overview_path = tmp_path("client_accumulate_overview.json"); + + { + let mut metrics = ClientMetrics::load(&path); + let mut b = SessionBuilder::new("connect", Duration::from_millis(30)); + b.set_message_roundtrip(Duration::from_millis(8)); + b.add_pipe_result(make_pr(64, 0, 1.0, 0.5)); + metrics.record_session(b.build()); + metrics.save(&path); + } + + { + let mut metrics = ClientMetrics::load(&path); + assert_eq!(metrics.sessions.len(), 1); + let mut b = SessionBuilder::new("register", Duration::from_millis(200)); + b.set_message_roundtrip(Duration::from_millis(12)); + b.add_pipe_result(make_pr(1024, 0, 4.0, 2.5)); + metrics.record_session(b.build()); + metrics.save(&path); + } + + let metrics = ClientMetrics::load(&path); + assert_eq!(metrics.sessions.len(), 2); + assert_eq!(metrics.sessions[0].auth_method, "connect"); + assert_eq!(metrics.sessions[1].auth_method, "register"); + + metrics.build_overview(&overview_path); + let overview_json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: ClientOverview = serde_json::from_str(&overview_json).unwrap(); + assert_eq!(overview.total_sessions, 2); + assert!((overview.aggregate.avg_auth_duration_ms - 115.0).abs() < 0.01); + assert!((overview.aggregate.avg_message_roundtrip_ms - 10.0).abs() < 0.01); + assert_eq!(overview.aggregate.total_pipe_bytes, 1088); + + let _ = std::fs::remove_file(&path); + let _ = std::fs::remove_file(&overview_path); + } +} diff --git a/example/client/src/pipes.rs b/example/client/src/pipes.rs index 5ae6edd..080a7c2 100644 --- a/example/client/src/pipes.rs +++ b/example/client/src/pipes.rs @@ -3,13 +3,16 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::sync::oneshot; use tokio::time::{Duration, Instant}; +use crate::metrics::PipeResult; + pub async fn run_pipe_demo( conn: &MTPConnection, iterations: usize, -) -> Result<(), Box> { +) -> Result, Box> { let sizes = [64, 256, 1024, 4096]; let mut all_elapsed = Vec::with_capacity(sizes.len() * iterations); let mut all_data_only = Vec::with_capacity(sizes.len() * iterations); + let mut pipe_results = Vec::with_capacity(sizes.len() * iterations); for (i, &size) in sizes.iter().enumerate() { let mut size_elapsed = Vec::with_capacity(iterations); @@ -97,6 +100,14 @@ pub async fn run_pipe_demo( data_only_elapsed.as_secs_f64() * 1000.0, ); + pipe_results.push(PipeResult { + size, + iteration: run, + total_ms: overall_elapsed.as_secs_f64() * 1000.0, + data_only_ms: data_only_elapsed.as_secs_f64() * 1000.0, + bytes_matched: matches, + }); + size_elapsed.push(overall_elapsed); size_data_only.push(data_only_elapsed); all_elapsed.push(overall_elapsed); @@ -123,7 +134,7 @@ pub async fn run_pipe_demo( all_elapsed.len() ); - Ok(()) + Ok(pipe_results) } /// Helper: average a slice of Durations without overflowing. diff --git a/example/server/Cargo.toml b/example/server/Cargo.toml index dcca160..a2586e6 100644 --- a/example/server/Cargo.toml +++ b/example/server/Cargo.toml @@ -15,3 +15,5 @@ serde_json = { version = "1" } hex = "0.4" base64 = "0.22" tracing-subscriber = "0.3.23" +serde = { version = "1", features = ["derive"] } +rand = "0.10.1" diff --git a/example/server/src/main.rs b/example/server/src/main.rs index 13f87a7..596c42d 100644 --- a/example/server/src/main.rs +++ b/example/server/src/main.rs @@ -1,6 +1,7 @@ mod clients; mod handlers; mod keys; +mod metrics; mod tls; #[path = "web-server.rs"] mod web_server; @@ -39,7 +40,7 @@ async fn handle_pipe_loopback( mtp::webserver::WebMtpSender, mtp::webserver::H3TransportReceiver, >, -) -> Result<(), Box> { +) -> Result> { let pipe_id = request.id(); println!(" [loopback] Accepting pipe {pipe_id} ..."); let mut reader = request.accept().await?; @@ -56,7 +57,7 @@ async fn handle_pipe_loopback( let copied = tokio::io::copy(&mut reader, &mut writer).await?; writer.finish_async().await?; println!(" [loopback] Pipe {pipe_id} complete ({copied} bytes)"); - Ok(()) + Ok(copied) } #[tokio::main] @@ -119,6 +120,10 @@ async fn main() -> Result<(), Box> { }, ); + let metrics = std::sync::Arc::new(metrics::ServerMetrics::load( + "metrics/server_sessions.json", + )); + println!("Starting integrated MTP web server on port 8080 ..."); let config = HostConfig::new( @@ -138,8 +143,22 @@ async fn main() -> Result<(), Box> { println!("TCP: HTTP/1.1 and HTTP/2"); println!("UDP: HTTP/3 and WebTransport"); - while let Some(conn) = host.accept().await? { + loop { + let conn = match host.accept().await { + Ok(Some(conn)) => conn, + Ok(None) => break, + Err(e) => { + let msg = e.to_string(); + eprintln!("Accept error: {msg}"); + metrics.record_accept_error(); + metrics.save("metrics/server_sessions.json"); + metrics.build_overview("metrics/server_overview.json"); + continue; + } + }; let decrypt_keyring = Arc::clone(&decrypt_keyring); + let metrics = Arc::clone(&metrics); + metrics.record_connection_version(&conn.version.to_string()); tokio::spawn(async move { let desc = conn.description.as_deref().unwrap_or("(no description)"); println!( @@ -151,12 +170,14 @@ async fn main() -> Result<(), Box> { ); println!("Client ID: {}", conn.client_id); + let mut session = metrics.start_session(conn.client_id, desc.to_string()); + let tm: &TypeMap = conn.codec.registry().get(&conn.version).unwrap(); println!("Waiting for messages / pipe requests ..."); let mut pipe_open = true; let mut message_open = true; - let mut messages_received = 0_u64; + let mut exit_reason = "normal".to_string(); while pipe_open || message_open { let activity = tokio::time::timeout(CONNECTION_IDLE_TIMEOUT, async { @@ -165,8 +186,17 @@ async fn main() -> Result<(), Box> { 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}"); + match handle_pipe_loopback(&conn, request).await { + Ok(bytes) => { + session.record_pipe(bytes); + } + Err(error) => { + let msg = error.to_string(); + if msg.contains("denied") { + session.record_pipe_denial(); + } + eprintln!(" [loopback] Pipe error: {msg}"); + } } } Err(mtp::common::CommunicationError::StreamClosed) @@ -183,18 +213,24 @@ async fn main() -> Result<(), Box> { message = conn.receive(), if message_open => { match message { Ok(message) => { - messages_received += 1; println!("Received: {message}"); - match handlers::process_and_respond( + let msg_start = std::time::Instant::now(); + let result = handlers::process_and_respond( &message, tm, conn.client_public_key.as_ref(), &decrypt_keyring, - ) { + ); + let latency = msg_start.elapsed(); + let ok = result.is_ok(); + session.record_message(latency, ok); + + match result { Ok(response) => { println!("Sending: {response}"); if let Err(error) = conn.sender.send(&response).await { eprintln!("Send error: {error}"); + session.record_send_error(); pipe_open = false; message_open = false; } @@ -220,16 +256,27 @@ async fn main() -> Result<(), Box> { .await; if activity.is_err() { + exit_reason = "idle timeout".to_string(); println!("Connection idle timeout reached"); break; } - if messages_received >= MAX_MESSAGES_PER_CONNECTION { + if session.messages_received() >= MAX_MESSAGES_PER_CONNECTION { + exit_reason = "message limit".to_string(); println!("Connection message limit reached"); break; } } - println!("Connection closed\n"); + let record = session.finish(exit_reason); + println!( + "Connection closed (messages: {}, pipes: {}, duration: {:.1}s)\n", + record.messages_received, + record.pipes_handled, + record.duration_secs + ); + + metrics.save("metrics/server_sessions.json"); + metrics.build_overview("metrics/server_overview.json"); }); } diff --git a/example/server/src/metrics.rs b/example/server/src/metrics.rs new file mode 100644 index 0000000..314a891 --- /dev/null +++ b/example/server/src/metrics.rs @@ -0,0 +1,883 @@ +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::path::Path; +use std::sync::Mutex; +use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH}; + +fn now_epoch_secs() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or_default() + .as_secs() +} + +fn generate_session_id() -> String { + let ts = now_epoch_secs(); + let rand_part: u32 = rand::random(); + format!("{ts}-{rand_part:08x}") +} + +// --------------------------------------------------------------------------- +// Persisted data types +// --------------------------------------------------------------------------- + +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct SessionRecord { + pub session_id: String, + pub client_id: u64, + pub description: String, + pub start_time: u64, + pub end_time: u64, + pub duration_secs: f64, + pub messages_received: u64, + pub messages_ok: u64, + pub messages_failed: u64, + pub pipes_handled: u64, + pub pipe_bytes_copied: u64, + pub pipe_denials: u64, + pub send_errors: u64, + pub avg_message_latency_ms: f64, + pub max_message_latency_ms: f64, + pub exit_reason: String, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct AggregateStats { + pub total_connections: u64, + pub total_messages: u64, + pub total_messages_ok: u64, + pub total_messages_failed: u64, + pub total_pipes: u64, + pub total_pipe_bytes: u64, + pub total_pipe_denials: u64, + pub total_send_errors: u64, + pub total_accept_errors: u64, + pub avg_session_duration_secs: f64, + pub avg_messages_per_session: f64, + pub avg_pipes_per_session: f64, + pub avg_message_latency_ms: f64, + pub max_message_latency_ms: f64, +} + +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct Overview { + pub total_sessions: u64, + pub first_session_timestamp: u64, + pub last_session_timestamp: u64, + pub aggregate: AggregateStats, + pub connection_versions: HashMap, + pub sessions: Vec, +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default)] +pub struct ServerMetricsFile { + pub total_connections: u64, + pub total_messages: u64, + pub total_messages_ok: u64, + pub total_messages_failed: u64, + pub total_pipes: u64, + pub total_pipe_bytes: u64, + pub total_pipe_denials: u64, + pub total_send_errors: u64, + pub total_accept_errors: u64, + pub connection_versions: HashMap, + pub sessions: Vec, +} + +// --------------------------------------------------------------------------- +// Live metrics state +// --------------------------------------------------------------------------- + +struct Inner { + total_connections: u64, + total_messages: u64, + total_messages_ok: u64, + total_messages_failed: u64, + total_pipes: u64, + total_pipe_bytes: u64, + total_pipe_denials: u64, + total_send_errors: u64, + total_accept_errors: u64, + connection_versions: HashMap, + active_connections: u64, + completed_sessions: Vec, +} + +pub struct ServerMetrics { + inner: Mutex, +} + +impl ServerMetrics { + pub fn new() -> Self { + Self { + inner: Mutex::new(Inner { + total_connections: 0, + total_messages: 0, + total_messages_ok: 0, + total_messages_failed: 0, + total_pipes: 0, + total_pipe_bytes: 0, + total_pipe_denials: 0, + total_send_errors: 0, + total_accept_errors: 0, + connection_versions: HashMap::new(), + active_connections: 0, + completed_sessions: Vec::new(), + }), + } + } + + pub fn load(path: &str) -> Self { + let file = std::fs::read_to_string(path) + .ok() + .and_then(|s| serde_json::from_str::(&s).ok()); + + let mut inner = Inner { + total_connections: 0, + total_messages: 0, + total_messages_ok: 0, + total_messages_failed: 0, + total_pipes: 0, + total_pipe_bytes: 0, + total_pipe_denials: 0, + total_send_errors: 0, + total_accept_errors: 0, + connection_versions: HashMap::new(), + active_connections: 0, + completed_sessions: Vec::new(), + }; + + if let Some(data) = file { + inner.total_connections = data.total_connections; + inner.total_messages = data.total_messages; + inner.total_messages_ok = data.total_messages_ok; + inner.total_messages_failed = data.total_messages_failed; + inner.total_pipes = data.total_pipes; + inner.total_pipe_bytes = data.total_pipe_bytes; + inner.total_pipe_denials = data.total_pipe_denials; + inner.total_send_errors = data.total_send_errors; + inner.total_accept_errors = data.total_accept_errors; + inner.connection_versions = data.connection_versions; + inner.completed_sessions = data.sessions; + } + + Self { + inner: Mutex::new(inner), + } + } + + pub fn save(&self, path: &str) { + let inner = self.inner.lock().unwrap(); + let data = self.to_file(&inner); + if let Some(parent) = Path::new(path).parent() { + let _ = std::fs::create_dir_all(parent); + } + let json = serde_json::to_string_pretty(&data).unwrap_or_default(); + let _ = std::fs::write(path, json); + } + + fn to_file(&self, inner: &Inner) -> ServerMetricsFile { + ServerMetricsFile { + total_connections: inner.total_connections, + total_messages: inner.total_messages, + total_messages_ok: inner.total_messages_ok, + total_messages_failed: inner.total_messages_failed, + total_pipes: inner.total_pipes, + total_pipe_bytes: inner.total_pipe_bytes, + total_pipe_denials: inner.total_pipe_denials, + total_send_errors: inner.total_send_errors, + total_accept_errors: inner.total_accept_errors, + connection_versions: inner.connection_versions.clone(), + sessions: inner.completed_sessions.clone(), + } + } + + pub fn start_session(&self, client_id: u64, description: String) -> SessionHandle<'_> { + let session_id = generate_session_id(); + let start = Instant::now(); + let start_time = now_epoch_secs(); + + self.inner.lock().unwrap().total_connections += 1; + self.inner.lock().unwrap().active_connections += 1; + + SessionHandle { + metrics: self, + session_id, + client_id, + description, + start, + start_time, + messages_received: 0, + messages_ok: 0, + messages_failed: 0, + pipes_handled: 0, + pipe_bytes: 0, + pipe_denials: 0, + send_errors: 0, + latencies: Vec::new(), + } + } + + pub fn snapshot(&self) -> ServerMetricsFile { + let inner = self.inner.lock().unwrap(); + self.to_file(&inner) + } + + pub fn record_accept_error(&self) { + self.inner.lock().unwrap().total_accept_errors += 1; + } + + pub fn record_connection_version(&self, version: &str) { + *self + .inner + .lock() + .unwrap() + .connection_versions + .entry(version.to_string()) + .or_insert(0) += 1; + } + + pub fn build_overview(&self, overview_path: &str) { + let inner = self.inner.lock().unwrap(); + let sessions = &inner.completed_sessions; + let total = sessions.len() as u64; + + if total == 0 { + let overview = Overview { + total_sessions: 0, + first_session_timestamp: 0, + last_session_timestamp: 0, + aggregate: AggregateStats::default(), + connection_versions: HashMap::new(), + sessions: Vec::new(), + }; + if let Some(parent) = Path::new(overview_path).parent() { + let _ = std::fs::create_dir_all(parent); + } + let json = serde_json::to_string_pretty(&overview).unwrap_or_default(); + let _ = std::fs::write(overview_path, json); + return; + } + + let first_ts = sessions.first().map(|s| s.start_time).unwrap_or(0); + let last_ts = sessions.last().map(|s| s.end_time).unwrap_or(0); + + let total_duration: f64 = sessions.iter().map(|s| s.duration_secs).sum(); + let total_msgs: u64 = sessions.iter().map(|s| s.messages_received).sum(); + let total_pipes: u64 = sessions.iter().map(|s| s.pipes_handled).sum(); + + let mut max_latency: f64 = 0.0; + let mut latency_sum: f64 = 0.0; + let mut latency_count: u64 = 0; + for s in sessions { + if s.avg_message_latency_ms > 0.0 { + latency_sum += s.avg_message_latency_ms * s.messages_ok as f64; + latency_count += s.messages_ok; + } + if s.max_message_latency_ms > max_latency { + max_latency = s.max_message_latency_ms; + } + } + + let aggregate = AggregateStats { + total_connections: inner.total_connections, + total_messages: inner.total_messages, + total_messages_ok: inner.total_messages_ok, + total_messages_failed: inner.total_messages_failed, + total_pipes: inner.total_pipes, + total_pipe_bytes: inner.total_pipe_bytes, + total_pipe_denials: inner.total_pipe_denials, + total_send_errors: inner.total_send_errors, + total_accept_errors: inner.total_accept_errors, + avg_session_duration_secs: total_duration / total as f64, + avg_messages_per_session: total_msgs as f64 / total as f64, + avg_pipes_per_session: total_pipes as f64 / total as f64, + avg_message_latency_ms: if latency_count > 0 { + latency_sum / latency_count as f64 + } else { + 0.0 + }, + max_message_latency_ms: max_latency, + }; + + let overview = Overview { + total_sessions: total, + first_session_timestamp: first_ts, + last_session_timestamp: last_ts, + aggregate, + connection_versions: inner.connection_versions.clone(), + sessions: sessions.clone(), + }; + + if let Some(parent) = Path::new(overview_path).parent() { + let _ = std::fs::create_dir_all(parent); + } + let json = serde_json::to_string_pretty(&overview).unwrap_or_default(); + let _ = std::fs::write(overview_path, json); + } + + fn finish_session(&self, record: SessionRecord) { + let mut inner = self.inner.lock().unwrap(); + inner.active_connections -= 1; + inner.total_messages += record.messages_received; + inner.total_messages_ok += record.messages_ok; + inner.total_messages_failed += record.messages_failed; + inner.total_pipes += record.pipes_handled; + inner.total_pipe_bytes += record.pipe_bytes_copied; + inner.total_pipe_denials += record.pipe_denials; + inner.total_send_errors += record.send_errors; + inner.completed_sessions.push(record); + } +} + +// --------------------------------------------------------------------------- +// Session handle — local accumulators, no mutex contention during connection +// --------------------------------------------------------------------------- + +pub struct SessionHandle<'a> { + metrics: &'a ServerMetrics, + session_id: String, + client_id: u64, + description: String, + start: Instant, + start_time: u64, + messages_received: u64, + messages_ok: u64, + messages_failed: u64, + pipes_handled: u64, + pipe_bytes: u64, + pipe_denials: u64, + send_errors: u64, + latencies: Vec, +} + +impl<'a> SessionHandle<'a> { + pub fn messages_received(&self) -> u64 { + self.messages_received + } + + pub fn record_message(&mut self, latency: Duration, ok: bool) { + self.messages_received += 1; + if ok { + self.messages_ok += 1; + } else { + self.messages_failed += 1; + } + self.latencies.push(latency.as_secs_f64() * 1000.0); + } + + pub fn record_pipe(&mut self, bytes: u64) { + self.pipes_handled += 1; + self.pipe_bytes += bytes; + } + + pub fn record_pipe_denial(&mut self) { + self.pipe_denials += 1; + } + + pub fn record_send_error(&mut self) { + self.send_errors += 1; + } + + pub fn finish(self, exit_reason: String) -> SessionRecord { + let elapsed = self.start.elapsed(); + let end_time = self.start_time + elapsed.as_secs(); + + let avg_latency = if self.latencies.is_empty() { + 0.0 + } else { + self.latencies.iter().sum::() / self.latencies.len() as f64 + }; + let max_latency = self.latencies.iter().copied().fold(0.0_f64, f64::max); + + let record = SessionRecord { + session_id: self.session_id, + client_id: self.client_id, + description: self.description, + start_time: self.start_time, + end_time, + duration_secs: elapsed.as_secs_f64(), + messages_received: self.messages_received, + messages_ok: self.messages_ok, + messages_failed: self.messages_failed, + pipes_handled: self.pipes_handled, + pipe_bytes_copied: self.pipe_bytes, + pipe_denials: self.pipe_denials, + send_errors: self.send_errors, + avg_message_latency_ms: avg_latency, + max_message_latency_ms: max_latency, + exit_reason, + }; + + self.metrics.finish_session(record.clone()); + record + } +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + fn tmp_path(name: &str) -> String { + let dir = std::env::temp_dir().join("mtp_server_metrics_test"); + let _ = std::fs::create_dir_all(&dir); + dir.join(name).to_str().unwrap().to_string() + } + + #[test] + fn test_session_record_roundtrip() { + let record = SessionRecord { + session_id: "test-123".into(), + client_id: 1000, + description: "test session".into(), + start_time: 1000, + end_time: 1010, + duration_secs: 10.0, + messages_received: 5, + messages_ok: 4, + messages_failed: 1, + pipes_handled: 2, + pipe_bytes_copied: 4096, + pipe_denials: 0, + send_errors: 0, + avg_message_latency_ms: 1.5, + max_message_latency_ms: 3.0, + exit_reason: "normal".into(), + }; + + let json = serde_json::to_string(&record).unwrap(); + let decoded: SessionRecord = serde_json::from_str(&json).unwrap(); + assert_eq!(record, decoded); + } + + #[test] + fn test_metrics_file_roundtrip() { + let file = ServerMetricsFile { + total_connections: 10, + total_messages: 50, + total_messages_ok: 48, + total_messages_failed: 2, + total_pipes: 5, + total_pipe_bytes: 20480, + total_pipe_denials: 1, + total_send_errors: 0, + total_accept_errors: 3, + connection_versions: HashMap::from([("2.0".into(), 8), ("1.0".into(), 2)]), + sessions: vec![ + SessionRecord { + session_id: "s1".into(), + client_id: 1000, + description: "first".into(), + start_time: 100, + end_time: 110, + duration_secs: 10.0, + messages_received: 3, + messages_ok: 3, + messages_failed: 0, + pipes_handled: 1, + pipe_bytes_copied: 1024, + pipe_denials: 0, + send_errors: 0, + avg_message_latency_ms: 0.5, + max_message_latency_ms: 1.0, + exit_reason: "normal".into(), + }, + SessionRecord { + session_id: "s2".into(), + client_id: 1001, + description: "second".into(), + start_time: 200, + end_time: 230, + duration_secs: 30.0, + messages_received: 7, + messages_ok: 6, + messages_failed: 1, + pipes_handled: 4, + pipe_bytes_copied: 19456, + pipe_denials: 1, + send_errors: 0, + avg_message_latency_ms: 2.0, + max_message_latency_ms: 5.0, + exit_reason: "idle timeout".into(), + }, + ], + }; + + let json = serde_json::to_string_pretty(&file).unwrap(); + let decoded: ServerMetricsFile = serde_json::from_str(&json).unwrap(); + assert_eq!(file.total_connections, decoded.total_connections); + assert_eq!(file.sessions.len(), decoded.sessions.len()); + assert_eq!(file.sessions[0], decoded.sessions[0]); + assert_eq!(file.sessions[1], decoded.sessions[1]); + } + + #[test] + fn test_session_handle_lifecycle() { + let metrics = ServerMetrics::new(); + let mut session = metrics.start_session(1000, "test".into()); + + session.record_message(Duration::from_millis(1), true); + session.record_message(Duration::from_millis(3), true); + session.record_message(Duration::from_millis(2), false); + session.record_pipe(512); + + let record = session.finish("test exit".into()); + + assert_eq!(record.client_id, 1000); + assert_eq!(record.messages_received, 3); + assert_eq!(record.messages_ok, 2); + assert_eq!(record.messages_failed, 1); + assert_eq!(record.pipes_handled, 1); + assert_eq!(record.pipe_bytes_copied, 512); + assert!(record.avg_message_latency_ms > 0.0); + assert_eq!(record.max_message_latency_ms, 3.0); + assert_eq!(record.exit_reason, "test exit"); + + let snap = metrics.snapshot(); + assert_eq!(snap.total_connections, 1); + assert_eq!(snap.total_messages, 3); + assert_eq!(snap.total_messages_ok, 2); + assert_eq!(snap.total_messages_failed, 1); + assert_eq!(snap.total_pipes, 1); + assert_eq!(snap.total_pipe_bytes, 512); + assert_eq!(snap.sessions.len(), 1); + } + + #[test] + fn test_overview_generation() { + let metrics = ServerMetrics::new(); + + for i in 0..3 { + let mut session = metrics.start_session(1000 + i as u64, format!("session {i}")); + for _ in 0..(i + 1) * 2 { + session.record_message(Duration::from_millis(1 + i), true); + } + session.record_pipe((i as u64 + 1) * 1000); + session.finish(format!("exit {i}")); + } + + let overview_path = tmp_path("overview_test.json"); + metrics.build_overview(&overview_path); + + let json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: Overview = serde_json::from_str(&json).unwrap(); + + assert_eq!(overview.total_sessions, 3); + assert!(overview.first_session_timestamp > 0); + assert!(overview.last_session_timestamp >= overview.first_session_timestamp); + assert_eq!(overview.aggregate.total_connections, 3); + assert_eq!(overview.aggregate.total_messages, 12); // 2+4+6 + assert_eq!(overview.aggregate.total_pipes, 3); + assert_eq!(overview.aggregate.total_pipe_bytes, 6000); // 1000+2000+3000 + assert!(overview.aggregate.avg_session_duration_secs >= 0.0); + assert_eq!(overview.sessions.len(), 3); + + let _ = std::fs::remove_file(&overview_path); + } + + #[test] + fn test_load_missing_file() { + let metrics = ServerMetrics::load("/nonexistent/path/metrics.json"); + let snap = metrics.snapshot(); + assert_eq!(snap.total_connections, 0); + assert!(snap.sessions.is_empty()); + } + + #[test] + fn test_multiple_sessions_accumulate() { + let path = tmp_path("accumulate_test.json"); + let metrics = ServerMetrics::load(&path); + + for i in 0..5 { + let mut session = metrics.start_session(1000, format!("s{i}")); + session.record_message(Duration::from_millis(1), true); + session.record_pipe(100); + session.finish(format!("done {i}")); + } + + metrics.save(&path); + + let metrics2 = ServerMetrics::load(&path); + let snap = metrics2.snapshot(); + assert_eq!(snap.total_connections, 5); + assert_eq!(snap.total_messages, 5); + assert_eq!(snap.total_messages_ok, 5); + assert_eq!(snap.total_pipes, 5); + assert_eq!(snap.total_pipe_bytes, 500); + assert_eq!(snap.sessions.len(), 5); + + let _ = std::fs::remove_file(&path); + } + + #[test] + fn test_overview_latencies() { + let metrics = ServerMetrics::new(); + + let mut s1 = metrics.start_session(1000, "s1".into()); + s1.record_message(Duration::from_millis(2), true); + s1.record_message(Duration::from_millis(4), true); + s1.finish("done".into()); + + let mut s2 = metrics.start_session(1001, "s2".into()); + s2.record_message(Duration::from_millis(1), true); + s2.finish("done".into()); + + let overview_path = tmp_path("latency_overview.json"); + metrics.build_overview(&overview_path); + let json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: Overview = serde_json::from_str(&json).unwrap(); + + // s1 avg = 3.0, s2 avg = 1.0 + // weighted avg = (3*2 + 1*1) / 3 = 7/3 ≈ 2.333 + assert!( + (overview.aggregate.avg_message_latency_ms - 7.0 / 3.0).abs() < 0.01, + "avg latency: {}", + overview.aggregate.avg_message_latency_ms + ); + assert_eq!(overview.aggregate.max_message_latency_ms, 4.0); + + let _ = std::fs::remove_file(&overview_path); + } + + #[test] + fn test_overview_empty() { + let metrics = ServerMetrics::new(); + let overview_path = tmp_path("empty_overview.json"); + metrics.build_overview(&overview_path); + + let json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: Overview = serde_json::from_str(&json).unwrap(); + assert_eq!(overview.total_sessions, 0); + assert!(overview.sessions.is_empty()); + + let _ = std::fs::remove_file(&overview_path); + } + + // ----------------------------------------------------------------------- + // Integration-style tests + // ----------------------------------------------------------------------- + + #[test] + fn test_full_session_lifecycle() { + let path = tmp_path("lifecycle.json"); + let overview_path = tmp_path("lifecycle_overview.json"); + + let metrics = ServerMetrics::load(&path); + + let mut s1 = metrics.start_session(1000, "first".into()); + s1.record_message(Duration::from_millis(1), true); + s1.record_message(Duration::from_millis(2), true); + let r1 = s1.finish("normal".into()); + + let mut s2 = metrics.start_session(1001, "second".into()); + s2.record_message(Duration::from_millis(5), true); + s2.record_message(Duration::from_millis(3), false); + s2.record_pipe(2048); + s2.record_pipe(4096); + let r2 = s2.finish("idle timeout".into()); + + let mut s3 = metrics.start_session(1002, "third".into()); + s3.record_pipe(1024); + let r3 = s3.finish("normal".into()); + + assert_eq!(r1.client_id, 1000); + assert_eq!(r1.messages_received, 2); + assert_eq!(r1.messages_ok, 2); + assert_eq!(r1.pipes_handled, 0); + + assert_eq!(r2.client_id, 1001); + assert_eq!(r2.messages_received, 2); + assert_eq!(r2.messages_ok, 1); + assert_eq!(r2.messages_failed, 1); + assert_eq!(r2.pipes_handled, 2); + assert_eq!(r2.pipe_bytes_copied, 6144); + assert_eq!(r2.exit_reason, "idle timeout"); + + assert_eq!(r3.client_id, 1002); + assert_eq!(r3.messages_received, 0); + assert_eq!(r3.pipes_handled, 1); + assert_eq!(r3.pipe_bytes_copied, 1024); + + let snap = metrics.snapshot(); + assert_eq!(snap.total_connections, 3); + assert_eq!(snap.total_messages, 4); + assert_eq!(snap.total_messages_ok, 3); + assert_eq!(snap.total_messages_failed, 1); + assert_eq!(snap.total_pipes, 3); + assert_eq!(snap.total_pipe_bytes, 7168); + assert_eq!(snap.sessions.len(), 3); + + metrics.save(&path); + let metrics2 = ServerMetrics::load(&path); + let snap2 = metrics2.snapshot(); + assert_eq!(snap2.total_connections, 3); + assert_eq!(snap2.total_messages, 4); + assert_eq!(snap2.sessions.len(), 3); + assert_eq!(snap2.sessions[1].exit_reason, "idle timeout"); + + metrics2.build_overview(&overview_path); + let overview_json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: Overview = serde_json::from_str(&overview_json).unwrap(); + assert_eq!(overview.total_sessions, 3); + assert_eq!(overview.aggregate.total_connections, 3); + assert_eq!(overview.aggregate.total_messages, 4); + assert_eq!(overview.aggregate.total_messages_ok, 3); + assert_eq!(overview.aggregate.total_messages_failed, 1); + assert_eq!(overview.aggregate.total_pipes, 3); + assert_eq!(overview.aggregate.total_pipe_bytes, 7168); + assert!(overview.aggregate.avg_session_duration_secs >= 0.0); + assert!(overview.aggregate.avg_messages_per_session > 0.0); + assert_eq!(overview.sessions.len(), 3); + + let _ = std::fs::remove_file(&path); + let _ = std::fs::remove_file(&overview_path); + } + + #[test] + fn test_overview_rebuild_accuracy() { + let path = tmp_path("accuracy.json"); + let overview_path = tmp_path("accuracy_overview.json"); + + let metrics = ServerMetrics::load(&path); + + for i in 0..10u32 { + let mut session = metrics.start_session(1000 + i as u64, format!("session {i}")); + let msg_count = (i + 1) * 2; + for j in 0..msg_count { + session.record_message(Duration::from_millis((j + 1) as u64), j % 3 != 0); + } + session.record_pipe((i as u64 + 1) * 512); + session.finish(format!("exit {i}")); + } + + metrics.save(&path); + let metrics2 = ServerMetrics::load(&path); + metrics2.build_overview(&overview_path); + + let overview_json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: Overview = serde_json::from_str(&overview_json).unwrap(); + + assert_eq!(overview.total_sessions, 10); + assert_eq!(overview.aggregate.total_connections, 10); + assert_eq!(overview.aggregate.total_messages, 110); + assert_eq!(overview.aggregate.total_pipes, 10); + assert_eq!(overview.aggregate.total_pipe_bytes, 28160); + assert!(overview.aggregate.avg_session_duration_secs >= 0.0); + assert!((overview.aggregate.avg_messages_per_session - 11.0).abs() < 0.01); + assert!((overview.aggregate.avg_pipes_per_session - 1.0).abs() < 0.01); + + let _ = std::fs::remove_file(&path); + let _ = std::fs::remove_file(&overview_path); + } + + #[test] + fn test_persistence_across_instances() { + let path = tmp_path("persistence.json"); + let overview_path = tmp_path("persistence_overview.json"); + + { + let metrics = ServerMetrics::load(&path); + let mut s1 = metrics.start_session(1000, "inst1-s1".into()); + s1.record_message(Duration::from_millis(10), true); + s1.record_pipe(100); + s1.finish("done".into()); + + let mut s2 = metrics.start_session(1001, "inst1-s2".into()); + s2.record_message(Duration::from_millis(20), true); + s2.finish("done".into()); + + metrics.save(&path); + metrics.build_overview(&overview_path); + } + + { + let metrics = ServerMetrics::load(&path); + let snap = metrics.snapshot(); + assert_eq!(snap.sessions.len(), 2); + assert_eq!(snap.total_connections, 2); + + let mut s3 = metrics.start_session(1002, "inst2-s1".into()); + s3.record_message(Duration::from_millis(5), true); + s3.record_pipe(200); + s3.record_pipe(300); + s3.finish("done".into()); + + metrics.save(&path); + metrics.build_overview(&overview_path); + } + + let metrics = ServerMetrics::load(&path); + let snap = metrics.snapshot(); + assert_eq!(snap.sessions.len(), 3); + assert_eq!(snap.total_connections, 3); + assert_eq!(snap.total_messages, 3); + assert_eq!(snap.total_messages_ok, 3); + assert_eq!(snap.total_pipes, 3); + assert_eq!(snap.total_pipe_bytes, 600); + + let overview_json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: Overview = serde_json::from_str(&overview_json).unwrap(); + assert_eq!(overview.total_sessions, 3); + assert_eq!(overview.sessions[0].description, "inst1-s1"); + assert_eq!(overview.sessions[1].description, "inst1-s2"); + assert_eq!(overview.sessions[2].description, "inst2-s1"); + + let _ = std::fs::remove_file(&path); + let _ = std::fs::remove_file(&overview_path); + } + + #[test] + fn test_accept_errors_and_versions() { + let path = tmp_path("accept_errors.json"); + let overview_path = tmp_path("accept_errors_overview.json"); + + let metrics = ServerMetrics::load(&path); + + // Simulate 5 accept errors + for _ in 0..5 { + metrics.record_accept_error(); + } + + // Simulate connection versions + metrics.record_connection_version("2.0"); + metrics.record_connection_version("2.0"); + metrics.record_connection_version("1.0"); + + // A normal session with pipe denials and send errors + let mut s1 = metrics.start_session(1000, "normal".into()); + s1.record_message(Duration::from_millis(1), true); + s1.record_pipe_denial(); + s1.record_send_error(); + s1.record_send_error(); + s1.finish("done".into()); + + metrics.save(&path); + let metrics2 = ServerMetrics::load(&path); + let snap = metrics2.snapshot(); + assert_eq!(snap.total_accept_errors, 5); + assert_eq!(snap.connection_versions["2.0"], 2); + assert_eq!(snap.connection_versions["1.0"], 1); + assert_eq!(snap.total_pipe_denials, 1); + assert_eq!(snap.total_send_errors, 2); + assert_eq!(snap.sessions.len(), 1); + assert_eq!(snap.sessions[0].pipe_denials, 1); + assert_eq!(snap.sessions[0].send_errors, 2); + + metrics2.build_overview(&overview_path); + let overview_json = std::fs::read_to_string(&overview_path).unwrap(); + let overview: Overview = serde_json::from_str(&overview_json).unwrap(); + assert_eq!(overview.aggregate.total_accept_errors, 5); + assert_eq!(overview.aggregate.total_pipe_denials, 1); + assert_eq!(overview.aggregate.total_send_errors, 2); + assert_eq!(overview.connection_versions["2.0"], 2); + assert_eq!(overview.connection_versions["1.0"], 1); + + let _ = std::fs::remove_file(&path); + let _ = std::fs::remove_file(&overview_path); + } +} diff --git a/host/src/connection.rs b/host/src/connection.rs index 3d09e28..ad7c330 100644 --- a/host/src/connection.rs +++ b/host/src/connection.rs @@ -3,7 +3,6 @@ use mtp_codec::{CommunicationType, DataType, DataValue}; use mtp_codec::{CommunicationValue, Version, registry::VersionedCodec}; use mtp_common::CommunicationError; use std::net::SocketAddr; -#[cfg(feature = "pipes")] use std::sync::Arc; #[cfg(feature = "pipes")] use tokio::sync::{Mutex, mpsc}; @@ -12,7 +11,6 @@ use tokio::sync::{Mutex, mpsc}; use crate::error::random_client_id; #[cfg(feature = "pipes")] use crate::pipe::{PipeDispatcher, PipeReceiver, PipeRequest, PipeSender, run_dispatcher}; -#[cfg(feature = "pipes")] use mtp_transport::Policy; mod connection_capability { @@ -161,6 +159,52 @@ where client_public_key: None, } } + + /// Construct an MTP connection with an explicit policy for pipe dispatch. + pub fn from_transport_parts_with_policy( + version: Version, + codec: VersionedCodec, + sender: S, + receiver: R, + path: String, + description: Option, + remote_addr: Option, + policy: Arc, + ) -> Self { + 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, + remote_addr, + app_rx: Mutex::new(app_rx), + pipe_req_rx: Mutex::new(pipe_req_rx), + 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"))] @@ -211,6 +255,35 @@ impl MTPConnection { client_public_key: None, } } + + pub fn from_transport_parts_with_policy( + version: Version, + codec: VersionedCodec, + sender: S, + receiver: R, + path: String, + description: Option, + remote_addr: Option, + _policy: Arc, + ) -> Self { + Self { + version, + codec, + sender, + receiver, + path, + remote_addr, + description, + _pipe_stream: std::marker::PhantomData, + _dispatcher_task: tokio::spawn(async {}), + #[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"))] diff --git a/host/src/engine.rs b/host/src/engine.rs new file mode 100644 index 0000000..ce94111 --- /dev/null +++ b/host/src/engine.rs @@ -0,0 +1,825 @@ +//! Transport-independent MTP handshake engine. +//! +//! This module contains the shared state machine used by both native `MTPHost` +//! and the web server's `MTPWebServer` to perform the MTP opening handshake, +//! version negotiation, authentication, and guest assignment. + +use crate::config::HostConfig; +use crate::error::AcceptError; +use mtp_codec::{ + CommunicationType, CommunicationValue, DataType, DataValue, Version, + registry::{Registry, VersionedCodec}, +}; +use mtp_common::{CommunicationError, RejectionReason}; +use std::sync::Arc; + +/// Trait for sending handshake messages during the opening exchange. +/// +/// Implemented by both the concrete `Sender` and `GenericSender`. +pub trait HandshakeSender: Send + Sync { + fn send( + &self, + msg: &CommunicationValue, + ) -> impl std::future::Future> + Send; + fn finish_stream( + &self, + ) -> impl std::future::Future> + Send; + fn close(&self); +} + +/// Trait for receiving handshake messages during the opening exchange. +/// +/// Implemented by both the concrete `Receiver` and `GenericReceiver`. +pub trait HandshakeReceiver: Send + Sync { + fn receive( + &self, + ) -> impl std::future::Future> + + Send; +} + +/// The result of a successful handshake, containing everything needed to +/// construct the final `MTPConnection`. +#[derive(Debug)] +pub struct HandshakeResult { + pub negotiated_version: Version, + pub codec: VersionedCodec, + pub description: Option, + #[cfg(feature = "crypto")] + pub auth_state: crate::error::AuthState, + #[cfg(feature = "crypto")] + pub client_id: u64, + #[cfg(feature = "crypto")] + pub client_public_key: Option, +} + +/// Transport-independent handshake state machine. +/// +/// Both `MTPHost` and `MTPWebServer` create a `HandshakeEngine` with the +/// shared `HostConfig` and delegate the full opening handshake to it. +pub struct HandshakeEngine { + registry: Registry, + #[cfg(feature = "crypto")] + config: Arc, +} + +impl HandshakeEngine { + #[cfg(feature = "crypto")] + pub fn new(registry: Registry, config: Arc) -> Self { + Self { registry, config } + } + + #[cfg(not(feature = "crypto"))] + pub fn new(registry: Registry, _config: Arc) -> Self { + Self { registry } + } + + /// Run the complete opening handshake with the given transport pair. + /// + /// This handles: + /// - Opening-frame timeout (when crypto is enabled) + /// - Opening-type classification (Identification, Register, or other) + /// - Version negotiation + /// - Authentication-policy selection (Unauthenticated, AllowAuthentication, ForceAuthentication) + /// - Guest allocation and collision avoidance + /// - Full challenge/response authentication when required + /// - PQ preflight checks and dual-signature verification + /// - Rejection response construction on failure + pub async fn accept( + &self, + sender: &S, + receiver: &R, + ) -> Result { + #[cfg(feature = "crypto")] + { + let timeout = self.config.auth_timeout; + tokio::time::timeout(timeout, self.accept_inner(sender, receiver)) + .await + .unwrap_or(Err(AcceptError::AuthenticationTimedOut)) + } + #[cfg(not(feature = "crypto"))] + self.accept_inner(sender, receiver).await + } + + async fn accept_inner( + &self, + sender: &S, + receiver: &R, + ) -> Result { + let first_msg = receiver.receive().await.map_err(AcceptError::Receive)?; + + let version_str = match first_msg.get_data(DataType::Version) { + DataValue::Str(s) => s.clone(), + _ => { + send_rejection_generic( + sender, + RejectionReason::AuthenticationFailed { + detail: "opening message omitted a valid protocol version".into(), + }, + ) + .await; + sender.close(); + return Err(AcceptError::MissingVersion); + } + }; + let client_version = match Version::parse(&version_str) { + Some(v) => v, + _ => { + send_rejection_generic( + sender, + RejectionReason::AuthenticationFailed { + detail: "opening message omitted a valid protocol version".into(), + }, + ) + .await; + sender.close(); + return Err(AcceptError::MissingVersion); + } + }; + + let negotiated = match self.registry.negotiate(std::slice::from_ref(&client_version)) { + Some(v) => v, + None => { + send_rejection_generic( + sender, + RejectionReason::BadVersion { + supported_versions: self + .registry + .versions() + .map(|v| v.to_string()) + .collect(), + }, + ) + .await; + sender.close(); + return Err(AcceptError::UnsupportedVersion(client_version)); + } + }; + + let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) + .ok_or_else(|| AcceptError::UnsupportedVersion(negotiated.clone()))?; + + let description = match first_msg.get_data(DataType::Description) { + DataValue::Str(s) => Some(s.clone()), + _ => None, + }; + + #[cfg(feature = "crypto")] + { + match self.config.authentication_policy { + crate::config::AuthenticationPolicy::ForceAuthentication => { + self.force_auth_handshake( + sender, + receiver, + first_msg, + negotiated, + codec, + description, + &version_str, + client_version, + ) + .await + } + crate::config::AuthenticationPolicy::AllowAuthentication => { + self.allow_auth_handshake( + sender, + receiver, + first_msg, + negotiated, + codec, + description, + &version_str, + client_version, + ) + .await + } + crate::config::AuthenticationPolicy::Unauthenticated => { + self.unauthenticated_handshake( + sender, + first_msg, + negotiated, + codec, + description, + ) + .await + } + } + } + + #[cfg(not(feature = "crypto"))] + { + let _ = sender; + let _ = receiver; + let _ = first_msg; + Ok(HandshakeResult { + negotiated_version: negotiated, + codec, + description, + }) + } + } + + #[cfg(feature = "crypto")] + async fn unauthenticated_handshake( + &self, + sender: &S, + first_msg: CommunicationValue, + negotiated: Version, + codec: VersionedCodec, + description: Option, + ) -> Result { + let tm = mtp_codec::TypeMap::latest(); + + // Reject Register frames on unauthenticated hosts + if Some(first_msg.get_type()) == CommunicationType::Register.try_to_id(&tm) { + send_rejection_generic( + sender, + RejectionReason::AuthenticationFailed { + detail: "authentication not allowed on this host".into(), + }, + ) + .await; + sender.close(); + return Err(AcceptError::AuthenticationFailed( + "authentication not allowed on this host".into(), + )); + } + + let guest_id = self.assign_guest_id().await?; + send_accepted_generic(sender, &negotiated, Some(guest_id)) + .await + .map_err(AcceptError::Send)?; + + Ok(HandshakeResult { + negotiated_version: negotiated, + codec, + description, + auth_state: crate::error::AuthState::Unauthenticated, + client_id: guest_id, + client_public_key: None, + }) + } + + #[cfg(feature = "crypto")] + #[allow(clippy::too_many_arguments)] + async fn allow_auth_handshake( + &self, + sender: &S, + receiver: &R, + first_msg: CommunicationValue, + negotiated: Version, + codec: VersionedCodec, + description: Option, + version_str: &str, + client_version: Version, + ) -> Result { + let tm = mtp_codec::TypeMap::latest(); + + // Register frames always go through full authentication + if Some(first_msg.get_type()) == CommunicationType::Register.try_to_id(&tm) { + let bundle = extract_register_bundle(&first_msg)?; + let pk_bytes = bundle.as_bytes(); + return self + .complete_auth_handshake( + sender, + receiver, + Flow::Register { bundle, pk_bytes }, + CommunicationType::RegisterResponse, + &negotiated, + &codec, + description, + version_str, + client_version, + ) + .await; + } + + // Identification: try lookup, fall back to guest + if Some(first_msg.get_type()) == CommunicationType::Identification.try_to_id(&tm) { + let cid = match first_msg.get_data(DataType::Id) { + DataValue::UnsignedNumber(n) => *n as u64, + _ => 0, + }; + + if cid > 0 { + if let Some(bundle) = + (self.config.get_existing_client)(cid, description.clone()).await + { + return self + .complete_auth_handshake( + sender, + receiver, + Flow::Login { id: cid, bundle }, + CommunicationType::IdentificationResponse, + &negotiated, + &codec, + description, + version_str, + client_version, + ) + .await; + } + } + + // Unknown or zero ID: fall back to guest + let guest_id = self.assign_guest_id().await?; + send_accepted_generic(sender, &negotiated, Some(guest_id)) + .await + .map_err(AcceptError::Send)?; + return Ok(HandshakeResult { + negotiated_version: negotiated, + codec, + description, + auth_state: crate::error::AuthState::Unauthenticated, + client_id: guest_id, + client_public_key: None, + }); + } + + sender.close(); + Err(AcceptError::AuthenticationFailed( + "unexpected message type".into(), + )) + } + + #[cfg(feature = "crypto")] + #[allow(clippy::too_many_arguments)] + async fn force_auth_handshake( + &self, + sender: &S, + receiver: &R, + first_msg: CommunicationValue, + negotiated: Version, + codec: VersionedCodec, + description: Option, + version_str: &str, + client_version: Version, + ) -> Result { + let tm = mtp_codec::TypeMap::latest(); + + let (flow, response_type) = if Some(first_msg.get_type()) + == CommunicationType::Identification.try_to_id(&tm) + { + let cid = match first_msg.get_data(DataType::Id) { + DataValue::UnsignedNumber(n) => *n as u64, + _ => { + sender.close(); + return Err(AcceptError::AuthenticationFailed( + "missing client id".into(), + )); + } + }; + let bundle = match (self.config.get_existing_client)(cid, description.clone()).await { + Some(b) => b, + None => { + let rejection = + CommunicationValue::new(CommunicationType::IdentificationResponse) + .add_typed_default(DataType::Connected, DataValue::BoolFalse) + .add_typed_default( + DataType::ErrorMessage, + DataValue::Str("unknown client id".into()), + ); + let _ = sender.send(&rejection).await; + sender.close(); + return Err(AcceptError::AuthenticationFailed( + "unknown client id".into(), + )); + } + }; + ( + Flow::Login { id: cid, bundle }, + CommunicationType::IdentificationResponse, + ) + } else if Some(first_msg.get_type()) == CommunicationType::Register.try_to_id(&tm) { + let bundle = extract_register_bundle(&first_msg)?; + let pk_bytes = bundle.as_bytes(); + ( + Flow::Register { bundle, pk_bytes }, + CommunicationType::RegisterResponse, + ) + } else { + sender.close(); + return Err(AcceptError::AuthenticationFailed( + "unexpected authentication message".into(), + )); + }; + + self.complete_auth_handshake( + sender, + receiver, + flow, + response_type, + &negotiated, + &codec, + description, + version_str, + client_version, + ) + .await + } + + #[cfg(feature = "crypto")] + #[allow(clippy::too_many_arguments)] + async fn complete_auth_handshake( + &self, + sender: &S, + receiver: &R, + flow: Flow, + response_type: CommunicationType, + negotiated: &Version, + codec: &VersionedCodec, + description: Option, + version_str: &str, + _client_version: Version, + ) -> Result { + use mtp_crypto::{Ed25519Signer, MlDsaSigner, SignatureScheme, auth, verify_ed25519}; + + let tm = mtp_codec::TypeMap::latest(); + + // PQ preflight: host requiring PQ must have a PQ key + let pq_enabled = !self + .config + .host_keyring + .sig_pq_secret_key + .as_bytes() + .is_empty(); + if self.config.require_pq + && (!pq_enabled + || self + .config + .host_keyring + .sig_pq_public_key + .as_bytes() + .is_empty()) + { + send_rejection_generic( + sender, + RejectionReason::AuthenticationFailed { + detail: "host requires PQ authentication but has no PQ signing key".into(), + }, + ) + .await; + sender.close(); + return Err(AcceptError::AuthenticationFailed( + "PQ authentication is required but the host PQ key is absent".into(), + )); + } + + // Initialize host signers + 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 + }; + + 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()))?; + 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 { + let sig = signer + .sign(&payload) + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?; + Ok((sig, Vec::new())) + } + }; + + // Sign and send challenge + let challenge_id = match &flow { + Flow::Login { id, .. } => *id, + Flow::Register { .. } => 0, + }; + + let server_challenge: u128 = rand::random(); + let (chal_sig, chal_pq_sig) = + host_sign(auth::challenge_payload(challenge_id, server_challenge)).await?; + + let mut challenge_msg = CommunicationValue::new(CommunicationType::Challenge) + .add_typed_default( + DataType::ServerNonce, + DataValue::UnsignedNumber(server_challenge), + ) + .add_typed_default(DataType::Signature, DataValue::Bytes(chal_sig)); + challenge_msg = challenge_msg.add_typed_default( + DataType::RequirePq, + if self.config.require_pq { + DataValue::BoolTrue + } else { + DataValue::BoolFalse + }, + ); + if pq_enabled { + challenge_msg = challenge_msg + .add_typed_default(DataType::PqSignature, DataValue::Bytes(chal_pq_sig)); + } + if let Err(e) = sender.send(&challenge_msg).await { + sender.close(); + return Err(AcceptError::Send(e)); + } + + // Receive and verify client proof + let proof = receiver.receive().await.map_err(|e| { + sender.close(); + AcceptError::Receive(e) + })?; + if Some(proof.get_type()) != CommunicationType::ChallengeResponse.try_to_id(&tm) { + sender.close(); + return Err(AcceptError::AuthenticationFailed( + "missing challenge response".into(), + )); + } + let client_nonce = match proof.get_data(DataType::ClientNonce) { + DataValue::UnsignedNumber(n) => *n, + _ => { + sender.close(); + return Err(AcceptError::AuthenticationFailed( + "missing client nonce".into(), + )); + } + }; + let sig_bytes = match proof.get_data(DataType::Signature) { + DataValue::Bytes(b) => b.clone(), + _ => { + sender.close(); + return Err(AcceptError::AuthenticationFailed( + "missing challenge signature".into(), + )); + } + }; + let pq_sig_bytes: Vec = match proof.get_data(DataType::PqSignature) { + DataValue::Bytes(b) => b.clone(), + _ => vec![], + }; + + let (proof_payload, bundle) = match &flow { + Flow::Login { id, bundle } => ( + auth::login_proof_payload(version_str, *id, server_challenge, client_nonce), + bundle, + ), + Flow::Register { + bundle, pk_bytes, .. + } => ( + auth::register_proof_payload(version_str, pk_bytes, server_challenge, client_nonce), + bundle, + ), + }; + + let has_client_pq_key = !bundle.sig_pq_public_key.as_bytes().is_empty(); + 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 { + false + }; + + if !proof_ok { + send_rejection_generic( + sender, + RejectionReason::AuthenticationFailed { + detail: "client proof signature invalid".into(), + }, + ) + .await; + sender.close(); + return Err(AcceptError::AuthenticationFailed( + "client proof signature invalid".into(), + )); + } + + // Register or login + let (assigned_id, client_bundle) = match flow { + Flow::Login { id, bundle } => (id, bundle), + Flow::Register { bundle, .. } => { + let new_id = + (self.config.complete_register)(bundle.clone(), description.clone()).await; + (new_id, bundle) + } + }; + + // Sign and send final response + let (host_sig, host_pq_sig) = host_sign(auth::host_final_payload( + assigned_id, + client_nonce, + server_challenge, + )) + .await?; + + let mut response = 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(client_nonce), + ) + .add_typed_default(DataType::Signature, DataValue::Bytes(host_sig)); + response = response + .add_typed_default(DataType::Version, DataValue::Str(negotiated.to_string())); + if pq_enabled { + response = + response.add_typed_default(DataType::PqSignature, DataValue::Bytes(host_pq_sig)); + } + + if let Err(e) = sender.send(&response).await { + sender.close(); + return Err(AcceptError::Send(e)); + } + if let Err(e) = sender.finish_stream().await { + sender.close(); + return Err(AcceptError::Send(e)); + } + + Ok(HandshakeResult { + negotiated_version: negotiated.clone(), + codec: codec.clone(), + description, + auth_state: crate::error::AuthState::Authenticated, + client_id: assigned_id, + client_public_key: Some(client_bundle), + }) + } +} + +#[cfg(feature = "crypto")] +enum Flow { + Login { + id: u64, + bundle: mtp_crypto::PublicKeyBundle, + }, + Register { + bundle: mtp_crypto::PublicKeyBundle, + pk_bytes: Vec, + }, +} + +// --------------------------------------------------------------------------- +// Guest ID allocation +// --------------------------------------------------------------------------- + +#[cfg(feature = "crypto")] +impl HandshakeEngine { + const GUEST_ID_MAX_RETRIES: u32 = 100; + + async fn assign_guest_id(&self) -> Result { + if let Some(ref generator) = self.config.guest_id_generator { + let id = generator().await.ok_or_else(|| { + AcceptError::AuthenticationFailed( + "guest id generator rejected the connection".into(), + ) + })?; + if id > mtp_codec::MAX_WIRE_ID { + return Err(AcceptError::AuthenticationFailed( + "guest id exceeds wire limit".into(), + )); + } + if (self.config.get_existing_client)(id, None).await.is_none() { + return Ok(id); + } + } + self.random_guest_id().await + } + + async fn random_guest_id(&self) -> Result { + for _ in 0..Self::GUEST_ID_MAX_RETRIES { + let id = rand::random::() & mtp_codec::MAX_WIRE_ID; + if (self.config.get_existing_client)(id, None).await.is_none() { + return Ok(id); + } + } + Err(AcceptError::AuthenticationFailed( + "failed to allocate a unique guest id after retries".into(), + )) + } +} + +// --------------------------------------------------------------------------- +// Helpers +// --------------------------------------------------------------------------- + +#[cfg(feature = "crypto")] +fn extract_register_bundle( + msg: &CommunicationValue, +) -> Result { + match msg.get_data(DataType::PublicKeys) { + DataValue::Bytes(b) => mtp_crypto::PublicKeyBundle::from_bytes(b).map_err(|_| { + AcceptError::AuthenticationFailed("invalid public key bundle".into()) + }), + _ => Err(AcceptError::AuthenticationFailed( + "missing public keys".into(), + )), + } +} + +async fn send_rejection_generic( + sender: &S, + reason: RejectionReason, +) { + let response = match &reason { + RejectionReason::BadVersion { supported_versions } => { + CommunicationValue::new(CommunicationType::ErrorBadVersion) + .add_typed_default( + DataType::Version, + DataValue::Str(supported_versions.join(",")), + ) + .add_typed_default(DataType::ErrorMessage, DataValue::Str(reason.to_string())) + } + _ => CommunicationValue::new(CommunicationType::IdentificationResponse) + .add_typed_default(DataType::Connected, DataValue::BoolFalse) + .add_typed_default(DataType::ErrorMessage, DataValue::Str(reason.to_string())), + }; + let _ = sender.send(&response).await; +} + +#[cfg(feature = "crypto")] +async fn send_accepted_generic( + sender: &S, + version: &Version, + assigned_id: Option, +) -> Result<(), CommunicationError> { + let mut response = CommunicationValue::new(CommunicationType::IdentificationResponse) + .add_typed_default(DataType::Connected, DataValue::BoolTrue) + .add_typed_default(DataType::Version, DataValue::Str(version.to_string())); + if let Some(id) = assigned_id { + response = response.add_typed_default(DataType::Id, DataValue::UnsignedNumber(id as u128)); + } + sender.send(&response).await?; + sender.finish_stream().await +} + +// --------------------------------------------------------------------------- +// Trait implementations for concrete transport types +// --------------------------------------------------------------------------- + +impl HandshakeSender for mtp_transport::Sender { + fn send( + &self, + msg: &CommunicationValue, + ) -> impl std::future::Future> + Send { + mtp_transport::Sender::send(self, msg) + } + fn finish_stream( + &self, + ) -> impl std::future::Future> + Send { + mtp_transport::Sender::finish_stream(self) + } + fn close(&self) { + let sender = self.clone(); + tokio::spawn(async move { sender.close().await }); + } +} + +impl HandshakeReceiver for mtp_transport::Receiver { + fn receive( + &self, + ) -> impl std::future::Future> + Send + { + mtp_transport::Receiver::receive(self) + } +} + +impl HandshakeSender for mtp_transport::GenericSender { + fn send( + &self, + msg: &CommunicationValue, + ) -> impl std::future::Future> + Send { + mtp_transport::GenericSender::send(self, msg) + } + fn finish_stream( + &self, + ) -> impl std::future::Future> + Send { + mtp_transport::GenericSender::finish_stream(self) + } + fn close(&self) { + mtp_transport::GenericSender::close(self); + } +} + +impl HandshakeReceiver for mtp_transport::GenericReceiver { + fn receive( + &self, + ) -> impl std::future::Future> + Send + { + mtp_transport::GenericReceiver::receive(self) + } +} diff --git a/host/src/error.rs b/host/src/error.rs index 6cba86b..7a7873e 100644 --- a/host/src/error.rs +++ b/host/src/error.rs @@ -1,45 +1,16 @@ -use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue, Version}; -use mtp_common::{CommunicationError, RejectionReason}; -use mtp_transport::Sender; +use mtp_codec::Version; +use mtp_common::CommunicationError; use std::{error::Error, fmt}; +#[cfg(test)] +use mtp_codec::{CommunicationValue, DataType, DataValue}; + #[cfg(feature = "crypto")] pub(crate) fn random_client_id() -> u64 { rand::random::() & mtp_codec::MAX_WIRE_ID } -pub(crate) async fn send_rejection(sender: &Sender, reason: RejectionReason) { - let response = match &reason { - RejectionReason::BadVersion { supported_versions } => { - CommunicationValue::new(CommunicationType::ErrorBadVersion) - .add_typed_default( - DataType::Version, - DataValue::Str(supported_versions.join(",")), - ) - .add_typed_default(DataType::ErrorMessage, DataValue::Str(reason.to_string())) - } - _ => CommunicationValue::new(CommunicationType::IdentificationResponse) - .add_typed_default(DataType::Connected, DataValue::BoolFalse) - .add_typed_default(DataType::ErrorMessage, DataValue::Str(reason.to_string())), - }; - let _ = sender.send(&response).await; -} - -pub(crate) async fn send_accepted( - sender: &Sender, - version: &Version, - assigned_id: Option, -) -> Result<(), CommunicationError> { - let mut response = CommunicationValue::new(CommunicationType::IdentificationResponse) - .add_typed_default(DataType::Connected, DataValue::BoolTrue) - .add_typed_default(DataType::Version, DataValue::Str(version.to_string())); - if let Some(id) = assigned_id { - response = response.add_typed_default(DataType::Id, DataValue::UnsignedNumber(id as u128)); - } - sender.send(&response).await?; - sender.finish_stream().await -} - +#[cfg(test)] pub(crate) fn extract_version(msg: &CommunicationValue) -> Option { let value = msg.get_data(DataType::Version); match value { diff --git a/host/src/handshake.rs b/host/src/handshake.rs index f188f54..328de89 100644 --- a/host/src/handshake.rs +++ b/host/src/handshake.rs @@ -1,23 +1,17 @@ +#[cfg(not(feature = "crypto"))] +use mtp_codec::{Version, registry::{Registry, VersionedCodec}}; #[cfg(feature = "crypto")] -use mtp_codec::{CommunicationType, CommunicationValue}; -use mtp_codec::{ - DataType, DataValue, Version, - registry::{Registry, VersionedCodec}, -}; -use mtp_common::RejectionReason; +use mtp_codec::registry::Registry; use mtp_transport::{Receiver, Sender}; use std::sync::Arc; use std::time::Instant; #[cfg(feature = "pipes")] use tokio::sync::mpsc; -#[cfg(feature = "crypto")] -use crate::config::AuthenticationPolicy; use crate::config::HostConfig; use crate::connection::MTPConnection; -#[cfg(feature = "crypto")] -use crate::error::AuthState; -use crate::error::{AcceptError, extract_version, send_accepted, send_rejection}; +use crate::engine::HandshakeEngine; +use crate::error::AcceptError; #[cfg(feature = "pipes")] use crate::pipe::PipeDispatcher; #[cfg(feature = "pipes")] @@ -137,182 +131,107 @@ impl HandshakeContext { sender: Sender, receiver: Receiver, ) -> Result, AcceptError> { + let engine = HandshakeEngine::new(self.registry.clone(), self.config.clone()); + let result = engine.accept(&sender, &receiver).await?; #[cfg(feature = "crypto")] { - return tokio::time::timeout( - self.config.auth_timeout, - self.accept_pair(sender, receiver), - ) - .await - .unwrap_or(Err(AcceptError::AuthenticationTimedOut)); + Ok(Some(self.connection_from_handshake_result( + sender, + receiver, + result, + ))) } - #[cfg(not(feature = "crypto"))] - self.accept_pair(sender, receiver).await - } - - async fn accept_pair( - &self, - sender: Sender, - receiver: Receiver, - ) -> Result, AcceptError> { - #[cfg(feature = "crypto")] - match self.config.authentication_policy { - AuthenticationPolicy::ForceAuthentication => { - let timeout = self.config.auth_timeout; - match tokio::time::timeout(timeout, self.accept_authenticated(sender, receiver)) - .await - { - Ok(result) => result, - Err(_) => Err(AcceptError::AuthenticationTimedOut), - } - } - AuthenticationPolicy::AllowAuthentication => { - return self.accept_allow_auth(sender, receiver).await; - } - AuthenticationPolicy::Unauthenticated => { - let first_msg = match receiver.receive().await { - Ok(m) => m, - Err(e) => return Err(AcceptError::Receive(e)), - }; - if Some(first_msg.get_type()) - == CommunicationType::Register.try_to_id(&mtp_codec::TypeMap::latest()) - { - send_rejection( - &sender, - RejectionReason::AuthenticationFailed { - detail: "authentication not allowed on this host".into(), - }, - ) - .await; - sender.close().await; - return Err(AcceptError::AuthenticationFailed( - "authentication not allowed on this host".into(), - )); - } - let client_version = match extract_version(&first_msg) { - Some(v) => v, - None => { - send_rejection( - &sender, - RejectionReason::AuthenticationFailed { - detail: "opening message omitted a valid protocol version".into(), - }, - ) - .await; - sender.close().await; - return Err(AcceptError::MissingVersion); - } - }; - let negotiated = match self - .registry - .negotiate(std::slice::from_ref(&client_version)) - { - Some(v) => v, - None => { - send_rejection( - &sender, - RejectionReason::BadVersion { - supported_versions: self - .registry - .versions() - .map(|v| v.to_string()) - .collect(), - }, - ) - .await; - sender.close().await; - return Err(AcceptError::UnsupportedVersion(client_version)); - } - }; - let codec = - match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) { - Some(codec) => codec, - None => { - return Err(AcceptError::UnsupportedVersion(negotiated)); - } - }; - let description = match first_msg.get_data(DataType::Description) { - DataValue::Str(s) => Some(s.clone()), - _ => None, - }; - let guest_id = self.assign_guest_id().await; - send_accepted(&sender, &negotiated, Some(guest_id)) - .await - .map_err(AcceptError::Send)?; - Ok(Some(self.connection_from_parts( - sender, - receiver, - negotiated, - codec, - description, - AuthState::Unauthenticated, - guest_id, - None, - ))) - } - } - #[cfg(not(feature = "crypto"))] { - let first_msg = match receiver.receive().await { - Ok(m) => m, - Err(e) => return Err(AcceptError::Receive(e)), - }; - let client_version = match extract_version(&first_msg) { - Some(v) => v, - None => { - send_rejection( - &sender, - RejectionReason::AuthenticationFailed { - detail: "opening message omitted a valid protocol version".into(), - }, - ) - .await; - sender.close().await; - return Err(AcceptError::MissingVersion); - } - }; - let negotiated = match self - .registry - .negotiate(std::slice::from_ref(&client_version)) - { - Some(v) => v, - None => { - send_rejection( - &sender, - RejectionReason::BadVersion { - supported_versions: vec![mtp_codec::PROTOCOL_VERSION.to_string()], - }, - ) - .await; - sender.close().await; - return Err(AcceptError::UnsupportedVersion(client_version)); - } - }; - let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) - { - Some(codec) => codec, - None => { - return Err(AcceptError::UnsupportedVersion(negotiated)); - } - }; - let description = match first_msg.get_data(DataType::Description) { - DataValue::Str(s) => Some(s.clone()), - _ => None, - }; - send_accepted(&sender, &negotiated, None) - .await - .map_err(AcceptError::Send)?; Ok(Some(self.connection_from_parts( sender, receiver, - negotiated, - codec, - description, + result.negotiated_version, + result.codec, + result.description, ))) } } + #[cfg(feature = "crypto")] + pub(crate) fn connection_from_handshake_result( + &self, + sender: Sender, + receiver: Receiver, + result: crate::engine::HandshakeResult, + ) -> MTPConnection { + let remote_addr = sender.handle().remote_addr(); + receiver.set_max_message_size(self.config.policy.max_message_size); + #[cfg(feature = "pipes")] + { + if self.config.send_pongs { + receiver.respond_to_pings(sender.clone()); + } + + let (app_tx, app_rx) = mpsc::channel(self.config.policy.receiver_queue_capacity); + let (pipe_req_tx, pipe_req_rx) = + mpsc::channel(self.config.policy.receiver_queue_capacity); + + let dispatcher = Arc::new(PipeDispatcher { + pending_creations: tokio::sync::Mutex::new(std::collections::HashMap::new()), + pending_pipes: tokio::sync::Mutex::new(std::collections::HashMap::new()), + policy: Arc::new(self.config.policy), + }); + + let dispatcher_clone = dispatcher.clone(); + let receiver_clone = receiver.clone(); + let sender_clone = sender.clone(); + let task = tokio::spawn(run_dispatcher( + receiver_clone, + sender_clone, + app_tx, + pipe_req_tx, + dispatcher_clone, + )); + + MTPConnection { + version: result.negotiated_version, + codec: result.codec, + sender, + receiver, + path: "/".to_string(), + remote_addr, + app_rx: tokio::sync::Mutex::new(app_rx), + pipe_req_rx: tokio::sync::Mutex::new(pipe_req_rx), + pipe_dispatcher: dispatcher, + description: result.description, + _dispatcher_task: task, + auth_state: result.auth_state, + client_id: result.client_id, + client_public_key: result.client_public_key, + } + } + + #[cfg(not(feature = "pipes"))] + { + if self.config.send_pongs { + receiver.respond_to_pings(sender.clone()); + } + + let task = tokio::spawn(async {}); + + MTPConnection { + version: result.negotiated_version, + codec: result.codec, + sender, + receiver, + path: "/".to_string(), + remote_addr, + _pipe_stream: std::marker::PhantomData, + description: result.description, + _dispatcher_task: task, + auth_state: result.auth_state, + client_id: result.client_id, + client_public_key: result.client_public_key, + } + } + } + #[cfg(not(feature = "crypto"))] #[allow(clippy::too_many_arguments)] pub(crate) fn connection_from_parts( @@ -388,645 +307,4 @@ impl HandshakeContext { } } } - - #[cfg(feature = "crypto")] - #[allow(clippy::too_many_arguments)] - pub(crate) fn connection_from_parts( - &self, - sender: Sender, - receiver: Receiver, - version: Version, - codec: VersionedCodec, - description: Option, - auth_state: AuthState, - client_id: u64, - client_public_key: Option, - ) -> MTPConnection { - let remote_addr = sender.handle().remote_addr(); - receiver.set_max_message_size(self.config.policy.max_message_size); - #[cfg(feature = "pipes")] - { - if self.config.send_pongs { - receiver.respond_to_pings(sender.clone()); - } - - let (app_tx, app_rx) = mpsc::channel(self.config.policy.receiver_queue_capacity); - let (pipe_req_tx, pipe_req_rx) = - mpsc::channel(self.config.policy.receiver_queue_capacity); - - let dispatcher = Arc::new(PipeDispatcher { - pending_creations: tokio::sync::Mutex::new(std::collections::HashMap::new()), - pending_pipes: tokio::sync::Mutex::new(std::collections::HashMap::new()), - policy: Arc::new(self.config.policy), - }); - - let dispatcher_clone = dispatcher.clone(); - let receiver_clone = receiver.clone(); - let sender_clone = sender.clone(); - let task = tokio::spawn(run_dispatcher( - receiver_clone, - sender_clone, - app_tx, - pipe_req_tx, - dispatcher_clone, - )); - - MTPConnection { - version, - codec, - sender, - receiver, - path: "/".to_string(), - remote_addr, - app_rx: tokio::sync::Mutex::new(app_rx), - pipe_req_rx: tokio::sync::Mutex::new(pipe_req_rx), - pipe_dispatcher: dispatcher, - description, - _dispatcher_task: task, - auth_state, - client_id, - client_public_key, - } - } - - #[cfg(not(feature = "pipes"))] - { - if self.config.send_pongs { - receiver.respond_to_pings(sender.clone()); - } - - let task = tokio::spawn(async {}); - - MTPConnection { - version, - codec, - sender, - receiver, - path: "/".to_string(), - remote_addr, - _pipe_stream: std::marker::PhantomData, - description, - _dispatcher_task: task, - auth_state, - client_id, - client_public_key, - } - } - } -} - -#[cfg(feature = "crypto")] -enum Flow { - Login { - id: u64, - bundle: mtp_crypto::PublicKeyBundle, - }, - Register { - bundle: mtp_crypto::PublicKeyBundle, - pk_bytes: Vec, - }, -} - -#[cfg(feature = "crypto")] -impl HandshakeContext { - const GUEST_ID_MAX_RETRIES: u32 = 100; - - async fn assign_guest_id(&self) -> u64 { - if let Some(ref generator) = self.config.guest_id_generator { - if let Some(id) = generator().await - && id <= mtp_codec::MAX_WIRE_ID - { - return id; - } - return self.random_guest_id().await; - } - self.random_guest_id().await - } - - async fn random_guest_id(&self) -> u64 { - for _ in 0..Self::GUEST_ID_MAX_RETRIES { - let id = rand::random::() & mtp_codec::MAX_WIRE_ID; - if (self.config.get_existing_client)(id, None).await.is_none() { - return id; - } - } - rand::random::() & mtp_codec::MAX_WIRE_ID - } - - async fn accept_authenticated( - &self, - sender: Sender, - receiver: Receiver, - ) -> Result, AcceptError> { - use mtp_crypto::PublicKeyBundle; - - let tm = mtp_codec::TypeMap::latest(); - - let hello = match receiver.receive().await { - Ok(m) => m, - Err(e) => { - sender.close().await; - return Err(AcceptError::Receive(e)); - } - }; - let version_str = match hello.get_data(DataType::Version) { - DataValue::Str(s) => s.clone(), - _ => { - sender.close().await; - return Err(AcceptError::MissingVersion); - } - }; - let client_version = match Version::parse(&version_str) { - Some(v) => v, - None => { - sender.close().await; - return Err(AcceptError::MissingVersion); - } - }; - - let description = match hello.get_data(DataType::Description) { - DataValue::Str(s) => Some(s.clone()), - _ => None, - }; - - let (flow, response_type) = if Some(hello.get_type()) - == CommunicationType::Identification.try_to_id(&tm) - { - let cid = match hello.get_data(DataType::Id) { - DataValue::UnsignedNumber(n) => *n as u64, - _ => { - sender.close().await; - return Err(AcceptError::AuthenticationFailed( - "missing client id".into(), - )); - } - }; - let bundle = match (self.config.get_existing_client)(cid, description.clone()).await { - Some(b) => b, - None => { - let rejection = - CommunicationValue::new(CommunicationType::IdentificationResponse) - .add_typed_default(DataType::Connected, DataValue::BoolFalse) - .add_typed_default( - DataType::ErrorMessage, - DataValue::Str("unknown client id".into()), - ); - let _ = sender.send(&rejection).await; - sender.close().await; - return Err(AcceptError::AuthenticationFailed( - "unknown client id".into(), - )); - } - }; - ( - Flow::Login { id: cid, bundle }, - CommunicationType::IdentificationResponse, - ) - } else 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(|_| { - AcceptError::AuthenticationFailed("invalid public key bundle".into()) - })?, - _ => { - sender.close().await; - return Err(AcceptError::AuthenticationFailed( - "missing public keys".into(), - )); - } - }; - let pk_bytes = bundle.as_bytes(); - ( - Flow::Register { bundle, pk_bytes }, - CommunicationType::RegisterResponse, - ) - } else { - sender.close().await; - return Err(AcceptError::AuthenticationFailed( - "unexpected authentication message".into(), - )); - }; - - self.complete_auth_handshake( - sender, - receiver, - flow, - response_type, - &version_str, - client_version, - description, - ) - .await - } - - #[allow(clippy::too_many_arguments)] - async fn complete_auth_handshake( - &self, - sender: Sender, - receiver: Receiver, - flow: Flow, - response_type: CommunicationType, - version_str: &str, - client_version: Version, - description: Option, - ) -> Result, AcceptError> { - 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)) - { - Some(version) => version, - None => { - send_rejection( - &sender, - RejectionReason::BadVersion { - supported_versions: vec![mtp_codec::PROTOCOL_VERSION.to_string()], - }, - ) - .await; - 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 - .sig_pq_secret_key - .as_bytes() - .is_empty(); - if self.config.require_pq - && (!pq_enabled - || self - .config - .host_keyring - .sig_pq_public_key - .as_bytes() - .is_empty()) - { - send_rejection( - &sender, - RejectionReason::AuthenticationFailed { - detail: "host requires PQ authentication but has no PQ signing key".into(), - }, - ) - .await; - sender.close().await; - return Err(AcceptError::AuthenticationFailed( - "PQ authentication is required but the host PQ key is absent".into(), - )); - } - - 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()))?; - 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 { - let sig = signer - .sign(&payload) - .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?; - Ok((sig, Vec::new())) - } - }; - - let challenge_id = match &flow { - Flow::Login { id, .. } => *id, - Flow::Register { .. } => 0, - }; - - 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)).await?; - tracing::debug!(elapsed = ?sign_challenge_started.elapsed(), "authentication handshake: sign challenge"); - - let mut challenge_msg = CommunicationValue::new(CommunicationType::Challenge) - .add_typed_default( - DataType::ServerNonce, - DataValue::UnsignedNumber(server_challenge), - ) - .add_typed_default(DataType::Signature, DataValue::Bytes(chal_sig)); - challenge_msg = challenge_msg.add_typed_default( - DataType::RequirePq, - if self.config.require_pq { - DataValue::BoolTrue - } else { - DataValue::BoolFalse - }, - ); - if pq_enabled { - 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().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().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().await; - return Err(AcceptError::AuthenticationFailed( - "missing challenge response".into(), - )); - } - let client_nonce = match proof.get_data(DataType::ClientNonce) { - DataValue::UnsignedNumber(n) => *n, - _ => { - sender.close().await; - return Err(AcceptError::AuthenticationFailed( - "missing client nonce".into(), - )); - } - }; - let sig_bytes = match proof.get_data(DataType::Signature) { - DataValue::Bytes(b) => b.clone(), - _ => { - sender.close().await; - return Err(AcceptError::AuthenticationFailed( - "missing challenge signature".into(), - )); - } - }; - let pq_sig_bytes: Vec = match proof.get_data(DataType::PqSignature) { - DataValue::Bytes(b) => b.clone(), - _ => vec![], - }; - - let (proof_payload, bundle) = match &flow { - Flow::Login { id, bundle } => ( - auth::login_proof_payload(version_str, *id, server_challenge, client_nonce), - bundle, - ), - Flow::Register { - bundle, pk_bytes, .. - } => ( - auth::register_proof_payload(version_str, pk_bytes, server_challenge, client_nonce), - bundle, - ), - }; - - let has_client_pq_key = !bundle.sig_pq_public_key.as_bytes().is_empty(); - 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 { - false - }; - tracing::debug!(elapsed = ?verify_proof_started.elapsed(), "authentication handshake: verify client proof"); - - if !proof_ok { - send_rejection( - &sender, - RejectionReason::AuthenticationFailed { - detail: "client proof signature invalid".into(), - }, - ) - .await; - 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, .. } => { - let new_id = - (self.config.complete_register)(bundle.clone(), description.clone()).await; - (new_id, bundle) - } - }; - tracing::debug!(elapsed = ?register_started.elapsed(), "authentication handshake: registration callback"); - - 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) - .add_typed_default(DataType::Id, DataValue::UnsignedNumber(assigned_id as u128)) - .add_typed_default( - DataType::ClientNonce, - DataValue::UnsignedNumber(client_nonce), - ) - .add_typed_default(DataType::Signature, DataValue::Bytes(host_sig)); - response = - response.add_typed_default(DataType::Version, DataValue::Str(negotiated.to_string())); - if pq_enabled { - response = - 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().await; - return Err(AcceptError::Send(e)); - } - if let Err(e) = sender.finish_stream().await { - 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, - None => { - return Err(AcceptError::UnsupportedVersion(negotiated)); - } - }; - - Ok(Some(self.connection_from_parts( - sender, - receiver, - negotiated, - codec, - description, - AuthState::Authenticated, - assigned_id, - Some(client_bundle), - ))) - } - - async fn accept_allow_auth( - &self, - sender: Sender, - receiver: Receiver, - ) -> Result, AcceptError> { - use mtp_crypto::PublicKeyBundle; - - let tm = mtp_codec::TypeMap::latest(); - - let hello = match receiver.receive().await { - Ok(m) => m, - Err(e) => return Err(AcceptError::Receive(e)), - }; - - let version_str = match hello.get_data(DataType::Version) { - DataValue::Str(s) => s.clone(), - _ => { - sender.close().await; - return Err(AcceptError::MissingVersion); - } - }; - let client_version = match Version::parse(&version_str) { - Some(v) => v, - None => { - sender.close().await; - return Err(AcceptError::MissingVersion); - } - }; - - let description = match hello.get_data(DataType::Description) { - DataValue::Str(s) => Some(s.clone()), - _ => None, - }; - - 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(|_| { - AcceptError::AuthenticationFailed("invalid public key bundle".into()) - })?, - _ => { - 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( - sender, - receiver, - Flow::Register { bundle, pk_bytes }, - CommunicationType::RegisterResponse, - &version_str, - client_version, - description, - ) - .await; - } - - if Some(hello.get_type()) == CommunicationType::Identification.try_to_id(&tm) { - let cid = match hello.get_data(DataType::Id) { - DataValue::UnsignedNumber(n) => *n as u64, - _ => 0, - }; - - if cid > 0 - && let Some(bundle) = - (self.config.get_existing_client)(cid, description.clone()).await - { - return self - .complete_auth_handshake( - sender, - receiver, - Flow::Login { id: cid, bundle }, - CommunicationType::IdentificationResponse, - &version_str, - client_version, - description, - ) - .await; - } - - let negotiated = match self - .registry - .negotiate(std::slice::from_ref(&client_version)) - { - Some(v) => v, - None => { - send_rejection( - &sender, - RejectionReason::BadVersion { - supported_versions: vec![mtp_codec::PROTOCOL_VERSION.to_string()], - }, - ) - .await; - sender.close().await; - return Err(AcceptError::UnsupportedVersion(client_version)); - } - }; - let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) - { - Some(codec) => codec, - None => { - return Err(AcceptError::UnsupportedVersion(negotiated)); - } - }; - let guest_id = self.assign_guest_id().await; - send_accepted(&sender, &negotiated, Some(guest_id)) - .await - .map_err(AcceptError::Send)?; - return Ok(Some(self.connection_from_parts( - sender, - receiver, - negotiated, - codec, - description, - AuthState::Unauthenticated, - guest_id, - None, - ))); - } - - sender.close().await; - Err(AcceptError::AuthenticationFailed( - "unexpected message type".into(), - )) - } } diff --git a/host/src/lib.rs b/host/src/lib.rs index df64b68..211b3a6 100644 --- a/host/src/lib.rs +++ b/host/src/lib.rs @@ -1,5 +1,6 @@ pub mod config; pub mod connection; +pub mod engine; pub mod error; pub mod handshake; #[cfg(feature = "pipes")] @@ -10,6 +11,7 @@ pub use MTPHost as Host; pub use config::HostConfig; pub use config::Policy; pub use connection::{MTPConnection, MtpReceiverLike, MtpSenderLike}; +pub use engine::{HandshakeEngine, HandshakeResult, HandshakeReceiver, HandshakeSender}; pub use error::AcceptError; pub use handshake::MTPHost; pub use mtp_transport::Receiver; diff --git a/mtp-webserver/src/h3.rs b/mtp-webserver/src/h3.rs index 6fce209..eb3539c 100644 --- a/mtp-webserver/src/h3.rs +++ b/mtp-webserver/src/h3.rs @@ -18,6 +18,7 @@ pub(crate) struct DriverConfig { pub(crate) policy: mtp_transport::Policy, pub(crate) host_config: Arc, pub(crate) metrics: Option>, + pub(crate) auth_semaphore: Arc, } pub(crate) async fn run_driver( @@ -37,6 +38,7 @@ pub(crate) async fn run_driver( policy, host_config, metrics, + auth_semaphore, } = config; let mut connection_tasks = tokio::task::JoinSet::new(); loop { @@ -63,6 +65,7 @@ pub(crate) async fn run_driver( let mtp_tx = mtp_tx.clone(); let metrics = metrics.clone(); let host_config = host_config.clone(); + let auth_semaphore = auth_semaphore.clone(); connection_tasks.spawn(async move { let _permit = permit; let connect_start = std::time::Instant::now(); @@ -137,10 +140,6 @@ pub(crate) async fn run_driver( return; } }; - // 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(), @@ -149,26 +148,29 @@ pub(crate) async fn run_driver( metrics.clone(), remote_addr, )); - let result = - accept_web_connection(session, mtp_path, connection, send_pongs, policy, host_config.clone()) - .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(); + let mtp_tx = mtp_tx.clone(); + let auth_semaphore = auth_semaphore.clone(); + let host_config = host_config.clone(); + let connection = connection.clone(); + tokio::spawn(async move { + let result = + accept_web_connection(session, mtp_path, connection, send_pongs, policy, host_config, auth_semaphore) + .await; + match mtp_tx.try_send(result) { + Ok(()) => {} + 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(); + } } } - 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/server.rs b/mtp-webserver/src/server.rs index 0dd27ea..b8f9dfd 100644 --- a/mtp-webserver/src/server.rs +++ b/mtp-webserver/src/server.rs @@ -228,6 +228,7 @@ impl MTPWebServer { 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 auth_semaphore = Arc::new(Semaphore::new(web_config.max_connections)); let router = web_config.router.clone(); let metrics = web_config.metrics.clone(); let driver_config = DriverConfig { @@ -240,6 +241,7 @@ impl MTPWebServer { policy: host_config.policy, host_config, metrics: web_config.metrics.clone(), + auth_semaphore, }; let quic_driver = tokio::spawn(run_driver( driver_endpoint, @@ -339,10 +341,23 @@ fn build_endpoint( .map_err(|_| CommunicationError::CertificateLoadFailed)?; tls.alpn_protocols = vec![b"h3".to_vec()]; - let server = quinn::ServerConfig::with_crypto(Arc::new( + let mut server = quinn::ServerConfig::with_crypto(Arc::new( quinn::crypto::rustls::QuicServerConfig::try_from(tls) .map_err(|error| CommunicationError::Other(error.to_string()))?, )); + // Apply policy keepalive and idle timeout settings to Quinn + server.transport_config({ + let mut transport = quinn::TransportConfig::default(); + if let Some(keep_alive) = config.policy.keep_alive_interval { + transport.keep_alive_interval(Some(keep_alive)); + } + if let Some(idle_timeout) = config.policy.max_idle_timeout { + transport.max_idle_timeout(Some(idle_timeout.try_into().map_err( + |error| CommunicationError::Other(format!("{error}")), + )?)); + } + Arc::new(transport) + }); quinn::Endpoint::server(server, SocketAddr::new(config.ip, config.port)) .map_err(|error| CommunicationError::Other(error.to_string())) } diff --git a/mtp-webserver/src/transport.rs b/mtp-webserver/src/transport.rs index f4928e4..1b2eb42 100644 --- a/mtp-webserver/src/transport.rs +++ b/mtp-webserver/src/transport.rs @@ -16,6 +16,9 @@ use std::time::Instant; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tracing::error; +#[cfg(feature = "crypto")] +const GUEST_ID_MAX_RETRIES: u32 = 100; + type Session = h3_webtransport::server::WebTransportSession; type H3SendStream = h3_webtransport::stream::SendStream, Bytes>; type H3RecvStream = h3_webtransport::stream::RecvStream; @@ -225,6 +228,38 @@ pub type WebMtpReceiver = GenericReceiver; pub type WebMTPConnection = mtp_host::MTPConnection; +#[cfg(feature = "crypto")] +impl H3TransportConnection { + /// Assign a unique guest ID, using the configured generator if present. + async fn assign_guest_id(host_config: &HostConfig) -> Result { + if let Some(ref generator) = host_config.guest_id_generator { + let id = generator().await.ok_or_else(|| { + AcceptError::AuthenticationFailed( + "guest id generator rejected the connection".into(), + ) + })?; + if id > mtp_codec::MAX_WIRE_ID { + return Err(AcceptError::AuthenticationFailed( + "guest id exceeds wire limit".into(), + )); + } + if (host_config.get_existing_client)(id, None).await.is_none() { + return Ok(id); + } + } + // Fall back to random ID with collision check + for _ in 0..GUEST_ID_MAX_RETRIES { + let id = rand::random::() & mtp_codec::MAX_WIRE_ID; + if (host_config.get_existing_client)(id, None).await.is_none() { + return Ok(id); + } + } + Err(AcceptError::AuthenticationFailed( + "failed to allocate a unique guest id after retries".into(), + )) + } +} + pub(crate) async fn accept_web_connection( session: Arc, path: String, @@ -232,15 +267,25 @@ pub(crate) async fn accept_web_connection( send_pongs: bool, policy: Policy, host_config: Arc, + auth_semaphore: Arc, ) -> Result { #[cfg(feature = "crypto")] { - tokio::time::timeout( + let permit = auth_semaphore + .clone() + .acquire_owned() + .await + .map_err(|_| { + AcceptError::AuthenticationFailed("authentication service stopped".into()) + })?; + let result = 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)) + .unwrap_or(Err(AcceptError::AuthenticationTimedOut)); + drop(permit); + result } #[cfg(not(feature = "crypto"))] @@ -278,12 +323,12 @@ async fn accept_web_connection_inner( DataValue::Str(value) => Some(value.clone()), _ => None, }; - let sender = WebMtpSender::new(transport, policy); + let sender = WebMtpSender::new(transport, policy.clone()); if send_pongs { receiver.respond_to_pings(sender.clone()).await; } let connection: WebMTPConnection = - mtp_host::MTPConnection::from_transport_parts_with_remote_addr( + mtp_host::MTPConnection::from_transport_parts_with_policy( negotiated, codec, sender, @@ -291,23 +336,101 @@ async fn accept_web_connection_inner( path, description.clone(), Some(remote_addr), + policy, ); + #[cfg(not(feature = "crypto"))] + { + connection.receiver.set_max_message_size(max_message_size); + return Ok(connection); + } #[cfg(feature = "crypto")] let mut connection = connection; #[cfg(feature = "crypto")] - if !matches!( - _host_config.authentication_policy, - mtp_host::AuthenticationPolicy::Unauthenticated - ) { + { + use mtp_codec::{CommunicationType, DataType, DataValue}; use mtp_crypto::{ Ed25519Signer, MlDsaSigner, PublicKeyBundle, SignatureScheme, auth, verify_ed25519, - verify_ml_dsa, }; + let tm = mtp_codec::TypeMap::latest(); + let is_allow_auth = matches!( + _host_config.authentication_policy, + mtp_host::AuthenticationPolicy::AllowAuthentication + ); + let is_force_auth = matches!( + _host_config.authentication_policy, + mtp_host::AuthenticationPolicy::ForceAuthentication + ); + + // Unauthenticated: send accepted response with guest ID (or ID 0) + if !is_allow_auth && !is_force_auth { + let response = mtp_codec::CommunicationValue::new(CommunicationType::IdentificationResponse) + .add_typed_default(DataType::Connected, DataValue::BoolTrue) + .add_typed_default( + DataType::Version, + DataValue::Str(connection.version.to_string()), + ) + .add_typed_default(DataType::Id, 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); + return Ok(connection); + } + + // AllowAuthentication / ForceAuthentication: perform authentication 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 first_type = first.get_type(); + + let id_type = CommunicationType::Identification.try_to_id(&tm); + let reg_type = CommunicationType::Register.try_to_id(&tm); + let first_type_opt = Some(first_type); + + let (client_id, client_bundle, response_type, is_guest) = if is_allow_auth + && first_type_opt == id_type { + // AllowAuthentication Identification: try lookup, fall back to guest + let id = match first.get_data(DataType::Id) { + DataValue::UnsignedNumber(value) => *value as u64, + _ => 0, + }; + if id > 0 { + if let Some(bundle) = + (_host_config.get_existing_client)(id, description.clone()).await + { + (id, Some(bundle), CommunicationType::IdentificationResponse, false) + } else { + // Unknown client: fall back to guest + let guest_id = H3TransportConnection::assign_guest_id(&_host_config).await?; + (guest_id, None, CommunicationType::IdentificationResponse, true) + } + } else { + // ID zero or missing: fall back to guest + let guest_id = H3TransportConnection::assign_guest_id(&_host_config).await?; + (guest_id, None, CommunicationType::IdentificationResponse, true) + } + } else if first_type_opt == reg_type { + // Registration: always authenticate (both AllowAuth and ForceAuth) + 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, Some(bundle), CommunicationType::RegisterResponse, false) + } else if first_type_opt == id_type { + // ForceAuthentication Identification: require lookup let id = match first.get_data(DataType::Id) { DataValue::UnsignedNumber(value) => *value as u64, _ => { @@ -319,70 +442,101 @@ async fn accept_web_connection_inner( 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) + (id, Some(bundle), CommunicationType::IdentificationResponse, false) } else { return Err(AcceptError::AuthenticationFailed( "unexpected authentication message".into(), )); }; - tracing::debug!(elapsed = ?client_lookup_started.elapsed(), "web authentication handshake: identify client"); + tracing::debug!(elapsed = ?client_lookup_started.elapsed(), is_guest, "web authentication handshake: identify client"); - let signer_init_started = Instant::now(); - let host_pq_signer = if !_host_config + // Guest path: skip challenge/response, send accepted with guest ID + if is_guest { + let guest_id = client_id; + let response = mtp_codec::CommunicationValue::new(CommunicationType::IdentificationResponse) + .add_typed_default(DataType::Connected, DataValue::BoolTrue) + .add_typed_default(DataType::Version, DataValue::Str(connection.version.to_string())) + .add_typed_default(DataType::Id, DataValue::UnsignedNumber(guest_id as u128)); + 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); + connection.auth_state = mtp_host::AuthState::Unauthenticated; + connection.client_id = guest_id; + return Ok(connection); + } + + let client_bundle = client_bundle.unwrap(); + + // PQ preflight: if host requires PQ, it must have a PQ key + let pq_enabled = !_host_config .host_keyring .sig_pq_secret_key .as_bytes() - .is_empty() + .is_empty(); + if _host_config.require_pq + && (!pq_enabled + || _host_config + .host_keyring + .sig_pq_public_key + .as_bytes() + .is_empty()) { - Some( + return Err(AcceptError::AuthenticationFailed( + "host requires PQ authentication but has no PQ signing key".into(), + )); + } + + let signer_init_started = Instant::now(); + let host_pq_signer = if pq_enabled { + Some(Arc::new( 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 host_sign = |payload: Vec| { + let host_config = _host_config.clone(); + let pq_signer = host_pq_signer.clone(); + async move { + let signer = + Ed25519Signer::new(&host_config.host_keyring.sig_cl_secret_key) + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?; + if let Some(pq_signer) = pq_signer { + mtp_crypto::sign_parallel::sign_dual_parallel_shared_pq( + signer, + pq_signer, + payload, + ) + .await + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string())) + } else { + let sig = signer + .sign(&payload) + .map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?; + Ok((sig, Vec::new())) + } + } }; let sign_challenge_started = Instant::now(); - let (sig, pq_sig) = host_sign(&auth::challenge_payload(client_id, server_challenge))?; + let (sig, pq_sig) = host_sign(auth::challenge_payload(client_id, server_challenge)).await?; tracing::debug!(elapsed = ?sign_challenge_started.elapsed(), "web authentication handshake: sign challenge"); let mut challenge = - mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Challenge) + mtp_codec::CommunicationValue::new(CommunicationType::Challenge) .add_typed_default( DataType::ServerNonce, DataValue::UnsignedNumber(server_challenge), @@ -396,7 +550,7 @@ async fn accept_web_connection_inner( DataValue::BoolFalse }, ); - if !pq_sig.is_empty() { + if pq_enabled { challenge = challenge.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_sig)); } @@ -425,7 +579,7 @@ async fn accept_web_connection_inner( } }; 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) + if Some(proof.get_type()) != CommunicationType::ChallengeResponse.try_to_id(&tm) { return Err(AcceptError::AuthenticationFailed( "missing challenge response".into(), @@ -452,7 +606,7 @@ async fn accept_web_connection_inner( _ => &[], }; let payload = if first.get_type() - == mtp_codec::CommunicationType::Register + == CommunicationType::Register .try_to_id(&tm) .unwrap() { @@ -465,36 +619,54 @@ async fn accept_web_connection_inner( } else { auth::login_proof_payload(&version.to_string(), client_id, server_challenge, nonce) }; + + // Verify client proof: classical is always required; PQ is verified + // when supplied (even if not required), matching native behavior. + let has_client_pq_key = !client_bundle.sig_pq_public_key.as_bytes().is_empty(); let verify_proof_started = Instant::now(); - if verify_ed25519(&client_bundle.sig_cl_public_key, &payload, signature).is_err() { + let proof_ok = if pq_signature.is_empty() { + !_host_config.require_pq + && verify_ed25519(&client_bundle.sig_cl_public_key, &payload, signature).is_ok() + } else if has_client_pq_key { + mtp_crypto::sign_parallel::verify_dual_parallel( + client_bundle.sig_cl_public_key.clone(), + client_bundle.sig_pq_public_key.clone(), + payload, + signature.to_vec(), + pq_signature.to_vec(), + ) + .await + .is_ok() + } else { + false + }; + tracing::debug!(elapsed = ?verify_proof_started.elapsed(), "web authentication handshake: verify client proof"); + + if !proof_ok { + let rejection = mtp_codec::CommunicationValue::new(CommunicationType::IdentificationResponse) + .add_typed_default(DataType::Connected, DataValue::BoolFalse) + .add_typed_default(DataType::ErrorMessage, DataValue::Str("client proof signature invalid".into())); + let _ = connection.sender.send(&rejection).await; + connection.sender.close(); 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 { + let assigned_id = if response_type == 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( + let (final_sig, final_pq) = host_sign(auth::host_final_payload( assigned_id, nonce, server_challenge, - ))?; + )) + .await?; 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) @@ -502,7 +674,7 @@ async fn accept_web_connection_inner( .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() { + if pq_enabled { response = response.add_typed_default(DataType::PqSignature, DataValue::Bytes(final_pq)); } @@ -523,37 +695,8 @@ async fn accept_web_connection_inner( connection.auth_state = mtp_host::AuthState::Authenticated; connection.client_id = assigned_id; connection.client_public_key = Some(client_bundle); - return Ok(connection); + 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/transport/src/connection.rs b/transport/src/connection.rs index 726ba1a..be04588 100644 --- a/transport/src/connection.rs +++ b/transport/src/connection.rs @@ -57,7 +57,7 @@ impl Default for Policy { application_close_code: 0, open_stream_timeout: Duration::from_millis(2_000), write_timeout: Duration::from_millis(2_000), - accept_stream_timeout: Duration::from_millis(10_000), + accept_stream_timeout: Duration::from_millis(500), read_timeout: Duration::from_millis(30_000), keep_alive_interval: Some(Duration::from_secs(3)), max_idle_timeout: Some(Duration::from_secs(30)), @@ -1212,28 +1212,6 @@ mod tests { assert_ne!(SendMode::PersistentStream, SendMode::SingleStreamPerMessage); } - #[test] - fn test_policy_default_values() { - let p = Policy::default(); - assert_eq!(p.send_mode, SendMode::PersistentStream); - assert_eq!(p.max_message_size, 16 * 1024 * 1024); - assert_eq!(p.handshake_max_message_size, 64 * 1024); - assert_eq!(p.close_frame_len, u32::MAX); - assert_eq!(p.application_close_code, 0); - assert_eq!(p.open_stream_timeout, Duration::from_millis(2_000)); - assert_eq!(p.write_timeout, Duration::from_millis(2_000)); - assert_eq!(p.accept_stream_timeout, Duration::from_millis(10_000)); - assert_eq!(p.read_timeout, Duration::from_millis(30_000)); - assert_eq!(p.keep_alive_interval, Some(Duration::from_secs(3))); - assert_eq!(p.max_idle_timeout, Some(Duration::from_secs(30))); - assert_eq!(p.force_close_delay, Duration::from_millis(300)); - assert_eq!(p.persistent_stream_max_retries, 4); - assert_eq!(p.persistent_stream_retry_backoff, Duration::from_millis(20)); - assert_eq!(p.receiver_queue_capacity, 1000); - assert_eq!(p.max_concurrent_stream_tasks, 128); - assert_eq!(p.max_frames_per_stream, None); - } - #[test] fn test_policy_clone() { let p = Policy::default(); diff --git a/transport/src/generic_connection.rs b/transport/src/generic_connection.rs index e389e30..7dfb960 100644 --- a/transport/src/generic_connection.rs +++ b/transport/src/generic_connection.rs @@ -166,6 +166,7 @@ pub struct GenericReceiver { connection: C, ping_sender: Arc>>>, max_message_size: Arc, + _accept_task: Arc>, } impl Clone for GenericReceiver { @@ -177,6 +178,15 @@ impl Clone for GenericReceiver { connection: self.connection.clone(), ping_sender: self.ping_sender.clone(), max_message_size: self.max_message_size.clone(), + _accept_task: self._accept_task.clone(), + } + } +} + +impl Drop for GenericReceiver { + fn drop(&mut self) { + if Arc::strong_count(&self._accept_task) == 1 { + self._accept_task.abort(); } } } @@ -187,17 +197,34 @@ 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 max_message_size = Arc::new(AtomicU64::new( + policy.handshake_max_message_size.min(policy.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 task_accept_task_tx = tx.clone(); + #[cfg(feature = "pipes")] + let task_accept_task_pipe_tx = pipe_tx.clone(); + let accept_task = tokio::spawn(async move { let limit = Arc::new(Semaphore::new( task_policy.max_concurrent_stream_tasks.max(1), )); loop { - let stream = match timeout( + // Backpressure: stop accepting new streams if the output queue is full. + #[cfg(feature = "pipes")] + let cap_full = task_accept_task_tx.capacity() == 0 + || task_accept_task_pipe_tx.capacity() == 0; + #[cfg(not(feature = "pipes"))] + let cap_full = task_accept_task_tx.capacity() == 0; + + if cap_full { + tokio::time::sleep(std::time::Duration::from_millis(1)).await; + continue; + } + + let stream = match tokio::time::timeout( task_policy.accept_stream_timeout, task_connection.accept_uni(), ) @@ -205,29 +232,34 @@ impl GenericReceiver { { Ok(Ok(stream)) => stream, Ok(Err(error)) => { - let _ = tx.send(Err(error)).await; + let _ = task_accept_task_tx.send(Err(error)).await; break; } Err(_) => { if task_connection.close_reason().is_some() { - let _ = tx.send(Err(CommunicationError::StreamClosed)).await; + let _ = task_accept_task_tx + .send(Err(CommunicationError::StreamClosed)) + .await; break; } else { continue; } } }; - let tx = tx.clone(); + // Acquire semaphore permit BEFORE spawning the task. + let permit = match limit.clone().acquire_owned().await { + Ok(permit) => permit, + Err(_) => break, + }; + let tx = task_accept_task_tx.clone(); #[cfg(feature = "pipes")] - let pipe_tx = pipe_tx.clone(); + let pipe_tx = task_accept_task_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(); + let connection = task_connection.clone(); tokio::spawn(async move { - let Ok(_permit) = permit.acquire_owned().await else { - return; - }; + let _permit = permit; let mut stream = stream; let mut frames = 0usize; loop { @@ -235,24 +267,36 @@ impl GenericReceiver { .max_frames_per_stream .is_some_and(|max| frames >= max) { + let close_error = CommunicationError::StreamError; + let _ = tx.send(Err(close_error.clone())).await; + connection.close( + policy.application_close_code, + b"max frames exceeded", + ); break; } let mut len = [0; 4]; - match timeout(policy.read_timeout, stream.read_exact(&mut len)).await { + match tokio::time::timeout(policy.read_timeout, stream.read_exact(&mut len)) + .await + { Ok(Ok(())) => {} 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"); + let _ = tx.send(Err(error)).await; + connection.close( + policy.application_close_code, + b"frame header read error", + ); break; } - Err(error) => { - tracing::error!( - "[mtp-transport] frame header read timed out: {error}" + Err(_) => { + tracing::warn!("MTP receive stream timed out while reading frame header"); + let _ = tx.send(Err(CommunicationError::StreamError)).await; + connection.close( + policy.application_close_code, + b"frame header timeout", ); - tracing::warn!(%error, "MTP receive stream timed out while reading frame header"); break; } } @@ -260,8 +304,14 @@ impl GenericReceiver { if len == policy.close_frame_len { break; } - if len as u64 > max_message_size.load(Ordering::Relaxed) { + let frame_limit = max_message_size.load(Ordering::Relaxed); + if len as u64 > frame_limit { tracing::warn!(len, "MTP receive stream frame is too large"); + let _ = tx.send(Err(CommunicationError::MessageTooLarge)).await; + connection.close( + policy.application_close_code, + b"frame too large", + ); break; } let target_len = len as usize; @@ -271,12 +321,17 @@ impl GenericReceiver { target_len, "MTP receive stream could not reserve frame body" ); + let _ = tx.send(Err(CommunicationError::MessageTooLarge)).await; + connection.close( + policy.application_close_code, + b"frame allocation failed", + ); break; } while body.len() < target_len { let chunk_len = (target_len - body.len()).min(16 * 1024); let mut chunk = [0u8; 16 * 1024]; - let body_read = timeout( + let body_read = tokio::time::timeout( policy.read_timeout, stream.read_exact(&mut chunk[..chunk_len]), ) @@ -284,16 +339,16 @@ impl GenericReceiver { if !matches!(&body_read, Ok(Ok(()))) || body.try_reserve(chunk_len).is_err() { - 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" ); + let _ = tx.send(Err(CommunicationError::StreamError)).await; + connection.close( + policy.application_close_code, + b"frame body read error", + ); break; } body.extend_from_slice(&chunk[..chunk_len]); @@ -306,6 +361,12 @@ impl GenericReceiver { Ok(message) => message, Err(_) => { tracing::warn!("MTP receive stream contained an invalid frame"); + let _ = tx.send(Err(CommunicationError::ParseCommunicationValue)) + .await; + connection.close( + policy.application_close_code, + b"invalid frame", + ); break; } }; @@ -367,6 +428,7 @@ impl GenericReceiver { connection, ping_sender, max_message_size, + _accept_task: Arc::new(accept_task), } } pub async fn respond_to_pings(&self, sender: GenericSender) { diff --git a/transport/src/host.rs b/transport/src/host.rs index e64ee11..a2f8f1f 100644 --- a/transport/src/host.rs +++ b/transport/src/host.rs @@ -126,31 +126,34 @@ pub async fn host_with_config( 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) => { - debug!("incoming WebTransport session failed: {e}"); - 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 - { - Ok(conn) => conn, - Err(e) => { - debug!("WebTransport request accept failed: {e}"); - 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())); + let policy = Arc::clone(&policy); + tokio::spawn(async move { + let session_started = Instant::now(); + let request = match incoming_session.await { + Ok(req) => req, + Err(e) => { + debug!("incoming WebTransport session failed: {e}"); + return; + } + }; + 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 + { + Ok(conn) => conn, + Err(e) => { + debug!("WebTransport request accept failed: {e}"); + return; + } + }; + tracing::debug!(elapsed = ?request_accept_started.elapsed(), "host accept loop: accept WebTransport request"); + + handle_connection(connection, incoming_tx, policy).await; + }); } });