[Clean] safer unwrap & except handling
Some checks failed
CI / checks (push) Failing after 1m51s

This commit is contained in:
Alex Emmet 2026-07-15 19:11:01 +02:00
commit 5f11d476b6
17 changed files with 475 additions and 348 deletions

View file

@ -3,18 +3,17 @@ use std::net::{IpAddr, Ipv4Addr};
use mtp_client::{ClientConfig, MTPClient}; use mtp_client::{ClientConfig, MTPClient};
use mtp_host::{HostConfig, MTPHost}; use mtp_host::{HostConfig, MTPHost};
fn generate_self_signed_cert() -> (Vec<u8>, Vec<u8>) { async fn generate_self_signed_cert() -> Result<(Vec<u8>, Vec<u8>), Box<dyn std::error::Error>> {
let key_pair = rcgen::KeyPair::generate().unwrap(); let key_pair = rcgen::KeyPair::generate()?;
let params = let params = rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()])?;
rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]).unwrap(); let cert = params.self_signed(&key_pair)?;
let cert = params.self_signed(&key_pair).unwrap();
let cert_pem = cert.pem(); let cert_pem = cert.pem();
let key_pem = key_pair.serialize_pem(); let key_pem = key_pair.serialize_pem();
(cert_pem.into_bytes(), key_pem.into_bytes()) Ok((cert_pem.into_bytes(), key_pem.into_bytes()))
} }
async fn start_host(send_pongs: bool) -> (MTPHost, Vec<u8>) { async fn start_host(send_pongs: bool) -> Result<(MTPHost, Vec<u8>), Box<dyn std::error::Error>> {
let (cert_pem, key_pem) = generate_self_signed_cert(); let (cert_pem, key_pem) = generate_self_signed_cert().await?;
let host = MTPHost::new( let host = MTPHost::new(
HostConfig::new( HostConfig::new(
IpAddr::V4(Ipv4Addr::LOCALHOST), IpAddr::V4(Ipv4Addr::LOCALHOST),
@ -24,14 +23,13 @@ async fn start_host(send_pongs: bool) -> (MTPHost, Vec<u8>) {
) )
.with_pongs(send_pongs), .with_pongs(send_pongs),
) )
.await .await?;
.unwrap(); Ok((host, cert_pem))
(host, cert_pem)
} }
#[tokio::test] #[tokio::test]
async fn test_ping_rtt_and_missed_ping_teardown() { async fn test_ping_rtt_and_missed_ping_teardown() -> Result<(), Box<dyn std::error::Error>> {
let (mut host, cert_pem) = start_host(true).await; let (mut host, cert_pem) = start_host(true).await?;
let url = format!("https://127.0.0.1:{}", host.local_addr().port()); let url = format!("https://127.0.0.1:{}", host.local_addr().port());
let client = MTPClient::connect( let client = MTPClient::connect(
@ -40,10 +38,9 @@ async fn test_ping_rtt_and_missed_ping_teardown() {
.with_ping_interval(std::time::Duration::from_millis(25)) .with_ping_interval(std::time::Duration::from_millis(25))
.with_max_missed_pings(3), .with_max_missed_pings(3),
) )
.await .await?;
.unwrap();
let _accepted = host.accept().await.unwrap().unwrap(); let _accepted = host.accept().await?;
let ping = tokio::time::timeout(std::time::Duration::from_secs(5), async { let ping = tokio::time::timeout(std::time::Duration::from_secs(5), async {
loop { loop {
@ -53,12 +50,11 @@ async fn test_ping_rtt_and_missed_ping_teardown() {
tokio::time::sleep(std::time::Duration::from_millis(10)).await; tokio::time::sleep(std::time::Duration::from_millis(10)).await;
} }
}) })
.await .await?;
.unwrap();
assert!(ping > std::time::Duration::ZERO); assert!(ping > std::time::Duration::ZERO);
let (mut silent_host, silent_cert_pem) = start_host(false).await; let (mut silent_host, silent_cert_pem) = start_host(false).await?;
let silent_url = format!("https://127.0.0.1:{}", silent_host.local_addr().port()); let silent_url = format!("https://127.0.0.1:{}", silent_host.local_addr().port());
let silent_client = MTPClient::connect( let silent_client = MTPClient::connect(
ClientConfig::new(silent_url) ClientConfig::new(silent_url)
@ -66,10 +62,9 @@ async fn test_ping_rtt_and_missed_ping_teardown() {
.with_ping_interval(std::time::Duration::from_millis(25)) .with_ping_interval(std::time::Duration::from_millis(25))
.with_max_missed_pings(2), .with_max_missed_pings(2),
) )
.await .await?;
.unwrap();
let _accepted = silent_host.accept().await.unwrap().unwrap(); let _accepted = silent_host.accept().await?;
let closed = tokio::time::timeout(std::time::Duration::from_secs(5), async { let closed = tokio::time::timeout(std::time::Duration::from_secs(5), async {
loop { loop {
@ -82,4 +77,5 @@ async fn test_ping_rtt_and_missed_ping_teardown() {
.await; .await;
assert!(closed.is_ok(), "client should close after missed pings"); assert!(closed.is_ok(), "client should close after missed pings");
Ok(())
} }

View file

@ -775,65 +775,67 @@ mod tests {
use super::*; use super::*;
use crate::data_value::DataValue; use crate::data_value::DataValue;
fn roundtrip(cv: CommunicationValue) -> CommunicationValue { fn roundtrip(cv: CommunicationValue) -> Result<CommunicationValue, Box<dyn std::error::Error>> {
let bytes = cv.to_bytes().expect("encode failed"); let bytes = cv.to_bytes()?;
let decoded = CommunicationValue::from_bytes(&bytes).expect("failed to deserialize"); let decoded = CommunicationValue::from_bytes(&bytes)?;
let bytes2 = decoded.to_bytes().expect("encode failed"); let bytes2 = decoded.to_bytes()?;
assert_eq!(bytes, bytes2); assert_eq!(bytes, bytes2);
decoded Ok(decoded)
} }
#[test] #[test]
fn test_flags_and_order_without_optional() { fn test_flags_and_order_without_optional() -> Result<(), Box<dyn std::error::Error>> {
let cv = CommunicationValue::new(CommunicationType::ErrorParsing).with_id(0); let cv = CommunicationValue::new(CommunicationType::ErrorParsing).with_id(0);
let bytes = cv.to_bytes().expect("encode failed"); let bytes = cv.to_bytes()?;
// [u32 len][u16 type][flags]... // [u32 len][u16 type][flags]...
assert!(bytes.len() >= 7); assert!(bytes.len() >= 7);
let mut c = Cursor::new(bytes.as_slice()); let mut c = Cursor::new(bytes.as_slice());
let total_len = c.read_u32::<BigEndian>().expect("read len"); let total_len = c.read_u32::<BigEndian>()?;
assert_eq!(total_len as usize + 4, bytes.len()); assert_eq!(total_len as usize + 4, bytes.len());
let typ = c.read_u16::<BigEndian>().expect("read type"); let typ = c.read_u16::<BigEndian>()?;
assert_eq!(typ, 12); assert_eq!(typ, 12);
let flags = c.read_u8().expect("read flags"); let flags = c.read_u8()?;
assert_eq!(flags & 0b0000_0111, 0); assert_eq!(flags & 0b0000_0111, 0);
Ok(())
} }
#[test] #[test]
fn test_flags_and_order_with_all_optional() { fn test_flags_and_order_with_all_optional() -> Result<(), Box<dyn std::error::Error>> {
let cv = CommunicationValue::new(CommunicationType::ErrorBadVersion) let cv = CommunicationValue::new(CommunicationType::ErrorBadVersion)
.with_id(0xAABBCCDD) .with_id(0xAABBCCDD)
.with_sender(0x0000_1122_3344_5566) .with_sender(0x0000_1122_3344_5566)
.with_receiver(0x0000_6677_8899_AABB); .with_receiver(0x0000_6677_8899_AABB);
let bytes = cv.to_bytes().expect("encode failed"); let bytes = cv.to_bytes()?;
let mut c = Cursor::new(bytes.as_slice()); let mut c = Cursor::new(bytes.as_slice());
let total_len = c.read_u32::<BigEndian>().expect("len"); let total_len = c.read_u32::<BigEndian>()?;
assert_eq!(total_len as usize + 4, bytes.len()); assert_eq!(total_len as usize + 4, bytes.len());
let typ = c.read_u16::<BigEndian>().expect("read type"); let typ = c.read_u16::<BigEndian>()?;
assert_eq!(typ, 13); assert_eq!(typ, 13);
let flags = c.read_u8().expect("read flags"); let flags = c.read_u8()?;
assert_eq!(flags & 0b0000_0111, 0b0000_0111); assert_eq!(flags & 0b0000_0111, 0b0000_0111);
let id = c.read_u32::<BigEndian>().expect("id"); let id = c.read_u32::<BigEndian>()?;
assert_eq!(id, 0xAABBCCDD); assert_eq!(id, 0xAABBCCDD);
let mut sender6 = [0u8; 6]; let mut sender6 = [0u8; 6];
c.read_exact(&mut sender6).expect("sender"); c.read_exact(&mut sender6)?;
assert_eq!(sender6, [0x11, 0x22, 0x33, 0x44, 0x55, 0x66]); assert_eq!(sender6, [0x11, 0x22, 0x33, 0x44, 0x55, 0x66]);
let mut receiver6 = [0u8; 6]; let mut receiver6 = [0u8; 6];
c.read_exact(&mut receiver6).expect("receiver"); c.read_exact(&mut receiver6)?;
assert_eq!(receiver6, [0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB]); assert_eq!(receiver6, [0x66, 0x77, 0x88, 0x99, 0xAA, 0xBB]);
Ok(())
} }
#[test] #[test]
fn test_roundtrip_complex() { fn test_roundtrip_complex() -> Result<(), Box<dyn std::error::Error>> {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
let cv = CommunicationValue::new(CommunicationType::Disconnect) let cv = CommunicationValue::new(CommunicationType::Disconnect)
.with_id(1234) .with_id(1234)
@ -847,7 +849,7 @@ mod tests {
DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]), DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]),
); );
let decoded = roundtrip(cv.clone()); let decoded = roundtrip(cv.clone())?;
assert_eq!(decoded.get_id(), 1234); assert_eq!(decoded.get_id(), 1234);
assert_eq!(decoded.get_sender(), 111); assert_eq!(decoded.get_sender(), 111);
@ -861,6 +863,7 @@ mod tests {
decoded.get_data(DataType::ClientNonce), decoded.get_data(DataType::ClientNonce),
&DataValue::SignedNumber(42) &DataValue::SignedNumber(42)
); );
Ok(())
} }
#[test] #[test]
@ -873,7 +876,7 @@ mod tests {
#[cfg(feature = "crypto")] #[cfg(feature = "crypto")]
#[test] #[test]
fn test_sign_verify_frame_roundtrip() { fn test_sign_verify_frame_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
use mtp_crypto::{Ed25519Signer, SigAlgorithm}; use mtp_crypto::{Ed25519Signer, SigAlgorithm};
let (signer, sk, _pk) = Ed25519Signer::generate(); let (signer, sk, _pk) = Ed25519Signer::generate();
@ -887,18 +890,20 @@ mod tests {
assert!(cv.sign_frame(SigAlgorithm::ED25519, &signer).is_some()); assert!(cv.sign_frame(SigAlgorithm::ED25519, &signer).is_some());
// Same in-memory value verifies (FLAG_SIGNED forced on both sides). // Same in-memory value verifies (FLAG_SIGNED forced on both sides).
let verifier = Ed25519Signer::new(&sk).unwrap(); let verifier = Ed25519Signer::new(&sk)?;
assert!(cv.verify_frame(&verifier).is_ok()); assert!(cv.verify_frame(&verifier).is_ok());
// Survives a wire round-trip. // Survives a wire round-trip.
let bytes = cv.to_bytes().expect("encode failed"); let bytes = cv.to_bytes()?;
let decoded = CommunicationValue::from_bytes(&bytes).expect("decode failed"); let decoded = CommunicationValue::from_bytes(&bytes)?;
assert!(decoded.verify_frame(&verifier).is_ok()); assert!(decoded.verify_frame(&verifier).is_ok());
Ok(())
} }
#[cfg(feature = "crypto")] #[cfg(feature = "crypto")]
#[test] #[test]
fn test_verify_frame_wrong_key_fails() { fn test_verify_frame_wrong_key_fails() -> Result<(), Box<dyn std::error::Error>> {
use mtp_crypto::{Ed25519Signer, SigAlgorithm}; use mtp_crypto::{Ed25519Signer, SigAlgorithm};
let (signer, _, _) = Ed25519Signer::generate(); let (signer, _, _) = Ed25519Signer::generate();
@ -908,7 +913,9 @@ mod tests {
.add_typed_default(DataType::PqSignature, DataValue::UnsignedNumber(42)); .add_typed_default(DataType::PqSignature, DataValue::UnsignedNumber(42));
assert!(cv.sign_frame(SigAlgorithm::ED25519, &signer).is_some()); assert!(cv.sign_frame(SigAlgorithm::ED25519, &signer).is_some());
let wrong = Ed25519Signer::new(&other_sk).unwrap(); let wrong = Ed25519Signer::new(&other_sk)?;
assert!(cv.verify_frame(&wrong).is_err()); assert!(cv.verify_frame(&wrong).is_err());
Ok(())
} }
} }

View file

@ -1200,27 +1200,30 @@ impl TryFrom<DataValue> for Vec<u8> {
mod tests { mod tests {
use super::*; use super::*;
fn container_roundtrip(values: Vec<(DataTypeId, DataValue)>) { fn container_roundtrip(values: Vec<(DataTypeId, DataValue)>) -> Result<(), Box<dyn std::error::Error>> {
let dv = DataValue::Container(values.clone()); let dv = DataValue::Container(values.clone());
let bytes = dv.to_bytes().expect("encode failed"); let bytes = dv.to_bytes()?;
let decoded = DataValue::from_bytes(&bytes).expect("roundtrip failed"); let decoded = DataValue::from_bytes(&bytes).ok_or("roundtrip failed")?;
assert_eq!(dv, decoded, "container roundtrip mismatch"); assert_eq!(dv, decoded, "container roundtrip mismatch");
Ok(())
} }
fn array_roundtrip(values: Vec<DataValue>) { fn array_roundtrip(values: Vec<DataValue>) -> Result<(), Box<dyn std::error::Error>> {
let dv = DataValue::Array(values.clone()); let dv = DataValue::Array(values.clone());
let bytes = dv.to_bytes().expect("encode failed"); let bytes = dv.to_bytes()?;
let decoded = DataValue::from_bytes(&bytes).expect("roundtrip failed"); let decoded = DataValue::from_bytes(&bytes).ok_or("roundtrip failed")?;
assert_eq!(dv, decoded, "array roundtrip mismatch"); assert_eq!(dv, decoded, "array roundtrip mismatch");
Ok(())
} }
#[test] #[test]
fn test_bool_in_container() { fn test_bool_in_container() -> Result<(), Box<dyn std::error::Error>> {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
container_roundtrip(vec![ container_roundtrip(vec![
(DataType::Id.to_id(&tm), DataValue::BoolTrue), (DataType::Id.to_id(&tm), DataValue::BoolTrue),
(DataType::ClientNonce.to_id(&tm), DataValue::BoolFalse), (DataType::ClientNonce.to_id(&tm), DataValue::BoolFalse),
]); ])?;
Ok(())
} }
#[test] #[test]
@ -1239,7 +1242,7 @@ mod tests {
} }
#[test] #[test]
fn test_signed_number_in_container() { fn test_signed_number_in_container() -> Result<(), Box<dyn std::error::Error>> {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
container_roundtrip(vec![ container_roundtrip(vec![
(DataType::Version.to_id(&tm), DataValue::SignedNumber(0)), (DataType::Version.to_id(&tm), DataValue::SignedNumber(0)),
@ -1256,11 +1259,12 @@ mod tests {
DataType::PublicKeys.to_id(&tm), DataType::PublicKeys.to_id(&tm),
DataValue::SignedNumber(i128::MIN), DataValue::SignedNumber(i128::MIN),
), ),
]); ])?;
Ok(())
} }
#[test] #[test]
fn test_unsigned_number_in_container() { fn test_unsigned_number_in_container() -> Result<(), Box<dyn std::error::Error>> {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
container_roundtrip(vec![ container_roundtrip(vec![
(DataType::Version.to_id(&tm), DataValue::UnsignedNumber(0)), (DataType::Version.to_id(&tm), DataValue::UnsignedNumber(0)),
@ -1269,11 +1273,12 @@ mod tests {
DataType::ClientNonce.to_id(&tm), DataType::ClientNonce.to_id(&tm),
DataValue::UnsignedNumber(u128::MAX), DataValue::UnsignedNumber(u128::MAX),
), ),
]); ])?;
Ok(())
} }
#[test] #[test]
fn test_float_in_container() { fn test_float_in_container() -> Result<(), Box<dyn std::error::Error>> {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
container_roundtrip(vec![ container_roundtrip(vec![
(DataType::Version.to_id(&tm), DataValue::Float(0, 0)), (DataType::Version.to_id(&tm), DataValue::Float(0, 0)),
@ -1282,11 +1287,12 @@ mod tests {
DataType::ClientNonce.to_id(&tm), DataType::ClientNonce.to_id(&tm),
DataValue::Float(255, 4294967295), DataValue::Float(255, 4294967295),
), ),
]); ])?;
Ok(())
} }
#[test] #[test]
fn test_str_in_container() { fn test_str_in_container() -> Result<(), Box<dyn std::error::Error>> {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
container_roundtrip(vec![ container_roundtrip(vec![
(DataType::Version.to_id(&tm), DataValue::Str(String::new())), (DataType::Version.to_id(&tm), DataValue::Str(String::new())),
@ -1295,11 +1301,12 @@ mod tests {
DataType::ClientNonce.to_id(&tm), DataType::ClientNonce.to_id(&tm),
DataValue::Str("a".repeat(1000)), DataValue::Str("a".repeat(1000)),
), ),
]); ])?;
Ok(())
} }
#[test] #[test]
fn test_bytes_in_container() { fn test_bytes_in_container() -> Result<(), Box<dyn std::error::Error>> {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
container_roundtrip(vec![ container_roundtrip(vec![
(DataType::Version.to_id(&tm), DataValue::Bytes(vec![])), (DataType::Version.to_id(&tm), DataValue::Bytes(vec![])),
@ -1311,40 +1318,45 @@ mod tests {
DataType::ClientNonce.to_id(&tm), DataType::ClientNonce.to_id(&tm),
DataValue::Bytes(vec![0x42; 100]), DataValue::Bytes(vec![0x42; 100]),
), ),
]); ])?;
Ok(())
} }
#[test] #[test]
fn test_null_in_container() { fn test_null_in_container() -> Result<(), Box<dyn std::error::Error>> {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
container_roundtrip(vec![(DataType::Version.to_id(&tm), DataValue::Null)]); container_roundtrip(vec![(DataType::Version.to_id(&tm), DataValue::Null)])?;
Ok(())
} }
#[test] #[test]
fn test_array_non_empty_roundtrip() { fn test_array_non_empty_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
array_roundtrip(vec![ array_roundtrip(vec![
DataValue::BoolTrue, DataValue::BoolTrue,
DataValue::SignedNumber(42), DataValue::SignedNumber(42),
DataValue::Str("hello".to_string()), DataValue::Str("hello".to_string()),
DataValue::Null, DataValue::Null,
]); ])?;
Ok(())
} }
#[test] #[test]
fn test_array_nested_roundtrip() { fn test_array_nested_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
array_roundtrip(vec![ array_roundtrip(vec![
DataValue::Array(vec![DataValue::BoolTrue, DataValue::BoolFalse]), DataValue::Array(vec![DataValue::BoolTrue, DataValue::BoolFalse]),
DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]), DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]),
]); ])?;
Ok(())
} }
#[test] #[test]
fn test_container_empty_roundtrip() { fn test_container_empty_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
container_roundtrip(vec![]); container_roundtrip(vec![])?;
Ok(())
} }
#[test] #[test]
fn test_container_mixed_roundtrip() { fn test_container_mixed_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
container_roundtrip(vec![ container_roundtrip(vec![
(DataType::Version.to_id(&tm), DataValue::BoolTrue), (DataType::Version.to_id(&tm), DataValue::BoolTrue),
@ -1358,11 +1370,12 @@ mod tests {
DataValue::UnsignedNumber(u128::MAX), DataValue::UnsignedNumber(u128::MAX),
), ),
(DataType::PublicKeys.to_id(&tm), DataValue::Null), (DataType::PublicKeys.to_id(&tm), DataValue::Null),
]); ])?;
Ok(())
} }
#[test] #[test]
fn test_container_nested_roundtrip() { fn test_container_nested_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
container_roundtrip(vec![ container_roundtrip(vec![
( (
@ -1373,7 +1386,8 @@ mod tests {
DataType::Id.to_id(&tm), DataType::Id.to_id(&tm),
DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]), DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]),
), ),
]); ])?;
Ok(())
} }
#[test] #[test]
@ -1592,24 +1606,16 @@ mod tests {
} }
#[test] #[test]
fn test_try_from_ok() { fn test_try_from_ok() -> Result<(), Box<dyn std::error::Error>> {
assert!(bool::try_from(DataValue::BoolTrue).unwrap()); assert!(bool::try_from(DataValue::BoolTrue)?);
assert!(!bool::try_from(DataValue::BoolFalse).unwrap()); assert!(!bool::try_from(DataValue::BoolFalse)?);
assert_eq!( assert_eq!(String::try_from(DataValue::Str("hi".to_string()))?, "hi");
String::try_from(DataValue::Str("hi".to_string())).unwrap(), assert_eq!(i128::try_from(DataValue::SignedNumber(-1))?, -1i128);
"hi" assert_eq!(i64::try_from(DataValue::SignedNumber(10))?, 10i64);
); assert_eq!(u128::try_from(DataValue::UnsignedNumber(99))?, 99u128);
assert_eq!(i128::try_from(DataValue::SignedNumber(-1)).unwrap(), -1i128); assert_eq!(u64::try_from(DataValue::UnsignedNumber(7))?, 7u64);
assert_eq!(i64::try_from(DataValue::SignedNumber(10)).unwrap(), 10i64); assert_eq!(Vec::<u8>::try_from(DataValue::Bytes(vec![0xAB]))?, vec![0xABu8]);
assert_eq!( Ok(())
u128::try_from(DataValue::UnsignedNumber(99)).unwrap(),
99u128
);
assert_eq!(u64::try_from(DataValue::UnsignedNumber(7)).unwrap(), 7u64);
assert_eq!(
Vec::<u8>::try_from(DataValue::Bytes(vec![0xAB])).unwrap(),
vec![0xABu8]
);
} }
#[test] #[test]
@ -1632,7 +1638,7 @@ mod tests {
#[cfg(feature = "crypto")] #[cfg(feature = "crypto")]
#[test] #[test]
fn test_encrypt_decrypt_container_roundtrip() { fn test_encrypt_decrypt_container_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
use mtp_crypto::{EncryptionType, Keyring}; use mtp_crypto::{EncryptionType, Keyring};
let tm = TypeMap::latest(); let tm = TypeMap::latest();
let keyring = Keyring::generate(); let keyring = Keyring::generate();
@ -1655,8 +1661,9 @@ mod tests {
assert!(dv.decrypt_into_container(&keyring, b"aad").is_some()); assert!(dv.decrypt_into_container(&keyring, b"aad").is_some());
assert!(matches!(dv, DataValue::Container(_))); assert!(matches!(dv, DataValue::Container(_)));
let entries = dv.as_container().unwrap(); let entries = dv.as_container().ok_or("expected container")?;
assert_eq!(entries.len(), 2); assert_eq!(entries.len(), 2);
Ok(())
} }
#[cfg(feature = "crypto")] #[cfg(feature = "crypto")]
@ -1725,7 +1732,7 @@ mod tests {
#[cfg(feature = "crypto")] #[cfg(feature = "crypto")]
#[test] #[test]
fn test_sign_verify_container_roundtrip() { fn test_sign_verify_container_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
use mtp_crypto::{Ed25519Signer, EncryptionType, Keyring, SigAlgorithm}; use mtp_crypto::{Ed25519Signer, EncryptionType, Keyring, SigAlgorithm};
let tm = TypeMap::latest(); let tm = TypeMap::latest();
@ -1755,23 +1762,24 @@ mod tests {
); );
assert!(matches!(dv, DataValue::SignedContainer(_))); assert!(matches!(dv, DataValue::SignedContainer(_)));
let verifier = Ed25519Signer::new(&sk).unwrap(); let verifier = Ed25519Signer::new(&sk)?;
assert!(dv.verify_into_container(&verifier).is_some()); assert!(dv.verify_into_container(&verifier).is_some());
assert!(matches!(dv, DataValue::Container(_))); assert!(matches!(dv, DataValue::Container(_)));
let entries = dv.as_container().unwrap(); let entries = dv.as_container().ok_or("expected container")?;
assert_eq!(entries.len(), 1); assert_eq!(entries.len(), 1);
Ok(())
} }
#[cfg(feature = "crypto")] #[cfg(feature = "crypto")]
#[test] #[test]
fn test_sign_container_wrong_key_fails() { fn test_sign_container_wrong_key_fails() -> Result<(), Box<dyn std::error::Error>> {
use mtp_crypto::{Ed25519Signer, SigAlgorithm}; use mtp_crypto::{Ed25519Signer, SigAlgorithm};
let tm = TypeMap::latest(); let tm = TypeMap::latest();
let (signer, _, _) = Ed25519Signer::generate(); let (signer, _, _) = Ed25519Signer::generate();
let (_, sk2, _) = Ed25519Signer::generate(); let (_, sk2, _) = Ed25519Signer::generate();
let wrong_verifier = Ed25519Signer::new(&sk2).unwrap(); let wrong_verifier = Ed25519Signer::new(&sk2)?;
let mut dv = DataValue::Container(vec![( let mut dv = DataValue::Container(vec![(
DataType::Version.to_id(&tm), DataType::Version.to_id(&tm),
@ -1780,5 +1788,6 @@ mod tests {
assert!(dv.sign_container(SigAlgorithm::ED25519, &signer).is_some()); assert!(dv.sign_container(SigAlgorithm::ED25519, &signer).is_some());
assert!(dv.verify_into_container(&wrong_verifier).is_none()); assert!(dv.verify_into_container(&wrong_verifier).is_none());
Ok(())
} }
} }

View file

@ -190,24 +190,24 @@ mod tests {
#[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))]
#[test] #[test]
fn encrypt_for_roundtrip() { fn encrypt_for_roundtrip() -> Result<(), CryptoError> {
let kr = Keyring::generate(); let kr = Keyring::generate();
let blob = encrypt_for( let blob = encrypt_for(
EncryptionType::MlKemChaCha20Poly1305, EncryptionType::MlKemChaCha20Poly1305,
&kr.public_key_bundle(), &kr.public_key_bundle(),
b"secret payload", b"secret payload",
b"aad", b"aad",
) )?;
.unwrap();
assert_eq!(blob[0], EncryptionType::ML_KEM_CHACHA20POLY1305); assert_eq!(blob[0], EncryptionType::ML_KEM_CHACHA20POLY1305);
let pt = decrypt_with(&blob, &kr, b"aad").unwrap(); let pt = decrypt_with(&blob, &kr, b"aad")?;
assert_eq!(pt, b"secret payload"); assert_eq!(pt, b"secret payload");
Ok(())
} }
#[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))]
#[test] #[test]
fn decrypt_with_wrong_keyring_fails() { fn decrypt_with_wrong_keyring_fails() -> Result<(), CryptoError> {
let kr = Keyring::generate(); let kr = Keyring::generate();
let other = Keyring::generate(); let other = Keyring::generate();
let blob = encrypt_for( let blob = encrypt_for(
@ -215,23 +215,23 @@ mod tests {
&kr.public_key_bundle(), &kr.public_key_bundle(),
b"secret", b"secret",
b"aad", b"aad",
) )?;
.unwrap();
assert!(decrypt_with(&blob, &other, b"aad").is_err()); assert!(decrypt_with(&blob, &other, b"aad").is_err());
Ok(())
} }
#[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))]
#[test] #[test]
fn decrypt_with_wrong_aad_fails() { fn decrypt_with_wrong_aad_fails() -> Result<(), CryptoError> {
let kr = Keyring::generate(); let kr = Keyring::generate();
let blob = encrypt_for( let blob = encrypt_for(
EncryptionType::MlKemChaCha20Poly1305, EncryptionType::MlKemChaCha20Poly1305,
&kr.public_key_bundle(), &kr.public_key_bundle(),
b"secret", b"secret",
b"right", b"right",
) )?;
.unwrap();
assert!(decrypt_with(&blob, &kr, b"wrong").is_err()); assert!(decrypt_with(&blob, &kr, b"wrong").is_err());
Ok(())
} }
#[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))]

View file

@ -498,13 +498,14 @@ impl Keyring {
use crate::error::CryptoError; use crate::error::CryptoError;
let mut offset = 0; let mut offset = 0;
let read_key = |offset: &mut usize| -> Result<Vec<u8>, CryptoError> { let read_key = |offset: &mut usize| -> Result<Vec<u8>, CryptoError> {
let len = u16::from_be_bytes( let slice = bytes
bytes .get(*offset..*offset + 2)
.get(*offset..*offset + 2) .ok_or(CryptoError::InvalidKeyLength)?;
.ok_or(CryptoError::InvalidKeyLength)? let len = if let Ok(arr) = <[u8; 2]>::try_from(slice) {
.try_into() u16::from_be_bytes(arr)
.expect("slice is 2 bytes, verified above"), } else {
) as usize; return Err(CryptoError::InvalidKeyLength);
} as usize;
*offset += 2; *offset += 2;
let key = bytes let key = bytes
.get(*offset..*offset + len) .get(*offset..*offset + len)
@ -712,14 +713,14 @@ mod tests {
use super::*; use super::*;
#[test] #[test]
fn public_key_bundle_roundtrip() { fn public_key_bundle_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let kem = KemPublicKey::new(vec![1u8; 32]); let kem = KemPublicKey::new(vec![1u8; 32]);
let pq = SignaturePqPublicKey::new(vec![2u8; 64]); let pq = SignaturePqPublicKey::new(vec![2u8; 64]);
let cl = SignaturePublicKey::new(vec![3u8; 32]); let cl = SignaturePublicKey::new(vec![3u8; 32]);
let bundle = PublicKeyBundle::new(kem, pq, cl); let bundle = PublicKeyBundle::new(kem, pq, cl);
let bytes = bundle.as_bytes(); let bytes = bundle.as_bytes();
let recovered = PublicKeyBundle::from_bytes(&bytes).unwrap(); let recovered = PublicKeyBundle::from_bytes(&bytes)?;
assert_eq!( assert_eq!(
bundle.kem_public_key.as_bytes(), bundle.kem_public_key.as_bytes(),
@ -733,22 +734,24 @@ mod tests {
bundle.sig_cl_public_key.as_bytes(), bundle.sig_cl_public_key.as_bytes(),
recovered.sig_cl_public_key.as_bytes() recovered.sig_cl_public_key.as_bytes()
); );
Ok(())
} }
#[test] #[test]
fn public_key_bundle_try_from_roundtrip() { fn public_key_bundle_try_from_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let bundle = PublicKeyBundle::new( let bundle = PublicKeyBundle::new(
KemPublicKey::new(vec![0xABu8; 48]), KemPublicKey::new(vec![0xABu8; 48]),
SignaturePqPublicKey::new(vec![0xCDu8; 96]), SignaturePqPublicKey::new(vec![0xCDu8; 96]),
SignaturePublicKey::new(vec![0xEFu8; 32]), SignaturePublicKey::new(vec![0xEFu8; 32]),
); );
let bytes: Vec<u8> = Vec::from(&bundle); let bytes: Vec<u8> = Vec::from(&bundle);
let recovered = PublicKeyBundle::try_from(bytes.as_slice()).unwrap(); let recovered = PublicKeyBundle::try_from(bytes.as_slice())?;
assert_eq!(bundle.as_bytes(), recovered.as_bytes()); assert_eq!(bundle.as_bytes(), recovered.as_bytes());
Ok(())
} }
#[test] #[test]
fn keyring_roundtrip() { fn keyring_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let keyring = Keyring::new( let keyring = Keyring::new(
KemPublicKey::new(vec![1u8; 32]), KemPublicKey::new(vec![1u8; 32]),
KemPrivateKey::new(vec![2u8; 32]), KemPrivateKey::new(vec![2u8; 32]),
@ -758,7 +761,7 @@ mod tests {
SignaturePrivateKey::new(vec![6u8; 32]), SignaturePrivateKey::new(vec![6u8; 32]),
); );
let bytes = keyring.to_bytes(); let bytes = keyring.to_bytes();
let recovered = Keyring::from_bytes(&bytes).unwrap(); let recovered = Keyring::from_bytes(&bytes)?;
assert_eq!( assert_eq!(
keyring.kem_public_key.as_bytes(), keyring.kem_public_key.as_bytes(),
recovered.kem_public_key.as_bytes() recovered.kem_public_key.as_bytes()
@ -771,10 +774,11 @@ mod tests {
keyring.sig_cl_public_key.as_bytes(), keyring.sig_cl_public_key.as_bytes(),
recovered.sig_cl_public_key.as_bytes() recovered.sig_cl_public_key.as_bytes()
); );
Ok(())
} }
#[test] #[test]
fn keyring_try_from_roundtrip() { fn keyring_try_from_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let keyring = Keyring::new( let keyring = Keyring::new(
KemPublicKey::new(vec![0u8; 16]), KemPublicKey::new(vec![0u8; 16]),
KemPrivateKey::new(vec![1u8; 16]), KemPrivateKey::new(vec![1u8; 16]),
@ -784,21 +788,23 @@ mod tests {
SignaturePrivateKey::new(vec![5u8; 16]), SignaturePrivateKey::new(vec![5u8; 16]),
); );
let bytes: Vec<u8> = Vec::from(&keyring); let bytes: Vec<u8> = Vec::from(&keyring);
let recovered = Keyring::try_from(bytes.as_slice()).unwrap(); let recovered = Keyring::try_from(bytes.as_slice())?;
assert_eq!(keyring.to_bytes(), recovered.to_bytes()); assert_eq!(keyring.to_bytes(), recovered.to_bytes());
Ok(())
} }
#[test] #[test]
fn hex_roundtrip() { fn hex_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let key = KemPublicKey::new(vec![0xDE, 0xAD, 0xBE, 0xEF]); let key = KemPublicKey::new(vec![0xDE, 0xAD, 0xBE, 0xEF]);
let hex = key.to_hex(); let hex = key.to_hex();
assert_eq!(hex, "deadbeef"); assert_eq!(hex, "deadbeef");
let recovered = KemPublicKey::from_hex(&hex).unwrap(); let recovered = KemPublicKey::from_hex(&hex)?;
assert_eq!(key.as_bytes(), recovered.as_bytes()); assert_eq!(key.as_bytes(), recovered.as_bytes());
Ok(())
} }
#[test] #[test]
fn keyring_hex_roundtrip() { fn keyring_hex_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let keyring = Keyring::new( let keyring = Keyring::new(
KemPublicKey::new(vec![1u8; 16]), KemPublicKey::new(vec![1u8; 16]),
KemPrivateKey::new(vec![2u8; 16]), KemPrivateKey::new(vec![2u8; 16]),
@ -808,12 +814,13 @@ mod tests {
SignaturePrivateKey::new(vec![6u8; 16]), SignaturePrivateKey::new(vec![6u8; 16]),
); );
let hex = keyring.to_hex(); let hex = keyring.to_hex();
let recovered = Keyring::from_hex(&hex).unwrap(); let recovered = Keyring::from_hex(&hex)?;
assert_eq!(keyring.to_bytes(), recovered.to_bytes()); assert_eq!(keyring.to_bytes(), recovered.to_bytes());
Ok(())
} }
#[test] #[test]
fn keyring_base64_roundtrip() { fn keyring_base64_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let keyring = Keyring::new( let keyring = Keyring::new(
KemPublicKey::new(vec![1u8; 16]), KemPublicKey::new(vec![1u8; 16]),
KemPrivateKey::new(vec![2u8; 16]), KemPrivateKey::new(vec![2u8; 16]),
@ -823,20 +830,22 @@ mod tests {
SignaturePrivateKey::new(vec![6u8; 16]), SignaturePrivateKey::new(vec![6u8; 16]),
); );
let b64 = keyring.to_base64(); let b64 = keyring.to_base64();
let recovered = Keyring::from_base64(&b64).unwrap(); let recovered = Keyring::from_base64(&b64)?;
assert_eq!(keyring.to_bytes(), recovered.to_bytes()); assert_eq!(keyring.to_bytes(), recovered.to_bytes());
Ok(())
} }
#[test] #[test]
fn public_key_bundle_base64_roundtrip() { fn public_key_bundle_base64_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let bundle = PublicKeyBundle::new( let bundle = PublicKeyBundle::new(
KemPublicKey::new(vec![1u8; 32]), KemPublicKey::new(vec![1u8; 32]),
SignaturePqPublicKey::new(vec![2u8; 64]), SignaturePqPublicKey::new(vec![2u8; 64]),
SignaturePublicKey::new(vec![3u8; 32]), SignaturePublicKey::new(vec![3u8; 32]),
); );
let b64 = bundle.to_base64(); let b64 = bundle.to_base64();
let recovered = PublicKeyBundle::from_base64(&b64).unwrap(); let recovered = PublicKeyBundle::from_base64(&b64)?;
assert_eq!(bundle.as_bytes(), recovered.as_bytes()); assert_eq!(bundle.as_bytes(), recovered.as_bytes());
Ok(())
} }
#[test] #[test]

View file

@ -69,13 +69,14 @@ mod tests {
#[cfg(feature = "chacha20poly1305")] #[cfg(feature = "chacha20poly1305")]
#[test] #[test]
fn aead_encrypt_decrypt() { fn aead_encrypt_decrypt() -> Result<(), CryptoError> {
use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305}; use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305};
let key = [0xAB; 32]; let key = [0xAB; 32];
let cipher = ChaCha20Poly1305::new(key); let cipher = ChaCha20Poly1305::new(key);
let ct = cipher.encrypt(b"hello world", b"aad").unwrap(); let ct = cipher.encrypt(b"hello world", b"aad")?;
let pt = cipher.decrypt(&ct, b"aad").unwrap(); let pt = cipher.decrypt(&ct, b"aad")?;
assert_eq!(pt, b"hello world"); assert_eq!(pt, b"hello world");
Ok(())
} }
#[cfg(feature = "chacha20poly1305")] #[cfg(feature = "chacha20poly1305")]
@ -84,7 +85,7 @@ mod tests {
use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305}; use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305};
let cipher_a = ChaCha20Poly1305::new([0xAB; 32]); let cipher_a = ChaCha20Poly1305::new([0xAB; 32]);
let cipher_b = ChaCha20Poly1305::new([0xCD; 32]); let cipher_b = ChaCha20Poly1305::new([0xCD; 32]);
let ct = cipher_a.encrypt(b"hello", b"").unwrap(); let ct = cipher_a.encrypt(b"hello", b"").expect("encryption should succeed");
assert!(cipher_b.decrypt(&ct, b"").is_err()); assert!(cipher_b.decrypt(&ct, b"").is_err());
} }
@ -93,7 +94,9 @@ mod tests {
fn aead_wrong_aad_fails() { fn aead_wrong_aad_fails() {
use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305}; use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305};
let cipher = ChaCha20Poly1305::new([0xAB; 32]); let cipher = ChaCha20Poly1305::new([0xAB; 32]);
let ct = cipher.encrypt(b"hello", b"correct-aad").unwrap(); let ct = cipher
.encrypt(b"hello", b"correct-aad")
.expect("encryption should succeed");
assert!(cipher.decrypt(&ct, b"wrong-aad").is_err()); assert!(cipher.decrypt(&ct, b"wrong-aad").is_err());
} }
@ -102,12 +105,12 @@ mod tests {
fn ed25519_sign_verify() { fn ed25519_sign_verify() {
let (signer, sk, pk) = Ed25519Signer::generate(); let (signer, sk, pk) = Ed25519Signer::generate();
let msg = b"test message"; let msg = b"test message";
let sig = signer.sign(msg).unwrap(); let sig = signer.sign(msg).expect("signing should succeed");
signer.verify(msg, &sig).unwrap(); signer.verify(msg, &sig).expect("verification should succeed");
verify_ed25519(&pk, msg, &sig).unwrap(); verify_ed25519(&pk, msg, &sig).expect("verification should succeed");
let loaded = Ed25519Signer::new(&sk).unwrap(); let loaded = Ed25519Signer::new(&sk).expect("signer loading should succeed");
loaded.verify(msg, &sig).unwrap(); loaded.verify(msg, &sig).expect("verification should succeed");
} }
#[cfg(feature = "ed25519-dalek")] #[cfg(feature = "ed25519-dalek")]
@ -115,7 +118,7 @@ mod tests {
fn ed25519_wrong_sig_fails() { fn ed25519_wrong_sig_fails() {
let (signer, _, pk) = Ed25519Signer::generate(); let (signer, _, pk) = Ed25519Signer::generate();
let msg = b"test message"; let msg = b"test message";
let sig = signer.sign(msg).unwrap(); let sig = signer.sign(msg).expect("signing should succeed");
assert!(verify_ed25519(&pk, b"wrong message", &sig).is_err()); assert!(verify_ed25519(&pk, b"wrong message", &sig).is_err());
} }
@ -124,12 +127,12 @@ mod tests {
fn mldsa_sign_verify() { fn mldsa_sign_verify() {
let (signer, sk, pk) = MlDsaSigner::generate(); let (signer, sk, pk) = MlDsaSigner::generate();
let msg = b"test message"; let msg = b"test message";
let sig = signer.sign(msg).unwrap(); let sig = signer.sign(msg).expect("signing should succeed");
signer.verify(msg, &sig).unwrap(); signer.verify(msg, &sig).expect("verification should succeed");
verify_ml_dsa(&pk, msg, &sig).unwrap(); verify_ml_dsa(&pk, msg, &sig).expect("verification should succeed");
let loaded = MlDsaSigner::new(&sk, &pk).unwrap(); let loaded = MlDsaSigner::new(&sk, &pk).expect("signer loading should succeed");
loaded.verify(msg, &sig).unwrap(); loaded.verify(msg, &sig).expect("verification should succeed");
} }
#[cfg(feature = "ml-dsa")] #[cfg(feature = "ml-dsa")]
@ -137,7 +140,7 @@ mod tests {
fn mldsa_wrong_sig_fails() { fn mldsa_wrong_sig_fails() {
let (signer, _, pk) = MlDsaSigner::generate(); let (signer, _, pk) = MlDsaSigner::generate();
let msg = b"test message"; let msg = b"test message";
let sig = signer.sign(msg).unwrap(); let sig = signer.sign(msg).expect("signing should succeed");
assert!(verify_ml_dsa(&pk, b"wrong message", &sig).is_err()); assert!(verify_ml_dsa(&pk, b"wrong message", &sig).is_err());
} }
@ -148,9 +151,10 @@ mod tests {
let (ed_signer, _, _) = Ed25519Signer::generate(); let (ed_signer, _, _) = Ed25519Signer::generate();
let (ml_signer, _, _) = MlDsaSigner::generate(); let (ml_signer, _, _) = MlDsaSigner::generate();
let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg").unwrap(); let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg")
.expect("dual signing should succeed");
dual.verify(ed_signer.verifying_key(), ml_signer.verifying_key(), b"msg") dual.verify(ed_signer.verifying_key(), ml_signer.verifying_key(), b"msg")
.unwrap(); .expect("dual verification should succeed");
} }
#[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))] #[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))]
@ -160,7 +164,8 @@ mod tests {
let (ed_signer, _, _) = Ed25519Signer::generate(); let (ed_signer, _, _) = Ed25519Signer::generate();
let (ml_signer, _, _) = MlDsaSigner::generate(); let (ml_signer, _, _) = MlDsaSigner::generate();
let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg").unwrap(); let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg")
.expect("dual signing should succeed");
assert!( assert!(
dual.verify( dual.verify(
ed_signer.verifying_key(), ed_signer.verifying_key(),
@ -174,18 +179,22 @@ mod tests {
#[cfg(feature = "hkdf")] #[cfg(feature = "hkdf")]
#[test] #[test]
fn hkdf_expand_produces_key() { fn hkdf_expand_produces_key() {
let key = derive_encryption_key(b"ikm", b"salt", b"context").unwrap(); let key = derive_encryption_key(b"ikm", b"salt", b"context")
.expect("key derivation should succeed");
assert_eq!(key.len(), 32); assert_eq!(key.len(), 32);
let expanded = hkdf_expand(b"ikm", b"salt", b"info", 64).unwrap(); let expanded = hkdf_expand(b"ikm", b"salt", b"info", 64)
.expect("HKDF expansion should succeed");
assert_eq!(expanded.len(), 64); assert_eq!(expanded.len(), 64);
} }
#[cfg(feature = "hkdf")] #[cfg(feature = "hkdf")]
#[test] #[test]
fn hkdf_different_info_different_key() { fn hkdf_different_info_different_key() {
let a = derive_encryption_key(b"ikm", b"salt", b"info-a").unwrap(); let a = derive_encryption_key(b"ikm", b"salt", b"info-a")
let b = derive_encryption_key(b"ikm", b"salt", b"info-b").unwrap(); .expect("key derivation should succeed");
let b = derive_encryption_key(b"ikm", b"salt", b"info-b")
.expect("key derivation should succeed");
assert_ne!(a, b); assert_ne!(a, b);
} }
@ -255,7 +264,7 @@ mod tests {
fn keyring_serialize_roundtrip() { fn keyring_serialize_roundtrip() {
let kr = Keyring::generate(); let kr = Keyring::generate();
let bytes = kr.to_bytes(); let bytes = kr.to_bytes();
let loaded = Keyring::from_bytes(&bytes).unwrap(); let loaded = Keyring::from_bytes(&bytes).expect("keyring roundtrip should succeed");
assert_eq!( assert_eq!(
kr.kem_public_key.as_bytes(), kr.kem_public_key.as_bytes(),
loaded.kem_public_key.as_bytes() loaded.kem_public_key.as_bytes()
@ -276,7 +285,8 @@ mod tests {
let kr = Keyring::generate(); let kr = Keyring::generate();
let bundle = kr.public_key_bundle(); let bundle = kr.public_key_bundle();
let bytes = bundle.as_bytes(); let bytes = bundle.as_bytes();
let loaded = PublicKeyBundle::from_bytes(&bytes).unwrap(); let loaded =
PublicKeyBundle::from_bytes(&bytes).expect("bundle roundtrip should succeed");
assert_eq!( assert_eq!(
bundle.kem_public_key.as_bytes(), bundle.kem_public_key.as_bytes(),
loaded.kem_public_key.as_bytes() loaded.kem_public_key.as_bytes()
@ -295,8 +305,8 @@ mod tests {
#[test] #[test]
fn hybrid_kem_roundtrip() { fn hybrid_kem_roundtrip() {
let (sk, pk) = HybridKem::generate_keypair(); let (sk, pk) = HybridKem::generate_keypair();
let enc = HybridKem::encapsulate(&pk).unwrap(); let enc = HybridKem::encapsulate(&pk).expect("encapsulation should succeed");
let ss = HybridKem::decapsulate(&sk, &enc.ciphertext).unwrap(); let ss = HybridKem::decapsulate(&sk, &enc.ciphertext).expect("decapsulation should succeed");
assert_eq!(enc.shared_secret, ss); assert_eq!(enc.shared_secret, ss);
} }
@ -309,8 +319,8 @@ mod tests {
let kr = Keyring::generate(); let kr = Keyring::generate();
let entities = vec![kr.public_key_bundle()]; let entities = vec![kr.public_key_bundle()];
let msg = b"secret data"; let msg = b"secret data";
let ct = encrypt_multi(msg, b"aad", &entities).unwrap(); let ct = encrypt_multi(msg, b"aad", &entities).expect("multi encrypt should succeed");
let pt = decrypt_multi(&ct, b"aad", &kr).unwrap(); let pt = decrypt_multi(&ct, b"aad", &kr).expect("multi decrypt should succeed");
assert_eq!(pt, msg); assert_eq!(pt, msg);
} }
} }

View file

@ -3,41 +3,48 @@ protocol_version: "0.0"
# Note that markers 0 to 31 are reserved for default use, manually working with them is not recommended # Note that markers 0 to 31 are reserved for default use, manually working with them is not recommended
# Fixed CommunicationType markers are: # Fixed CommunicationType markers are:
# Error: 0 # Identification: 0
# ErrorParsing: 1 # IdentificationResponse: 1
# ErrorBadVersion: 2 # Register: 2
# Disconnect: 3 # RegisterResponse: 3
# Redirect: 4 # Challenge: 4
# Shutdown: 5 # ChallengeResponse: 5
# BadRequest: 6 # Ping: 6
# Unauthorized: 7 # Pong: 7
# Forbidden: 8 # Disconnect: 8
# NotFound: 9 # Redirect: 9
# TooManyRequests: 10 # Shutdown: 10
# InternalServerError: 11 # Error: 11
# BadGateway: 12 # ErrorParsing: 12
# ServiceUnavailable: 13 # ErrorBadVersion: 13
# GatewayTimeout: 14 # BadRequest: 14
# Identification: 15 # Unauthorized: 15
# IdentificationResponse: 16 # Forbidden: 16
# Register: 17 # NotFound: 17
# RegisterResponse: 18 # TooManyRequests: 18
# Ping: 19 # InternalServerError: 19
# Pong: 20 # BadGateway: 20
# ServiceUnavailable: 21
# GatewayTimeout: 22
# PipeRequest: 23
# PipeResponse: 24
# PipeAbort: 25
# #
# Fixed Data Type markers are: # Fixed Data Type markers are:
# Error: 0 # Version: 0
# ErrorParsing: 1 # id: 1
# ErrorMessage: 2 # ClientNonce: 2
# Version: 3 # ServerNonce: 3
# Description: 4 # PublicKeys: 4
# Timestamp: 5 # Signature: 5
# Id: 6 # PqSignature: 6
# ClientNonce: 7 # Description: 7
# ServerNonce: 8 # Connected: 8
# PublicKeys: 9 # Timestamp: 9
# Signature: 10 # Error: 10
# Connected: 11 # ErrorParsing: 11
# ErrorMessage: 12
# Accepted: 13,
# #
# If a Type can't be used it will be mapped to 0 # If a Type can't be used it will be mapped to 0

View file

@ -26,8 +26,15 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
"Missing TLS certificate at {cert_path}: enter the Nix shell first or run the server to generate it: {e}" "Missing TLS certificate at {cert_path}: enter the Nix shell first or run the server to generate it: {e}"
) )
}); });
let host_public_key = load_public_key_bundle("host.mpkb") let host_public_key = match load_public_key_bundle("host.mpkb") {
.expect("Missing host.mpkb: run the server first to export it"); Ok(bundle) => bundle,
Err(e) => {
return Err(format!(
"Missing host.mpkb: run the server first to export it ({e})"
)
.into());
}
};
println!("Connecting to 127.0.0.1:8080 ..."); println!("Connecting to 127.0.0.1:8080 ...");

View file

@ -7,11 +7,10 @@ pub fn build_demo_message(
client_id: u64, client_id: u64,
keyring: &Keyring, keyring: &Keyring,
server_bundle: &PublicKeyBundle, server_bundle: &PublicKeyBundle,
) -> CommunicationValue { ) -> Result<CommunicationValue, Box<dyn std::error::Error>> {
// Encrypt to the server's KEM public key; the server decrypts with its keyring. // Encrypt to the server's KEM public key; the server decrypts with its keyring.
let enc_type = EncryptionType::MlKemChaCha20Poly1305; let enc_type = EncryptionType::MlKemChaCha20Poly1305;
let signer = let signer = Ed25519Signer::new(&keyring.sig_cl_secret_key)?;
Ed25519Signer::new(&keyring.sig_cl_secret_key).expect("Ed25519 signer from keyring");
let tm = TypeMap::latest(); let tm = TypeMap::latest();
@ -52,8 +51,7 @@ pub fn build_demo_message(
); );
let timestamp = std::time::SystemTime::now() let timestamp = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH) .duration_since(std::time::UNIX_EPOCH)?
.unwrap()
.as_secs(); .as_secs();
let msg = CommunicationValue::new(CommunicationType::Ping) let msg = CommunicationValue::new(CommunicationType::Ping)
@ -84,7 +82,7 @@ pub fn build_demo_message(
.add_typed_default(DataType::SignedPayload, dv_sig) .add_typed_default(DataType::SignedPayload, dv_sig)
.add_typed_default(DataType::SecurePayload, dv_sec) .add_typed_default(DataType::SecurePayload, dv_sec)
.with_sender(client_id); .with_sender(client_id);
msg Ok(msg)
} }
pub async fn send_and_receive( pub async fn send_and_receive(
@ -92,7 +90,7 @@ pub async fn send_and_receive(
keyring: &Keyring, keyring: &Keyring,
server_bundle: &PublicKeyBundle, server_bundle: &PublicKeyBundle,
) -> Result<(), Box<dyn std::error::Error>> { ) -> Result<(), Box<dyn std::error::Error>> {
let msg = build_demo_message(conn.client_id, keyring, server_bundle); let msg = build_demo_message(conn.client_id, keyring, server_bundle)?;
println!("Sending: {msg}"); println!("Sending: {msg}");
conn.sender.send(&msg).await?; conn.sender.send(&msg).await?;

View file

@ -19,7 +19,12 @@ pub fn load_client_db(
Err(_) => HashMap::new(), Err(_) => HashMap::new(),
}; };
let clients: Arc<Mutex<HashMap<u64, PublicKeyBundle>>> = Arc::new(Mutex::new(clients_map)); let clients: Arc<Mutex<HashMap<u64, PublicKeyBundle>>> = Arc::new(Mutex::new(clients_map));
let next_value = clients.lock().unwrap().keys().max().unwrap_or(&999) + 1; let next_value = {
let guard = clients
.lock()
.map_err(|_| std::io::Error::other("client database mutex poisoned"))?;
guard.keys().max().copied().unwrap_or(999) + 1
};
let next_id = Arc::new(Mutex::new(next_value)); let next_id = Arc::new(Mutex::new(next_value));
Ok((clients, next_id)) Ok((clients, next_id))
} }

View file

@ -17,17 +17,47 @@ pub fn process_and_respond(
tm: &TypeMap, tm: &TypeMap,
client_pk: Option<&mtp::crypto::PublicKeyBundle>, client_pk: Option<&mtp::crypto::PublicKeyBundle>,
host_keyring: &Keyring, host_keyring: &Keyring,
) -> CommunicationValue { ) -> Result<CommunicationValue, String> {
let desc_id = DataTypeId(tm.data_id_enum(DataType::Description).unwrap()); let desc_id = DataTypeId(
let ts_id = DataTypeId(tm.data_id_enum(DataType::Timestamp).unwrap()); tm.data_id_enum(DataType::Description)
let data_id = DataTypeId(tm.data_id_enum(DataType::Data).unwrap()); .ok_or("missing Description type mapping")?,
let flags_id = DataTypeId(tm.data_id_enum(DataType::Flags).unwrap()); );
let value_id = DataTypeId(tm.data_id_enum(DataType::Value).unwrap()); let ts_id = DataTypeId(
let bin_id = DataTypeId(tm.data_id_enum(DataType::BinaryData).unwrap()); tm.data_id_enum(DataType::Timestamp)
let items_id = DataTypeId(tm.data_id_enum(DataType::Items).unwrap()); .ok_or("missing Timestamp type mapping")?,
let _enc_id = DataTypeId(tm.data_id_enum(DataType::EncryptedPayload).unwrap()); );
let _sig_id = DataTypeId(tm.data_id_enum(DataType::SignedPayload).unwrap()); let data_id = DataTypeId(
let _secure_id = DataTypeId(tm.data_id_enum(DataType::SecurePayload).unwrap()); tm.data_id_enum(DataType::Data)
.ok_or("missing Data type mapping")?,
);
let flags_id = DataTypeId(
tm.data_id_enum(DataType::Flags)
.ok_or("missing Flags type mapping")?,
);
let value_id = DataTypeId(
tm.data_id_enum(DataType::Value)
.ok_or("missing Value type mapping")?,
);
let bin_id = DataTypeId(
tm.data_id_enum(DataType::BinaryData)
.ok_or("missing BinaryData type mapping")?,
);
let items_id = DataTypeId(
tm.data_id_enum(DataType::Items)
.ok_or("missing Items type mapping")?,
);
let _enc_id = DataTypeId(
tm.data_id_enum(DataType::EncryptedPayload)
.ok_or("missing EncryptedPayload type mapping")?,
);
let _sig_id = DataTypeId(
tm.data_id_enum(DataType::SignedPayload)
.ok_or("missing SignedPayload type mapping")?,
);
let _secure_id = DataTypeId(
tm.data_id_enum(DataType::SecurePayload)
.ok_or("missing SecurePayload type mapping")?,
);
let description = msg.get_data(DataType::Description); let description = msg.get_data(DataType::Description);
let timestamp = msg.get_data(DataType::Timestamp); let timestamp = msg.get_data(DataType::Timestamp);
@ -118,10 +148,10 @@ pub fn process_and_respond(
let now = std::time::SystemTime::now() let now = std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH) .duration_since(std::time::UNIX_EPOCH)
.unwrap() .map_err(|e| e.to_string())?
.as_secs(); .as_secs();
CommunicationValue::from_comm(CommunicationType::Pong, tm) Ok(CommunicationValue::from_comm(CommunicationType::Pong, tm)
.add_data(desc_id, description.clone()) .add_data(desc_id, description.clone())
.add_data(ts_id, DataValue::UnsignedNumber(now as u128)) .add_data(ts_id, DataValue::UnsignedNumber(now as u128))
.add_data( .add_data(
@ -134,5 +164,5 @@ pub fn process_and_respond(
.add_data(flags_id, flags.clone()) .add_data(flags_id, flags.clone())
.add_data(value_id, value.clone()) .add_data(value_id, value.clone())
.add_data(bin_id, binary.clone()) .add_data(bin_id, binary.clone())
.add_data(items_id, items.clone()) .add_data(items_id, items.clone()))
} }

View file

@ -62,7 +62,10 @@ async fn handle_pipe_loopback(
tokio::io::AsyncWriteExt::write_all(&mut writer, &buf[..n]).await?; tokio::io::AsyncWriteExt::write_all(&mut writer, &buf[..n]).await?;
} }
writer.finish().await?; writer.finish().await?;
println!(" [loopback] Pipe {pipe_id} loopback complete ({} bytes)", total); println!(
" [loopback] Pipe {pipe_id} loopback complete ({} bytes)",
total
);
} }
None => { None => {
eprintln!(" [loopback] Return pipe denied by client for pipe {pipe_id}"); eprintln!(" [loopback] Return pipe denied by client for pipe {pipe_id}");
@ -83,10 +86,13 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let (_host_id, host_keyring) = keys::load_or_generate_host_keys("host.mk")?; let (_host_id, host_keyring) = keys::load_or_generate_host_keys("host.mk")?;
keys::export_host_public_keys(&host_keyring)?; keys::export_host_public_keys(&host_keyring)?;
let decrypt_keyring = Arc::new( let decrypt_keyring = Arc::new(match mtp::crypto::Keyring::from_bytes(&host_keyring.to_bytes())
mtp::crypto::Keyring::from_bytes(&host_keyring.to_bytes()) {
.expect("re-load host keyring for decryption"), Ok(keyring) => keyring,
); Err(e) => {
return Err(format!("failed to re-load host keyring for decryption: {e}").into());
}
});
let (clients, next_id) = clients::load_client_db("clients.json")?; let (clients, next_id) = clients::load_client_db("clients.json")?;
@ -94,7 +100,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let get_existing_user = move |id: u64, _description: Option<String>| { let get_existing_user = move |id: u64, _description: Option<String>| {
let clients = clients_for_get.clone(); let clients = clients_for_get.clone();
Box::pin(async move { Box::pin(async move {
let result = clients.lock().unwrap().get(&id).cloned(); let result = clients.lock()?.get(&id).cloned();
if result.is_some() { if result.is_some() {
println!("Auth lookup: client ID {id} found"); println!("Auth lookup: client ID {id} found");
} else { } else {
@ -113,8 +119,8 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
let nid_arc = next_id_for_register.clone(); let nid_arc = next_id_for_register.clone();
let path = clients_path.clone(); let path = clients_path.clone();
Box::pin(async move { Box::pin(async move {
let mut db = db_arc.lock().unwrap(); let mut db = db_arc.lock()?;
let mut nid = nid_arc.lock().unwrap(); let mut nid = nid_arc.lock()?;
let id = *nid; let id = *nid;
*nid += 1; *nid += 1;
db.insert(id, bundle); db.insert(id, bundle);
@ -160,7 +166,7 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
); );
println!("Client ID: {}", conn.client_id); println!("Client ID: {}", conn.client_id);
let tm: &TypeMap = conn.codec.registry().get(&conn.version).unwrap(); let tm: &TypeMap = conn.codec.registry().get(&conn.version)?;
println!("Waiting for messages / pipe requests ..."); println!("Waiting for messages / pipe requests ...");
loop { loop {
@ -185,12 +191,18 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
match msg { match msg {
Ok(msg) => { Ok(msg) => {
println!("Received: {msg}"); println!("Received: {msg}");
let response = handlers::process_and_respond( let response = match handlers::process_and_respond(
&msg, &msg,
tm, tm,
conn.client_public_key.as_ref(), conn.client_public_key.as_ref(),
&decrypt_keyring, &decrypt_keyring,
); ) {
Ok(response) => response,
Err(e) => {
eprintln!("Failed to build response: {e}");
continue;
}
};
println!("Sending: {response}"); println!("Sending: {response}");
if let Err(e) = conn.sender.send(&response).await { if let Err(e) = conn.sender.send(&response).await {
eprintln!("Send error: {e}"); eprintln!("Send error: {e}");

View file

@ -146,52 +146,57 @@ mod tests {
} }
#[test] #[test]
fn keyring_save_load_roundtrip() { fn keyring_save_load_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let path = temp_path(KEYRING_EXTENSION); let path = temp_path(KEYRING_EXTENSION);
let keyring = sample_keyring(); let keyring = sample_keyring();
save_keyring(&keyring, &path).unwrap(); save_keyring(&keyring, &path)?;
let loaded = load_keyring(&path).unwrap(); let loaded = load_keyring(&path)?;
assert_eq!(keyring.to_bytes(), loaded.to_bytes()); assert_eq!(keyring.to_bytes(), loaded.to_bytes());
let _ = fs::remove_file(&path); let _ = fs::remove_file(&path);
Ok(())
} }
#[test] #[test]
fn bundle_save_load_roundtrip() { fn bundle_save_load_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let path = temp_path(BUNDLE_EXTENSION); let path = temp_path(BUNDLE_EXTENSION);
let bundle = sample_keyring().public_key_bundle(); let bundle = sample_keyring().public_key_bundle();
save_public_key_bundle(&bundle, &path).unwrap(); save_public_key_bundle(&bundle, &path)?;
let loaded = load_public_key_bundle(&path).unwrap(); let loaded = load_public_key_bundle(&path)?;
assert_eq!(bundle.as_bytes(), loaded.as_bytes()); assert_eq!(bundle.as_bytes(), loaded.as_bytes());
let _ = fs::remove_file(&path); let _ = fs::remove_file(&path);
Ok(())
} }
#[test] #[test]
fn loading_bundle_as_keyring_fails_on_magic() { fn loading_bundle_as_keyring_fails_on_magic() -> Result<(), Box<dyn std::error::Error>> {
let path = temp_path(BUNDLE_EXTENSION); let path = temp_path(BUNDLE_EXTENSION);
save_public_key_bundle(&sample_keyring().public_key_bundle(), &path).unwrap(); save_public_key_bundle(&sample_keyring().public_key_bundle(), &path)?;
assert!(matches!( assert!(matches!(
load_keyring(&path), load_keyring(&path),
Err(FileError::BadMagic { .. }) Err(FileError::BadMagic { .. })
)); ));
let _ = fs::remove_file(&path); let _ = fs::remove_file(&path);
Ok(())
} }
#[test] #[test]
fn truncated_file_is_rejected() { fn truncated_file_is_rejected() -> Result<(), Box<dyn std::error::Error>> {
let path = temp_path(KEYRING_EXTENSION); let path = temp_path(KEYRING_EXTENSION);
fs::write(&path, b"MT").unwrap(); fs::write(&path, b"MT")?;
assert!(matches!(load_keyring(&path), Err(FileError::Truncated(2)))); assert!(matches!(load_keyring(&path), Err(FileError::Truncated(2))));
let _ = fs::remove_file(&path); let _ = fs::remove_file(&path);
Ok(())
} }
#[cfg(unix)] #[cfg(unix)]
#[test] #[test]
fn keyring_file_is_owner_only() { fn keyring_file_is_owner_only() -> Result<(), Box<dyn std::error::Error>> {
use std::os::unix::fs::PermissionsExt; use std::os::unix::fs::PermissionsExt;
let path = temp_path(KEYRING_EXTENSION); let path = temp_path(KEYRING_EXTENSION);
save_keyring(&sample_keyring(), &path).unwrap(); save_keyring(&sample_keyring(), &path)?;
let mode = fs::metadata(&path).unwrap().permissions().mode(); let mode = fs::metadata(&path)?.permissions().mode();
assert_eq!(mode & 0o777, 0o600); assert_eq!(mode & 0o777, 0o600);
let _ = fs::remove_file(&path); let _ = fs::remove_file(&path);
Ok(())
} }
} }

View file

@ -225,7 +225,10 @@ async fn run_dispatcher(
// The WebTransport client opens the pipe request on a fresh stream. // The WebTransport client opens the pipe request on a fresh stream.
// That stream arrives here as a pipe event, not a message event, so we // That stream arrives here as a pipe event, not a message event, so we
// reconstruct the host-side PipeRequest here and hand it to receive_pipe(). // reconstruct the host-side PipeRequest here and hand it to receive_pipe().
println!("Dispatcher treating pipe stream as PipeRequest: id={}", pipe_id); println!(
"Dispatcher treating pipe stream as PipeRequest: id={}",
pipe_id
);
let req = PipeRequest { let req = PipeRequest {
pipe_id, pipe_id,
description: reader.description().to_string(), description: reader.description().to_string(),
@ -471,8 +474,12 @@ impl MTPHost {
Some(v) => v, Some(v) => v,
None => return Err(AcceptError::UnsupportedVersion(client_version)), None => return Err(AcceptError::UnsupportedVersion(client_version)),
}; };
let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) {
.expect("negotiated version must be registered"); Some(codec) => codec,
None => {
return Err(AcceptError::UnsupportedVersion(negotiated));
}
};
let description = match first_msg.get_data(DataType::Description) { let description = match first_msg.get_data(DataType::Description) {
DataValue::Str(s) => Some(s.clone()), DataValue::Str(s) => Some(s.clone()),
_ => None, _ => None,
@ -508,8 +515,12 @@ impl MTPHost {
Some(v) => v, Some(v) => v,
None => return Err(AcceptError::UnsupportedVersion(client_version)), None => return Err(AcceptError::UnsupportedVersion(client_version)),
}; };
let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) {
.expect("negotiated version must be registered"); Some(codec) => codec,
None => {
return Err(AcceptError::UnsupportedVersion(negotiated));
}
};
let description = match first_msg.get_data(DataType::Description) { let description = match first_msg.get_data(DataType::Description) {
DataValue::Str(s) => Some(s.clone()), DataValue::Str(s) => Some(s.clone()),
_ => None, _ => None,
@ -1047,8 +1058,12 @@ impl MTPHost {
return Err(AcceptError::Send(e)); return Err(AcceptError::Send(e));
} }
let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) {
.expect("negotiated version must be registered"); Some(codec) => codec,
None => {
return Err(AcceptError::UnsupportedVersion(negotiated));
}
};
Ok(Some(self.connection_from_parts( Ok(Some(self.connection_from_parts(
sender, sender,
@ -1162,8 +1177,12 @@ impl MTPHost {
Some(v) => v, Some(v) => v,
None => return Err(AcceptError::UnsupportedVersion(client_version)), None => return Err(AcceptError::UnsupportedVersion(client_version)),
}; };
let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) {
.expect("negotiated version must be registered"); Some(codec) => codec,
None => {
return Err(AcceptError::UnsupportedVersion(negotiated));
}
};
return Ok(Some(self.connection_from_parts( return Ok(Some(self.connection_from_parts(
sender, sender,
receiver, receiver,
@ -1246,9 +1265,10 @@ mod tests {
} }
#[test] #[test]
fn host_config_pongs_default_to_enabled() { fn host_config_pongs_default_to_enabled() -> Result<(), Box<dyn std::error::Error>> {
let config = HostConfig::new("127.0.0.1".parse().unwrap(), 4433, Vec::new(), Vec::new()); let config = HostConfig::new("127.0.0.1".parse()?, 4433, Vec::new(), Vec::new());
assert!(config.send_pongs); assert!(config.send_pongs);
assert!(!config.with_pongs(false).send_pongs); assert!(!config.with_pongs(false).send_pongs);
Ok(())
} }
} }

View file

@ -6,10 +6,18 @@ use std::sync::Arc;
use wtransport::{Connection as WTConnection, Endpoint, ServerConfig}; use wtransport::{Connection as WTConnection, Endpoint, ServerConfig};
fn generate_self_signed_cert() -> (Vec<u8>, Vec<u8>) { fn generate_self_signed_cert() -> (Vec<u8>, Vec<u8>) {
let key_pair = rcgen::KeyPair::generate().unwrap(); let key_pair = match rcgen::KeyPair::generate() {
let params = Ok(key_pair) => key_pair,
rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]).unwrap(); Err(e) => panic!("failed to generate self-signed key pair: {e}"),
let cert = params.self_signed(&key_pair).unwrap(); };
let params = match rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]) {
Ok(params) => params,
Err(e) => panic!("failed to build self-signed certificate params: {e}"),
};
let cert = match params.self_signed(&key_pair) {
Ok(cert) => cert,
Err(e) => panic!("failed to self-sign certificate: {e}"),
};
let cert_pem = cert.pem(); let cert_pem = cert.pem();
let key_pem = key_pair.serialize_pem(); let key_pem = key_pair.serialize_pem();
(cert_pem.into_bytes(), key_pem.into_bytes()) (cert_pem.into_bytes(), key_pem.into_bytes())

View file

@ -4,40 +4,39 @@ use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue, Type
use mtp_transport::{Host, Policy, Receiver, Sender, connect, host}; use mtp_transport::{Host, Policy, Receiver, Sender, connect, host};
fn generate_self_signed_cert() -> (Vec<u8>, Vec<u8>) { fn generate_self_signed_cert() -> (Vec<u8>, Vec<u8>) {
let key_pair = rcgen::KeyPair::generate().unwrap(); let key_pair = rcgen::KeyPair::generate().expect("failed to generate self-signed key pair");
let params = let params = rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()])
rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]).unwrap(); .expect("failed to build self-signed certificate params");
let cert = params.self_signed(&key_pair).unwrap(); let cert = params
.self_signed(&key_pair)
.expect("failed to self-sign certificate");
let cert_pem = cert.pem(); let cert_pem = cert.pem();
let key_pem = key_pair.serialize_pem(); let key_pem = key_pair.serialize_pem();
(cert_pem.into_bytes(), key_pem.into_bytes()) (cert_pem.into_bytes(), key_pem.into_bytes())
} }
async fn start_test_host(cert_pem: Vec<u8>, key_pem: Vec<u8>) -> Host { async fn start_test_host(cert_pem: Vec<u8>, key_pem: Vec<u8>) -> Result<Host, Box<dyn std::error::Error>> {
host( Ok(host(
IpAddr::V4(Ipv4Addr::LOCALHOST), IpAddr::V4(Ipv4Addr::LOCALHOST),
0, 0,
cert_pem, cert_pem,
key_pem, key_pem,
Policy::default(), Policy::default(),
) )
.await .await?)
.unwrap()
} }
async fn connect_to_host(h: &Host, cert_pem: Vec<u8>) -> (Sender, Receiver) { async fn connect_to_host(h: &Host, cert_pem: Vec<u8>) -> Result<(Sender, Receiver), Box<dyn std::error::Error>> {
let url = format!("https://127.0.0.1:{}", h.local_addr().port()); let url = format!("https://127.0.0.1:{}", h.local_addr().port());
connect(&url, Some(cert_pem), Policy::default()) Ok(connect(&url, Some(cert_pem), Policy::default()).await?)
.await
.unwrap()
} }
async fn connected_pair() -> (Host, Sender, Receiver, Sender, Receiver) { async fn connected_pair() -> Result<(Host, Sender, Receiver, Sender, Receiver), Box<dyn std::error::Error>> {
let (cert_pem, key_pem) = generate_self_signed_cert(); let (cert_pem, key_pem) = generate_self_signed_cert();
let mut h = start_test_host(cert_pem.clone(), key_pem).await; let mut h = start_test_host(cert_pem.clone(), key_pem).await?;
let (client_tx, client_rx) = connect_to_host(&h, cert_pem).await; let (client_tx, client_rx) = connect_to_host(&h, cert_pem).await?;
let (host_tx, host_rx) = h.next().await.unwrap(); let (host_tx, host_rx) = h.next().await.ok_or("host did not accept connection")?;
(h, client_tx, client_rx, host_tx, host_rx) Ok((h, client_tx, client_rx, host_tx, host_rx))
} }
fn numbered_message(comm_type: CommunicationType, value: u128, tm: &TypeMap) -> CommunicationValue { fn numbered_message(comm_type: CommunicationType, value: u128, tm: &TypeMap) -> CommunicationValue {
@ -61,103 +60,105 @@ fn assert_numbered_message(
} }
#[tokio::test] #[tokio::test]
async fn test_host_start_and_stop() { async fn test_host_start_and_stop() -> Result<(), Box<dyn std::error::Error>> {
let (cert_pem, key_pem) = generate_self_signed_cert(); let (cert_pem, key_pem) = generate_self_signed_cert();
let h = start_test_host(cert_pem, key_pem).await; let h = start_test_host(cert_pem, key_pem).await?;
let addr = h.local_addr(); let addr = h.local_addr();
// Port should be non-zero (OS-assigned) // Port should be non-zero (OS-assigned)
assert!(addr.port() > 0); assert!(addr.port() > 0);
Ok(())
} }
#[tokio::test] #[tokio::test]
async fn test_send_receive_roundtrip() { async fn test_send_receive_roundtrip() -> Result<(), Box<dyn std::error::Error>> {
let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await; let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await?;
let tm = TypeMap::latest(); let tm = TypeMap::latest();
// Client sends a simple message // Client sends a simple message
let msg = numbered_message(CommunicationType::Ping, 42, &tm); let msg = numbered_message(CommunicationType::Ping, 42, &tm);
client_tx.send(&msg).await.unwrap(); client_tx.send(&msg).await?;
// Host receives it // Host receives it
let received = host_rx.receive().await.unwrap(); let received = host_rx.receive().await?;
assert_numbered_message(&received, CommunicationType::Ping, 42, &tm); assert_numbered_message(&received, CommunicationType::Ping, 42, &tm);
// Host sends a response // Host sends a response
let resp = numbered_message(CommunicationType::Pong, 99, &tm); let resp = numbered_message(CommunicationType::Pong, 99, &tm);
host_tx.send(&resp).await.unwrap(); host_tx.send(&resp).await?;
// Client receives it // Client receives it
let client_received = client_rx.receive().await.unwrap(); let client_received = client_rx.receive().await?;
assert_numbered_message(&client_received, CommunicationType::Pong, 99, &tm); assert_numbered_message(&client_received, CommunicationType::Pong, 99, &tm);
// Close both sides // Close both sides
client_tx.close(); client_tx.close();
host_tx.close(); host_tx.close();
Ok(())
} }
#[tokio::test] #[tokio::test]
async fn test_concurrent_messages() { async fn test_concurrent_messages() -> Result<(), Box<dyn std::error::Error>> {
let (_h, client_tx, _client_rx, _host_tx, host_rx) = connected_pair().await; let (_h, client_tx, _client_rx, _host_tx, host_rx) = connected_pair().await?;
let tm = TypeMap::latest(); let tm = TypeMap::latest();
// Send 5 messages in sequence // Send 5 messages in sequence
for i in 0..5u128 { for i in 0..5u128 {
let msg = numbered_message(CommunicationType::Ping, i, &tm); let msg = numbered_message(CommunicationType::Ping, i, &tm);
client_tx.send(&msg).await.unwrap(); client_tx.send(&msg).await?;
} }
// Receive all 5 in order // Receive all 5 in order
for i in 0..5u128 { for i in 0..5u128 {
let received = host_rx.receive().await.unwrap(); let received = host_rx.receive().await?;
assert_numbered_message(&received, CommunicationType::Ping, i, &tm); assert_numbered_message(&received, CommunicationType::Ping, i, &tm);
} }
// Send 3 responses back // Send 3 responses back
for i in 0..3u128 { for i in 0..3u128 {
let msg = numbered_message(CommunicationType::Pong, i * 10, &tm); let msg = numbered_message(CommunicationType::Pong, i * 10, &tm);
client_tx.send(&msg).await.unwrap(); client_tx.send(&msg).await?;
} }
for i in 0..3u128 { for i in 0..3u128 {
let received = host_rx.receive().await.unwrap(); let received = host_rx.receive().await?;
assert_numbered_message(&received, CommunicationType::Pong, i * 10, &tm); assert_numbered_message(&received, CommunicationType::Pong, i * 10, &tm);
} }
client_tx.close(); client_tx.close();
Ok(())
} }
#[tokio::test] #[tokio::test]
async fn test_close_detection() { async fn test_close_detection() -> Result<(), Box<dyn std::error::Error>> {
let (_h, client_tx, _client_rx, _host_tx, host_rx) = connected_pair().await; let (_h, client_tx, _client_rx, _host_tx, host_rx) = connected_pair().await?;
// Send a message then close // Send a message then close
let msg = CommunicationValue::new(CommunicationType::Ping); let msg = CommunicationValue::new(CommunicationType::Ping);
client_tx.send(&msg).await.unwrap(); client_tx.send(&msg).await?;
client_tx.close(); client_tx.close();
// Host should still receive the message // Host should still receive the message
let tm = TypeMap::latest(); let tm = TypeMap::latest();
let received = host_rx.receive().await.unwrap(); let received = host_rx.receive().await?;
assert_eq!(received.get_type(), CommunicationType::Ping.to_id(&tm)); assert_eq!(received.get_type(), CommunicationType::Ping.to_id(&tm));
// Host should get an error or closed signal on next receive // Host should get an error or closed signal on next receive
let result = host_rx.receive().await; let result = host_rx.receive().await;
assert!(result.is_err()); assert!(result.is_err());
Ok(())
} }
#[tokio::test] #[tokio::test]
async fn test_host_shutdown_stops_accepting() { async fn test_host_shutdown_stops_accepting() -> Result<(), Box<dyn std::error::Error>> {
let (cert_pem, key_pem) = generate_self_signed_cert(); let (cert_pem, key_pem) = generate_self_signed_cert();
let mut h = start_test_host(cert_pem.clone(), key_pem).await; let mut h = start_test_host(cert_pem.clone(), key_pem).await?;
let url = format!("https://127.0.0.1:{}", h.local_addr().port()); let url = format!("https://127.0.0.1:{}", h.local_addr().port());
// A connection succeeds while the host is accepting. // A connection succeeds while the host is accepting.
let (_c_tx, _c_rx) = connect(&url, Some(cert_pem.clone()), Policy::default()) let (_c_tx, _c_rx) = connect(&url, Some(cert_pem.clone()), Policy::default()).await?;
.await let _accepted = h.next().await.ok_or("host did not accept connection")?;
.unwrap();
let _accepted = h.next().await.unwrap();
// After shutdown the accept task is aborted and its endpoint is dropped, so // After shutdown the accept task is aborted and its endpoint is dropped, so
// new connections no longer succeed. Guard with a timeout so a hung connect // new connections no longer succeed. Guard with a timeout so a hung connect
@ -173,16 +174,17 @@ async fn test_host_shutdown_stops_accepting() {
matches!(result, Err(_) | Ok(Err(_))), matches!(result, Err(_) | Ok(Err(_))),
"connect should not succeed after host shutdown" "connect should not succeed after host shutdown"
); );
Ok(())
} }
#[tokio::test] #[tokio::test]
async fn test_drop_receiver_keeps_sender_alive() { async fn test_drop_receiver_keeps_sender_alive() -> Result<(), Box<dyn std::error::Error>> {
let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await; let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await?;
// Client sends a message the host receives. // Client sends a message the host receives.
let msg = CommunicationValue::new(CommunicationType::Ping); let msg = CommunicationValue::new(CommunicationType::Ping);
client_tx.send(&msg).await.unwrap(); client_tx.send(&msg).await?;
let _ = host_rx.receive().await.unwrap(); let _ = host_rx.receive().await?;
// Dropping the host Receiver aborts only its accept task; the Sender shares // Dropping the host Receiver aborts only its accept task; the Sender shares
// the same connection and must keep working. // the same connection and must keep working.
@ -191,69 +193,71 @@ async fn test_drop_receiver_keeps_sender_alive() {
let tm = TypeMap::latest(); let tm = TypeMap::latest();
let resp = numbered_message(CommunicationType::Pong, 7, &tm); let resp = numbered_message(CommunicationType::Pong, 7, &tm);
host_tx.send(&resp).await.unwrap(); host_tx.send(&resp).await?;
let got = client_rx.receive().await.unwrap(); let got = client_rx.receive().await?;
assert_numbered_message(&got, CommunicationType::Pong, 7, &tm); assert_numbered_message(&got, CommunicationType::Pong, 7, &tm);
client_tx.close(); client_tx.close();
host_tx.close(); host_tx.close();
Ok(())
} }
#[tokio::test] #[tokio::test]
async fn test_persistent_stream_reopens_after_local_finish() { async fn test_persistent_stream_reopens_after_local_finish() -> Result<(), Box<dyn std::error::Error>> {
let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await; let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await?;
let tm = TypeMap::latest(); let tm = TypeMap::latest();
let msg1 = numbered_message(CommunicationType::Ping, 11, &tm); let msg1 = numbered_message(CommunicationType::Ping, 11, &tm);
client_tx.send(&msg1).await.unwrap(); client_tx.send(&msg1).await?;
let received1 = host_rx.receive().await.unwrap(); let received1 = host_rx.receive().await?;
assert_numbered_message(&received1, CommunicationType::Ping, 11, &tm); assert_numbered_message(&received1, CommunicationType::Ping, 11, &tm);
client_tx.finish_stream().await.unwrap(); client_tx.finish_stream().await?;
let msg2 = numbered_message(CommunicationType::Pong, 22, &tm); let msg2 = numbered_message(CommunicationType::Pong, 22, &tm);
client_tx.send(&msg2).await.unwrap(); client_tx.send(&msg2).await?;
let received2 = host_rx.receive().await.unwrap(); let received2 = host_rx.receive().await?;
assert_numbered_message(&received2, CommunicationType::Pong, 22, &tm); assert_numbered_message(&received2, CommunicationType::Pong, 22, &tm);
client_tx.close(); client_tx.close();
host_tx.close(); host_tx.close();
drop(client_rx); drop(client_rx);
Ok(())
} }
#[tokio::test] #[tokio::test]
async fn test_receiver_backpressure_with_small_queue() { async fn test_receiver_backpressure_with_small_queue() -> Result<(), Box<dyn std::error::Error>> {
let (cert_pem, key_pem) = generate_self_signed_cert(); let (cert_pem, key_pem) = generate_self_signed_cert();
let mut h = start_test_host(cert_pem.clone(), key_pem).await; let mut h = start_test_host(cert_pem.clone(), key_pem).await?;
let url = format!("https://127.0.0.1:{}", h.local_addr().port()); let url = format!("https://127.0.0.1:{}", h.local_addr().port());
let policy = Policy::default().with_receiver_queue_capacity(1); let policy = Policy::default().with_receiver_queue_capacity(1);
let (client_tx, client_rx) = connect(&url, Some(cert_pem), policy).await.unwrap(); let (client_tx, client_rx) = connect(&url, Some(cert_pem), policy).await?;
let (_host_tx, host_rx) = h.next().await.unwrap(); let (_host_tx, host_rx) = h.next().await.ok_or("host did not accept connection")?;
let tm = TypeMap::latest(); let tm = TypeMap::latest();
for i in 0..8u128 { for i in 0..8u128 {
client_tx client_tx
.send(&numbered_message(CommunicationType::Ping, i, &tm)) .send(&numbered_message(CommunicationType::Ping, i, &tm))
.await .await
.unwrap(); ?;
} }
for i in 0..8u128 { for i in 0..8u128 {
let received = tokio::time::timeout(std::time::Duration::from_secs(5), host_rx.receive()) let received = tokio::time::timeout(std::time::Duration::from_secs(5), host_rx.receive())
.await .await?
.unwrap() ?;
.unwrap();
assert_numbered_message(&received, CommunicationType::Ping, i, &tm); assert_numbered_message(&received, CommunicationType::Ping, i, &tm);
} }
client_tx.close(); client_tx.close();
drop(client_rx); drop(client_rx);
h.shutdown(); h.shutdown();
Ok(())
} }
#[tokio::test] #[tokio::test]
async fn test_max_frames_per_stream_enforced() { async fn test_max_frames_per_stream_enforced() -> Result<(), Box<dyn std::error::Error>> {
let (cert_pem, key_pem) = generate_self_signed_cert(); let (cert_pem, key_pem) = generate_self_signed_cert();
let policy = Policy::default().with_max_frames_per_stream(Some(1)); let policy = Policy::default().with_max_frames_per_stream(Some(1));
let mut h = host( let mut h = host(
@ -263,35 +267,33 @@ async fn test_max_frames_per_stream_enforced() {
key_pem, key_pem,
policy, policy,
) )
.await .await?;
.unwrap();
let url = format!("https://127.0.0.1:{}", h.local_addr().port()); let url = format!("https://127.0.0.1:{}", h.local_addr().port());
let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default()) let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default())
.await .await
.unwrap(); ?;
let (_host_tx, host_rx) = h.next().await.unwrap(); let (_host_tx, host_rx) = h.next().await.ok_or("host did not accept connection")?;
let tm = TypeMap::latest(); let tm = TypeMap::latest();
client_tx client_tx
.send(&numbered_message(CommunicationType::Ping, 1, &tm)) .send(&numbered_message(CommunicationType::Ping, 1, &tm))
.await .await?;
.unwrap(); let first = host_rx.receive().await?;
let first = host_rx.receive().await.unwrap();
assert_numbered_message(&first, CommunicationType::Ping, 1, &tm); assert_numbered_message(&first, CommunicationType::Ping, 1, &tm);
client_tx client_tx
.send(&numbered_message(CommunicationType::Ping, 2, &tm)) .send(&numbered_message(CommunicationType::Ping, 2, &tm))
.await .await?;
.unwrap();
let second = host_rx.receive().await; let second = host_rx.receive().await;
assert!(second.is_err(), "stream should be closed after frame limit"); assert!(second.is_err(), "stream should be closed after frame limit");
client_tx.close(); client_tx.close();
h.shutdown(); h.shutdown();
Ok(())
} }
#[tokio::test] #[tokio::test]
async fn test_semaphore_saturation_with_concurrent_streams() { async fn test_semaphore_saturation_with_concurrent_streams() -> Result<(), Box<dyn std::error::Error>> {
let (cert_pem, key_pem) = generate_self_signed_cert(); let (cert_pem, key_pem) = generate_self_signed_cert();
let policy = Policy::default() let policy = Policy::default()
.with_send_mode(mtp_transport::SendMode::SingleStreamPerMessage) .with_send_mode(mtp_transport::SendMode::SingleStreamPerMessage)
@ -304,13 +306,11 @@ async fn test_semaphore_saturation_with_concurrent_streams() {
key_pem, key_pem,
policy, policy,
) )
.await .await?;
.unwrap();
let url = format!("https://127.0.0.1:{}", h.local_addr().port()); let url = format!("https://127.0.0.1:{}", h.local_addr().port());
let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default()) let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default())
.await .await?;
.unwrap(); let (_host_tx, host_rx) = h.next().await.ok_or("host did not accept connection")?;
let (_host_tx, host_rx) = h.next().await.unwrap();
let tm = TypeMap::latest(); let tm = TypeMap::latest();
let mut joins = Vec::new(); let mut joins = Vec::new();
@ -323,17 +323,17 @@ async fn test_semaphore_saturation_with_concurrent_streams() {
tokio::time::sleep(std::time::Duration::from_millis(100)).await; tokio::time::sleep(std::time::Duration::from_millis(100)).await;
for join in joins { for join in joins {
join.await.unwrap().unwrap(); join.await??;
} }
for i in 0..6u128 { for i in 0..6u128 {
let received = tokio::time::timeout(std::time::Duration::from_secs(5), host_rx.receive()) let received = tokio::time::timeout(std::time::Duration::from_secs(5), host_rx.receive())
.await .await?
.unwrap() ?;
.unwrap();
assert_numbered_message(&received, CommunicationType::Ping, i, &tm); assert_numbered_message(&received, CommunicationType::Ping, i, &tm);
} }
client_tx.close(); client_tx.close();
h.shutdown(); h.shutdown();
Ok(())
} }

View file

@ -37,26 +37,30 @@ mod tests {
#[wasm_bindgen_test] #[wasm_bindgen_test]
fn from_codec_error_invalid_encoding() { fn from_codec_error_invalid_encoding() {
let err = from_codec_error(mtp_common::CodecError::InvalidEncoding); let err = from_codec_error(mtp_common::CodecError::InvalidEncoding);
assert!(err.as_string().unwrap().contains("Invalid encoding")); let msg = err.as_string().unwrap_or_default();
assert!(msg.contains("Invalid encoding"));
} }
#[wasm_bindgen_test] #[wasm_bindgen_test]
fn from_codec_error_unknown_version() { fn from_codec_error_unknown_version() {
let err = from_codec_error(mtp_common::CodecError::UnknownVersion); let err = from_codec_error(mtp_common::CodecError::UnknownVersion);
assert!(err.as_string().unwrap().contains("Unknown version")); let msg = err.as_string().unwrap_or_default();
assert!(msg.contains("Unknown version"));
} }
#[wasm_bindgen_test] #[wasm_bindgen_test]
fn from_communication_error_renders() { fn from_communication_error_renders() {
use mtp_common::CommunicationError; use mtp_common::CommunicationError;
let err = from_communication_error(CommunicationError::ConnectionLost); let err = from_communication_error(CommunicationError::ConnectionLost);
assert!(err.as_string().unwrap().contains("Connection terminated")); let msg = err.as_string().unwrap_or_default();
assert!(msg.contains("Connection terminated"));
} }
#[wasm_bindgen_test] #[wasm_bindgen_test]
fn from_crypto_error_renders() { fn from_crypto_error_renders() {
use mtp_crypto::CryptoError; use mtp_crypto::CryptoError;
let err = from_crypto_error(CryptoError::InvalidKeyLength); let err = from_crypto_error(CryptoError::InvalidKeyLength);
assert!(err.as_string().unwrap().contains("invalid key length")); let msg = err.as_string().unwrap_or_default();
assert!(msg.contains("invalid key length"));
} }
} }