From 195df85172a22fe710e6ce082dbe82db5f6c8d19 Mon Sep 17 00:00:00 2001 From: Stan Bondi Date: Mon, 3 Oct 2022 14:00:50 +0200 Subject: [PATCH] fix(dht/encryption): greatly reduce heap allocations for encrypted messaging (#4753) Description --- - Encrypt, decrypt and message padding mutate a single buffer for encrypted messages Motivation and Context --- Encrypted message handling should be as efficient as possible. The previous implementation performed allocations of the full padded message size twice for encryption and twice for decryption. Increasing memory usage, and negating the performance benefits of using an encryption keystream. This PR allocates a single buffer for the message to be de/encrypted and de/encrypts the contents in-place using the BytesMut type from the `bytes` crate. How Has This Been Tested? --- This change is backwards compatible, tested on current esme network and updated existing tests as required. Discovery: OK Memorynet: OK PingPong: OK InteractiveTransactions: OK SafTransactions: OK --- Cargo.lock | 1 - comms/core/src/lib.rs | 2 +- comms/core/src/message/mod.rs | 13 + comms/core/src/protocol/rpc/body.rs | 2 +- comms/dht/Cargo.toml | 1 - comms/dht/src/crypt.rs | 299 ++++++++++-------- comms/dht/src/dedup/mod.rs | 8 +- comms/dht/src/dht.rs | 11 +- comms/dht/src/envelope.rs | 5 +- comms/dht/src/inbound/decryption.rs | 51 ++- comms/dht/src/inbound/deserialize.rs | 4 +- comms/dht/src/inbound/dht_handler/task.rs | 2 +- comms/dht/src/inbound/forward.rs | 20 +- comms/dht/src/outbound/broadcast.rs | 24 +- comms/dht/src/outbound/message.rs | 9 +- comms/dht/src/outbound/mock.rs | 37 ++- comms/dht/src/outbound/requester.rs | 20 +- comms/dht/src/outbound/serialize.rs | 2 +- .../dht/src/store_forward/saf_handler/task.rs | 51 ++- comms/dht/src/store_forward/store.rs | 8 +- comms/dht/src/test_utils/makers.rs | 58 +++- comms/dht/src/test_utils/mod.rs | 2 + comms/dht/tests/dht.rs | 25 +- 23 files changed, 383 insertions(+), 272 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 577db90300..9cc0708252 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4805,7 +4805,6 @@ version = "0.38.4" dependencies = [ "anyhow", "bitflags 1.3.2", - "bytes 0.5.6", "chacha20 0.7.3", "chacha20poly1305 0.9.1", "chrono", diff --git a/comms/core/src/lib.rs b/comms/core/src/lib.rs index d72b52585a..570794feb7 100644 --- a/comms/core/src/lib.rs +++ b/comms/core/src/lib.rs @@ -66,6 +66,6 @@ pub mod multiaddr { } pub use async_trait::async_trait; -pub use bytes::{Bytes, BytesMut}; +pub use bytes::{Buf, BufMut, Bytes, BytesMut}; #[cfg(feature = "rpc")] pub use tower::make::MakeService; diff --git a/comms/core/src/message/mod.rs b/comms/core/src/message/mod.rs index d5a2718797..99172105fc 100644 --- a/comms/core/src/message/mod.rs +++ b/comms/core/src/message/mod.rs @@ -26,6 +26,8 @@ #[macro_use] mod envelope; + +use bytes::BytesMut; pub use envelope::EnvelopeBody; mod error; @@ -52,5 +54,16 @@ pub trait MessageExt: prost::Message { ); buf } + + /// Encodes a message into a BytesMut, allocating the buffer on the heap as necessary. + fn encode_into_bytes_mut(&self) -> BytesMut + where Self: Sized { + let mut buf = BytesMut::with_capacity(self.encoded_len()); + self.encode(&mut buf).expect( + "prost::Message::encode documentation says it is infallible unless the buffer has insufficient capacity. \ + This buffer's capacity was set with encoded_len", + ); + buf + } } impl MessageExt for T {} diff --git a/comms/core/src/protocol/rpc/body.rs b/comms/core/src/protocol/rpc/body.rs index 0790712e0e..a93ce23e8d 100644 --- a/comms/core/src/protocol/rpc/body.rs +++ b/comms/core/src/protocol/rpc/body.rs @@ -177,7 +177,7 @@ impl BodyBytes { } pub fn into_vec(self) -> Vec { - self.0.map(|bytes| bytes.to_vec()).unwrap_or_else(Vec::new) + self.0.map(|bytes| bytes.into()).unwrap_or_else(Vec::new) } pub fn into_bytes(self) -> Option { diff --git a/comms/dht/Cargo.toml b/comms/dht/Cargo.toml index 25d0590575..1fe961ee8e 100644 --- a/comms/dht/Cargo.toml +++ b/comms/dht/Cargo.toml @@ -21,7 +21,6 @@ tari_common_sqlite = { path = "../../common_sqlite" } anyhow = "1.0.53" bitflags = "1.2.0" -bytes = "0.5" chacha20 = "0.7.1" chacha20poly1305 = "0.9.1" chrono = { version = "0.4.19", default-features = false } diff --git a/comms/dht/src/crypt.rs b/comms/dht/src/crypt.rs index 518cace315..e37012db45 100644 --- a/comms/dht/src/crypt.rs +++ b/comms/dht/src/crypt.rs @@ -20,7 +20,7 @@ // WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE // USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. -use std::mem::size_of; +use std::{iter, mem::size_of}; use chacha20::{ cipher::{NewCipher, StreamCipher}, @@ -34,8 +34,13 @@ use chacha20poly1305::{ ChaCha20Poly1305, }; use digest::Digest; +use prost::bytes::BytesMut; use rand::{rngs::OsRng, RngCore}; -use tari_comms::types::{CommsPublicKey, CommsSecretKey}; +use tari_comms::{ + message::MessageExt, + types::{CommsPublicKey, CommsSecretKey}, + BufMut, +}; use tari_crypto::{ keys::DiffieHellmanSharedSecret, tari_utilities::{epoch_time::EpochTime, ByteArray}, @@ -69,31 +74,39 @@ pub fn generate_ecdh_secret(secret_key: &CommsSecretKey, public_key: &CommsPubli output } -fn pad_message_to_base_length_multiple(message: &[u8]) -> Result, DhtEncryptError> { - // We require a 32-bit length representation, and also don't want to overflow after including this encoding - if message.len() > ((u32::max_value() - (size_of::() as u32)) as usize) { - return Err(DhtEncryptError::PaddingError("Message is too long".to_string())); +fn get_message_padding_length(message_length: usize) -> usize { + if message_length == 0 { + return MESSAGE_BASE_LENGTH; } - let message_length = message.len(); - let encoded_length = (message_length as u32).to_le_bytes(); - // Pad the message (if needed) to the next multiple of the base length - let padding_length = if ((message_length + size_of::()) % MESSAGE_BASE_LENGTH) == 0 { + if message_length % MESSAGE_BASE_LENGTH == 0 { 0 } else { - MESSAGE_BASE_LENGTH - ((message_length + size_of::()) % MESSAGE_BASE_LENGTH) - }; - - // The padded message is the encoded length, message, and zero padding - let mut padded_message = Vec::with_capacity(size_of::() + message_length + padding_length); - padded_message.extend_from_slice(&encoded_length); - padded_message.extend_from_slice(message); - padded_message.extend(std::iter::repeat(0u8).take(padding_length)); + MESSAGE_BASE_LENGTH - (message_length % MESSAGE_BASE_LENGTH) + } +} - Ok(padded_message) +/// Pads a message to a multiple of MESSAGE_BASE_LENGTH excluding the additional prefix space +fn pad_message_to_base_length_multiple( + message: &mut BytesMut, + additional_prefix_space: usize, +) -> Result<(), DhtEncryptError> { + // We require a 32-bit length representation, and also don't want to overflow after including this encoding + if message.len() > u32::MAX as usize { + return Err(DhtEncryptError::PaddingError("Message is too long".to_string())); + } + let padding_length = + get_message_padding_length(message.len().checked_sub(additional_prefix_space).ok_or_else(|| { + DhtEncryptError::PaddingError("Message length shorter than the additional_prefix_space".to_string()) + })?); + message.reserve(message.len() + padding_length); + message.extend(iter::repeat(0u8).take(padding_length)); + + Ok(()) } -fn get_original_message_from_padded_text(padded_message: &[u8]) -> Result, DhtEncryptError> { +/// Returns the unpadded message. The messages must have the length prefixed to it and the nonce is removec. +fn get_original_message_from_padded_text(padded_message: &mut BytesMut) -> Result<(), DhtEncryptError> { // NOTE: This function can return errors relating to message length // It is important not to leak error types to an adversary, or to have timing differences @@ -112,25 +125,22 @@ fn get_original_message_from_padded_text(padded_message: &[u8]) -> Result()); let mut encoded_length = [0u8; size_of::()]; - encoded_length.copy_from_slice(&padded_message[0..size_of::()]); + encoded_length.copy_from_slice(&len[..]); let message_length = u32::from_le_bytes(encoded_length) as usize; // The padded message is too short for the decoded length - let end = message_length - .checked_add(size_of::()) - .ok_or_else(|| DhtEncryptError::PaddingError("Claimed unpadded message length is too large".to_string()))?; - if end > padded_message.len() { + if message_length > padded_message.len() { return Err(DhtEncryptError::CipherError( "Claimed unpadded message length is too large".to_string(), )); } // Remove the padding (we don't check for valid padding, as this is offloaded to authentication) - let start = size_of::(); - let unpadded_message = &padded_message[start..end]; + padded_message.truncate(message_length); - Ok(unpadded_message.to_vec()) + Ok(()) } pub fn generate_key_message(data: &[u8]) -> CipherKey { @@ -150,21 +160,20 @@ pub fn generate_key_signature_for_authenticated_encryption(data: &[u8]) -> Authe } /// Decrypts cipher text using ChaCha20 stream cipher given the cipher key and cipher text with integral nonce. -pub fn decrypt(cipher_key: &CipherKey, cipher_text: &[u8]) -> Result, DhtEncryptError> { +pub fn decrypt(cipher_key: &CipherKey, cipher_text: &mut BytesMut) -> Result<(), DhtEncryptError> { if cipher_text.len() < size_of::() { return Err(DhtEncryptError::InvalidDecryptionNonceNotIncluded); } - let (nonce, cipher_text) = cipher_text.split_at(size_of::()); - let nonce = Nonce::from_slice(nonce); - let mut cipher_text = cipher_text.to_vec(); + let nonce = cipher_text.split_to(size_of::()); + let nonce = Nonce::from_slice(&nonce); let mut cipher = ChaCha20::new(&cipher_key.0, nonce); - cipher.apply_keystream(cipher_text.as_mut_slice()); + cipher.apply_keystream(cipher_text); // get original message, from decrypted padded cipher text - let cipher_text = get_original_message_from_padded_text(cipher_text.as_slice())?; - Ok(cipher_text) + get_original_message_from_padded_text(cipher_text)?; + Ok(()) } pub fn decrypt_with_chacha20_poly1305( @@ -183,23 +192,49 @@ pub fn decrypt_with_chacha20_poly1305( Ok(decrypted_signature) } -/// Encrypt the plain text using the ChaCha20 stream cipher -pub fn encrypt(cipher_key: &CipherKey, plain_text: &[u8]) -> Result, DhtEncryptError> { - // pad plain_text to avoid message length leaks - let plain_text = pad_message_to_base_length_multiple(plain_text)?; - +/// Encrypt the plain text using the ChaCha20 stream cipher. The message is assumed to have a 32-bit length prepended +/// onto it. +pub fn encrypt(cipher_key: &CipherKey, plain_text: &mut BytesMut) -> Result<(), DhtEncryptError> { + if plain_text.len() < size_of::() { + return Err(DhtEncryptError::PaddingError( + "Message is not long enough to include a nonce".to_string(), + )); + } + // add nonce let mut nonce = [0u8; size_of::()]; OsRng.fill_bytes(&mut nonce); + let nonce = Nonce::from(nonce); + plain_text[..size_of::()].copy_from_slice(&nonce[..]); + + // pad plain_text to avoid message length leaks + // Excludes the nonce in the padded message length - this is mostly for backwards compatibility + pad_message_to_base_length_multiple(plain_text, size_of::())?; - let nonce_ga = Nonce::from_slice(&nonce); - let mut cipher = ChaCha20::new(&cipher_key.0, nonce_ga); + let mut cipher = ChaCha20::new(&cipher_key.0, &nonce); - let mut buf = vec![0u8; plain_text.len() + nonce.len()]; - buf[..nonce.len()].copy_from_slice(&nonce[..]); + cipher.apply_keystream(&mut plain_text[size_of::()..]); + Ok(()) +} - buf[nonce.len()..].copy_from_slice(plain_text.as_slice()); - cipher.apply_keystream(&mut buf[nonce.len()..]); - Ok(buf) +/// Encodes a prost Message, efficiently prepending the little-endian 32-bit length to the encoding +fn encode_with_prepended_length(msg: &T, additional_prefix_space: usize) -> BytesMut { + let len = msg.encoded_len(); + let mut buf = BytesMut::with_capacity(size_of::() + additional_prefix_space + len); + buf.extend(iter::repeat(0).take(additional_prefix_space)); + buf.put_u32_le(len as u32); + msg.encode(&mut buf).expect( + "prost::Message::encode documentation says it is infallible unless the buffer has insufficient capacity. This \ + buffer's capacity was set with encoded_len", + ); + buf +} + +pub fn prepare_message(is_encrypted: bool, message: &T) -> BytesMut { + if is_encrypted { + encode_with_prepended_length(message, size_of::()) + } else { + message.encode_into_bytes_mut() + } } /// Produces authenticated encryption of the signature using the ChaCha20-Poly1305 stream cipher, @@ -276,6 +311,8 @@ pub fn create_message_domain_separated_hash_parts( #[cfg(test)] mod test { + use prost::Message; + use tari_comms::message::MessageExt; use tari_crypto::keys::PublicKey; use tari_utilities::hex::from_hex; @@ -285,10 +322,11 @@ mod test { fn encrypt_decrypt() { let pk = CommsPublicKey::default(); let key = CipherKey(*chacha20::Key::from_slice(pk.as_bytes())); - let plain_text = "Last enemy position 0830h AJ 9863".as_bytes().to_vec(); - let encrypted = encrypt(&key, &plain_text).unwrap(); - let decrypted = decrypt(&key, &encrypted).unwrap(); - assert_eq!(decrypted, plain_text); + let plain_text = "Last enemy position 0830h AJ 9863".to_string(); + let mut msg = prepare_message(true, &plain_text); + encrypt(&key, &mut msg).unwrap(); + decrypt(&key, &mut msg).unwrap(); + assert_eq!(String::decode(&msg[..]).unwrap(), plain_text); } #[test] @@ -299,9 +337,11 @@ mod test { "6063cd49c7b871c0fc9785e9b959fda553fadbb10bcaaced0958e88eb6858e05fe310b4401a78d03b52a81be49db2bffcce13765e1a64460063d33289b1a3527af3df8e292c79abca71aa9a87baa1a0a6c23532a3297dda9e0c22d4b60606db1ed02a75e7a7d21fafe1214cbf8a3a66ec319a6aafeeb0e7b06375370c52b2abe63170ce50552a697f1ff87dc03ae1df574ed8e7abf915aec6959808ec526d6da78f08f2bed24268028baeba3ebd52d0fde34b145267ced68a08c4d480c213d0ab8b9c55a1630e956ed9531d6db600537f3997d612bec5905bd3ce72f5eace475e9793e6e316349f9cbd49022e401870af357605a1c7d279d5f414a5cae13e378711f345eabf46eb7fabc9465376f027a8d2d69448243cf2d70223c2430f6ce1ff55b8c2f54e27ccf77040f70c9eb84c9da9f8176a867ebcf8cbb9dfbb9256d688a76ec02af3afe3fa8221a4876462a754bc65a15c584b8c132a48dd955821961e47d5fdce62668b8f11a42b9127d5d98414bbd026eed0b5511628b96eb8435880f38d28bad6573f90611397b7f0df46d371e29a1de3591a3aa221623a540693941d6fc22d4fa57bc612282868f024613c2f69224141ad623eb49cb2b151ff79b36ec30f9a2842c696fd94f3092989ba3f8f5136850ed1f970a559e3fe19c24b84f650da4d46a5abc3514a242e3088cae15b6012e2a3cd878a7b988cc6183e81a27fc263d3be1ef07403262b22972b4f78a38d930f3e4ec1f5c9f1a1639cd138f24c665fc39b808dfbed81f8888e8e23aac1cdd980dcd0f7d95dc71fb02abcf532be17208b8d86a5fd501c18f2b18234956617797ea78523907c46a7b558cde76d1d034d6ab36ca33c4c17ea617929e3e40af011417f9b3983912fd0ed60685261532bcca5f3b3863d92cf2e4eb33a97613d9ca55112d1bd63df9d591042f972e2da3bde7b5a572f9f4cd2633330fdb3250430f27d3a40a30a996d5a41d61dafc3fe96a1fb63e2cab5c3ec0f1084e778f303da498fc970fe3117b8513166a6c8798e00a82a6c96a61b419e12717fac57dd1c989d3a10b3ab798ee0c15a5cd04dbe83667523ad4b7587ff331513b63f15c72f1844d67c7830f723fa754969d3f4254895e71d7087617ded071a797ee5791b7a95abcb360ef504bdaa85191b09b2345c0d0096fafa85b10d1675f0c4a1fe45231fd88c715c32c38d9697cebfb712e1ce57645c8faa08b4983ab8f537b017c9898b7907c06b25c01ea0f9d736d573ec7e3b14efa84d84258131ddc11a9696e12234ab65fb4e653d8dcdf4c6ce51e0104aedf2b8089593c4b5d665ed885c2d798843cf80041657ba9d800bbb9eee7076c212cfdc56df1d63e3eb4de5c6f132376c39b0272fc35e4729988eeaf9f7748142b68b10a72a948f6db24baebe3323c11002edfbe878972f02920a9ca536ff917c37a20eb0d984ee230885950fb271e56a8640d5fff75cef501cc82a46244eca5bfed4ab180b4c4d61abb209f38dabddcb4a82581bf4990631e11cc670507f8ce88c10a7fcc06ed824d7a1aab0b4dc4a2b984f45c7447b077262dc6b90f04516a5598193e95a40e982f092a2b61fcf5203379e6b042715792915adfd0edc66a69bb7fbacf64b9c86f3b15da66f8d8eecfc373204ad10f2c666181f0facd972be208c620bb0539e396c6254f2efcc2934ae6265fdc0b8172577435758e7da46df6f175f2867e6bd734fe70416b70b643b9e4b0e0e844ccfa019c699f85313f94cc6a1bb371bb14fd65a1599355ec8cdd71a555d98a900b8cbd2b1de3378f42a137eb95b473d62be3559ed687583d6963b6857d3be5f7acc12f5f5b04b26b5c582ce6dd8d1bee6322b02c2c1dc29fcba20899d529df4fd6c1edbfd1081d6cf10b20b9451ad935da2c4cef66c160550b180ba1b668029ed15448cd288427aca7f6e6505fdfc69b8111a2a071601d78484e857705a4bc7f9423800ded4eba46e0f22ee85c48fc9e8a191355edc0868df350d627a7f1120d79ba4aa1dde1ec97f8daeb0a07be914d5f6d2e74270666d03e4ca92845957b85982761dc1ee6f7603e31681dd323a045c0ac3b06b515d7bd485bfe7f6abe31e35aac7d8536b3f9c572121fcdd44c505ccfffe514e498732cab4e70524a5281b0942f5ae861b535764f056df6a1951b3c1c261f21b3b5f0a276ed05e32879ede428683b34ac8e7ebc88c9c767bf5e7cfb0cf444c1f9fd5be9041f69f6ae9772b0e035c6a2a7d9c1c614858a506a4a4bc00cee0577561b96e98c973edfa43d52471db9c716699e52260a20150aa99f8adea872c999b66fb4395d5b8c717a2c97eb7638a1d92da2ef8b2ec80db3afa3ce83445aaccae09f38c0b84c85a8983ba4c3b9a13fed4c65fd8899333da4dbca549cd2a487eb58841881f3571dfa4821bc522b56993d657bce51dfb41caf6c2cb78e8b6beceddc44febdea144da13ae9ccd9465b3ac96b06dfe79baced35ad51763d05090dc7620c89f448134507f41828be8703fd2ab1f53370e75e55366eba1e903311313707279d5965e3343476c0a8aeef2001ad88d5e452d648dd2029a6f549809c4177d1871c88abcd1404d52ebee2dd97dc52ad1a9c018428a1a64fda6773a6ea967d4124a6cf98c7e6dc4c4d9c051a376d3e3fe2e17f6cd044dd60ee32e9d6bdbdfdcbdecc4e7306092186a7ad8ab87328f9fedb6ee8ab9417968fbaa0e582205a660fa55e1ba3c5b0c84b67017f250338125894d162c400be8d563c9f0416dc5641d31bad577543cba8c6c9a7c04064e412597d47c6272d8e087bc11397533cb1bd7feebea9feee44e1b6a5f49b937594da3b719e1982a90594277f43798a39e419c204f18a6920e5ac1a751eddeaef9392a6f84d68d73aabc6ba68750d47ad4da8bd842662226225a764661ea11ff9f13d328e0242a0b513aa5ad9fbe9d484b3d28a41890e4fc62820ef2342a90c0837b30c831eb78213e7e2cd6dfbda26a7e6103ab8b4219462ca70ca57c79638b2c49f0469ea6f68335071294257c5337ccf452ca1bfedf81610f353e7576f02a2b32aba64a4252946fda330de11990f51207817860e0d8b7c9cb58a5858155db61376a01c02aaedb7017fd3c36adf4f3c07f29f352330c6d78ab6bbb7d4aabf3725833e86523b755094273465ba57545162623036a7786f426d0a63e13bebf2205a6b488bd6da3c93469a4df4b3811e9c63d62c61e0cdd263df821adc0d1b751c1314be9fd93761b447931e425db7e09baac9083aed472de5fe6172c8e8f729ade8faa96d131e86204462e14e0411b4b7629de25a0c5dbf848c9ca8c42376f5d54bff34bf36074136bbe98228745dbc9d411d891553f0af00240e1729ce7757fba2775fa5b700e95460910008584a833fb9edc073cd4d8333643631e193040d850f87cd50d9ca2e2e5c3943787dc4a4677ac7e130c2d6739945fd3b059ebe040abb38a20d73a7669516cf8503f40642217c8580a27b127f1f33eaa7adff44c922afac813c870795563fac79d139d5b5233a26728328f88f1f9daaaea1c4e1ee64ded0b006ce46015d512e8c4a411ab788a5383563949c95846202250c5b9e0baab0bc8620327ed2aacd36e1bdc9d3a4d6b4e22627d75bd088cdd47ca204f1ce44357d1b471b37581c820f6bbdfe3da1f4f90dc353833731703b7b9bb87ff2d0cae1e2f0321994759d1a21b2075a620b58b814cd65812092891261dd7e879b65843480382f59e20d6b6c67b2fb750ff0cfce897891f976b0fae7ac31e02384b251bcbecce6ff98819cf0cd6d41fdab9ba6907742394732ed5e74bcd13aad1a188855c020f09e62540be9b2992a397b30107ad730ebe183504226b303f30032f4c0a683812d05be57961430504866bd2ee6993423ebe34ba4d2d022ce6d5b2345bbed34d6807aec473ad0701b9b8fe2db1cef57748dfcb29ddb3b253a865dd7383d04253cd70c350d02ccd2371cecfd74aae820fa91eddd89d27925c33183e03d44c7f88f8068c64d223d2d5f4ab18fae6d209e1e267395576f4f48ae056da7d6e91f94991659b4c07f44aa1c45aaf75b7274b7668753f968d5e6635f4abf238e5d44ffc38e68cad8237f7e7a25d5fc0dcd5afc2bedbac6b42e8bc8064118c9042d1159f70dfac73d65c8a9782c264445af11c878591d49d49ad46f4e6d086d55232afd234c3bceab2eef0e22e5c2875670c5125e8a172f5f2168e59fe0cb5e9e1a81bf645a2c45d115b9a3efe9fe2d1799f12b0c11f50ae5540ff4e90e6220eb62451e10ce1418929e03c751d9019d47b87847595333feb6ab4af40662d04c3ece4f93b4c2c2f2ee2078724090336f16a4f33801095036a31b557960b5d8d2552f0aadfa3dc9dcfe8f1dd6a61631b6a69ee6ce8433153f8b1ea99a9a5ac688026d6ef408f2aa958ada8baf0193b3989f359c7a913fcb9eec230568584bcda3a759c824884c9febff518c7cf312360d2c1ffd2bbdd0b2e9346cbe1bf383446bd2fec431475ec509474ef9eb06817f53d3c4ca74fba08c3b434eabf3ae9fcc2287c588fc5574bff37066705ca9a39d088cd5cbb83b385b5cf647ced0c23885295d2b24f37e4098be82edccc23e1c973b1855e2009de63408c78e570b3cec65c6d236d81adb1bc298436a1e125b99bb995a5c6df5b2a4e70b8cf1db5de38120134527ce349c32f8e35fa43837aa38cdb1d5695a34d12d27bd5ee4536d9a20e62b55e59cdc7ecca1f4398dac7a4b756d9e131a7d2c8bde32c20ef0424154c88c8276fdf3c75f08f3cd423bd648ff3520680a1f1dd956451881f6d31238c11c99a20e1d9170410c8d8eb88ce90e179fc80e23e36a28b1810383a4d0d1ef0f2db94206aa1fb25498b425e5ad1f0f0bd3eed22ca5545ef541880f37f8fea82fecd59d8c94765d3a454e81775844701412e3c01a6dcdbf277428969a7f08d67313cdd2ce3b531addee28733552ee1bf4124ad8b3e40e04b94599e04cce60f5676307b0605ad7dc73b03cc26227eab60196d37c312a01858f5ad6a901e0f1c796c52cb9690da5c712a2d36c74e65ec9a60ea41387b8a0f79697cdfd93e40ab569d6a55361be97fb7ac8d80b5a5482908d44af94df2fb09a777978f4d911008d528ff44aef960cfd25fb56e26c341850721f020f9fd112cd52fc28dd129ffc2f9a12a829dcb69b54a894d4b3d1ac3b63bc9bcd39e30a00e419c8f4d2b630c224880a7d3af9c19c8a79262818b368589e7ad03b722021306fbcbcf7bc87ad418a3eb6616e7d4ce286264554be6040e8e4cd0c5a9bbdd2367e47d1fe0a9c3eeaf2455c4f6f779bab3d5bea5284a244fc3e804fb6d0e50fec91f85b71c6ba91f43a240fa48900229e5f3038b0806f70a1cb72fdea58b664f06c04bf688183a4f22255d6976f2102aafb669ee117fa1e44ae325ad52001469fed9d26e4f8592f56e42bf5e7195f521c0beaf891e47a703075fa1948ee07add55a765346b94ae498fa96145ad8460f23248222e329398fec6ad7f323c448ce82bb706b24e07adc0681901a63d5d1c7b871a9df8009ed7bb10be4e39a987c1bf039554a016ac8693284a7248fb8a9aa440dde213c2414447727c1556d25f1fbce057652044e2350b9ef5627584d403a934dd33e8c26e20799f1dbf915705b70d66256d31ca7c407307fa18e163917635d67f742828deba4b942b5f0d916b5e737b5811d3c3b4ac386c7ebaad1a6c465ce9fb229bc6ce7ae62f8efd8632e5312db8ba213d28d19843ac7fbae105a1433921b34c216c3c2ab247080a629c7ac5507129b27ce0d38ddde06722a5a0a979894d6140c31a82bfc517adf0b57c761f75bc14d65d8701e0dea92a06584f2d877dc5fb0b32496754e6b0115e99a9623ad631ea0a76b4e7893bf0982151e1c5ed6d64a305393f6f715de333653ace204c2f03de8f36c463c937f7f23326a88337624fc606317d7c0ea2badf69e40602c2ff1e2dcc9cfca1ccef566381712af157c5e458335c8a283733a617a75fd7cef52c515c754443c8e9e1930994805e6f0b2a9a2ccbd848f6580896317dd9dbfda17d00e80d35bc58a704fbe7d6d6d45752811130f682ea9471903c9af9b5c95d074ec87b32c86dc5b29a186a60a7a03e630c7a5cd38e6ab1a2f561642d5662658fc20239233505727575e75dbbc5be630f9ebc9485f03bbf0569554a87bfb5cfd397daf5d8d92a38fc4b24b1a433edad26a12c4e29506362dec83fb9b1f31158e834cde319d40bd283f0a1f3995b3ade08bcc01c794d656583b928300a6f2be57e5bf1586a123cf28ac8e2287e0b7ab67419dc4f527f714fdee8c47088cc1857a39825524dc3f5a3777c1f906cf496dac43e3f8304ebe5d5696da5b7e5d79f176a391736ba46bc718356e1713a00ea754a52b5899ba7eb71b10bbd211cead7d1890f2d8bb981a2549e2cd53bf895f96f628c4d00061275c87f4dcbadcb5944f912d27aaa4124cadd0e2a1d82ee4d3a8b977bd1a03fc6f79caf4c306addea0bd72b754c113350655324dec3dcbb1f1de66e3e7a9f06ba0e04de0cae7af7d6e31298bf5be706038e0d8477a79ea5f8e21decdf6f5ff71090d8cb2dabb9d1a87ee526b0ec84be81ad9585b09f165cfff7a4a63e30ac7341a3a42e3e02bfc34486a2a5492a39dd31dc233c74f454584e5bd2524382a08357ba2d3ef46833551ed3fa6f672d5ba0ec75258430738c16989840b6ae6909f10340d1845bd975cf2933047a1c22c332606f76681dae8727921d4f1345f0457700b8622ea72a50c17cda201f7019d4af9dc0b67bd95317d98c2fe38ab8e12dbefe236b463014caf9c3cdd390dfcff034e90e4f51e1233bb8b341bba6f922d1b0629e261844018f39d054b26cb82592e33466354f552a14f9e6175d418cb9724fa045e723b5ab9ece5aa45800f1202b3d174fe4e129e7320a9063039f8bdea8601762ff45933503e0bd10944893e565d641a39289da67e5269dc7cfc22dc3d5f5011b66b25340cea4055bd66a7752c69e624bfc12e5cc68cd0b5cab3242860b7f40541303e228e666a6500e1e739b0b6d853b5715cf3a668facc135d133d4eefa035f36c16838f25ab4a0d2f20d18f05c0fae1a705370162cfe6f7fcf1c69654c2ce73e3820b48568c25a6d9d036c2386dee6ef14e5fc967ffccde38bf263c8c0f924bcfcdf54669dcf872724280cc4f81bb2aa993998f6312d0c6084ed823e5e6bff0ed25cc4e82b749dc11f4cf55290344d9c307d634793e81b9d3d457765dc6f81b66f1a6aaaa1079558a4892edfc342fe24856200b5dcc65c9e7809b655d3cc7bb26bcda91933f590bc61099fd0b83e04bad2174150645afd7c3ccac5417234e30da4e7574af953f8d9b7a5029417d439f1d13c4390bed2bc05d73821ff2355c33da5f95623c73abd826614572841e4777a9a0b538cef4a2c6327c75116977322a8c488f466178cdcaf3f0e10df86dbd1827ba2cc4c8fba90a1d64ad783a77704c5b1262cd11cb010f09ab04377d6e5ebed4d5dfe8eaa0cc2535a0be69bdf5e1987167b8135428ef84287aa4424c35c7a7bc94cabd553df4840121403b2ba3479e1a7f86085cce49c245af944a2ed78b77784309d05d5f5587a6e589baa6e7d279b0ab43afb5497b6d5954b3a8f66dd547d3b72565437ead511c9d5342406aace95e5cc31f2c6d618a24c219d0298f980a571ce29b999d20b3ed94a60e286ed7ddc647c439e3d421814da6d91c8b7d5b3fa70ec6a3c261ab9ee4ac779545edcc7db6df3345db26b91c5a997dac5e62b75dc05358821bab4fe65a049fe9ce8e537ae81a10dfcca0ef97c6cb95e3ff1573e5461f7b505e1678fbe97a41ab696b53f6ea09038ecda09eed34247251424b766306c0c64fd836274cf85fe0e0c19638127c2210b580c9194fe0cccac7ac80e3dde38e9cccd6a194ee923f4e73800bec0c77f60553f9c7c8413ea87d20c114d7b415fdd87fc55f273b1b3a9f9c71c4462d5b3f300daf0fc6c338278e5991e0c6de07a3c288d237df00325230be204f7b2bb7a127ac28b001e4225e910eeb9521f5af6cfeae1f18c08bf8ac9d1513c3794ba5b8ea9fb9a57825cb154fc1e9a9dddf809dd6bb11a207625b23b274344e7e0b7ca666e456735d5901f1341aca42e749183823b3debbd563aeebc68f9b15dce13d0fc1acef47d38d5967c2b6b3fe8ed69b180dbcbf17455ee6825641202ccd145c0a0a0f4091622338f48474e5838d8915f814eb87ad45e710b07f79f662c2120278ce05978d8a7aee20fc5661a08c072977ed878092e7183332b70c9c54db307c705e527f6fd2076e39c216b0490f552d52a109652958c62fc6bf7f913818dbdf5d92550779aae541d54d059d5844658422c17a24e374fa6f92e5a9fda87eee249747b9cd292043c9731d2c1d08d06eab030fb49e779cb58bf4f776d6aa0185db860007d8b2d0f7205dacb9201ac9538d2c37062f736b6b44e971e11500", ) .unwrap(); - let plain_text = decrypt(&key, &cipher_text).unwrap(); + + let mut text = BytesMut::from(&cipher_text[..]); + decrypt(&key, &mut text).unwrap(); let secret_msg = "Last enemy position 0830h AJ 9863".as_bytes().to_vec(); - assert_eq!(plain_text, secret_msg); + assert_eq!(text, secret_msg); } #[test] @@ -399,106 +439,86 @@ mod test { #[test] fn pad_message_correctness() { // test for small message - let message = &[0u8, 10, 22, 11, 38, 74, 59, 91, 73, 82, 75, 23, 59]; - let prepend_message = (message.len() as u32).to_le_bytes(); - let pad = std::iter::repeat(0u8) - .take(MESSAGE_BASE_LENGTH - message.len() - prepend_message.len()) + let message = [0u8, 10, 22, 11, 38, 74, 59, 91, 73, 82, 75, 23, 59].as_slice(); + let pad = iter::repeat(0u8) + .take(MESSAGE_BASE_LENGTH - message.len()) .collect::>(); - let pad_message = pad_message_to_base_length_multiple(message).unwrap(); + let mut pad_message = BytesMut::from(message); + pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap(); // padded message is of correct length assert_eq!(pad_message.len(), MESSAGE_BASE_LENGTH); - // prepend message is well specified - assert_eq!(prepend_message, pad_message[..prepend_message.len()]); // message body is well specified - assert_eq!( - *message, - pad_message[prepend_message.len()..prepend_message.len() + message.len()] - ); + assert_eq!(*message, pad_message[..message.len()]); // pad is well specified - assert_eq!(pad, pad_message[prepend_message.len() + message.len()..]); + assert_eq!(pad, pad_message[message.len()..]); // test for large message - let message = &[100u8; MESSAGE_BASE_LENGTH * 8 - 100]; - let prepend_message = (message.len() as u32).to_le_bytes(); - let pad_message = pad_message_to_base_length_multiple(message).unwrap(); - let pad = std::iter::repeat(0u8) - .take((8 * MESSAGE_BASE_LENGTH) - message.len() - prepend_message.len()) + let message = encode_with_prepended_length(&vec![100u8; MESSAGE_BASE_LENGTH * 8 - 100], 0); + let mut pad_message = message.clone(); + pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap(); + let pad = iter::repeat(0u8) + .take((8 * MESSAGE_BASE_LENGTH) - message.len()) .collect::>(); // padded message is of correct length assert_eq!(pad_message.len(), 8 * MESSAGE_BASE_LENGTH); - // prepend message is well specified - assert_eq!(prepend_message, pad_message[..prepend_message.len()]); // message body is well specified - assert_eq!( - *message, - pad_message[prepend_message.len()..prepend_message.len() + message.len()] - ); + assert_eq!(*message, pad_message[..message.len()]); // pad is well specified - assert_eq!(pad, pad_message[prepend_message.len() + message.len()..]); + assert_eq!(pad, pad_message[message.len()..]); // test for base message of multiple base length - let message = &[100u8; MESSAGE_BASE_LENGTH * 9 - 123]; - let prepend_message = (message.len() as u32).to_le_bytes(); + let message = encode_with_prepended_length(&vec![100u8; MESSAGE_BASE_LENGTH * 9 - 123], 0); let pad = std::iter::repeat(0u8) - .take((9 * MESSAGE_BASE_LENGTH) - message.len() - prepend_message.len()) + .take((9 * MESSAGE_BASE_LENGTH) - message.len()) .collect::>(); - let pad_message = pad_message_to_base_length_multiple(message).unwrap(); + let mut pad_message = message.clone(); + pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap(); // padded message is of correct length assert_eq!(pad_message.len(), 9 * MESSAGE_BASE_LENGTH); - // prepend message is well specified - assert_eq!(prepend_message, pad_message[..prepend_message.len()]); // message body is well specified - assert_eq!( - *message, - pad_message[prepend_message.len()..prepend_message.len() + message.len()] - ); + assert_eq!(*message, pad_message[..message.len()]); // pad is well specified - assert_eq!(pad, pad_message[prepend_message.len() + message.len()..]); + assert_eq!(pad, pad_message[message.len()..]); // test for empty message - let message: [u8; 0] = []; - let prepend_message = (message.len() as u32).to_le_bytes(); - let pad_message = pad_message_to_base_length_multiple(&message).unwrap(); + let message = encode_with_prepended_length(&vec![], 0); + let mut pad_message = message.clone(); + pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap(); let pad = [0u8; MESSAGE_BASE_LENGTH - 4]; // padded message is of correct length assert_eq!(pad_message.len(), MESSAGE_BASE_LENGTH); - // prepend message is well specified - assert_eq!(prepend_message, pad_message[..prepend_message.len()]); // message body is well specified - assert_eq!( - message, - pad_message[prepend_message.len()..prepend_message.len() + message.len()] - ); + assert_eq!(message, pad_message[..message.len()]); // pad is well specified - assert_eq!(pad, pad_message[prepend_message.len() + message.len()..]); + assert_eq!(pad, pad_message[message.len()..]); } #[test] fn unpadding_failure_modes() { // The padded message is empty - let message: [u8; 0] = []; - assert!(get_original_message_from_padded_text(&message) + let mut message = BytesMut::new(); + assert!(get_original_message_from_padded_text(&mut message) .unwrap_err() .to_string() .contains("Padded message is not long enough for length extraction")); // We cannot extract the message length - let message = [0u8; size_of::() - 1]; - assert!(get_original_message_from_padded_text(&message) + let mut message = BytesMut::from([0u8; size_of::() - 1].as_slice()); + assert!(get_original_message_from_padded_text(&mut message) .unwrap_err() .to_string() .contains("Padded message is not long enough for length extraction")); // The padded message is not a multiple of the base length - let message = [0u8; 2 * MESSAGE_BASE_LENGTH + 1]; - assert!(get_original_message_from_padded_text(&message) + let mut message = BytesMut::from([0u8; 2 * MESSAGE_BASE_LENGTH + 1].as_slice()); + assert!(get_original_message_from_padded_text(&mut message) .unwrap_err() .to_string() .contains("Padded message must be a multiple of the base length")); @@ -508,44 +528,56 @@ mod test { fn get_original_message_from_padded_text_successful() { // test for short message let message = vec![0u8, 10, 22, 11, 38, 74, 59, 91, 73, 82, 75, 23, 59]; - let pad_message = pad_message_to_base_length_multiple(message.as_slice()).unwrap(); + let mut pad_message = encode_with_prepended_length(&message, 0); + pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap(); - let output_message = get_original_message_from_padded_text(pad_message.as_slice()).unwrap(); - assert_eq!(message, output_message); + // + let mut output_message = pad_message.clone(); + get_original_message_from_padded_text(&mut output_message).unwrap(); + assert_eq!(message.to_encoded_bytes(), output_message); // test for large message let message = vec![100u8; 1024]; - let pad_message = pad_message_to_base_length_multiple(message.as_slice()).unwrap(); + let mut pad_message = encode_with_prepended_length(&message, 0); + pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap(); - let output_message = get_original_message_from_padded_text(pad_message.as_slice()).unwrap(); - assert_eq!(message, output_message); + let mut output_message = pad_message.clone(); + get_original_message_from_padded_text(&mut output_message).unwrap(); + assert_eq!(message.to_encoded_bytes(), output_message); // test for base message of base length let message = vec![100u8; 984]; - let pad_message = pad_message_to_base_length_multiple(message.as_slice()).unwrap(); + let mut pad_message = encode_with_prepended_length(&message, 0); + pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap(); - let output_message = get_original_message_from_padded_text(pad_message.as_slice()).unwrap(); - assert_eq!(message, output_message); + let mut output_message = pad_message.clone(); + get_original_message_from_padded_text(&mut output_message).unwrap(); + assert_eq!(message.to_encoded_bytes(), output_message); // test for empty message let message: Vec = vec![]; - let pad_message = pad_message_to_base_length_multiple(message.as_slice()).unwrap(); + let mut pad_message = encode_with_prepended_length(&message, 0); + pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap(); - let output_message = get_original_message_from_padded_text(pad_message.as_slice()).unwrap(); - assert_eq!(message, output_message); + let mut output_message = pad_message.clone(); + get_original_message_from_padded_text(&mut output_message).unwrap(); + assert_eq!(message.to_encoded_bytes(), output_message); } #[test] fn padding_fails_if_pad_message_prepend_length_is_bigger_than_plaintext_length() { - let message = "This is my secret message, keep it secret !".as_bytes(); - let mut pad_message = pad_message_to_base_length_multiple(message).unwrap(); + let message = "This is my secret message, keep it secret !".as_bytes().to_vec(); + let mut pad_message = encode_with_prepended_length(&message, 0); + pad_message_to_base_length_multiple(&mut pad_message, 0).unwrap(); + let mut pad_message = pad_message.to_vec(); // we modify the prepend length, in order to assert that the get original message // method will output a different length message pad_message[0] = 1; - let modified_message = get_original_message_from_padded_text(pad_message.as_slice()).unwrap(); - assert!(message.len() != modified_message.len()); + let mut modified_message = BytesMut::from(pad_message.as_slice()); + get_original_message_from_padded_text(&mut modified_message).unwrap(); + assert_ne!(message.len(), modified_message.len()); // add big number from le bytes of prepend bytes pad_message[0] = 255; @@ -553,7 +585,8 @@ mod test { pad_message[2] = 255; pad_message[3] = 255; - assert!(get_original_message_from_padded_text(pad_message.as_slice()) + let mut pad_message = BytesMut::from(pad_message.as_slice()); + assert!(get_original_message_from_padded_text(&mut pad_message) .unwrap_err() .to_string() .contains("Claimed unpadded message length is too large")); @@ -565,24 +598,30 @@ mod test { // in any way the value of the decrypted content, by applying a cipher stream let pk = CommsPublicKey::default(); let key = CipherKey(*chacha20::Key::from_slice(pk.as_bytes())); - let message = "My secret message, keep it secret !".as_bytes().to_vec(); - let mut encrypted = encrypt(&key, &message).unwrap(); + let message = "My secret message, keep it secret !".to_string(); + let mut msg = encode_with_prepended_length(&message, size_of::()); + encrypt(&key, &mut msg).unwrap(); - let n = encrypted.len(); - encrypted[n - 1] += 1; + let n = msg.len(); + msg[n - 1] += 1; - assert!(decrypt(&key, &encrypted).unwrap() == message); + decrypt(&key, &mut msg).unwrap(); + assert_eq!(String::decode(&msg[..]).unwrap(), message); } #[test] fn decryption_fails_if_message_body_is_modified() { let pk = CommsPublicKey::default(); let key = CipherKey(*chacha20::Key::from_slice(pk.as_bytes())); - let message = "My secret message, keep it secret !".as_bytes().to_vec(); - let mut encrypted = encrypt(&key, &message).unwrap(); + let message = "My secret message, keep it secret !".to_string(); + let mut msg = encode_with_prepended_length(&message, size_of::()); + encrypt(&key, &mut msg).unwrap(); - encrypted[size_of::() + size_of::() + 1] += 1; + msg[size_of::() + size_of::() + 1] += 1; - assert!(decrypt(&key, &encrypted).unwrap() != message); + // TODO: decryption does not "fail" is this intended? + decrypt(&key, &mut msg).unwrap(); + eprintln!("msg = {:?}", msg); + assert_ne!(msg, message); } } diff --git a/comms/dht/src/dedup/mod.rs b/comms/dht/src/dedup/mod.rs index fa7cf3f20e..f49e48a4a4 100644 --- a/comms/dht/src/dedup/mod.rs +++ b/comms/dht/src/dedup/mod.rs @@ -197,7 +197,7 @@ mod test { assert!(dedup.poll_ready(&mut cx).is_ready()); let node_identity = make_node_identity(); let inbound_message = - make_dht_inbound_message(&node_identity, vec![], DhtMessageFlags::empty(), false, false).unwrap(); + make_dht_inbound_message(&node_identity, &vec![], DhtMessageFlags::empty(), false, false).unwrap(); let decrypted_msg = DecryptedDhtMessage::succeeded(wrap_in_envelope_body!(vec![]), None, inbound_message); rt.block_on(dedup.call(decrypted_msg.clone())).unwrap(); @@ -213,12 +213,12 @@ mod test { #[test] fn deterministic_hash() { const TEST_MSG: &[u8] = b"test123"; - const EXPECTED_HASH: &str = "d6333668f259f677703fbe4e89152ee41c7c01f6dec502befc63120246523ffe"; + const EXPECTED_HASH: &str = "1c2bb1bcff443af4441b789bd1d6984bb8d7bed2c9f85e8cf4f45615fdd9e47d"; let node_identity = make_node_identity(); let dht_message = make_dht_inbound_message( &node_identity, - TEST_MSG.to_vec(), + &TEST_MSG.to_vec(), DhtMessageFlags::empty(), false, false, @@ -229,7 +229,7 @@ mod test { let node_identity = make_node_identity(); let dht_message = make_dht_inbound_message( &node_identity, - TEST_MSG.to_vec(), + &TEST_MSG.to_vec(), DhtMessageFlags::empty(), false, false, diff --git a/comms/dht/src/dht.rs b/comms/dht/src/dht.rs index 5361665f42..e3acdae98a 100644 --- a/comms/dht/src/dht.rs +++ b/comms/dht/src/dht.rs @@ -494,7 +494,7 @@ mod test { let msg = wrap_in_envelope_body!(b"secret".to_vec()); let dht_envelope = make_dht_envelope( &node_identity, - msg.to_encoded_bytes(), + &msg, DhtMessageFlags::empty(), false, MessageTag::new(), @@ -546,7 +546,7 @@ mod test { // Encrypt for self let dht_envelope = make_dht_envelope( &node_identity, - msg.to_encoded_bytes(), + &msg, DhtMessageFlags::ENCRYPTED, true, MessageTag::new(), @@ -602,10 +602,11 @@ mod test { let node_identity2 = make_node_identity(); let ecdh_key = crypt::generate_ecdh_secret(node_identity2.secret_key(), node_identity2.public_key()); let key_message = crypt::generate_key_message(&ecdh_key); - let encrypted_bytes = crypt::encrypt(&key_message, &msg.to_encoded_bytes()).unwrap(); + let mut encrypted_bytes = msg.encode_into_bytes_mut(); + crypt::encrypt(&key_message, &mut encrypted_bytes).unwrap(); let dht_envelope = make_dht_envelope( &node_identity2, - encrypted_bytes, + &encrypted_bytes.to_vec(), DhtMessageFlags::ENCRYPTED, true, MessageTag::new(), @@ -667,7 +668,7 @@ mod test { let msg = wrap_in_envelope_body!(b"secret".to_vec()); let mut dht_envelope = make_dht_envelope( &node_identity, - msg.to_encoded_bytes(), + &msg, DhtMessageFlags::empty(), false, MessageTag::new(), diff --git a/comms/dht/src/envelope.rs b/comms/dht/src/envelope.rs index 27038803af..3f4f2ef06e 100644 --- a/comms/dht/src/envelope.rs +++ b/comms/dht/src/envelope.rs @@ -28,7 +28,6 @@ use std::{ }; use bitflags::bitflags; -use bytes::Bytes; use chrono::{DateTime, NaiveDateTime, Utc}; use prost_types::Timestamp; use serde::{Deserialize, Serialize}; @@ -249,10 +248,10 @@ impl From for DhtHeader { } impl DhtEnvelope { - pub fn new(header: DhtHeader, body: &Bytes) -> Self { + pub fn new(header: DhtHeader, body: Vec) -> Self { Self { header: Some(header), - body: body.to_vec(), + body, } } } diff --git a/comms/dht/src/inbound/decryption.rs b/comms/dht/src/inbound/decryption.rs index 3c9c9e634c..419baf00c0 100644 --- a/comms/dht/src/inbound/decryption.rs +++ b/comms/dht/src/inbound/decryption.rs @@ -30,6 +30,7 @@ use tari_comms::{ message::EnvelopeBody, peer_manager::NodeIdentity, pipeline::PipelineError, + BytesMut, }; use thiserror::Error; use tower::{layer::Layer, Service, ServiceExt}; @@ -406,11 +407,11 @@ where S: Service message_body: &[u8], ) -> Result { let key_message = crypt::generate_key_message(shared_secret); - let decrypted = - crypt::decrypt(&key_message, message_body).map_err(DecryptionError::DecryptionFailedMalformedCipher)?; + let mut decrypted = BytesMut::from(message_body); + crypt::decrypt(&key_message, &mut decrypted).map_err(DecryptionError::DecryptionFailedMalformedCipher)?; // Deserialization into an EnvelopeBody is done here to determine if the // decryption produced valid bytes or not. - EnvelopeBody::decode(decrypted.as_slice()) + EnvelopeBody::decode(decrypted.freeze()) .and_then(|body| { // Check if we received a body length of zero // @@ -477,10 +478,11 @@ mod test { use futures::{executor::block_on, future}; use tari_comms::{ - message::{MessageExt, MessageTag}, + message::MessageTag, runtime, test_utils::mocks::create_connectivity_mock, wrap_in_envelope_body, + BytesMut, }; use tari_test_utils::{counter_context, unpack_enum}; use tokio::time::sleep; @@ -492,6 +494,7 @@ mod test { test_utils::{ make_dht_header, make_dht_inbound_message, + make_dht_inbound_message_raw, make_keypair, make_node_identity, make_valid_message_signature, @@ -527,14 +530,8 @@ mod test { let mut service = DecryptionService::new(Default::default(), node_identity.clone(), connectivity, service); let plain_text_msg = wrap_in_envelope_body!(b"Secret plans".to_vec()); - let inbound_msg = make_dht_inbound_message( - &node_identity, - plain_text_msg.to_encoded_bytes(), - DhtMessageFlags::ENCRYPTED, - true, - true, - ) - .unwrap(); + let inbound_msg = + make_dht_inbound_message(&node_identity, &plain_text_msg, DhtMessageFlags::ENCRYPTED, true, true).unwrap(); block_on(service.call(inbound_msg)).unwrap(); let decrypted = result.lock().unwrap().take().unwrap(); @@ -560,7 +557,7 @@ mod test { let some_other_node_identity = make_node_identity(); let inbound_msg = make_dht_inbound_message( &some_other_node_identity, - some_secret, + &some_secret, DhtMessageFlags::ENCRYPTED, true, true, @@ -591,7 +588,7 @@ mod test { let nonsense = b"Cannot Decrypt this".to_vec(); let inbound_msg = - make_dht_inbound_message(&node_identity, nonsense.clone(), DhtMessageFlags::ENCRYPTED, true, true).unwrap(); + make_dht_inbound_message_raw(&node_identity, nonsense, DhtMessageFlags::ENCRYPTED, true, true).unwrap(); let err = service.call(inbound_msg).await.unwrap_err(); let err = err.downcast::().unwrap(); @@ -615,14 +612,8 @@ mod test { let mut service = DecryptionService::new(Default::default(), node_identity.clone(), connectivity, service); let plain_text_msg = b"Secret message to nowhere".to_vec(); - let inbound_msg = make_dht_inbound_message( - &node_identity, - plain_text_msg.to_encoded_bytes(), - DhtMessageFlags::ENCRYPTED, - true, - false, - ) - .unwrap(); + let inbound_msg = + make_dht_inbound_message(&node_identity, &plain_text_msg, DhtMessageFlags::ENCRYPTED, true, false).unwrap(); let err = service.call(inbound_msg).await.unwrap_err(); let err = err.downcast::().unwrap(); @@ -645,13 +636,15 @@ mod test { let node_identity = make_node_identity(); let mut service = DecryptionService::new(Default::default(), node_identity.clone(), connectivity, service); - let plain_text_msg = b"Secret message".to_vec(); + let plain_text_msg = BytesMut::from(b"Secret message".as_slice()); let (e_secret_key, e_public_key) = make_keypair(); let shared_secret = crypt::generate_ecdh_secret(&e_secret_key, node_identity.public_key()); let key_message = crypt::generate_key_message(&shared_secret); let msg_tag = MessageTag::new(); - let message = crypt::encrypt(&key_message, &plain_text_msg).unwrap(); + let mut message = plain_text_msg.clone(); + crypt::encrypt(&key_message, &mut message).unwrap(); + let message = message.freeze(); let header = make_dht_header( &node_identity, &e_public_key, @@ -663,7 +656,7 @@ mod test { true, ) .unwrap(); - let envelope = DhtEnvelope::new(header.into(), &message.into()); + let envelope = DhtEnvelope::new(header.into(), message.into()); let msg_tag = MessageTag::new(); let mut inbound_msg = DhtInboundMessage::new( msg_tag, @@ -706,13 +699,15 @@ mod test { let node_identity = make_node_identity(); let mut service = DecryptionService::new(Default::default(), node_identity.clone(), connectivity, service); - let plain_text_msg = b"Public message".to_vec(); + let plain_text_msg = BytesMut::from(b"Public message".as_slice()); let (e_secret_key, e_public_key) = make_keypair(); let shared_secret = crypt::generate_ecdh_secret(&e_secret_key, node_identity.public_key()); let key_message = crypt::generate_key_message(&shared_secret); let msg_tag = MessageTag::new(); - let message = crypt::encrypt(&key_message, &plain_text_msg).unwrap(); + let mut message = plain_text_msg.clone(); + crypt::encrypt(&key_message, &mut message).unwrap(); + let message = message.freeze(); let header = make_dht_header( &node_identity, &e_public_key, @@ -724,7 +719,7 @@ mod test { true, ) .unwrap(); - let envelope = DhtEnvelope::new(header.into(), &message.into()); + let envelope = DhtEnvelope::new(header.into(), message.into()); let msg_tag = MessageTag::new(); let mut inbound_msg = DhtInboundMessage::new( msg_tag, diff --git a/comms/dht/src/inbound/deserialize.rs b/comms/dht/src/inbound/deserialize.rs index 7f50c7317a..23537899ff 100644 --- a/comms/dht/src/inbound/deserialize.rs +++ b/comms/dht/src/inbound/deserialize.rs @@ -161,7 +161,7 @@ mod test { let dht_envelope = make_dht_envelope( &node_identity, - b"A".to_vec(), + &b"A".to_vec(), DhtMessageFlags::empty(), false, MessageTag::new(), @@ -181,7 +181,7 @@ mod test { .unwrap(); let msg = spy.pop_request().unwrap(); - assert_eq!(msg.body, b"A".to_vec()); + assert_eq!(msg.body, b"A".to_vec().to_encoded_bytes()); assert_eq!(msg.dht_header, dht_envelope.header.unwrap().try_into().unwrap()); } } diff --git a/comms/dht/src/inbound/dht_handler/task.rs b/comms/dht/src/inbound/dht_handler/task.rs index 1760b47295..e6ee3c7a5d 100644 --- a/comms/dht/src/inbound/dht_handler/task.rs +++ b/comms/dht/src/inbound/dht_handler/task.rs @@ -234,7 +234,7 @@ where S: Service .with_debug_info("Propagating join message".to_string()) .with_dht_header(dht_header) .finish(), - body.to_encoded_bytes(), + body.encode_into_bytes_mut(), ) .await?; } diff --git a/comms/dht/src/inbound/forward.rs b/comms/dht/src/inbound/forward.rs index 7ddd9e4fa7..e687eff8a1 100644 --- a/comms/dht/src/inbound/forward.rs +++ b/comms/dht/src/inbound/forward.rs @@ -24,7 +24,8 @@ use std::task::Poll; use futures::{future::BoxFuture, task::Context}; use log::*; -use tari_comms::{peer_manager::Peer, pipeline::PipelineError}; +use prost::bytes::BufMut; +use tari_comms::{peer_manager::Peer, pipeline::PipelineError, BytesMut}; use tari_utilities::epoch_time::EpochTime; use tower::{layer::Layer, Service, ServiceExt}; @@ -204,12 +205,11 @@ where S: Service return Ok(()); } } - - let body = decryption_result + let err_body = decryption_result .as_ref() - .err() - .cloned() - .expect("previous check that decryption failed"); + .expect_err("previous check that decryption failed"); + let mut body = BytesMut::with_capacity(err_body.len()); + body.put(err_body.as_slice()); let excluded_peers = vec![source_peer.node_id.clone()]; let dest_node_id = dht_header.destination.to_derived_node_id(); @@ -259,7 +259,7 @@ where S: Service mod test { use std::time::Duration; - use tari_comms::{runtime, runtime::task, wrap_in_envelope_body}; + use tari_comms::{message::MessageExt, runtime, runtime::task, wrap_in_envelope_body}; use tokio::sync::mpsc; use super::*; @@ -278,7 +278,7 @@ mod test { let node_identity = make_node_identity(); let inbound_msg = - make_dht_inbound_message(&node_identity, b"".to_vec(), DhtMessageFlags::empty(), false, false).unwrap(); + make_dht_inbound_message(&node_identity, &b"".to_vec(), DhtMessageFlags::empty(), false, false).unwrap(); let msg = DecryptedDhtMessage::succeeded( wrap_in_envelope_body!(Vec::new()), Some(node_identity.public_key().clone()), @@ -300,7 +300,7 @@ mod test { let sample_body = b"Lorem ipsum"; let inbound_msg = make_dht_inbound_message( &make_node_identity(), - sample_body.to_vec(), + &sample_body.to_vec(), DhtMessageFlags::empty(), false, false, @@ -318,7 +318,7 @@ mod test { let (params, body) = oms_mock_state.pop_call().await.unwrap(); // Header and body are preserved when forwarding - assert_eq!(&body.to_vec(), &sample_body); + assert_eq!(&body.to_vec(), &sample_body.to_vec().to_encoded_bytes()); assert_eq!(params.dht_header.unwrap(), header); } } diff --git a/comms/dht/src/outbound/broadcast.rs b/comms/dht/src/outbound/broadcast.rs index 8999d2fd41..1dd2b49649 100644 --- a/comms/dht/src/outbound/broadcast.rs +++ b/comms/dht/src/outbound/broadcast.rs @@ -22,7 +22,6 @@ use std::{sync::Arc, task::Poll}; -use bytes::Bytes; use chrono::{DateTime, Utc}; use futures::{ future, @@ -37,6 +36,8 @@ use tari_comms::{ peer_manager::{NodeId, NodeIdentity, Peer}, pipeline::PipelineError, types::CommsPublicKey, + Bytes, + BytesMut, }; use tari_crypto::{keys::PublicKey, tari_utilities::epoch_time::EpochTime}; use tari_utilities::{hex::Hex, ByteArray}; @@ -238,7 +239,7 @@ where S: Service async fn handle_send_message( &mut self, params: FinalSendMessageParams, - body: Bytes, + body: BytesMut, reply_tx: oneshot::Sender, ) -> Result, DhtOutboundError> { trace!(target: LOG_TARGET, "Send params: {:?}", params); @@ -405,7 +406,7 @@ where S: Service extra_flags: DhtMessageFlags, force_origin: bool, is_broadcast: bool, - body: Bytes, + body: BytesMut, expires: Option>, tag: Option, ) -> Result<(Vec, Vec), DhtOutboundError> { @@ -485,7 +486,7 @@ where S: Service message_type: DhtMessageType, flags: DhtMessageFlags, expires: Option, - body: Bytes, + mut body: BytesMut, ) -> Result { match encryption { OutboundEncryption::EncryptFor(public_key) => { @@ -497,7 +498,8 @@ where S: Service // Generate key message for encryption of message let key_message = crypt::generate_key_message(&shared_ephemeral_secret); // Encrypt the message with the body with key message above - let encrypted_body = crypt::encrypt(&key_message, &body)?; + crypt::encrypt(&key_message, &mut body)?; + let encrypted_body = body.freeze(); // Produce domain separated signature signature let mac_signature = crypt::create_message_domain_separated_hash_parts( @@ -525,7 +527,7 @@ where S: Service Ok(( Some(Arc::new(e_public_key)), Some(encrypted_message_signature.into()), - encrypted_body.into(), + encrypted_body, )) }, OutboundEncryption::ClearText => { @@ -546,9 +548,9 @@ where S: Service &binding_message_representation, ) .to_proto(); - Ok((None, Some(signature.to_encoded_bytes().into()), body)) + Ok((None, Some(signature.to_encoded_bytes().into()), body.freeze())) } else { - Ok((None, None, body)) + Ok((None, None, body.freeze())) } }, } @@ -633,7 +635,7 @@ mod test { service .call(DhtOutboundRequest::SendMessage( Box::new(SendMessageParams::new().flood(vec![]).finish()), - b"custom_msg".to_vec().into(), + b"custom_msg".as_slice().into(), reply_tx, )) .await @@ -680,7 +682,7 @@ mod test { .with_discovery(false) .finish(), ), - Bytes::from_static(b"custom_msg"), + BytesMut::from(b"custom_msg".as_slice()), reply_tx, )) .await @@ -728,7 +730,7 @@ mod test { .with_discovery(true) .finish(), ), - b"custom_msg".to_vec().into(), + b"custom_msg".as_slice().into(), reply_tx, )) .await diff --git a/comms/dht/src/outbound/message.rs b/comms/dht/src/outbound/message.rs index 544287e090..588cbb6929 100644 --- a/comms/dht/src/outbound/message.rs +++ b/comms/dht/src/outbound/message.rs @@ -22,11 +22,12 @@ use std::{fmt, fmt::Display, sync::Arc}; -use bytes::Bytes; use tari_comms::{ message::{MessageTag, MessagingReplyTx}, peer_manager::NodeId, types::CommsPublicKey, + Bytes, + BytesMut, }; use tari_utilities::hex::Hex; use thiserror::Error; @@ -145,7 +146,11 @@ impl SendMessageResponse { #[derive(Debug)] pub enum DhtOutboundRequest { /// Send a message using the given broadcast strategy - SendMessage(Box, Bytes, oneshot::Sender), + SendMessage( + Box, + BytesMut, + oneshot::Sender, + ), } impl fmt::Display for DhtOutboundRequest { diff --git a/comms/dht/src/outbound/mock.rs b/comms/dht/src/outbound/mock.rs index 7d7b58d926..c36640f13c 100644 --- a/comms/dht/src/outbound/mock.rs +++ b/comms/dht/src/outbound/mock.rs @@ -21,15 +21,17 @@ // USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. use std::{ + mem, sync::Arc, time::{Duration, Instant}, }; -use bytes::Bytes; +use chacha20::Nonce; use log::*; use tari_comms::{ message::{MessageTag, MessagingReplyTx}, protocol::messaging::SendFailReason, + BytesMut, }; use tokio::{ sync::{mpsc, oneshot, watch, Mutex, RwLock}, @@ -61,7 +63,7 @@ pub fn create_outbound_service_mock(size: usize) -> (OutboundMessageRequester, O #[derive(Clone)] pub struct OutboundServiceMockState { #[allow(clippy::type_complexity)] - calls: Arc>>, + calls: Arc>>, next_response: Arc>>, notif_sender: Arc>, notif_reciever: watch::Receiver<()>, @@ -121,17 +123,36 @@ impl OutboundServiceMockState { self.next_response.write().await.take() } - pub async fn add_call(&self, req: (FinalSendMessageParams, Bytes)) { + async fn add_call(&self, req: (FinalSendMessageParams, BytesMut)) { self.calls.lock().await.push(req); let _r = self.notif_sender.send(()); } - pub async fn take_calls(&self) -> Vec<(FinalSendMessageParams, Bytes)> { - self.calls.lock().await.drain(..).collect() + pub async fn take_calls(&self) -> Vec<(FinalSendMessageParams, BytesMut)> { + self.calls + .lock() + .await + .drain(..) + .map(|(p, mut b)| { + if p.encryption.is_encrypt() { + // Remove prefix data + (p, b.split_off(mem::size_of::() + mem::size_of::())) + } else { + (p, b) + } + }) + .collect() } - pub async fn pop_call(&self) -> Option<(FinalSendMessageParams, Bytes)> { - self.calls.lock().await.pop() + pub async fn pop_call(&self) -> Option<(FinalSendMessageParams, BytesMut)> { + self.calls.lock().await.pop().map(|(p, mut b)| { + if p.encryption.is_encrypt() { + // Remove prefix data + (p, b.split_off(mem::size_of::() + mem::size_of::())) + } else { + (p, b) + } + }) } pub async fn set_behaviour(&self, behaviour: MockBehaviour) { @@ -232,7 +253,7 @@ impl OutboundServiceMock { async fn add_call( &mut self, params: FinalSendMessageParams, - body: Bytes, + body: BytesMut, ) -> (SendMessageResponse, MessagingReplyTx) { self.mock_state.add_call((params, body)).await; let (inner_reply_tx, inner_reply_rx) = oneshot::channel(); diff --git a/comms/dht/src/outbound/requester.rs b/comms/dht/src/outbound/requester.rs index 0b1e38e9ee..0ac3e2e619 100644 --- a/comms/dht/src/outbound/requester.rs +++ b/comms/dht/src/outbound/requester.rs @@ -21,11 +21,12 @@ // USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. use log::*; -use tari_comms::{message::MessageExt, peer_manager::NodeId, types::CommsPublicKey, wrap_in_envelope_body}; +use tari_comms::{peer_manager::NodeId, types::CommsPublicKey, wrap_in_envelope_body, BytesMut}; use tokio::sync::{mpsc, oneshot}; use super::message::DhtOutboundRequest; use crate::{ + crypt::prepare_message, domain_message::OutboundDomainMessage, envelope::NodeDestination, outbound::{ @@ -259,7 +260,8 @@ impl OutboundMessageRequester { } else { message.to_propagation_header() }; - let body = wrap_in_envelope_body!(header, message.into_inner()).to_encoded_bytes(); + let msg = wrap_in_envelope_body!(header, message.into_inner()); + let body = prepare_message(params.encryption.is_encrypt(), &msg); self.send_raw(params, body).await } @@ -275,7 +277,8 @@ impl OutboundMessageRequester { if cfg!(debug_assertions) { trace!(target: LOG_TARGET, "Send Message: {} {:?}", params, message); } - let body = wrap_in_envelope_body!(message).to_encoded_bytes(); + let msg = wrap_in_envelope_body!(message); + let body = prepare_message(params.encryption.is_encrypt(), &msg); self.send_raw(params, body).await } @@ -291,7 +294,8 @@ impl OutboundMessageRequester { if cfg!(debug_assertions) { trace!(target: LOG_TARGET, "Send Message: {} {:?}", params, message); } - let body = wrap_in_envelope_body!(message).to_encoded_bytes(); + let msg = wrap_in_envelope_body!(message); + let body = prepare_message(params.encryption.is_encrypt(), &msg); self.send_raw_no_wait(params, body).await } @@ -299,11 +303,11 @@ impl OutboundMessageRequester { pub async fn send_raw( &mut self, params: FinalSendMessageParams, - body: Vec, + body: BytesMut, ) -> Result { let (reply_tx, reply_rx) = oneshot::channel(); self.sender - .send(DhtOutboundRequest::SendMessage(Box::new(params), body.into(), reply_tx)) + .send(DhtOutboundRequest::SendMessage(Box::new(params), body, reply_tx)) .await?; reply_rx @@ -315,11 +319,11 @@ impl OutboundMessageRequester { pub async fn send_raw_no_wait( &mut self, params: FinalSendMessageParams, - body: Vec, + body: BytesMut, ) -> Result<(), DhtOutboundError> { let (reply_tx, _) = oneshot::channel(); self.sender - .send(DhtOutboundRequest::SendMessage(Box::new(params), body.into(), reply_tx)) + .send(DhtOutboundRequest::SendMessage(Box::new(params), body, reply_tx)) .await?; Ok(()) } diff --git a/comms/dht/src/outbound/serialize.rs b/comms/dht/src/outbound/serialize.rs index 0d47a64cc5..ff0e8fe745 100644 --- a/comms/dht/src/outbound/serialize.rs +++ b/comms/dht/src/outbound/serialize.rs @@ -97,7 +97,7 @@ where message_tag: tag.as_value(), expires, }); - let envelope = DhtEnvelope::new(dht_header, &body); + let envelope = DhtEnvelope::new(dht_header, body.into()); let body = Bytes::from(envelope.to_encoded_bytes()); diff --git a/comms/dht/src/store_forward/saf_handler/task.rs b/comms/dht/src/store_forward/saf_handler/task.rs index 0aada15e4e..7f5390d382 100644 --- a/comms/dht/src/store_forward/saf_handler/task.rs +++ b/comms/dht/src/store_forward/saf_handler/task.rs @@ -34,6 +34,7 @@ use tari_comms::{ peer_manager::{NodeId, NodeIdentity, Peer, PeerFeatures, PeerManagerError}, pipeline::PipelineError, types::CommsPublicKey, + BytesMut, }; use tari_utilities::{convert::try_convert_all, ByteArray}; use tokio::sync::mpsc; @@ -563,9 +564,10 @@ where S: Service ); let key_message = crypt::generate_key_message(&shared_secret); - let decrypted_bytes = crypt::decrypt(&key_message, body)?; + let mut decrypted_bytes = BytesMut::from(body); + crypt::decrypt(&key_message, &mut decrypted_bytes)?; let envelope_body = - EnvelopeBody::decode(decrypted_bytes.as_slice()).map_err(|_| StoreAndForwardError::DecryptionFailed)?; + EnvelopeBody::decode(decrypted_bytes.freeze()).map_err(|_| StoreAndForwardError::DecryptionFailed)?; if envelope_body.is_empty() { return Err(StoreAndForwardError::InvalidEnvelopeBody); } @@ -702,7 +704,7 @@ mod test { None, make_dht_inbound_message( &node_identity, - b"Keep this for others please".to_vec(), + &b"Keep this for others please".to_vec(), DhtMessageFlags::ENCRYPTED, true, false, @@ -793,10 +795,9 @@ mod test { sleep(Duration::from_secs(5)).await; } assert_eq!(oms_mock_state.call_count().await, 1); - let call = oms_mock_state.pop_call().await.unwrap(); + let (_, body) = oms_mock_state.pop_call().await.unwrap(); - let body = call.1.to_vec(); - let body = EnvelopeBody::decode(body.as_slice()).unwrap(); + let body = EnvelopeBody::decode(body).unwrap(); let msg = body.decode_part::(0).unwrap().unwrap(); assert_eq!(msg.messages().len(), 1); @@ -827,19 +828,19 @@ mod test { let node_identity = make_node_identity(); - let msg_a = wrap_in_envelope_body!(&b"A".to_vec()).to_encoded_bytes(); + let msg_a = wrap_in_envelope_body!(&b"A".to_vec()); let inbound_msg_a = - make_dht_inbound_message(&node_identity, msg_a.clone(), DhtMessageFlags::ENCRYPTED, true, false).unwrap(); + make_dht_inbound_message(&node_identity, &msg_a, DhtMessageFlags::ENCRYPTED, true, false).unwrap(); // Need to know the peer to process a stored message peer_manager .add_peer(Clone::clone(&*inbound_msg_a.source_peer)) .await .unwrap(); - let msg_b = &wrap_in_envelope_body!(b"B".to_vec()).to_encoded_bytes(); + let msg_b = wrap_in_envelope_body!(b"B".to_vec()); let inbound_msg_b = - make_dht_inbound_message(&node_identity, msg_b.clone(), DhtMessageFlags::ENCRYPTED, true, false).unwrap(); + make_dht_inbound_message(&node_identity, &msg_b, DhtMessageFlags::ENCRYPTED, true, false).unwrap(); // Need to know the peer to process a stored message peer_manager .add_peer(Clone::clone(&*inbound_msg_b.source_peer)) @@ -856,20 +857,14 @@ mod test { let msg2 = ProtoStoredMessage::new(0, inbound_msg_b.dht_header, inbound_msg_b.body, msg2_time); // Cleartext message - let clear_msg = wrap_in_envelope_body!(b"Clear".to_vec()).to_encoded_bytes(); - let clear_header = make_dht_inbound_message( - &node_identity, - clear_msg.clone(), - DhtMessageFlags::empty(), - false, - false, - ) - .unwrap() - .dht_header; + let clear_msg = wrap_in_envelope_body!(b"Clear".to_vec()); + let clear_header = make_dht_inbound_message(&node_identity, &clear_msg, DhtMessageFlags::empty(), false, false) + .unwrap() + .dht_header; let msg_clear_time = Utc::now() .checked_sub_signed(chrono::Duration::from_std(Duration::from_secs(120)).unwrap()) .unwrap(); - let msg_clear = ProtoStoredMessage::new(0, clear_header, clear_msg, msg_clear_time); + let msg_clear = ProtoStoredMessage::new(0, clear_header, clear_msg.to_encoded_bytes(), msg_clear_time); let mut message = DecryptedDhtMessage::succeeded( wrap_in_envelope_body!(StoredMessagesResponse { messages: vec![msg1.clone(), msg2, msg_clear], @@ -879,7 +874,7 @@ mod test { None, make_dht_inbound_message( &node_identity, - b"Stored message".to_vec(), + &b"Stored message".to_vec(), DhtMessageFlags::ENCRYPTED, true, false, @@ -950,9 +945,9 @@ mod test { let node_identity = make_node_identity(); - let msg_a = wrap_in_envelope_body!(&b"A".to_vec()).to_encoded_bytes(); + let msg_a = wrap_in_envelope_body!(&b"A".to_vec()); let inbound_msg_a = - make_dht_inbound_message(&node_identity, msg_a, DhtMessageFlags::ENCRYPTED, true, false).unwrap(); + make_dht_inbound_message(&node_identity, &msg_a, DhtMessageFlags::ENCRYPTED, true, false).unwrap(); peer_manager .add_peer(Clone::clone(&*inbound_msg_a.source_peer)) .await @@ -973,7 +968,7 @@ mod test { None, make_dht_inbound_message( &node_identity, - b"Stored message".to_vec(), + &b"Stored message".to_vec(), DhtMessageFlags::ENCRYPTED, true, false, @@ -1023,9 +1018,9 @@ mod test { let node_identity = make_node_identity(); - let msg_a = wrap_in_envelope_body!(&b"A".to_vec()).to_encoded_bytes(); + let msg_a = wrap_in_envelope_body!(&b"A".to_vec()); let inbound_msg_a = - make_dht_inbound_message(&node_identity, msg_a, DhtMessageFlags::ENCRYPTED, true, false).unwrap(); + make_dht_inbound_message(&node_identity, &msg_a, DhtMessageFlags::ENCRYPTED, true, false).unwrap(); peer_manager .add_peer(Clone::clone(&*inbound_msg_a.source_peer)) .await @@ -1046,7 +1041,7 @@ mod test { None, make_dht_inbound_message( &node_identity, - b"Stored message".to_vec(), + &b"Stored message".to_vec(), DhtMessageFlags::ENCRYPTED, true, false, diff --git a/comms/dht/src/store_forward/store.rs b/comms/dht/src/store_forward/store.rs index 61519cd8ca..c0d2b8d224 100644 --- a/comms/dht/src/store_forward/store.rs +++ b/comms/dht/src/store_forward/store.rs @@ -483,7 +483,7 @@ mod test { let inbound_msg = make_dht_inbound_message( &make_node_identity(), - b"".to_vec(), + &b"".to_vec(), DhtMessageFlags::empty(), false, false, @@ -509,7 +509,7 @@ mod test { let msg_node_identity = make_node_identity(); let inbound_msg = make_dht_inbound_message( &msg_node_identity, - b"This shouldnt be stored".to_vec(), + &b"This shouldnt be stored".to_vec(), DhtMessageFlags::ENCRYPTED, true, false, @@ -539,7 +539,7 @@ mod test { let mut inbound_msg = make_dht_inbound_message( &origin_node_identity, - b"Will you keep this for me?".to_vec(), + &b"Will you keep this for me?".to_vec(), DhtMessageFlags::ENCRYPTED, true, false, @@ -582,7 +582,7 @@ mod test { let mut inbound_msg = make_dht_inbound_message( &origin_node_identity, - b"Will you keep this for me?".to_vec(), + &b"Will you keep this for me?".to_vec(), DhtMessageFlags::ENCRYPTED, true, false, diff --git a/comms/dht/src/test_utils/makers.rs b/comms/dht/src/test_utils/makers.rs index 7646346b3a..c9ba154d98 100644 --- a/comms/dht/src/test_utils/makers.rs +++ b/comms/dht/src/test_utils/makers.rs @@ -36,6 +36,7 @@ use tari_test_utils::{paths::create_temporary_data_path, random}; use crate::{ crypt, + crypt::prepare_message, envelope::{DhtMessageFlags, DhtMessageHeader, NodeDestination}, inbound::DhtInboundMessage, message_signature::MessageSignature, @@ -123,9 +124,9 @@ pub fn make_valid_message_signature(node_identity: &NodeIdentity, message: &[u8] .to_encoded_bytes() } -pub fn make_dht_inbound_message( +pub fn make_dht_inbound_message( node_identity: &NodeIdentity, - body: Vec, + body: &T, flags: DhtMessageFlags, include_origin: bool, include_destination: bool, @@ -148,24 +149,65 @@ pub fn make_dht_inbound_message( )) } +pub fn make_dht_inbound_message_raw( + node_identity: &NodeIdentity, + body: Vec, + flags: DhtMessageFlags, + include_origin: bool, + include_destination: bool, +) -> Result { + let msg_tag = MessageTag::new(); + let (e_secret_key, e_public_key) = make_keypair(); + let header = make_dht_header( + node_identity, + &e_public_key, + &e_secret_key, + &body, + flags, + include_origin, + msg_tag, + include_destination, + )? + .into(); + let envelope = DhtEnvelope::new(header, body); + Ok(DhtInboundMessage::new( + msg_tag, + envelope.header.unwrap().try_into().unwrap(), + Arc::new(Peer::new( + node_identity.public_key().clone(), + node_identity.node_id().clone(), + Vec::::new().into(), + PeerFlags::empty(), + PeerFeatures::COMMUNICATION_NODE, + Default::default(), + Default::default(), + )), + envelope.body, + )) +} + pub fn make_keypair() -> (CommsSecretKey, CommsPublicKey) { CommsPublicKey::random_keypair(&mut OsRng) } -pub fn make_dht_envelope( +pub fn make_dht_envelope( node_identity: &NodeIdentity, - mut message: Vec, + message: &T, flags: DhtMessageFlags, include_origin: bool, trace: MessageTag, include_destination: bool, ) -> Result { let (e_secret_key, e_public_key) = make_keypair(); - if flags.is_encrypted() { + let message = if flags.is_encrypted() { let shared_secret = crypt::generate_ecdh_secret(&e_secret_key, node_identity.public_key()); let key_message = crypt::generate_key_message(&shared_secret); - message = crypt::encrypt(&key_message, &message).unwrap(); - } + let mut message = prepare_message(true, message); + crypt::encrypt(&key_message, &mut message).unwrap(); + message.freeze() + } else { + prepare_message(false, message).freeze() + }; let header = make_dht_header( node_identity, &e_public_key, @@ -177,7 +219,7 @@ pub fn make_dht_envelope( include_destination, )? .into(); - Ok(DhtEnvelope::new(header, &message.into())) + Ok(DhtEnvelope::new(header, message.into())) } pub fn build_peer_manager() -> Arc { diff --git a/comms/dht/src/test_utils/mod.rs b/comms/dht/src/test_utils/mod.rs index 03f531c075..39e8fff377 100644 --- a/comms/dht/src/test_utils/mod.rs +++ b/comms/dht/src/test_utils/mod.rs @@ -45,7 +45,9 @@ pub use dht_actor_mock::{create_dht_actor_mock, DhtMockState}; mod dht_discovery_mock; pub use dht_discovery_mock::{create_dht_discovery_mock, DhtDiscoveryMockState}; +#[cfg(test)] mod makers; +#[cfg(test)] pub use makers::*; mod service; diff --git a/comms/dht/tests/dht.rs b/comms/dht/tests/dht.rs index 9928c1df79..55e9da8e72 100644 --- a/comms/dht/tests/dht.rs +++ b/comms/dht/tests/dht.rs @@ -33,7 +33,6 @@ use tari_comms::{ protocol::messaging::{MessagingEvent, MessagingEventSender, MessagingProtocolExtension}, transports::MemoryTransport, types::CommsDatabase, - wrap_in_envelope_body, CommsBuilder, CommsNode, }; @@ -419,16 +418,13 @@ async fn dht_store_forward() { node_A .dht .outbound_requester() - .send_raw( - params.clone(), - wrap_in_envelope_body!(secret_msg1.to_vec()).to_encoded_bytes(), - ) + .send_message_no_header(params.clone(), secret_msg1.to_vec()) .await .unwrap(); node_A .dht .outbound_requester() - .send_raw(params, wrap_in_envelope_body!(secret_msg2.to_vec()).to_encoded_bytes()) + .send_message_no_header(params, secret_msg2.to_vec()) .await .unwrap(); @@ -722,7 +718,7 @@ async fn dht_do_not_store_invalid_message_in_dedup() { // Get the message that was received by Node B let mut msg = node_B.next_inbound_message(Duration::from_secs(10)).await.unwrap(); - let bytes = msg.decryption_result.unwrap().to_encoded_bytes(); + let bytes = msg.decryption_result.unwrap().encode_into_bytes_mut(); // Clone header without modification let header_unmodified = msg.dht_header.clone(); @@ -972,9 +968,9 @@ async fn dht_propagate_message_contents_not_malleable_ban() { let msg = node_B.next_inbound_message(Duration::from_secs(10)).await.unwrap(); - let mut bytes = msg.decryption_result.unwrap().to_encoded_bytes(); + let mut envelope = msg.decryption_result.unwrap(); // Change the message - bytes.push(0x42); + envelope.push_part([0x42].to_vec()); let mut connectivity_events = node_C.comms.connectivity().get_event_subscription(); @@ -982,7 +978,7 @@ async fn dht_propagate_message_contents_not_malleable_ban() { node_B .dht .outbound_requester() - .send_raw( + .send_message_no_header( SendMessageParams::new() .propagate(node_B.node_identity().public_key().clone().into(), vec![msg .source_peer @@ -990,7 +986,7 @@ async fn dht_propagate_message_contents_not_malleable_ban() { .clone()]) .with_dht_header(msg.dht_header) .finish(), - bytes, + envelope, ) .await .unwrap(); @@ -1016,7 +1012,6 @@ async fn dht_propagate_message_contents_not_malleable_ban() { #[tokio::test] #[allow(non_snake_case)] async fn dht_header_not_malleable() { - env_logger::init(); let node_C = make_node("node_C", PeerFeatures::COMMUNICATION_NODE, dht_config(), None).await; // Node B knows about Node C let mut node_B = make_node( @@ -1081,14 +1076,14 @@ async fn dht_header_not_malleable() { // Modify the header msg.dht_header.message_type = DhtMessageType::from_i32(21i32).unwrap(); - let bytes = msg.decryption_result.unwrap().to_encoded_bytes(); + let envelope = msg.decryption_result.unwrap(); let mut connectivity_events = node_C.comms.connectivity().get_event_subscription(); // Propagate the changed message (to node C) node_B .dht .outbound_requester() - .send_raw( + .send_message_no_header( SendMessageParams::new() .propagate(node_B.node_identity().public_key().clone().into(), vec![msg .source_peer @@ -1096,7 +1091,7 @@ async fn dht_header_not_malleable() { .clone()]) .with_dht_header(msg.dht_header) .finish(), - bytes, + envelope, ) .await .unwrap();