methanium/mtp
Watch
3
0
Fork
You've already forked mtp
 0
Code Issues Pull requests 3 Projects Releases 22 Packages Wiki Activity Actions

General & Crypto

Browse source
This commit is contained in:
Alex Emmet 2026-06-20 15:25:36 +02:00
commit 02f94993c7
27 changed files with 1881 additions and 47 deletions
Download patch file Download diff file Expand all files Collapse all files

2
client/Cargo.toml
View file

@ -1,5 +1,5 @@
[package]
name = "client"
name = "mtp-client"
version = "0.1.0"
edition = "2024"

4
codec/Cargo.toml
View file

@ -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 = "*"

4
codec/src/communication_value.rs
View file

@ -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();

18
codec/src/data_value.rs
View file

@ -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;

4
codec/src/lib.rs
View file

@ -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<Vec<u8>, CodecError> {
// write header using typemap.data_id(), serialize value

13
common/Cargo.toml
View file

@ -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",
] }

80
common/src/lib.rs
View file

@ -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),
}

17
crypto/Cargo.toml
View file

@ -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"]

113
crypto/README.md Normal file
View file

@ -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
```

171
crypto/src/aead.rs Normal file
View file

@ -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<Vec<u8>, CryptoError>;
}
pub trait AeadDecrypt {
fn decrypt(&self, ciphertext: &[u8], aad: &[u8]) -> Result<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, 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<Vec<u8>, CryptoError> {
use aes_gcm::aead::{Aead, KeyInit, Payload};
use aes_gcm::Aes256Gcm as AesGcmInner;
use aes_gcm::Nonce;
let key = aes_gcm::Key::<AesGcmInner>::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<Vec<u8>, 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::<AesGcmInner>::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
}
}

38
crypto/src/error.rs Normal file
View file

@ -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 {}

29
crypto/src/hash.rs Normal file
View file

@ -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()
}
}

34
crypto/src/kdf.rs Normal file
View file

@ -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<Vec<u8>, CryptoError> {
let hk = Hkdf::<Sha256>::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::<Sha256>::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)
}

45
crypto/src/kem.rs Normal file
View file

@ -0,0 +1,45 @@
use crate::error::CryptoError;
use crate::keypair::{KemPrivateKey, KemPublicKey};
pub struct Encapsulated {
pub ciphertext: Vec<u8>,
pub shared_secret: Vec<u8>,
}
#[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<Encapsulated, CryptoError> {
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<Vec<u8>, 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())
}
}

212
crypto/src/keypair.rs Normal file
View file

@ -0,0 +1,212 @@
use zeroize::{Zeroize, ZeroizeOnDrop};
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct EncryptionPrivateKey(Vec<u8>);
impl EncryptionPrivateKey {
pub fn new(bytes: Vec<u8>) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
}
impl From<Vec<u8>> for EncryptionPrivateKey {
fn from(bytes: Vec<u8>) -> Self {
Self(bytes)
}
}
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct SignaturePrivateKey(Vec<u8>);
impl SignaturePrivateKey {
pub fn new(bytes: Vec<u8>) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
}
impl From<Vec<u8>> for SignaturePrivateKey {
fn from(bytes: Vec<u8>) -> Self {
Self(bytes)
}
}
#[derive(Clone)]
pub struct EncryptionPublicKey(Vec<u8>);
impl EncryptionPublicKey {
pub fn new(bytes: Vec<u8>) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
}
impl From<Vec<u8>> for EncryptionPublicKey {
fn from(bytes: Vec<u8>) -> Self {
Self(bytes)
}
}
#[derive(Clone)]
pub struct SignaturePublicKey(Vec<u8>);
impl SignaturePublicKey {
pub fn new(bytes: Vec<u8>) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
}
impl From<Vec<u8>> for SignaturePublicKey {
fn from(bytes: Vec<u8>) -> 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<u8>);
impl KemPrivateKey {
pub fn new(bytes: Vec<u8>) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
}
impl From<Vec<u8>> for KemPrivateKey {
fn from(bytes: Vec<u8>) -> Self {
Self(bytes)
}
}
#[derive(Clone)]
pub struct KemPublicKey(Vec<u8>);
impl KemPublicKey {
pub fn new(bytes: Vec<u8>) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
}
impl From<Vec<u8>> for KemPublicKey {
fn from(bytes: Vec<u8>) -> Self {
Self(bytes)
}
}
#[derive(Clone)]
pub struct SignaturePqPublicKey(Vec<u8>);
impl SignaturePqPublicKey {
pub fn new(bytes: Vec<u8>) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
}
impl From<Vec<u8>> for SignaturePqPublicKey {
fn from(bytes: Vec<u8>) -> Self {
Self(bytes)
}
}
#[derive(Zeroize, ZeroizeOnDrop)]
pub struct SignaturePqPrivateKey(Vec<u8>);
impl SignaturePqPrivateKey {
pub fn new(bytes: Vec<u8>) -> Self {
Self(bytes)
}
pub fn as_bytes(&self) -> &[u8] {
&self.0
}
}
impl From<Vec<u8>> for SignaturePqPrivateKey {
fn from(bytes: Vec<u8>) -> 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,
}
}
}

56
crypto/src/lib.rs
View file

@ -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;

232
crypto/src/sign.rs Normal file
View file

@ -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<Vec<u8>, 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<Self, CryptoError> {
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<Vec<u8>, 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<ml_dsa::MlDsa65>,
public: ml_dsa::VerifyingKey<ml_dsa::MlDsa65>,
}
#[cfg(feature = "ml-dsa")]
impl MlDsaSigner {
pub fn new(
secret_key: &SignaturePqPrivateKey,
public_key: &SignaturePqPublicKey,
) -> Result<Self, CryptoError> {
let encoded_sk =
ml_dsa::EncodedSigningKey::<ml_dsa::MlDsa65>::try_from(secret_key.as_bytes())
.map_err(|_| CryptoError::KeyGenerationFailed)?;
let secret = ml_dsa::SigningKey::<ml_dsa::MlDsa65>::decode(&encoded_sk);
let encoded_pk =
ml_dsa::EncodedVerifyingKey::<ml_dsa::MlDsa65>::try_from(public_key.as_bytes())
.map_err(|_| CryptoError::KeyGenerationFailed)?;
let public = ml_dsa::VerifyingKey::<ml_dsa::MlDsa65>::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<ml_dsa::MlDsa65> {
&self.public
}
pub fn signing_key(&self) -> &ml_dsa::SigningKey<ml_dsa::MlDsa65> {
&self.secret
}
}
#[cfg(feature = "ml-dsa")]
impl SignatureScheme for MlDsaSigner {
fn sign(&self, msg: &[u8]) -> Result<Vec<u8>, 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::<ml_dsa::MlDsa65>::try_from(signature)
.map_err(|_| CryptoError::InvalidSignature)?;
let sig = ml_dsa::Signature::<ml_dsa::MlDsa65>::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::<ml_dsa::MlDsa65>::try_from(public_key.as_bytes())
.map_err(|_| CryptoError::InvalidSignature)?;
let public = ml_dsa::VerifyingKey::<ml_dsa::MlDsa65>::decode(&encoded_pk);
let encoded_sig =
ml_dsa::EncodedSignature::<ml_dsa::MlDsa65>::try_from(signature)
.map_err(|_| CryptoError::InvalidSignature)?;
let sig = ml_dsa::Signature::<ml_dsa::MlDsa65>::decode(&encoded_sig)
.ok_or(CryptoError::InvalidSignature)?;
public
.verify(msg, &sig)
.map_err(|_| CryptoError::VerificationFailed)
}
pub struct DualSignature {
pub ed25519: Vec<u8>,
pub mldsa: Vec<u8>,
}
#[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))]
pub fn sign_dual(
ed25519_sk: &ed25519_dalek::SigningKey,
mldsa_sk: &ml_dsa::SigningKey<ml_dsa::MlDsa65>,
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<ml_dsa::MlDsa65>,
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::<ml_dsa::MlDsa65>::try_from(self.mldsa.as_slice())
.map_err(|_| CryptoError::InvalidSignature)?;
let ml_sig = ml_dsa::Signature::<ml_dsa::MlDsa65>::decode(&encoded_sig)
.ok_or(CryptoError::InvalidSignature)?;
mldsa_vk
.verify(message, &ml_sig)
.map_err(|_| CryptoError::VerificationFailed)?;
Ok(())
}
}

2
host/Cargo.toml
View file

@ -1,5 +1,5 @@
[package]
name = "host"
name = "mtp-host"
version = "0.1.0"
edition = "2024"

2
registry/Cargo.toml
View file

@ -4,4 +4,4 @@ version = "0.1.0"
edition = "2024"
[dependencies]
common = { path = "../common" }
mtp-common = { path = "../common" }

2
registry/src/lib.rs
View file

@ -1,4 +1,4 @@
use common::RegistryError;
use mtp_common::RegistryError;
use std::collections::HashMap;
#[repr(transparent)]

22
transport/Cargo.toml
View file

@ -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 = []

93
transport/src/client.rs Normal file
View file

@ -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<Vec<u8>>,
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<u8>,
policy: &Policy,
) -> Result<ClientConfig, CommunicationError> {
let mut root_store = RootCertStore::empty();
let certs = rustls::pki_types::CertificateDer::pem_slice_iter(&server_cert)
.collect::<Result<Vec<_>, _>>()
.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<ClientConfig, CommunicationError> {
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())
}

519
transport/src/connection.rs Normal file
View file

@ -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<Duration>,
pub max_idle_timeout: Option<Duration>,
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<Option<wtransport::SendStream>>,
handle: Arc<ConnectionHandle>,
connection: Connection,
policy: Arc<Policy>,
}
impl Sender {
pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> 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<wtransport::SendStream, CommunicationError> {
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<wtransport::SendStream>,
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<wtransport::SendStream>,
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<ConnectionHandle> {
&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<CommunicationError> {
self.handle.close_reason()
}
}
pub struct Receiver {
rx: Mutex<mpsc::Receiver<Result<CommunicationValue, CommunicationError>>>,
_accept_task: tokio::task::JoinHandle<()>,
handle: Arc<ConnectionHandle>,
}
impl Receiver {
pub fn new(connection: Connection, handle: Arc<ConnectionHandle>, policy: Arc<Policy>) -> Self {
let (tx, rx) = mpsc::channel::<Result<CommunicationValue, CommunicationError>>(
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<ReceivedFrame, CommunicationError> {
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<CommunicationValue, CommunicationError> {
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<ConnectionHandle> {
&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<CommunicationError> {
self.handle.close_reason()
}
}

65
transport/src/connection_handle.rs Normal file
View file

@ -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<Option<CommunicationError>>,
close_rx: watch::Receiver<Option<CommunicationError>>,
}
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<CommunicationError>) {
if !self.closed.swap(true, Ordering::SeqCst) {
let _ = self.close_tx.send(reason);
}
}
pub fn close_reason(&self) -> Option<CommunicationError> {
self.close_rx.borrow().clone()
}
pub fn subscribe_close(&self) -> watch::Receiver<Option<CommunicationError>> {
self.close_rx.clone()
}
pub fn close_with_error(&self, error: CommunicationError) {
self.close(Some(error));
}
pub async fn wait_closed(self: Arc<Self>) -> Option<CommunicationError> {
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()
}
}

123
transport/src/host.rs Normal file
View file

@ -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<u8>,
key_pem: Vec<u8>,
policy: Policy,
) -> Result<Host, CommunicationError> {
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<Policy>,
) {
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<u8>,
key_pem: Vec<u8>,
policy: &Policy,
) -> Result<ServerConfig, CommunicationError> {
let cert_chain = rustls::pki_types::CertificateDer::pem_slice_iter(&cert_pem)
.collect::<Result<Vec<_>, _>>()
.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::<IpAddr>()?;
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)
}

22
transport/src/lib.rs
View file

@ -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};

2
type-map/Cargo.toml
View file

@ -1,5 +1,5 @@
[package]
name = "type-map"
name = "mtp-type-map"
version = "0.1.0"
edition = "2024"