diff --git a/client/src/lib.rs b/client/src/lib.rs index 0df6e7e..1edc985 100644 --- a/client/src/lib.rs +++ b/client/src/lib.rs @@ -1,6 +1,4 @@ -#[cfg(feature = "crypto")] -use mtp_codec::DataType; -use mtp_codec::{CommunicationValue, DataTypeId, DataValue, PROTOCOL_VERSION, Version}; +use mtp_codec::{CommunicationValue, DataType, DataValue, PROTOCOL_VERSION, Version}; use mtp_common::CommunicationError; use mtp_transport::{Policy, Receiver, Sender}; @@ -45,9 +43,9 @@ impl MTPClient { // Build the initial identification message with the protocol version. 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), + .add_typed_default(DataType::Version, DataValue::Str(version_str)) + .add_typed_default( + DataType::Id, DataValue::UnsignedNumber(config.client_id.into()), ); @@ -124,7 +122,9 @@ impl MTPClient { // 2. Receive host response (single message) let response = receiver.receive().await?; - let connected = response.get_data(DataTypeId(11)); + let tm = mtp_codec::TypeMap::latest(); + + let connected = response.get_data(DataType::Connected.to_id(&tm)); match connected { DataValue::BoolTrue => {} DataValue::BoolFalse => { @@ -139,7 +139,7 @@ impl MTPClient { } } - let echo_nonce = response.get_data(DataTypeId(7)); + let echo_nonce = response.get_data(DataType::ClientNonce.to_id(&tm)); match echo_nonce { DataValue::UnsignedNumber(n) if *n == client_nonce as u128 => {} _ => { @@ -149,7 +149,7 @@ impl MTPClient { } } - let host_new_nonce = match response.get_data(DataTypeId(5)) { + let host_new_nonce = match response.get_data(DataType::Timestamp.to_id(&tm)) { DataValue::UnsignedNumber(n) => *n as u128, _ => { return Err(CommunicationError::AuthenticationFailed( @@ -158,7 +158,7 @@ impl MTPClient { } }; - let host_sig = match response.get_data(DataTypeId(10)) { + let host_sig = match response.get_data(DataType::Signature.to_id(&tm)) { DataValue::Bytes(b) => b.clone(), _ => { return Err(CommunicationError::AuthenticationFailed( @@ -167,7 +167,7 @@ impl MTPClient { } }; - let host_pq_sig = match response.get_data(DataTypeId(12)) { + let host_pq_sig = match response.get_data(DataType::PqSignature.to_id(&tm)) { DataValue::Bytes(b) => b.clone(), _ => vec![], }; @@ -264,7 +264,9 @@ impl MTPClient { // 2. Receive host response (single message) let response = receiver.receive().await?; - let connected = response.get_data(DataTypeId(11)); + let tm = mtp_codec::TypeMap::latest(); + + let connected = response.get_data(DataType::Connected.to_id(&tm)); match connected { DataValue::BoolTrue => {} DataValue::BoolFalse => { @@ -279,7 +281,7 @@ impl MTPClient { } } - let assigned_id = match response.get_data(DataTypeId(6)) { + let assigned_id = match response.get_data(DataType::Id.to_id(&tm)) { DataValue::UnsignedNumber(n) => *n as u128, _ => { return Err(CommunicationError::AuthenticationFailed( @@ -288,7 +290,7 @@ impl MTPClient { } }; - let echo_nonce = response.get_data(DataTypeId(7)); + let echo_nonce = response.get_data(DataType::ClientNonce.to_id(&tm)); match echo_nonce { DataValue::UnsignedNumber(n) if *n == client_nonce as u128 => {} _ => { @@ -298,7 +300,7 @@ impl MTPClient { } } - let host_new_nonce = match response.get_data(DataTypeId(5)) { + let host_new_nonce = match response.get_data(DataType::Timestamp.to_id(&tm)) { DataValue::UnsignedNumber(n) => *n as u128, _ => { return Err(CommunicationError::AuthenticationFailed( @@ -307,7 +309,7 @@ impl MTPClient { } }; - let host_sig = match response.get_data(DataTypeId(10)) { + let host_sig = match response.get_data(DataType::Signature.to_id(&tm)) { DataValue::Bytes(b) => b.clone(), _ => { return Err(CommunicationError::AuthenticationFailed( @@ -316,7 +318,7 @@ impl MTPClient { } }; - let host_pq_sig = match response.get_data(DataTypeId(12)) { + let host_pq_sig = match response.get_data(DataType::PqSignature.to_id(&tm)) { DataValue::Bytes(b) => b.clone(), _ => vec![], }; diff --git a/codec/Cargo.toml b/codec/Cargo.toml index 6e91169..6dc27bd 100644 --- a/codec/Cargo.toml +++ b/codec/Cargo.toml @@ -14,4 +14,4 @@ rand = { version = "0.8", features = ["std", "std_rng"] } [features] default = [] registry = ["mtp-type-map/registry"] -crypto = ["dep:mtp-crypto"] +crypto = ["dep:mtp-crypto", "mtp-crypto/mlkem-tls"] diff --git a/codec/src/communication_value.rs b/codec/src/communication_value.rs index 0721398..f0776d4 100644 --- a/codec/src/communication_value.rs +++ b/codec/src/communication_value.rs @@ -146,18 +146,35 @@ impl CommunicationValue { * bit3 => is data encrypted If so data bytes will be an encrypted container * bit4 => is communication value signed */ - pub fn to_bytes(&self) -> Vec { + /* + * Build the canonical metadata header and data payload shared by both + * `to_bytes` and `build_signed_payload`. Keeping a single source here + * guarantees the serialized frame and the signed-over bytes stay in sync. + * + * Returns `(metadata, data_bytes)` where + * metadata = comm_type || flags || id? || sender? || receiver? + * + * `force_signed` forces the `FLAG_SIGNED` bit on regardless of whether a + * signature is currently attached. The signed-payload path passes `true` so + * that the bytes signed by `sign_frame` (before the signature is stored) and + * the bytes verified by `verify_frame` (after it is stored) are identical. + */ + fn build_metadata_and_data( + &self, + force_signed: bool, + ) -> Result<(Vec, Vec), CodecError> { 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(_) - ) - }); + 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; @@ -179,7 +196,7 @@ impl CommunicationValue { if is_encrypted { flags |= FLAG_ENCRYPTED; } - if has_frame_sig { + if has_frame_sig || force_signed { flags |= FLAG_SIGNED; } @@ -212,34 +229,39 @@ impl CommunicationValue { }) .unwrap_or_default() } else { - let container_value = DataValue::container_from_map(&self.data); - container_value.to_bytes() + DataValue::container_from_map(&self.data).to_bytes()? }; #[cfg(not(feature = "crypto"))] - let data_bytes = { - let container_value = DataValue::container_from_map(&self.data); - container_value.to_bytes() - }; + let data_bytes = DataValue::container_from_map(&self.data).to_bytes()?; + + Ok((metadata, data_bytes)) + } + + pub fn to_bytes(&self) -> Result, CodecError> { + let (metadata, data_bytes) = self.build_metadata_and_data(false)?; let mut payload = Vec::new(); payload.extend_from_slice(&metadata); #[cfg(feature = "crypto")] - if let Some((_alg, _sig)) = &self.frame_signature { + 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.push(*alg); + payload.extend_from_slice(sig); } payload.extend_from_slice(&data_bytes); + let len = u32::try_from(payload.len()).map_err(|_| CodecError::TooManyEntries)?; let mut frame = Vec::with_capacity(4 + payload.len()); - let _ = frame.write_u32::(payload.len() as u32); + frame + .write_u32::(len) + .map_err(|_| CodecError::InvalidEncoding)?; frame.extend_from_slice(&payload); - frame + Ok(frame) } pub fn from_bytes(bytes: &[u8]) -> Result { @@ -372,7 +394,7 @@ impl CommunicationValue { algorithm: u8, signer: &impl SignatureScheme, ) -> Option<()> { - let signed_payload = self.build_signed_payload(); + let signed_payload = self.build_signed_payload().ok()?; let sig = signer.sign(&signed_payload).ok()?; self.frame_signature = Some((algorithm, sig)); Some(()) @@ -389,7 +411,7 @@ impl CommunicationValue { .as_ref() .ok_or(CodecError::InvalidEncoding)?; - let signed_payload = self.build_signed_payload(); + let signed_payload = self.build_signed_payload()?; verifier .verify(&signed_payload, sig) .map_err(|_| CodecError::InvalidEncoding) @@ -400,75 +422,11 @@ impl CommunicationValue { * 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() + fn build_signed_payload(&self) -> Result, CodecError> { + // Force FLAG_SIGNED on so the signed bytes match whether or not the + // signature has been attached yet (sign_frame runs before storing it). + let (metadata, data_bytes) = self.build_metadata_and_data(true)?; + Ok([metadata, data_bytes].concat()) } #[cfg(feature = "crypto")] @@ -610,9 +568,9 @@ mod tests { use crate::data_value::DataValue; fn roundtrip(cv: CommunicationValue) -> CommunicationValue { - let bytes = cv.to_bytes(); + let bytes = cv.to_bytes().expect("encode failed"); let decoded = CommunicationValue::from_bytes(&bytes).expect("failed to deserialize"); - let bytes2 = decoded.to_bytes(); + let bytes2 = decoded.to_bytes().expect("encode failed"); assert_eq!(bytes, bytes2); decoded } @@ -620,7 +578,7 @@ mod tests { #[test] fn test_flags_and_order_without_optional() { let cv = CommunicationValue::new(CommunicationType::ErrorParsing).with_id(0); - let bytes = cv.to_bytes(); + let bytes = cv.to_bytes().expect("encode failed"); // [u32 len][u16 type][flags]... assert!(bytes.len() >= 7); @@ -642,7 +600,7 @@ mod tests { .with_sender(0x0000_1122_3344_5566) .with_receiver(0x0000_6677_8899_AABB); - let bytes = cv.to_bytes(); + let bytes = cv.to_bytes().expect("encode failed"); let mut c = Cursor::new(bytes.as_slice()); let total_len = c.read_u32::().expect("len"); @@ -703,4 +661,45 @@ mod tests { bad[0..4].copy_from_slice(&(1000u32.to_be_bytes())); assert!(CommunicationValue::from_bytes(&bad).is_err()); } + + #[cfg(feature = "crypto")] + #[test] + fn test_sign_verify_frame_roundtrip() { + use mtp_crypto::{Ed25519Signer, SigAlgorithm}; + + let (signer, sk, _pk) = Ed25519Signer::generate(); + + let mut cv = CommunicationValue::new(CommunicationType::Ping) + .with_id(7) + .with_sender(1) + .with_receiver(2) + .add_data(DataTypeId(6), DataValue::UnsignedNumber(42)); + + assert!(cv.sign_frame(SigAlgorithm::ED25519, &signer).is_some()); + + // Same in-memory value verifies (FLAG_SIGNED forced on both sides). + let verifier = Ed25519Signer::new(&sk).unwrap(); + assert!(cv.verify_frame(&verifier).is_ok()); + + // Survives a wire round-trip. + let bytes = cv.to_bytes().expect("encode failed"); + let decoded = CommunicationValue::from_bytes(&bytes).expect("decode failed"); + assert!(decoded.verify_frame(&verifier).is_ok()); + } + + #[cfg(feature = "crypto")] + #[test] + fn test_verify_frame_wrong_key_fails() { + use mtp_crypto::{Ed25519Signer, SigAlgorithm}; + + let (signer, _, _) = Ed25519Signer::generate(); + let (_, other_sk, _) = Ed25519Signer::generate(); + + let mut cv = CommunicationValue::new(CommunicationType::Ping) + .add_data(DataTypeId(6), DataValue::UnsignedNumber(42)); + assert!(cv.sign_frame(SigAlgorithm::ED25519, &signer).is_some()); + + let wrong = Ed25519Signer::new(&other_sk).unwrap(); + assert!(cv.verify_frame(&wrong).is_err()); + } } diff --git a/codec/src/data_value.rs b/codec/src/data_value.rs index 7ccff8d..f8de1fb 100644 --- a/codec/src/data_value.rs +++ b/codec/src/data_value.rs @@ -6,10 +6,11 @@ use std::fmt; use std::hash::{Hash, Hasher}; use std::io::Cursor; +use mtp_common::CodecError; use mtp_type_map::DataTypeId; #[cfg(feature = "crypto")] -use mtp_crypto::{AeadDecrypt, AeadEncrypt, SigAlgorithm, SignatureScheme}; +use mtp_crypto::{EncryptionType, Keyring, PublicKeyBundle, SigAlgorithm, SignatureScheme}; #[derive(Debug, Clone, PartialEq, Eq)] pub enum DataKind { @@ -124,9 +125,9 @@ impl DataValue { * 0x07 => Bytes * 0x08 => Array * 0x09 => Container - * 0x0A => EncryptedContainer (4 bytes u32 len + encrypted bytes) - * 0x0B => SignedContainer (4 bytes u32 len + 3373 bytes signature) - * 0x0C => SignedEncryptedContainer (4 bytes u32 len + 3373 bytes signature + encrypted bytes) + * 0x0A => EncryptedContainer (1 byte EncryptionType + KEM ciphertext + AEAD payload) + * 0x0B => SignedContainer (1 byte SigAlgorithm + signature + serialized container) + * 0x0C => SignedEncryptedContainer (encrypted blob that decrypts to a SignedContainer) * 0xFF => Null */ const KIND_BOOL_TRUE: u8 = 0x01; @@ -272,13 +273,14 @@ impl DataValue { /* * Decrypt an `EncryptedContainer` in-place, replacing it with the - * deserialized `Container`. Returns `None` if decryption or - * deserialization fails. + * deserialized `Container`. The algorithm (and which keypair to use) is read + * from the blob's leading `EncryptionType` byte; the matching key is taken + * from `keyring`. 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, keyring: &Keyring, aad: &[u8]) -> Option<()> { let data = self.as_encrypted_container()?; - let plaintext = cipher.decrypt(&data, aad).ok()?; + let plaintext = mtp_crypto::decrypt_with(&data, keyring, aad).ok()?; let dv = DataValue::from_bytes(&plaintext)?; match dv { DataValue::Container(entries) => { @@ -291,13 +293,21 @@ impl DataValue { /* * Encrypt a `Container` into an `EncryptedContainer` in-place. + * `enc_type` selects the algorithm and `recipient` provides the public key + * encapsulated to. The resulting blob is self-describing: its leading byte + * is `enc_type`, so `decrypt_into_container` needs only a `Keyring`. * 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, + enc_type: EncryptionType, + recipient: &PublicKeyBundle, + aad: &[u8], + ) -> Option<()> { let entries = self.as_container()?; - let plaintext = DataValue::Container(entries).to_bytes(); - let ct = cipher.encrypt(&plaintext, aad).ok()?; + let plaintext = DataValue::Container(entries).to_bytes().ok()?; + let ct = mtp_crypto::encrypt_for(enc_type, recipient, &plaintext, aad).ok()?; *self = DataValue::EncryptedContainer(ct); Some(()) } @@ -311,7 +321,7 @@ impl DataValue { #[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 container_bytes = Self::encode_container(&entries).ok()?; let sig = signer.sign(&container_bytes).ok()?; @@ -354,35 +364,40 @@ impl DataValue { /* * Encrypt a `Container` into a `SignedEncryptedContainer` in-place. - * The result is an opaque ciphertext that decrypts to a `SignedContainer`. + * The container is first signed (with `algorithm`/`signer`), then the signed + * blob is encrypted with `enc_type` to `recipient`. 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, + enc_type: EncryptionType, + recipient: &PublicKeyBundle, aad: &[u8], ) -> Option<()> { self.sign_container(algorithm, signer)?; let blob = self.as_signed_container()?; - let ct = cipher.encrypt(&blob, aad).ok()?; + let ct = mtp_crypto::encrypt_for(enc_type, recipient, &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. + * `SignedContainer`. The algorithm and keypair are resolved from the blob's + * leading byte and `keyring`. Does NOT verify; call `verify_into_container` + * next. */ #[cfg(feature = "crypto")] pub fn decrypt_signed_encrypted_container( &mut self, - cipher: &impl AeadDecrypt, + keyring: &Keyring, aad: &[u8], ) -> Option<()> { let data = self.as_signed_encrypted_container()?; - let plaintext = cipher.decrypt(&data, aad).ok()?; + let plaintext = mtp_crypto::decrypt_with(&data, keyring, aad).ok()?; *self = DataValue::SignedContainer(plaintext); Some(()) } @@ -400,16 +415,14 @@ impl DataValue { } } - pub fn to_bytes(&self) -> Vec { + pub fn to_bytes(&self) -> Result, CodecError> { match self { DataValue::Container(entries) => Self::encode_container(entries), DataValue::Array(arr) => Self::encode_array(arr), _ => { let mut out = Vec::new(); - if Self::write_value_payload(&mut out, self).is_none() { - return Vec::new(); - } - out + Self::write_value_payload(&mut out, self)?; + Ok(out) } } } @@ -423,8 +436,8 @@ impl DataValue { Some(value) } - pub fn to_base64(&self) -> String { - general_purpose::STANDARD.encode(self.to_bytes()) + pub fn to_base64(&self) -> Result { + Ok(general_purpose::STANDARD.encode(self.to_bytes()?)) } pub fn from_base64(base64_str: &str) -> Option { @@ -432,144 +445,141 @@ impl DataValue { Self::from_bytes(&bytes) } - fn encode_container(entries: &[(DataTypeId, DataValue)]) -> Vec { + fn encode_container(entries: &[(DataTypeId, DataValue)]) -> Result, CodecError> { let mut out = Vec::new(); - if out - .write_u16::(u16::try_from(entries.len()).ok().unwrap_or(0)) - .is_err() - { - return Vec::new(); - } + let count = u16::try_from(entries.len()).map_err(|_| CodecError::TooManyEntries)?; + out.write_u16::(count) + .map_err(|_| CodecError::InvalidEncoding)?; for (key, value) in entries { - if !Self::write_container_entry(&mut out, key.clone(), value) { - return Vec::new(); - } + Self::write_container_entry(&mut out, key.clone(), value)?; } - out + Ok(out) } - fn write_container_entry(buf: &mut Vec, key: DataTypeId, value: &DataValue) -> bool { + fn write_container_entry( + buf: &mut Vec, + key: DataTypeId, + value: &DataValue, + ) -> Result<(), CodecError> { let kind = Self::kind_marker(value); buf.push(kind); if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL { - let _ = buf.write_u16::(key.0); - return true; + buf.write_u16::(key.0) + .map_err(|_| CodecError::InvalidEncoding)?; + return Ok(()); } let mut payload = Vec::new(); - if Self::write_value_payload(&mut payload, value).is_none() { - return false; - } + Self::write_value_payload(&mut payload, value)?; - if buf.write_u32::(payload.len() as u32).is_err() { - return false; - } - let _ = buf.write_u16::(key.0); + let len = u32::try_from(payload.len()).map_err(|_| CodecError::TooManyEntries)?; + buf.write_u32::(len) + .map_err(|_| CodecError::InvalidEncoding)?; + buf.write_u16::(key.0) + .map_err(|_| CodecError::InvalidEncoding)?; buf.extend_from_slice(&payload); - true + Ok(()) } - fn encode_array(arr: &[DataValue]) -> Vec { + fn encode_array(arr: &[DataValue]) -> Result, CodecError> { let mut out = Vec::new(); - if out - .write_u16::(u16::try_from(arr.len()).ok().unwrap_or(0)) - .is_err() - { - return Vec::new(); - } + let count = u16::try_from(arr.len()).map_err(|_| CodecError::TooManyEntries)?; + out.write_u16::(count) + .map_err(|_| CodecError::InvalidEncoding)?; for value in arr { - if !Self::write_array_entry(&mut out, value) { - return Vec::new(); - } + Self::write_array_entry(&mut out, value)?; } - out + Ok(out) } - fn write_array_entry(buf: &mut Vec, value: &DataValue) -> bool { + fn write_array_entry(buf: &mut Vec, value: &DataValue) -> Result<(), CodecError> { let kind = Self::kind_marker(value); buf.push(kind); if kind == Self::KIND_BOOL_TRUE || kind == Self::KIND_BOOL_FALSE || kind == Self::KIND_NULL { - return true; + return Ok(()); } let mut payload = Vec::new(); - if Self::write_value_payload(&mut payload, value).is_none() { - return false; - } + Self::write_value_payload(&mut payload, value)?; - if buf.write_u32::(payload.len() as u32).is_err() { - return false; - } + let len = u32::try_from(payload.len()).map_err(|_| CodecError::TooManyEntries)?; + buf.write_u32::(len) + .map_err(|_| CodecError::InvalidEncoding)?; buf.extend_from_slice(&payload); - true + Ok(()) } - fn write_value_payload(buf: &mut Vec, value: &DataValue) -> Option<()> { + fn write_value_payload(buf: &mut Vec, value: &DataValue) -> Result<(), CodecError> { match value { - DataValue::BoolTrue => Some(()), - DataValue::BoolFalse => Some(()), + DataValue::BoolTrue => Ok(()), + DataValue::BoolFalse => Ok(()), DataValue::Bool(v) => { + // Kept intentionally: the kind marker already encodes the boolean, + // so both arms carry no payload. Retained for clear compatibility. if *v { - Some(()) + Ok(()) } else { - Some(()) + Ok(()) } } DataValue::SignedNumber(n) => { - buf.write_i128::(*n).ok()?; - Some(()) + buf.write_i128::(*n) + .map_err(|_| CodecError::InvalidEncoding)?; + Ok(()) } DataValue::UnsignedNumber(n) => { - buf.write_u128::(*n).ok()?; - Some(()) + buf.write_u128::(*n) + .map_err(|_| CodecError::InvalidEncoding)?; + Ok(()) } DataValue::Float(a, b) => { - buf.write_u8(*a).ok()?; - buf.write_u32::(*b).ok()?; - Some(()) + buf.write_u8(*a).map_err(|_| CodecError::InvalidEncoding)?; + buf.write_u32::(*b) + .map_err(|_| CodecError::InvalidEncoding)?; + Ok(()) } DataValue::Str(s) => { buf.extend_from_slice(s.as_bytes()); - Some(()) + Ok(()) } DataValue::Array(arr) => { - let bytes = Self::encode_array(arr); + let bytes = Self::encode_array(arr)?; buf.extend_from_slice(&bytes); - Some(()) + Ok(()) } DataValue::Bytes(b) => { buf.extend_from_slice(b); - Some(()) + Ok(()) } DataValue::Container(entries) => { - let bytes = Self::encode_container(entries); + let bytes = Self::encode_container(entries)?; buf.extend_from_slice(&bytes); - Some(()) + Ok(()) } #[cfg(feature = "crypto")] DataValue::EncryptedContainer(data) => { buf.extend_from_slice(data); - Some(()) + Ok(()) } #[cfg(feature = "crypto")] DataValue::SignedContainer(data) => { buf.extend_from_slice(data); - Some(()) + Ok(()) } #[cfg(feature = "crypto")] DataValue::SignedEncryptedContainer(data) => { buf.extend_from_slice(data); - Some(()) + Ok(()) } - DataValue::Null => Some(()), + DataValue::Null => Ok(()), } } @@ -901,62 +911,28 @@ impl PartialEq for DataValue { impl Hash for DataValue { fn hash(&self, state: &mut H) { use DataValue::*; + // Use the wire kind marker as the per-variant discriminant. It is unique + // per kind and maps BoolTrue/Bool(true) (and BoolFalse/Bool(false)) to the + // same marker, keeping the hash consistent with the Eq bool equivalence. + Self::kind_marker(self).hash(state); match self { - BoolTrue | Bool(true) => { - 0u8.hash(state); - true.hash(state); - } - BoolFalse | Bool(false) => { - 0u8.hash(state); - false.hash(state); - } - SignedNumber(n) => { - 1u8.hash(state); - n.hash(state); - } - UnsignedNumber(n) => { - 2u8.hash(state); - n.hash(state); - } + BoolTrue | BoolFalse | Bool(_) | Null => {} + SignedNumber(n) => n.hash(state), + UnsignedNumber(n) => n.hash(state), Float(n, m) => { - 3u8.hash(state); n.hash(state); m.hash(state); } - Str(s) => { - 2u8.hash(state); - s.hash(state); - } - Array(a) => { - 3u8.hash(state); - a.hash(state); - } - Bytes(a) => { - 4u8.hash(state); - a.hash(state); - } - Container(c) => { - 5u8.hash(state); - c.hash(state); - } + Str(s) => s.hash(state), + Array(a) => a.hash(state), + Bytes(a) => a.hash(state), + Container(c) => c.hash(state), #[cfg(feature = "crypto")] - EncryptedContainer(c) => { - 6u8.hash(state); - c.hash(state); - } + EncryptedContainer(c) => c.hash(state), #[cfg(feature = "crypto")] - SignedContainer(c) => { - 7u8.hash(state); - c.hash(state); - } + SignedContainer(c) => c.hash(state), #[cfg(feature = "crypto")] - SignedEncryptedContainer(c) => { - 8u8.hash(state); - c.hash(state); - } - Null => { - 9u8.hash(state); - } + SignedEncryptedContainer(c) => c.hash(state), } } } @@ -970,14 +946,14 @@ mod tests { /// 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 bytes = dv.to_bytes().expect("encode failed"); 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 bytes = dv.to_bytes().expect("encode failed"); let decoded = DataValue::from_bytes(&bytes).expect("roundtrip failed"); assert_eq!(dv, decoded, "array roundtrip mismatch"); } @@ -1111,7 +1087,7 @@ mod tests { DataTypeId(7), DataValue::Bytes(vec![0xDE, 0xAD, 0xBE, 0xEF]), )]); - let b64 = dv.to_base64(); + let b64 = dv.to_base64().expect("encode failed"); let decoded = DataValue::from_base64(&b64).expect("base64 roundtrip failed"); assert_eq!(dv, decoded); } @@ -1203,7 +1179,7 @@ mod tests { #[test] fn test_truncated_container_rejected() { let dv = DataValue::Container(vec![(DataTypeId(1), DataValue::Str("hello".to_string()))]); - let bytes = dv.to_bytes(); + let bytes = dv.to_bytes().expect("encode failed"); // 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()); @@ -1263,19 +1239,22 @@ mod tests { #[cfg(feature = "crypto")] #[test] fn test_encrypt_decrypt_container_roundtrip() { - use mtp_crypto::ChaCha20Poly1305; - let key = [0xAB; 32]; - let cipher = ChaCha20Poly1305::new(key); + use mtp_crypto::{EncryptionType, Keyring}; + let keyring = Keyring::generate(); + let bundle = keyring.public_key_bundle(); 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!( + dv.encrypt_container(EncryptionType::MlKemChaCha20Poly1305, &bundle, b"aad") + .is_some() + ); assert!(matches!(dv, DataValue::EncryptedContainer(_))); - assert!(dv.decrypt_into_container(&cipher, b"aad").is_some()); + assert!(dv.decrypt_into_container(&keyring, b"aad").is_some()); assert!(matches!(dv, DataValue::Container(_))); let entries = dv.as_container().unwrap(); @@ -1285,46 +1264,68 @@ mod tests { #[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]); + use mtp_crypto::{EncryptionType, Keyring}; + let keyring_a = Keyring::generate(); + let keyring_b = Keyring::generate(); 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()); + assert!( + dv.encrypt_container( + EncryptionType::MlKemChaCha20Poly1305, + &keyring_a.public_key_bundle(), + b"aad" + ) + .is_some() + ); + assert!(dv.decrypt_into_container(&keyring_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]); + use mtp_crypto::{EncryptionType, Keyring}; + let keyring = Keyring::generate(); 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()); + assert!( + dv.encrypt_container( + EncryptionType::MlKemChaCha20Poly1305, + &keyring.public_key_bundle(), + b"correct-aad" + ) + .is_some() + ); + assert!(dv.decrypt_into_container(&keyring, b"wrong-aad").is_none()); } #[cfg(feature = "crypto")] #[test] fn test_encrypt_non_container_fails() { - let cipher = mtp_crypto::ChaCha20Poly1305::new([0xAB; 32]); + use mtp_crypto::{EncryptionType, Keyring}; + let keyring = Keyring::generate(); let mut dv = DataValue::Str("not a container".to_string()); - assert!(dv.encrypt_container(&cipher, b"aad").is_none()); + assert!( + dv.encrypt_container( + EncryptionType::MlKemChaCha20Poly1305, + &keyring.public_key_bundle(), + b"aad" + ) + .is_none() + ); } #[cfg(feature = "crypto")] #[test] fn test_sign_verify_container_roundtrip() { - use mtp_crypto::{ChaCha20Poly1305, Ed25519Signer, SigAlgorithm}; + use mtp_crypto::{Ed25519Signer, EncryptionType, Keyring, SigAlgorithm}; + let keyring = Keyring::generate(); let (signer, sk, _pk) = Ed25519Signer::generate(); - let cipher = ChaCha20Poly1305::new([0xAB; 32]); let mut dv = DataValue::Container(vec![( DataTypeId(1), @@ -1332,13 +1333,19 @@ mod tests { )]); assert!( - dv.sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, &cipher, b"aad") - .is_some() + dv.sign_and_encrypt_container( + SigAlgorithm::ED25519, + &signer, + EncryptionType::MlKemChaCha20Poly1305, + &keyring.public_key_bundle(), + b"aad" + ) + .is_some() ); assert!(matches!(dv, DataValue::SignedEncryptedContainer(_))); assert!( - dv.decrypt_signed_encrypted_container(&cipher, b"aad") + dv.decrypt_signed_encrypted_container(&keyring, b"aad") .is_some() ); assert!(matches!(dv, DataValue::SignedContainer(_))); diff --git a/common/src/lib.rs b/common/src/lib.rs index 7637d37..5be1d93 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -12,6 +12,8 @@ pub enum CodecError { ReservedCommunicationType(u16), #[error("Invalid encoding")] InvalidEncoding, + #[error("Too many entries to encode")] + TooManyEntries, #[error("Crypto failed: {0}")] CryptoFailed(String), } @@ -69,6 +71,9 @@ pub enum CommunicationError { #[error("ParseCommunicationValue error")] ParseCommunicationValue, + #[error("Encode error")] + Encode, + #[error("Parse Certificate error")] CertificateParseFailed, @@ -139,6 +144,9 @@ pub enum CommunicationError { #[error("ParseCommunicationValue error")] ParseCommunicationValue, + #[error("Encode error")] + Encode, + #[error("Parse Certificate error")] CertificateParseFailed, @@ -182,6 +190,7 @@ impl PartialEq for CommunicationError { (Self::ConnectionLost, Self::ConnectionLost) => true, (Self::Quinn(_), Self::Quinn(_)) => true, (Self::ParseCommunicationValue, Self::ParseCommunicationValue) => true, + (Self::Encode, Self::Encode) => true, (Self::CertificateParseFailed, Self::CertificateParseFailed) => true, (Self::CertificateLoadFailed, Self::CertificateLoadFailed) => true, (Self::ParseError(a), Self::ParseError(b)) => a == b, @@ -213,6 +222,7 @@ impl PartialEq for CommunicationError { (Self::ClosedByPeer, Self::ClosedByPeer) => true, (Self::ConnectionLost, Self::ConnectionLost) => true, (Self::ParseCommunicationValue, Self::ParseCommunicationValue) => true, + (Self::Encode, Self::Encode) => true, (Self::CertificateParseFailed, Self::CertificateParseFailed) => true, (Self::CertificateLoadFailed, Self::CertificateLoadFailed) => true, (Self::ParseError(a), Self::ParseError(b)) => a == b, diff --git a/crypto/src/enc.rs b/crypto/src/enc.rs new file mode 100644 index 0000000..3b7a44b --- /dev/null +++ b/crypto/src/enc.rs @@ -0,0 +1,246 @@ +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +use crate::error::CryptoError; + +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +use crate::kdf::derive_encryption_key; +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +use crate::kem::HybridKem; +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +use crate::keypair::{Keyring, PublicKeyBundle}; + +/* + * Algorithm selector for encrypted containers. + * + * Mirrors `SigAlgorithm` for signatures: a single marking byte identifies the + * key-encapsulation mechanism and the AEAD used to seal a container. The byte + * is stored as the first byte of every encrypted blob so the decryptor can pick + * the matching algorithm (and the matching keypair from a `Keyring`) without + * any out-of-band agreement. + * + * All variants currently use ML-KEM (X25519MlKem768) for key encapsulation and + * differ only in the AEAD. AES-256-GCM variants require the `aes-gcm` feature + * (enabled via the crate's `full` feature); sealing/opening with a variant whose + * AEAD feature is not compiled in returns an error. + */ +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum EncryptionType { + /// ML-KEM (X25519MlKem768) key encapsulation with XChaCha20-Poly1305 AEAD. + MlKemChaCha20Poly1305, + /// ML-KEM (X25519MlKem768) key encapsulation with AES-256-GCM AEAD. + MlKemAes256Gcm, +} + +impl EncryptionType { + pub const ML_KEM_CHACHA20POLY1305: u8 = 0x01; + pub const ML_KEM_AES256_GCM: u8 = 0x02; + + /// The marking byte written at the front of an encrypted blob. + pub const fn to_byte(self) -> u8 { + match self { + Self::MlKemChaCha20Poly1305 => Self::ML_KEM_CHACHA20POLY1305, + Self::MlKemAes256Gcm => Self::ML_KEM_AES256_GCM, + } + } + + /// Recover an `EncryptionType` from its marking byte, or `None` if unknown. + pub const fn from_byte(b: u8) -> Option { + match b { + Self::ML_KEM_CHACHA20POLY1305 => Some(Self::MlKemChaCha20Poly1305), + Self::ML_KEM_AES256_GCM => Some(Self::MlKemAes256Gcm), + _ => None, + } + } +} + +/* + * Seal `plaintext` with a 32-byte AEAD key chosen by `enc_type`. + * + * Returns `EncryptionFailed` when the selected AEAD's feature is not compiled in. + */ +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +#[allow(unused_variables)] +fn aead_seal( + enc_type: EncryptionType, + key: [u8; 32], + plaintext: &[u8], + aad: &[u8], +) -> Result, CryptoError> { + #[allow(unused_imports)] + use crate::aead::AeadEncrypt; + match enc_type { + #[cfg(feature = "chacha20poly1305")] + EncryptionType::MlKemChaCha20Poly1305 => { + crate::aead::ChaCha20Poly1305::new(key).encrypt(plaintext, aad) + } + #[cfg(feature = "aes-gcm")] + EncryptionType::MlKemAes256Gcm => crate::aead::Aes256Gcm::new(key).encrypt(plaintext, aad), + #[allow(unreachable_patterns)] + _ => Err(CryptoError::EncryptionFailed), + } +} + +/* + * Open `ciphertext` with a 32-byte AEAD key chosen by `enc_type`. + * + * Returns `DecryptionFailed` when the selected AEAD's feature is not compiled in. + */ +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +#[allow(unused_variables)] +fn aead_open( + enc_type: EncryptionType, + key: [u8; 32], + ciphertext: &[u8], + aad: &[u8], +) -> Result, CryptoError> { + #[allow(unused_imports)] + use crate::aead::AeadDecrypt; + match enc_type { + #[cfg(feature = "chacha20poly1305")] + EncryptionType::MlKemChaCha20Poly1305 => { + crate::aead::ChaCha20Poly1305::new(key).decrypt(ciphertext, aad) + } + #[cfg(feature = "aes-gcm")] + EncryptionType::MlKemAes256Gcm => crate::aead::Aes256Gcm::new(key).decrypt(ciphertext, aad), + #[allow(unreachable_patterns)] + _ => Err(CryptoError::DecryptionFailed), + } +} + +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +const ENC_KDF_SALT: &[u8] = b"mtp-container-enc"; +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +const ENC_KDF_CONTEXT: &[u8] = b"single-recipient"; + +/* + * Encrypt `plaintext` for a single recipient, selecting the algorithm with + * `enc_type` and the recipient's KEM public key from `recipient`. + * + * The returned, self-describing blob is laid out as: + * [1 byte EncryptionType] [2 bytes u16 kem_ct_len] [kem_ciphertext] [aead_payload] + * where `aead_payload` is the AEAD output (nonce + ciphertext + tag). The AEAD + * key is derived from the KEM shared secret via HKDF, so no separate content key + * is transmitted. + * + * Requires the `mlkem-tls` and `hkdf` features, plus the AEAD feature backing + * `enc_type`. + */ +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +pub fn encrypt_for( + enc_type: EncryptionType, + recipient: &PublicKeyBundle, + plaintext: &[u8], + aad: &[u8], +) -> Result, CryptoError> { + let enc = HybridKem::encapsulate(&recipient.kem_public_key)?; + let key = derive_encryption_key(&enc.shared_secret, ENC_KDF_SALT, ENC_KDF_CONTEXT)?; + let aead_payload = aead_seal(enc_type, key, plaintext, aad)?; + + let kem_ct = enc.ciphertext; + let mut out = Vec::with_capacity(1 + 2 + kem_ct.len() + aead_payload.len()); + out.push(enc_type.to_byte()); + out.extend_from_slice(&(kem_ct.len() as u16).to_be_bytes()); + out.extend_from_slice(&kem_ct); + out.extend_from_slice(&aead_payload); + Ok(out) +} + +/* + * Decrypt a blob produced by [`encrypt_for`] using `keyring`. + * + * The leading byte selects the `EncryptionType` (and thus which keypair to use + * from the keyring); for the current ML-KEM variants that is `kem_secret_key`. + * Returns `DecryptionFailed` on any malformed input or authentication failure. + * + * Requires the `mlkem-tls` and `hkdf` features, plus the AEAD feature backing + * the blob's algorithm. + */ +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +pub fn decrypt_with(blob: &[u8], keyring: &Keyring, aad: &[u8]) -> Result, CryptoError> { + if blob.len() < 3 { + return Err(CryptoError::DecryptionFailed); + } + let enc_type = EncryptionType::from_byte(blob[0]).ok_or(CryptoError::DecryptionFailed)?; + let kem_ct_len = u16::from_be_bytes([blob[1], blob[2]]) as usize; + let kem_end = 3usize + .checked_add(kem_ct_len) + .ok_or(CryptoError::DecryptionFailed)?; + let kem_ct = blob.get(3..kem_end).ok_or(CryptoError::DecryptionFailed)?; + let aead_payload = blob.get(kem_end..).ok_or(CryptoError::DecryptionFailed)?; + + let shared_secret = HybridKem::decapsulate(&keyring.kem_secret_key, kem_ct)?; + let key = derive_encryption_key(&shared_secret, ENC_KDF_SALT, ENC_KDF_CONTEXT)?; + aead_open(enc_type, key, aead_payload, aad) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn encryption_type_byte_roundtrip() { + for t in [ + EncryptionType::MlKemChaCha20Poly1305, + EncryptionType::MlKemAes256Gcm, + ] { + assert_eq!(EncryptionType::from_byte(t.to_byte()), Some(t)); + } + assert_eq!(EncryptionType::from_byte(0x00), None); + assert_eq!(EncryptionType::from_byte(0xFF), None); + } + + #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] + #[test] + fn encrypt_for_roundtrip() { + let kr = Keyring::generate(); + let blob = encrypt_for( + EncryptionType::MlKemChaCha20Poly1305, + &kr.public_key_bundle(), + b"secret payload", + b"aad", + ) + .unwrap(); + assert_eq!(blob[0], EncryptionType::ML_KEM_CHACHA20POLY1305); + + let pt = decrypt_with(&blob, &kr, b"aad").unwrap(); + assert_eq!(pt, b"secret payload"); + } + + #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] + #[test] + fn decrypt_with_wrong_keyring_fails() { + let kr = Keyring::generate(); + let other = Keyring::generate(); + let blob = encrypt_for( + EncryptionType::MlKemChaCha20Poly1305, + &kr.public_key_bundle(), + b"secret", + b"aad", + ) + .unwrap(); + assert!(decrypt_with(&blob, &other, b"aad").is_err()); + } + + #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] + #[test] + fn decrypt_with_wrong_aad_fails() { + let kr = Keyring::generate(); + let blob = encrypt_for( + EncryptionType::MlKemChaCha20Poly1305, + &kr.public_key_bundle(), + b"secret", + b"right", + ) + .unwrap(); + assert!(decrypt_with(&blob, &kr, b"wrong").is_err()); + } + + #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] + #[test] + fn decrypt_with_malformed_fails() { + let kr = Keyring::generate(); + assert!(decrypt_with(b"", &kr, b"").is_err()); + assert!(decrypt_with(&[0x01, 0x00], &kr, b"").is_err()); + // Unknown algorithm byte. + assert!(decrypt_with(&[0x7F, 0x00, 0x00], &kr, b"").is_err()); + } +} diff --git a/crypto/src/kdf.rs b/crypto/src/kdf.rs index 113880c..52baf4d 100644 --- a/crypto/src/kdf.rs +++ b/crypto/src/kdf.rs @@ -16,10 +16,12 @@ pub fn hkdf_expand( } pub fn hkdf_extract(ikm: &[u8], salt: &[u8]) -> [u8; 32] { - let (_, hk) = Hkdf::::extract(Some(salt), ikm); - let mut okm = [0u8; 32]; - hk.expand(&[], &mut okm).expect("hkdf expand failed"); - okm + // Return the pseudo-random key (PRK) produced by HKDF-Extract directly. + // Extract cannot fail, so this avoids the panicking expand step entirely. + let (prk, _) = Hkdf::::extract(Some(salt), ikm); + let mut out = [0u8; 32]; + out.copy_from_slice(&prk); + out } pub fn derive_encryption_key( diff --git a/crypto/src/lib.rs b/crypto/src/lib.rs index d132b63..fc5fc48 100644 --- a/crypto/src/lib.rs +++ b/crypto/src/lib.rs @@ -17,6 +17,8 @@ pub use sign::SigAlgorithm; #[cfg(feature = "mlkem-tls")] pub mod kem; +pub mod enc; + pub mod helper; pub use aead::{AeadCipher, AeadDecrypt, AeadEncrypt}; @@ -51,6 +53,11 @@ pub use kdf::{derive_encryption_key, hkdf_expand, hkdf_extract}; #[cfg(feature = "mlkem-tls")] pub use kem::{Encapsulated, HybridKem}; +pub use enc::EncryptionType; + +#[cfg(all(feature = "mlkem-tls", feature = "hkdf"))] +pub use enc::{decrypt_with, encrypt_for}; + #[cfg(all(feature = "mlkem-tls", feature = "chacha20poly1305", feature = "hkdf"))] pub use helper::{decrypt_multi, encrypt_multi, MultiEncryptedMessage, RecipientEntry}; @@ -140,7 +147,7 @@ mod tests { let (ed_signer, _, _) = Ed25519Signer::generate(); let (ml_signer, _, _) = MlDsaSigner::generate(); - let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg"); + let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg").unwrap(); dual .verify( ed_signer.verifying_key(), @@ -157,7 +164,7 @@ mod tests { let (ed_signer, _, _) = Ed25519Signer::generate(); let (ml_signer, _, _) = MlDsaSigner::generate(); - let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg"); + let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg").unwrap(); assert!(dual .verify(ed_signer.verifying_key(), ml_signer.verifying_key(), b"wrong") .is_err()); diff --git a/crypto/src/sign.rs b/crypto/src/sign.rs index c39bb55..5c57f46 100644 --- a/crypto/src/sign.rs +++ b/crypto/src/sign.rs @@ -218,19 +218,20 @@ pub fn sign_dual( ed25519_sk: &ed25519_dalek::SigningKey, mldsa_sk: &ml_dsa::SigningKey, message: &[u8], -) -> DualSignature { +) -> Result { let ed25519 = { use ed25519_dalek::Signer; ed25519_sk.sign(message).to_bytes().to_vec() }; let mldsa = { use ml_dsa::Signer; - mldsa_sk.try_sign(message) - .expect("ML-DSA signing failed") + mldsa_sk + .try_sign(message) + .map_err(|_| CryptoError::SigningFailed)? .encode() .to_vec() }; - DualSignature { ed25519, mldsa } + Ok(DualSignature { ed25519, mldsa }) } impl DualSignature { diff --git a/example-usage/client/src/main.rs b/example-usage/client/src/main.rs index 5f3c9b4..405f5af 100644 --- a/example-usage/client/src/main.rs +++ b/example-usage/client/src/main.rs @@ -31,9 +31,10 @@ async fn main() -> Result<(), Box> { client_id: 0, }; + let server_bundle = host_public_key.clone(); let (conn, keyring) = auth::connect_or_register(config, host_public_key, "client_keys.json").await?; - messages::send_and_receive(&conn, &keyring).await?; + messages::send_and_receive(&conn, &keyring, &server_bundle).await?; println!("\nDone"); Ok(()) diff --git a/example-usage/client/src/messages.rs b/example-usage/client/src/messages.rs index 1c17e5a..7e11efe 100644 --- a/example-usage/client/src/messages.rs +++ b/example-usage/client/src/messages.rs @@ -1,18 +1,14 @@ use mtp::client::MTPConnection; use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue}; -use mtp::crypto::{ChaCha20Poly1305, Ed25519Signer, Keyring, SigAlgorithm}; +use mtp::crypto::{Ed25519Signer, EncryptionType, Keyring, PublicKeyBundle, SigAlgorithm}; -fn derive_demo_key() -> [u8; 32] { - mtp::crypto::derive_encryption_key( - b"MTP-demo-shared-secret", - b"MTP-demo-salt", - b"encrypted-container-demo", - ) - .expect("key derivation must succeed") -} - -pub fn build_demo_message(client_id: u64, keyring: &Keyring) -> CommunicationValue { - let cipher = ChaCha20Poly1305::new(derive_demo_key()); +pub fn build_demo_message( + client_id: u64, + keyring: &Keyring, + server_bundle: &PublicKeyBundle, +) -> CommunicationValue { + // Encrypt to the server's KEM public key; the server decrypts with its keyring. + let enc_type = EncryptionType::MlKemChaCha20Poly1305; let signer = Ed25519Signer::new(&keyring.sig_cl_secret_key) .expect("Ed25519 signer from keyring"); @@ -21,7 +17,7 @@ pub fn build_demo_message(client_id: u64, keyring: &Keyring) -> CommunicationVal (DataTypeId(2), DataValue::UnsignedNumber(42)), ]); let mut dv_enc = inner_enc; - dv_enc.encrypt_container(&cipher, b"demo-aad"); + dv_enc.encrypt_container(enc_type, server_bundle, b"demo-aad"); let inner_sig = DataValue::Container(vec![ (DataTypeId(1), DataValue::Str("signed by client".into())), @@ -35,7 +31,7 @@ pub fn build_demo_message(client_id: u64, keyring: &Keyring) -> CommunicationVal (DataTypeId(2), DataValue::UnsignedNumber(7)), ]); let mut dv_sec = inner_sec; - dv_sec.sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, &cipher, b"demo-aad"); + dv_sec.sign_and_encrypt_container(SigAlgorithm::ED25519, &signer, enc_type, server_bundle, b"demo-aad"); let timestamp = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) @@ -66,8 +62,9 @@ pub fn build_demo_message(client_id: u64, keyring: &Keyring) -> CommunicationVal pub async fn send_and_receive( conn: &MTPConnection, keyring: &Keyring, + server_bundle: &PublicKeyBundle, ) -> Result<(), Box> { - let msg = build_demo_message(conn.client_id, keyring); + let msg = build_demo_message(conn.client_id, keyring, server_bundle); println!("Sending: {msg}"); conn.sender.send(&msg).await?; diff --git a/example-usage/server/src/handlers.rs b/example-usage/server/src/handlers.rs index 768293b..e706c6a 100644 --- a/example-usage/server/src/handlers.rs +++ b/example-usage/server/src/handlers.rs @@ -1,5 +1,7 @@ use mtp::codec::{CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue, TypeMap}; -use mtp::crypto::{ChaCha20Poly1305, CryptoError, SignatureScheme, SignaturePublicKey, verify_ed25519}; +use mtp::crypto::{ + CryptoError, Keyring, SignaturePublicKey, SignatureScheme, verify_ed25519, +}; struct Ed25519Verifier(SignaturePublicKey); @@ -12,19 +14,11 @@ impl SignatureScheme for Ed25519Verifier { } } -fn derive_demo_key() -> [u8; 32] { - mtp::crypto::derive_encryption_key( - b"MTP-demo-shared-secret", - b"MTP-demo-salt", - b"encrypted-container-demo", - ) - .expect("key derivation must succeed") -} - pub fn process_and_respond( msg: &CommunicationValue, tm: &TypeMap, client_pk: Option<&mtp::crypto::PublicKeyBundle>, + host_keyring: &Keyring, ) -> CommunicationValue { let desc_id = DataTypeId(tm.data_id_enum(DataType::Description).unwrap()); let ts_id = DataTypeId(tm.data_id_enum(DataType::Timestamp).unwrap()); @@ -56,8 +50,6 @@ pub fn process_and_respond( println!(" Binary: {:?}", binary.as_bytes()); println!(" Items: {:?}", items.as_array()); - let cipher = ChaCha20Poly1305::new(derive_demo_key()); - let mut enc_status = String::from("EncryptedPayload: not present"); let mut sig_status = String::from("SignedPayload: not present"); let mut secure_status = String::from("SecurePayload: not present"); @@ -65,7 +57,7 @@ pub fn process_and_respond( let enc = msg.get_data(enc_id); if matches!(enc, DataValue::EncryptedContainer(_)) { let mut dv = enc.clone(); - if dv.decrypt_into_container(&cipher, b"demo-aad").is_some() { + if dv.decrypt_into_container(host_keyring, b"demo-aad").is_some() { if let Some(entries) = dv.as_container() { println!(" Decrypted EncryptedPayload: {:?}", entries); enc_status = format!("EncryptedPayload decrypted OK ({} entries)", entries.len()); @@ -99,7 +91,7 @@ pub fn process_and_respond( if let Some(pk_bundle) = client_pk { let verifier = Ed25519Verifier(pk_bundle.sig_cl_public_key.clone()); let mut dv = secure.clone(); - if dv.decrypt_signed_encrypted_container(&cipher, b"demo-aad").is_some() + if dv.decrypt_signed_encrypted_container(host_keyring, b"demo-aad").is_some() && dv.verify_into_container(&verifier).is_some() { if let Some(entries) = dv.as_container() { diff --git a/example-usage/server/src/main.rs b/example-usage/server/src/main.rs index 414bbe1..1a4c2f4 100644 --- a/example-usage/server/src/main.rs +++ b/example-usage/server/src/main.rs @@ -12,6 +12,11 @@ async fn main() -> Result<(), Box> { let (host_id, host_keyring) = keys::load_or_generate_host_keys("host_keys.json")?; keys::export_host_public_keys(&host_keyring)?; + // The keyring is moved into the host config; keep a copy for decrypting the + // demo payloads clients encrypt to our KEM public key. + let decrypt_keyring = mtp::crypto::Keyring::from_bytes(&host_keyring.to_bytes()) + .expect("re-load host keyring for decryption"); + let (clients, next_id) = clients::load_client_db("clients.json")?; let clients_for_get = clients.clone(); @@ -62,8 +67,12 @@ async fn main() -> Result<(), Box> { match conn.receiver.receive().await { Ok(msg) => { println!("Received: {msg}"); - let response = - handlers::process_and_respond(&msg, tm, conn.client_public_key.as_ref()); + let response = handlers::process_and_respond( + &msg, + tm, + conn.client_public_key.as_ref(), + &decrypt_keyring, + ); println!("Sending: {response}"); conn.sender.send(&response).await?; } diff --git a/host/src/lib.rs b/host/src/lib.rs index cbf0348..9d725be 100644 --- a/host/src/lib.rs +++ b/host/src/lib.rs @@ -1,7 +1,5 @@ -#[cfg(feature = "crypto")] -use mtp_codec::DataType; use mtp_codec::{ - CommunicationValue, DataTypeId, DataValue, Version, + CommunicationValue, DataType, DataValue, TypeMap, Version, registry::{Registry, VersionedCodec}, }; use mtp_common::CommunicationError; @@ -143,9 +141,11 @@ impl MTPHost { Ed25519Signer, PublicKeyBundle, SignatureScheme, verify_ed25519, verify_ml_dsa, }; + let tm = TypeMap::latest(); + // 1. Receive client message first (no host greeting) let msg = receiver.receive().await.ok()?; - let version_str = match msg.get_data(DataTypeId(3)) { + let version_str = match msg.get_data(DataType::Version.to_id(&tm)) { DataValue::Str(s) => s.clone(), _ => { sender.close(); @@ -155,7 +155,7 @@ impl MTPHost { let client_version = Version::parse(&version_str)?; - let client_nonce = match msg.get_data(DataTypeId(7)) { + let client_nonce = match msg.get_data(DataType::ClientNonce.to_id(&tm)) { DataValue::UnsignedNumber(n) => *n, _ => { sender.close(); @@ -163,7 +163,7 @@ impl MTPHost { } }; - let sig_bytes = match msg.get_data(DataTypeId(10)) { + let sig_bytes = match msg.get_data(DataType::Signature.to_id(&tm)) { DataValue::Bytes(b) => b.clone(), _ => { sender.close(); @@ -171,14 +171,15 @@ impl MTPHost { } }; - let pq_sig_bytes: Vec = match msg.get_data(DataTypeId(12)) { + let pq_sig_bytes: Vec = match msg.get_data(DataType::PqSignature.to_id(&tm)) { DataValue::Bytes(b) => b.clone(), _ => vec![], }; - let (assigned_id, client_bundle) = if msg.get_type() == mtp_codec::CommunicationTypeId(15) { + let (assigned_id, client_bundle) = + if msg.get_type() == mtp_codec::CommunicationType::Identification.to_id(&tm) { // LOGIN - let cid = match msg.get_data(DataTypeId(6)) { + let cid = match msg.get_data(DataType::Id.to_id(&tm)) { DataValue::UnsignedNumber(n) => *n as u64, _ => { sender.close(); @@ -238,9 +239,9 @@ impl MTPHost { /* ===== End Signature ===== */ (cid, bundle) - } else if msg.get_type() == mtp_codec::CommunicationTypeId(17) { + } else if msg.get_type() == mtp_codec::CommunicationType::Register.to_id(&tm) { // REGISTER - let bundle = match msg.get_data(DataTypeId(9)) { + let bundle = match msg.get_data(DataType::PublicKeys.to_id(&tm)) { DataValue::Bytes(b) => PublicKeyBundle::from_bytes(b).ok()?, _ => { sender.close(); @@ -359,7 +360,8 @@ impl MTPHost { * (reserved ID 3) mapping to `DataValue::Str("major.minor")`. */ fn extract_version(msg: &CommunicationValue) -> Option { - let value = msg.get_data(DataTypeId(3)); + let tm = TypeMap::latest(); + let value = msg.get_data(DataType::Version.to_id(&tm)); match value { DataValue::Str(s) => Version::parse(s.as_str()), _ => None, @@ -378,7 +380,7 @@ mod tests { mtp_codec::CommunicationType::Identification, &tm, ) - .add_data(DataTypeId(3), DataValue::Str("2.0".to_string())); + .add_data(DataType::Version.to_id(&tm), DataValue::Str("2.0".to_string())); let version = extract_version(&msg); assert_eq!(version, Some(Version(2, 0))); } @@ -400,7 +402,7 @@ mod tests { mtp_codec::CommunicationType::Identification, &tm, ) - .add_data(DataTypeId(3), DataValue::UnsignedNumber(42)); + .add_data(DataType::Version.to_id(&tm), DataValue::UnsignedNumber(42)); assert!(extract_version(&msg).is_none()); } } diff --git a/transport/src/connection.rs b/transport/src/connection.rs index ed6bbd8..91abc39 100644 --- a/transport/src/connection.rs +++ b/transport/src/connection.rs @@ -83,7 +83,7 @@ impl Sender { data: &CommunicationValue, policy: &Policy, ) -> Result<(), CommunicationError> { - let bytes = data.to_bytes(); + let bytes = data.to_bytes().map_err(|_| CommunicationError::Encode)?; if bytes.len() as u64 > policy.max_message_size || bytes.len() as u64 >= policy.close_frame_len as u64 { diff --git a/wasm/src/client.rs b/wasm/src/client.rs index 18951af..e73b10e 100644 --- a/wasm/src/client.rs +++ b/wasm/src/client.rs @@ -97,7 +97,10 @@ impl WasmClient { let ident = CommunicationValue::new(CommunicationType::Identification) .add_typed_default(DataType::Version, DataValue::Str(version_str)) .add_typed_default(DataType::Id, DataValue::UnsignedNumber(config.client_id as u128)); - transport.send_frame(&ident.to_bytes()).await?; + let ident_bytes = ident + .to_bytes() + .map_err(|e| js_error(&format!("encode failed: {}", e)))?; + transport.send_frame(&ident_bytes).await?; self.transport = Some(transport); self.set_state(ConnectionState::Connected); @@ -155,7 +158,8 @@ impl WasmClient { .add_typed_default(DataType::Id, DataValue::UnsignedNumber(client_id as u128)) .add_typed_default(DataType::ClientNonce, DataValue::UnsignedNumber(client_nonce)) .add_typed_default(DataType::Signature, DataValue::Bytes(signature)) - .to_bytes(); + .to_bytes() + .map_err(|e| js_error(&format!("encode failed: {}", e)))?; let transport = WasmTransport::connect(&config.url, config.server_certificate_hashes.clone()).await?; let inner = transport.inner().clone(); @@ -247,7 +251,8 @@ impl WasmClient { .add_typed_default(DataType::ClientNonce, DataValue::UnsignedNumber(client_nonce)) .add_typed_default(DataType::PublicKeys, DataValue::Bytes(pk_bytes)) .add_typed_default(DataType::Signature, DataValue::Bytes(signature)) - .to_bytes(); + .to_bytes() + .map_err(|e| js_error(&format!("encode failed: {}", e)))?; let transport = WasmTransport::connect(&config.url, config.server_certificate_hashes.clone()).await?; let inner = transport.inner().clone(); diff --git a/wasm/src/message.rs b/wasm/src/message.rs index 52699e8..908dbab 100644 --- a/wasm/src/message.rs +++ b/wasm/src/message.rs @@ -1,13 +1,10 @@ use wasm_bindgen::prelude::*; use mtp_codec::{ - CommunicationType, CommunicationValue, DataType, DataTypeId, DataValue, + CommunicationType, CommunicationTypeId, CommunicationValue, DataType, DataTypeId, DataValue, }; use mtp_type_map::communication_type_name; -use mtp_crypto::{ - ChaCha20Poly1305, Ed25519Signer, Keyring, SigAlgorithm, - derive_encryption_key, -}; +use mtp_crypto::{Ed25519Signer, EncryptionType, Keyring, SigAlgorithm}; use crate::error::js_error; @@ -18,7 +15,7 @@ pub fn build_ping_frame( description: &str, timestamp: u64, data: &[u8], -) -> Vec { +) -> Result, JsValue> { let mut msg = CommunicationValue::new(CommunicationType::Ping) .add_typed_default(DataType::Description, DataValue::Str(description.to_string())) .add_typed_default(DataType::Timestamp, DataValue::UnsignedNumber(timestamp as u128)) @@ -29,6 +26,7 @@ pub fn build_ping_frame( } msg.to_bytes() + .map_err(|e| js_error(&format!("encode failed: {}", e))) } /// Build a demo Ping frame with encrypted and signed containers @@ -38,14 +36,10 @@ pub fn build_demo_message(client_id: u64, keyring_bytes: &[u8]) -> Result Result Result Result Result< let frame = CommunicationValue::new(comm_type_enum) .with_id(id) .add_data(DataTypeId(32), DataValue::Str(json_data.to_string())) - .to_bytes(); + .to_bytes() + .map_err(|e| js_error(&format!("encode failed: {}", e)))?; Ok(frame) } @@ -214,7 +210,7 @@ mod tests { #[wasm_bindgen_test] fn build_ping_frame_roundtrip() { - let bytes = build_ping_frame(42, "test-ping", 1234567890, &[]); + let bytes = build_ping_frame(42, "test-ping", 1234567890, &[]).expect("encode failed"); let cv = CommunicationValue::from_bytes(&bytes).expect("decode failed"); assert_eq!(cv.get_type(), CommunicationTypeId(19)); // Ping @@ -226,7 +222,7 @@ mod tests { #[wasm_bindgen_test] fn build_ping_frame_with_data() { let payload = b"attachment-data"; - let bytes = build_ping_frame(99, "with-data", 555, payload); + let bytes = build_ping_frame(99, "with-data", 555, payload).expect("encode failed"); let cv = CommunicationValue::from_bytes(&bytes).expect("decode failed"); assert_eq!(cv.get_type(), CommunicationTypeId(19)); @@ -238,22 +234,16 @@ mod tests { #[wasm_bindgen_test] fn build_ping_frame_client_id_zero() { - let bytes = build_ping_frame(0, "zero-id", 0, &[]); + let bytes = build_ping_frame(0, "zero-id", 0, &[]).expect("encode failed"); let cv = CommunicationValue::from_bytes(&bytes).expect("decode failed"); assert_eq!(cv.get_sender(), 0); } #[wasm_bindgen_test] fn build_demo_message_roundtrip() { - let (_signer, sk, pk) = Ed25519Signer::generate(); - let keyring = Keyring::new( - mtp_crypto::KemPublicKey::new(vec![]), - mtp_crypto::KemPrivateKey::new(vec![]), - mtp_crypto::SignaturePqPublicKey::new(vec![]), - mtp_crypto::SignaturePqPrivateKey::new(vec![]), - pk, - sk, - ); + // A full keyring is required: the demo now KEM-encrypts to its own + // public key, so the KEM keypair must be real. + let keyring = Keyring::generate(); let keyring_bytes = keyring.to_bytes(); let result = build_demo_message(7, &keyring_bytes); @@ -282,7 +272,8 @@ mod tests { .add_typed_default(DataType::ClientNonce, DataValue::UnsignedNumber(999)) .add_typed_default(DataType::Id, DataValue::UnsignedNumber(42)) .add_typed_default(DataType::Timestamp, DataValue::UnsignedNumber(12345)) - .to_bytes(); + .to_bytes() + .expect("encode failed"); let result = parse_auth_response(&resp).expect("parse failed"); @@ -299,7 +290,8 @@ mod tests { fn parse_auth_response_rejected() { let resp = CommunicationValue::new(CommunicationType::IdentificationResponse) .add_typed_default(DataType::Connected, DataValue::BoolFalse) - .to_bytes(); + .to_bytes() + .expect("encode failed"); let result = parse_auth_response(&resp).expect("parse failed"); @@ -318,7 +310,8 @@ mod tests { let resp = CommunicationValue::new(CommunicationType::IdentificationResponse) .add_typed_default(DataType::Connected, DataValue::BoolTrue) .add_typed_default(DataType::Signature, DataValue::Bytes(sig_bytes.clone())) - .to_bytes(); + .to_bytes() + .expect("encode failed"); let result = parse_auth_response(&resp).expect("parse failed");