From 5f11d476b6cdbde0851088e1087995c9baef6653 Mon Sep 17 00:00:00 2001 From: Alex Emmet <111742636+Alex-Emmet@users.noreply.github.com> Date: Wed, 15 Jul 2026 19:11:01 +0200 Subject: [PATCH] [Clean] safer unwrap & except handling --- client/tests/ping.rs | 40 ++++---- codec/src/communication_value.rs | 59 ++++++----- codec/src/data_value.rs | 119 +++++++++++----------- crypto/src/enc.rs | 20 ++-- crypto/src/keypair.rs | 55 ++++++----- crypto/src/lib.rs | 70 +++++++------ example-type-maps.yaml | 73 +++++++------- example/client/src/main.rs | 11 ++- example/client/src/messages.rs | 12 +-- example/server/src/clients.rs | 7 +- example/server/src/handlers.rs | 58 ++++++++--- example/server/src/main.rs | 34 ++++--- files/src/lib.rs | 31 +++--- host/src/lib.rs | 42 +++++--- transport/src/host.rs | 16 ++- transport/tests/integration.rs | 164 +++++++++++++++---------------- wasm/src/error.rs | 12 ++- 17 files changed, 475 insertions(+), 348 deletions(-) diff --git a/client/tests/ping.rs b/client/tests/ping.rs index e497db9..6addcaa 100644 --- a/client/tests/ping.rs +++ b/client/tests/ping.rs @@ -3,18 +3,17 @@ use std::net::{IpAddr, Ipv4Addr}; use mtp_client::{ClientConfig, MTPClient}; use mtp_host::{HostConfig, MTPHost}; -fn generate_self_signed_cert() -> (Vec, Vec) { - let key_pair = rcgen::KeyPair::generate().unwrap(); - let params = - rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]).unwrap(); - let cert = params.self_signed(&key_pair).unwrap(); +async fn generate_self_signed_cert() -> Result<(Vec, Vec), Box> { + let key_pair = rcgen::KeyPair::generate()?; + let params = rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()])?; + let cert = params.self_signed(&key_pair)?; let cert_pem = cert.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) { - let (cert_pem, key_pem) = generate_self_signed_cert(); +async fn start_host(send_pongs: bool) -> Result<(MTPHost, Vec), Box> { + let (cert_pem, key_pem) = generate_self_signed_cert().await?; let host = MTPHost::new( HostConfig::new( IpAddr::V4(Ipv4Addr::LOCALHOST), @@ -24,14 +23,13 @@ async fn start_host(send_pongs: bool) -> (MTPHost, Vec) { ) .with_pongs(send_pongs), ) - .await - .unwrap(); - (host, cert_pem) + .await?; + Ok((host, cert_pem)) } #[tokio::test] -async fn test_ping_rtt_and_missed_ping_teardown() { - let (mut host, cert_pem) = start_host(true).await; +async fn test_ping_rtt_and_missed_ping_teardown() -> Result<(), Box> { + let (mut host, cert_pem) = start_host(true).await?; let url = format!("https://127.0.0.1:{}", host.local_addr().port()); 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_max_missed_pings(3), ) - .await - .unwrap(); + .await?; - let _accepted = host.accept().await.unwrap().unwrap(); + let _accepted = host.accept().await?; let ping = tokio::time::timeout(std::time::Duration::from_secs(5), async { loop { @@ -53,12 +50,11 @@ async fn test_ping_rtt_and_missed_ping_teardown() { tokio::time::sleep(std::time::Duration::from_millis(10)).await; } }) - .await - .unwrap(); + .await?; 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_client = MTPClient::connect( 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_max_missed_pings(2), ) - .await - .unwrap(); + .await?; - 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 { loop { @@ -82,4 +77,5 @@ async fn test_ping_rtt_and_missed_ping_teardown() { .await; assert!(closed.is_ok(), "client should close after missed pings"); + Ok(()) } diff --git a/codec/src/communication_value.rs b/codec/src/communication_value.rs index 96acccc..be64cda 100644 --- a/codec/src/communication_value.rs +++ b/codec/src/communication_value.rs @@ -775,65 +775,67 @@ mod tests { use super::*; use crate::data_value::DataValue; - fn roundtrip(cv: CommunicationValue) -> CommunicationValue { - let bytes = cv.to_bytes().expect("encode failed"); - let decoded = CommunicationValue::from_bytes(&bytes).expect("failed to deserialize"); - let bytes2 = decoded.to_bytes().expect("encode failed"); + fn roundtrip(cv: CommunicationValue) -> Result> { + let bytes = cv.to_bytes()?; + let decoded = CommunicationValue::from_bytes(&bytes)?; + let bytes2 = decoded.to_bytes()?; assert_eq!(bytes, bytes2); - decoded + Ok(decoded) } #[test] - fn test_flags_and_order_without_optional() { + fn test_flags_and_order_without_optional() -> Result<(), Box> { 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]... assert!(bytes.len() >= 7); let mut c = Cursor::new(bytes.as_slice()); - let total_len = c.read_u32::().expect("read len"); + let total_len = c.read_u32::()?; assert_eq!(total_len as usize + 4, bytes.len()); - let typ = c.read_u16::().expect("read type"); + let typ = c.read_u16::()?; assert_eq!(typ, 12); - let flags = c.read_u8().expect("read flags"); + let flags = c.read_u8()?; assert_eq!(flags & 0b0000_0111, 0); + Ok(()) } #[test] - fn test_flags_and_order_with_all_optional() { + fn test_flags_and_order_with_all_optional() -> Result<(), Box> { let cv = CommunicationValue::new(CommunicationType::ErrorBadVersion) .with_id(0xAABBCCDD) .with_sender(0x0000_1122_3344_5566) .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 total_len = c.read_u32::().expect("len"); + let total_len = c.read_u32::()?; assert_eq!(total_len as usize + 4, bytes.len()); - let typ = c.read_u16::().expect("read type"); + let typ = c.read_u16::()?; assert_eq!(typ, 13); - let flags = c.read_u8().expect("read flags"); + let flags = c.read_u8()?; assert_eq!(flags & 0b0000_0111, 0b0000_0111); - let id = c.read_u32::().expect("id"); + let id = c.read_u32::()?; assert_eq!(id, 0xAABBCCDD); 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]); 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]); + Ok(()) } #[test] - fn test_roundtrip_complex() { + fn test_roundtrip_complex() -> Result<(), Box> { let tm = TypeMap::latest(); let cv = CommunicationValue::new(CommunicationType::Disconnect) .with_id(1234) @@ -847,7 +849,7 @@ mod tests { 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_sender(), 111); @@ -861,6 +863,7 @@ mod tests { decoded.get_data(DataType::ClientNonce), &DataValue::SignedNumber(42) ); + Ok(()) } #[test] @@ -873,7 +876,7 @@ mod tests { #[cfg(feature = "crypto")] #[test] - fn test_sign_verify_frame_roundtrip() { + fn test_sign_verify_frame_roundtrip() -> Result<(), Box> { use mtp_crypto::{Ed25519Signer, SigAlgorithm}; let (signer, sk, _pk) = Ed25519Signer::generate(); @@ -887,18 +890,20 @@ mod tests { assert!(cv.sign_frame(SigAlgorithm::ED25519, &signer).is_some()); // Same in-memory value verifies (FLAG_SIGNED forced on both sides). - let verifier = Ed25519Signer::new(&sk).unwrap(); + let verifier = Ed25519Signer::new(&sk)?; assert!(cv.verify_frame(&verifier).is_ok()); // Survives a wire round-trip. - let bytes = cv.to_bytes().expect("encode failed"); - let decoded = CommunicationValue::from_bytes(&bytes).expect("decode failed"); + let bytes = cv.to_bytes()?; + let decoded = CommunicationValue::from_bytes(&bytes)?; assert!(decoded.verify_frame(&verifier).is_ok()); + + Ok(()) } #[cfg(feature = "crypto")] #[test] - fn test_verify_frame_wrong_key_fails() { + fn test_verify_frame_wrong_key_fails() -> Result<(), Box> { use mtp_crypto::{Ed25519Signer, SigAlgorithm}; let (signer, _, _) = Ed25519Signer::generate(); @@ -908,7 +913,9 @@ mod tests { .add_typed_default(DataType::PqSignature, DataValue::UnsignedNumber(42)); 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()); + + Ok(()) } } diff --git a/codec/src/data_value.rs b/codec/src/data_value.rs index 8b79d75..65f0f93 100644 --- a/codec/src/data_value.rs +++ b/codec/src/data_value.rs @@ -1200,27 +1200,30 @@ impl TryFrom for Vec { mod tests { use super::*; - fn container_roundtrip(values: Vec<(DataTypeId, DataValue)>) { + fn container_roundtrip(values: Vec<(DataTypeId, DataValue)>) -> Result<(), Box> { let dv = DataValue::Container(values.clone()); - let bytes = dv.to_bytes().expect("encode failed"); - let decoded = DataValue::from_bytes(&bytes).expect("roundtrip failed"); + let bytes = dv.to_bytes()?; + let decoded = DataValue::from_bytes(&bytes).ok_or("roundtrip failed")?; assert_eq!(dv, decoded, "container roundtrip mismatch"); + Ok(()) } - fn array_roundtrip(values: Vec) { + fn array_roundtrip(values: Vec) -> Result<(), Box> { let dv = DataValue::Array(values.clone()); - let bytes = dv.to_bytes().expect("encode failed"); - let decoded = DataValue::from_bytes(&bytes).expect("roundtrip failed"); + let bytes = dv.to_bytes()?; + let decoded = DataValue::from_bytes(&bytes).ok_or("roundtrip failed")?; assert_eq!(dv, decoded, "array roundtrip mismatch"); + Ok(()) } #[test] - fn test_bool_in_container() { + fn test_bool_in_container() -> Result<(), Box> { let tm = TypeMap::latest(); container_roundtrip(vec![ (DataType::Id.to_id(&tm), DataValue::BoolTrue), (DataType::ClientNonce.to_id(&tm), DataValue::BoolFalse), - ]); + ])?; + Ok(()) } #[test] @@ -1239,7 +1242,7 @@ mod tests { } #[test] - fn test_signed_number_in_container() { + fn test_signed_number_in_container() -> Result<(), Box> { let tm = TypeMap::latest(); container_roundtrip(vec![ (DataType::Version.to_id(&tm), DataValue::SignedNumber(0)), @@ -1256,11 +1259,12 @@ mod tests { DataType::PublicKeys.to_id(&tm), DataValue::SignedNumber(i128::MIN), ), - ]); + ])?; + Ok(()) } #[test] - fn test_unsigned_number_in_container() { + fn test_unsigned_number_in_container() -> Result<(), Box> { let tm = TypeMap::latest(); container_roundtrip(vec![ (DataType::Version.to_id(&tm), DataValue::UnsignedNumber(0)), @@ -1269,11 +1273,12 @@ mod tests { DataType::ClientNonce.to_id(&tm), DataValue::UnsignedNumber(u128::MAX), ), - ]); + ])?; + Ok(()) } #[test] - fn test_float_in_container() { + fn test_float_in_container() -> Result<(), Box> { let tm = TypeMap::latest(); container_roundtrip(vec![ (DataType::Version.to_id(&tm), DataValue::Float(0, 0)), @@ -1282,11 +1287,12 @@ mod tests { DataType::ClientNonce.to_id(&tm), DataValue::Float(255, 4294967295), ), - ]); + ])?; + Ok(()) } #[test] - fn test_str_in_container() { + fn test_str_in_container() -> Result<(), Box> { let tm = TypeMap::latest(); container_roundtrip(vec![ (DataType::Version.to_id(&tm), DataValue::Str(String::new())), @@ -1295,11 +1301,12 @@ mod tests { DataType::ClientNonce.to_id(&tm), DataValue::Str("a".repeat(1000)), ), - ]); + ])?; + Ok(()) } #[test] - fn test_bytes_in_container() { + fn test_bytes_in_container() -> Result<(), Box> { let tm = TypeMap::latest(); container_roundtrip(vec![ (DataType::Version.to_id(&tm), DataValue::Bytes(vec![])), @@ -1311,40 +1318,45 @@ mod tests { DataType::ClientNonce.to_id(&tm), DataValue::Bytes(vec![0x42; 100]), ), - ]); + ])?; + Ok(()) } #[test] - fn test_null_in_container() { + fn test_null_in_container() -> Result<(), Box> { 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] - fn test_array_non_empty_roundtrip() { + fn test_array_non_empty_roundtrip() -> Result<(), Box> { array_roundtrip(vec![ DataValue::BoolTrue, DataValue::SignedNumber(42), DataValue::Str("hello".to_string()), DataValue::Null, - ]); + ])?; + Ok(()) } #[test] - fn test_array_nested_roundtrip() { + fn test_array_nested_roundtrip() -> Result<(), Box> { array_roundtrip(vec![ DataValue::Array(vec![DataValue::BoolTrue, DataValue::BoolFalse]), DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]), - ]); + ])?; + Ok(()) } #[test] - fn test_container_empty_roundtrip() { - container_roundtrip(vec![]); + fn test_container_empty_roundtrip() -> Result<(), Box> { + container_roundtrip(vec![])?; + Ok(()) } #[test] - fn test_container_mixed_roundtrip() { + fn test_container_mixed_roundtrip() -> Result<(), Box> { let tm = TypeMap::latest(); container_roundtrip(vec![ (DataType::Version.to_id(&tm), DataValue::BoolTrue), @@ -1358,11 +1370,12 @@ mod tests { DataValue::UnsignedNumber(u128::MAX), ), (DataType::PublicKeys.to_id(&tm), DataValue::Null), - ]); + ])?; + Ok(()) } #[test] - fn test_container_nested_roundtrip() { + fn test_container_nested_roundtrip() -> Result<(), Box> { let tm = TypeMap::latest(); container_roundtrip(vec![ ( @@ -1373,7 +1386,8 @@ mod tests { DataType::Id.to_id(&tm), DataValue::Array(vec![DataValue::SignedNumber(1), DataValue::SignedNumber(2)]), ), - ]); + ])?; + Ok(()) } #[test] @@ -1592,24 +1606,16 @@ mod tests { } #[test] - fn test_try_from_ok() { - assert!(bool::try_from(DataValue::BoolTrue).unwrap()); - assert!(!bool::try_from(DataValue::BoolFalse).unwrap()); - assert_eq!( - String::try_from(DataValue::Str("hi".to_string())).unwrap(), - "hi" - ); - assert_eq!(i128::try_from(DataValue::SignedNumber(-1)).unwrap(), -1i128); - assert_eq!(i64::try_from(DataValue::SignedNumber(10)).unwrap(), 10i64); - assert_eq!( - u128::try_from(DataValue::UnsignedNumber(99)).unwrap(), - 99u128 - ); - assert_eq!(u64::try_from(DataValue::UnsignedNumber(7)).unwrap(), 7u64); - assert_eq!( - Vec::::try_from(DataValue::Bytes(vec![0xAB])).unwrap(), - vec![0xABu8] - ); + fn test_try_from_ok() -> Result<(), Box> { + assert!(bool::try_from(DataValue::BoolTrue)?); + assert!(!bool::try_from(DataValue::BoolFalse)?); + assert_eq!(String::try_from(DataValue::Str("hi".to_string()))?, "hi"); + assert_eq!(i128::try_from(DataValue::SignedNumber(-1))?, -1i128); + assert_eq!(i64::try_from(DataValue::SignedNumber(10))?, 10i64); + assert_eq!(u128::try_from(DataValue::UnsignedNumber(99))?, 99u128); + assert_eq!(u64::try_from(DataValue::UnsignedNumber(7))?, 7u64); + assert_eq!(Vec::::try_from(DataValue::Bytes(vec![0xAB]))?, vec![0xABu8]); + Ok(()) } #[test] @@ -1632,7 +1638,7 @@ mod tests { #[cfg(feature = "crypto")] #[test] - fn test_encrypt_decrypt_container_roundtrip() { + fn test_encrypt_decrypt_container_roundtrip() -> Result<(), Box> { use mtp_crypto::{EncryptionType, Keyring}; let tm = TypeMap::latest(); let keyring = Keyring::generate(); @@ -1655,8 +1661,9 @@ mod tests { assert!(dv.decrypt_into_container(&keyring, b"aad").is_some()); 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); + Ok(()) } #[cfg(feature = "crypto")] @@ -1725,7 +1732,7 @@ mod tests { #[cfg(feature = "crypto")] #[test] - fn test_sign_verify_container_roundtrip() { + fn test_sign_verify_container_roundtrip() -> Result<(), Box> { use mtp_crypto::{Ed25519Signer, EncryptionType, Keyring, SigAlgorithm}; let tm = TypeMap::latest(); @@ -1755,23 +1762,24 @@ mod tests { ); 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!(matches!(dv, DataValue::Container(_))); - let entries = dv.as_container().unwrap(); + let entries = dv.as_container().ok_or("expected container")?; assert_eq!(entries.len(), 1); + Ok(()) } #[cfg(feature = "crypto")] #[test] - fn test_sign_container_wrong_key_fails() { + fn test_sign_container_wrong_key_fails() -> Result<(), Box> { use mtp_crypto::{Ed25519Signer, SigAlgorithm}; let tm = TypeMap::latest(); let (signer, _, _) = 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![( DataType::Version.to_id(&tm), @@ -1780,5 +1788,6 @@ mod tests { assert!(dv.sign_container(SigAlgorithm::ED25519, &signer).is_some()); assert!(dv.verify_into_container(&wrong_verifier).is_none()); + Ok(()) } } diff --git a/crypto/src/enc.rs b/crypto/src/enc.rs index 3b7a44b..e03e716 100644 --- a/crypto/src/enc.rs +++ b/crypto/src/enc.rs @@ -190,24 +190,24 @@ mod tests { #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] #[test] - fn encrypt_for_roundtrip() { + fn encrypt_for_roundtrip() -> Result<(), CryptoError> { let kr = Keyring::generate(); let blob = encrypt_for( EncryptionType::MlKemChaCha20Poly1305, &kr.public_key_bundle(), b"secret payload", b"aad", - ) - .unwrap(); + )?; assert_eq!(blob[0], EncryptionType::ML_KEM_CHACHA20POLY1305); - let pt = decrypt_with(&blob, &kr, b"aad").unwrap(); + let pt = decrypt_with(&blob, &kr, b"aad")?; assert_eq!(pt, b"secret payload"); + Ok(()) } #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] #[test] - fn decrypt_with_wrong_keyring_fails() { + fn decrypt_with_wrong_keyring_fails() -> Result<(), CryptoError> { let kr = Keyring::generate(); let other = Keyring::generate(); let blob = encrypt_for( @@ -215,23 +215,23 @@ mod tests { &kr.public_key_bundle(), b"secret", b"aad", - ) - .unwrap(); + )?; assert!(decrypt_with(&blob, &other, b"aad").is_err()); + Ok(()) } #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] #[test] - fn decrypt_with_wrong_aad_fails() { + fn decrypt_with_wrong_aad_fails() -> Result<(), CryptoError> { let kr = Keyring::generate(); let blob = encrypt_for( EncryptionType::MlKemChaCha20Poly1305, &kr.public_key_bundle(), b"secret", b"right", - ) - .unwrap(); + )?; assert!(decrypt_with(&blob, &kr, b"wrong").is_err()); + Ok(()) } #[cfg(all(feature = "mlkem-tls", feature = "hkdf", feature = "chacha20poly1305"))] diff --git a/crypto/src/keypair.rs b/crypto/src/keypair.rs index 3454cc9..350a7fc 100644 --- a/crypto/src/keypair.rs +++ b/crypto/src/keypair.rs @@ -498,13 +498,14 @@ impl Keyring { use crate::error::CryptoError; let mut offset = 0; let read_key = |offset: &mut usize| -> Result, CryptoError> { - let len = u16::from_be_bytes( - bytes - .get(*offset..*offset + 2) - .ok_or(CryptoError::InvalidKeyLength)? - .try_into() - .expect("slice is 2 bytes, verified above"), - ) as usize; + let slice = bytes + .get(*offset..*offset + 2) + .ok_or(CryptoError::InvalidKeyLength)?; + let len = if let Ok(arr) = <[u8; 2]>::try_from(slice) { + u16::from_be_bytes(arr) + } else { + return Err(CryptoError::InvalidKeyLength); + } as usize; *offset += 2; let key = bytes .get(*offset..*offset + len) @@ -712,14 +713,14 @@ mod tests { use super::*; #[test] - fn public_key_bundle_roundtrip() { + fn public_key_bundle_roundtrip() -> Result<(), Box> { let kem = KemPublicKey::new(vec![1u8; 32]); let pq = SignaturePqPublicKey::new(vec![2u8; 64]); let cl = SignaturePublicKey::new(vec![3u8; 32]); let bundle = PublicKeyBundle::new(kem, pq, cl); let bytes = bundle.as_bytes(); - let recovered = PublicKeyBundle::from_bytes(&bytes).unwrap(); + let recovered = PublicKeyBundle::from_bytes(&bytes)?; assert_eq!( bundle.kem_public_key.as_bytes(), @@ -733,22 +734,24 @@ mod tests { bundle.sig_cl_public_key.as_bytes(), recovered.sig_cl_public_key.as_bytes() ); + Ok(()) } #[test] - fn public_key_bundle_try_from_roundtrip() { + fn public_key_bundle_try_from_roundtrip() -> Result<(), Box> { let bundle = PublicKeyBundle::new( KemPublicKey::new(vec![0xABu8; 48]), SignaturePqPublicKey::new(vec![0xCDu8; 96]), SignaturePublicKey::new(vec![0xEFu8; 32]), ); let bytes: Vec = 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()); + Ok(()) } #[test] - fn keyring_roundtrip() { + fn keyring_roundtrip() -> Result<(), Box> { let keyring = Keyring::new( KemPublicKey::new(vec![1u8; 32]), KemPrivateKey::new(vec![2u8; 32]), @@ -758,7 +761,7 @@ mod tests { SignaturePrivateKey::new(vec![6u8; 32]), ); let bytes = keyring.to_bytes(); - let recovered = Keyring::from_bytes(&bytes).unwrap(); + let recovered = Keyring::from_bytes(&bytes)?; assert_eq!( keyring.kem_public_key.as_bytes(), recovered.kem_public_key.as_bytes() @@ -771,10 +774,11 @@ mod tests { keyring.sig_cl_public_key.as_bytes(), recovered.sig_cl_public_key.as_bytes() ); + Ok(()) } #[test] - fn keyring_try_from_roundtrip() { + fn keyring_try_from_roundtrip() -> Result<(), Box> { let keyring = Keyring::new( KemPublicKey::new(vec![0u8; 16]), KemPrivateKey::new(vec![1u8; 16]), @@ -784,21 +788,23 @@ mod tests { SignaturePrivateKey::new(vec![5u8; 16]), ); let bytes: Vec = 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()); + Ok(()) } #[test] - fn hex_roundtrip() { + fn hex_roundtrip() -> Result<(), Box> { let key = KemPublicKey::new(vec![0xDE, 0xAD, 0xBE, 0xEF]); let hex = key.to_hex(); 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()); + Ok(()) } #[test] - fn keyring_hex_roundtrip() { + fn keyring_hex_roundtrip() -> Result<(), Box> { let keyring = Keyring::new( KemPublicKey::new(vec![1u8; 16]), KemPrivateKey::new(vec![2u8; 16]), @@ -808,12 +814,13 @@ mod tests { SignaturePrivateKey::new(vec![6u8; 16]), ); 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()); + Ok(()) } #[test] - fn keyring_base64_roundtrip() { + fn keyring_base64_roundtrip() -> Result<(), Box> { let keyring = Keyring::new( KemPublicKey::new(vec![1u8; 16]), KemPrivateKey::new(vec![2u8; 16]), @@ -823,20 +830,22 @@ mod tests { SignaturePrivateKey::new(vec![6u8; 16]), ); 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()); + Ok(()) } #[test] - fn public_key_bundle_base64_roundtrip() { + fn public_key_bundle_base64_roundtrip() -> Result<(), Box> { let bundle = PublicKeyBundle::new( KemPublicKey::new(vec![1u8; 32]), SignaturePqPublicKey::new(vec![2u8; 64]), SignaturePublicKey::new(vec![3u8; 32]), ); 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()); + Ok(()) } #[test] diff --git a/crypto/src/lib.rs b/crypto/src/lib.rs index d093a7a..df87616 100644 --- a/crypto/src/lib.rs +++ b/crypto/src/lib.rs @@ -69,13 +69,14 @@ mod tests { #[cfg(feature = "chacha20poly1305")] #[test] - fn aead_encrypt_decrypt() { + fn aead_encrypt_decrypt() -> Result<(), CryptoError> { use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305}; let key = [0xAB; 32]; let cipher = ChaCha20Poly1305::new(key); - let ct = cipher.encrypt(b"hello world", b"aad").unwrap(); - let pt = cipher.decrypt(&ct, b"aad").unwrap(); + let ct = cipher.encrypt(b"hello world", b"aad")?; + let pt = cipher.decrypt(&ct, b"aad")?; assert_eq!(pt, b"hello world"); + Ok(()) } #[cfg(feature = "chacha20poly1305")] @@ -84,7 +85,7 @@ mod tests { use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305}; let cipher_a = ChaCha20Poly1305::new([0xAB; 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()); } @@ -93,7 +94,9 @@ mod tests { fn aead_wrong_aad_fails() { use crate::aead::{AeadDecrypt, AeadEncrypt, ChaCha20Poly1305}; 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()); } @@ -102,12 +105,12 @@ mod tests { fn ed25519_sign_verify() { let (signer, sk, pk) = Ed25519Signer::generate(); let msg = b"test message"; - let sig = signer.sign(msg).unwrap(); - signer.verify(msg, &sig).unwrap(); - verify_ed25519(&pk, msg, &sig).unwrap(); + let sig = signer.sign(msg).expect("signing should succeed"); + signer.verify(msg, &sig).expect("verification should succeed"); + verify_ed25519(&pk, msg, &sig).expect("verification should succeed"); - let loaded = Ed25519Signer::new(&sk).unwrap(); - loaded.verify(msg, &sig).unwrap(); + let loaded = Ed25519Signer::new(&sk).expect("signer loading should succeed"); + loaded.verify(msg, &sig).expect("verification should succeed"); } #[cfg(feature = "ed25519-dalek")] @@ -115,7 +118,7 @@ mod tests { fn ed25519_wrong_sig_fails() { let (signer, _, pk) = Ed25519Signer::generate(); 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()); } @@ -124,12 +127,12 @@ mod tests { fn mldsa_sign_verify() { let (signer, sk, pk) = MlDsaSigner::generate(); let msg = b"test message"; - let sig = signer.sign(msg).unwrap(); - signer.verify(msg, &sig).unwrap(); - verify_ml_dsa(&pk, msg, &sig).unwrap(); + let sig = signer.sign(msg).expect("signing should succeed"); + signer.verify(msg, &sig).expect("verification should succeed"); + verify_ml_dsa(&pk, msg, &sig).expect("verification should succeed"); - let loaded = MlDsaSigner::new(&sk, &pk).unwrap(); - loaded.verify(msg, &sig).unwrap(); + let loaded = MlDsaSigner::new(&sk, &pk).expect("signer loading should succeed"); + loaded.verify(msg, &sig).expect("verification should succeed"); } #[cfg(feature = "ml-dsa")] @@ -137,7 +140,7 @@ mod tests { fn mldsa_wrong_sig_fails() { let (signer, _, pk) = MlDsaSigner::generate(); 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()); } @@ -148,9 +151,10 @@ mod tests { let (ed_signer, _, _) = Ed25519Signer::generate(); let (ml_signer, _, _) = MlDsaSigner::generate(); - let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg").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") - .unwrap(); + .expect("dual verification should succeed"); } #[cfg(all(feature = "ed25519-dalek", feature = "ml-dsa"))] @@ -160,7 +164,8 @@ mod tests { let (ed_signer, _, _) = Ed25519Signer::generate(); let (ml_signer, _, _) = MlDsaSigner::generate(); - let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg").unwrap(); + let dual = sign_dual(ed_signer.signing_key(), ml_signer.signing_key(), b"msg") + .expect("dual signing should succeed"); assert!( dual.verify( ed_signer.verifying_key(), @@ -174,18 +179,22 @@ mod tests { #[cfg(feature = "hkdf")] #[test] 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); - 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); } #[cfg(feature = "hkdf")] #[test] fn hkdf_different_info_different_key() { - let a = derive_encryption_key(b"ikm", b"salt", b"info-a").unwrap(); - let b = derive_encryption_key(b"ikm", b"salt", b"info-b").unwrap(); + let a = derive_encryption_key(b"ikm", b"salt", b"info-a") + .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); } @@ -255,7 +264,7 @@ mod tests { fn keyring_serialize_roundtrip() { let kr = Keyring::generate(); 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!( kr.kem_public_key.as_bytes(), loaded.kem_public_key.as_bytes() @@ -276,7 +285,8 @@ mod tests { let kr = Keyring::generate(); let bundle = kr.public_key_bundle(); 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!( bundle.kem_public_key.as_bytes(), loaded.kem_public_key.as_bytes() @@ -295,8 +305,8 @@ mod tests { #[test] fn hybrid_kem_roundtrip() { let (sk, pk) = HybridKem::generate_keypair(); - let enc = HybridKem::encapsulate(&pk).unwrap(); - let ss = HybridKem::decapsulate(&sk, &enc.ciphertext).unwrap(); + let enc = HybridKem::encapsulate(&pk).expect("encapsulation should succeed"); + let ss = HybridKem::decapsulate(&sk, &enc.ciphertext).expect("decapsulation should succeed"); assert_eq!(enc.shared_secret, ss); } @@ -309,8 +319,8 @@ mod tests { let kr = Keyring::generate(); let entities = vec![kr.public_key_bundle()]; let msg = b"secret data"; - let ct = encrypt_multi(msg, b"aad", &entities).unwrap(); - let pt = decrypt_multi(&ct, b"aad", &kr).unwrap(); + let ct = encrypt_multi(msg, b"aad", &entities).expect("multi encrypt should succeed"); + let pt = decrypt_multi(&ct, b"aad", &kr).expect("multi decrypt should succeed"); assert_eq!(pt, msg); } } diff --git a/example-type-maps.yaml b/example-type-maps.yaml index 0041f48..16cd948 100644 --- a/example-type-maps.yaml +++ b/example-type-maps.yaml @@ -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 # Fixed CommunicationType markers are: -# Error: 0 -# ErrorParsing: 1 -# ErrorBadVersion: 2 -# Disconnect: 3 -# Redirect: 4 -# Shutdown: 5 -# BadRequest: 6 -# Unauthorized: 7 -# Forbidden: 8 -# NotFound: 9 -# TooManyRequests: 10 -# InternalServerError: 11 -# BadGateway: 12 -# ServiceUnavailable: 13 -# GatewayTimeout: 14 -# Identification: 15 -# IdentificationResponse: 16 -# Register: 17 -# RegisterResponse: 18 -# Ping: 19 -# Pong: 20 +# Identification: 0 +# IdentificationResponse: 1 +# Register: 2 +# RegisterResponse: 3 +# Challenge: 4 +# ChallengeResponse: 5 +# Ping: 6 +# Pong: 7 +# Disconnect: 8 +# Redirect: 9 +# Shutdown: 10 +# Error: 11 +# ErrorParsing: 12 +# ErrorBadVersion: 13 +# BadRequest: 14 +# Unauthorized: 15 +# Forbidden: 16 +# NotFound: 17 +# TooManyRequests: 18 +# InternalServerError: 19 +# BadGateway: 20 +# ServiceUnavailable: 21 +# GatewayTimeout: 22 +# PipeRequest: 23 +# PipeResponse: 24 +# PipeAbort: 25 # # Fixed Data Type markers are: -# Error: 0 -# ErrorParsing: 1 -# ErrorMessage: 2 -# Version: 3 -# Description: 4 -# Timestamp: 5 -# Id: 6 -# ClientNonce: 7 -# ServerNonce: 8 -# PublicKeys: 9 -# Signature: 10 -# Connected: 11 +# Version: 0 +# id: 1 +# ClientNonce: 2 +# ServerNonce: 3 +# PublicKeys: 4 +# Signature: 5 +# PqSignature: 6 +# Description: 7 +# Connected: 8 +# Timestamp: 9 +# Error: 10 +# ErrorParsing: 11 +# ErrorMessage: 12 +# Accepted: 13, # # If a Type can't be used it will be mapped to 0 diff --git a/example/client/src/main.rs b/example/client/src/main.rs index 23a0174..037449e 100644 --- a/example/client/src/main.rs +++ b/example/client/src/main.rs @@ -26,8 +26,15 @@ async fn main() -> Result<(), Box> { "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") - .expect("Missing host.mpkb: run the server first to export it"); + let host_public_key = match load_public_key_bundle("host.mpkb") { + 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 ..."); diff --git a/example/client/src/messages.rs b/example/client/src/messages.rs index 08501fa..4bab1fb 100644 --- a/example/client/src/messages.rs +++ b/example/client/src/messages.rs @@ -7,11 +7,10 @@ pub fn build_demo_message( client_id: u64, keyring: &Keyring, server_bundle: &PublicKeyBundle, -) -> CommunicationValue { +) -> Result> { // Encrypt to the server's KEM public key; the server decrypts with its keyring. let enc_type = EncryptionType::MlKemChaCha20Poly1305; - let signer = - Ed25519Signer::new(&keyring.sig_cl_secret_key).expect("Ed25519 signer from keyring"); + let signer = Ed25519Signer::new(&keyring.sig_cl_secret_key)?; let tm = TypeMap::latest(); @@ -52,8 +51,7 @@ pub fn build_demo_message( ); let timestamp = std::time::SystemTime::now() - .duration_since(std::time::UNIX_EPOCH) - .unwrap() + .duration_since(std::time::UNIX_EPOCH)? .as_secs(); 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::SecurePayload, dv_sec) .with_sender(client_id); - msg + Ok(msg) } pub async fn send_and_receive( @@ -92,7 +90,7 @@ pub async fn send_and_receive( keyring: &Keyring, server_bundle: &PublicKeyBundle, ) -> Result<(), Box> { - 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}"); conn.sender.send(&msg).await?; diff --git a/example/server/src/clients.rs b/example/server/src/clients.rs index d9ad714..4dfed94 100644 --- a/example/server/src/clients.rs +++ b/example/server/src/clients.rs @@ -19,7 +19,12 @@ pub fn load_client_db( Err(_) => HashMap::new(), }; let clients: Arc>> = 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)); Ok((clients, next_id)) } diff --git a/example/server/src/handlers.rs b/example/server/src/handlers.rs index 789cd4a..3c08ffa 100644 --- a/example/server/src/handlers.rs +++ b/example/server/src/handlers.rs @@ -17,17 +17,47 @@ pub fn process_and_respond( tm: &TypeMap, client_pk: Option<&mtp::crypto::PublicKeyBundle>, host_keyring: &Keyring, -) -> CommunicationValue { - let desc_id = DataTypeId(tm.data_id_enum(DataType::Description).unwrap()); - let ts_id = DataTypeId(tm.data_id_enum(DataType::Timestamp).unwrap()); - let data_id = DataTypeId(tm.data_id_enum(DataType::Data).unwrap()); - let flags_id = DataTypeId(tm.data_id_enum(DataType::Flags).unwrap()); - let value_id = DataTypeId(tm.data_id_enum(DataType::Value).unwrap()); - let bin_id = DataTypeId(tm.data_id_enum(DataType::BinaryData).unwrap()); - let items_id = DataTypeId(tm.data_id_enum(DataType::Items).unwrap()); - let _enc_id = DataTypeId(tm.data_id_enum(DataType::EncryptedPayload).unwrap()); - let _sig_id = DataTypeId(tm.data_id_enum(DataType::SignedPayload).unwrap()); - let _secure_id = DataTypeId(tm.data_id_enum(DataType::SecurePayload).unwrap()); +) -> Result { + let desc_id = DataTypeId( + tm.data_id_enum(DataType::Description) + .ok_or("missing Description type mapping")?, + ); + let ts_id = DataTypeId( + tm.data_id_enum(DataType::Timestamp) + .ok_or("missing Timestamp type mapping")?, + ); + let data_id = DataTypeId( + 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 timestamp = msg.get_data(DataType::Timestamp); @@ -118,10 +148,10 @@ pub fn process_and_respond( let now = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) - .unwrap() + .map_err(|e| e.to_string())? .as_secs(); - CommunicationValue::from_comm(CommunicationType::Pong, tm) + Ok(CommunicationValue::from_comm(CommunicationType::Pong, tm) .add_data(desc_id, description.clone()) .add_data(ts_id, DataValue::UnsignedNumber(now as u128)) .add_data( @@ -134,5 +164,5 @@ pub fn process_and_respond( .add_data(flags_id, flags.clone()) .add_data(value_id, value.clone()) .add_data(bin_id, binary.clone()) - .add_data(items_id, items.clone()) + .add_data(items_id, items.clone())) } diff --git a/example/server/src/main.rs b/example/server/src/main.rs index 0cb82db..40b064a 100644 --- a/example/server/src/main.rs +++ b/example/server/src/main.rs @@ -62,7 +62,10 @@ async fn handle_pipe_loopback( tokio::io::AsyncWriteExt::write_all(&mut writer, &buf[..n]).await?; } writer.finish().await?; - println!(" [loopback] Pipe {pipe_id} loopback complete ({} bytes)", total); + println!( + " [loopback] Pipe {pipe_id} loopback complete ({} bytes)", + total + ); } None => { eprintln!(" [loopback] Return pipe denied by client for pipe {pipe_id}"); @@ -83,10 +86,13 @@ async fn main() -> Result<(), Box> { let (_host_id, host_keyring) = keys::load_or_generate_host_keys("host.mk")?; keys::export_host_public_keys(&host_keyring)?; - let decrypt_keyring = Arc::new( - mtp::crypto::Keyring::from_bytes(&host_keyring.to_bytes()) - .expect("re-load host keyring for decryption"), - ); + let decrypt_keyring = Arc::new(match mtp::crypto::Keyring::from_bytes(&host_keyring.to_bytes()) + { + 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")?; @@ -94,7 +100,7 @@ async fn main() -> Result<(), Box> { let get_existing_user = move |id: u64, _description: Option| { let clients = clients_for_get.clone(); Box::pin(async move { - let result = clients.lock().unwrap().get(&id).cloned(); + let result = clients.lock()?.get(&id).cloned(); if result.is_some() { println!("Auth lookup: client ID {id} found"); } else { @@ -113,8 +119,8 @@ async fn main() -> Result<(), Box> { let nid_arc = next_id_for_register.clone(); let path = clients_path.clone(); Box::pin(async move { - let mut db = db_arc.lock().unwrap(); - let mut nid = nid_arc.lock().unwrap(); + let mut db = db_arc.lock()?; + let mut nid = nid_arc.lock()?; let id = *nid; *nid += 1; db.insert(id, bundle); @@ -160,7 +166,7 @@ async fn main() -> Result<(), Box> { ); 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 ..."); loop { @@ -185,12 +191,18 @@ async fn main() -> Result<(), Box> { match msg { Ok(msg) => { println!("Received: {msg}"); - let response = handlers::process_and_respond( + let response = match handlers::process_and_respond( &msg, tm, conn.client_public_key.as_ref(), &decrypt_keyring, - ); + ) { + Ok(response) => response, + Err(e) => { + eprintln!("Failed to build response: {e}"); + continue; + } + }; println!("Sending: {response}"); if let Err(e) = conn.sender.send(&response).await { eprintln!("Send error: {e}"); diff --git a/files/src/lib.rs b/files/src/lib.rs index 1e40f26..79ad20c 100644 --- a/files/src/lib.rs +++ b/files/src/lib.rs @@ -146,52 +146,57 @@ mod tests { } #[test] - fn keyring_save_load_roundtrip() { + fn keyring_save_load_roundtrip() -> Result<(), Box> { let path = temp_path(KEYRING_EXTENSION); let keyring = sample_keyring(); - save_keyring(&keyring, &path).unwrap(); - let loaded = load_keyring(&path).unwrap(); + save_keyring(&keyring, &path)?; + let loaded = load_keyring(&path)?; assert_eq!(keyring.to_bytes(), loaded.to_bytes()); let _ = fs::remove_file(&path); + Ok(()) } #[test] - fn bundle_save_load_roundtrip() { + fn bundle_save_load_roundtrip() -> Result<(), Box> { let path = temp_path(BUNDLE_EXTENSION); let bundle = sample_keyring().public_key_bundle(); - save_public_key_bundle(&bundle, &path).unwrap(); - let loaded = load_public_key_bundle(&path).unwrap(); + save_public_key_bundle(&bundle, &path)?; + let loaded = load_public_key_bundle(&path)?; assert_eq!(bundle.as_bytes(), loaded.as_bytes()); let _ = fs::remove_file(&path); + Ok(()) } #[test] - fn loading_bundle_as_keyring_fails_on_magic() { + fn loading_bundle_as_keyring_fails_on_magic() -> Result<(), Box> { 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!( load_keyring(&path), Err(FileError::BadMagic { .. }) )); let _ = fs::remove_file(&path); + Ok(()) } #[test] - fn truncated_file_is_rejected() { + fn truncated_file_is_rejected() -> Result<(), Box> { 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)))); let _ = fs::remove_file(&path); + Ok(()) } #[cfg(unix)] #[test] - fn keyring_file_is_owner_only() { + fn keyring_file_is_owner_only() -> Result<(), Box> { use std::os::unix::fs::PermissionsExt; let path = temp_path(KEYRING_EXTENSION); - save_keyring(&sample_keyring(), &path).unwrap(); - let mode = fs::metadata(&path).unwrap().permissions().mode(); + save_keyring(&sample_keyring(), &path)?; + let mode = fs::metadata(&path)?.permissions().mode(); assert_eq!(mode & 0o777, 0o600); let _ = fs::remove_file(&path); + Ok(()) } } diff --git a/host/src/lib.rs b/host/src/lib.rs index 5e2d541..fb3295f 100644 --- a/host/src/lib.rs +++ b/host/src/lib.rs @@ -225,7 +225,10 @@ async fn run_dispatcher( // 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 // 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 { pipe_id, description: reader.description().to_string(), @@ -471,8 +474,12 @@ impl MTPHost { Some(v) => v, None => return Err(AcceptError::UnsupportedVersion(client_version)), }; - let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) - .expect("negotiated version must be registered"); + let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) { + Some(codec) => codec, + None => { + return Err(AcceptError::UnsupportedVersion(negotiated)); + } + }; let description = match first_msg.get_data(DataType::Description) { DataValue::Str(s) => Some(s.clone()), _ => None, @@ -508,8 +515,12 @@ impl MTPHost { Some(v) => v, None => return Err(AcceptError::UnsupportedVersion(client_version)), }; - let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) - .expect("negotiated version must be registered"); + let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) { + Some(codec) => codec, + None => { + return Err(AcceptError::UnsupportedVersion(negotiated)); + } + }; let description = match first_msg.get_data(DataType::Description) { DataValue::Str(s) => Some(s.clone()), _ => None, @@ -1047,8 +1058,12 @@ impl MTPHost { return Err(AcceptError::Send(e)); } - let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) - .expect("negotiated version must be registered"); + let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) { + Some(codec) => codec, + None => { + return Err(AcceptError::UnsupportedVersion(negotiated)); + } + }; Ok(Some(self.connection_from_parts( sender, @@ -1162,8 +1177,12 @@ impl MTPHost { Some(v) => v, None => return Err(AcceptError::UnsupportedVersion(client_version)), }; - let codec = VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) - .expect("negotiated version must be registered"); + let codec = match VersionedCodec::for_version(self.registry.clone(), negotiated.clone()) { + Some(codec) => codec, + None => { + return Err(AcceptError::UnsupportedVersion(negotiated)); + } + }; return Ok(Some(self.connection_from_parts( sender, receiver, @@ -1246,9 +1265,10 @@ mod tests { } #[test] - fn host_config_pongs_default_to_enabled() { - let config = HostConfig::new("127.0.0.1".parse().unwrap(), 4433, Vec::new(), Vec::new()); + fn host_config_pongs_default_to_enabled() -> Result<(), Box> { + let config = HostConfig::new("127.0.0.1".parse()?, 4433, Vec::new(), Vec::new()); assert!(config.send_pongs); assert!(!config.with_pongs(false).send_pongs); + Ok(()) } } diff --git a/transport/src/host.rs b/transport/src/host.rs index ddc74ce..e103132 100644 --- a/transport/src/host.rs +++ b/transport/src/host.rs @@ -6,10 +6,18 @@ use std::sync::Arc; use wtransport::{Connection as WTConnection, Endpoint, ServerConfig}; fn generate_self_signed_cert() -> (Vec, Vec) { - let key_pair = rcgen::KeyPair::generate().unwrap(); - let params = - rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]).unwrap(); - let cert = params.self_signed(&key_pair).unwrap(); + let key_pair = match rcgen::KeyPair::generate() { + Ok(key_pair) => key_pair, + Err(e) => panic!("failed to generate self-signed key pair: {e}"), + }; + 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 key_pem = key_pair.serialize_pem(); (cert_pem.into_bytes(), key_pem.into_bytes()) diff --git a/transport/tests/integration.rs b/transport/tests/integration.rs index d9551d4..201e366 100644 --- a/transport/tests/integration.rs +++ b/transport/tests/integration.rs @@ -4,40 +4,39 @@ use mtp_codec::{CommunicationType, CommunicationValue, DataType, DataValue, Type use mtp_transport::{Host, Policy, Receiver, Sender, connect, host}; fn generate_self_signed_cert() -> (Vec, Vec) { - let key_pair = rcgen::KeyPair::generate().unwrap(); - let params = - rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]).unwrap(); - let cert = params.self_signed(&key_pair).unwrap(); + let key_pair = rcgen::KeyPair::generate().expect("failed to generate self-signed key pair"); + let params = rcgen::CertificateParams::new(vec!["localhost".into(), "127.0.0.1".into()]) + .expect("failed to build self-signed certificate params"); + let cert = params + .self_signed(&key_pair) + .expect("failed to self-sign certificate"); let cert_pem = cert.pem(); let key_pem = key_pair.serialize_pem(); (cert_pem.into_bytes(), key_pem.into_bytes()) } -async fn start_test_host(cert_pem: Vec, key_pem: Vec) -> Host { - host( +async fn start_test_host(cert_pem: Vec, key_pem: Vec) -> Result> { + Ok(host( IpAddr::V4(Ipv4Addr::LOCALHOST), 0, cert_pem, key_pem, Policy::default(), ) - .await - .unwrap() + .await?) } -async fn connect_to_host(h: &Host, cert_pem: Vec) -> (Sender, Receiver) { +async fn connect_to_host(h: &Host, cert_pem: Vec) -> Result<(Sender, Receiver), Box> { let url = format!("https://127.0.0.1:{}", h.local_addr().port()); - connect(&url, Some(cert_pem), Policy::default()) - .await - .unwrap() + Ok(connect(&url, Some(cert_pem), Policy::default()).await?) } -async fn connected_pair() -> (Host, Sender, Receiver, Sender, Receiver) { +async fn connected_pair() -> Result<(Host, Sender, Receiver, Sender, Receiver), Box> { let (cert_pem, key_pem) = generate_self_signed_cert(); - 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 (host_tx, host_rx) = h.next().await.unwrap(); - (h, client_tx, client_rx, host_tx, host_rx) + 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 (host_tx, host_rx) = h.next().await.ok_or("host did not accept connection")?; + Ok((h, client_tx, client_rx, host_tx, host_rx)) } fn numbered_message(comm_type: CommunicationType, value: u128, tm: &TypeMap) -> CommunicationValue { @@ -61,103 +60,105 @@ fn assert_numbered_message( } #[tokio::test] -async fn test_host_start_and_stop() { +async fn test_host_start_and_stop() -> Result<(), Box> { 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(); // Port should be non-zero (OS-assigned) assert!(addr.port() > 0); + Ok(()) } #[tokio::test] -async fn test_send_receive_roundtrip() { - let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await; +async fn test_send_receive_roundtrip() -> Result<(), Box> { + let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await?; let tm = TypeMap::latest(); // Client sends a simple message let msg = numbered_message(CommunicationType::Ping, 42, &tm); - client_tx.send(&msg).await.unwrap(); + client_tx.send(&msg).await?; // Host receives it - let received = host_rx.receive().await.unwrap(); + let received = host_rx.receive().await?; assert_numbered_message(&received, CommunicationType::Ping, 42, &tm); // Host sends a response let resp = numbered_message(CommunicationType::Pong, 99, &tm); - host_tx.send(&resp).await.unwrap(); + host_tx.send(&resp).await?; // 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); // Close both sides client_tx.close(); host_tx.close(); + Ok(()) } #[tokio::test] -async fn test_concurrent_messages() { - let (_h, client_tx, _client_rx, _host_tx, host_rx) = connected_pair().await; +async fn test_concurrent_messages() -> Result<(), Box> { + let (_h, client_tx, _client_rx, _host_tx, host_rx) = connected_pair().await?; let tm = TypeMap::latest(); // Send 5 messages in sequence for i in 0..5u128 { let msg = numbered_message(CommunicationType::Ping, i, &tm); - client_tx.send(&msg).await.unwrap(); + client_tx.send(&msg).await?; } // Receive all 5 in order 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); } // Send 3 responses back for i in 0..3u128 { 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 { - let received = host_rx.receive().await.unwrap(); + let received = host_rx.receive().await?; assert_numbered_message(&received, CommunicationType::Pong, i * 10, &tm); } client_tx.close(); + Ok(()) } #[tokio::test] -async fn test_close_detection() { - let (_h, client_tx, _client_rx, _host_tx, host_rx) = connected_pair().await; +async fn test_close_detection() -> Result<(), Box> { + let (_h, client_tx, _client_rx, _host_tx, host_rx) = connected_pair().await?; // Send a message then close let msg = CommunicationValue::new(CommunicationType::Ping); - client_tx.send(&msg).await.unwrap(); + client_tx.send(&msg).await?; client_tx.close(); // Host should still receive the message 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)); // Host should get an error or closed signal on next receive let result = host_rx.receive().await; assert!(result.is_err()); + Ok(()) } #[tokio::test] -async fn test_host_shutdown_stops_accepting() { +async fn test_host_shutdown_stops_accepting() -> Result<(), Box> { 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()); // A connection succeeds while the host is accepting. - let (_c_tx, _c_rx) = connect(&url, Some(cert_pem.clone()), Policy::default()) - .await - .unwrap(); - let _accepted = h.next().await.unwrap(); + let (_c_tx, _c_rx) = connect(&url, Some(cert_pem.clone()), Policy::default()).await?; + let _accepted = h.next().await.ok_or("host did not accept connection")?; // 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 @@ -173,16 +174,17 @@ async fn test_host_shutdown_stops_accepting() { matches!(result, Err(_) | Ok(Err(_))), "connect should not succeed after host shutdown" ); + Ok(()) } #[tokio::test] -async fn test_drop_receiver_keeps_sender_alive() { - let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await; +async fn test_drop_receiver_keeps_sender_alive() -> Result<(), Box> { + let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await?; // Client sends a message the host receives. let msg = CommunicationValue::new(CommunicationType::Ping); - client_tx.send(&msg).await.unwrap(); - let _ = host_rx.receive().await.unwrap(); + client_tx.send(&msg).await?; + let _ = host_rx.receive().await?; // Dropping the host Receiver aborts only its accept task; the Sender shares // the same connection and must keep working. @@ -191,69 +193,71 @@ async fn test_drop_receiver_keeps_sender_alive() { let tm = TypeMap::latest(); 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); client_tx.close(); host_tx.close(); + Ok(()) } #[tokio::test] -async fn test_persistent_stream_reopens_after_local_finish() { - let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await; +async fn test_persistent_stream_reopens_after_local_finish() -> Result<(), Box> { + let (_h, client_tx, client_rx, host_tx, host_rx) = connected_pair().await?; let tm = TypeMap::latest(); let msg1 = numbered_message(CommunicationType::Ping, 11, &tm); - client_tx.send(&msg1).await.unwrap(); - let received1 = host_rx.receive().await.unwrap(); + client_tx.send(&msg1).await?; + let received1 = host_rx.receive().await?; 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); - client_tx.send(&msg2).await.unwrap(); - let received2 = host_rx.receive().await.unwrap(); + client_tx.send(&msg2).await?; + let received2 = host_rx.receive().await?; assert_numbered_message(&received2, CommunicationType::Pong, 22, &tm); client_tx.close(); host_tx.close(); drop(client_rx); + Ok(()) } #[tokio::test] -async fn test_receiver_backpressure_with_small_queue() { +async fn test_receiver_backpressure_with_small_queue() -> Result<(), Box> { 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 policy = Policy::default().with_receiver_queue_capacity(1); - let (client_tx, client_rx) = connect(&url, Some(cert_pem), policy).await.unwrap(); - let (_host_tx, host_rx) = h.next().await.unwrap(); + let (client_tx, client_rx) = connect(&url, Some(cert_pem), policy).await?; + let (_host_tx, host_rx) = h.next().await.ok_or("host did not accept connection")?; let tm = TypeMap::latest(); for i in 0..8u128 { client_tx .send(&numbered_message(CommunicationType::Ping, i, &tm)) .await - .unwrap(); + ?; } for i in 0..8u128 { let received = tokio::time::timeout(std::time::Duration::from_secs(5), host_rx.receive()) - .await - .unwrap() - .unwrap(); + .await? + ?; assert_numbered_message(&received, CommunicationType::Ping, i, &tm); } client_tx.close(); drop(client_rx); h.shutdown(); + Ok(()) } #[tokio::test] -async fn test_max_frames_per_stream_enforced() { +async fn test_max_frames_per_stream_enforced() -> Result<(), Box> { let (cert_pem, key_pem) = generate_self_signed_cert(); let policy = Policy::default().with_max_frames_per_stream(Some(1)); let mut h = host( @@ -263,35 +267,33 @@ async fn test_max_frames_per_stream_enforced() { key_pem, policy, ) - .await - .unwrap(); + .await?; let url = format!("https://127.0.0.1:{}", h.local_addr().port()); let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default()) .await - .unwrap(); - let (_host_tx, host_rx) = h.next().await.unwrap(); + ?; + let (_host_tx, host_rx) = h.next().await.ok_or("host did not accept connection")?; let tm = TypeMap::latest(); client_tx .send(&numbered_message(CommunicationType::Ping, 1, &tm)) - .await - .unwrap(); - let first = host_rx.receive().await.unwrap(); + .await?; + let first = host_rx.receive().await?; assert_numbered_message(&first, CommunicationType::Ping, 1, &tm); client_tx .send(&numbered_message(CommunicationType::Ping, 2, &tm)) - .await - .unwrap(); + .await?; let second = host_rx.receive().await; assert!(second.is_err(), "stream should be closed after frame limit"); client_tx.close(); h.shutdown(); + Ok(()) } #[tokio::test] -async fn test_semaphore_saturation_with_concurrent_streams() { +async fn test_semaphore_saturation_with_concurrent_streams() -> Result<(), Box> { let (cert_pem, key_pem) = generate_self_signed_cert(); let policy = Policy::default() .with_send_mode(mtp_transport::SendMode::SingleStreamPerMessage) @@ -304,13 +306,11 @@ async fn test_semaphore_saturation_with_concurrent_streams() { key_pem, policy, ) - .await - .unwrap(); + .await?; let url = format!("https://127.0.0.1:{}", h.local_addr().port()); let (client_tx, _client_rx) = connect(&url, Some(cert_pem), Policy::default()) - .await - .unwrap(); - let (_host_tx, host_rx) = h.next().await.unwrap(); + .await?; + let (_host_tx, host_rx) = h.next().await.ok_or("host did not accept connection")?; let tm = TypeMap::latest(); 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; for join in joins { - join.await.unwrap().unwrap(); + join.await??; } for i in 0..6u128 { let received = tokio::time::timeout(std::time::Duration::from_secs(5), host_rx.receive()) - .await - .unwrap() - .unwrap(); + .await? + ?; assert_numbered_message(&received, CommunicationType::Ping, i, &tm); } client_tx.close(); h.shutdown(); + Ok(()) } diff --git a/wasm/src/error.rs b/wasm/src/error.rs index c3ed596..2fb7d2a 100644 --- a/wasm/src/error.rs +++ b/wasm/src/error.rs @@ -37,26 +37,30 @@ mod tests { #[wasm_bindgen_test] fn from_codec_error_invalid_encoding() { 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] fn from_codec_error_unknown_version() { 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] fn from_communication_error_renders() { use mtp_common::CommunicationError; 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] fn from_crypto_error_renders() { use mtp_crypto::CryptoError; 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")); } }