From 6ef1293603049dde289b4cb04d4f75e45d0a0def Mon Sep 17 00:00:00 2001 From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com> Date: Sun, 28 Jun 2026 03:26:07 +0200 Subject: [PATCH] [Add] Ease of use functions --- Cargo.lock | 2 + client/Cargo.toml | 1 + client/src/lib.rs | 204 +++++++++++++++----- codec/src/communication_value.rs | 220 +++++++++++++++++++-- codec/src/data_value.rs | 309 ++++++++++++++++++++++++++++++ common/src/lib.rs | 2 + crypto/src/keypair.rs | 48 +++++ example/server/src/handlers.rs | 20 +- host/Cargo.toml | 1 + host/src/lib.rs | 149 +++++++++++---- transport/tests/integration.rs | 2 +- type-map/build.rs | 319 +++++++++++++++++++++++++++++++ type-map/src/lib.rs | 10 + wasm/src/client.rs | 20 +- wasm/src/frame.rs | 20 +- 15 files changed, 1203 insertions(+), 124 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 02d8a70..2e208b1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -999,6 +999,7 @@ dependencies = [ "mtp-crypto", "mtp-transport", "rand 0.8.6", + "tokio", ] [[package]] @@ -1049,6 +1050,7 @@ dependencies = [ "mtp-crypto", "mtp-transport", "rand 0.8.6", + "tokio", ] [[package]] diff --git a/client/Cargo.toml b/client/Cargo.toml index 561715b..df69b8c 100644 --- a/client/Cargo.toml +++ b/client/Cargo.toml @@ -9,6 +9,7 @@ mtp-codec = { path = "../codec" } mtp-transport = { path = "../transport" } mtp-crypto = { path = "../crypto", optional = true } rand = "0.8" +tokio = { version = "1", features = ["time"] } [features] crypto = ["dep:mtp-crypto", "mtp-codec/crypto"] diff --git a/client/src/lib.rs b/client/src/lib.rs index e679df2..ac063d7 100644 --- a/client/src/lib.rs +++ b/client/src/lib.rs @@ -1,6 +1,7 @@ use mtp_codec::{CommunicationValue, DataType, DataValue, PROTOCOL_VERSION, Version}; use mtp_common::CommunicationError; use mtp_transport::{Policy, Receiver, Sender}; +use tokio::time::Duration; #[cfg(feature = "crypto")] fn unexpected_response_type_error( @@ -20,6 +21,8 @@ pub struct ClientConfig { pub url: String, pub tls: ClientTlsConfig, pub client_id: u64, + #[cfg(feature = "crypto")] + pub auth_timeout: Duration, } #[derive(Debug, Clone, PartialEq, Eq)] @@ -34,6 +37,8 @@ impl ClientConfig { url: url.into(), tls: ClientTlsConfig::SystemRoots, client_id: 0, + #[cfg(feature = "crypto")] + auth_timeout: Duration::from_secs(30), } } @@ -51,6 +56,12 @@ impl ClientConfig { self } + #[cfg(feature = "crypto")] + pub fn with_auth_timeout(mut self, timeout: Duration) -> Self { + self.auth_timeout = timeout; + self + } + fn server_cert(&self) -> Option> { match &self.tls { ClientTlsConfig::SystemRoots => None, @@ -73,6 +84,7 @@ pub struct MTPConnection { #[cfg(feature = "crypto")] #[derive(Debug, Clone, PartialEq, Eq)] pub enum AuthState { + Unauthenticated, Pending, Authenticated, Failed, @@ -91,7 +103,6 @@ impl MTPClient { let (sender, receiver) = mtp_transport::connect(&config.url, config.server_cert(), Policy::default()).await?; - // Build the initial identification message with the protocol version. let version_str = format!("{}", PROTOCOL_VERSION); let ident = CommunicationValue::new(mtp_codec::CommunicationType::Identification) .add_typed_default(DataType::Version, DataValue::Str(version_str)) @@ -107,7 +118,7 @@ impl MTPClient { sender, receiver, #[cfg(feature = "crypto")] - auth_state: AuthState::Authenticated, + auth_state: AuthState::Unauthenticated, #[cfg(feature = "crypto")] client_id: config.client_id, }) @@ -126,14 +137,13 @@ impl MTPClient { #[cfg(feature = "crypto")] fn verify_host_challenge( challenge: &CommunicationValue, - tm: &mtp_codec::TypeMap, host_pk: &mtp_crypto::PublicKeyBundle, id: u64, server_challenge: u128, ) -> Result<(), CommunicationError> { use mtp_crypto::{auth, verify_ed25519, verify_ml_dsa}; - let sig = match challenge.get_data(DataType::Signature.to_id(tm)) { + let sig = match challenge.get_data(DataType::Signature) { DataValue::Bytes(b) => b.clone(), _ => { return Err(CommunicationError::AuthenticationFailed( @@ -141,7 +151,7 @@ fn verify_host_challenge( )); } }; - let pq_sig = match challenge.get_data(DataType::PqSignature.to_id(tm)) { + let pq_sig = match challenge.get_data(DataType::PqSignature) { DataValue::Bytes(b) => b.clone(), _ => vec![], }; @@ -165,7 +175,6 @@ fn verify_host_challenge( #[cfg(feature = "crypto")] fn verify_host_final( response: &CommunicationValue, - tm: &mtp_codec::TypeMap, host_pk: &mtp_crypto::PublicKeyBundle, id: u64, client_nonce: u128, @@ -173,7 +182,7 @@ fn verify_host_final( ) -> Result<(), CommunicationError> { use mtp_crypto::{auth, verify_ed25519, verify_ml_dsa}; - match response.get_data(DataType::ClientNonce.to_id(tm)) { + match response.get_data(DataType::ClientNonce) { DataValue::UnsignedNumber(n) if *n == client_nonce => {} _ => { return Err(CommunicationError::AuthenticationFailed( @@ -182,7 +191,7 @@ fn verify_host_final( } } - let sig = match response.get_data(DataType::Signature.to_id(tm)) { + let sig = match response.get_data(DataType::Signature) { DataValue::Bytes(b) => b.clone(), _ => { return Err(CommunicationError::AuthenticationFailed( @@ -190,7 +199,7 @@ fn verify_host_final( )); } }; - let pq_sig = match response.get_data(DataType::PqSignature.to_id(tm)) { + let pq_sig = match response.get_data(DataType::PqSignature) { DataValue::Bytes(b) => b.clone(), _ => vec![], }; @@ -210,10 +219,9 @@ fn verify_host_final( #[cfg(feature = "crypto")] fn check_connected( response: &CommunicationValue, - tm: &mtp_codec::TypeMap, reject_msg: &str, ) -> Result<(), CommunicationError> { - match response.get_data(DataType::Connected.to_id(tm)) { + match response.get_data(DataType::Connected) { DataValue::BoolTrue => Ok(()), DataValue::BoolFalse => Err(CommunicationError::AuthenticationFailed(reject_msg.into())), _ => Err(CommunicationError::AuthenticationFailed( @@ -271,7 +279,7 @@ async fn receive_verified_challenge( )); } - let server_challenge = match challenge.get_data(DataType::ServerNonce.to_id(tm)) { + let server_challenge = match challenge.get_data(DataType::ServerNonce) { DataValue::UnsignedNumber(n) => *n, _ => { return Err(CommunicationError::AuthenticationFailed( @@ -280,13 +288,7 @@ async fn receive_verified_challenge( } }; - verify_host_challenge( - &challenge, - tm, - host_public_key_bundle, - bound_id, - server_challenge, - )?; + verify_host_challenge(&challenge, host_public_key_bundle, bound_id, server_challenge)?; Ok(server_challenge) } @@ -297,6 +299,25 @@ impl MTPClient { config: ClientConfig, keys: &mtp_crypto::Keyring, host_public_key_bundle: &mtp_crypto::PublicKeyBundle, + ) -> Result { + let timeout = config.auth_timeout; + match tokio::time::timeout( + timeout, + Self::auth_connect_inner(config, keys, host_public_key_bundle), + ) + .await + { + Ok(result) => result, + Err(_) => Err(CommunicationError::AuthenticationFailed( + "authentication timed out".into(), + )), + } + } + + async fn auth_connect_inner( + config: ClientConfig, + keys: &mtp_crypto::Keyring, + host_public_key_bundle: &mtp_crypto::PublicKeyBundle, ) -> Result { use mtp_crypto::auth; @@ -313,17 +334,27 @@ impl MTPClient { DataType::Id, DataValue::UnsignedNumber(config.client_id as u128), ); - sender.send(&ident).await?; + if let Err(e) = sender.send(&ident).await { + sender.close(); + return Err(e); + } // 2. Receive and verify the host's challenge. - let server_challenge = receive_verified_challenge( + let server_challenge = match receive_verified_challenge( &receiver, &tm, host_public_key_bundle, config.client_id, "auth_connect challenge", ) - .await?; + .await + { + Ok(c) => c, + Err(e) => { + sender.close(); + return Err(e); + } + }; // 3. Sign the host's challenge and send the proof. let client_nonce: u128 = rand::random(); @@ -334,28 +365,49 @@ impl MTPClient { client_nonce, ); - let proof = signed_challenge_response(keys, &proof_payload, client_nonce)?; - sender.send(&proof).await?; + let proof = match signed_challenge_response(keys, &proof_payload, client_nonce) { + Ok(p) => p, + Err(e) => { + sender.close(); + return Err(e); + } + }; + if let Err(e) = sender.send(&proof).await { + sender.close(); + return Err(e); + } // 4. Receive and verify the host's final confirmation. - let response = receiver.receive().await?; + let response = match receiver.receive().await { + Ok(r) => r, + Err(e) => { + sender.close(); + return Err(e); + } + }; let expected_type = mtp_codec::CommunicationType::IdentificationResponse.to_id(&tm); if response.get_type() != expected_type { + sender.close(); return Err(unexpected_response_type_error( "auth_connect", expected_type, &response, )); } - check_connected(&response, &tm, "Server rejected authentication")?; - verify_host_final( + if let Err(e) = check_connected(&response, "Server rejected authentication") { + sender.close(); + return Err(e); + } + if let Err(e) = verify_host_final( &response, - &tm, host_public_key_bundle, config.client_id, client_nonce, server_challenge, - )?; + ) { + sender.close(); + return Err(e); + } Ok(MTPConnection { version: PROTOCOL_VERSION, @@ -370,6 +422,25 @@ impl MTPClient { config: ClientConfig, keys: &mtp_crypto::Keyring, host_public_key_bundle: &mtp_crypto::PublicKeyBundle, + ) -> Result { + let timeout = config.auth_timeout; + match tokio::time::timeout( + timeout, + Self::auth_register_inner(config, keys, host_public_key_bundle), + ) + .await + { + Ok(result) => result, + Err(_) => Err(CommunicationError::AuthenticationFailed( + "authentication timed out".into(), + )), + } + } + + async fn auth_register_inner( + config: ClientConfig, + keys: &mtp_crypto::Keyring, + host_public_key_bundle: &mtp_crypto::PublicKeyBundle, ) -> Result { use mtp_crypto::auth; @@ -385,54 +456,86 @@ impl MTPClient { let register = CommunicationValue::new(mtp_codec::CommunicationType::Register) .add_typed_default(DataType::Version, DataValue::Str(version_str.clone())) .add_typed_default(DataType::PublicKeys, DataValue::Bytes(pk_bytes.clone())); - sender.send(®ister).await?; + if let Err(e) = sender.send(®ister).await { + sender.close(); + return Err(e); + } // 2. Receive and verify the host's challenge (register binds id = 0). - let server_challenge = receive_verified_challenge( + let server_challenge = match receive_verified_challenge( &receiver, &tm, host_public_key_bundle, 0, "auth_register challenge", ) - .await?; + .await + { + Ok(c) => c, + Err(e) => { + sender.close(); + return Err(e); + } + }; // 3. Sign the host's challenge over the bundle and send the proof. let client_nonce: u128 = rand::random(); let proof_payload = auth::register_proof_payload(&version_str, &pk_bytes, server_challenge, client_nonce); - let proof = signed_challenge_response(keys, &proof_payload, client_nonce)?; - sender.send(&proof).await?; + let proof = match signed_challenge_response(keys, &proof_payload, client_nonce) { + Ok(p) => p, + Err(e) => { + sender.close(); + return Err(e); + } + }; + if let Err(e) = sender.send(&proof).await { + sender.close(); + return Err(e); + } // 4. Receive the host's final confirmation; extract the assigned id and // verify the host signature binds to it. - let response = receiver.receive().await?; + let response = match receiver.receive().await { + Ok(r) => r, + Err(e) => { + sender.close(); + return Err(e); + } + }; let expected_type = mtp_codec::CommunicationType::RegisterResponse.to_id(&tm); if response.get_type() != expected_type { + sender.close(); return Err(unexpected_response_type_error( "auth_register", expected_type, &response, )); } - check_connected(&response, &tm, "Server rejected registration")?; - let assigned_id = match response.get_data(DataType::Id.to_id(&tm)) { + if let Err(e) = check_connected(&response, "Server rejected registration") { + sender.close(); + return Err(e); + } + let assigned_id = match response.get_data(DataType::Id) { DataValue::UnsignedNumber(n) => *n as u64, _ => { + sender.close(); return Err(CommunicationError::AuthenticationFailed( "Missing assigned ID".into(), )); } }; - verify_host_final( + if let Err(e) = verify_host_final( &response, - &tm, host_public_key_bundle, assigned_id, client_nonce, server_challenge, - )?; + ) { + sender.close(); + return Err(e); + } Ok(MTPConnection { version: PROTOCOL_VERSION, @@ -470,8 +573,23 @@ mod tests { #[cfg(feature = "crypto")] #[test] - fn test_auth_state_derive() { - assert_eq!(AuthState::Pending, AuthState::Pending); - assert_ne!(AuthState::Authenticated, AuthState::Failed); + fn test_auth_state_unauthenticated_is_not_authenticated() { + assert_ne!(AuthState::Unauthenticated, AuthState::Authenticated); + assert_ne!(AuthState::Pending, AuthState::Authenticated); + } + + #[cfg(feature = "crypto")] + #[test] + fn test_auth_timeout_default() { + let config = ClientConfig::new("https://localhost:4433"); + assert_eq!(config.auth_timeout, Duration::from_secs(30)); + } + + #[cfg(feature = "crypto")] + #[test] + fn test_auth_timeout_custom() { + let config = ClientConfig::new("https://localhost:4433") + .with_auth_timeout(Duration::from_secs(10)); + assert_eq!(config.auth_timeout, Duration::from_secs(10)); } } diff --git a/codec/src/communication_value.rs b/codec/src/communication_value.rs index dc03e70..fafc549 100644 --- a/codec/src/communication_value.rs +++ b/codec/src/communication_value.rs @@ -3,7 +3,7 @@ use std::collections::BTreeMap; use std::fmt; use std::io::{Cursor, Read}; -use crate::data_value::DataValue; +use crate::data_value::{DataKind, DataValue}; use crate::rand_u32; use mtp_common::CodecError; use mtp_type_map::{ @@ -12,9 +12,7 @@ use mtp_type_map::{ }; #[cfg(feature = "crypto")] -use mtp_crypto::SigAlgorithm; -#[cfg(feature = "crypto")] -use mtp_crypto::SignatureScheme; +use mtp_crypto::{PublicKeyBundle, SigAlgorithm, SignatureScheme}; const FLAG_HAS_SENDER: u8 = 0b0000_0001; const FLAG_HAS_RECEIVER: u8 = 0b0000_0010; @@ -121,13 +119,147 @@ impl CommunicationValue { self } - pub fn get_data(&self, data_type: DataTypeId) -> &DataValue { - self.data.get(&data_type).unwrap_or(&DataValue::Null) + pub fn get_data(&self, data_type: DataType) -> &DataValue { + let tm_owned; + let tm = match &self.type_map { + Some(tm) => tm, + None => { + tm_owned = TypeMap::latest(); + &tm_owned + } + }; + match tm.data_id_enum(data_type) { + Some(raw_id) => self.data.get(&DataTypeId(raw_id)).unwrap_or(&DataValue::Null), + None => &DataValue::Null, + } + } + + pub fn get_data_opt(&self, data_type: DataType) -> Option<&DataValue> { + let tm_owned; + let tm = match &self.type_map { + Some(tm) => tm, + None => { + tm_owned = TypeMap::latest(); + &tm_owned + } + }; + let raw_id = tm.data_id_enum(data_type)?; + self.data.get(&DataTypeId(raw_id)) + } + + pub fn has_data(&self, data_type: DataType) -> Option { + self.get_data_opt(data_type).map(|v| v.kind()) + } + + pub fn get_comm_type_enum(&self) -> Option { + let tm_owned; + let tm = match &self.type_map { + Some(tm) => tm, + None => { + tm_owned = TypeMap::latest(); + &tm_owned + } + }; + tm.comm_enum_id(self.comm_type.0) } pub fn data(&self) -> &BTreeMap { &self.data } + + pub fn data_len(&self) -> usize { + self.data.len() + } + + // ── type checks ────────────────────────────────────────────────────────── + + pub fn is_type(&self, comm_type: CommunicationType) -> bool { + self.get_comm_type_enum() == Some(comm_type) + } + + pub fn get_type_name(&self) -> Option<&'static str> { + communication_type_name(self.comm_type.0) + } + + // ── mutation ───────────────────────────────────────────────────────────── + + pub fn set_data(&mut self, data_type: DataType, value: DataValue) { + let tm = self.type_map.clone().unwrap_or_else(TypeMap::latest); + if let Some(raw_id) = tm.data_id_enum(data_type) { + self.data.insert(DataTypeId(raw_id), value); + } + } + + #[must_use] + pub fn with_data(mut self, data_type: DataType, value: DataValue) -> Self { + self.set_data(data_type, value); + self + } + + pub fn remove_data(&mut self, data_type: DataType) -> Option { + let tm_owned; + let tm = match &self.type_map { + Some(tm) => tm, + None => { + tm_owned = TypeMap::latest(); + &tm_owned + } + }; + let raw_id = tm.data_id_enum(data_type)?; + self.data.remove(&DataTypeId(raw_id)) + } + + #[must_use] + pub fn reply_to(&self, comm_type: CommunicationType) -> Self { + Self::new(comm_type) + .with_sender(self.receiver) + .with_receiver(self.sender) + } + + pub fn merge(&mut self, other: &CommunicationValue) { + for (id, value) in &other.data { + self.data.insert(*id, value.clone()); + } + } + + // ── typed iteration ────────────────────────────────────────────────────── + + pub fn iter_typed_data(&self) -> impl Iterator, &DataValue)> + '_ { + let tm = self.type_map.clone().unwrap_or_else(TypeMap::latest); + self.data + .iter() + .map(move |(id, val)| (tm.data_enum_id(id.0), val)) + } + + // ── typed field accessors ───────────────────────────────────────────────── + + pub fn get_bool(&self, data_type: DataType) -> Option { + self.get_data_opt(data_type)?.as_bool() + } + + pub fn get_str(&self, data_type: DataType) -> Option<&str> { + self.get_data_opt(data_type)?.as_str() + } + + pub fn get_u128(&self, data_type: DataType) -> Option { + self.get_data_opt(data_type)?.as_unsigned_number() + } + + pub fn get_i128(&self, data_type: DataType) -> Option { + self.get_data_opt(data_type)?.as_signed_number() + } + + pub fn get_float(&self, data_type: DataType) -> Option<(u8, u32)> { + self.get_data_opt(data_type)?.as_float() + } + + pub fn get_bytes(&self, data_type: DataType) -> Option<&[u8]> { + self.get_data_opt(data_type)?.as_bytes_slice() + } + + pub fn get_array(&self, data_type: DataType) -> Option<&[DataValue]> { + self.get_data_opt(data_type)?.as_array_slice() + } } impl CommunicationValue { @@ -433,6 +565,64 @@ impl CommunicationValue { self.frame_signature.as_ref() } + /* + * Verify the frame signature using a `PublicKeyBundle`. Dispatches to + * Ed25519, ML-DSA-65, or both (DUAL) based on the stored algorithm byte. + * Returns `false` if the frame has no signature or verification fails. + */ + #[cfg(feature = "crypto")] + pub fn validate_signature(&self, pk: &PublicKeyBundle) -> bool { + let Some((alg, _)) = &self.frame_signature else { + return false; + }; + struct Ed25519Verifier<'a>(&'a mtp_crypto::SignaturePublicKey); + impl SignatureScheme for Ed25519Verifier<'_> { + fn sign(&self, _: &[u8]) -> Result, mtp_crypto::CryptoError> { + Err(mtp_crypto::CryptoError::SigningFailed) + } + fn verify(&self, msg: &[u8], sig: &[u8]) -> Result<(), mtp_crypto::CryptoError> { + mtp_crypto::verify_ed25519(self.0, msg, sig) + } + } + struct MlDsaVerifier<'a>(&'a mtp_crypto::SignaturePqPublicKey); + impl SignatureScheme for MlDsaVerifier<'_> { + fn sign(&self, _: &[u8]) -> Result, mtp_crypto::CryptoError> { + Err(mtp_crypto::CryptoError::SigningFailed) + } + fn verify(&self, msg: &[u8], sig: &[u8]) -> Result<(), mtp_crypto::CryptoError> { + mtp_crypto::verify_ml_dsa(self.0, msg, sig) + } + } + match *alg { + SigAlgorithm::ED25519 => { + self.verify_frame(&Ed25519Verifier(&pk.sig_cl_public_key)).is_ok() + } + SigAlgorithm::ML_DSA_65 => { + self.verify_frame(&MlDsaVerifier(&pk.sig_pq_public_key)).is_ok() + } + SigAlgorithm::DUAL => { + // For DUAL, verify_frame passes the full combined sig to the verifier. + // We wrap a verifier that splits and checks both halves. + struct DualVerifier<'a>(&'a mtp_crypto::SignaturePublicKey, &'a mtp_crypto::SignaturePqPublicKey); + impl SignatureScheme for DualVerifier<'_> { + fn sign(&self, _: &[u8]) -> Result, mtp_crypto::CryptoError> { + Err(mtp_crypto::CryptoError::SigningFailed) + } + fn verify(&self, msg: &[u8], sig: &[u8]) -> Result<(), mtp_crypto::CryptoError> { + const ED_LEN: usize = 64; + if sig.len() < ED_LEN { + return Err(mtp_crypto::CryptoError::InvalidSignature); + } + mtp_crypto::verify_ed25519(self.0, msg, &sig[..ED_LEN])?; + mtp_crypto::verify_ml_dsa(self.1, msg, &sig[ED_LEN..]) + } + } + self.verify_frame(&DualVerifier(&pk.sig_cl_public_key, &pk.sig_pq_public_key)).is_ok() + } + _ => false, + } + } + #[cfg(feature = "registry")] pub fn migrate(&self, target_tm: &TypeMap) -> Result { let comm_name = communication_type_name(self.comm_type.0) @@ -544,15 +734,23 @@ impl fmt::Display for CommunicationValue { write!(f, ", R:{}{:X}{}", ORANGE, self.receiver, RESET)?; } - let name = communication_type_name(self.comm_type.0).unwrap_or("?"); + let name = self + .get_comm_type_enum() + .map(|t| t.name()) + .unwrap_or_else(|| communication_type_name(self.comm_type.0).unwrap_or("?")); write!(f, ", {}: ", name)?; + let tm = self.type_map.clone().unwrap_or_else(TypeMap::latest); write!(f, "{{")?; - for (i, (key, value)) in self.data.iter().enumerate() { + for (i, (raw_id, value)) in self.data.iter().enumerate() { if i > 0 { write!(f, ", ")?; } - let dname = data_type_name(key.0).unwrap_or("?"); + let dname = tm + .data_enum_id(raw_id.0) + .map(|t| t.name()) + .or_else(|| data_type_name(raw_id.0)) + .unwrap_or("?"); write!(f, "{}: ", dname)?; fmt_data_value(value, f)?; } @@ -645,11 +843,11 @@ mod tests { assert_eq!(decoded.get_receiver(), 222); assert_eq!(decoded.get_type(), CommunicationType::Disconnect.to_id(&tm)); assert_eq!( - decoded.get_data(DataType::Id.to_id(&tm)), + decoded.get_data(DataType::Id), &DataValue::Str("alice".to_string()) ); assert_eq!( - decoded.get_data(DataType::ClientNonce.to_id(&tm)), + decoded.get_data(DataType::ClientNonce), &DataValue::SignedNumber(42) ); } diff --git a/codec/src/data_value.rs b/codec/src/data_value.rs index 5f4aced..99fbe35 100644 --- a/codec/src/data_value.rs +++ b/codec/src/data_value.rs @@ -39,6 +39,28 @@ pub enum DataKind { Null, } +impl fmt::Display for DataKind { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + DataKind::Bool => f.write_str("Bool"), + DataKind::SignedNumber => f.write_str("SignedNumber"), + DataKind::UnsignedNumber => f.write_str("UnsignedNumber"), + DataKind::Float => f.write_str("Float"), + DataKind::Str => f.write_str("Str"), + DataKind::Bytes => f.write_str("Bytes"), + DataKind::Array(inner) => write!(f, "Array<{}>", inner), + DataKind::Container => f.write_str("Container"), + #[cfg(feature = "crypto")] + DataKind::EncryptedContainer => f.write_str("EncryptedContainer"), + #[cfg(feature = "crypto")] + DataKind::SignedContainer => f.write_str("SignedContainer"), + #[cfg(feature = "crypto")] + DataKind::SignedEncryptedContainer => f.write_str("SignedEncryptedContainer"), + DataKind::Null => f.write_str("Null"), + } + } +} + #[derive(Debug, Clone, Eq)] pub enum DataValue { BoolTrue, @@ -259,6 +281,77 @@ impl DataValue { } } + pub fn as_number(&self) -> Option { + match self { + DataValue::SignedNumber(n) => Some(*n), + DataValue::UnsignedNumber(n) => Some(*n as i128), + _ => None, + } + } + + pub fn is_null(&self) -> bool { + matches!(self, DataValue::Null) + } + + pub fn is_truthy(&self) -> bool { + match self { + DataValue::BoolTrue | DataValue::Bool(true) => true, + DataValue::BoolFalse | DataValue::Bool(false) | DataValue::Null => false, + DataValue::UnsignedNumber(0) | DataValue::SignedNumber(0) => false, + _ => true, + } + } + + pub fn get_field(&self, key: DataTypeId) -> Option<&DataValue> { + match self { + DataValue::Container(entries) => { + entries.iter().find(|(k, _)| *k == key).map(|(_, v)| v) + } + _ => None, + } + } + + pub fn as_container_map(&self) -> Option> { + match self { + DataValue::Container(entries) => Some(entries.iter().cloned().collect()), + _ => None, + } + } + + pub fn as_bytes_slice(&self) -> Option<&[u8]> { + match self { + DataValue::Bytes(b) => Some(b), + _ => None, + } + } + + pub fn as_array_slice(&self) -> Option<&[DataValue]> { + match self { + DataValue::Array(a) => Some(a), + _ => None, + } + } + + pub fn type_name(&self) -> &'static str { + match self { + DataValue::Bool(_) | DataValue::BoolTrue | DataValue::BoolFalse => "Bool", + DataValue::SignedNumber(_) => "SignedNumber", + DataValue::UnsignedNumber(_) => "UnsignedNumber", + DataValue::Float(_, _) => "Float", + DataValue::Str(_) => "Str", + DataValue::Bytes(_) => "Bytes", + DataValue::Array(_) => "Array", + DataValue::Container(_) => "Container", + #[cfg(feature = "crypto")] + DataValue::EncryptedContainer(_) => "EncryptedContainer", + #[cfg(feature = "crypto")] + DataValue::SignedContainer(_) => "SignedContainer", + #[cfg(feature = "crypto")] + DataValue::SignedEncryptedContainer(_) => "SignedEncryptedContainer", + DataValue::Null => "Null", + } + } + #[cfg(feature = "crypto")] pub fn as_encrypted_container(&self) -> Option> { match self { @@ -374,6 +467,60 @@ impl DataValue { Some(()) } + /* + * Verify a `SignedContainer` without mutating self. Dispatches to + * Ed25519, ML-DSA-65, or both (DUAL) based on the algorithm byte + * embedded in the blob. Returns `false` for any other variant. + */ + #[cfg(feature = "crypto")] + pub fn validate_signature(&self, pk: &PublicKeyBundle) -> bool { + let blob = match self { + DataValue::SignedContainer(b) => b, + _ => return false, + }; + if blob.is_empty() { + return false; + } + let alg = blob[0]; + let sig_len = match SigAlgorithm::length(alg) { + Some(n) => n, + None => return false, + }; + if blob.len() < 1 + sig_len + 2 { + return false; + } + let signature = &blob[1..1 + sig_len]; + let container_bytes = &blob[1 + sig_len..]; + match alg { + SigAlgorithm::ED25519 => { + mtp_crypto::verify_ed25519(&pk.sig_cl_public_key, container_bytes, signature).is_ok() + } + SigAlgorithm::ML_DSA_65 => { + mtp_crypto::verify_ml_dsa(&pk.sig_pq_public_key, container_bytes, signature).is_ok() + } + SigAlgorithm::DUAL => { + const ED_LEN: usize = 64; + if signature.len() < ED_LEN { + return false; + } + let ed_ok = mtp_crypto::verify_ed25519( + &pk.sig_cl_public_key, + container_bytes, + &signature[..ED_LEN], + ) + .is_ok(); + let ml_ok = mtp_crypto::verify_ml_dsa( + &pk.sig_pq_public_key, + container_bytes, + &signature[ED_LEN..], + ) + .is_ok(); + ed_ok && ml_ok + } + _ => false, + } + } + /* * Encrypt a `Container` into a `SignedEncryptedContainer` in-place. * The container is first signed (with `algorithm`/`signer`), then the signed @@ -895,6 +1042,133 @@ impl Hash for DataValue { } } +/* ================================ FROM / TRY-FROM ================================ */ + +impl From for DataValue { + fn from(v: bool) -> Self { + if v { DataValue::BoolTrue } else { DataValue::BoolFalse } + } +} + +impl From<&str> for DataValue { + fn from(s: &str) -> Self { + DataValue::Str(s.to_string()) + } +} + +impl From for DataValue { + fn from(s: String) -> Self { + DataValue::Str(s) + } +} + +impl From for DataValue { + fn from(n: i64) -> Self { + DataValue::SignedNumber(n as i128) + } +} + +impl From for DataValue { + fn from(n: i128) -> Self { + DataValue::SignedNumber(n) + } +} + +impl From for DataValue { + fn from(n: u64) -> Self { + DataValue::UnsignedNumber(n as u128) + } +} + +impl From for DataValue { + fn from(n: u128) -> Self { + DataValue::UnsignedNumber(n) + } +} + +impl From> for DataValue { + fn from(b: Vec) -> Self { + DataValue::Bytes(b) + } +} + +impl From<&[u8]> for DataValue { + fn from(b: &[u8]) -> Self { + DataValue::Bytes(b.to_vec()) + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DataValueTypeMismatch { + pub expected: &'static str, + pub got: &'static str, +} + +impl fmt::Display for DataValueTypeMismatch { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "expected {}, got {}", self.expected, self.got) + } +} + +impl std::error::Error for DataValueTypeMismatch {} + +impl TryFrom for bool { + type Error = DataValueTypeMismatch; + fn try_from(v: DataValue) -> Result { + v.as_bool().ok_or(DataValueTypeMismatch { expected: "Bool", got: v.type_name() }) + } +} + +impl TryFrom for String { + type Error = DataValueTypeMismatch; + fn try_from(v: DataValue) -> Result { + match v { + DataValue::Str(s) => Ok(s), + other => Err(DataValueTypeMismatch { expected: "Str", got: other.type_name() }), + } + } +} + +impl TryFrom for i128 { + type Error = DataValueTypeMismatch; + fn try_from(v: DataValue) -> Result { + v.as_signed_number().ok_or(DataValueTypeMismatch { expected: "SignedNumber", got: v.type_name() }) + } +} + +impl TryFrom for i64 { + type Error = DataValueTypeMismatch; + fn try_from(v: DataValue) -> Result { + let n = v.as_signed_number().ok_or(DataValueTypeMismatch { expected: "SignedNumber", got: v.type_name() })?; + Ok(n as i64) + } +} + +impl TryFrom for u128 { + type Error = DataValueTypeMismatch; + fn try_from(v: DataValue) -> Result { + v.as_unsigned_number().ok_or(DataValueTypeMismatch { expected: "UnsignedNumber", got: v.type_name() }) + } +} + +impl TryFrom for u64 { + type Error = DataValueTypeMismatch; + fn try_from(v: DataValue) -> Result { + let n = v.as_unsigned_number().ok_or(DataValueTypeMismatch { expected: "UnsignedNumber", got: v.type_name() })?; + Ok(n as u64) + } +} + +impl TryFrom for Vec { + type Error = DataValueTypeMismatch; + fn try_from(v: DataValue) -> Result { + match v { + DataValue::Bytes(b) => Ok(b), + other => Err(DataValueTypeMismatch { expected: "Bytes", got: other.type_name() }), + } + } +} + /* ================================ TESTS ================================ */ #[cfg(test)] mod tests { @@ -1265,6 +1539,41 @@ mod tests { assert!(s.contains("6:")); } + #[test] + fn test_from_primitives() { + assert_eq!(DataValue::from(true), DataValue::BoolTrue); + assert_eq!(DataValue::from(false), DataValue::BoolFalse); + assert_eq!(DataValue::from("hello"), DataValue::Str("hello".to_string())); + assert_eq!(DataValue::from("hello".to_string()), DataValue::Str("hello".to_string())); + assert_eq!(DataValue::from(42i64), DataValue::SignedNumber(42)); + assert_eq!(DataValue::from(42i128), DataValue::SignedNumber(42)); + assert_eq!(DataValue::from(42u64), DataValue::UnsignedNumber(42)); + assert_eq!(DataValue::from(42u128), DataValue::UnsignedNumber(42)); + assert_eq!(DataValue::from(vec![1u8, 2, 3]), DataValue::Bytes(vec![1, 2, 3])); + assert_eq!(DataValue::from([1u8, 2, 3].as_ref()), DataValue::Bytes(vec![1, 2, 3])); + } + + #[test] + fn test_try_from_ok() { + assert_eq!(bool::try_from(DataValue::BoolTrue).unwrap(), true); + assert_eq!(bool::try_from(DataValue::BoolFalse).unwrap(), false); + assert_eq!(String::try_from(DataValue::Str("hi".to_string())).unwrap(), "hi"); + assert_eq!(i128::try_from(DataValue::SignedNumber(-1)).unwrap(), -1i128); + assert_eq!(i64::try_from(DataValue::SignedNumber(10)).unwrap(), 10i64); + assert_eq!(u128::try_from(DataValue::UnsignedNumber(99)).unwrap(), 99u128); + assert_eq!(u64::try_from(DataValue::UnsignedNumber(7)).unwrap(), 7u64); + assert_eq!(Vec::::try_from(DataValue::Bytes(vec![0xAB])).unwrap(), vec![0xABu8]); + } + + #[test] + fn test_try_from_err() { + assert!(bool::try_from(DataValue::Null).is_err()); + assert!(String::try_from(DataValue::SignedNumber(1)).is_err()); + assert!(i128::try_from(DataValue::BoolTrue).is_err()); + assert!(u128::try_from(DataValue::Str("x".to_string())).is_err()); + assert!(Vec::::try_from(DataValue::Null).is_err()); + } + #[test] fn test_array_display() { let dv = DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]); diff --git a/common/src/lib.rs b/common/src/lib.rs index 0a239ee..71dea88 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -16,6 +16,8 @@ pub enum CodecError { TooManyEntries, #[error("Crypto failed: {0}")] CryptoFailed(String), + #[error("Missing required field: {0}")] + MissingField(String), } /* ================================ TESTS ================================ */ diff --git a/crypto/src/keypair.rs b/crypto/src/keypair.rs index adea54a..a5e80c5 100644 --- a/crypto/src/keypair.rs +++ b/crypto/src/keypair.rs @@ -21,6 +21,12 @@ impl From> for EncryptionPrivateKey { } } +impl From<&[u8]> for EncryptionPrivateKey { + fn from(bytes: &[u8]) -> Self { + Self(bytes.to_vec()) + } +} + #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[cfg_attr(feature = "serde", serde(transparent))] #[derive(Zeroize, ZeroizeOnDrop)] @@ -42,6 +48,12 @@ impl From> for SignaturePrivateKey { } } +impl From<&[u8]> for SignaturePrivateKey { + fn from(bytes: &[u8]) -> Self { + Self(bytes.to_vec()) + } +} + #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[cfg_attr(feature = "serde", serde(transparent))] #[derive(Clone)] @@ -63,6 +75,12 @@ impl From> for EncryptionPublicKey { } } +impl From<&[u8]> for EncryptionPublicKey { + fn from(bytes: &[u8]) -> Self { + Self(bytes.to_vec()) + } +} + #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[cfg_attr(feature = "serde", serde(transparent))] #[derive(Clone)] @@ -84,6 +102,12 @@ impl From> for SignaturePublicKey { } } +impl From<&[u8]> for SignaturePublicKey { + fn from(bytes: &[u8]) -> Self { + Self(bytes.to_vec()) + } +} + #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[cfg_attr(feature = "serde", serde(transparent))] #[derive(Zeroize, ZeroizeOnDrop)] @@ -105,6 +129,12 @@ impl From> for KemPrivateKey { } } +impl From<&[u8]> for KemPrivateKey { + fn from(bytes: &[u8]) -> Self { + Self(bytes.to_vec()) + } +} + #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[cfg_attr(feature = "serde", serde(transparent))] #[derive(Clone)] @@ -126,6 +156,12 @@ impl From> for KemPublicKey { } } +impl From<&[u8]> for KemPublicKey { + fn from(bytes: &[u8]) -> Self { + Self(bytes.to_vec()) + } +} + #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[cfg_attr(feature = "serde", serde(transparent))] #[derive(Clone)] @@ -147,6 +183,12 @@ impl From> for SignaturePqPublicKey { } } +impl From<&[u8]> for SignaturePqPublicKey { + fn from(bytes: &[u8]) -> Self { + Self(bytes.to_vec()) + } +} + #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[cfg_attr(feature = "serde", serde(transparent))] #[derive(Zeroize, ZeroizeOnDrop)] @@ -168,6 +210,12 @@ impl From> for SignaturePqPrivateKey { } } +impl From<&[u8]> for SignaturePqPrivateKey { + fn from(bytes: &[u8]) -> Self { + Self(bytes.to_vec()) + } +} + #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[derive(ZeroizeOnDrop)] pub struct Keyring { diff --git a/example/server/src/handlers.rs b/example/server/src/handlers.rs index 2a35595..941d468 100644 --- a/example/server/src/handlers.rs +++ b/example/server/src/handlers.rs @@ -31,13 +31,13 @@ pub fn process_and_respond( let sig_id = DataTypeId(tm.data_id_enum(DataType::SignedPayload).unwrap()); let secure_id = DataTypeId(tm.data_id_enum(DataType::SecurePayload).unwrap()); - let description = msg.get_data(desc_id); - let timestamp = msg.get_data(ts_id); - let data = msg.get_data(data_id); - let flags = msg.get_data(flags_id); - let value = msg.get_data(value_id); - let binary = msg.get_data(bin_id); - let items = msg.get_data(items_id); + let description = msg.get_data(DataType::Description); + let timestamp = msg.get_data(DataType::Timestamp); + let data = msg.get_data(DataType::Data); + let flags = msg.get_data(DataType::Flags); + let value = msg.get_data(DataType::Value); + let binary = msg.get_data(DataType::BinaryData); + let items = msg.get_data(DataType::Items); println!( " Description: {}", @@ -54,7 +54,7 @@ pub fn process_and_respond( let mut sig_status = String::from("SignedPayload: not present"); let mut secure_status = String::from("SecurePayload: not present"); - let enc = msg.get_data(enc_id); + let enc = msg.get_data(DataType::EncryptedPayload); if matches!(enc, DataValue::EncryptedContainer(_)) { let mut dv = enc.clone(); if dv.decrypt_into_container(host_keyring, b"demo-aad").is_some() { @@ -68,7 +68,7 @@ pub fn process_and_respond( } } - let sig = msg.get_data(sig_id); + let sig = msg.get_data(DataType::SignedPayload); if matches!(sig, DataValue::SignedContainer(_)) { if let Some(pk_bundle) = client_pk { let verifier = Ed25519Verifier(pk_bundle.sig_cl_public_key.clone()); @@ -88,7 +88,7 @@ pub fn process_and_respond( } } - let secure = msg.get_data(secure_id); + let secure = msg.get_data(DataType::SecurePayload); if matches!(secure, DataValue::SignedEncryptedContainer(_)) { if let Some(pk_bundle) = client_pk { let verifier = Ed25519Verifier(pk_bundle.sig_cl_public_key.clone()); diff --git a/host/Cargo.toml b/host/Cargo.toml index 9750284..5c7ab76 100644 --- a/host/Cargo.toml +++ b/host/Cargo.toml @@ -9,6 +9,7 @@ mtp-codec = { path = "../codec", features = ["registry"] } mtp-transport = { path = "../transport", features = ["host"] } mtp-crypto = { path = "../crypto", optional = true } rand = "0.8" +tokio = { version = "1", features = ["time"] } [features] crypto = ["dep:mtp-crypto", "mtp-codec/crypto"] diff --git a/host/src/lib.rs b/host/src/lib.rs index 54e58a1..58c79ae 100644 --- a/host/src/lib.rs +++ b/host/src/lib.rs @@ -5,7 +5,24 @@ use mtp_codec::{ use mtp_common::CommunicationError; use mtp_transport::{Policy, Receiver, Sender}; use std::net::IpAddr; +use std::pin::Pin; use std::{error::Error, fmt}; +use tokio::time::Duration; + +/* ---- async callback type aliases ---- */ +#[cfg(feature = "crypto")] +type GetExistingUser = Box< + dyn Fn(u64) -> Pin> + Send>> + + Send + + Sync, +>; + +#[cfg(feature = "crypto")] +type CompleteRegister = Box< + dyn Fn(mtp_crypto::PublicKeyBundle) -> Pin + Send>> + + Send + + Sync, +>; /* Host configuration. */ pub struct HostConfig { @@ -17,11 +34,13 @@ pub struct HostConfig { #[cfg(feature = "crypto")] pub require_authentication: bool, #[cfg(feature = "crypto")] + pub auth_timeout: Duration, + #[cfg(feature = "crypto")] pub host_keyring: mtp_crypto::Keyring, #[cfg(feature = "crypto")] - pub get_existing_user: Box Option + Send + Sync>, + pub get_existing_user: GetExistingUser, #[cfg(feature = "crypto")] - pub complete_register: Box u64 + Send + Sync>, + pub complete_register: CompleteRegister, } impl HostConfig { @@ -34,6 +53,8 @@ impl HostConfig { #[cfg(feature = "crypto")] require_authentication: false, #[cfg(feature = "crypto")] + auth_timeout: Duration::from_secs(30), + #[cfg(feature = "crypto")] host_keyring: mtp_crypto::Keyring::new( mtp_crypto::KemPublicKey::new(Vec::new()), mtp_crypto::KemPrivateKey::new(Vec::new()), @@ -43,9 +64,9 @@ impl HostConfig { mtp_crypto::SignaturePrivateKey::new(Vec::new()), ), #[cfg(feature = "crypto")] - get_existing_user: Box::new(|_| None), + get_existing_user: Box::new(|_| Box::pin(async { None })), #[cfg(feature = "crypto")] - complete_register: Box::new(|_| 0), + complete_register: Box::new(|_| Box::pin(async { 0 })), } } @@ -53,8 +74,14 @@ impl HostConfig { pub fn with_authentication( mut self, host_keyring: mtp_crypto::Keyring, - get_existing_user: impl Fn(u64) -> Option + Send + Sync + 'static, - complete_register: impl Fn(mtp_crypto::PublicKeyBundle) -> u64 + Send + Sync + 'static, + get_existing_user: impl Fn(u64) -> Pin> + Send>> + + Send + + Sync + + 'static, + complete_register: impl Fn(mtp_crypto::PublicKeyBundle) -> Pin + Send>> + + Send + + Sync + + 'static, ) -> Self { self.require_authentication = true; self.host_keyring = host_keyring; @@ -62,6 +89,12 @@ impl HostConfig { self.complete_register = Box::new(complete_register); self } + + #[cfg(feature = "crypto")] + pub fn with_auth_timeout(mut self, timeout: Duration) -> Self { + self.auth_timeout = timeout; + self + } } #[derive(Debug, Clone, PartialEq, Eq)] @@ -70,6 +103,7 @@ pub enum AcceptError { MissingVersion, UnsupportedVersion(Version), AuthenticationFailed(String), + AuthenticationTimedOut, Send(CommunicationError), } @@ -85,6 +119,7 @@ impl fmt::Display for AcceptError { write!(f, "unsupported protocol version: {version}") } Self::AuthenticationFailed(reason) => write!(f, "authentication failed: {reason}"), + Self::AuthenticationTimedOut => write!(f, "authentication handshake timed out"), Self::Send(error) => write!(f, "failed to send handshake message: {error}"), } } @@ -95,6 +130,7 @@ impl Error for AcceptError {} #[cfg(feature = "crypto")] #[derive(Debug, Clone, PartialEq, Eq)] pub enum AuthState { + Unauthenticated, Pending, Authenticated, Failed, @@ -158,7 +194,16 @@ impl MTPHost { #[cfg(feature = "crypto")] if self.config.require_authentication { - return self.accept_authenticated(sender, receiver).await; + let timeout = self.config.auth_timeout; + return match tokio::time::timeout( + timeout, + self.accept_authenticated(sender, receiver), + ) + .await + { + Ok(result) => result, + Err(_) => Err(AcceptError::AuthenticationTimedOut), + }; } // Read the first message (always encoded with reserved types). @@ -167,13 +212,6 @@ impl MTPHost { Err(e) => return Err(AcceptError::Receive(e)), }; - /* - * Extract the client's version from the first message. - * The client is expected to send DataType::Version (reserved ID 3) - * as a DataValue::Str("X.Y"). - * - * Then negotiate the version for single-version clients - */ let client_version = match extract_version(&first_msg) { Some(v) => v, None => return Err(AcceptError::MissingVersion), @@ -195,7 +233,7 @@ impl MTPHost { sender, receiver, #[cfg(feature = "crypto")] - auth_state: AuthState::Authenticated, + auth_state: AuthState::Unauthenticated, #[cfg(feature = "crypto")] client_id: 0, #[cfg(feature = "crypto")] @@ -280,20 +318,31 @@ impl MTPHost { }; // ===== Step 1: receive the client's unsigned hello ===== - let hello = receiver.receive().await.map_err(AcceptError::Receive)?; - let version_str = match hello.get_data(DataType::Version.to_id(&tm)) { + let hello = match receiver.receive().await { + Ok(m) => m, + Err(e) => { + sender.close(); + return Err(AcceptError::Receive(e)); + } + }; + let version_str = match hello.get_data(DataType::Version) { DataValue::Str(s) => s.clone(), _ => { sender.close(); return Err(AcceptError::MissingVersion); } }; - let client_version = Version::parse(&version_str).ok_or(AcceptError::MissingVersion)?; + let client_version = match Version::parse(&version_str) { + Some(v) => v, + None => { + sender.close(); + return Err(AcceptError::MissingVersion); + } + }; let (flow, response_type) = if hello.get_type() == mtp_codec::CommunicationType::Identification.to_id(&tm) { - // LOGIN: look up the claimed user before issuing a challenge. - let cid = match hello.get_data(DataType::Id.to_id(&tm)) { + let cid = match hello.get_data(DataType::Id) { DataValue::UnsignedNumber(n) => *n as u64, _ => { sender.close(); @@ -302,7 +351,7 @@ impl MTPHost { )); } }; - let bundle = match (self.config.get_existing_user)(cid) { + let bundle = match (self.config.get_existing_user)(cid).await { Some(b) => b, None => { let rejection = CommunicationValue::new( @@ -321,8 +370,7 @@ impl MTPHost { mtp_codec::CommunicationType::IdentificationResponse, ) } else if hello.get_type() == mtp_codec::CommunicationType::Register.to_id(&tm) { - // REGISTER: the client presents the bundle it wants to register. - let bundle = match hello.get_data(DataType::PublicKeys.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()) })?, @@ -345,7 +393,6 @@ impl MTPHost { )); }; - // The id bound into the challenge (0 for register: none assigned yet). let challenge_id = match &flow { Flow::Login { id, .. } => *id, Flow::Register { .. } => 0, @@ -366,20 +413,26 @@ impl MTPHost { challenge_msg = challenge_msg .add_typed_default(DataType::PqSignature, DataValue::Bytes(chal_pq_sig)); } - sender - .send(&challenge_msg) - .await - .map_err(AcceptError::Send)?; + if let Err(e) = sender.send(&challenge_msg).await { + sender.close(); + return Err(AcceptError::Send(e)); + } // ===== Step 3: receive and verify the client's proof ===== - let proof = receiver.receive().await.map_err(AcceptError::Receive)?; + let proof = match receiver.receive().await { + Ok(m) => m, + Err(e) => { + sender.close(); + return Err(AcceptError::Receive(e)); + } + }; if proof.get_type() != mtp_codec::CommunicationType::ChallengeResponse.to_id(&tm) { sender.close(); return Err(AcceptError::AuthenticationFailed( "missing challenge response".into(), )); } - let client_nonce = match proof.get_data(DataType::ClientNonce.to_id(&tm)) { + let client_nonce = match proof.get_data(DataType::ClientNonce) { DataValue::UnsignedNumber(n) => *n, _ => { sender.close(); @@ -388,7 +441,7 @@ impl MTPHost { )); } }; - let sig_bytes = match proof.get_data(DataType::Signature.to_id(&tm)) { + let sig_bytes = match proof.get_data(DataType::Signature) { DataValue::Bytes(b) => b.clone(), _ => { sender.close(); @@ -397,7 +450,7 @@ impl MTPHost { )); } }; - let pq_sig_bytes: Vec = match proof.get_data(DataType::PqSignature.to_id(&tm)) { + let pq_sig_bytes: Vec = match proof.get_data(DataType::PqSignature) { DataValue::Bytes(b) => b.clone(), _ => vec![], }; @@ -439,7 +492,7 @@ impl MTPHost { 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()); + let new_id = (self.config.complete_register)(bundle.clone()).await; (new_id, bundle) } }; @@ -464,14 +517,26 @@ impl MTPHost { response.add_typed_default(DataType::PqSignature, DataValue::Bytes(host_pq_sig)); } - sender.send(&response).await.map_err(AcceptError::Send)?; - sender.finish_stream().await.map_err(AcceptError::Send)?; + 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)); + } // ===== Version negotiation ===== - let negotiated = self + let negotiated = match self .registry .negotiate(std::slice::from_ref(&client_version)) - .ok_or(AcceptError::UnsupportedVersion(client_version))?; + { + Some(v) => v, + None => { + sender.close(); + return Err(AcceptError::UnsupportedVersion(client_version)); + } + }; let codec = VersionedCodec::new(self.registry.clone()); Ok(Some(MTPConnection { @@ -493,8 +558,7 @@ impl MTPHost { * (reserved ID 3) mapping to `DataValue::Str("major.minor")`. */ fn extract_version(msg: &CommunicationValue) -> Option { - let tm = TypeMap::latest(); - let value = msg.get_data(DataType::Version.to_id(&tm)); + let value = msg.get_data(DataType::Version); match value { DataValue::Str(s) => Version::parse(s.as_str()), _ => None, @@ -541,4 +605,11 @@ mod tests { .add_data(DataType::Version.to_id(&tm), DataValue::UnsignedNumber(42)); assert!(extract_version(&msg).is_none()); } + + #[cfg(feature = "crypto")] + #[test] + fn auth_state_unauthenticated_is_not_authenticated() { + assert_ne!(AuthState::Unauthenticated, AuthState::Authenticated); + assert_ne!(AuthState::Pending, AuthState::Authenticated); + } } diff --git a/transport/tests/integration.rs b/transport/tests/integration.rs index 73fc824..1099b9e 100644 --- a/transport/tests/integration.rs +++ b/transport/tests/integration.rs @@ -55,7 +55,7 @@ fn assert_numbered_message( ) { assert_eq!(message.get_type(), comm_type.to_id(tm)); assert_eq!( - message.get_data(DataType::PqSignature.to_id(tm)).clone(), + message.get_data(DataType::PqSignature).clone(), DataValue::UnsignedNumber(value) ); } diff --git a/type-map/build.rs b/type-map/build.rs index 9c4c318..ea8e09c 100755 --- a/type-map/build.rs +++ b/type-map/build.rs @@ -261,12 +261,15 @@ fn generate(config: &Config, multi_version: bool) -> String { generate_versioned_constructors(&mut out, &sorted); generate_builtin_type_maps(&mut out, &sorted); generate_lookup_methods(&mut out, config, &sorted); + generate_all_types_methods(&mut out, config, &sorted); } else { generate_single_version_lookup(&mut out, config); + generate_all_types_methods_single(&mut out, config); } generate_enum_conversion_methods(&mut out); generate_reverse_lookups(&mut out, config, &sorted, multi_version); + generate_id_display_impls(&mut out); out } @@ -564,6 +567,78 @@ fn generate_lookup_methods( writeln!(out, " _ => None,").unwrap(); writeln!(out, " }}").unwrap(); writeln!(out, " }}").unwrap(); + writeln!(out).unwrap(); + + // Reverse lookups: wire ID → enum variant (version-aware) + writeln!( + out, + " pub fn comm_enum_id(&self, id: u16) -> Option {{" + ) + .unwrap(); + writeln!(out, " match self.version {{").unwrap(); + for (version_key, major, minor) in sorted_versions { + let tm_cfg = &config.type_maps[version_key]; + writeln!( + out, + " Version({}, {}) => match id {{", + major, minor + ) + .unwrap(); + for entry in RESERVED_COMM_TYPES { + writeln!( + out, + " {} => Some(CommunicationType::{}),", + entry.id, entry.name + ) + .unwrap(); + } + for (name, id) in &tm_cfg.communication_types { + writeln!( + out, + " {} => Some(CommunicationType::{}),", + id, name + ) + .unwrap(); + } + writeln!(out, " _ => None,").unwrap(); + writeln!(out, " }},").unwrap(); + } + writeln!(out, " _ => None,").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out).unwrap(); + + writeln!( + out, + " pub fn data_enum_id(&self, id: u16) -> Option {{" + ) + .unwrap(); + writeln!(out, " match self.version {{").unwrap(); + for (version_key, major, minor) in sorted_versions { + let tm_cfg = &config.type_maps[version_key]; + writeln!( + out, + " Version({}, {}) => match id {{", + major, minor + ) + .unwrap(); + for entry in RESERVED_DATA_TYPES { + writeln!( + out, + " {} => Some(DataType::{}),", + entry.id, entry.name + ) + .unwrap(); + } + for (name, id) in &tm_cfg.data_types { + writeln!(out, " {} => Some(DataType::{}),", id, name).unwrap(); + } + writeln!(out, " _ => None,").unwrap(); + writeln!(out, " }},").unwrap(); + } + writeln!(out, " _ => None,").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); writeln!(out, "}}").unwrap(); writeln!(out).unwrap(); } @@ -631,6 +706,66 @@ fn generate_single_version_lookup(out: &mut String, config: &Config) { writeln!(out, " _ => None,").unwrap(); writeln!(out, " }}").unwrap(); writeln!(out, " }}").unwrap(); + writeln!(out).unwrap(); + + // Reverse lookups: wire ID → enum variant (single version) + writeln!( + out, + " pub fn comm_enum_id(&self, id: u16) -> Option {{" + ) + .unwrap(); + writeln!(out, " match self.version {{").unwrap(); + writeln!(out, " PROTOCOL_VERSION => match id {{").unwrap(); + for entry in RESERVED_COMM_TYPES { + writeln!( + out, + " {} => Some(CommunicationType::{}),", + entry.id, entry.name + ) + .unwrap(); + } + if let Some(tm_cfg) = tm_cfg { + for (name, id) in &tm_cfg.communication_types { + writeln!( + out, + " {} => Some(CommunicationType::{}),", + id, name + ) + .unwrap(); + } + } + writeln!(out, " _ => None,").unwrap(); + writeln!(out, " }},").unwrap(); + writeln!(out, " _ => None,").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out).unwrap(); + + writeln!( + out, + " pub fn data_enum_id(&self, id: u16) -> Option {{" + ) + .unwrap(); + writeln!(out, " match self.version {{").unwrap(); + writeln!(out, " PROTOCOL_VERSION => match id {{").unwrap(); + for entry in RESERVED_DATA_TYPES { + writeln!( + out, + " {} => Some(DataType::{}),", + entry.id, entry.name + ) + .unwrap(); + } + if let Some(tm_cfg) = tm_cfg { + for (name, id) in &tm_cfg.data_types { + writeln!(out, " {} => Some(DataType::{}),", id, name).unwrap(); + } + } + writeln!(out, " _ => None,").unwrap(); + writeln!(out, " }},").unwrap(); + writeln!(out, " _ => None,").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); writeln!(out, "}}").unwrap(); writeln!(out).unwrap(); } @@ -745,3 +880,187 @@ fn generate_enum_conversion_methods(out: &mut String) { writeln!(out, "}}").unwrap(); writeln!(out).unwrap(); } + +fn generate_all_types_methods( + out: &mut String, + config: &Config, + sorted_versions: &[(String, u16, u16)], +) { + writeln!(out, "#[allow(unreachable_patterns)]").unwrap(); + writeln!(out, "impl TypeMap {{").unwrap(); + writeln!( + out, + " pub fn all_comm_types(&self) -> &'static [CommunicationType] {{" + ) + .unwrap(); + writeln!(out, " match self.version {{").unwrap(); + for (version_key, major, minor) in sorted_versions { + let tm_cfg = &config.type_maps[version_key]; + write!(out, " Version({}, {}) => &[", major, minor).unwrap(); + let mut first = true; + for entry in RESERVED_COMM_TYPES { + if !first { + write!(out, ", ").unwrap(); + } + write!(out, "CommunicationType::{}", entry.name).unwrap(); + first = false; + } + for (name, _) in &tm_cfg.communication_types { + if !first { + write!(out, ", ").unwrap(); + } + write!(out, "CommunicationType::{}", name).unwrap(); + first = false; + } + writeln!(out, "],").unwrap(); + } + writeln!(out, " _ => &[],").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out).unwrap(); + + writeln!( + out, + " pub fn all_data_types(&self) -> &'static [DataType] {{" + ) + .unwrap(); + writeln!(out, " match self.version {{").unwrap(); + for (version_key, major, minor) in sorted_versions { + let tm_cfg = &config.type_maps[version_key]; + write!(out, " Version({}, {}) => &[", major, minor).unwrap(); + let mut first = true; + for entry in RESERVED_DATA_TYPES { + if !first { + write!(out, ", ").unwrap(); + } + write!(out, "DataType::{}", entry.name).unwrap(); + first = false; + } + for (name, _) in &tm_cfg.data_types { + if !first { + write!(out, ", ").unwrap(); + } + write!(out, "DataType::{}", name).unwrap(); + first = false; + } + writeln!(out, "],").unwrap(); + } + writeln!(out, " _ => &[],").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, "}}").unwrap(); + writeln!(out).unwrap(); +} + +fn generate_all_types_methods_single(out: &mut String, config: &Config) { + let tm_cfg = config.type_maps.get(&config.protocol_version); + + writeln!(out, "#[allow(unreachable_patterns)]").unwrap(); + writeln!(out, "impl TypeMap {{").unwrap(); + writeln!( + out, + " pub fn all_comm_types(&self) -> &'static [CommunicationType] {{" + ) + .unwrap(); + writeln!(out, " match self.version {{").unwrap(); + write!(out, " PROTOCOL_VERSION => &[").unwrap(); + let mut first = true; + for entry in RESERVED_COMM_TYPES { + if !first { + write!(out, ", ").unwrap(); + } + write!(out, "CommunicationType::{}", entry.name).unwrap(); + first = false; + } + if let Some(tm_cfg) = tm_cfg { + for (name, _) in &tm_cfg.communication_types { + if !first { + write!(out, ", ").unwrap(); + } + write!(out, "CommunicationType::{}", name).unwrap(); + first = false; + } + } + writeln!(out, "],").unwrap(); + writeln!(out, " _ => &[],").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out).unwrap(); + + writeln!( + out, + " pub fn all_data_types(&self) -> &'static [DataType] {{" + ) + .unwrap(); + writeln!(out, " match self.version {{").unwrap(); + write!(out, " PROTOCOL_VERSION => &[").unwrap(); + let mut first = true; + for entry in RESERVED_DATA_TYPES { + if !first { + write!(out, ", ").unwrap(); + } + write!(out, "DataType::{}", entry.name).unwrap(); + first = false; + } + if let Some(tm_cfg) = tm_cfg { + for (name, _) in &tm_cfg.data_types { + if !first { + write!(out, ", ").unwrap(); + } + write!(out, "DataType::{}", name).unwrap(); + first = false; + } + } + writeln!(out, "],").unwrap(); + writeln!(out, " _ => &[],").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, "}}").unwrap(); + writeln!(out).unwrap(); +} + +fn generate_id_display_impls(out: &mut String) { + writeln!(out, "impl std::fmt::Display for CommunicationTypeId {{").unwrap(); + writeln!( + out, + " fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {{" + ) + .unwrap(); + writeln!( + out, + " match communication_type_name(self.0) {{" + ) + .unwrap(); + writeln!(out, " Some(name) => f.write_str(name),").unwrap(); + writeln!( + out, + " None => write!(f, \"CommTypeId({{}})\", self.0)," + ) + .unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, "}}").unwrap(); + writeln!(out).unwrap(); + + writeln!(out, "impl std::fmt::Display for DataTypeId {{").unwrap(); + writeln!( + out, + " fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {{" + ) + .unwrap(); + writeln!( + out, + " match data_type_name(self.0) {{" + ) + .unwrap(); + writeln!(out, " Some(name) => f.write_str(name),").unwrap(); + writeln!( + out, + " None => write!(f, \"DataTypeId({{}})\", self.0)," + ) + .unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, " }}").unwrap(); + writeln!(out, "}}").unwrap(); + writeln!(out).unwrap(); +} diff --git a/type-map/src/lib.rs b/type-map/src/lib.rs index 854fe38..91bf33e 100644 --- a/type-map/src/lib.rs +++ b/type-map/src/lib.rs @@ -62,6 +62,16 @@ impl std::fmt::Display for Version { } } +impl Version { + pub fn is_newer_than(&self, other: &Version) -> bool { + self > other + } + + pub fn is_compatible_with(&self, other: &Version) -> bool { + self.0 == other.0 + } +} + /* * A single protocol version's type dictionary. * Compiled into the client, or loaded by the host via the registry. diff --git a/wasm/src/client.rs b/wasm/src/client.rs index 3d60139..8b7603b 100644 --- a/wasm/src/client.rs +++ b/wasm/src/client.rs @@ -159,11 +159,11 @@ fn verify_host_challenge( id: u64, server_challenge: u128, ) -> Result<(), JsValue> { - let sig = match challenge.get_data(DataType::Signature.to_id(tm)) { + let sig = match challenge.get_data(DataType::Signature) { DataValue::Bytes(b) => b.clone(), _ => return Err(js_error("missing host challenge signature")), }; - let pq_sig = match challenge.get_data(DataType::PqSignature.to_id(tm)) { + let pq_sig = match challenge.get_data(DataType::PqSignature) { DataValue::Bytes(b) => b.clone(), _ => vec![], }; @@ -191,14 +191,14 @@ fn verify_host_final( client_nonce: u128, server_challenge: u128, ) -> Result<(), JsValue> { - if *resp.get_data(DataType::ClientNonce.to_id(tm)) != DataValue::UnsignedNumber(client_nonce) { + if *resp.get_data(DataType::ClientNonce) != DataValue::UnsignedNumber(client_nonce) { return Err(js_error("nonce mismatch")); } - let host_sig = match resp.get_data(DataType::Signature.to_id(tm)) { + let host_sig = match resp.get_data(DataType::Signature) { DataValue::Bytes(b) => b.clone(), _ => return Err(js_error("missing host signature")), }; - let host_pq_sig = match resp.get_data(DataType::PqSignature.to_id(tm)) { + let host_pq_sig = match resp.get_data(DataType::PqSignature) { DataValue::Bytes(b) => b.clone(), _ => vec![], }; @@ -393,7 +393,7 @@ impl WasmClient { )); } - if resp_comm.get_data(DataType::Connected.to_id(&tm)) != &DataValue::BoolTrue { + if resp_comm.get_data(DataType::Connected) != &DataValue::BoolTrue { self.set_state(ConnectionState::Disconnected); return Err(js_error("host rejected authentication")); } @@ -412,7 +412,7 @@ impl WasmClient { } // Extract assigned ID - let assigned_id = match resp_comm.get_data(DataType::Id.to_id(&tm)) { + let assigned_id = match resp_comm.get_data(DataType::Id) { DataValue::UnsignedNumber(n) => *n as u64, _ => { self.set_state(ConnectionState::Disconnected); @@ -495,12 +495,12 @@ impl WasmClient { )); } - if resp_comm.get_data(DataType::Connected.to_id(&tm)) != &DataValue::BoolTrue { + if resp_comm.get_data(DataType::Connected) != &DataValue::BoolTrue { self.set_state(ConnectionState::Disconnected); return Err(js_error("host rejected registration")); } - let assigned_id = match resp_comm.get_data(DataType::Id.to_id(&tm)) { + let assigned_id = match resp_comm.get_data(DataType::Id) { DataValue::UnsignedNumber(n) => *n as u64, _ => { self.set_state(ConnectionState::Disconnected); @@ -731,7 +731,7 @@ impl WasmClient { )); } - let server_challenge = match challenge.get_data(DataType::ServerNonce.to_id(tm)) { + let server_challenge = match challenge.get_data(DataType::ServerNonce) { DataValue::UnsignedNumber(n) => *n, _ => { self.set_state(ConnectionState::Disconnected); diff --git a/wasm/src/frame.rs b/wasm/src/frame.rs index e0315b7..04cbcd8 100644 --- a/wasm/src/frame.rs +++ b/wasm/src/frame.rs @@ -254,26 +254,26 @@ pub fn parse_auth_response(response: &[u8]) -> Result { .map_err(|e| js_error(&format!("parse failed: {}", e)))?; let connected = matches!( - comm.get_data(DataType::Connected.to_id(&TypeMap::latest())), + comm.get_data(DataType::Connected)), DataValue::BoolTrue ); - let client_nonce = match comm.get_data(DataType::ClientNonce.to_id(&TypeMap::latest())) { + let client_nonce = match comm.get_data(DataType::ClientNonce)) { DataValue::UnsignedNumber(n) => Some(*n), _ => None, }; - let assigned_id = match comm.get_data(DataType::Id.to_id(&TypeMap::latest())) { + let assigned_id = match comm.get_data(DataType::Id)) { DataValue::UnsignedNumber(n) => Some(*n as u64), _ => None, }; - let timestamp = match comm.get_data(DataType::Timestamp.to_id(&TypeMap::latest())) { + let timestamp = match comm.get_data(DataType::Timestamp)) { DataValue::UnsignedNumber(n) => Some(*n), _ => None, }; - let signature = match comm.get_data(DataType::Signature.to_id(&TypeMap::latest())) { + let signature = match comm.get_data(DataType::Signature)) { DataValue::Bytes(b) => Some(b.clone()), _ => None, }; @@ -376,11 +376,11 @@ mod tests { assert_eq!(cv.get_type(), CommunicationType::Ping.to_id(&tm)); assert_eq!(cv.get_sender(), 42); assert_eq!( - cv.get_data(DataType::Description.to_id(&tm)), + cv.get_data(DataType::Description), &DataValue::Str("test-ping".into()) ); assert_eq!( - cv.get_data(DataType::Timestamp.to_id(&tm)), + cv.get_data(DataType::Timestamp), &DataValue::UnsignedNumber(1234567890) ); } @@ -395,15 +395,15 @@ mod tests { assert_eq!(cv.get_type(), CommunicationType::Ping.to_id(&tm)); assert_eq!(cv.get_sender(), 99); assert_eq!( - cv.get_data(DataType::Description.to_id(&tm)), + cv.get_data(DataType::Description), &DataValue::Str("with-data".into()) ); assert_eq!( - cv.get_data(DataType::Timestamp.to_id(&tm)), + cv.get_data(DataType::Timestamp), &DataValue::UnsignedNumber(555) ); assert_eq!( - cv.get_data(DataType::Id.to_id(&tm)), + cv.get_data(DataType::Id), &DataValue::Bytes(payload.to_vec()) ); }