diff --git a/CONNECTOR.md b/CONNECTOR.md index f91196f..03b98b8 100644 --- a/CONNECTOR.md +++ b/CONNECTOR.md @@ -1,10 +1,10 @@ # Connector -This file documents the Connection and Version Negotiation logic. +This file documents the connection and version negotiation logic. ## Registry -The `registry` crate provides a multi-version `Registry` used by the host for version negotiation. Accessed through the `mtp` facade (requires the `host` feature): +The `registry` module provides a multi-version `Registry` used by the host for version negotiation. Accessed through the `mtp` facade (requires the `host` feature): ```rust use mtp::codec::registry::Registry; @@ -33,11 +33,12 @@ The host creates a QUIC server, manages the registry, and handles version negoti ### Initialization +The host binds to the address from the `mtp_BIND` environment variable (defaults to `::`) on the specified port: + ```rust use mtp::host::{MTPHost, HostConfig}; let config = HostConfig { - ip: "::".into(), port: 4433, tls_fullchain: std::fs::read("cert.pem")?, tls_key: std::fs::read("key.pem")?, @@ -50,35 +51,31 @@ let mut host = MTPHost::new(config).await?; ```rust while let Some(conn) = host.accept().await { - /* - * conn.version is the negotiated version - * conn.codec is a VersionedCodec scoped to that version - * conn.sender / conn.receiver for raw CommunicationValue I/O - */ + // conn.version is the negotiated version + // conn.codec is a VersionedCodec scoped to that version + // conn.sender / conn.receiver for raw CommunicationValue I/O let msg = conn.receiver.receive().await?; - // ... } ``` The host's `accept()` method: 1. Accepts a QUIC connection -2. Reads the first `CommunicationValue` (always encoded with reserved type IDs) -3. Extracts the client's protocol version from `DataType::Version` (wire ID 3) -4. Calls `registry.negotiate(&[client_version])` -5. Returns `None` if the version is unsupported (caller sends `ErrorBadVersion` and disconnects) -6. Returns an `MTPConnection` with the negotiated version otherwise +2. If authentication is required (crypto feature): performs login/register handshake +3. Reads the first `CommunicationValue` (always encoded with reserved type IDs) +4. Extracts the client's protocol version from `DataType::Version` (wire ID 3) +5. Calls `registry.negotiate(&[client_version])` +6. Returns `None` if the version is unsupported +7. Returns an `MTPConnection` with the negotiated version otherwise -### Login/Register Flow +### Login/Register Handshake (crypto feature) -The complete login/register handshake (see design docs) builds on top of `MTPConnection`: +When `require_authentication` is set, the host sends a **greeting** first (host ID, public keys, nonce). The client then responds with either: -1. Client sends `Identification` with version, ID, nonce, signature -2. Host verifies signature via `get_key` callback -3. Host responds with approval + nonces + signature -4. Client verifies response +- **Login** (`CommunicationType::Identification`, ID 15): client ID, nonce, signature +- **Register** (`CommunicationType::Register`, ID 17): public keys, nonce, signature -New clients use the `Register` variant instead, presenting their public key for registration. +The host verifies the client's signature, sends a signed response, and the client verifies the host's signature. --- @@ -94,13 +91,14 @@ let config = ClientConfig { server_cert: None, // or Some(cert_pem_bytes) }; -// Connect (existing client) +// Connect (unauthenticated, existing client) let conn = MTPClient::connect(config, 8765).await?; -/* - * conn.version is the compiled-in PROTOCOL_VERSION - * conn.sender / conn.receiver for I/O - */ +// Authenticated login +let conn = MTPClient::auth_connect(config, 8765, keys, host_pk).await?; + +// Registration (new client) +let conn = MTPClient::auth_register(config, keys, host_pk).await?; ``` The client's `PROTOCOL_VERSION` constant is set by `protocol_version` in `type-maps.yaml` and baked in at compile time. The client never imports the `registry` crate; it only uses `mtp::type_map` for enum types and `mtp::codec` for encoding. @@ -111,26 +109,26 @@ The client's `PROTOCOL_VERSION` constant is set by `protocol_version` in `type-m ``` Client (v2.0) Host (v0.0, v1.0, v2.0) - │ │ - │ QUIC connect │ - │───────────────────────→│ - │ │ - │ CommValue{ Ident. } │ - │ Version → "2.0" │ - │ Id → 8765 │ - │ Nonce → ... │ - │ Signature → ... │ - │───────────────────────→│ - │ │ registry.negotiate(&[Version(2,0)]) - │ │ → Some(Version(2,0)) - │ │ - │ Response │ - │←───────────────────────│ (uses v2.0 TypeMap for encoding) - │ Status, Nonces, │ - │ Signature │ - │ │ - │ (subsequent messages │ - │ use v2.0 TypeMap) │ + | | + | QUIC connect | + |----------------------->| + | | + | CommValue{ Ident. } | + | Version -> "2.0" | + | Id -> 8765 | + | Nonce -> ... | + | Signature -> ... | + |----------------------->| + | | registry.negotiate(&[Version(2,0)]) + | | -> Some(Version(2,0)) + | | + | Response | + |<-----------------------| (uses v2.0 TypeMap for encoding) + | Status, Nonces, | + | Signature | + | | + | (subsequent messages | + | use v2.0 TypeMap) | ``` -If the client sends an unsupported version (e.g. v3.0 when the host only knows up to v2.0), `negotiate` returns `None` and the host sends `ErrorBadVersion` using reserved types. +If the client sends an unsupported version (e.g. v3.0 when the host only knows up to v2.0), `negotiate` returns `None` and the connection is closed. diff --git a/Cargo.toml b/Cargo.toml index 2627783..b09bee7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -33,13 +33,19 @@ mtp-codec = { path = "codec" } mtp-transport = { path = "transport" } # --- optional, behind features --- -mtp-crypto = { path = "crypto", optional = true } +mtp-crypto = { path = "crypto", optional = true, features = [ + "serde", + "mlkem-tls", +] } mtp-host = { path = "host", optional = true } mtp-client = { path = "client", optional = true } [features] default = [] +# Serialization +serde = ["mtp-crypto/serde"] + # Message encryption (AEAD, signatures, KEM, KDF, hashing). # When combined with `host` or `client`, also enables connection authentication. crypto = [ diff --git a/README.md b/README.md index 57e2412..3c4b139 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ -# **M**ethanium **T**ransport **P**rotocol +# Methanium Transport Protocol -**MTP** is a **m**odular **t**ransport **p**rotocol by Methanium. +MTP is a modular transport protocol built on QUIC. It provides version-negotiable type maps, a binary codec, cryptographic primitives (classical and post-quantum), and host/client connection management with mutual authentication. ## Getting Started @@ -22,9 +22,9 @@ The `mtp` umbrella crate re-exports all sub-crates behind feature flags: Core crates (`codec`, `transport`, `common`, `type_map`) are always available. ```rust -use mtp::codec::{encode, decode, DataValue}; -use mtp::type_map::TypeMap; -use mtp::transport::Sender; +use mtp::codec::{CommunicationValue, DataValue}; +use mtp::type_map::{CommunicationType, DataType, TypeMap}; +use mtp::transport::{Sender, Receiver}; #[cfg(feature = "crypto")] use mtp::crypto::ChaCha20Poly1305; @@ -32,38 +32,47 @@ use mtp::crypto::ChaCha20Poly1305; ## Sub-crates -All sub-crates are re-exported through the `mtp` facade and can be referenced -as `mtp::codec`, `mtp::transport`, `mtp::common`, `mtp::type_map`, -`mtp::crypto`, `mtp::host`, `mtp::client`. +All sub-crates are re-exported through the `mtp` facade and can be referenced as `mtp::codec`, `mtp::transport`, `mtp::common`, `mtp::type_map`, `mtp::crypto`, `mtp::host`, `mtp::client`. ### Codec -The Codec crate handles encoding and decoding of MTP packets using Communication Types and Data Types from the Registry. +The codec crate handles binary encoding and decoding of MTP packets using Communication Types and Data Types resolved through the type-map registry. -**Data Values:** -- Container -- Encrypted Container (requires `crypto` feature) -- Signed Integer -- Unsigned Integer +**Data Value types:** +- Container (key-value map of typed entries) +- Encrypted Container (requires `crypto`) +- Signed Container (requires `crypto`) +- SignedEncrypted Container (requires `crypto`) +- Signed Integer (i128) +- Unsigned Integer (u128) - Boolean -- Signed Float +- Float (exponent + mantissa) - String - Array -- Binary (List of Bytes) +- Bytes +- Null -Encoding/decoding uses a `TypeMap` to resolve type names to wire IDs. - -The Codec uses the Crypto crate to encrypt and decrypt Encrypted Containers. - ---- +Encoding and decoding use a `TypeMap` to resolve type names to wire IDs. The `CommunicationValue` struct provides the frame format (type, flags, optional id/sender/receiver, data payload, optional signature). ### Transport -The Transport crate wraps QUIC using `wtransport`. It provides `Sender`/`Receiver` for bidirectional message passing over QUIC streams. +The transport crate wraps QUIC using `wtransport`. It provides `Sender`/`Receiver` for bidirectional message passing over uni-directional QUIC streams. Supports two send modes: persistent stream and single-stream-per-message. + +### Host + +The host crate provides `MTPHost` with built-in version negotiation and optional authenticated login/registration (requires `crypto`). Accepts connections, negotiates protocol version, and returns `MTPConnection` handles. + +### Client + +The client crate provides `MTPClient` that connects to an MTP host. Supports `connect` (unauthenticated), `auth_connect` (login), and `auth_register` (registration) when built with `crypto`. ### Common -Common handles logging and error handling. Personal information reported to Common will be anonymized. +Common defines shared error types (`CodecError`, `CommunicationError`) used across all crates. + +### Type Map + +The type-map build script reads a YAML configuration to generate `CommunicationType` and `DataType` enums at compile time. The runtime crate provides `TypeMap`, `Version`, `CommunicationTypeId`, `DataTypeId`, and the multi-version `Registry` (requires `registry` feature). ### Crypto Stack @@ -77,3 +86,16 @@ Common handles logging and error handling. Personal information reported to Comm | sha2 | No | Standard construction; widely reviewed | | zeroize | No | Simple; widely used | | mlkem-tls | No | mlkem-rs backend unaudited | + +## Test Coverage + +| Crate | Tests | Notes | +| -------------- | ----- | ------------------------------------------ | +| common | 7 | Error type Display, Clone, format | +| type-map | 3 | Registry builtin, negotiate, latest | +| codec | 32 | DataValue roundtrip, accessors, Display, Hash, base64; CommunicationValue frame encode/decode | +| crypto | 20 | AEAD, signatures, KEM, KDF, hash, key types, multi-encrypt | +| transport | 12 | ConnectionHandle state transitions; Policy defaults; SendMode | +| host | 3 | Version extraction from messages | +| client | 7 | ClientConfig; AuthState derives | +| **Total** | **84** | | diff --git a/client/src/lib.rs b/client/src/lib.rs index 4d030d2..0df6e7e 100644 --- a/client/src/lib.rs +++ b/client/src/lib.rs @@ -7,6 +7,7 @@ use mtp_transport::{Policy, Receiver, Sender}; pub struct ClientConfig { pub url: String, pub server_cert: Option>, + pub client_id: u64, } // Established MTP connection with a single negotiated version. @@ -37,10 +38,7 @@ impl MTPClient { * The first message includes the client's protocol version * (a reserved `Version` data entry) so the host can negotiate. */ - pub async fn connect( - config: ClientConfig, - client_id: u64, - ) -> Result { + pub async fn connect(config: ClientConfig) -> Result { let (sender, receiver) = mtp_transport::connect(&config.url, config.server_cert, Policy::default()).await?; @@ -48,7 +46,10 @@ impl MTPClient { let version_str = format!("{}", PROTOCOL_VERSION); let ident = CommunicationValue::new(mtp_codec::CommunicationType::Identification) .add_data(DataTypeId(3), DataValue::Str(version_str)) - .add_data(DataTypeId(6), DataValue::UnsignedNumber(client_id.into())); + .add_data( + DataTypeId(6), + DataValue::UnsignedNumber(config.client_id.into()), + ); sender.send(&ident).await?; @@ -59,7 +60,7 @@ impl MTPClient { #[cfg(feature = "crypto")] auth_state: AuthState::Authenticated, #[cfg(feature = "crypto")] - client_id, + client_id: config.client_id, }) } } @@ -69,50 +70,58 @@ impl MTPClient { impl MTPClient { pub async fn auth_connect( config: ClientConfig, - client_id: u64, - keys: mtp_crypto::KeyGroup, - host_public_key: mtp_crypto::SignaturePublicKey, + keys: &mtp_crypto::Keyring, + host_public_key_bundle: &mtp_crypto::PublicKeyBundle, ) -> Result { - use mtp_crypto::{Ed25519Signer, SignatureScheme, verify_ed25519}; + use mtp_crypto::{ + Ed25519Signer, MlDsaSigner, SignatureScheme, verify_ed25519, verify_ml_dsa, + }; let (sender, receiver) = mtp_transport::connect(&config.url, config.server_cert, Policy::default()).await?; - // 1. Receive host greeting (IdentificationResponse with host ID, public keys, nonce) - let _greeting = receiver.receive().await?; - - // 2. Build and send Identification message with auth data + // 1. Build and send Identification message immediately (no greeting) let client_nonce: u128 = rand::random(); let version_str = format!("{}", PROTOCOL_VERSION); let mut sig_payload = Vec::new(); sig_payload.extend_from_slice(version_str.as_bytes()); - sig_payload.extend_from_slice(&client_id.to_be_bytes()); + sig_payload.extend_from_slice(&config.client_id.to_be_bytes()); sig_payload.extend_from_slice(&client_nonce.to_be_bytes()); - let signer = Ed25519Signer::new(&keys.signature_private_key) + let signer = Ed25519Signer::new(&keys.sig_cl_secret_key) .map_err(|e| CommunicationError::Other(e.to_string()))?; /* ===== Signature ===== */ - // Current: signing concatenated bytes - // Future: sign entire CommunicationValue let signature = signer .sign(&sig_payload) .map_err(|e| CommunicationError::Other(e.to_string()))?; - /* ===== End Signature ===== */ - let ident = CommunicationValue::new(mtp_codec::CommunicationType::Identification) + let mut ident = CommunicationValue::new(mtp_codec::CommunicationType::Identification) .add_typed_default(DataType::Version, DataValue::Str(version_str)) - .add_typed_default(DataType::Id, DataValue::UnsignedNumber(client_id as u128)) + .add_typed_default( + DataType::Id, + DataValue::UnsignedNumber(config.client_id as u128), + ) .add_typed_default( DataType::ClientNonce, DataValue::UnsignedNumber(client_nonce), ) .add_typed_default(DataType::Signature, DataValue::Bytes(signature)); + if !keys.sig_pq_secret_key.as_bytes().is_empty() { + let pq_signer = MlDsaSigner::new(&keys.sig_pq_secret_key, &keys.sig_pq_public_key) + .map_err(|e| CommunicationError::Other(e.to_string()))?; + let pq_signature = pq_signer + .sign(&sig_payload) + .map_err(|e| CommunicationError::Other(e.to_string()))?; + ident = ident.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_signature)); + } + /* ===== End Signature ===== */ + sender.send(&ident).await?; - // 3. Receive host response + // 2. Receive host response (single message) let response = receiver.receive().await?; let connected = response.get_data(DataTypeId(11)); @@ -158,15 +167,36 @@ impl MTPClient { } }; + let host_pq_sig = match response.get_data(DataTypeId(12)) { + DataValue::Bytes(b) => b.clone(), + _ => vec![], + }; + let mut host_sig_payload = Vec::new(); - host_sig_payload.extend_from_slice(b"true"); + host_sig_payload.push(0x01); + host_sig_payload.extend_from_slice(&config.client_id.to_be_bytes()); host_sig_payload.extend_from_slice(&client_nonce.to_be_bytes()); host_sig_payload.extend_from_slice(&host_new_nonce.to_be_bytes()); /* ===== Signature ===== */ - verify_ed25519(&host_public_key, &host_sig_payload, &host_sig).map_err(|_| { - CommunicationError::AuthenticationFailed("Host signature invalid".into()) - })?; + verify_ed25519( + &host_public_key_bundle.sig_cl_public_key, + &host_sig_payload, + &host_sig, + ) + .map_err(|_| CommunicationError::AuthenticationFailed("Host signature invalid".into()))?; + if !host_pq_sig.is_empty() + && verify_ml_dsa( + &host_public_key_bundle.sig_pq_public_key, + &host_sig_payload, + &host_pq_sig, + ) + .is_err() + { + return Err(CommunicationError::AuthenticationFailed( + "Host PQ signature invalid".into(), + )); + } /* ===== End Signature ===== */ Ok(MTPConnection { @@ -174,24 +204,23 @@ impl MTPClient { sender, receiver, auth_state: AuthState::Authenticated, - client_id, + client_id: config.client_id, }) } pub async fn auth_register( config: ClientConfig, - keys: mtp_crypto::Keyring, - host_public_key: mtp_crypto::SignaturePublicKey, + keys: &mtp_crypto::Keyring, + host_public_key_bundle: &mtp_crypto::PublicKeyBundle, ) -> Result { - use mtp_crypto::{Ed25519Signer, SignatureScheme, verify_ed25519}; + use mtp_crypto::{ + Ed25519Signer, MlDsaSigner, SignatureScheme, verify_ed25519, verify_ml_dsa, + }; let (sender, receiver) = mtp_transport::connect(&config.url, config.server_cert, Policy::default()).await?; - // 1. Receive host greeting (IdentificationResponse with host ID, public keys, nonce) - let _greeting = receiver.receive().await?; - - // 2. Build and send Register message with public keys and auth data + // 1. Build and send Register message immediately let client_nonce: u128 = rand::random(); let version_str = format!("{}", PROTOCOL_VERSION); let pk_bundle = keys.public_key_bundle(); @@ -209,9 +238,8 @@ impl MTPClient { let signature = signer .sign(&sig_payload) .map_err(|e| CommunicationError::Other(e.to_string()))?; - /* ===== End Signature ===== */ - let register = CommunicationValue::new(mtp_codec::CommunicationType::Register) + let mut register = CommunicationValue::new(mtp_codec::CommunicationType::Register) .add_typed_default(DataType::Version, DataValue::Str(version_str)) .add_typed_default( DataType::ClientNonce, @@ -220,9 +248,20 @@ impl MTPClient { .add_typed_default(DataType::PublicKeys, DataValue::Bytes(pk_bytes)) .add_typed_default(DataType::Signature, DataValue::Bytes(signature)); + if !keys.sig_pq_secret_key.as_bytes().is_empty() { + let pq_signer = MlDsaSigner::new(&keys.sig_pq_secret_key, &keys.sig_pq_public_key) + .map_err(|e| CommunicationError::Other(e.to_string()))?; + let pq_signature = pq_signer + .sign(&sig_payload) + .map_err(|e| CommunicationError::Other(e.to_string()))?; + register = + register.add_typed_default(DataType::PqSignature, DataValue::Bytes(pq_signature)); + } + /* ===== End Signature ===== */ + sender.send(®ister).await?; - // 3. Receive host response + // 2. Receive host response (single message) let response = receiver.receive().await?; let connected = response.get_data(DataTypeId(11)); @@ -277,16 +316,36 @@ impl MTPClient { } }; + let host_pq_sig = match response.get_data(DataTypeId(12)) { + DataValue::Bytes(b) => b.clone(), + _ => vec![], + }; + let mut host_sig_payload = Vec::new(); - host_sig_payload.extend_from_slice(b"true"); - host_sig_payload.extend_from_slice(&assigned_id.to_be_bytes()); + host_sig_payload.push(0x01); + host_sig_payload.extend_from_slice(&(assigned_id as u64).to_be_bytes()); host_sig_payload.extend_from_slice(&client_nonce.to_be_bytes()); host_sig_payload.extend_from_slice(&host_new_nonce.to_be_bytes()); /* ===== Signature ===== */ - verify_ed25519(&host_public_key, &host_sig_payload, &host_sig).map_err(|_| { - CommunicationError::AuthenticationFailed("Host signature invalid".into()) - })?; + verify_ed25519( + &host_public_key_bundle.sig_cl_public_key, + &host_sig_payload, + &host_sig, + ) + .map_err(|_| CommunicationError::AuthenticationFailed("Host signature invalid".into()))?; + if !host_pq_sig.is_empty() + && verify_ml_dsa( + &host_public_key_bundle.sig_pq_public_key, + &host_sig_payload, + &host_pq_sig, + ) + .is_err() + { + return Err(CommunicationError::AuthenticationFailed( + "Host PQ signature invalid".into(), + )); + } /* ===== End Signature ===== */ Ok(MTPConnection { @@ -298,4 +357,38 @@ impl MTPClient { }) } } -/* ===== End Authentication ===== */ + +/* ================================ TESTS ================================ */ +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_client_config_url() { + let config = ClientConfig { + url: "https://example.com:4433".into(), + server_cert: None, + client_id: 0, + }; + assert_eq!(config.url, "https://example.com:4433"); + assert!(config.server_cert.is_none()); + } + + #[test] + fn test_client_config_with_cert() { + let config = ClientConfig { + url: "https://localhost:4433".into(), + server_cert: Some(vec![0x01, 0x02, 0x03]), + client_id: 42, + }; + assert_eq!(config.server_cert, Some(vec![0x01, 0x02, 0x03])); + assert_eq!(config.client_id, 42); + } + + #[cfg(feature = "crypto")] + #[test] + fn test_auth_state_derive() { + assert_eq!(AuthState::Pending, AuthState::Pending); + assert_ne!(AuthState::Authenticated, AuthState::Failed); + } +} diff --git a/codec/src/communication_value.rs b/codec/src/communication_value.rs index db3733c..0721398 100644 --- a/codec/src/communication_value.rs +++ b/codec/src/communication_value.rs @@ -7,10 +7,21 @@ use crate::data_value::DataValue; use crate::rand_u32; use mtp_common::CodecError; use mtp_type_map::{ - communication_type_name, data_type_name, CommunicationType, CommunicationTypeId, DataType, - DataTypeId, TypeMap, PROTOCOL_VERSION, + CommunicationType, CommunicationTypeId, DataType, DataTypeId, PROTOCOL_VERSION, TypeMap, + communication_type_name, data_type_name, }; +#[cfg(feature = "crypto")] +use mtp_crypto::SigAlgorithm; +#[cfg(feature = "crypto")] +use mtp_crypto::SignatureScheme; + +const FLAG_HAS_SENDER: u8 = 0b0000_0001; +const FLAG_HAS_RECEIVER: u8 = 0b0000_0010; +const FLAG_HAS_ID: u8 = 0b0000_0100; +const FLAG_ENCRYPTED: u8 = 0b0000_1000; +const FLAG_SIGNED: u8 = 0b0001_0000; + #[derive(Debug, Clone, PartialEq, Eq)] pub struct CommunicationValue { id: u32, @@ -19,6 +30,8 @@ pub struct CommunicationValue { receiver: u64, data: BTreeMap, type_map: Option, + #[cfg(feature = "crypto")] + frame_signature: Option<(u8, Vec)>, } impl CommunicationValue { @@ -33,6 +46,8 @@ impl CommunicationValue { receiver: 0, data: BTreeMap::new(), type_map: Some(tm), + #[cfg(feature = "crypto")] + frame_signature: None, } } @@ -47,6 +62,8 @@ impl CommunicationValue { receiver: 0, data: BTreeMap::new(), type_map: Some(tm.clone()), + #[cfg(feature = "crypto")] + frame_signature: None, } } @@ -107,63 +124,116 @@ impl CommunicationValue { pub fn get_data(&self, data_type: DataTypeId) -> &DataValue { self.data.get(&data_type).unwrap_or(&DataValue::Null) } - } impl CommunicationValue { /* * Frame format (strict new format): - * [4 bytes u32 total_length] // number of bytes after this field + * [4 bytes u32 total_length] // number of bytes after this field * [2 bytes u16 communication_type] * [1 byte flags] - * [optional 4 bytes id] // if flags bit2 set - * [optional 6 bytes sender] // if flags bit0 set - * [optional 6 bytes receiver]// if flags bit1 set + * [optional 4 bytes id] // if flags bit2 set + * [optional 6 bytes sender] // if flags bit0 set + * [optional 6 bytes receiver] // if flags bit1 set + * [optional 1 byte signature type] // if flags bit4 set; Type defines length of signature + * [optional signature] // if flags bit4 set * [data container bytes...] * * Flags: * bit0 => has sender * bit1 => has receiver * bit2 => has id + * bit3 => is data encrypted If so data bytes will be an encrypted container + * bit4 => is communication value signed */ pub fn to_bytes(&self) -> Vec { - let mut payload = Vec::new(); - let has_sender = self.sender != 0; let has_receiver = self.receiver != 0; let has_id = self.id != 0; + #[cfg(feature = "crypto")] + let is_encrypted = self.data.len() == 1 && self.data.values().any(|v| { + matches!( + v, + DataValue::EncryptedContainer(_) | DataValue::SignedEncryptedContainer(_) + ) + }); + #[cfg(not(feature = "crypto"))] + let is_encrypted = false; + + #[cfg(feature = "crypto")] + let has_frame_sig = self.frame_signature.is_some(); + #[cfg(not(feature = "crypto"))] + let has_frame_sig = false; + let mut flags: u8 = 0; if has_sender { - flags |= 0b0000_0001; + flags |= FLAG_HAS_SENDER; } if has_receiver { - flags |= 0b0000_0010; + flags |= FLAG_HAS_RECEIVER; } if has_id { - flags |= 0b0000_0100; + flags |= FLAG_HAS_ID; + } + if is_encrypted { + flags |= FLAG_ENCRYPTED; + } + if has_frame_sig { + flags |= FLAG_SIGNED; } - let _ = payload.write_u16::(self.comm_type.0); - payload.push(flags); + let mut metadata = Vec::new(); + let _ = metadata.write_u16::(self.comm_type.0); + metadata.push(flags); if has_id { - let _ = payload.write_u32::(self.id); + let _ = metadata.write_u32::(self.id); } if has_sender { let sender_be = self.sender.to_be_bytes(); - payload.extend_from_slice(&sender_be[2..]); // 6 bytes + metadata.extend_from_slice(&sender_be[2..]); } if has_receiver { let receiver_be = self.receiver.to_be_bytes(); - payload.extend_from_slice(&receiver_be[2..]); // 6 bytes + metadata.extend_from_slice(&receiver_be[2..]); } - let container_value = DataValue::container_from_map(&self.data); - let container_bytes = container_value.to_bytes(); - payload.extend_from_slice(&container_bytes); + #[cfg(feature = "crypto")] + let data_bytes = if is_encrypted { + self.data + .values() + .find_map(|v| match v { + DataValue::EncryptedContainer(ct) => Some(ct.clone()), + DataValue::SignedEncryptedContainer(ct) => Some(ct.clone()), + _ => None, + }) + .unwrap_or_default() + } else { + let container_value = DataValue::container_from_map(&self.data); + container_value.to_bytes() + }; + + #[cfg(not(feature = "crypto"))] + let data_bytes = { + let container_value = DataValue::container_from_map(&self.data); + container_value.to_bytes() + }; + + let mut payload = Vec::new(); + payload.extend_from_slice(&metadata); + + #[cfg(feature = "crypto")] + if let Some((_alg, _sig)) = &self.frame_signature { + // algorithm and signature are computed by sign_frame() and stored. + // The frame bytes are built by using the pre-computed signature. + payload.push(*_alg); + payload.extend_from_slice(_sig); + } + + payload.extend_from_slice(&data_bytes); let mut frame = Vec::with_capacity(4 + payload.len()); let _ = frame.write_u32::(payload.len() as u32); @@ -189,12 +259,17 @@ impl CommunicationValue { .map_err(|_| CodecError::InvalidEncoding)?; let comm_type = CommunicationTypeId(comm_type_num); - let flags = cursor - .read_u8() - .map_err(|_| CodecError::InvalidEncoding)?; - let has_sender = (flags & 0b0000_0001) != 0; - let has_receiver = (flags & 0b0000_0010) != 0; - let has_id = (flags & 0b0000_0100) != 0; + let flags = cursor.read_u8().map_err(|_| CodecError::InvalidEncoding)?; + let has_sender = (flags & FLAG_HAS_SENDER) != 0; + let has_receiver = (flags & FLAG_HAS_RECEIVER) != 0; + let has_id = (flags & FLAG_HAS_ID) != 0; + let is_encrypted = (flags & FLAG_ENCRYPTED) != 0; + let is_signed = (flags & FLAG_SIGNED) != 0; + + #[cfg(not(feature = "crypto"))] + if is_signed || is_encrypted { + return Err(CodecError::InvalidEncoding); + } let id = if has_id { cursor @@ -224,14 +299,47 @@ impl CommunicationValue { 0 }; + #[cfg(feature = "crypto")] + let frame_signature = if is_signed { + let alg = cursor.read_u8().map_err(|_| CodecError::InvalidEncoding)?; + let sig_len = + SigAlgorithm::length(alg).ok_or(CodecError::InvalidEncoding)?; + let mut sig = vec![0u8; sig_len]; + cursor + .read_exact(&mut sig) + .map_err(|_| CodecError::InvalidEncoding)?; + Some((alg, sig)) + } else { + None + }; + let pos = cursor.position() as usize; if pos > frame_end { return Err(CodecError::InvalidEncoding); } let data_bytes = &bytes[pos..frame_end]; - let data_value = DataValue::from_bytes(data_bytes).ok_or(CodecError::InvalidEncoding)?; - let data = data_value.as_map().ok_or(CodecError::InvalidEncoding)?; + + #[cfg(feature = "crypto")] + let data = if is_encrypted { + let mut map = BTreeMap::new(); + map.insert( + DataTypeId(0), + DataValue::EncryptedContainer(data_bytes.to_vec()), + ); + map + } else { + let data_value = + DataValue::from_bytes(data_bytes).ok_or(CodecError::InvalidEncoding)?; + data_value.as_map().ok_or(CodecError::InvalidEncoding)? + }; + + #[cfg(not(feature = "crypto"))] + let data = { + let data_value = + DataValue::from_bytes(data_bytes).ok_or(CodecError::InvalidEncoding)?; + data_value.as_map().ok_or(CodecError::InvalidEncoding)? + }; Ok(Self { id, @@ -240,6 +348,8 @@ impl CommunicationValue { receiver, data, type_map: None, + #[cfg(feature = "crypto")] + frame_signature, }) } @@ -249,15 +359,129 @@ impl CommunicationValue { Ok(val) } + /* + * Sign the frame. Computes a signature over the canonical form: + * comm_type || flags || id? || sender? || receiver? || data_bytes + * + * After calling this, `to_bytes()` will embed the algorithm and + * signature before the data payload. + */ + #[cfg(feature = "crypto")] + pub fn sign_frame( + &mut self, + algorithm: u8, + signer: &impl SignatureScheme, + ) -> Option<()> { + let signed_payload = self.build_signed_payload(); + let sig = signer.sign(&signed_payload).ok()?; + self.frame_signature = Some((algorithm, sig)); + Some(()) + } + + /* + * Verify the frame signature. Reconstructs the signed payload from + * current state and checks it against the stored signature. + */ + #[cfg(feature = "crypto")] + pub fn verify_frame(&self, verifier: &impl SignatureScheme) -> Result<(), CodecError> { + let (_algorithm, sig) = self + .frame_signature + .as_ref() + .ok_or(CodecError::InvalidEncoding)?; + + let signed_payload = self.build_signed_payload(); + verifier + .verify(&signed_payload, sig) + .map_err(|_| CodecError::InvalidEncoding) + } + + /* + * Reconstruct the signed payload that the frame signature covers: + * comm_type || flags || id? || sender? || receiver? || data_bytes + */ + #[cfg(feature = "crypto")] + fn build_signed_payload(&self) -> Vec { + let has_sender = self.sender != 0; + let has_receiver = self.receiver != 0; + let has_id = self.id != 0; + + let is_encrypted = self.data.len() == 1 && self.data.values().any(|v| { + matches!( + v, + DataValue::EncryptedContainer(_) | DataValue::SignedEncryptedContainer(_) + ) + }); + + let mut flags: u8 = 0; + if has_sender { + flags |= FLAG_HAS_SENDER; + } + if has_receiver { + flags |= FLAG_HAS_RECEIVER; + } + if has_id { + flags |= FLAG_HAS_ID; + } + if is_encrypted { + flags |= FLAG_ENCRYPTED; + } + if self.frame_signature.is_some() { + flags |= FLAG_SIGNED; + } + + let mut metadata = Vec::new(); + let _ = metadata.write_u16::(self.comm_type.0); + metadata.push(flags); + + if has_id { + let _ = metadata.write_u32::(self.id); + } + + if has_sender { + let sender_be = self.sender.to_be_bytes(); + metadata.extend_from_slice(&sender_be[2..]); + } + + if has_receiver { + let receiver_be = self.receiver.to_be_bytes(); + metadata.extend_from_slice(&receiver_be[2..]); + } + + #[cfg(feature = "crypto")] + let data_bytes = if is_encrypted { + self.data + .values() + .find_map(|v| match v { + DataValue::EncryptedContainer(ct) => Some(ct.clone()), + DataValue::SignedEncryptedContainer(ct) => Some(ct.clone()), + _ => None, + }) + .unwrap_or_default() + } else { + let container_value = DataValue::container_from_map(&self.data); + container_value.to_bytes() + }; + + #[cfg(not(feature = "crypto"))] + let data_bytes = { + let container_value = DataValue::container_from_map(&self.data); + container_value.to_bytes() + }; + + [metadata, data_bytes].concat() + } + + #[cfg(feature = "crypto")] + pub fn get_frame_signature(&self) -> Option<&(u8, Vec)> { + self.frame_signature.as_ref() + } + #[cfg(feature = "registry")] pub fn migrate(&self, target_tm: &TypeMap) -> Result { - let comm_name = - communication_type_name(self.comm_type.0).ok_or_else(|| { - CodecError::UnknownCommunicationType(self.comm_type.0.to_string()) - })?; - let comm_variant = CommunicationType::from_name(comm_name).ok_or_else(|| { - CodecError::UnknownCommunicationType(comm_name.to_string()) - })?; + let comm_name = communication_type_name(self.comm_type.0) + .ok_or_else(|| CodecError::UnknownCommunicationType(self.comm_type.0.to_string()))?; + let comm_variant = CommunicationType::from_name(comm_name) + .ok_or_else(|| CodecError::UnknownCommunicationType(comm_name.to_string()))?; let new_comm_id = CommunicationTypeId( target_tm .comm_id_enum(comm_variant) @@ -285,6 +509,8 @@ impl CommunicationValue { receiver: self.receiver, data: new_data, type_map: Some(target_tm.clone()), + #[cfg(feature = "crypto")] + frame_signature: self.frame_signature.clone(), }) } } diff --git a/codec/src/data_value.rs b/codec/src/data_value.rs index bbe9fe4..9ad52c0 100644 --- a/codec/src/data_value.rs +++ b/codec/src/data_value.rs @@ -9,7 +9,7 @@ use std::io::Cursor; use mtp_type_map::DataTypeId; #[cfg(feature = "crypto")] -use mtp_crypto::{AeadDecrypt, AeadEncrypt}; +use mtp_crypto::{AeadDecrypt, AeadEncrypt, SigAlgorithm, SignatureScheme}; #[derive(Debug, Clone, PartialEq, Eq)] pub enum DataKind { @@ -27,6 +27,10 @@ pub enum DataKind { #[cfg(feature = "crypto")] EncryptedContainer, + #[cfg(feature = "crypto")] + SignedContainer, + #[cfg(feature = "crypto")] + SignedEncryptedContainer, Null, } @@ -44,12 +48,47 @@ pub enum DataValue { Str(String), Bytes(Vec), Array(Vec), - + /* + * Container format: + * [2 bytes u16 entry_count] // number of entries + * [1 byte kind] // DataValue kind marker + * [if kind == BOOL_TRUE or BOOL_FALSE:] + * [2 bytes u16 key] // DataTypeId discriminant + * [else:] + * [4 bytes u32 payload_len] // length of the value payload + * [2 bytes u16 key] // DataTypeId discriminant + * [payload_len bytes payload] // value data (interpreted based on kind) + */ Container(Vec<(DataTypeId, DataValue)>), + /* + * Container format: + * [4 bytes u32 entry_count] // length of the container + * [binary data] + * -> After decryption, the container is parsed as a regular container + */ #[cfg(feature = "crypto")] EncryptedContainer(Vec), + /* + * Container format: + * [4 bytes u32 entry_count] // length of the container + * [binary data] + * -> Can be turned into Container + * -> Can be used with a public key to verify integrity + */ + #[cfg(feature = "crypto")] + SignedContainer(Vec), + + /* + * Container format: + * [4 bytes u32 entry_count] // length of the container + * [binary data] + * -> After decryption, the container is parsed as a signed container + */ + #[cfg(feature = "crypto")] + SignedEncryptedContainer(Vec), + Null, } @@ -86,7 +125,9 @@ impl DataValue { * 0x08 => Array * 0x09 => Container * 0x0A => EncryptedContainer (4 bytes u32 len + encrypted bytes) - * 0x0B => Null + * 0x0B => SignedContainer (4 bytes u32 len + 3373 bytes signature) + * 0x0C => SignedEncryptedContainer (4 bytes u32 len + 3373 bytes signature + encrypted bytes) + * 0xFF => Null */ const KIND_BOOL_TRUE: u8 = 0x01; const KIND_BOOL_FALSE: u8 = 0x02; @@ -102,8 +143,12 @@ impl DataValue { const KIND_CONTAINER: u8 = 0x09; #[cfg(feature = "crypto")] const KIND_ENCRYPTED_CONTAINER: u8 = 0x0A; + #[cfg(feature = "crypto")] + const KIND_SIGNED_CONTAINER: u8 = 0x0B; + #[cfg(feature = "crypto")] + const KIND_SIGNED_ENCRYPTED_CONTAINER: u8 = 0x0C; - const KIND_NULL: u8 = 0x0B; + const KIND_NULL: u8 = 0xFF; pub fn container_from_map(map: &BTreeMap) -> DataValue { let mut container = Vec::new(); @@ -131,6 +176,10 @@ impl DataValue { DataValue::Container(_) => DataKind::Container, #[cfg(feature = "crypto")] DataValue::EncryptedContainer(_) => DataKind::EncryptedContainer, + #[cfg(feature = "crypto")] + DataValue::SignedContainer(_) => DataKind::SignedContainer, + #[cfg(feature = "crypto")] + DataValue::SignedEncryptedContainer(_) => DataKind::SignedEncryptedContainer, DataValue::Null => DataKind::Null, } } @@ -205,17 +254,29 @@ impl DataValue { } } + #[cfg(feature = "crypto")] + pub fn as_signed_container(&self) -> Option> { + match self { + DataValue::SignedContainer(b) => Some(b.clone()), + _ => None, + } + } + + #[cfg(feature = "crypto")] + pub fn as_signed_encrypted_container(&self) -> Option> { + match self { + DataValue::SignedEncryptedContainer(c) => Some(c.clone()), + _ => None, + } + } + /* * Decrypt an `EncryptedContainer` in-place, replacing it with the * deserialized `Container`. Returns `None` if decryption or * deserialization fails. */ #[cfg(feature = "crypto")] - pub fn decrypt_into_container( - &mut self, - cipher: &impl AeadDecrypt, - aad: &[u8], - ) -> Option<()> { + pub fn decrypt_into_container(&mut self, cipher: &impl AeadDecrypt, aad: &[u8]) -> Option<()> { let data = self.as_encrypted_container()?; let plaintext = cipher.decrypt(&data, aad).ok()?; let dv = DataValue::from_bytes(&plaintext)?; @@ -233,11 +294,7 @@ impl DataValue { * Returns `None` if the value is not a `Container` or encryption fails. */ #[cfg(feature = "crypto")] - pub fn encrypt_container( - &mut self, - cipher: &impl AeadEncrypt, - aad: &[u8], - ) -> Option<()> { + pub fn encrypt_container(&mut self, cipher: &impl AeadEncrypt, aad: &[u8]) -> Option<()> { let entries = self.as_container()?; let plaintext = DataValue::Container(entries).to_bytes(); let ct = cipher.encrypt(&plaintext, aad).ok()?; @@ -245,6 +302,98 @@ impl DataValue { Some(()) } + /* + * Sign a `Container` in-place, replacing it with a `SignedContainer`. + * The wire blob is: [1 byte alg] [N bytes sig] [serialized container bytes]. + * The signature covers only the serialized container bytes (not the alg byte). + * Returns `None` if the value is not a `Container` or signing fails. + */ + #[cfg(feature = "crypto")] + pub fn sign_container( + &mut self, + algorithm: u8, + signer: &impl SignatureScheme, + ) -> Option<()> { + let entries = self.as_container()?; + let container_bytes = Self::encode_container(&entries); + + let sig = signer.sign(&container_bytes).ok()?; + + let mut blob = Vec::with_capacity(1 + sig.len() + container_bytes.len()); + blob.push(algorithm); + blob.extend_from_slice(&sig); + blob.extend_from_slice(&container_bytes); + + *self = DataValue::SignedContainer(blob); + Some(()) + } + + /* + * Verify a `SignedContainer` in-place, replacing it with the deserialized + * `Container` on success. Returns `None` if verification fails or the + * blob is malformed. + */ + #[cfg(feature = "crypto")] + pub fn verify_into_container( + &mut self, + verifier: &impl SignatureScheme, + ) -> Option<()> { + let blob = self.as_signed_container()?; + if blob.len() < 1 + 64 + 2 { + return None; + } + + let algorithm = blob[0]; + let sig_len = SigAlgorithm::length(algorithm)?; + if blob.len() < 1 + sig_len + 2 { + return None; + } + + let signature = &blob[1..1 + sig_len]; + let container_bytes = &blob[1 + sig_len..]; + + verifier.verify(container_bytes, signature).ok()?; + + let entries = DataValue::from_bytes(container_bytes)?.as_container()?; + *self = DataValue::Container(entries); + Some(()) + } + + /* + * Encrypt a `Container` into a `SignedEncryptedContainer` in-place. + * The result is an opaque ciphertext that decrypts to a `SignedContainer`. + */ + #[cfg(feature = "crypto")] + pub fn sign_and_encrypt_container( + &mut self, + algorithm: u8, + signer: &impl SignatureScheme, + cipher: &impl AeadEncrypt, + aad: &[u8], + ) -> Option<()> { + self.sign_container(algorithm, signer)?; + let blob = self.as_signed_container()?; + let ct = cipher.encrypt(&blob, aad).ok()?; + *self = DataValue::SignedEncryptedContainer(ct); + Some(()) + } + + /* + * Decrypt a `SignedEncryptedContainer` in-place, replacing it with a + * `SignedContainer`. Does NOT verify — call `verify_into_container` next. + */ + #[cfg(feature = "crypto")] + pub fn decrypt_signed_encrypted_container( + &mut self, + cipher: &impl AeadDecrypt, + aad: &[u8], + ) -> Option<()> { + let data = self.as_signed_encrypted_container()?; + let plaintext = cipher.decrypt(&data, aad).ok()?; + *self = DataValue::SignedContainer(plaintext); + Some(()) + } + pub fn as_map(&self) -> Option> { match self { DataValue::Container(c) => { @@ -311,7 +460,8 @@ impl DataValue { let kind = Self::kind_marker(value); buf.push(kind); - if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL { + if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL + { let _ = buf.write_u16::(key.0); return true; } @@ -351,7 +501,8 @@ impl DataValue { let kind = Self::kind_marker(value); buf.push(kind); - if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL { + if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL + { return true; } @@ -414,6 +565,16 @@ impl DataValue { buf.extend_from_slice(data); Some(()) } + #[cfg(feature = "crypto")] + DataValue::SignedContainer(data) => { + buf.extend_from_slice(data); + Some(()) + } + #[cfg(feature = "crypto")] + DataValue::SignedEncryptedContainer(data) => { + buf.extend_from_slice(data); + Some(()) + } DataValue::Null => Some(()), } @@ -448,7 +609,10 @@ impl DataValue { for _ in 0..count { let kind = cursor.read_u8().ok()?; - if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL { + if kind == Self::KIND_BOOL_TRUE + || kind == Self::KIND_BOOL_FALSE + || kind == Self::KIND_NULL + { let key = DataTypeId(cursor.read_u16::().ok()?); let value = if kind == Self::KIND_BOOL_TRUE { DataValue::BoolTrue @@ -491,7 +655,10 @@ impl DataValue { for _ in 0..count { let kind = cursor.read_u8().ok()?; - if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL { + if kind == Self::KIND_BOOL_TRUE + || kind == Self::KIND_BOOL_FALSE + || kind == Self::KIND_NULL + { let value = if kind == Self::KIND_BOOL_TRUE { DataValue::BoolTrue } else if kind == Self::KIND_BOOL_FALSE { @@ -609,9 +776,33 @@ impl DataValue { cursor.set_position(end as u64); Some(DataValue::EncryptedContainer(data)) } + #[cfg(feature = "crypto")] + Self::KIND_SIGNED_CONTAINER => { + let len = payload_len?; + let start = cursor.position() as usize; + let end = start.checked_add(len)?; + if end > cursor.get_ref().len() { + return None; + } + let data = cursor.get_ref()[start..end].to_vec(); + cursor.set_position(end as u64); + Some(DataValue::SignedContainer(data)) + } + #[cfg(feature = "crypto")] + Self::KIND_SIGNED_ENCRYPTED_CONTAINER => { + let len = payload_len?; + let start = cursor.position() as usize; + let end = start.checked_add(len)?; + if end > cursor.get_ref().len() { + return None; + } + let data = cursor.get_ref()[start..end].to_vec(); + cursor.set_position(end as u64); + Some(DataValue::SignedEncryptedContainer(data)) + } Self::KIND_NULL => Some(DataValue::Null), #[cfg(not(feature = "crypto"))] - 0x0A => None, + 0x0A | 0x0B | 0x0C => None, _ => None, } } @@ -636,6 +827,10 @@ impl DataValue { DataValue::Container(_) => Self::KIND_CONTAINER, #[cfg(feature = "crypto")] DataValue::EncryptedContainer(_) => Self::KIND_ENCRYPTED_CONTAINER, + #[cfg(feature = "crypto")] + DataValue::SignedContainer(_) => Self::KIND_SIGNED_CONTAINER, + #[cfg(feature = "crypto")] + DataValue::SignedEncryptedContainer(_) => Self::KIND_SIGNED_ENCRYPTED_CONTAINER, DataValue::Null => Self::KIND_NULL, } } @@ -674,6 +869,10 @@ impl fmt::Display for DataValue { DataValue::Bytes(_) => write!(f, "(Binary)"), #[cfg(feature = "crypto")] DataValue::EncryptedContainer(_) => write!(f, "(Secure)"), + #[cfg(feature = "crypto")] + DataValue::SignedContainer(_) => write!(f, "(Signed)"), + #[cfg(feature = "crypto")] + DataValue::SignedEncryptedContainer(_) => write!(f, "(SignedSecure)"), DataValue::Null => write!(f, "null"), } } @@ -696,6 +895,10 @@ impl PartialEq for DataValue { (Container(a), Container(b)) => a == b, #[cfg(feature = "crypto")] (EncryptedContainer(a), EncryptedContainer(b)) => a == b, + #[cfg(feature = "crypto")] + (SignedContainer(a), SignedContainer(b)) => a == b, + #[cfg(feature = "crypto")] + (SignedEncryptedContainer(a), SignedEncryptedContainer(b)) => a == b, (Null, Null) => true, _ => false, } @@ -748,9 +951,408 @@ impl Hash for DataValue { 6u8.hash(state); c.hash(state); } - Null => { + #[cfg(feature = "crypto")] + SignedContainer(c) => { 7u8.hash(state); + c.hash(state); + } + #[cfg(feature = "crypto")] + SignedEncryptedContainer(c) => { + 8u8.hash(state); + c.hash(state); + } + Null => { + 9u8.hash(state); } } } } + +/* ================================ TESTS ================================ */ +#[cfg(test)] +mod tests { + use super::*; + + /// Only Container and Array can be top-level serialized forms. + /// Scalars must be tested inside a container. + fn container_roundtrip(values: Vec<(DataTypeId, DataValue)>) { + let dv = DataValue::Container(values.clone()); + let bytes = dv.to_bytes(); + let decoded = DataValue::from_bytes(&bytes).expect("roundtrip failed"); + assert_eq!(dv, decoded, "container roundtrip mismatch"); + } + + fn array_roundtrip(values: Vec) { + let dv = DataValue::Array(values.clone()); + let bytes = dv.to_bytes(); + let decoded = DataValue::from_bytes(&bytes).expect("roundtrip failed"); + assert_eq!(dv, decoded, "array roundtrip mismatch"); + } + + #[test] + fn test_bool_in_container() { + container_roundtrip(vec![ + (DataTypeId(1), DataValue::BoolTrue), + (DataTypeId(2), DataValue::BoolFalse), + ]); + } + + #[test] + fn test_bool_true_eq() { + assert_eq!(DataValue::BoolTrue, DataValue::Bool(true)); + assert_eq!(DataValue::BoolFalse, DataValue::Bool(false)); + assert_ne!(DataValue::BoolTrue, DataValue::Bool(false)); + } + + #[test] + fn test_bool_as_bool() { + assert_eq!(DataValue::BoolTrue.as_bool(), Some(true)); + assert_eq!(DataValue::BoolFalse.as_bool(), Some(false)); + assert_eq!(DataValue::Bool(true).as_bool(), Some(true)); + assert_eq!(DataValue::Null.as_bool(), None); + } + + #[test] + fn test_signed_number_in_container() { + container_roundtrip(vec![ + (DataTypeId(1), DataValue::SignedNumber(0)), + (DataTypeId(2), DataValue::SignedNumber(42)), + (DataTypeId(3), DataValue::SignedNumber(-42)), + (DataTypeId(4), DataValue::SignedNumber(i128::MAX)), + (DataTypeId(5), DataValue::SignedNumber(i128::MIN)), + ]); + } + + #[test] + fn test_unsigned_number_in_container() { + container_roundtrip(vec![ + (DataTypeId(1), DataValue::UnsignedNumber(0)), + (DataTypeId(2), DataValue::UnsignedNumber(42)), + (DataTypeId(3), DataValue::UnsignedNumber(u128::MAX)), + ]); + } + + #[test] + fn test_float_in_container() { + container_roundtrip(vec![ + (DataTypeId(1), DataValue::Float(0, 0)), + (DataTypeId(2), DataValue::Float(2, 12345)), + (DataTypeId(3), DataValue::Float(255, 4294967295)), + ]); + } + + #[test] + fn test_str_in_container() { + container_roundtrip(vec![ + (DataTypeId(1), DataValue::Str(String::new())), + (DataTypeId(2), DataValue::Str("hello".to_string())), + (DataTypeId(3), DataValue::Str("a".repeat(1000))), + ]); + } + + #[test] + fn test_bytes_in_container() { + container_roundtrip(vec![ + (DataTypeId(1), DataValue::Bytes(vec![])), + (DataTypeId(2), DataValue::Bytes(vec![0x00, 0xFF, 0xAB])), + (DataTypeId(3), DataValue::Bytes(vec![0x42; 100])), + ]); + } + + #[test] + fn test_null_in_container() { + container_roundtrip(vec![ + (DataTypeId(1), DataValue::Null), + ]); + } + + #[test] + fn test_array_non_empty_roundtrip() { + array_roundtrip(vec![ + DataValue::BoolTrue, + DataValue::SignedNumber(42), + DataValue::Str("hello".to_string()), + DataValue::Null, + ]); + } + + #[test] + fn test_array_nested_roundtrip() { + array_roundtrip(vec![ + DataValue::Array(vec![DataValue::BoolTrue, DataValue::BoolFalse]), + DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]), + ]); + } + + #[test] + fn test_container_empty_roundtrip() { + container_roundtrip(vec![]); + } + + #[test] + fn test_container_mixed_roundtrip() { + container_roundtrip(vec![ + (DataTypeId(1), DataValue::BoolTrue), + (DataTypeId(2), DataValue::SignedNumber(-100)), + (DataTypeId(3), DataValue::Str("test".to_string())), + (DataTypeId(4), DataValue::UnsignedNumber(u128::MAX)), + (DataTypeId(5), DataValue::Null), + ]); + } + + #[test] + fn test_container_nested_roundtrip() { + container_roundtrip(vec![ + (DataTypeId(1), DataValue::Container(vec![ + (DataTypeId(10), DataValue::BoolTrue), + ])), + (DataTypeId(2), DataValue::Array(vec![ + DataValue::SignedNumber(1), + DataValue::SignedNumber(2), + ])), + ]); + } + + #[test] + fn test_container_base64_roundtrip() { + let dv = DataValue::Container(vec![ + (DataTypeId(7), DataValue::Bytes(vec![0xDE, 0xAD, 0xBE, 0xEF])), + ]); + let b64 = dv.to_base64(); + let decoded = DataValue::from_base64(&b64).expect("base64 roundtrip failed"); + assert_eq!(dv, decoded); + } + + #[test] + fn test_kind_classification() { + assert_eq!(DataValue::BoolTrue.kind(), DataKind::Bool); + assert_eq!(DataValue::Bool(false).kind(), DataKind::Bool); + assert_eq!(DataValue::SignedNumber(0).kind(), DataKind::SignedNumber); + assert_eq!(DataValue::UnsignedNumber(0).kind(), DataKind::UnsignedNumber); + assert_eq!(DataValue::Float(0, 0).kind(), DataKind::Float); + assert_eq!(DataValue::Str(String::new()).kind(), DataKind::Str); + assert_eq!(DataValue::Bytes(vec![]).kind(), DataKind::Bytes); + assert_eq!(DataValue::Array(vec![]).kind(), DataKind::Array(Box::new(DataKind::Null))); + assert_eq!(DataValue::Container(vec![]).kind(), DataKind::Container); + assert_eq!(DataValue::Null.kind(), DataKind::Null); + } + + #[test] + fn test_as_accessors() { + let dv = DataValue::Container(vec![ + (DataTypeId(1), DataValue::Str("alice".to_string())), + (DataTypeId(2), DataValue::SignedNumber(42)), + (DataTypeId(3), DataValue::Bytes(vec![0x01, 0x02])), + (DataTypeId(4), DataValue::Array(vec![DataValue::BoolTrue])), + ]); + + let map = dv.as_map().expect("should be a container"); + assert_eq!(map.get(&DataTypeId(1)).and_then(|v| v.as_str()), Some("alice")); + assert_eq!(map.get(&DataTypeId(2)).and_then(|v| v.as_signed_number()), Some(42)); + assert_eq!(map.get(&DataTypeId(3)).and_then(|v| v.as_bytes()), Some(vec![0x01, 0x02])); + assert_eq!(map.get(&DataTypeId(4)).and_then(|v| v.as_array()), Some(vec![DataValue::BoolTrue])); + } + + #[test] + fn test_as_string() { + let dv = DataValue::Str("hello".to_string()); + assert_eq!(dv.as_string(), Some("hello".to_string())); + assert_eq!(dv.as_str(), Some("hello")); + assert_eq!(DataValue::Null.as_string(), None); + } + + #[test] + fn test_as_float() { + assert_eq!(DataValue::Float(3, 14).as_float(), Some((3, 14))); + assert_eq!(DataValue::Null.as_float(), None); + } + + #[test] + fn test_container_from_map() { + let mut map = BTreeMap::new(); + map.insert(DataTypeId(1), DataValue::BoolTrue); + map.insert(DataTypeId(2), DataValue::SignedNumber(99)); + let dv = DataValue::container_from_map(&map); + let container = dv.as_container().expect("should be container"); + assert_eq!(container.len(), 2); + } + + #[test] + fn test_invalid_short_input() { + assert!(DataValue::from_bytes(&[]).is_none()); + assert!(DataValue::from_bytes(&[0x01]).is_none()); + } + + #[test] + fn test_invalid_kind_rejected() { + let bytes = vec![0x00, 0x01, 0x0D, 0x00, 0x00, 0x00, 0x01, 0x00, 0x01, 0x41]; + assert!(DataValue::from_bytes(&bytes).is_none()); + } + + #[test] + fn test_truncated_container_rejected() { + let dv = DataValue::Container(vec![ + (DataTypeId(1), DataValue::Str("hello".to_string())), + ]); + let bytes = dv.to_bytes(); + // Truncate to fewer than 2 bytes so neither container nor array can be read + assert!(DataValue::from_bytes(&bytes[..1]).is_none()); + assert!(DataValue::from_bytes(&bytes[..0]).is_none()); + } + + #[test] + fn test_display_basic() { + assert_eq!(format!("{}", DataValue::BoolTrue), "true"); + assert_eq!(format!("{}", DataValue::BoolFalse), "false"); + assert_eq!(format!("{}", DataValue::Null), "null"); + assert_eq!(format!("{}", DataValue::SignedNumber(42)), "42"); + assert_eq!(format!("{}", DataValue::UnsignedNumber(42)), "42"); + assert_eq!(format!("{}", DataValue::Str("hi".to_string())), "\"hi\""); + assert_eq!(format!("{}", DataValue::Bytes(vec![])), "(Binary)"); + } + + #[test] + fn test_hash_consistency() { + use std::collections::HashSet; + let mut set = HashSet::new(); + set.insert(DataValue::BoolTrue); + set.insert(DataValue::BoolFalse); + set.insert(DataValue::Null); + set.insert(DataValue::SignedNumber(1)); + set.insert(DataValue::UnsignedNumber(1)); + assert_eq!(set.len(), 5); + set.insert(DataValue::Bool(true)); + assert_eq!(set.len(), 5); + } + + #[test] + fn test_float_display() { + let s = format!("{}", DataValue::Float(2, 12345)); + assert_eq!(s, "12345e2"); + } + + #[test] + fn test_container_display() { + let dv = DataValue::Container(vec![ + (DataTypeId(3), DataValue::Str("v2.0".to_string())), + (DataTypeId(6), DataValue::UnsignedNumber(42)), + ]); + let s = format!("{}", dv); + assert!(s.contains("3:")); + assert!(s.contains("6:")); + } + + #[test] + fn test_array_display() { + let dv = DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]); + let s = format!("{}", dv); + assert_eq!(s, "[1, 2]"); + } + + /* ===== Crypto container tests ===== */ + + #[cfg(feature = "crypto")] + #[test] + fn test_encrypt_decrypt_container_roundtrip() { + use mtp_crypto::ChaCha20Poly1305; + let key = [0xAB; 32]; + let cipher = ChaCha20Poly1305::new(key); + + let mut dv = DataValue::Container(vec![ + (DataTypeId(1), DataValue::Str("secret".to_string())), + (DataTypeId(2), DataValue::UnsignedNumber(42)), + ]); + + assert!(dv.encrypt_container(&cipher, b"aad").is_some()); + assert!(matches!(dv, DataValue::EncryptedContainer(_))); + + assert!(dv.decrypt_into_container(&cipher, b"aad").is_some()); + assert!(matches!(dv, DataValue::Container(_))); + + let entries = dv.as_container().unwrap(); + assert_eq!(entries.len(), 2); + } + + #[cfg(feature = "crypto")] + #[test] + fn test_encrypt_container_wrong_key_fails() { + use mtp_crypto::ChaCha20Poly1305; + let cipher_a = ChaCha20Poly1305::new([0xAB; 32]); + let cipher_b = ChaCha20Poly1305::new([0xCD; 32]); + + let mut dv = DataValue::Container(vec![ + (DataTypeId(1), DataValue::Str("secret".to_string())), + ]); + + assert!(dv.encrypt_container(&cipher_a, b"aad").is_some()); + assert!(dv.decrypt_into_container(&cipher_b, b"aad").is_none()); + } + + #[cfg(feature = "crypto")] + #[test] + fn test_encrypt_container_wrong_aad_fails() { + use mtp_crypto::ChaCha20Poly1305; + let cipher = ChaCha20Poly1305::new([0xAB; 32]); + + let mut dv = DataValue::Container(vec![ + (DataTypeId(1), DataValue::Str("secret".to_string())), + ]); + + assert!(dv.encrypt_container(&cipher, b"correct-aad").is_some()); + assert!(dv.decrypt_into_container(&cipher, b"wrong-aad").is_none()); + } + + #[cfg(feature = "crypto")] + #[test] + fn test_encrypt_non_container_fails() { + let cipher = mtp_crypto::ChaCha20Poly1305::new([0xAB; 32]); + + let mut dv = DataValue::Str("not a container".to_string()); + assert!(dv.encrypt_container(&cipher, b"aad").is_none()); + } + + #[cfg(feature = "crypto")] + #[test] + fn test_sign_verify_container_roundtrip() { + use mtp_crypto::{ChaCha20Poly1305, Ed25519Signer, SigAlgorithm}; + + let (signer, sk, _pk) = Ed25519Signer::generate(); + let cipher = ChaCha20Poly1305::new([0xAB; 32]); + + let mut dv = DataValue::Container(vec![ + (DataTypeId(1), DataValue::Str("signed data".to_string())), + ]); + + assert!(dv + .sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, &cipher, b"aad") + .is_some()); + assert!(matches!(dv, DataValue::SignedEncryptedContainer(_))); + + assert!(dv.decrypt_signed_encrypted_container(&cipher, b"aad").is_some()); + assert!(matches!(dv, DataValue::SignedContainer(_))); + + let verifier = Ed25519Signer::new(&sk).unwrap(); + assert!(dv.verify_into_container(&verifier).is_some()); + assert!(matches!(dv, DataValue::Container(_))); + + let entries = dv.as_container().unwrap(); + assert_eq!(entries.len(), 1); + } + + #[cfg(feature = "crypto")] + #[test] + fn test_sign_container_wrong_key_fails() { + use mtp_crypto::{Ed25519Signer, SigAlgorithm}; + + let (signer, _, _) = Ed25519Signer::generate(); + let (_, sk2, _) = Ed25519Signer::generate(); + let wrong_verifier = Ed25519Signer::new(&sk2).unwrap(); + + let mut dv = DataValue::Container(vec![ + (DataTypeId(1), DataValue::Str("signed data".to_string())), + ]); + + assert!(dv.sign_container(SigAlgorithm::ED25519, &signer).is_some()); + assert!(dv.verify_into_container(&wrong_verifier).is_none()); + } +} diff --git a/common/src/lib.rs b/common/src/lib.rs index 339d426..5cc7acc 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -18,8 +18,65 @@ pub enum CodecError { CryptoFailed(String), } +/* ================================ TESTS ================================ */ +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_codec_error_display() { + let e = CodecError::InvalidEncoding; + assert_eq!(format!("{}", e), "Invalid encoding"); + } + + #[test] + fn test_codec_error_unknown_version() { + assert_eq!(format!("{}", CodecError::UnknownVersion), "Unknown version"); + } + + #[test] + fn test_codec_error_clone_eq() { + let a = CodecError::InvalidEncoding; + let b = a.clone(); + assert_eq!(a, b); + } + + #[test] + fn test_communication_error_display() { + assert_eq!(format!("{}", CommunicationError::UseAfterClosed), "Use after Closed"); + assert_eq!(format!("{}", CommunicationError::StreamClosed), "Stream Closed"); + assert_eq!(format!("{}", CommunicationError::StreamError), "Stream Error"); + } + + #[test] + fn test_communication_error_clone() { + let a = CommunicationError::UseAfterClosed; + let b = a.clone(); + assert_eq!(format!("{:?}", a), format!("{:?}", b)); + } + + #[test] + fn test_communication_error_authentication_failed() { + let e = CommunicationError::AuthenticationFailed("bad key".into()); + assert!(format!("{}", e).contains("bad key")); + } + + #[test] + fn test_communication_error_other() { + let e = CommunicationError::Other("custom error".into()); + assert!(format!("{}", e).contains("custom error")); + } +} + #[derive(Debug, Error, Clone)] pub enum CommunicationError { + /* ====== NOTE on PartialEq ====== + * Manual PartialEq implementation below compares some inner + * values (String, Vec) structurally and others only by + * variant discriminant (types from quinn / wtransport that + * don't themselves implement PartialEq). Keep this impl in + * sync if variants are added or changed. + */ #[error("Use after Closed")] UseAfterClosed, @@ -89,3 +146,42 @@ pub enum CommunicationError { #[error("Other: {0}")] Other(String), } + +/* ===== Manual PartialEq ===== + * + * Compare by variant, and structurally for String-holding variants. + * Third-party error types (quinn, wtransport) are compared only by + * discriminant since they don't implement PartialEq themselves. + */ +impl PartialEq for CommunicationError { + fn eq(&self, other: &Self) -> bool { + match (self, other) { + (Self::UseAfterClosed, Self::UseAfterClosed) => true, + (Self::ClosedLocally, Self::ClosedLocally) => true, + (Self::ClosedByPeer, Self::ClosedByPeer) => true, + (Self::ConnectionLost, Self::ConnectionLost) => true, + (Self::Quinn(_), Self::Quinn(_)) => true, + (Self::ParseCommunicationValue, Self::ParseCommunicationValue) => true, + (Self::CertificateParseFailed, Self::CertificateParseFailed) => true, + (Self::CertificateLoadFailed, Self::CertificateLoadFailed) => true, + (Self::ParseError(a), Self::ParseError(b)) => a == b, + (Self::ConnectionError(_), Self::ConnectionError(_)) => true, + (Self::ConnectingError(a), Self::ConnectingError(b)) => a == b, + (Self::ReadToEndError(_), Self::ReadToEndError(_)) => true, + (Self::WriteError(_), Self::WriteError(_)) => true, + (Self::ClosedError(_), Self::ClosedError(_)) => true, + (Self::MessageTooLarge, Self::MessageTooLarge) => true, + (Self::ReadExactError(_), Self::ReadExactError(_)) => true, + (Self::StreamClosed, Self::StreamClosed) => true, + (Self::StreamError, Self::StreamError) => true, + (Self::StreamWriteError(_), Self::StreamWriteError(_)) => true, + (Self::StreamReadExactError(_), Self::StreamReadExactError(_)) => true, + (Self::CryptoProviderInstallFailed, Self::CryptoProviderInstallFailed) => true, + (Self::AuthenticationFailed(a), Self::AuthenticationFailed(b)) => a == b, + (Self::Other(a), Self::Other(b)) => a == b, + _ => false, + } + } +} + +impl Eq for CommunicationError {} diff --git a/crypto/Cargo.toml b/crypto/Cargo.toml index b00a2fa..6a6cbd4 100644 --- a/crypto/Cargo.toml +++ b/crypto/Cargo.toml @@ -20,7 +20,9 @@ ml-dsa = { version = "0.1.1", optional = true } serde = { version = "1", optional = true, features = ["derive"] } [features] -default = ["chacha20poly1305", "ed25519-dalek", "hkdf", "sha2"] +default = ["chacha20poly1305", "ed25519-dalek", "hkdf", "sha2", "ml-dsa"] +# Enabling ml-dsa by default ensures dual-signature support in the handshake +# without requiring a separate PQC feature flag in host/client crates. full = ["default", "aes-gcm"] pqc = ["mlkem-tls", "ml-dsa"] serde = ["dep:serde"] diff --git a/crypto/src/keypair.rs b/crypto/src/keypair.rs index e7d87a5..bf41106 100644 --- a/crypto/src/keypair.rs +++ b/crypto/src/keypair.rs @@ -84,32 +84,6 @@ impl From> for SignaturePublicKey { } } -#[derive(ZeroizeOnDrop)] -pub struct KeyGroup { - #[zeroize(skip)] - pub encryption_public_key: EncryptionPublicKey, - pub encryption_private_key: EncryptionPrivateKey, - #[zeroize(skip)] - pub signature_public_key: SignaturePublicKey, - pub signature_private_key: SignaturePrivateKey, -} - -impl KeyGroup { - pub fn new( - encryption_public_key: EncryptionPublicKey, - encryption_private_key: EncryptionPrivateKey, - signature_public_key: SignaturePublicKey, - signature_private_key: SignaturePrivateKey, - ) -> Self { - Self { - encryption_public_key, - encryption_private_key, - signature_public_key, - signature_private_key, - } - } -} - #[cfg_attr(feature = "serde", derive(serde::Serialize, serde::Deserialize))] #[cfg_attr(feature = "serde", serde(transparent))] #[derive(Zeroize, ZeroizeOnDrop)] diff --git a/crypto/src/lib.rs b/crypto/src/lib.rs index 073d6e2..d132b63 100644 --- a/crypto/src/lib.rs +++ b/crypto/src/lib.rs @@ -11,6 +11,9 @@ pub mod kdf; #[cfg(any(feature = "ed25519-dalek", feature = "ml-dsa"))] pub mod sign; +#[cfg(any(feature = "ed25519-dalek", feature = "ml-dsa"))] +pub use sign::SigAlgorithm; + #[cfg(feature = "mlkem-tls")] pub mod kem; @@ -19,7 +22,7 @@ pub mod helper; pub use aead::{AeadCipher, AeadDecrypt, AeadEncrypt}; pub use error::CryptoError; pub use keypair::{ - EncryptionPrivateKey, EncryptionPublicKey, KemPrivateKey, KemPublicKey, KeyGroup, Keyring, + EncryptionPrivateKey, EncryptionPublicKey, KemPrivateKey, KemPublicKey, Keyring, PublicKeyBundle, SignaturePqPrivateKey, SignaturePqPublicKey, SignaturePrivateKey, SignaturePublicKey, }; diff --git a/crypto/src/sign.rs b/crypto/src/sign.rs index aeb593f..c39bb55 100644 --- a/crypto/src/sign.rs +++ b/crypto/src/sign.rs @@ -3,6 +3,24 @@ use crate::error::CryptoError; #[cfg(feature = "ed25519-dalek")] use crate::keypair::{SignaturePrivateKey, SignaturePublicKey}; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct SigAlgorithm; + +impl SigAlgorithm { + pub const ED25519: u8 = 0x01; + pub const ML_DSA_65: u8 = 0x02; + pub const DUAL: u8 = 0x03; + + pub const fn length(alg: u8) -> Option { + match alg { + Self::ED25519 => Some(64), + Self::ML_DSA_65 => Some(3309), + Self::DUAL => Some(3373), + _ => None, + } + } +} + #[cfg(feature = "ed25519-dalek")] use rand_core::RngCore; diff --git a/example-usage/.gitignore b/example-usage/.gitignore index cf5e6e3..11b7871 100644 --- a/example-usage/.gitignore +++ b/example-usage/.gitignore @@ -1,4 +1,6 @@ client_keys.json host_keys.json host_sig_pk.bin +host_sig_pq_pk.bin +host_enc_kem_pk.bin clients.json diff --git a/example-usage/client/Cargo.toml b/example-usage/client/Cargo.toml index 1f2a829..b2c8193 100644 --- a/example-usage/client/Cargo.toml +++ b/example-usage/client/Cargo.toml @@ -1,16 +1,14 @@ [package] -name = "mtp-example-client" +name = "client" version = "0.1.0" edition = "2024" [[bin]] -name = "mtp-example-client" +name = "client" path = "src/main.rs" [dependencies] -mtp-client = { path = "../../client", features = ["crypto"] } -mtp-codec = { path = "../../codec" } -mtp-crypto = { path = "../../crypto" } +mtp = { path = "../../", features = ["client", "crypto"] } tokio = { version = "1", features = ["full"] } serde_json = "1" hex = "0.4" diff --git a/example-usage/client/src/main.rs b/example-usage/client/src/main.rs index 88ed12e..fff8006 100644 --- a/example-usage/client/src/main.rs +++ b/example-usage/client/src/main.rs @@ -1,11 +1,10 @@ use std::fs; -use mtp_client::{ClientConfig, MTPClient, MTPConnection}; -use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue}; -use mtp_crypto::{ - Ed25519Signer, EncryptionPrivateKey, EncryptionPublicKey, KeyGroup, Keyring, - KemPrivateKey, KemPublicKey, SignaturePqPrivateKey, SignaturePqPublicKey, - SignaturePrivateKey, SignaturePublicKey, +use mtp::client::{ClientConfig, MTPClient, MTPConnection}; +use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataValue}; +use mtp::crypto::{ + Ed25519Signer, KemPublicKey, Keyring, MlDsaSigner, PublicKeyBundle, SignaturePqPublicKey, + SignaturePublicKey, }; /* @@ -14,45 +13,39 @@ use mtp_crypto::{ * connection and the local keyring. */ async fn connect_or_register( - config: ClientConfig, - host_public_key: SignaturePublicKey, + mut config: ClientConfig, + host_public_key: PublicKeyBundle, client_key_path: &str, ) -> Result<(MTPConnection, Keyring), Box> { if let Ok(data) = fs::read_to_string(client_key_path) { let json: serde_json::Value = serde_json::from_str(&data)?; let client_id = json["client_id"].as_u64().expect("Invalid client_id"); - let keyring = Keyring::from_bytes( - &hex::decode(json["keyring"].as_str().expect("Missing keyring"))?, - )?; + let keyring = Keyring::from_bytes(&hex::decode( + json["keyring"].as_str().expect("Missing keyring"), + )?)?; println!("Loaded client keys (ID: {})", client_id); - let keygroup = KeyGroup::new( - EncryptionPublicKey::new(vec![]), - EncryptionPrivateKey::new(vec![]), - keyring.sig_cl_public_key.clone(), - SignaturePrivateKey::new(keyring.sig_cl_secret_key.as_bytes().to_vec()), - ); - - let conn = MTPClient::auth_connect(config, client_id, keygroup, host_public_key).await?; + config.client_id = client_id; + let conn = MTPClient::auth_connect(config, &keyring, &host_public_key).await?; println!("Authenticated (version {})", conn.version); Ok((conn, keyring)) } else { - println!("No existing keys found - registering new client"); + println!("No existing keys found: registering new client"); - let (_signer, sig_sk, sig_pk) = Ed25519Signer::generate(); + let (_ed_signer, sig_sk, sig_pk) = Ed25519Signer::generate(); + let (_pq_signer, sig_pq_sk, sig_pq_pk) = MlDsaSigner::generate(); let keyring = Keyring::new( - KemPublicKey::new(vec![]), - KemPrivateKey::new(vec![]), - SignaturePqPublicKey::new(vec![]), - SignaturePqPrivateKey::new(vec![]), + mtp::crypto::KemPublicKey::new(vec![]), + mtp::crypto::KemPrivateKey::new(vec![]), + sig_pq_pk, + sig_pq_sk, sig_pk, sig_sk, ); - // Serialise before the move so we can persist and return the keyring let keyring_bytes = keyring.to_bytes(); - 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 json = serde_json::json!({ @@ -79,7 +72,10 @@ async fn send_ping_and_receive(conn: &MTPConnection) -> Result<(), Box Result<(), Box Result<(), Box> { - let cert_pem = fs::read("server.pem").expect("Missing server.pem - run server first"); - let host_public_key = SignaturePublicKey::new( - fs::read("host_sig_pk.bin").expect("Missing host_sig_pk.bin - run server first"), + let cert_pem = fs::read("server.pem").expect("Missing server.pem: run server first"); + let host_public_key = PublicKeyBundle::new( + KemPublicKey::new( + fs::read("host_enc_kem_pk.bin").expect("Missing host_enc_kem_pk.bin: run server first"), + ), + SignaturePqPublicKey::new( + fs::read("host_sig_pq_pk.bin").expect("Missing host_sig_pq_pk.bin: run server first"), + ), + SignaturePublicKey::new( + fs::read("host_sig_pk.bin").expect("Missing host_sig_pk.bin: run server first"), + ), ); println!("Connecting to 127.0.0.1:8080 ..."); @@ -107,6 +111,7 @@ async fn main() -> Result<(), Box> { let config = ClientConfig { url: "https://127.0.0.1:8080".into(), server_cert: Some(cert_pem), + client_id: 0, }; let (conn, _keyring) = connect_or_register(config, host_public_key, "client_keys.json").await?; diff --git a/example-usage/server/Cargo.toml b/example-usage/server/Cargo.toml index 19f3031..79d33be 100644 --- a/example-usage/server/Cargo.toml +++ b/example-usage/server/Cargo.toml @@ -1,17 +1,16 @@ [package] -name = "mtp-example-server" +name = "server" version = "0.1.0" edition = "2024" [[bin]] -name = "mtp-example-server" +name = "server" path = "src/main.rs" [dependencies] -mtp-host = { path = "../../host", features = ["crypto"] } -mtp-codec = { path = "../../codec" } -mtp-crypto = { path = "../../crypto", features = ["serde"] } +mtp = { path = "../../", features = ["crypto", "host"] } rcgen = "0.14" tokio = { version = "1", features = ["full"] } -serde_json = "1" +serde_json = { version = "1" } hex = "0.4" +serde_core = "1.0.228" diff --git a/example-usage/server/src/main.rs b/example-usage/server/src/main.rs index b8aea9c..2d3b7a7 100644 --- a/example-usage/server/src/main.rs +++ b/example-usage/server/src/main.rs @@ -2,19 +2,19 @@ use std::collections::HashMap; use std::fs; use std::sync::{Arc, Mutex}; -use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue, TypeMap}; -use mtp_host::{HostConfig, MTPHost}; -use mtp_crypto::{ - Ed25519Signer, KemPrivateKey, KemPublicKey, Keyring, PublicKeyBundle, - SignaturePqPrivateKey, SignaturePqPublicKey, -}; +use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue, TypeMap}; +use mtp::crypto::{Ed25519Signer, Keyring, MlDsaSigner, PublicKeyBundle, kem::HybridKem}; +use mtp::host::{HostConfig, MTPHost}; /* * Load an existing TLS certificate and key pair from disk. If neither * file exists, generate a self-signed certificate so the server is * immediately usable without external tooling. */ -fn load_or_generate_tls(cert_path: &str, key_path: &str) -> Result<(Vec, Vec), Box> { +fn load_or_generate_tls( + cert_path: &str, + key_path: &str, +) -> Result<(Vec, Vec), Box> { if let (Ok(c), Ok(k)) = (fs::read(cert_path), fs::read(key_path)) { println!("Using existing TLS cert from {cert_path}"); return Ok((c, k)); @@ -50,15 +50,10 @@ fn load_or_generate_host_keys(path: &str) -> Result<(u64, Keyring), Box Result<(u64, Keyring), Box Result< - ( - Arc>>, - Arc>, - ), - Box, -> { - let clients: Arc>> = Arc::new(Mutex::new( - if let Ok(data) = fs::read_to_string(path) { +) -> Result<(Arc>>, Arc>), Box> +{ + let clients: Arc>> = + Arc::new(Mutex::new(if let Ok(data) = fs::read_to_string(path) { serde_json::from_str(&data).unwrap_or_default() } else { HashMap::new() - }, - )); + })); let next_id = Arc::new(Mutex::new( clients.lock().unwrap().keys().max().unwrap_or(&999) + 1, )); @@ -126,10 +115,15 @@ async fn main() -> Result<(), Box> { let (cert_pem, key_pem) = load_or_generate_tls("server.pem", "server.key")?; let (host_id, host_keyring) = load_or_generate_host_keys("host_keys.json")?; - // Export the host's Ed25519 public key so clients can verify it + // Export the host's public keys so clients can verify it fs::write( - "host_sig_pk.bin", - host_keyring.sig_cl_public_key.as_bytes(), + "host_enc_kem_pk.bin", + host_keyring.kem_public_key.as_bytes(), + )?; + fs::write("host_sig_pk.bin", host_keyring.sig_cl_public_key.as_bytes())?; + fs::write( + "host_sig_pq_pk.bin", + host_keyring.sig_pq_public_key.as_bytes(), )?; let (clients, next_id) = load_client_db("clients.json")?; @@ -157,6 +151,7 @@ async fn main() -> Result<(), Box> { println!("Starting MTP server on port 8080 ..."); let config = HostConfig { + ip: std::net::IpAddr::V4(std::net::Ipv4Addr::UNSPECIFIED), port: 8080, tls_fullchain: cert_pem, tls_key: key_pem, diff --git a/host/src/lib.rs b/host/src/lib.rs index 6d6b24e..cbf0348 100644 --- a/host/src/lib.rs +++ b/host/src/lib.rs @@ -6,9 +6,11 @@ use mtp_codec::{ }; use mtp_common::CommunicationError; use mtp_transport::{Policy, Receiver, Sender}; +use std::net::IpAddr; // Host configuration. pub struct HostConfig { + pub ip: IpAddr, pub port: u16, pub tls_fullchain: Vec, pub tls_key: Vec, @@ -60,6 +62,7 @@ impl MTPHost { let registry = Registry::builtin(); let transport = mtp_transport::host( + config.ip, config.port, config.tls_fullchain.clone(), config.tls_key.clone(), @@ -90,7 +93,7 @@ impl MTPHost { return self.accept_authenticated(sender, receiver).await; } - /* Read the first message (always encoded with reserved types). */ + // Read the first message (always encoded with reserved types). let first_msg = receiver.receive().await.ok()?; /* @@ -136,30 +139,12 @@ impl MTPHost { sender: Sender, receiver: Receiver, ) -> Option { - use mtp_crypto::{Ed25519Signer, PublicKeyBundle, SignatureScheme, verify_ed25519}; + use mtp_crypto::{ + Ed25519Signer, PublicKeyBundle, SignatureScheme, verify_ed25519, verify_ml_dsa, + }; - // 1. Send host greeting - let host_nonce: u128 = rand::random(); - let host_bundle = self.config.host_keyring.public_key_bundle(); - - let greeting = - CommunicationValue::new(mtp_codec::CommunicationType::IdentificationResponse) - .add_typed_default( - DataType::Id, - DataValue::UnsignedNumber(self.config.host_id as u128), - ) - .add_typed_default( - DataType::PublicKeys, - DataValue::Bytes(host_bundle.as_bytes()), - ) - .add_typed_default(DataType::ServerNonce, DataValue::UnsignedNumber(host_nonce)); - - sender.send(&greeting).await.ok()?; - - // 2. Receive client message + // 1. Receive client message first (no host greeting) let msg = receiver.receive().await.ok()?; - - // Extract common fields let version_str = match msg.get_data(DataTypeId(3)) { DataValue::Str(s) => s.clone(), _ => { @@ -186,6 +171,11 @@ impl MTPHost { } }; + let pq_sig_bytes: Vec = match msg.get_data(DataTypeId(12)) { + DataValue::Bytes(b) => b.clone(), + _ => vec![], + }; + let (assigned_id, client_bundle) = if msg.get_type() == mtp_codec::CommunicationTypeId(15) { // LOGIN let cid = match msg.get_data(DataTypeId(6)) { @@ -231,12 +221,26 @@ impl MTPHost { sender.close(); return None; } + if !pq_sig_bytes.is_empty() + && verify_ml_dsa(&bundle.sig_pq_public_key, &sig_payload, &pq_sig_bytes).is_err() + { + let rejection = + CommunicationValue::new(mtp_codec::CommunicationType::IdentificationResponse) + .add_typed_default(DataType::Connected, DataValue::BoolFalse) + .add_typed_default( + DataType::ClientNonce, + DataValue::UnsignedNumber(client_nonce), + ); + let _ = sender.send(&rejection).await; + sender.close(); + return None; + } /* ===== End Signature ===== */ (cid, bundle) - } else if msg.get_type() == mtp_codec::CommunicationTypeId(14) { + } else if msg.get_type() == mtp_codec::CommunicationTypeId(17) { // REGISTER - let bundle = match msg.get_data(DataTypeId(8)) { + let bundle = match msg.get_data(DataTypeId(9)) { DataValue::Bytes(b) => PublicKeyBundle::from_bytes(b).ok()?, _ => { sender.close(); @@ -263,6 +267,20 @@ impl MTPHost { sender.close(); return None; } + if !pq_sig_bytes.is_empty() + && verify_ml_dsa(&bundle.sig_pq_public_key, &sig_payload, &pq_sig_bytes).is_err() + { + let rejection = + CommunicationValue::new(mtp_codec::CommunicationType::RegisterResponse) + .add_typed_default(DataType::Connected, DataValue::BoolFalse) + .add_typed_default( + DataType::ClientNonce, + DataValue::UnsignedNumber(client_nonce), + ); + let _ = sender.send(&rejection).await; + sender.close(); + return None; + } /* ===== End Signature ===== */ let new_id = (self.config.complete_register)(bundle.clone()); @@ -272,11 +290,11 @@ impl MTPHost { return None; }; - // 4. Send success response + // 2. Send success response (single host message) let new_nonce: u128 = rand::random(); let mut host_sig_payload = Vec::new(); - host_sig_payload.extend_from_slice(b"true"); + host_sig_payload.push(0x01); host_sig_payload.extend_from_slice(&assigned_id.to_be_bytes()); host_sig_payload.extend_from_slice(&client_nonce.to_be_bytes()); host_sig_payload.extend_from_slice(&new_nonce.to_be_bytes()); @@ -285,9 +303,8 @@ impl MTPHost { /* ===== Signature ===== */ let host_sig = host_signer.sign(&host_sig_payload).ok()?; - /* ===== End Signature ===== */ - let response = + let mut response = CommunicationValue::new(mtp_codec::CommunicationType::IdentificationResponse) .add_typed_default(DataType::Connected, DataValue::BoolTrue) .add_typed_default( @@ -298,9 +315,28 @@ impl MTPHost { .add_typed_default(DataType::Signature, DataValue::Bytes(host_sig)) .add_typed_default(DataType::Id, DataValue::UnsignedNumber(assigned_id as u128)); + if !self + .config + .host_keyring + .sig_pq_secret_key + .as_bytes() + .is_empty() + { + use mtp_crypto::MlDsaSigner; + let host_pq_signer = MlDsaSigner::new( + &self.config.host_keyring.sig_pq_secret_key, + &self.config.host_keyring.sig_pq_public_key, + ) + .ok()?; + let host_pq_sig = host_pq_signer.sign(&host_sig_payload).ok()?; + response = + response.add_typed_default(DataType::PqSignature, DataValue::Bytes(host_pq_sig)); + } + /* ===== End Signature ===== */ + sender.send(&response).await.ok()?; - // 5. Version negotiation + // 3. Version negotiation let negotiated = self.registry.negotiate(&[client_version])?; let codec = VersionedCodec::new(self.registry.clone()); diff --git a/transport/Cargo.toml b/transport/Cargo.toml index d7ccda3..7e4d972 100644 --- a/transport/Cargo.toml +++ b/transport/Cargo.toml @@ -21,6 +21,9 @@ thiserror = "2.0.18" rustls-native-certs = "0.8.4" log = "0.4" +[dev-dependencies] +rcgen = "0.14" + [features] default = [] # Enables hosting a MTP server diff --git a/transport/src/connection.rs b/transport/src/connection.rs index 1738999..ed6bbd8 100644 --- a/transport/src/connection.rs +++ b/transport/src/connection.rs @@ -56,7 +56,6 @@ impl Default for Policy { enum ReceivedFrame { Message(CommunicationValue), ClosedByPeer, - #[allow(dead_code)] Idle, } @@ -517,3 +516,48 @@ impl Receiver { self.handle.close_reason() } } + +/* ================================ TESTS ================================ */ +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_send_mode_derive() { + assert_eq!(SendMode::PersistentStream, SendMode::PersistentStream); + 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, 1_000_000_000); + 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.max_transient_recv_errors, 20); + assert_eq!(p.transient_recv_backoff, Duration::from_millis(100)); + assert_eq!(p.receiver_queue_capacity, 1000); + } + + #[test] + fn test_policy_clone() { + let p = Policy::default(); + let cloned = p.clone(); + assert_eq!(p.send_mode, cloned.send_mode); + } + + #[test] + fn test_policy_debug() { + let p = Policy::default(); + let debug_str = format!("{:?}", p); + assert!(debug_str.contains("Policy")); + } +} diff --git a/transport/src/connection_handle.rs b/transport/src/connection_handle.rs index 4ebe200..cbbbf2e 100644 --- a/transport/src/connection_handle.rs +++ b/transport/src/connection_handle.rs @@ -63,3 +63,70 @@ impl Default for ConnectionHandle { Self::new() } } + +/* ================================ TESTS ================================ */ +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_new_is_open() { + let h = ConnectionHandle::new(); + assert!(h.is_open()); + assert!(!h.is_closed()); + } + + #[test] + fn test_close_transitions_state() { + let h = ConnectionHandle::new(); + h.close(Some(CommunicationError::StreamClosed)); + assert!(!h.is_open()); + assert!(h.is_closed()); + } + + #[test] + fn test_close_reason_some() { + let h = ConnectionHandle::new(); + h.close(Some(CommunicationError::UseAfterClosed)); + assert!(h.close_reason().is_some()); + } + + #[test] + fn test_close_reason_none() { + let h = ConnectionHandle::new(); + h.close(None); + assert!(h.close_reason().is_none()); + } + + #[test] + fn test_default_is_new() { + let h = ConnectionHandle::default(); + assert!(h.is_open()); + } + + #[test] + fn test_close_with_error() { + let h = ConnectionHandle::new(); + h.close_with_error(CommunicationError::MessageTooLarge); + assert!(h.is_closed()); + assert!(h.close_reason().is_some()); + } + + #[test] + fn test_multiple_close_first_wins() { + let h = ConnectionHandle::new(); + h.close(Some(CommunicationError::StreamClosed)); + h.close(Some(CommunicationError::UseAfterClosed)); + // First close reason is preserved + assert!(h.close_reason().is_some()); + } + + #[test] + fn test_close_sends_reason() { + let h = ConnectionHandle::new(); + let mut rx = h.subscribe_close(); + h.close(Some(CommunicationError::ClosedLocally)); + // After close, the watch channel is updated + assert!(rx.borrow_and_update().is_some()); + } +} diff --git a/transport/src/host.rs b/transport/src/host.rs index d963fab..2274344 100644 --- a/transport/src/host.rs +++ b/transport/src/host.rs @@ -1,4 +1,5 @@ use crate::{ConnectionHandle, Policy, Receiver, Sender}; +use log; use mtp_common::CommunicationError; use rustls::pki_types::{PrivateKeyDer, pem::PemObject}; use std::net::{IpAddr, SocketAddr}; @@ -13,7 +14,13 @@ pub struct Host { impl Host { pub async fn next(&mut self) -> Option<(Sender, Receiver)> { - self.incoming.recv().await + log::warn!("[transport Host::next] waiting on recv..."); + let result = self.incoming.recv().await; + match &result { + Some(_) => log::warn!("[transport Host::next] received connection"), + None => log::warn!("[transport Host::next] incoming channel closed - sender dropped"), + } + result } pub fn local_addr(&self) -> std::net::SocketAddr { @@ -22,6 +29,7 @@ impl Host { } pub async fn host( + ip: IpAddr, port: u16, cert_pem: Vec, key_pem: Vec, @@ -29,7 +37,7 @@ pub async fn host( ) -> Result { let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); - let server_config = configure_server(port, cert_pem, key_pem, &policy).await?; + let server_config = configure_server(ip, port, cert_pem, key_pem, &policy).await?; let endpoint = Endpoint::server(server_config) .map_err(|e| CommunicationError::Other(format!("Endpoint creation failed: {}", e)))?; @@ -42,25 +50,37 @@ pub async fn host( let policy = Arc::new(policy); let task = tokio::spawn(async move { + log::warn!("[transport bg task] started"); loop { + log::warn!("[transport bg task] waiting for connection..."); let incoming_session = endpoint.accept().await; + log::warn!("[transport bg task] got incoming session"); let request = match incoming_session.await { - Ok(req) => req, - Err(_) => { + Ok(req) => { + log::warn!("[transport bg task] got request"); + req + } + Err(e) => { + log::warn!("[transport bg task] incoming session error: {e}"); continue; } }; let connection = match request.accept().await { - Ok(conn) => conn, - Err(_) => { + Ok(conn) => { + log::warn!("[transport bg task] connection accepted"); + conn + } + Err(e) => { + log::warn!("[transport bg task] accept error: {e}"); continue; } }; let incoming_tx = incoming_tx.clone(); let policy = policy.clone(); + eprintln!("[transport bg task] spawning handle_connection"); tokio::spawn(handle_connection(connection, incoming_tx, policy)); } }); @@ -85,6 +105,7 @@ async fn handle_connection( } async fn configure_server( + bind_ip: IpAddr, port: u16, cert_pem: Vec, key_pem: Vec, @@ -104,12 +125,6 @@ async fn configure_server( tls_config.alpn_protocols = vec![b"h3".to_vec()]; - let bind_ip = std::env::var("mtp_BIND") - .ok() - .and_then(|s| if s.is_empty() { None } else { Some(s) }) - .unwrap_or_else(|| "::".to_string()) - .parse::() - .map_err(|e| CommunicationError::ParseError(e.to_string()))?; let bind_addr = SocketAddr::new(bind_ip, port); let server_config = ServerConfig::builder() diff --git a/transport/tests/integration.rs b/transport/tests/integration.rs new file mode 100644 index 0000000..93377b3 --- /dev/null +++ b/transport/tests/integration.rs @@ -0,0 +1,174 @@ +use std::net::{IpAddr, Ipv4Addr}; + +use mtp_transport::{Policy, host, connect}; + +fn generate_self_signed_cert() -> (Vec, Vec) { + let key_pair = rcgen::KeyPair::generate().unwrap(); + let params = + rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]).unwrap(); + let cert = params.self_signed(&key_pair).unwrap(); + let cert_pem = cert.pem(); + let key_pem = key_pair.serialize_pem(); + (cert_pem.into_bytes(), key_pem.into_bytes()) +} + +#[tokio::test] +async fn test_host_start_and_stop() { + let (cert_pem, key_pem) = generate_self_signed_cert(); + let h = host( + IpAddr::V4(Ipv4Addr::LOCALHOST), + 0, + cert_pem, + key_pem, + Policy::default(), + ) + .await + .unwrap(); + let addr = h.local_addr(); + // Port should be non-zero (OS-assigned) + assert!(addr.port() > 0); +} + +#[tokio::test] +async fn test_send_receive_roundtrip() { + let (cert_pem, key_pem) = generate_self_signed_cert(); + let mut h = host( + IpAddr::V4(Ipv4Addr::LOCALHOST), + 0, + cert_pem.clone(), + key_pem, + Policy::default(), + ) + .await + .unwrap(); + let addr = h.local_addr(); + + let url = format!("https://127.0.0.1:{}", addr.port()); + let (client_tx, client_rx) = + connect(&url, Some(cert_pem), Policy::default()).await.unwrap(); + + // Accept on host side + let (host_tx, host_rx) = h.next().await.unwrap(); + + // Client sends a simple message + let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping) + .add_data( + mtp_codec::DataTypeId(6), + mtp_codec::DataValue::UnsignedNumber(42), + ); + client_tx.send(&msg).await.unwrap(); + + // Host receives it + let received = host_rx.receive().await.unwrap(); + assert_eq!(received.get_type(), mtp_codec::CommunicationTypeId(19)); // Ping + let val = received.get_data(mtp_codec::DataTypeId(6)).clone(); + assert_eq!(val, mtp_codec::DataValue::UnsignedNumber(42)); + + // Host sends a response + let resp = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Pong) + .add_data( + mtp_codec::DataTypeId(6), + mtp_codec::DataValue::UnsignedNumber(99), + ); + host_tx.send(&resp).await.unwrap(); + + // Client receives it + let client_received = client_rx.receive().await.unwrap(); + assert_eq!(client_received.get_type(), mtp_codec::CommunicationTypeId(20)); // Pong + let client_val = client_received.get_data(mtp_codec::DataTypeId(6)).clone(); + assert_eq!(client_val, mtp_codec::DataValue::UnsignedNumber(99)); + + // Close both sides + client_tx.close(); + host_tx.close(); +} + +#[tokio::test] +async fn test_concurrent_messages() { + let (cert_pem, key_pem) = generate_self_signed_cert(); + let mut h = host( + IpAddr::V4(Ipv4Addr::LOCALHOST), + 0, + cert_pem.clone(), + key_pem, + Policy::default(), + ) + .await + .unwrap(); + let addr = h.local_addr(); + + let url = format!("https://127.0.0.1:{}", addr.port()); + let (client_tx, _client_rx) = + connect(&url, Some(cert_pem), Policy::default()).await.unwrap(); + + let (_host_tx, host_rx) = h.next().await.unwrap(); + + // Send 5 messages in sequence + for i in 0..5u128 { + let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping) + .add_data( + mtp_codec::DataTypeId(6), + mtp_codec::DataValue::UnsignedNumber(i), + ); + client_tx.send(&msg).await.unwrap(); + } + + // Receive all 5 in order + for i in 0..5u128 { + let received = host_rx.receive().await.unwrap(); + let val = received.get_data(mtp_codec::DataTypeId(6)).clone(); + assert_eq!(val, mtp_codec::DataValue::UnsignedNumber(i)); + } + + // Send 3 responses back + for i in 0..3u128 { + let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Pong) + .add_data( + mtp_codec::DataTypeId(6), + mtp_codec::DataValue::UnsignedNumber(i * 10), + ); + client_tx.send(&msg).await.unwrap(); + } + + for i in 0..3u128 { + let received = host_rx.receive().await.unwrap(); + let val = received.get_data(mtp_codec::DataTypeId(6)).clone(); + assert_eq!(val, mtp_codec::DataValue::UnsignedNumber(i * 10)); + } + + client_tx.close(); +} + +#[tokio::test] +async fn test_close_detection() { + let (cert_pem, key_pem) = generate_self_signed_cert(); + let mut h = host( + IpAddr::V4(Ipv4Addr::LOCALHOST), + 0, + cert_pem.clone(), + key_pem, + Policy::default(), + ) + .await + .unwrap(); + let addr = h.local_addr(); + + let url = format!("https://127.0.0.1:{}", addr.port()); + let (client_tx, _client_rx) = + connect(&url, Some(cert_pem), Policy::default()).await.unwrap(); + + let (_host_tx, host_rx) = h.next().await.unwrap(); + + // Send a message then close + let msg = mtp_codec::CommunicationValue::new(mtp_codec::CommunicationType::Ping); + client_tx.send(&msg).await.unwrap(); + client_tx.close(); + + // Host should still receive the message + let received = host_rx.receive().await.unwrap(); + assert_eq!(received.get_type(), mtp_codec::CommunicationTypeId(19)); // Ping + + // Host should get an error or closed signal on next receive + let result = host_rx.receive().await; + assert!(result.is_err()); +} diff --git a/type-map/build.rs b/type-map/build.rs index 8930db0..97b1a22 100644 --- a/type-map/build.rs +++ b/type-map/build.rs @@ -156,22 +156,28 @@ const RESERVED_DATA_TYPES: &[ReservedEntry] = &[ name: "Connected", id: 11, }, + ReservedEntry { + name: "PqSignature", + id: 12, + }, ]; fn main() { let out = std::path::PathBuf::from(std::env::var("OUT_DIR").unwrap()); let multi_version = std::env::var("CARGO_FEATURE_REGISTRY").is_ok(); + println!("cargo:rerun-if-env-changed=MTP_TYPE_MAPS"); + let config = match std::env::var("MTP_TYPE_MAPS") { Ok(config_path) => { + println!("cargo:rerun-if-changed={}", config_path); + let content = std::fs::read_to_string(&config_path).expect("Failed to read type-maps.yaml"); serde_yaml::from_str(&content).expect("Failed to parse type-maps.yaml") } Err(_) => { - eprintln!( - "warning: MTP_TYPE_MAPS not set; generating types with reserved entries only" - ); + eprint!("warning: MTP_TYPE_MAPS not set; generating types with reserved entries only"); Config { protocol_version: String::new(), type_maps: BTreeMap::new(), @@ -441,12 +447,7 @@ fn generate_data_type_enum(out: &mut String, user_names: &BTreeSet<&str>) { .unwrap(); } for name in user_names { - writeln!( - out, - " \"{}\" => Some(DataType::{}),", - name, name - ) - .unwrap(); + writeln!(out, " \"{}\" => Some(DataType::{}),", name, name).unwrap(); } writeln!(out, " _ => None,").unwrap(); diff --git a/type-map/src/lib.rs b/type-map/src/lib.rs index 29f5429..6bc5563 100644 --- a/type-map/src/lib.rs +++ b/type-map/src/lib.rs @@ -81,8 +81,8 @@ pub use registry::*; #[cfg(feature = "registry")] mod registry { + use crate::{TypeMap, Version, builtin_type_maps}; use std::collections::BTreeMap; - use crate::{Version, TypeMap, builtin_type_maps}; /* * Multi-version type-map registry. @@ -142,6 +142,7 @@ mod registry { } } + /* ================================ TESTS ================================ */ #[cfg(test)] mod tests { use super::*;