Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
1af9732554 |
|||
|
cab2cd7a52 |
22 changed files with 2912 additions and 1011 deletions
1
example/.gitignore
vendored
1
example/.gitignore
vendored
|
|
@ -14,3 +14,4 @@ web-client/dist/
|
||||||
client.id
|
client.id
|
||||||
*.mk
|
*.mk
|
||||||
*.mpkb
|
*.mpkb
|
||||||
|
metrics/
|
||||||
|
|
|
||||||
4
example/Cargo.lock
generated
4
example/Cargo.lock
generated
|
|
@ -237,6 +237,8 @@ version = "0.2.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"mtp",
|
"mtp",
|
||||||
"rand 0.10.2",
|
"rand 0.10.2",
|
||||||
|
"serde",
|
||||||
|
"serde_json",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tracing-subscriber",
|
"tracing-subscriber",
|
||||||
]
|
]
|
||||||
|
|
@ -1988,6 +1990,8 @@ dependencies = [
|
||||||
"hex",
|
"hex",
|
||||||
"http",
|
"http",
|
||||||
"mtp",
|
"mtp",
|
||||||
|
"rand 0.10.2",
|
||||||
|
"serde",
|
||||||
"serde_json",
|
"serde_json",
|
||||||
"tokio",
|
"tokio",
|
||||||
"tracing-subscriber",
|
"tracing-subscriber",
|
||||||
|
|
|
||||||
|
|
@ -12,3 +12,5 @@ mtp = { version = "0.2.0", path = "../../", features = ["client", "crypto", "fil
|
||||||
tokio = { version = "1", features = ["full"] }
|
tokio = { version = "1", features = ["full"] }
|
||||||
rand = "0.10.1"
|
rand = "0.10.1"
|
||||||
tracing-subscriber = "0.3.23"
|
tracing-subscriber = "0.3.23"
|
||||||
|
serde = { version = "1", features = ["derive"] }
|
||||||
|
serde_json = "1"
|
||||||
|
|
|
||||||
|
|
@ -1,4 +1,4 @@
|
||||||
use std::time::Instant;
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use tokio::fs;
|
use tokio::fs;
|
||||||
|
|
||||||
|
|
@ -12,7 +12,7 @@ pub async fn connect_or_register(
|
||||||
mut config: ClientConfig,
|
mut config: ClientConfig,
|
||||||
host_public_key: PublicKeyBundle,
|
host_public_key: PublicKeyBundle,
|
||||||
key_prefix: &str,
|
key_prefix: &str,
|
||||||
) -> Result<(MTPConnection, Keyring), Box<dyn std::error::Error>> {
|
) -> Result<(MTPConnection, Keyring, String, Duration), Box<dyn std::error::Error>> {
|
||||||
let keyring_path = format!("{key_prefix}.mk");
|
let keyring_path = format!("{key_prefix}.mk");
|
||||||
let id_path = format!("{key_prefix}.id");
|
let id_path = format!("{key_prefix}.id");
|
||||||
|
|
||||||
|
|
@ -30,12 +30,12 @@ pub async fn connect_or_register(
|
||||||
config.client_id = client_id;
|
config.client_id = client_id;
|
||||||
let auth_started = Instant::now();
|
let auth_started = Instant::now();
|
||||||
let conn = MTPClient::auth_connect(config, &keyring, &host_public_key).await?;
|
let conn = MTPClient::auth_connect(config, &keyring, &host_public_key).await?;
|
||||||
|
let auth_duration = auth_started.elapsed();
|
||||||
println!(
|
println!(
|
||||||
"Authenticated (version {}) in {:?}",
|
"Authenticated (version {}) in {:?}",
|
||||||
conn.version,
|
conn.version, auth_duration
|
||||||
auth_started.elapsed()
|
|
||||||
);
|
);
|
||||||
return Ok((conn, keyring));
|
return Ok((conn, keyring, "connect".into(), auth_duration));
|
||||||
}
|
}
|
||||||
|
|
||||||
println!("No existing keys found: registering new client");
|
println!("No existing keys found: registering new client");
|
||||||
|
|
@ -52,12 +52,14 @@ pub async fn connect_or_register(
|
||||||
sig_sk,
|
sig_sk,
|
||||||
);
|
);
|
||||||
|
|
||||||
|
let reg_started = Instant::now();
|
||||||
let conn = MTPClient::auth_register(config, &keyring, &host_public_key).await?;
|
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)?;
|
save_keyring_raw(&keyring, &keyring_path)?;
|
||||||
fs::write(&id_path, conn.client_id.to_string()).await?;
|
fs::write(&id_path, conn.client_id.to_string()).await?;
|
||||||
println!("Saved client keys -> {keyring_path}");
|
println!("Saved client keys -> {keyring_path}");
|
||||||
|
|
||||||
Ok((conn, keyring))
|
Ok((conn, keyring, "register".into(), reg_duration))
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,9 +1,11 @@
|
||||||
mod auth;
|
mod auth;
|
||||||
|
mod metrics;
|
||||||
mod messages;
|
mod messages;
|
||||||
mod pipes;
|
mod pipes;
|
||||||
|
|
||||||
use std::fs;
|
use std::fs;
|
||||||
use std::path::Path;
|
use std::path::Path;
|
||||||
|
use std::time::Duration;
|
||||||
|
|
||||||
use mtp::client::ClientConfig;
|
use mtp::client::ClientConfig;
|
||||||
use mtp::files::load_public_key_bundle;
|
use mtp::files::load_public_key_bundle;
|
||||||
|
|
@ -36,6 +38,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
let mut client_metrics = metrics::ClientMetrics::load("metrics/client_sessions.json");
|
||||||
|
|
||||||
println!("Connecting to 127.0.0.1:8080 ...");
|
println!("Connecting to 127.0.0.1:8080 ...");
|
||||||
|
|
||||||
let config = ClientConfig::new("https://127.0.0.1:8080")
|
let config = ClientConfig::new("https://127.0.0.1:8080")
|
||||||
|
|
@ -43,11 +47,43 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
.with_description("MTP example client");
|
.with_description("MTP example client");
|
||||||
|
|
||||||
let server_bundle = host_public_key.clone();
|
let server_bundle = host_public_key.clone();
|
||||||
let (conn, keyring) = auth::connect_or_register(config, host_public_key, "client").await?;
|
let (conn, keyring, auth_method, auth_duration) =
|
||||||
messages::send_and_receive(&conn, &keyring, &server_bundle).await?;
|
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 ---");
|
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;
|
conn.sender.close().await;
|
||||||
println!("\nDone");
|
println!("\nDone");
|
||||||
|
|
|
||||||
|
|
@ -1,3 +1,5 @@
|
||||||
|
use std::time::{Duration, Instant};
|
||||||
|
|
||||||
use mtp::client::MTPConnection;
|
use mtp::client::MTPConnection;
|
||||||
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
|
use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue};
|
||||||
use mtp::crypto::{Ed25519Signer, EncryptionType, Keyring, PublicKeyBundle, SigAlgorithm};
|
use mtp::crypto::{Ed25519Signer, EncryptionType, Keyring, PublicKeyBundle, SigAlgorithm};
|
||||||
|
|
@ -89,17 +91,22 @@ pub async fn send_and_receive(
|
||||||
conn: &MTPConnection,
|
conn: &MTPConnection,
|
||||||
keyring: &Keyring,
|
keyring: &Keyring,
|
||||||
server_bundle: &PublicKeyBundle,
|
server_bundle: &PublicKeyBundle,
|
||||||
) -> Result<(), Box<dyn std::error::Error>> {
|
) -> Result<Duration, Box<dyn std::error::Error>> {
|
||||||
let msg = build_demo_message(conn.client_id, keyring, server_bundle)?;
|
let msg = build_demo_message(conn.client_id, keyring, server_bundle)?;
|
||||||
println!("Sending: {msg}");
|
println!("Sending: {msg}");
|
||||||
|
let start = Instant::now();
|
||||||
conn.sender.send(&msg).await?;
|
conn.sender.send(&msg).await?;
|
||||||
|
|
||||||
match conn.receive().await {
|
match conn.receive().await {
|
||||||
Ok(resp) => {
|
Ok(resp) => {
|
||||||
|
let roundtrip = start.elapsed();
|
||||||
println!("Received: {resp}");
|
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(())
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
554
example/client/src/metrics.rs
Normal file
554
example/client/src/metrics.rs
Normal file
|
|
@ -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<String>,
|
||||||
|
pub message_roundtrip_ms: f64,
|
||||||
|
pub pipe_results: Vec<PipeResult>,
|
||||||
|
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<ClientSessionRecord>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
|
||||||
|
pub struct ClientMetricsFile {
|
||||||
|
pub sessions: Vec<ClientSessionRecord>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Live metrics state
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
pub struct ClientMetrics {
|
||||||
|
sessions: Vec<ClientSessionRecord>,
|
||||||
|
}
|
||||||
|
|
||||||
|
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::<ClientMetricsFile>(&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<String>,
|
||||||
|
message_roundtrip_ms: f64,
|
||||||
|
pipe_results: Vec<PipeResult>,
|
||||||
|
}
|
||||||
|
|
||||||
|
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::<f64>()
|
||||||
|
/ 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::<f64>()
|
||||||
|
/ 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -3,13 +3,16 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||||
use tokio::sync::oneshot;
|
use tokio::sync::oneshot;
|
||||||
use tokio::time::{Duration, Instant};
|
use tokio::time::{Duration, Instant};
|
||||||
|
|
||||||
|
use crate::metrics::PipeResult;
|
||||||
|
|
||||||
pub async fn run_pipe_demo(
|
pub async fn run_pipe_demo(
|
||||||
conn: &MTPConnection,
|
conn: &MTPConnection,
|
||||||
iterations: usize,
|
iterations: usize,
|
||||||
) -> Result<(), Box<dyn std::error::Error>> {
|
) -> Result<Vec<PipeResult>, Box<dyn std::error::Error>> {
|
||||||
let sizes = [64, 256, 1024, 4096];
|
let sizes = [64, 256, 1024, 4096];
|
||||||
let mut all_elapsed = Vec::with_capacity(sizes.len() * iterations);
|
let mut all_elapsed = Vec::with_capacity(sizes.len() * iterations);
|
||||||
let mut all_data_only = 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() {
|
for (i, &size) in sizes.iter().enumerate() {
|
||||||
let mut size_elapsed = Vec::with_capacity(iterations);
|
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,
|
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_elapsed.push(overall_elapsed);
|
||||||
size_data_only.push(data_only_elapsed);
|
size_data_only.push(data_only_elapsed);
|
||||||
all_elapsed.push(overall_elapsed);
|
all_elapsed.push(overall_elapsed);
|
||||||
|
|
@ -123,7 +134,7 @@ pub async fn run_pipe_demo(
|
||||||
all_elapsed.len()
|
all_elapsed.len()
|
||||||
);
|
);
|
||||||
|
|
||||||
Ok(())
|
Ok(pipe_results)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Helper: average a slice of Durations without overflowing.
|
/// Helper: average a slice of Durations without overflowing.
|
||||||
|
|
|
||||||
|
|
@ -15,3 +15,5 @@ serde_json = { version = "1" }
|
||||||
hex = "0.4"
|
hex = "0.4"
|
||||||
base64 = "0.22"
|
base64 = "0.22"
|
||||||
tracing-subscriber = "0.3.23"
|
tracing-subscriber = "0.3.23"
|
||||||
|
serde = { version = "1", features = ["derive"] }
|
||||||
|
rand = "0.10.1"
|
||||||
|
|
|
||||||
|
|
@ -1,6 +1,7 @@
|
||||||
mod clients;
|
mod clients;
|
||||||
mod handlers;
|
mod handlers;
|
||||||
mod keys;
|
mod keys;
|
||||||
|
mod metrics;
|
||||||
mod tls;
|
mod tls;
|
||||||
#[path = "web-server.rs"]
|
#[path = "web-server.rs"]
|
||||||
mod web_server;
|
mod web_server;
|
||||||
|
|
@ -39,7 +40,7 @@ async fn handle_pipe_loopback(
|
||||||
mtp::webserver::WebMtpSender,
|
mtp::webserver::WebMtpSender,
|
||||||
mtp::webserver::H3TransportReceiver,
|
mtp::webserver::H3TransportReceiver,
|
||||||
>,
|
>,
|
||||||
) -> Result<(), Box<dyn std::error::Error>> {
|
) -> Result<u64, Box<dyn std::error::Error>> {
|
||||||
let pipe_id = request.id();
|
let pipe_id = request.id();
|
||||||
println!(" [loopback] Accepting pipe {pipe_id} ...");
|
println!(" [loopback] Accepting pipe {pipe_id} ...");
|
||||||
let mut reader = request.accept().await?;
|
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?;
|
let copied = tokio::io::copy(&mut reader, &mut writer).await?;
|
||||||
writer.finish_async().await?;
|
writer.finish_async().await?;
|
||||||
println!(" [loopback] Pipe {pipe_id} complete ({copied} bytes)");
|
println!(" [loopback] Pipe {pipe_id} complete ({copied} bytes)");
|
||||||
Ok(())
|
Ok(copied)
|
||||||
}
|
}
|
||||||
|
|
||||||
#[tokio::main]
|
#[tokio::main]
|
||||||
|
|
@ -119,6 +120,10 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
|
|
||||||
|
let metrics = std::sync::Arc::new(metrics::ServerMetrics::load(
|
||||||
|
"metrics/server_sessions.json",
|
||||||
|
));
|
||||||
|
|
||||||
println!("Starting integrated MTP web server on port 8080 ...");
|
println!("Starting integrated MTP web server on port 8080 ...");
|
||||||
|
|
||||||
let config = HostConfig::new(
|
let config = HostConfig::new(
|
||||||
|
|
@ -138,8 +143,22 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
println!("TCP: HTTP/1.1 and HTTP/2");
|
println!("TCP: HTTP/1.1 and HTTP/2");
|
||||||
println!("UDP: HTTP/3 and WebTransport");
|
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 decrypt_keyring = Arc::clone(&decrypt_keyring);
|
||||||
|
let metrics = Arc::clone(&metrics);
|
||||||
|
metrics.record_connection_version(&conn.version.to_string());
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let desc = conn.description.as_deref().unwrap_or("(no description)");
|
let desc = conn.description.as_deref().unwrap_or("(no description)");
|
||||||
println!(
|
println!(
|
||||||
|
|
@ -151,12 +170,14 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
);
|
);
|
||||||
println!("Client ID: {}", conn.client_id);
|
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();
|
let tm: &TypeMap = conn.codec.registry().get(&conn.version).unwrap();
|
||||||
|
|
||||||
println!("Waiting for messages / pipe requests ...");
|
println!("Waiting for messages / pipe requests ...");
|
||||||
let mut pipe_open = true;
|
let mut pipe_open = true;
|
||||||
let mut message_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 {
|
while pipe_open || message_open {
|
||||||
let activity = tokio::time::timeout(CONNECTION_IDLE_TIMEOUT, async {
|
let activity = tokio::time::timeout(CONNECTION_IDLE_TIMEOUT, async {
|
||||||
|
|
@ -165,8 +186,17 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
pipe_request = conn.receive_pipe(), if pipe_open => {
|
pipe_request = conn.receive_pipe(), if pipe_open => {
|
||||||
match pipe_request {
|
match pipe_request {
|
||||||
Ok(request) => {
|
Ok(request) => {
|
||||||
if let Err(error) = handle_pipe_loopback(&conn, request).await {
|
match handle_pipe_loopback(&conn, request).await {
|
||||||
eprintln!(" [loopback] Pipe error: {error}");
|
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)
|
Err(mtp::common::CommunicationError::StreamClosed)
|
||||||
|
|
@ -183,18 +213,24 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
message = conn.receive(), if message_open => {
|
message = conn.receive(), if message_open => {
|
||||||
match message {
|
match message {
|
||||||
Ok(message) => {
|
Ok(message) => {
|
||||||
messages_received += 1;
|
|
||||||
println!("Received: {message}");
|
println!("Received: {message}");
|
||||||
match handlers::process_and_respond(
|
let msg_start = std::time::Instant::now();
|
||||||
|
let result = handlers::process_and_respond(
|
||||||
&message,
|
&message,
|
||||||
tm,
|
tm,
|
||||||
conn.client_public_key.as_ref(),
|
conn.client_public_key.as_ref(),
|
||||||
&decrypt_keyring,
|
&decrypt_keyring,
|
||||||
) {
|
);
|
||||||
|
let latency = msg_start.elapsed();
|
||||||
|
let ok = result.is_ok();
|
||||||
|
session.record_message(latency, ok);
|
||||||
|
|
||||||
|
match result {
|
||||||
Ok(response) => {
|
Ok(response) => {
|
||||||
println!("Sending: {response}");
|
println!("Sending: {response}");
|
||||||
if let Err(error) = conn.sender.send(&response).await {
|
if let Err(error) = conn.sender.send(&response).await {
|
||||||
eprintln!("Send error: {error}");
|
eprintln!("Send error: {error}");
|
||||||
|
session.record_send_error();
|
||||||
pipe_open = false;
|
pipe_open = false;
|
||||||
message_open = false;
|
message_open = false;
|
||||||
}
|
}
|
||||||
|
|
@ -220,16 +256,27 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||||||
.await;
|
.await;
|
||||||
|
|
||||||
if activity.is_err() {
|
if activity.is_err() {
|
||||||
|
exit_reason = "idle timeout".to_string();
|
||||||
println!("Connection idle timeout reached");
|
println!("Connection idle timeout reached");
|
||||||
break;
|
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");
|
println!("Connection message limit reached");
|
||||||
break;
|
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");
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
883
example/server/src/metrics.rs
Normal file
883
example/server/src/metrics.rs
Normal file
|
|
@ -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<String, u64>,
|
||||||
|
pub sessions: Vec<SessionRecord>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[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<String, u64>,
|
||||||
|
pub sessions: Vec<SessionRecord>,
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// 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<String, u64>,
|
||||||
|
active_connections: u64,
|
||||||
|
completed_sessions: Vec<SessionRecord>,
|
||||||
|
}
|
||||||
|
|
||||||
|
pub struct ServerMetrics {
|
||||||
|
inner: Mutex<Inner>,
|
||||||
|
}
|
||||||
|
|
||||||
|
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::<ServerMetricsFile>(&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<f64>,
|
||||||
|
}
|
||||||
|
|
||||||
|
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::<f64>() / 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);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -3,7 +3,6 @@ use mtp_codec::{CommunicationType, DataType, DataValue};
|
||||||
use mtp_codec::{CommunicationValue, Version, registry::VersionedCodec};
|
use mtp_codec::{CommunicationValue, Version, registry::VersionedCodec};
|
||||||
use mtp_common::CommunicationError;
|
use mtp_common::CommunicationError;
|
||||||
use std::net::SocketAddr;
|
use std::net::SocketAddr;
|
||||||
#[cfg(feature = "pipes")]
|
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
#[cfg(feature = "pipes")]
|
#[cfg(feature = "pipes")]
|
||||||
use tokio::sync::{Mutex, mpsc};
|
use tokio::sync::{Mutex, mpsc};
|
||||||
|
|
@ -12,7 +11,6 @@ use tokio::sync::{Mutex, mpsc};
|
||||||
use crate::error::random_client_id;
|
use crate::error::random_client_id;
|
||||||
#[cfg(feature = "pipes")]
|
#[cfg(feature = "pipes")]
|
||||||
use crate::pipe::{PipeDispatcher, PipeReceiver, PipeRequest, PipeSender, run_dispatcher};
|
use crate::pipe::{PipeDispatcher, PipeReceiver, PipeRequest, PipeSender, run_dispatcher};
|
||||||
#[cfg(feature = "pipes")]
|
|
||||||
use mtp_transport::Policy;
|
use mtp_transport::Policy;
|
||||||
|
|
||||||
mod connection_capability {
|
mod connection_capability {
|
||||||
|
|
@ -161,6 +159,52 @@ where
|
||||||
client_public_key: None,
|
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<String>,
|
||||||
|
remote_addr: Option<SocketAddr>,
|
||||||
|
policy: Arc<Policy>,
|
||||||
|
) -> 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"))]
|
#[cfg(not(feature = "pipes"))]
|
||||||
|
|
@ -211,6 +255,35 @@ impl<S, R, P> MTPConnection<S, R, P> {
|
||||||
client_public_key: None,
|
client_public_key: None,
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn from_transport_parts_with_policy(
|
||||||
|
version: Version,
|
||||||
|
codec: VersionedCodec,
|
||||||
|
sender: S,
|
||||||
|
receiver: R,
|
||||||
|
path: String,
|
||||||
|
description: Option<String>,
|
||||||
|
remote_addr: Option<SocketAddr>,
|
||||||
|
_policy: Arc<Policy>,
|
||||||
|
) -> 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"))]
|
#[cfg(not(feature = "pipes"))]
|
||||||
|
|
|
||||||
825
host/src/engine.rs
Normal file
825
host/src/engine.rs
Normal file
|
|
@ -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<C>`.
|
||||||
|
pub trait HandshakeSender: Send + Sync {
|
||||||
|
fn send(
|
||||||
|
&self,
|
||||||
|
msg: &CommunicationValue,
|
||||||
|
) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send;
|
||||||
|
fn finish_stream(
|
||||||
|
&self,
|
||||||
|
) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send;
|
||||||
|
fn close(&self);
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Trait for receiving handshake messages during the opening exchange.
|
||||||
|
///
|
||||||
|
/// Implemented by both the concrete `Receiver` and `GenericReceiver<C>`.
|
||||||
|
pub trait HandshakeReceiver: Send + Sync {
|
||||||
|
fn receive(
|
||||||
|
&self,
|
||||||
|
) -> impl std::future::Future<Output = Result<CommunicationValue, CommunicationError>>
|
||||||
|
+ 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<String>,
|
||||||
|
#[cfg(feature = "crypto")]
|
||||||
|
pub auth_state: crate::error::AuthState,
|
||||||
|
#[cfg(feature = "crypto")]
|
||||||
|
pub client_id: u64,
|
||||||
|
#[cfg(feature = "crypto")]
|
||||||
|
pub client_public_key: Option<mtp_crypto::PublicKeyBundle>,
|
||||||
|
}
|
||||||
|
|
||||||
|
/// 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<HostConfig>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl HandshakeEngine {
|
||||||
|
#[cfg(feature = "crypto")]
|
||||||
|
pub fn new(registry: Registry, config: Arc<HostConfig>) -> Self {
|
||||||
|
Self { registry, config }
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(not(feature = "crypto"))]
|
||||||
|
pub fn new(registry: Registry, _config: Arc<HostConfig>) -> 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<S: HandshakeSender, R: HandshakeReceiver>(
|
||||||
|
&self,
|
||||||
|
sender: &S,
|
||||||
|
receiver: &R,
|
||||||
|
) -> Result<HandshakeResult, AcceptError> {
|
||||||
|
#[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<S: HandshakeSender, R: HandshakeReceiver>(
|
||||||
|
&self,
|
||||||
|
sender: &S,
|
||||||
|
receiver: &R,
|
||||||
|
) -> Result<HandshakeResult, AcceptError> {
|
||||||
|
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<S: HandshakeSender>(
|
||||||
|
&self,
|
||||||
|
sender: &S,
|
||||||
|
first_msg: CommunicationValue,
|
||||||
|
negotiated: Version,
|
||||||
|
codec: VersionedCodec,
|
||||||
|
description: Option<String>,
|
||||||
|
) -> Result<HandshakeResult, AcceptError> {
|
||||||
|
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<S: HandshakeSender, R: HandshakeReceiver>(
|
||||||
|
&self,
|
||||||
|
sender: &S,
|
||||||
|
receiver: &R,
|
||||||
|
first_msg: CommunicationValue,
|
||||||
|
negotiated: Version,
|
||||||
|
codec: VersionedCodec,
|
||||||
|
description: Option<String>,
|
||||||
|
version_str: &str,
|
||||||
|
client_version: Version,
|
||||||
|
) -> Result<HandshakeResult, AcceptError> {
|
||||||
|
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<S: HandshakeSender, R: HandshakeReceiver>(
|
||||||
|
&self,
|
||||||
|
sender: &S,
|
||||||
|
receiver: &R,
|
||||||
|
first_msg: CommunicationValue,
|
||||||
|
negotiated: Version,
|
||||||
|
codec: VersionedCodec,
|
||||||
|
description: Option<String>,
|
||||||
|
version_str: &str,
|
||||||
|
client_version: Version,
|
||||||
|
) -> Result<HandshakeResult, AcceptError> {
|
||||||
|
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<S: HandshakeSender, R: HandshakeReceiver>(
|
||||||
|
&self,
|
||||||
|
sender: &S,
|
||||||
|
receiver: &R,
|
||||||
|
flow: Flow,
|
||||||
|
response_type: CommunicationType,
|
||||||
|
negotiated: &Version,
|
||||||
|
codec: &VersionedCodec,
|
||||||
|
description: Option<String>,
|
||||||
|
version_str: &str,
|
||||||
|
_client_version: Version,
|
||||||
|
) -> Result<HandshakeResult, AcceptError> {
|
||||||
|
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<u8>| 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<u8> = 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<u8>,
|
||||||
|
},
|
||||||
|
}
|
||||||
|
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
// Guest ID allocation
|
||||||
|
// ---------------------------------------------------------------------------
|
||||||
|
|
||||||
|
#[cfg(feature = "crypto")]
|
||||||
|
impl HandshakeEngine {
|
||||||
|
const GUEST_ID_MAX_RETRIES: u32 = 100;
|
||||||
|
|
||||||
|
async fn assign_guest_id(&self) -> Result<u64, AcceptError> {
|
||||||
|
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<u64, AcceptError> {
|
||||||
|
for _ in 0..Self::GUEST_ID_MAX_RETRIES {
|
||||||
|
let id = rand::random::<u64>() & 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<mtp_crypto::PublicKeyBundle, AcceptError> {
|
||||||
|
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<S: HandshakeSender>(
|
||||||
|
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<S: HandshakeSender>(
|
||||||
|
sender: &S,
|
||||||
|
version: &Version,
|
||||||
|
assigned_id: Option<u64>,
|
||||||
|
) -> 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<Output = Result<(), CommunicationError>> + Send {
|
||||||
|
mtp_transport::Sender::send(self, msg)
|
||||||
|
}
|
||||||
|
fn finish_stream(
|
||||||
|
&self,
|
||||||
|
) -> impl std::future::Future<Output = Result<(), CommunicationError>> + 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<Output = Result<CommunicationValue, CommunicationError>> + Send
|
||||||
|
{
|
||||||
|
mtp_transport::Receiver::receive(self)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<C: mtp_transport::TransportConnection> HandshakeSender for mtp_transport::GenericSender<C> {
|
||||||
|
fn send(
|
||||||
|
&self,
|
||||||
|
msg: &CommunicationValue,
|
||||||
|
) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send {
|
||||||
|
mtp_transport::GenericSender::send(self, msg)
|
||||||
|
}
|
||||||
|
fn finish_stream(
|
||||||
|
&self,
|
||||||
|
) -> impl std::future::Future<Output = Result<(), CommunicationError>> + Send {
|
||||||
|
mtp_transport::GenericSender::finish_stream(self)
|
||||||
|
}
|
||||||
|
fn close(&self) {
|
||||||
|
mtp_transport::GenericSender::close(self);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<C: mtp_transport::TransportConnection> HandshakeReceiver for mtp_transport::GenericReceiver<C> {
|
||||||
|
fn receive(
|
||||||
|
&self,
|
||||||
|
) -> impl std::future::Future<Output = Result<CommunicationValue, CommunicationError>> + Send
|
||||||
|
{
|
||||||
|
mtp_transport::GenericReceiver::receive(self)
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
@ -1,45 +1,16 @@
|
||||||
use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue, Version};
|
use mtp_codec::Version;
|
||||||
use mtp_common::{CommunicationError, RejectionReason};
|
use mtp_common::CommunicationError;
|
||||||
use mtp_transport::Sender;
|
|
||||||
use std::{error::Error, fmt};
|
use std::{error::Error, fmt};
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
use mtp_codec::{CommunicationValue, DataType, DataValue};
|
||||||
|
|
||||||
#[cfg(feature = "crypto")]
|
#[cfg(feature = "crypto")]
|
||||||
pub(crate) fn random_client_id() -> u64 {
|
pub(crate) fn random_client_id() -> u64 {
|
||||||
rand::random::<u64>() & mtp_codec::MAX_WIRE_ID
|
rand::random::<u64>() & mtp_codec::MAX_WIRE_ID
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn send_rejection(sender: &Sender, reason: RejectionReason) {
|
#[cfg(test)]
|
||||||
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<u64>,
|
|
||||||
) -> 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
|
|
||||||
}
|
|
||||||
|
|
||||||
pub(crate) fn extract_version(msg: &CommunicationValue) -> Option<Version> {
|
pub(crate) fn extract_version(msg: &CommunicationValue) -> Option<Version> {
|
||||||
let value = msg.get_data(DataType::Version);
|
let value = msg.get_data(DataType::Version);
|
||||||
match value {
|
match value {
|
||||||
|
|
|
||||||
|
|
@ -1,23 +1,17 @@
|
||||||
|
#[cfg(not(feature = "crypto"))]
|
||||||
|
use mtp_codec::{Version, registry::{Registry, VersionedCodec}};
|
||||||
#[cfg(feature = "crypto")]
|
#[cfg(feature = "crypto")]
|
||||||
use mtp_codec::{CommunicationType, CommunicationValue};
|
use mtp_codec::registry::Registry;
|
||||||
use mtp_codec::{
|
|
||||||
DataType, DataValue, Version,
|
|
||||||
registry::{Registry, VersionedCodec},
|
|
||||||
};
|
|
||||||
use mtp_common::RejectionReason;
|
|
||||||
use mtp_transport::{Receiver, Sender};
|
use mtp_transport::{Receiver, Sender};
|
||||||
use std::sync::Arc;
|
use std::sync::Arc;
|
||||||
use std::time::Instant;
|
use std::time::Instant;
|
||||||
#[cfg(feature = "pipes")]
|
#[cfg(feature = "pipes")]
|
||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
|
|
||||||
#[cfg(feature = "crypto")]
|
|
||||||
use crate::config::AuthenticationPolicy;
|
|
||||||
use crate::config::HostConfig;
|
use crate::config::HostConfig;
|
||||||
use crate::connection::MTPConnection;
|
use crate::connection::MTPConnection;
|
||||||
#[cfg(feature = "crypto")]
|
use crate::engine::HandshakeEngine;
|
||||||
use crate::error::AuthState;
|
use crate::error::AcceptError;
|
||||||
use crate::error::{AcceptError, extract_version, send_accepted, send_rejection};
|
|
||||||
#[cfg(feature = "pipes")]
|
#[cfg(feature = "pipes")]
|
||||||
use crate::pipe::PipeDispatcher;
|
use crate::pipe::PipeDispatcher;
|
||||||
#[cfg(feature = "pipes")]
|
#[cfg(feature = "pipes")]
|
||||||
|
|
@ -137,182 +131,107 @@ impl HandshakeContext {
|
||||||
sender: Sender,
|
sender: Sender,
|
||||||
receiver: Receiver,
|
receiver: Receiver,
|
||||||
) -> Result<Option<MTPConnection>, AcceptError> {
|
) -> Result<Option<MTPConnection>, AcceptError> {
|
||||||
|
let engine = HandshakeEngine::new(self.registry.clone(), self.config.clone());
|
||||||
|
let result = engine.accept(&sender, &receiver).await?;
|
||||||
#[cfg(feature = "crypto")]
|
#[cfg(feature = "crypto")]
|
||||||
{
|
{
|
||||||
return tokio::time::timeout(
|
Ok(Some(self.connection_from_handshake_result(
|
||||||
self.config.auth_timeout,
|
sender,
|
||||||
self.accept_pair(sender, receiver),
|
receiver,
|
||||||
)
|
result,
|
||||||
.await
|
)))
|
||||||
.unwrap_or(Err(AcceptError::AuthenticationTimedOut));
|
|
||||||
}
|
}
|
||||||
#[cfg(not(feature = "crypto"))]
|
|
||||||
self.accept_pair(sender, receiver).await
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn accept_pair(
|
|
||||||
&self,
|
|
||||||
sender: Sender,
|
|
||||||
receiver: Receiver,
|
|
||||||
) -> Result<Option<MTPConnection>, 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"))]
|
#[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(
|
Ok(Some(self.connection_from_parts(
|
||||||
sender,
|
sender,
|
||||||
receiver,
|
receiver,
|
||||||
negotiated,
|
result.negotiated_version,
|
||||||
codec,
|
result.codec,
|
||||||
description,
|
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"))]
|
#[cfg(not(feature = "crypto"))]
|
||||||
#[allow(clippy::too_many_arguments)]
|
#[allow(clippy::too_many_arguments)]
|
||||||
pub(crate) fn connection_from_parts(
|
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<String>,
|
|
||||||
auth_state: AuthState,
|
|
||||||
client_id: u64,
|
|
||||||
client_public_key: Option<mtp_crypto::PublicKeyBundle>,
|
|
||||||
) -> 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<u8>,
|
|
||||||
},
|
|
||||||
}
|
|
||||||
|
|
||||||
#[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::<u64>() & mtp_codec::MAX_WIRE_ID;
|
|
||||||
if (self.config.get_existing_client)(id, None).await.is_none() {
|
|
||||||
return id;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
rand::random::<u64>() & mtp_codec::MAX_WIRE_ID
|
|
||||||
}
|
|
||||||
|
|
||||||
async fn accept_authenticated(
|
|
||||||
&self,
|
|
||||||
sender: Sender,
|
|
||||||
receiver: Receiver,
|
|
||||||
) -> Result<Option<MTPConnection>, 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<String>,
|
|
||||||
) -> Result<Option<MTPConnection>, 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<u8>| 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<u8> = 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<Option<MTPConnection>, 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(),
|
|
||||||
))
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -1,5 +1,6 @@
|
||||||
pub mod config;
|
pub mod config;
|
||||||
pub mod connection;
|
pub mod connection;
|
||||||
|
pub mod engine;
|
||||||
pub mod error;
|
pub mod error;
|
||||||
pub mod handshake;
|
pub mod handshake;
|
||||||
#[cfg(feature = "pipes")]
|
#[cfg(feature = "pipes")]
|
||||||
|
|
@ -10,6 +11,7 @@ pub use MTPHost as Host;
|
||||||
pub use config::HostConfig;
|
pub use config::HostConfig;
|
||||||
pub use config::Policy;
|
pub use config::Policy;
|
||||||
pub use connection::{MTPConnection, MtpReceiverLike, MtpSenderLike};
|
pub use connection::{MTPConnection, MtpReceiverLike, MtpSenderLike};
|
||||||
|
pub use engine::{HandshakeEngine, HandshakeResult, HandshakeReceiver, HandshakeSender};
|
||||||
pub use error::AcceptError;
|
pub use error::AcceptError;
|
||||||
pub use handshake::MTPHost;
|
pub use handshake::MTPHost;
|
||||||
pub use mtp_transport::Receiver;
|
pub use mtp_transport::Receiver;
|
||||||
|
|
|
||||||
|
|
@ -18,6 +18,7 @@ pub(crate) struct DriverConfig {
|
||||||
pub(crate) policy: mtp_transport::Policy,
|
pub(crate) policy: mtp_transport::Policy,
|
||||||
pub(crate) host_config: Arc<HostConfig>,
|
pub(crate) host_config: Arc<HostConfig>,
|
||||||
pub(crate) metrics: Option<Arc<dyn WebServerMetrics>>,
|
pub(crate) metrics: Option<Arc<dyn WebServerMetrics>>,
|
||||||
|
pub(crate) auth_semaphore: Arc<Semaphore>,
|
||||||
}
|
}
|
||||||
|
|
||||||
pub(crate) async fn run_driver(
|
pub(crate) async fn run_driver(
|
||||||
|
|
@ -37,6 +38,7 @@ pub(crate) async fn run_driver(
|
||||||
policy,
|
policy,
|
||||||
host_config,
|
host_config,
|
||||||
metrics,
|
metrics,
|
||||||
|
auth_semaphore,
|
||||||
} = config;
|
} = config;
|
||||||
let mut connection_tasks = tokio::task::JoinSet::new();
|
let mut connection_tasks = tokio::task::JoinSet::new();
|
||||||
loop {
|
loop {
|
||||||
|
|
@ -63,6 +65,7 @@ pub(crate) async fn run_driver(
|
||||||
let mtp_tx = mtp_tx.clone();
|
let mtp_tx = mtp_tx.clone();
|
||||||
let metrics = metrics.clone();
|
let metrics = metrics.clone();
|
||||||
let host_config = host_config.clone();
|
let host_config = host_config.clone();
|
||||||
|
let auth_semaphore = auth_semaphore.clone();
|
||||||
connection_tasks.spawn(async move {
|
connection_tasks.spawn(async move {
|
||||||
let _permit = permit;
|
let _permit = permit;
|
||||||
let connect_start = std::time::Instant::now();
|
let connect_start = std::time::Instant::now();
|
||||||
|
|
@ -137,10 +140,6 @@ pub(crate) async fn run_driver(
|
||||||
return;
|
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(
|
tokio::spawn(run_session_requests(
|
||||||
session.clone(),
|
session.clone(),
|
||||||
router.clone(),
|
router.clone(),
|
||||||
|
|
@ -149,26 +148,29 @@ pub(crate) async fn run_driver(
|
||||||
metrics.clone(),
|
metrics.clone(),
|
||||||
remote_addr,
|
remote_addr,
|
||||||
));
|
));
|
||||||
let result =
|
let mtp_tx = mtp_tx.clone();
|
||||||
accept_web_connection(session, mtp_path, connection, send_pongs, policy, host_config.clone())
|
let auth_semaphore = auth_semaphore.clone();
|
||||||
.await;
|
let host_config = host_config.clone();
|
||||||
match mtp_tx.try_send(result) {
|
let connection = connection.clone();
|
||||||
Ok(()) => {
|
tokio::spawn(async move {
|
||||||
// The detached session driver remains active while the
|
let result =
|
||||||
// delivered MTP connection keeps the session alive.
|
accept_web_connection(session, mtp_path, connection, send_pongs, policy, host_config, auth_semaphore)
|
||||||
}
|
.await;
|
||||||
Err(tokio::sync::mpsc::error::TrySendError::Full(result)) => {
|
match mtp_tx.try_send(result) {
|
||||||
tracing::warn!("MTP connection backlog is full; dropping connection");
|
Ok(()) => {}
|
||||||
if let Ok(connection) = result {
|
Err(tokio::sync::mpsc::error::TrySendError::Full(result)) => {
|
||||||
connection.sender.close();
|
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;
|
return;
|
||||||
}
|
}
|
||||||
let router = router.clone();
|
let router = router.clone();
|
||||||
|
|
|
||||||
|
|
@ -228,6 +228,7 @@ impl MTPWebServer {
|
||||||
let (mtp_tx, mtp_incoming) = tokio::sync::mpsc::channel(web_config.max_connections.max(1));
|
let (mtp_tx, mtp_incoming) = tokio::sync::mpsc::channel(web_config.max_connections.max(1));
|
||||||
let (shutdown_tx, shutdown_rx) = watch::channel(());
|
let (shutdown_tx, shutdown_rx) = watch::channel(());
|
||||||
let connection_semaphore = Arc::new(Semaphore::new(web_config.max_connections));
|
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 router = web_config.router.clone();
|
||||||
let metrics = web_config.metrics.clone();
|
let metrics = web_config.metrics.clone();
|
||||||
let driver_config = DriverConfig {
|
let driver_config = DriverConfig {
|
||||||
|
|
@ -240,6 +241,7 @@ impl MTPWebServer {
|
||||||
policy: host_config.policy,
|
policy: host_config.policy,
|
||||||
host_config,
|
host_config,
|
||||||
metrics: web_config.metrics.clone(),
|
metrics: web_config.metrics.clone(),
|
||||||
|
auth_semaphore,
|
||||||
};
|
};
|
||||||
let quic_driver = tokio::spawn(run_driver(
|
let quic_driver = tokio::spawn(run_driver(
|
||||||
driver_endpoint,
|
driver_endpoint,
|
||||||
|
|
@ -339,10 +341,23 @@ fn build_endpoint(
|
||||||
.map_err(|_| CommunicationError::CertificateLoadFailed)?;
|
.map_err(|_| CommunicationError::CertificateLoadFailed)?;
|
||||||
tls.alpn_protocols = vec![b"h3".to_vec()];
|
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)
|
quinn::crypto::rustls::QuicServerConfig::try_from(tls)
|
||||||
.map_err(|error| CommunicationError::Other(error.to_string()))?,
|
.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))
|
quinn::Endpoint::server(server, SocketAddr::new(config.ip, config.port))
|
||||||
.map_err(|error| CommunicationError::Other(error.to_string()))
|
.map_err(|error| CommunicationError::Other(error.to_string()))
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -16,6 +16,9 @@ use std::time::Instant;
|
||||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||||
use tracing::error;
|
use tracing::error;
|
||||||
|
|
||||||
|
#[cfg(feature = "crypto")]
|
||||||
|
const GUEST_ID_MAX_RETRIES: u32 = 100;
|
||||||
|
|
||||||
type Session = h3_webtransport::server::WebTransportSession<h3_quinn::Connection, Bytes>;
|
type Session = h3_webtransport::server::WebTransportSession<h3_quinn::Connection, Bytes>;
|
||||||
type H3SendStream = h3_webtransport::stream::SendStream<h3_quinn::SendStream<Bytes>, Bytes>;
|
type H3SendStream = h3_webtransport::stream::SendStream<h3_quinn::SendStream<Bytes>, Bytes>;
|
||||||
type H3RecvStream = h3_webtransport::stream::RecvStream<h3_quinn::RecvStream, Bytes>;
|
type H3RecvStream = h3_webtransport::stream::RecvStream<h3_quinn::RecvStream, Bytes>;
|
||||||
|
|
@ -225,6 +228,38 @@ pub type WebMtpReceiver = GenericReceiver<H3TransportConnection>;
|
||||||
pub type WebMTPConnection =
|
pub type WebMTPConnection =
|
||||||
mtp_host::MTPConnection<WebMtpSender, WebMtpReceiver, H3TransportReceiver>;
|
mtp_host::MTPConnection<WebMtpSender, WebMtpReceiver, H3TransportReceiver>;
|
||||||
|
|
||||||
|
#[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<u64, AcceptError> {
|
||||||
|
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::<u64>() & 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(
|
pub(crate) async fn accept_web_connection(
|
||||||
session: Arc<Session>,
|
session: Arc<Session>,
|
||||||
path: String,
|
path: String,
|
||||||
|
|
@ -232,15 +267,25 @@ pub(crate) async fn accept_web_connection(
|
||||||
send_pongs: bool,
|
send_pongs: bool,
|
||||||
policy: Policy,
|
policy: Policy,
|
||||||
host_config: Arc<HostConfig>,
|
host_config: Arc<HostConfig>,
|
||||||
|
auth_semaphore: Arc<tokio::sync::Semaphore>,
|
||||||
) -> Result<WebMTPConnection, AcceptError> {
|
) -> Result<WebMTPConnection, AcceptError> {
|
||||||
#[cfg(feature = "crypto")]
|
#[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,
|
host_config.auth_timeout,
|
||||||
accept_web_connection_inner(session, path, quinn, send_pongs, policy, host_config),
|
accept_web_connection_inner(session, path, quinn, send_pongs, policy, host_config),
|
||||||
)
|
)
|
||||||
.await
|
.await
|
||||||
.unwrap_or(Err(AcceptError::AuthenticationTimedOut))
|
.unwrap_or(Err(AcceptError::AuthenticationTimedOut));
|
||||||
|
drop(permit);
|
||||||
|
result
|
||||||
}
|
}
|
||||||
|
|
||||||
#[cfg(not(feature = "crypto"))]
|
#[cfg(not(feature = "crypto"))]
|
||||||
|
|
@ -278,12 +323,12 @@ async fn accept_web_connection_inner(
|
||||||
DataValue::Str(value) => Some(value.clone()),
|
DataValue::Str(value) => Some(value.clone()),
|
||||||
_ => None,
|
_ => None,
|
||||||
};
|
};
|
||||||
let sender = WebMtpSender::new(transport, policy);
|
let sender = WebMtpSender::new(transport, policy.clone());
|
||||||
if send_pongs {
|
if send_pongs {
|
||||||
receiver.respond_to_pings(sender.clone()).await;
|
receiver.respond_to_pings(sender.clone()).await;
|
||||||
}
|
}
|
||||||
let connection: WebMTPConnection =
|
let connection: WebMTPConnection =
|
||||||
mtp_host::MTPConnection::from_transport_parts_with_remote_addr(
|
mtp_host::MTPConnection::from_transport_parts_with_policy(
|
||||||
negotiated,
|
negotiated,
|
||||||
codec,
|
codec,
|
||||||
sender,
|
sender,
|
||||||
|
|
@ -291,23 +336,101 @@ async fn accept_web_connection_inner(
|
||||||
path,
|
path,
|
||||||
description.clone(),
|
description.clone(),
|
||||||
Some(remote_addr),
|
Some(remote_addr),
|
||||||
|
policy,
|
||||||
);
|
);
|
||||||
|
#[cfg(not(feature = "crypto"))]
|
||||||
|
{
|
||||||
|
connection.receiver.set_max_message_size(max_message_size);
|
||||||
|
return Ok(connection);
|
||||||
|
}
|
||||||
#[cfg(feature = "crypto")]
|
#[cfg(feature = "crypto")]
|
||||||
let mut connection = connection;
|
let mut connection = connection;
|
||||||
#[cfg(feature = "crypto")]
|
#[cfg(feature = "crypto")]
|
||||||
if !matches!(
|
{
|
||||||
_host_config.authentication_policy,
|
use mtp_codec::{CommunicationType, DataType, DataValue};
|
||||||
mtp_host::AuthenticationPolicy::Unauthenticated
|
|
||||||
) {
|
|
||||||
use mtp_crypto::{
|
use mtp_crypto::{
|
||||||
Ed25519Signer, MlDsaSigner, PublicKeyBundle, SignatureScheme, auth, verify_ed25519,
|
Ed25519Signer, MlDsaSigner, PublicKeyBundle, SignatureScheme, auth, verify_ed25519,
|
||||||
verify_ml_dsa,
|
|
||||||
};
|
};
|
||||||
|
|
||||||
let tm = mtp_codec::TypeMap::latest();
|
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_lookup_started = Instant::now();
|
||||||
let (client_id, client_bundle, response_type) = if Some(first.get_type())
|
let first_type = first.get_type();
|
||||||
== mtp_codec::CommunicationType::Identification.try_to_id(&tm)
|
|
||||||
|
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) {
|
let id = match first.get_data(DataType::Id) {
|
||||||
DataValue::UnsignedNumber(value) => *value as u64,
|
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())
|
let bundle = (_host_config.get_existing_client)(id, description.clone())
|
||||||
.await
|
.await
|
||||||
.ok_or_else(|| AcceptError::AuthenticationFailed("unknown client id".into()))?;
|
.ok_or_else(|| AcceptError::AuthenticationFailed("unknown client id".into()))?;
|
||||||
(
|
(id, Some(bundle), CommunicationType::IdentificationResponse, false)
|
||||||
id,
|
|
||||||
bundle,
|
|
||||||
mtp_codec::CommunicationType::IdentificationResponse,
|
|
||||||
)
|
|
||||||
} else if Some(first.get_type()) == mtp_codec::CommunicationType::Register.try_to_id(&tm) {
|
|
||||||
let bundle = match first.get_data(DataType::PublicKeys) {
|
|
||||||
DataValue::Bytes(bytes) => PublicKeyBundle::from_bytes(bytes).map_err(|_| {
|
|
||||||
AcceptError::AuthenticationFailed("invalid public key bundle".into())
|
|
||||||
})?,
|
|
||||||
_ => {
|
|
||||||
return Err(AcceptError::AuthenticationFailed(
|
|
||||||
"missing public keys".into(),
|
|
||||||
));
|
|
||||||
}
|
|
||||||
};
|
|
||||||
(0, bundle, mtp_codec::CommunicationType::RegisterResponse)
|
|
||||||
} else {
|
} else {
|
||||||
return Err(AcceptError::AuthenticationFailed(
|
return Err(AcceptError::AuthenticationFailed(
|
||||||
"unexpected authentication message".into(),
|
"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();
|
// Guest path: skip challenge/response, send accepted with guest ID
|
||||||
let host_pq_signer = if !_host_config
|
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
|
.host_keyring
|
||||||
.sig_pq_secret_key
|
.sig_pq_secret_key
|
||||||
.as_bytes()
|
.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(
|
MlDsaSigner::new(
|
||||||
&_host_config.host_keyring.sig_pq_secret_key,
|
&_host_config.host_keyring.sig_pq_secret_key,
|
||||||
&_host_config.host_keyring.sig_pq_public_key,
|
&_host_config.host_keyring.sig_pq_public_key,
|
||||||
)
|
)
|
||||||
.map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?,
|
.map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?,
|
||||||
)
|
))
|
||||||
} else {
|
} else {
|
||||||
None
|
None
|
||||||
};
|
};
|
||||||
tracing::debug!(elapsed = ?signer_init_started.elapsed(), "web authentication handshake: signer initialization");
|
tracing::debug!(elapsed = ?signer_init_started.elapsed(), "web authentication handshake: signer initialization");
|
||||||
|
|
||||||
let server_challenge: u128 = rand::random();
|
let server_challenge: u128 = rand::random();
|
||||||
let host_sign = |payload: &[u8]| -> Result<(Vec<u8>, Vec<u8>), AcceptError> {
|
let host_sign = |payload: Vec<u8>| {
|
||||||
let signer = Ed25519Signer::new(&_host_config.host_keyring.sig_cl_secret_key)
|
let host_config = _host_config.clone();
|
||||||
.map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?;
|
let pq_signer = host_pq_signer.clone();
|
||||||
let sig = signer
|
async move {
|
||||||
.sign(payload)
|
let signer =
|
||||||
.map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?;
|
Ed25519Signer::new(&host_config.host_keyring.sig_cl_secret_key)
|
||||||
let pq = if let Some(pq_signer) = host_pq_signer.as_ref() {
|
.map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?;
|
||||||
pq_signer
|
if let Some(pq_signer) = pq_signer {
|
||||||
.sign(payload)
|
mtp_crypto::sign_parallel::sign_dual_parallel_shared_pq(
|
||||||
.map_err(|e| AcceptError::AuthenticationFailed(e.to_string()))?
|
signer,
|
||||||
} else {
|
pq_signer,
|
||||||
Vec::new()
|
payload,
|
||||||
};
|
)
|
||||||
Ok((sig, pq))
|
.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 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");
|
tracing::debug!(elapsed = ?sign_challenge_started.elapsed(), "web authentication handshake: sign challenge");
|
||||||
let mut challenge =
|
let mut challenge =
|
||||||
mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Challenge)
|
mtp_codec::CommunicationValue::new(CommunicationType::Challenge)
|
||||||
.add_typed_default(
|
.add_typed_default(
|
||||||
DataType::ServerNonce,
|
DataType::ServerNonce,
|
||||||
DataValue::UnsignedNumber(server_challenge),
|
DataValue::UnsignedNumber(server_challenge),
|
||||||
|
|
@ -396,7 +550,7 @@ async fn accept_web_connection_inner(
|
||||||
DataValue::BoolFalse
|
DataValue::BoolFalse
|
||||||
},
|
},
|
||||||
);
|
);
|
||||||
if !pq_sig.is_empty() {
|
if pq_enabled {
|
||||||
challenge =
|
challenge =
|
||||||
challenge.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_sig));
|
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");
|
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(
|
return Err(AcceptError::AuthenticationFailed(
|
||||||
"missing challenge response".into(),
|
"missing challenge response".into(),
|
||||||
|
|
@ -452,7 +606,7 @@ async fn accept_web_connection_inner(
|
||||||
_ => &[],
|
_ => &[],
|
||||||
};
|
};
|
||||||
let payload = if first.get_type()
|
let payload = if first.get_type()
|
||||||
== mtp_codec::CommunicationType::Register
|
== CommunicationType::Register
|
||||||
.try_to_id(&tm)
|
.try_to_id(&tm)
|
||||||
.unwrap()
|
.unwrap()
|
||||||
{
|
{
|
||||||
|
|
@ -465,36 +619,54 @@ async fn accept_web_connection_inner(
|
||||||
} else {
|
} else {
|
||||||
auth::login_proof_payload(&version.to_string(), client_id, server_challenge, nonce)
|
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();
|
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(
|
return Err(AcceptError::AuthenticationFailed(
|
||||||
"client proof signature invalid".into(),
|
"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 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
|
(_host_config.complete_register)(client_bundle.clone(), description.clone()).await
|
||||||
} else {
|
} else {
|
||||||
client_id
|
client_id
|
||||||
};
|
};
|
||||||
tracing::debug!(elapsed = ?register_started.elapsed(), "web authentication handshake: registration callback");
|
tracing::debug!(elapsed = ?register_started.elapsed(), "web authentication handshake: registration callback");
|
||||||
let sign_final_started = Instant::now();
|
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,
|
assigned_id,
|
||||||
nonce,
|
nonce,
|
||||||
server_challenge,
|
server_challenge,
|
||||||
))?;
|
))
|
||||||
|
.await?;
|
||||||
tracing::debug!(elapsed = ?sign_final_started.elapsed(), "web authentication handshake: sign final response");
|
tracing::debug!(elapsed = ?sign_final_started.elapsed(), "web authentication handshake: sign final response");
|
||||||
let mut response = mtp_codec::CommunicationValue::new(response_type)
|
let mut response = mtp_codec::CommunicationValue::new(response_type)
|
||||||
.add_typed_default(DataType::Connected, DataValue::BoolTrue)
|
.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::ClientNonce, DataValue::UnsignedNumber(nonce))
|
||||||
.add_typed_default(DataType::Signature, DataValue::Bytes(final_sig))
|
.add_typed_default(DataType::Signature, DataValue::Bytes(final_sig))
|
||||||
.add_typed_default(DataType::Version, DataValue::Str(version.to_string()));
|
.add_typed_default(DataType::Version, DataValue::Str(version.to_string()));
|
||||||
if !final_pq.is_empty() {
|
if pq_enabled {
|
||||||
response =
|
response =
|
||||||
response.add_typed_default(DataType::PqSignature, DataValue::Bytes(final_pq));
|
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.auth_state = mtp_host::AuthState::Authenticated;
|
||||||
connection.client_id = assigned_id;
|
connection.client_id = assigned_id;
|
||||||
connection.client_public_key = Some(client_bundle);
|
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)
|
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -57,7 +57,7 @@ impl Default for Policy {
|
||||||
application_close_code: 0,
|
application_close_code: 0,
|
||||||
open_stream_timeout: Duration::from_millis(2_000),
|
open_stream_timeout: Duration::from_millis(2_000),
|
||||||
write_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),
|
read_timeout: Duration::from_millis(30_000),
|
||||||
keep_alive_interval: Some(Duration::from_secs(3)),
|
keep_alive_interval: Some(Duration::from_secs(3)),
|
||||||
max_idle_timeout: Some(Duration::from_secs(30)),
|
max_idle_timeout: Some(Duration::from_secs(30)),
|
||||||
|
|
@ -1212,28 +1212,6 @@ mod tests {
|
||||||
assert_ne!(SendMode::PersistentStream, SendMode::SingleStreamPerMessage);
|
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]
|
#[test]
|
||||||
fn test_policy_clone() {
|
fn test_policy_clone() {
|
||||||
let p = Policy::default();
|
let p = Policy::default();
|
||||||
|
|
|
||||||
|
|
@ -166,6 +166,7 @@ pub struct GenericReceiver<C: TransportConnection> {
|
||||||
connection: C,
|
connection: C,
|
||||||
ping_sender: Arc<RwLock<Option<GenericSender<C>>>>,
|
ping_sender: Arc<RwLock<Option<GenericSender<C>>>>,
|
||||||
max_message_size: Arc<AtomicU64>,
|
max_message_size: Arc<AtomicU64>,
|
||||||
|
_accept_task: Arc<tokio::task::JoinHandle<()>>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl<C: TransportConnection> Clone for GenericReceiver<C> {
|
impl<C: TransportConnection> Clone for GenericReceiver<C> {
|
||||||
|
|
@ -177,6 +178,15 @@ impl<C: TransportConnection> Clone for GenericReceiver<C> {
|
||||||
connection: self.connection.clone(),
|
connection: self.connection.clone(),
|
||||||
ping_sender: self.ping_sender.clone(),
|
ping_sender: self.ping_sender.clone(),
|
||||||
max_message_size: self.max_message_size.clone(),
|
max_message_size: self.max_message_size.clone(),
|
||||||
|
_accept_task: self._accept_task.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl<C: TransportConnection> Drop for GenericReceiver<C> {
|
||||||
|
fn drop(&mut self) {
|
||||||
|
if Arc::strong_count(&self._accept_task) == 1 {
|
||||||
|
self._accept_task.abort();
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -187,17 +197,34 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
||||||
#[cfg(feature = "pipes")]
|
#[cfg(feature = "pipes")]
|
||||||
let (pipe_tx, pipe_rx) = mpsc::channel(policy.receiver_queue_capacity);
|
let (pipe_tx, pipe_rx) = mpsc::channel(policy.receiver_queue_capacity);
|
||||||
let ping_sender: Arc<RwLock<Option<GenericSender<C>>>> = Arc::new(RwLock::new(None));
|
let ping_sender: Arc<RwLock<Option<GenericSender<C>>>> = Arc::new(RwLock::new(None));
|
||||||
let max_message_size = Arc::new(AtomicU64::new(policy.handshake_max_message_size));
|
let 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_ping_sender = ping_sender.clone();
|
||||||
let task_connection = connection.clone();
|
let task_connection = connection.clone();
|
||||||
let task_policy = policy.clone();
|
let task_policy = policy.clone();
|
||||||
let task_max_message_size = max_message_size.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(
|
let limit = Arc::new(Semaphore::new(
|
||||||
task_policy.max_concurrent_stream_tasks.max(1),
|
task_policy.max_concurrent_stream_tasks.max(1),
|
||||||
));
|
));
|
||||||
loop {
|
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_policy.accept_stream_timeout,
|
||||||
task_connection.accept_uni(),
|
task_connection.accept_uni(),
|
||||||
)
|
)
|
||||||
|
|
@ -205,29 +232,34 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
||||||
{
|
{
|
||||||
Ok(Ok(stream)) => stream,
|
Ok(Ok(stream)) => stream,
|
||||||
Ok(Err(error)) => {
|
Ok(Err(error)) => {
|
||||||
let _ = tx.send(Err(error)).await;
|
let _ = task_accept_task_tx.send(Err(error)).await;
|
||||||
break;
|
break;
|
||||||
}
|
}
|
||||||
Err(_) => {
|
Err(_) => {
|
||||||
if task_connection.close_reason().is_some() {
|
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;
|
break;
|
||||||
} else {
|
} else {
|
||||||
continue;
|
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")]
|
#[cfg(feature = "pipes")]
|
||||||
let pipe_tx = pipe_tx.clone();
|
let pipe_tx = task_accept_task_pipe_tx.clone();
|
||||||
let policy = task_policy.clone();
|
let policy = task_policy.clone();
|
||||||
let max_message_size = task_max_message_size.clone();
|
let max_message_size = task_max_message_size.clone();
|
||||||
let permit = limit.clone();
|
|
||||||
let ping_sender = task_ping_sender.clone();
|
let ping_sender = task_ping_sender.clone();
|
||||||
|
let connection = task_connection.clone();
|
||||||
tokio::spawn(async move {
|
tokio::spawn(async move {
|
||||||
let Ok(_permit) = permit.acquire_owned().await else {
|
let _permit = permit;
|
||||||
return;
|
|
||||||
};
|
|
||||||
let mut stream = stream;
|
let mut stream = stream;
|
||||||
let mut frames = 0usize;
|
let mut frames = 0usize;
|
||||||
loop {
|
loop {
|
||||||
|
|
@ -235,24 +267,36 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
||||||
.max_frames_per_stream
|
.max_frames_per_stream
|
||||||
.is_some_and(|max| frames >= max)
|
.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;
|
break;
|
||||||
}
|
}
|
||||||
let mut len = [0; 4];
|
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(Ok(())) => {}
|
||||||
Ok(Err(CommunicationError::StreamClosed)) => break,
|
Ok(Err(CommunicationError::StreamClosed)) => break,
|
||||||
Ok(Err(error)) => {
|
Ok(Err(error)) => {
|
||||||
tracing::error!(
|
|
||||||
"[mtp-transport] frame header read failed: {error}"
|
|
||||||
);
|
|
||||||
tracing::warn!(%error, "MTP receive stream failed while reading frame header");
|
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;
|
break;
|
||||||
}
|
}
|
||||||
Err(error) => {
|
Err(_) => {
|
||||||
tracing::error!(
|
tracing::warn!("MTP receive stream timed out while reading frame header");
|
||||||
"[mtp-transport] frame header read timed out: {error}"
|
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;
|
break;
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -260,8 +304,14 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
||||||
if len == policy.close_frame_len {
|
if len == policy.close_frame_len {
|
||||||
break;
|
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");
|
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;
|
break;
|
||||||
}
|
}
|
||||||
let target_len = len as usize;
|
let target_len = len as usize;
|
||||||
|
|
@ -271,12 +321,17 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
||||||
target_len,
|
target_len,
|
||||||
"MTP receive stream could not reserve frame body"
|
"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;
|
break;
|
||||||
}
|
}
|
||||||
while body.len() < target_len {
|
while body.len() < target_len {
|
||||||
let chunk_len = (target_len - body.len()).min(16 * 1024);
|
let chunk_len = (target_len - body.len()).min(16 * 1024);
|
||||||
let mut chunk = [0u8; 16 * 1024];
|
let mut chunk = [0u8; 16 * 1024];
|
||||||
let body_read = timeout(
|
let body_read = tokio::time::timeout(
|
||||||
policy.read_timeout,
|
policy.read_timeout,
|
||||||
stream.read_exact(&mut chunk[..chunk_len]),
|
stream.read_exact(&mut chunk[..chunk_len]),
|
||||||
)
|
)
|
||||||
|
|
@ -284,16 +339,16 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
||||||
if !matches!(&body_read, Ok(Ok(())))
|
if !matches!(&body_read, Ok(Ok(())))
|
||||||
|| body.try_reserve(chunk_len).is_err()
|
|| body.try_reserve(chunk_len).is_err()
|
||||||
{
|
{
|
||||||
tracing::error!(
|
|
||||||
"[mtp-transport] frame body read failed ({} bytes): {:?}",
|
|
||||||
chunk_len,
|
|
||||||
body_read
|
|
||||||
);
|
|
||||||
tracing::warn!(
|
tracing::warn!(
|
||||||
pipe_chunk_len = chunk_len,
|
pipe_chunk_len = chunk_len,
|
||||||
?body_read,
|
?body_read,
|
||||||
"MTP receive stream failed while reading frame body"
|
"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;
|
break;
|
||||||
}
|
}
|
||||||
body.extend_from_slice(&chunk[..chunk_len]);
|
body.extend_from_slice(&chunk[..chunk_len]);
|
||||||
|
|
@ -306,6 +361,12 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
||||||
Ok(message) => message,
|
Ok(message) => message,
|
||||||
Err(_) => {
|
Err(_) => {
|
||||||
tracing::warn!("MTP receive stream contained an invalid frame");
|
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;
|
break;
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
@ -367,6 +428,7 @@ impl<C: TransportConnection> GenericReceiver<C> {
|
||||||
connection,
|
connection,
|
||||||
ping_sender,
|
ping_sender,
|
||||||
max_message_size,
|
max_message_size,
|
||||||
|
_accept_task: Arc::new(accept_task),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
pub async fn respond_to_pings(&self, sender: GenericSender<C>) {
|
pub async fn respond_to_pings(&self, sender: GenericSender<C>) {
|
||||||
|
|
|
||||||
|
|
@ -126,31 +126,34 @@ pub async fn host_with_config(
|
||||||
let incoming_session = endpoint.accept().await;
|
let incoming_session = endpoint.accept().await;
|
||||||
tracing::debug!(elapsed = ?accept_started.elapsed(), "host accept loop: received QUIC connection");
|
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();
|
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;
|
||||||
|
});
|
||||||
}
|
}
|
||||||
});
|
});
|
||||||
|
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue