diff --git a/client/Cargo.toml b/client/Cargo.toml index 09c50c3..d48a31b 100644 --- a/client/Cargo.toml +++ b/client/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "client" +name = "mtp-client" version = "0.1.0" edition = "2024" diff --git a/codec/Cargo.toml b/codec/Cargo.toml index 652fbe6..f5e514a 100644 --- a/codec/Cargo.toml +++ b/codec/Cargo.toml @@ -4,8 +4,8 @@ version = "0.1.0" edition = "2024" [dependencies] -type-map = { path = "../type-map" } -common = { path = "../common" } +mtp-type-map = { path = "../type-map" } +mtp-common = { path = "../common" } mtp-crypto = { path = "../crypto", optional = true } base64 = "*" byteorder = "*" diff --git a/codec/src/communication_value.rs b/codec/src/communication_value.rs index 354efc2..255617c 100644 --- a/codec/src/communication_value.rs +++ b/codec/src/communication_value.rs @@ -4,7 +4,7 @@ use std::io::{Cursor, Read}; use crate::data_value::DataValue; use crate::rand_u32; -use type_map::{CommTypeId, DataTypeId}; +use mtp_type_map::{CommTypeId, DataTypeId}; #[derive(Debug, Clone, PartialEq, Eq)] pub struct CommunicationValue { @@ -201,7 +201,7 @@ impl CommunicationValue { mod tests { use super::*; use crate::data_value::DataValue; - use type_map::{CommTypeId, DataTypeId}; + use mtp_type_map::{CommTypeId, DataTypeId}; fn roundtrip(cv: CommunicationValue) -> CommunicationValue { let bytes = cv.to_bytes(); diff --git a/codec/src/data_value.rs b/codec/src/data_value.rs index f283dd2..cd57d84 100644 --- a/codec/src/data_value.rs +++ b/codec/src/data_value.rs @@ -5,7 +5,7 @@ use std::collections::BTreeMap; use std::hash::{Hash, Hasher}; use std::io::Cursor; -use type_map::DataTypeId; +use mtp_type_map::DataTypeId; #[derive(Debug, Clone, PartialEq, Eq)] pub enum DataKind { @@ -52,7 +52,7 @@ pub enum DataValue { impl DataValue { /* * Container format: - * [2 bytes u16 entry_count] // number of key-value pairs + * [2 bytes u16 entry_count] // length of the container * [1 byte kind] // DataValue kind marker * [if kind == BOOL_TRUE or BOOL_FALSE:] * [1 byte key] // DataTypes discriminant @@ -64,11 +64,15 @@ impl DataValue { * Kind markers: * 0x01 => BoolTrue * 0x02 => BoolFalse - * 0x03 => Number (i64, 8 bytes big-endian) - * 0x04 => Str (UTF-8 bytes) - * 0x05 => Array (nested container format) - * 0x06 => Container (nested container format) - * 0x07 => Null + * 0x03 => Signed Number (i64, 8 bytes big-endian) + * 0x04 => Unsigned Number (u64, 8 bytes big-endian) + * 0x05 => Float (1 byte exponent, 3 bytes mantissa) + * 0x06 => Str (UTF-8 bytes) + * 0x07 => Bytes + * 0x08 => Array (nested container format) + * 0x09 => Container (nested container format) + * 0x0A => EncryptedContainer (nested container format) + * 0x0B => Null */ const KIND_BOOL_TRUE: u8 = 0x01; const KIND_BOOL_FALSE: u8 = 0x02; diff --git a/codec/src/lib.rs b/codec/src/lib.rs index 071dbe6..3f53233 100644 --- a/codec/src/lib.rs +++ b/codec/src/lib.rs @@ -8,9 +8,9 @@ pub use data_value::{DataKind, DataValue}; pub use util::rand_u32; -pub use type_map::{CommTypeId, DataTypeId, TypeMap, Version}; +pub use mtp_type_map::{CommTypeId, DataTypeId, TypeMap, Version}; -use common::CodecError; +use mtp_common::CodecError; pub fn encode(_value: &DataValue, _typemap: &TypeMap) -> Result, CodecError> { // write header using typemap.data_id(), serialize value diff --git a/common/Cargo.toml b/common/Cargo.toml index fa02773..e4e2bcf 100644 --- a/common/Cargo.toml +++ b/common/Cargo.toml @@ -1,6 +1,17 @@ [package] -name = "common" +name = "mtp-common" version = "0.1.0" edition = "2024" [dependencies] +thiserror = "2.0.18" +wtransport = { version = "0.7.1", default-features = false, features = [ + "aws-lc-rs", + "quinn", + "self-signed", +] } +rustls = { version = "0.23.40" } +quinn = { version = "0.11.9", default-features = false, features = [ + "rustls-aws-lc-rs", + "rustls", +] } diff --git a/common/src/lib.rs b/common/src/lib.rs index 0b3c97b..0ec7f70 100644 --- a/common/src/lib.rs +++ b/common/src/lib.rs @@ -1,3 +1,5 @@ +use thiserror::Error; + pub enum RegistryError { ReservedCommId(u8, String), } @@ -10,3 +12,81 @@ pub enum CodecError { ReservedCommunicationType(u8), InvalidEncoding, } + +#[derive(Debug, Error, Clone)] +pub enum CommunicationError { + #[error("Use after Closed")] + UseAfterClosed, + + #[error("Connection closed by local shutdown")] + ClosedLocally, + + #[error("Connection closed by peer")] + ClosedByPeer, + + #[error("Connection terminated unexpectedly")] + ConnectionLost, + + #[error("QUIC error: {0}")] + Quinn(#[from] quinn::ConnectionError), + + #[error("ParseCommunicationValue error")] + ParseCommunicationValue, + + #[error("Parse Certificate error")] + CertificateParseFailed, + + #[error("Loading Certificate error")] + CertificateLoadFailed, + + #[error("ParseBool error: {0}")] + ParseBool(#[from] std::str::ParseBoolError), + + #[error("ParseInt error: {0}")] + ParseInt(#[from] std::num::ParseIntError), + + #[error("ParseFloat error: {0}")] + ParseFloat(#[from] std::num::ParseFloatError), + + #[error("ParseAddr error: {0}")] + ParseAddr(#[from] std::net::AddrParseError), + + #[error("Connection error: {0}")] + ConnectionError(#[from] wtransport::error::ConnectionError), + + #[error("Connecting error: {0}")] + ConnectingError(String), + + #[error("ReadToEnd error: {0}")] + ReadToEndError(#[from] quinn::ReadToEndError), + + #[error("Write error: {0}")] + WriteError(#[from] quinn::WriteError), + + #[error("Closed error: {0}")] + ClosedError(#[from] quinn::ClosedStream), + + #[error("Message too large")] + MessageTooLarge, + + #[error("ReadExactError: {0}")] + ReadExactError(#[from] quinn::ReadExactError), + + #[error("Stream Closed")] + StreamClosed, + + #[error("Stream Error")] + StreamError, + + #[error("Stream Error: {0}")] + StreamWriteError(#[from] wtransport::error::StreamWriteError), + + #[error("Read Exact Error: {0}")] + StreamReadExactError(#[from] wtransport::error::StreamReadExactError), + + #[error("Crypto Provider Install Error")] + CryptoProviderInstallFailed, + + #[error("Other: {0}")] + Other(String), +} diff --git a/crypto/Cargo.toml b/crypto/Cargo.toml index d882c57..bb3004e 100644 --- a/crypto/Cargo.toml +++ b/crypto/Cargo.toml @@ -1,6 +1,21 @@ [package] name = "mtp-crypto" -version = "0.1.0" +version = "0.2.0" edition = "2024" [dependencies] +chacha20poly1305 = { version = "0.10", optional = true } +aes-gcm = { version = "0.10", optional = true } +ed25519-dalek = { version = "2.1", optional = true, features = ["pkcs8", "pem"] } +hkdf = { version = "0.12", optional = true } +sha2 = { version = "0.10", optional = true } +zeroize = { version = "1.7", features = ["derive"] } +rand_core = { version = "0.6", features = ["getrandom"] } +getrandom = "0.2" +mlkem-tls = { version = "0.2", optional = true } +ml-dsa = { version = "0.0.4", optional = true } + +[features] +default = ["chacha20poly1305", "ed25519-dalek", "hkdf", "sha2"] +full = ["chacha20poly1305", "aes-gcm", "ed25519-dalek", "hkdf", "sha2"] +pqc = ["mlkem-tls", "ml-dsa"] diff --git a/crypto/README.md b/crypto/README.md new file mode 100644 index 0000000..9c1c593 --- /dev/null +++ b/crypto/README.md @@ -0,0 +1,113 @@ +# mtp-crypto + +Cryptographic primitives for the MTP protocol. Classical and post-quantum. + +## Features + +| Feature | Primitives | Status | +|---------|-----------|--------| +| `default` | XChaCha20-Poly1305, Ed25519, HKDF-SHA-256, SHA-256 | Classical | +| `full` | default + AES-256-GCM | Classical | +| `pqc` | ML-KEM-768+X25519 hybrid KEM, ML-DSA-65 | Post-quantum | + +## AEAD + +XChaCha20-Poly1305 (default) and AES-256-GCM (`full` feature). Nonce is prepended to ciphertext. + +```rust +use mtp_crypto::{ChaCha20Poly1305, AeadEncrypt, AeadDecrypt}; + +let cipher = ChaCha20Poly1305::new([0u8; 32]); +let ct = cipher.encrypt(b"hello", b"aad")?; +let pt = cipher.decrypt(&ct, b"aad")?; +``` + +## Signatures + +### Ed25519 (classical) + +```rust +use mtp_crypto::{Ed25519Signer, SignatureScheme}; + +let (signer, sk, pk) = Ed25519Signer::generate(); +let sig = signer.sign(b"message")?; +signer.verify(b"message", &sig)?; +``` + +### ML-DSA-65 (post-quantum, requires `pqc`) + +```rust +use mtp_crypto::{MlDsaSigner, SignatureScheme}; + +let (signer, sk, pk) = MlDsaSigner::generate(); +let sig = signer.sign(b"message")?; +signer.verify(b"message", &sig)?; + +// Load from stored bytes +let signer = MlDsaSigner::new(&sk, &pk)?; +``` + +### Dual signatures (requires `pqc`) + +```rust +use mtp_crypto::{sign_dual, DualSignature, Ed25519Signer, MlDsaSigner}; + +let (ed_signer, _, _) = Ed25519Signer::generate(); +let (ml_signer, _, _) = MlDsaSigner::generate(); +let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg"); +dual.verify(ed_signer.verifying_key(), ml_signer.verifying_key(), b"msg")?; +``` + +## Hybrid KEM (requires `pqc`) + +X25519 + ML-KEM-768. 64-byte shared secret. Feed into HKDF before use. + +```rust +use mtp_crypto::HybridKem; + +let (sk, pk) = HybridKem::generate_keypair(); +let enc = HybridKem::encapsulate(&pk)?; +let ss = HybridKem::decapsulate(&sk, &enc.ciphertext)?; +assert_eq!(enc.shared_secret, ss); +``` + +## KDF + +```rust +use mtp_crypto::{hkdf_expand, derive_encryption_key}; + +let key = derive_encryption_key(b"ikm", b"salt", b"context")?; +``` + +## Hashing + +```rust +use mtp_crypto::{sha256, sha256_double}; + +let h = sha256(b"data"); +let h2 = sha256_double(b"data"); +``` + +## Key types + +| Type | Secret | Zeroized | +|------|--------|----------| +| `EncryptionPrivateKey` | KEM/ECDH secret | Yes | +| `EncryptionPublicKey` | KEM/ECDH public | No | +| `SignaturePrivateKey` | Classical signing key | Yes | +| `SignaturePublicKey` | Classical verifying key | No | +| `KemPrivateKey` | Hybrid KEM secret | Yes | +| `KemPublicKey` | Hybrid KEM public | No | +| `SignaturePqPrivateKey` | PQC signing key | Yes | +| `SignaturePqPublicKey` | PQC verifying key | No | + +`KeyGroup` holds classical keys; `Keyring` holds all six (hybrid KEM + PQ sig + classical sig). + +## Feature flags + +```toml +[dependencies] +mtp-crypto = { path = "../crypto" } # classical +mtp-crypto = { path = "../crypto", features = ["pqc"] } # post-quantum +mtp-crypto = { path = "../crypto", features = ["full", "pqc"] } # all +``` diff --git a/crypto/src/aead.rs b/crypto/src/aead.rs new file mode 100644 index 0000000..4e65064 --- /dev/null +++ b/crypto/src/aead.rs @@ -0,0 +1,171 @@ +use crate::error::CryptoError; + +#[cfg(any(feature = "chacha20poly1305", feature = "aes-gcm"))] +use rand_core::OsRng; + +#[cfg(any(feature = "chacha20poly1305", feature = "aes-gcm"))] +use rand_core::RngCore; + +pub trait AeadEncrypt { + fn encrypt(&self, plaintext: &[u8], aad: &[u8]) -> Result, CryptoError>; +} + +pub trait AeadDecrypt { + fn decrypt(&self, ciphertext: &[u8], aad: &[u8]) -> Result, CryptoError>; +} + +pub trait AeadCipher: AeadEncrypt + AeadDecrypt { + fn key_size() -> usize; +} + +#[cfg(feature = "chacha20poly1305")] +pub struct ChaCha20Poly1305 { + key: [u8; 32], +} + +#[cfg(feature = "chacha20poly1305")] +impl ChaCha20Poly1305 { + pub fn new(key: [u8; 32]) -> Self { + Self { key } + } +} + +#[cfg(feature = "chacha20poly1305")] +impl AeadEncrypt for ChaCha20Poly1305 { + fn encrypt(&self, plaintext: &[u8], aad: &[u8]) -> Result, CryptoError> { + use chacha20poly1305::aead::{Aead, KeyInit, Payload}; + use chacha20poly1305::XChaCha20Poly1305; + use chacha20poly1305::XNonce; + + let key = chacha20poly1305::Key::from_slice(&self.key); + let cipher = XChaCha20Poly1305::new(key); + + let mut nonce = [0u8; 24]; + OsRng.fill_bytes(&mut nonce); + let nonce_ref = XNonce::from_slice(&nonce); + + let payload = Payload { + msg: plaintext, + aad, + }; + + let mut ciphertext = cipher + .encrypt(nonce_ref, payload) + .map_err(|_| CryptoError::EncryptionFailed)?; + + let mut out = Vec::with_capacity(nonce.len() + ciphertext.len()); + out.extend_from_slice(&nonce); + out.append(&mut ciphertext); + Ok(out) + } +} + +#[cfg(feature = "chacha20poly1305")] +impl AeadDecrypt for ChaCha20Poly1305 { + fn decrypt(&self, ciphertext: &[u8], aad: &[u8]) -> Result, CryptoError> { + use chacha20poly1305::aead::{Aead, KeyInit, Payload}; + use chacha20poly1305::XChaCha20Poly1305; + use chacha20poly1305::XNonce; + + if ciphertext.len() < 24 { + return Err(CryptoError::InvalidNonceLength); + } + + let (nonce, ct) = ciphertext.split_at(24); + let key = chacha20poly1305::Key::from_slice(&self.key); + let cipher = XChaCha20Poly1305::new(key); + let nonce_ref = XNonce::from_slice(nonce); + + let payload = Payload { + msg: ct, + aad, + }; + + cipher + .decrypt(nonce_ref, payload) + .map_err(|_| CryptoError::DecryptionFailed) + } +} + +#[cfg(feature = "chacha20poly1305")] +impl AeadCipher for ChaCha20Poly1305 { + fn key_size() -> usize { + 32 + } +} + +#[cfg(feature = "aes-gcm")] +pub struct Aes256Gcm { + key: [u8; 32], +} + +#[cfg(feature = "aes-gcm")] +impl Aes256Gcm { + pub fn new(key: [u8; 32]) -> Self { + Self { key } + } +} + +#[cfg(feature = "aes-gcm")] +impl AeadEncrypt for Aes256Gcm { + fn encrypt(&self, plaintext: &[u8], aad: &[u8]) -> Result, CryptoError> { + use aes_gcm::aead::{Aead, KeyInit, Payload}; + use aes_gcm::Aes256Gcm as AesGcmInner; + use aes_gcm::Nonce; + + let key = aes_gcm::Key::::from_slice(&self.key); + let cipher = AesGcmInner::new(key); + + let mut nonce = [0u8; 12]; + OsRng.fill_bytes(&mut nonce); + let nonce_ref = Nonce::from_slice(&nonce); + + let payload = Payload { + msg: plaintext, + aad, + }; + + let mut ciphertext = cipher + .encrypt(nonce_ref, payload) + .map_err(|_| CryptoError::EncryptionFailed)?; + + let mut out = Vec::with_capacity(nonce.len() + ciphertext.len()); + out.extend_from_slice(&nonce); + out.append(&mut ciphertext); + Ok(out) + } +} + +#[cfg(feature = "aes-gcm")] +impl AeadDecrypt for Aes256Gcm { + fn decrypt(&self, ciphertext: &[u8], aad: &[u8]) -> Result, CryptoError> { + use aes_gcm::aead::{Aead, KeyInit, Payload}; + use aes_gcm::Aes256Gcm as AesGcmInner; + use aes_gcm::Nonce; + + if ciphertext.len() < 12 { + return Err(CryptoError::InvalidNonceLength); + } + + let (nonce, ct) = ciphertext.split_at(12); + let key = aes_gcm::Key::::from_slice(&self.key); + let cipher = AesGcmInner::new(key); + let nonce_ref = Nonce::from_slice(nonce); + + let payload = Payload { + msg: ct, + aad, + }; + + cipher + .decrypt(nonce_ref, payload) + .map_err(|_| CryptoError::DecryptionFailed) + } +} + +#[cfg(feature = "aes-gcm")] +impl AeadCipher for Aes256Gcm { + fn key_size() -> usize { + 32 + } +} diff --git a/crypto/src/error.rs b/crypto/src/error.rs new file mode 100644 index 0000000..53ad891 --- /dev/null +++ b/crypto/src/error.rs @@ -0,0 +1,38 @@ +use std::fmt; + +#[derive(Debug, Clone)] +pub enum CryptoError { + EncryptionFailed, + DecryptionFailed, + InvalidKeyLength, + InvalidNonceLength, + InvalidSignature, + SigningFailed, + VerificationFailed, + KeyGenerationFailed, + KdfError, + KemEncapsulationFailed, + KemDecapsulationFailed, + UnknownAlgorithm, +} + +impl fmt::Display for CryptoError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + CryptoError::EncryptionFailed => write!(f, "encryption failed"), + CryptoError::DecryptionFailed => write!(f, "decryption failed"), + CryptoError::InvalidKeyLength => write!(f, "invalid key length"), + CryptoError::InvalidNonceLength => write!(f, "invalid nonce length"), + CryptoError::InvalidSignature => write!(f, "invalid signature"), + CryptoError::SigningFailed => write!(f, "signing failed"), + CryptoError::VerificationFailed => write!(f, "verification failed"), + CryptoError::KeyGenerationFailed => write!(f, "key generation failed"), + CryptoError::KdfError => write!(f, "KDF error"), + CryptoError::KemEncapsulationFailed => write!(f, "KEM encapsulation failed"), + CryptoError::KemDecapsulationFailed => write!(f, "KEM decapsulation failed"), + CryptoError::UnknownAlgorithm => write!(f, "unknown algorithm"), + } + } +} + +impl std::error::Error for CryptoError {} diff --git a/crypto/src/hash.rs b/crypto/src/hash.rs new file mode 100644 index 0000000..39765b3 --- /dev/null +++ b/crypto/src/hash.rs @@ -0,0 +1,29 @@ +use sha2::Digest; + +pub fn sha256(data: &[u8]) -> [u8; 32] { + let mut hasher = sha2::Sha256::new(); + hasher.update(data); + let result = hasher.finalize(); + result.into() +} + +pub fn sha256_double(data: &[u8]) -> [u8; 32] { + sha256(&sha256(data)) +} + +pub struct Sha256Hasher(sha2::Sha256); + +impl Sha256Hasher { + pub fn new() -> Self { + Self(sha2::Sha256::new()) + } + + pub fn update(&mut self, data: &[u8]) { + self.0.update(data); + } + + pub fn finalize(self) -> [u8; 32] { + let result = self.0.finalize(); + result.into() + } +} diff --git a/crypto/src/kdf.rs b/crypto/src/kdf.rs new file mode 100644 index 0000000..113880c --- /dev/null +++ b/crypto/src/kdf.rs @@ -0,0 +1,34 @@ +use crate::error::CryptoError; +use hkdf::Hkdf; +use sha2::Sha256; + +pub fn hkdf_expand( + ikm: &[u8], + salt: &[u8], + info: &[u8], + okm_len: usize, +) -> Result, CryptoError> { + let hk = Hkdf::::new(Some(salt), ikm); + let mut okm = vec![0u8; okm_len]; + hk.expand(info, &mut okm) + .map_err(|_| CryptoError::KdfError)?; + Ok(okm) +} + +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 +} + +pub fn derive_encryption_key( + ikm: &[u8], + salt: &[u8], + context: &[u8], +) -> Result<[u8; 32], CryptoError> { + let key = hkdf_expand(ikm, salt, context, 32)?; + let mut out = [0u8; 32]; + out.copy_from_slice(&key); + Ok(out) +} diff --git a/crypto/src/kem.rs b/crypto/src/kem.rs new file mode 100644 index 0000000..9e0c1c6 --- /dev/null +++ b/crypto/src/kem.rs @@ -0,0 +1,45 @@ +use crate::error::CryptoError; +use crate::keypair::{KemPrivateKey, KemPublicKey}; + +pub struct Encapsulated { + pub ciphertext: Vec, + pub shared_secret: Vec, +} + +#[cfg(feature = "mlkem-tls")] +pub struct HybridKem; + +#[cfg(feature = "mlkem-tls")] +impl HybridKem { + pub fn generate_keypair() -> (KemPrivateKey, KemPublicKey) { + let (dk, ek) = + mlkem_tls::X25519MlKem768::keygen(&mut rand_core::OsRng); + ( + KemPrivateKey::new(dk.as_bytes().to_vec()), + KemPublicKey::new(ek.as_bytes().to_vec()), + ) + } + + pub fn encapsulate(recipient_pk: &KemPublicKey) -> Result { + let ek = mlkem_tls::EncapsKey768::try_from(recipient_pk.as_bytes()) + .map_err(|_| CryptoError::KemEncapsulationFailed)?; + let (ct, ss) = + mlkem_tls::X25519MlKem768::encapsulate(&ek, &mut rand_core::OsRng); + Ok(Encapsulated { + ciphertext: ct.as_bytes().to_vec(), + shared_secret: ss.as_bytes().to_vec(), + }) + } + + pub fn decapsulate( + recipient_sk: &KemPrivateKey, + ciphertext: &[u8], + ) -> Result, CryptoError> { + let dk = mlkem_tls::DecapsKey768::try_from(recipient_sk.as_bytes()) + .map_err(|_| CryptoError::KemDecapsulationFailed)?; + let ct = mlkem_tls::Ciphertext768Hybrid::try_from(ciphertext) + .map_err(|_| CryptoError::KemDecapsulationFailed)?; + let ss = mlkem_tls::X25519MlKem768::decapsulate(&dk, &ct); + Ok(ss.as_bytes().to_vec()) + } +} diff --git a/crypto/src/keypair.rs b/crypto/src/keypair.rs new file mode 100644 index 0000000..0fcb29f --- /dev/null +++ b/crypto/src/keypair.rs @@ -0,0 +1,212 @@ +use zeroize::{Zeroize, ZeroizeOnDrop}; + +#[derive(Zeroize, ZeroizeOnDrop)] +pub struct EncryptionPrivateKey(Vec); + +impl EncryptionPrivateKey { + pub fn new(bytes: Vec) -> Self { + Self(bytes) + } + + pub fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + +impl From> for EncryptionPrivateKey { + fn from(bytes: Vec) -> Self { + Self(bytes) + } +} + +#[derive(Zeroize, ZeroizeOnDrop)] +pub struct SignaturePrivateKey(Vec); + +impl SignaturePrivateKey { + pub fn new(bytes: Vec) -> Self { + Self(bytes) + } + + pub fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + +impl From> for SignaturePrivateKey { + fn from(bytes: Vec) -> Self { + Self(bytes) + } +} + +#[derive(Clone)] +pub struct EncryptionPublicKey(Vec); + +impl EncryptionPublicKey { + pub fn new(bytes: Vec) -> Self { + Self(bytes) + } + + pub fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + +impl From> for EncryptionPublicKey { + fn from(bytes: Vec) -> Self { + Self(bytes) + } +} + +#[derive(Clone)] +pub struct SignaturePublicKey(Vec); + +impl SignaturePublicKey { + pub fn new(bytes: Vec) -> Self { + Self(bytes) + } + + pub fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + +impl From> for SignaturePublicKey { + fn from(bytes: Vec) -> Self { + Self(bytes) + } +} + +#[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, + } + } +} + +#[derive(Zeroize, ZeroizeOnDrop)] +pub struct KemPrivateKey(Vec); + +impl KemPrivateKey { + pub fn new(bytes: Vec) -> Self { + Self(bytes) + } + + pub fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + +impl From> for KemPrivateKey { + fn from(bytes: Vec) -> Self { + Self(bytes) + } +} + +#[derive(Clone)] +pub struct KemPublicKey(Vec); + +impl KemPublicKey { + pub fn new(bytes: Vec) -> Self { + Self(bytes) + } + + pub fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + +impl From> for KemPublicKey { + fn from(bytes: Vec) -> Self { + Self(bytes) + } +} + +#[derive(Clone)] +pub struct SignaturePqPublicKey(Vec); + +impl SignaturePqPublicKey { + pub fn new(bytes: Vec) -> Self { + Self(bytes) + } + + pub fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + +impl From> for SignaturePqPublicKey { + fn from(bytes: Vec) -> Self { + Self(bytes) + } +} + +#[derive(Zeroize, ZeroizeOnDrop)] +pub struct SignaturePqPrivateKey(Vec); + +impl SignaturePqPrivateKey { + pub fn new(bytes: Vec) -> Self { + Self(bytes) + } + + pub fn as_bytes(&self) -> &[u8] { + &self.0 + } +} + +impl From> for SignaturePqPrivateKey { + fn from(bytes: Vec) -> Self { + Self(bytes) + } +} + +#[derive(ZeroizeOnDrop)] +pub struct Keyring { + #[zeroize(skip)] + pub kem_public_key: KemPublicKey, + pub kem_secret_key: KemPrivateKey, + #[zeroize(skip)] + pub sig_pq_public_key: SignaturePqPublicKey, + pub sig_pq_secret_key: SignaturePqPrivateKey, + #[zeroize(skip)] + pub sig_cl_public_key: SignaturePublicKey, + pub sig_cl_secret_key: SignaturePrivateKey, +} + +impl Keyring { + pub fn new( + kem_public_key: KemPublicKey, + kem_secret_key: KemPrivateKey, + sig_pq_public_key: SignaturePqPublicKey, + sig_pq_secret_key: SignaturePqPrivateKey, + sig_cl_public_key: SignaturePublicKey, + sig_cl_secret_key: SignaturePrivateKey, + ) -> Self { + Self { + kem_public_key, + kem_secret_key, + sig_pq_public_key, + sig_pq_secret_key, + sig_cl_public_key, + sig_cl_secret_key, + } + } +} diff --git a/crypto/src/lib.rs b/crypto/src/lib.rs index b93cf3f..d726960 100644 --- a/crypto/src/lib.rs +++ b/crypto/src/lib.rs @@ -1,14 +1,46 @@ -pub fn add(left: u64, right: u64) -> u64 { - left + right -} +pub mod aead; +pub mod error; +pub mod keypair; -#[cfg(test)] -mod tests { - use super::*; +#[cfg(feature = "sha2")] +pub mod hash; - #[test] - fn it_works() { - let result = add(2, 2); - assert_eq!(result, 4); - } -} +#[cfg(feature = "hkdf")] +pub mod kdf; + +#[cfg(any(feature = "ed25519-dalek", feature = "ml-dsa"))] +pub mod sign; + +#[cfg(feature = "mlkem-tls")] +pub mod kem; + +pub use aead::{AeadCipher, AeadDecrypt, AeadEncrypt}; +pub use error::CryptoError; +pub use keypair::{ + EncryptionPrivateKey, EncryptionPublicKey, KemPrivateKey, KemPublicKey, KeyGroup, Keyring, + SignaturePqPrivateKey, SignaturePqPublicKey, SignaturePrivateKey, SignaturePublicKey, +}; + +#[cfg(feature = "chacha20poly1305")] +pub use aead::ChaCha20Poly1305; + +#[cfg(feature = "aes-gcm")] +pub use aead::Aes256Gcm; + +#[cfg(feature = "ed25519-dalek")] +pub use sign::{verify_ed25519, Ed25519Signer, SignatureScheme}; + +#[cfg(feature = "ml-dsa")] +pub use sign::{verify_ml_dsa, MlDsaSigner}; + +#[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))] +pub use sign::{sign_dual, DualSignature}; + +#[cfg(feature = "sha2")] +pub use hash::{sha256, sha256_double, Sha256Hasher}; + +#[cfg(feature = "hkdf")] +pub use kdf::{derive_encryption_key, hkdf_expand, hkdf_extract}; + +#[cfg(feature = "mlkem-tls")] +pub use kem::HybridKem; diff --git a/crypto/src/sign.rs b/crypto/src/sign.rs new file mode 100644 index 0000000..fbc641a --- /dev/null +++ b/crypto/src/sign.rs @@ -0,0 +1,232 @@ +use crate::error::CryptoError; + +#[cfg(feature = "ed25519-dalek")] +use crate::keypair::{SignaturePrivateKey, SignaturePublicKey}; + +#[cfg(feature = "ed25519-dalek")] +use rand_core::RngCore; + +#[cfg(feature = "ml-dsa")] +use crate::keypair::{SignaturePqPrivateKey, SignaturePqPublicKey}; + +pub trait SignatureScheme { + fn sign(&self, msg: &[u8]) -> Result, CryptoError>; + fn verify(&self, msg: &[u8], signature: &[u8]) -> Result<(), CryptoError>; +} + +#[cfg(feature = "ed25519-dalek")] +pub struct Ed25519Signer { + secret: ed25519_dalek::SigningKey, + public: ed25519_dalek::VerifyingKey, +} + +#[cfg(feature = "ed25519-dalek")] +impl Ed25519Signer { + pub fn new(secret_key: &SignaturePrivateKey) -> Result { + let bytes: [u8; 32] = secret_key + .as_bytes() + .try_into() + .map_err(|_| CryptoError::KeyGenerationFailed)?; + let secret = ed25519_dalek::SigningKey::from_bytes(&bytes); + let public = secret.verifying_key(); + Ok(Self { secret, public }) + } + + pub fn generate() -> (Self, SignaturePrivateKey, SignaturePublicKey) { + let mut bytes = [0u8; 32]; + rand_core::OsRng.fill_bytes(&mut bytes); + let secret = ed25519_dalek::SigningKey::from_bytes(&bytes); + let public = secret.verifying_key(); + let priv_key = SignaturePrivateKey::new(secret.to_bytes().to_vec()); + let pub_key = SignaturePublicKey::new(public.to_bytes().to_vec()); + let signer = Self { secret, public }; + (signer, priv_key, pub_key) + } + + pub fn public_key(&self) -> SignaturePublicKey { + SignaturePublicKey::new(self.public.to_bytes().to_vec()) + } +} + +#[cfg(feature = "ed25519-dalek")] +impl SignatureScheme for Ed25519Signer { + fn sign(&self, msg: &[u8]) -> Result, CryptoError> { + use ed25519_dalek::Signer; + let signature = self.secret.sign(msg).to_bytes().to_vec(); + Ok(signature) + } + + fn verify(&self, msg: &[u8], signature: &[u8]) -> Result<(), CryptoError> { + use ed25519_dalek::Verifier; + let sig_bytes: [u8; 64] = signature + .try_into() + .map_err(|_| CryptoError::InvalidSignature)?; + let sig = ed25519_dalek::Signature::from_bytes(&sig_bytes); + self.public + .verify(msg, &sig) + .map_err(|_| CryptoError::VerificationFailed) + } +} + +#[cfg(feature = "ed25519-dalek")] +pub fn verify_ed25519( + public_key: &SignaturePublicKey, + msg: &[u8], + signature: &[u8], +) -> Result<(), CryptoError> { + use ed25519_dalek::Verifier; + + let pub_bytes: [u8; 32] = public_key + .as_bytes() + .try_into() + .map_err(|_| CryptoError::InvalidSignature)?; + let public = ed25519_dalek::VerifyingKey::from_bytes(&pub_bytes) + .map_err(|_| CryptoError::InvalidSignature)?; + let sig_bytes: [u8; 64] = signature + .try_into() + .map_err(|_| CryptoError::InvalidSignature)?; + let sig = ed25519_dalek::Signature::from_bytes(&sig_bytes); + + public + .verify(msg, &sig) + .map_err(|_| CryptoError::VerificationFailed) +} + +#[cfg(feature = "ml-dsa")] +pub struct MlDsaSigner { + secret: ml_dsa::SigningKey, + public: ml_dsa::VerifyingKey, +} + +#[cfg(feature = "ml-dsa")] +impl MlDsaSigner { + pub fn new( + secret_key: &SignaturePqPrivateKey, + public_key: &SignaturePqPublicKey, + ) -> Result { + let encoded_sk = + ml_dsa::EncodedSigningKey::::try_from(secret_key.as_bytes()) + .map_err(|_| CryptoError::KeyGenerationFailed)?; + let secret = ml_dsa::SigningKey::::decode(&encoded_sk); + let encoded_pk = + ml_dsa::EncodedVerifyingKey::::try_from(public_key.as_bytes()) + .map_err(|_| CryptoError::KeyGenerationFailed)?; + let public = ml_dsa::VerifyingKey::::decode(&encoded_pk); + Ok(Self { secret, public }) + } + + pub fn generate() -> (Self, SignaturePqPrivateKey, SignaturePqPublicKey) { + use ml_dsa::KeyGen; + let kp = ml_dsa::MlDsa65::key_gen(&mut rand_core::OsRng); + let secret = kp.signing_key().clone(); + let public = kp.verifying_key().clone(); + let priv_key = SignaturePqPrivateKey::new(secret.encode().to_vec()); + let pub_key = SignaturePqPublicKey::new(public.encode().to_vec()); + let signer = Self { secret, public }; + (signer, priv_key, pub_key) + } + + pub fn public_key(&self) -> SignaturePqPublicKey { + SignaturePqPublicKey::new(self.public.encode().to_vec()) + } + + pub fn verifying_key(&self) -> &ml_dsa::VerifyingKey { + &self.public + } + + pub fn signing_key(&self) -> &ml_dsa::SigningKey { + &self.secret + } +} + +#[cfg(feature = "ml-dsa")] +impl SignatureScheme for MlDsaSigner { + fn sign(&self, msg: &[u8]) -> Result, CryptoError> { + use ml_dsa::signature::Signer; + let signature = self.secret.sign(msg); + Ok(signature.encode().to_vec()) + } + + fn verify(&self, msg: &[u8], signature: &[u8]) -> Result<(), CryptoError> { + use ml_dsa::signature::Verifier; + let encoded_sig = + ml_dsa::EncodedSignature::::try_from(signature) + .map_err(|_| CryptoError::InvalidSignature)?; + let sig = ml_dsa::Signature::::decode(&encoded_sig) + .ok_or(CryptoError::InvalidSignature)?; + self.public + .verify(msg, &sig) + .map_err(|_| CryptoError::VerificationFailed) + } +} + +#[cfg(feature = "ml-dsa")] +pub fn verify_ml_dsa( + public_key: &SignaturePqPublicKey, + msg: &[u8], + signature: &[u8], +) -> Result<(), CryptoError> { + use ml_dsa::signature::Verifier; + + let encoded_pk = + ml_dsa::EncodedVerifyingKey::::try_from(public_key.as_bytes()) + .map_err(|_| CryptoError::InvalidSignature)?; + let public = ml_dsa::VerifyingKey::::decode(&encoded_pk); + let encoded_sig = + ml_dsa::EncodedSignature::::try_from(signature) + .map_err(|_| CryptoError::InvalidSignature)?; + let sig = ml_dsa::Signature::::decode(&encoded_sig) + .ok_or(CryptoError::InvalidSignature)?; + + public + .verify(msg, &sig) + .map_err(|_| CryptoError::VerificationFailed) +} + +pub struct DualSignature { + pub ed25519: Vec, + pub mldsa: Vec, +} + +#[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))] +pub fn sign_dual( + ed25519_sk: &ed25519_dalek::SigningKey, + mldsa_sk: &ml_dsa::SigningKey, + message: &[u8], +) -> DualSignature { + use ed25519_dalek::Signer; + + DualSignature { + ed25519: ed25519_sk.sign(message).to_bytes().to_vec(), + mldsa: mldsa_sk.sign(message).encode().to_vec(), + } +} + +impl DualSignature { + #[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))] + pub fn verify( + &self, + ed25519_vk: &ed25519_dalek::VerifyingKey, + mldsa_vk: &ml_dsa::VerifyingKey, + message: &[u8], + ) -> Result<(), CryptoError> { + use ed25519_dalek::Verifier; + + let ed_sig = ed25519_dalek::Signature::from_slice(&self.ed25519) + .map_err(|_| CryptoError::InvalidSignature)?; + ed25519_vk + .verify(message, &ed_sig) + .map_err(|_| CryptoError::VerificationFailed)?; + + let encoded_sig = + ml_dsa::EncodedSignature::::try_from(self.mldsa.as_slice()) + .map_err(|_| CryptoError::InvalidSignature)?; + let ml_sig = ml_dsa::Signature::::decode(&encoded_sig) + .ok_or(CryptoError::InvalidSignature)?; + mldsa_vk + .verify(message, &ml_sig) + .map_err(|_| CryptoError::VerificationFailed)?; + + Ok(()) + } +} diff --git a/host/Cargo.toml b/host/Cargo.toml index b35bfb6..0d5513f 100644 --- a/host/Cargo.toml +++ b/host/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "host" +name = "mtp-host" version = "0.1.0" edition = "2024" diff --git a/registry/Cargo.toml b/registry/Cargo.toml index 20e847b..b6d6d57 100644 --- a/registry/Cargo.toml +++ b/registry/Cargo.toml @@ -4,4 +4,4 @@ version = "0.1.0" edition = "2024" [dependencies] -common = { path = "../common" } +mtp-common = { path = "../common" } diff --git a/registry/src/lib.rs b/registry/src/lib.rs index aef951f..bd007b0 100644 --- a/registry/src/lib.rs +++ b/registry/src/lib.rs @@ -1,4 +1,4 @@ -use common::RegistryError; +use mtp_common::RegistryError; use std::collections::HashMap; #[repr(transparent)] diff --git a/transport/Cargo.toml b/transport/Cargo.toml index 4ac083f..1958e5a 100644 --- a/transport/Cargo.toml +++ b/transport/Cargo.toml @@ -1,6 +1,26 @@ [package] -name = "transport" +name = "mtp-transport" version = "0.1.0" edition = "2024" [dependencies] +mtp-codec = { path = "../codec" } +mtp-common = { path = "../common" } +wtransport = { version = "0.7.1", default-features = false, features = [ + "aws-lc-rs", + "quinn", + "self-signed", +] } +rustls = { version = "0.23.40" } +quinn = { version = "0.11.9", default-features = false, features = [ + "rustls-aws-lc-rs", + "rustls", +] } +tokio = { version = "1", features = ["full"] } +thiserror = "2.0.18" +rustls-native-certs = "0.8.4" + +[features] +default = [] +# Enables hosting a MTP server +host = [] diff --git a/transport/src/client.rs b/transport/src/client.rs new file mode 100644 index 0000000..811f5d7 --- /dev/null +++ b/transport/src/client.rs @@ -0,0 +1,93 @@ +use std::sync::Arc; + +use mtp_common::CommunicationError; +use rustls::{ClientConfig as RustlsClientConfig, RootCertStore, pki_types::pem::PemObject}; +use wtransport::{ClientConfig, Endpoint}; + +use crate::{ConnectionHandle, Policy, Receiver, Sender}; + +pub async fn connect( + url: &str, + server_cert: Option>, + policy: Policy, +) -> Result<(Sender, Receiver), CommunicationError> { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + let client_config = if let Some(cert_pem) = server_cert { + configure_client_with_cert(cert_pem, &policy)? + } else { + configure_client_system_roots(&policy)? + }; + + let endpoint = Endpoint::client(client_config) + .map_err(|e| CommunicationError::Other(format!("Endpoint creation failed: {}", e)))?; + + let connection = endpoint + .connect(url) + .await + .map_err(|e| CommunicationError::ConnectingError(e.to_string()))?; + + let handle = Arc::new(ConnectionHandle::new()); + let policy = Arc::new(policy); + + let sender = Sender::new(connection.clone(), handle.clone(), policy.clone()); + let receiver = Receiver::new(connection, handle, policy); + + Ok((sender, receiver)) +} + +fn configure_client_with_cert( + server_cert: Vec, + policy: &Policy, +) -> Result { + let mut root_store = RootCertStore::empty(); + + let certs = rustls::pki_types::CertificateDer::pem_slice_iter(&server_cert) + .collect::, _>>() + .map_err(|_| CommunicationError::CertificateParseFailed)?; + + for cert in certs { + root_store + .add(cert) + .map_err(|_| CommunicationError::CertificateParseFailed)?; + } + + let mut tls_config = RustlsClientConfig::builder() + .with_root_certificates(root_store) + .with_no_client_auth(); + + tls_config.alpn_protocols = vec![b"h3".to_vec()]; + + Ok(ClientConfig::builder() + .with_bind_default() + .with_custom_tls(tls_config) + .keep_alive_interval(policy.keep_alive_interval) + .max_idle_timeout(policy.max_idle_timeout) + .map_err(|e| CommunicationError::Other(e.to_string()))? + .build()) +} + +fn configure_client_system_roots(policy: &Policy) -> Result { + let mut root_store = RootCertStore::empty(); + + // Load native certs + let certs = rustls_native_certs::load_native_certs().certs; + + for cert in certs { + root_store.add(cert).ok(); + } + + let mut tls_config = RustlsClientConfig::builder() + .with_root_certificates(root_store) + .with_no_client_auth(); + + tls_config.alpn_protocols = vec![b"h3".to_vec()]; + + Ok(ClientConfig::builder() + .with_bind_default() + .with_custom_tls(tls_config) + .keep_alive_interval(policy.keep_alive_interval) + .max_idle_timeout(policy.max_idle_timeout) + .map_err(|e| CommunicationError::Other(e.to_string()))? + .build()) +} diff --git a/transport/src/connection.rs b/transport/src/connection.rs new file mode 100644 index 0000000..3f1ac96 --- /dev/null +++ b/transport/src/connection.rs @@ -0,0 +1,519 @@ +use crate::ConnectionHandle; +use mtp_codec::CommunicationValue; +use mtp_common::CommunicationError; +use std::sync::Arc; +use tokio::sync::{Mutex, mpsc}; +use tokio::time::{Duration, sleep, timeout}; +use wtransport::Connection; + +const APPLICATION_CLOSE_REASON: &str = "mtp-close"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum SendMode { + PersistentStream, + SingleStreamPerMessage, +} + +#[derive(Debug, Clone)] +pub struct Policy { + pub send_mode: SendMode, + pub max_message_size: u64, + pub close_frame_len: u32, + pub application_close_code: u32, + pub open_stream_timeout: Duration, + pub write_timeout: Duration, + pub accept_stream_timeout: Duration, + pub read_timeout: Duration, + pub keep_alive_interval: Option, + pub max_idle_timeout: Option, + pub force_close_delay: Duration, + pub max_transient_recv_errors: usize, + pub transient_recv_backoff: Duration, + pub receiver_queue_capacity: usize, +} + +impl Default for Policy { + fn default() -> Self { + Self { + send_mode: SendMode::PersistentStream, + max_message_size: 1_000_000_000, + close_frame_len: u32::MAX, + application_close_code: 0, + open_stream_timeout: Duration::from_millis(2_000), + write_timeout: Duration::from_millis(2_000), + accept_stream_timeout: Duration::from_millis(10_000), + read_timeout: Duration::from_millis(30_000), + keep_alive_interval: Some(Duration::from_secs(3)), + max_idle_timeout: Some(Duration::from_secs(30)), + force_close_delay: Duration::from_millis(300), + max_transient_recv_errors: 20, + transient_recv_backoff: Duration::from_millis(100), + receiver_queue_capacity: 1000, + } + } +} + +#[allow(unused)] +enum ReceivedFrame { + Message(CommunicationValue), + ClosedByPeer, + Idle, +} + +pub struct Sender { + send_guard: Mutex<()>, + stream_guard: Mutex>, + handle: Arc, + connection: Connection, + policy: Arc, +} + +impl Sender { + pub fn new(connection: Connection, handle: Arc, policy: Arc) -> Self { + Self { + send_guard: Mutex::new(()), + stream_guard: Mutex::new(None), + handle, + connection, + policy, + } + } + + async fn write_frame( + stream: &mut wtransport::SendStream, + data: &CommunicationValue, + policy: &Policy, + ) -> Result<(), CommunicationError> { + let bytes = data.to_bytes(); + if bytes.len() as u64 > policy.max_message_size + || bytes.len() as u64 >= policy.close_frame_len as u64 + { + return Err(CommunicationError::MessageTooLarge); + } + + use tokio::io::AsyncWriteExt; + + timeout(policy.write_timeout, stream.write_u32(bytes.len() as u32)) + .await + .map_err(|_| CommunicationError::StreamError)? + .map_err(|_| CommunicationError::StreamError)?; + + timeout(policy.write_timeout, stream.write_all(&bytes)) + .await + .map_err(|_| CommunicationError::StreamError)? + .map_err(CommunicationError::from)?; + + Ok(()) + } + + fn normalize_send_error(error: CommunicationError) -> CommunicationError { + match error { + CommunicationError::ConnectionError(_) + | CommunicationError::ReadExactError(_) + | CommunicationError::ClosedError(_) + | CommunicationError::StreamReadExactError(_) + | CommunicationError::StreamError => CommunicationError::StreamClosed, + other => other, + } + } + + async fn open_uni_stream( + conn: &Connection, + policy: &Policy, + ) -> Result { + let opening = timeout(policy.open_stream_timeout, conn.open_uni()) + .await + .map_err(|_| CommunicationError::StreamError)? + .map_err(CommunicationError::ConnectionError)?; + + let stream = timeout(policy.open_stream_timeout, opening) + .await + .map_err(|_| CommunicationError::StreamError)? + .map_err(|_| CommunicationError::StreamError)?; + + Ok(stream) + } + + async fn ensure_stream<'a>( + conn: &Connection, + stream_opt: &'a mut Option, + policy: &Policy, + ) -> Result<&'a mut wtransport::SendStream, CommunicationError> { + if stream_opt.is_none() { + *stream_opt = Some(Self::open_uni_stream(conn, policy).await?); + } + + match stream_opt.as_mut() { + Some(stream) => Ok(stream), + _ => Err(CommunicationError::StreamError), + } + } + + async fn send_on_persistent_stream( + conn: &Connection, + stream_opt: &mut Option, + data: &CommunicationValue, + policy: &Policy, + ) -> Result<(), CommunicationError> { + let mut tries = 0usize; + loop { + if conn.quic_connection().close_reason().is_some() { + return Err(CommunicationError::StreamClosed); + } + + let res = { + let stream = Self::ensure_stream(conn, stream_opt, policy).await?; + Self::write_frame(stream, data, policy).await + }; + + if res.is_ok() { + return Ok(()); + } + + *stream_opt = None; + tries += 1; + if tries >= 4 { + let stream = Self::ensure_stream(conn, stream_opt, policy).await?; + return Self::write_frame(stream, data, policy).await; + } + + tokio::time::sleep(std::time::Duration::from_millis(20 * tries as u64)).await; + } + } + + async fn send_on_single_stream( + conn: &Connection, + data: &CommunicationValue, + policy: &Policy, + ) -> Result<(), CommunicationError> { + let mut stream = Self::open_uni_stream(conn, policy).await?; + Self::write_frame(&mut stream, data, policy).await?; + + timeout(policy.write_timeout, stream.finish()) + .await + .map_err(|_| CommunicationError::StreamError)? + .map_err(|_| CommunicationError::StreamError)?; + + Ok(()) + } + async fn send_close_frame( + conn: &Connection, + policy: &Policy, + ) -> Result<(), CommunicationError> { + let mut stream = Self::open_uni_stream(conn, policy).await?; + + use tokio::io::AsyncWriteExt; + timeout( + policy.write_timeout, + stream.write_u32(policy.close_frame_len), + ) + .await + .map_err(|_| CommunicationError::StreamError)? + .map_err(|_| CommunicationError::StreamError)?; + + if let Err(e) = timeout(policy.write_timeout, stream.finish()) + .await + .map_err(|_| CommunicationError::StreamError)? + { + println!("[Sender] close frame finish failed: {e}"); + } + + Ok(()) + } + + pub async fn send(&self, data: &CommunicationValue) -> Result<(), CommunicationError> { + if self.handle.is_closed() { + return Err(self + .handle + .close_reason() + .unwrap_or(CommunicationError::UseAfterClosed)); + } + + let _send_lock = self.send_guard.lock().await; + + if self.connection.quic_connection().close_reason().is_some() { + let reason = self + .handle + .close_reason() + .unwrap_or(CommunicationError::StreamClosed); + self.handle.close(Some(reason.clone())); + return Err(reason); + } + + let res = match self.policy.send_mode { + SendMode::PersistentStream => { + let mut stream_opt = self.stream_guard.lock().await; + let r = Self::send_on_persistent_stream( + &self.connection, + &mut stream_opt, + data, + &self.policy, + ) + .await; + if r.is_err() { + *stream_opt = None; + } + r + } + SendMode::SingleStreamPerMessage => { + Self::send_on_single_stream(&self.connection, data, &self.policy).await + } + }; + + match res { + Ok(()) => Ok(()), + Err(e) => { + let normalized = Self::normalize_send_error(e); + + if self.connection.quic_connection().close_reason().is_some() + || matches!(normalized, CommunicationError::StreamClosed) + { + self.handle.close(Some(normalized.clone())); + } + + Err(normalized) + } + } + } + + pub fn handle(&self) -> &Arc { + &self.handle + } + + pub fn close(&self) { + let connection = self.connection.clone(); + let handle = self.handle.clone(); + let policy = self.policy.clone(); + + tokio::spawn(async move { + if connection.quic_connection().close_reason().is_some() || handle.is_closed() { + handle.close(Some(CommunicationError::StreamClosed)); + return; + } + + let _ = Self::send_close_frame(&connection, &policy).await; + + handle.close(Some(CommunicationError::StreamClosed)); + + sleep(policy.force_close_delay).await; + if connection.quic_connection().close_reason().is_none() { + connection.quic_connection().close( + policy.application_close_code.into(), + APPLICATION_CLOSE_REASON.as_bytes(), + ); + } + }); + } + + pub fn is_open(&self) -> bool { + self.handle.is_open() + } + + pub fn is_closed(&self) -> bool { + self.handle.is_closed() + } + + pub fn close_reason(&self) -> Option { + self.handle.close_reason() + } +} + +pub struct Receiver { + rx: Mutex>>, + _accept_task: tokio::task::JoinHandle<()>, + handle: Arc, +} + +impl Receiver { + pub fn new(connection: Connection, handle: Arc, policy: Arc) -> Self { + let (tx, rx) = mpsc::channel::>( + policy.receiver_queue_capacity, + ); + + let conn_handle = handle.clone(); + let accept_connection = connection.clone(); + let accept_policy = policy.clone(); + + let accept_task = tokio::spawn(async move { + let mut close_rx = conn_handle.subscribe_close(); + + loop { + tokio::select! { + _ = close_rx.changed() => { + if close_rx.borrow().is_some() { + break; + } + } + + accepted = timeout( + accept_policy.accept_stream_timeout, + accept_connection.accept_uni() + ) => { + match accepted { + Ok(Ok(stream)) => { + let tx_stream = tx.clone(); + let stream_handle = conn_handle.clone(); + let stream_policy = accept_policy.clone(); + + tokio::spawn(async move { + let mut s = stream; + loop { + match Self::read_one_frame(&mut s, &stream_policy).await { + Ok(ReceivedFrame::Message(msg)) => { + if tx_stream.send(Ok(msg)).await.is_err() { + break; + } + } + Ok(ReceivedFrame::ClosedByPeer) => { + let close_error = CommunicationError::StreamClosed; + let _ = tx_stream.send(Err(close_error.clone())).await; + stream_handle.close(Some(close_error)); + break; + } + Ok(ReceivedFrame::Idle) => { + break; + } + Err(e) => { + let close_error = match e { + CommunicationError::ConnectionError(_) + | CommunicationError::ReadExactError(_) + | CommunicationError::ClosedError(_) + | CommunicationError::StreamReadExactError(_) + | CommunicationError::StreamError => CommunicationError::StreamClosed, + other => other, + }; + + let _ = tx_stream.send(Err(close_error.clone())).await; + stream_handle.close(Some(close_error)); + break; + } + } + } + }); + } + + Ok(Err(_e)) => { + // A connection error from accept_uni means the connection is permanently closed. + let close_error = CommunicationError::StreamClosed; + let _ = tx.send(Err(close_error.clone())).await; + conn_handle.close(Some(close_error)); + break; + } + + Err(_) => { + if accept_connection.quic_connection().close_reason().is_some() { + let close_error = CommunicationError::StreamClosed; + let _ = tx.send(Err(close_error.clone())).await; + conn_handle.close(Some(close_error)); + break; + } + } + } + } + } + + if conn_handle.is_closed() { + break; + } + } + }); + + Self { + rx: Mutex::new(rx), + _accept_task: accept_task, + handle, + } + } + + async fn read_one_frame( + stream: &mut wtransport::RecvStream, + policy: &Policy, + ) -> Result { + use std::io::ErrorKind; + use tokio::io::AsyncReadExt; + + let mut attempts = 0; + let len = loop { + match stream.read_u32().await { + Ok(len) => break len, + Err(e) => { + if e.kind() == ErrorKind::Interrupted && attempts < 3 { + attempts += 1; + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + continue; + } + if e.kind() == ErrorKind::UnexpectedEof { + return Ok(ReceivedFrame::Idle); + } + println!("[Receiver] read_u32 failed: {e}"); + return Err(CommunicationError::StreamError); + } + } + }; + + if len == policy.close_frame_len { + return Ok(ReceivedFrame::ClosedByPeer); + } + + let len = len as usize; + if len as u64 > policy.max_message_size { + return Err(CommunicationError::MessageTooLarge); + } + + let mut buf = vec![0u8; len]; + match timeout(policy.read_timeout, stream.read_exact(&mut buf)).await { + Ok(Ok(())) => {} + Ok(Err(e)) => match timeout(policy.read_timeout, stream.read_exact(&mut buf)).await { + Ok(Ok(())) => {} + _ => return Err(e.into()), + }, + Err(_) => { + println!("[Receiver] read_exact timed out (len={})", len); + return Err(CommunicationError::StreamError); + } + } + + let message = CommunicationValue::from_bytes(&buf) + .ok_or(CommunicationError::ParseCommunicationValue)?; + + Ok(ReceivedFrame::Message(message)) + } + + pub async fn receive(&self) -> Result { + if self.handle.is_closed() { + return Err(self + .handle + .close_reason() + .unwrap_or(CommunicationError::StreamClosed)); + } + + let mut rx = self.rx.lock().await; + match rx.recv().await { + Some(result) => result, + _ => Err(self + .handle + .close_reason() + .unwrap_or(CommunicationError::StreamClosed)), + } + } + + pub fn handle(&self) -> &Arc { + &self.handle + } + + pub fn close(&self) { + self.handle.close(None); + } + + pub fn is_open(&self) -> bool { + self.handle.is_open() + } + + pub fn is_closed(&self) -> bool { + self.handle.is_closed() + } + + pub fn close_reason(&self) -> Option { + self.handle.close_reason() + } +} diff --git a/transport/src/connection_handle.rs b/transport/src/connection_handle.rs new file mode 100644 index 0000000..4ebe200 --- /dev/null +++ b/transport/src/connection_handle.rs @@ -0,0 +1,65 @@ +use mtp_common::CommunicationError; +use std::sync::{ + Arc, + atomic::{AtomicBool, Ordering}, +}; +use tokio::sync::watch; + +#[derive(Debug)] +pub struct ConnectionHandle { + closed: AtomicBool, + close_tx: watch::Sender>, + close_rx: watch::Receiver>, +} + +impl ConnectionHandle { + pub fn new() -> Self { + let (close_tx, close_rx) = watch::channel(None); + Self { + closed: AtomicBool::new(false), + close_tx, + close_rx, + } + } + + pub fn is_open(&self) -> bool { + !self.closed.load(Ordering::SeqCst) + } + + pub fn is_closed(&self) -> bool { + self.closed.load(Ordering::SeqCst) + } + + pub fn close(&self, reason: Option) { + if !self.closed.swap(true, Ordering::SeqCst) { + let _ = self.close_tx.send(reason); + } + } + + pub fn close_reason(&self) -> Option { + self.close_rx.borrow().clone() + } + + pub fn subscribe_close(&self) -> watch::Receiver> { + self.close_rx.clone() + } + + pub fn close_with_error(&self, error: CommunicationError) { + self.close(Some(error)); + } + + pub async fn wait_closed(self: Arc) -> Option { + let mut rx = self.subscribe_close(); + if self.is_closed() { + return rx.borrow().clone(); + } + let _ = rx.changed().await.ok()?; + rx.borrow().clone() + } +} + +impl Default for ConnectionHandle { + fn default() -> Self { + Self::new() + } +} diff --git a/transport/src/host.rs b/transport/src/host.rs new file mode 100644 index 0000000..b6da6dd --- /dev/null +++ b/transport/src/host.rs @@ -0,0 +1,123 @@ +use crate::{ConnectionHandle, Policy, Receiver, Sender}; +use mtp_common::CommunicationError; +use rustls::pki_types::{PrivateKeyDer, pem::PemObject}; +use std::net::{IpAddr, SocketAddr}; +use std::sync::Arc; +use wtransport::{Connection as WTConnection, Endpoint, ServerConfig}; + +pub struct Host { + incoming: tokio::sync::mpsc::Receiver<(Sender, Receiver)>, + local_addr: std::net::SocketAddr, + _task: tokio::task::JoinHandle<()>, +} + +impl Host { + pub async fn next(&mut self) -> Option<(Sender, Receiver)> { + self.incoming.recv().await + } + + pub fn local_addr(&self) -> std::net::SocketAddr { + self.local_addr + } +} + +pub async fn host( + port: u16, + cert_pem: Vec, + key_pem: Vec, + policy: Policy, +) -> Result { + let _ = rustls::crypto::aws_lc_rs::default_provider().install_default(); + + let server_config = configure_server(port, cert_pem, key_pem, &policy).await?; + let endpoint = Endpoint::server(server_config) + .map_err(|e| CommunicationError::Other(format!("Endpoint creation failed: {}", e)))?; + + let local_addr = endpoint + .local_addr() + .map_err(|e| CommunicationError::Other(e.to_string()))?; + + let (incoming_tx, incoming_rx) = tokio::sync::mpsc::channel(16); + + let policy = Arc::new(policy); + + let task = tokio::spawn(async move { + loop { + let incoming_session = endpoint.accept().await; + + let request = match incoming_session.await { + Ok(req) => req, + Err(_) => { + continue; + } + }; + + let connection = match request.accept().await { + Ok(conn) => conn, + Err(_) => { + continue; + } + }; + + let incoming_tx = incoming_tx.clone(); + let policy = policy.clone(); + tokio::spawn(handle_connection(connection, incoming_tx, policy)); + } + }); + + Ok(Host { + incoming: incoming_rx, + local_addr, + _task: task, + }) +} + +async fn handle_connection( + connection: WTConnection, + tx: tokio::sync::mpsc::Sender<(Sender, Receiver)>, + policy: Arc, +) { + let handle = Arc::new(ConnectionHandle::new()); + + let sender = Sender::new(connection.clone(), handle.clone(), policy.clone()); + let receiver = Receiver::new(connection, handle, policy); + let _ = tx.send((sender, receiver)).await; +} + +async fn configure_server( + port: u16, + cert_pem: Vec, + key_pem: Vec, + policy: &Policy, +) -> Result { + let cert_chain = rustls::pki_types::CertificateDer::pem_slice_iter(&cert_pem) + .collect::, _>>() + .map_err(|_| CommunicationError::CertificateLoadFailed)?; + + let key = PrivateKeyDer::from_pem_slice(&key_pem) + .map_err(|_| CommunicationError::CertificateParseFailed)?; + + let mut tls_config = rustls::ServerConfig::builder() + .with_no_client_auth() + .with_single_cert(cert_chain, key) + .map_err(|_| CommunicationError::CertificateLoadFailed)?; + + 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::()?; + let bind_addr = SocketAddr::new(bind_ip, port); + + let server_config = ServerConfig::builder() + .with_bind_address(bind_addr) + .with_custom_tls(tls_config) + .keep_alive_interval(policy.keep_alive_interval) + .max_idle_timeout(policy.max_idle_timeout) + .map_err(|e| CommunicationError::Other(e.to_string()))? + .build(); + + Ok(server_config) +} diff --git a/transport/src/lib.rs b/transport/src/lib.rs index b93cf3f..f4be9ee 100644 --- a/transport/src/lib.rs +++ b/transport/src/lib.rs @@ -1,14 +1,12 @@ -pub fn add(left: u64, right: u64) -> u64 { - left + right -} +pub mod client; +pub mod connection; +pub mod connection_handle; -#[cfg(test)] -mod tests { - use super::*; +pub use client::connect; +pub use connection::{Policy, Receiver, SendMode, Sender}; +pub use connection_handle::ConnectionHandle; - #[test] - fn it_works() { - let result = add(2, 2); - assert_eq!(result, 4); - } -} +#[cfg(feature = "host")] +pub mod host; +#[cfg(feature = "host")] +pub use host::{Host, host}; diff --git a/type-map/Cargo.toml b/type-map/Cargo.toml index 66f916d..a565677 100644 --- a/type-map/Cargo.toml +++ b/type-map/Cargo.toml @@ -1,5 +1,5 @@ [package] -name = "type-map" +name = "mtp-type-map" version = "0.1.0" edition = "2024"