From a0154853b8f53504fb7fd26e03ea5ab1b471585b Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Wed, 20 Nov 2024 19:07:15 +0900
Subject: [PATCH] use itertools::chain! instead of concat_bytes

---
 nomos-mix/message/Cargo.toml                   |  1 +
 nomos-mix/message/src/sphinx/layered_cipher.rs | 17 +++++++++++------
 nomos-mix/message/src/sphinx/mod.rs            |  8 --------
 nomos-mix/message/src/sphinx/packet.rs         | 14 +++++++++-----
 nomos-mix/message/src/sphinx/routing.rs        |  5 +++--
 5 files changed, 24 insertions(+), 21 deletions(-)

diff --git a/nomos-mix/message/Cargo.toml b/nomos-mix/message/Cargo.toml
index 71a30d20..46571fff 100644
--- a/nomos-mix/message/Cargo.toml
+++ b/nomos-mix/message/Cargo.toml
@@ -4,6 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+itertools = "0.13"
 rand_chacha = "0.3"
 serde = { version = "1", features = ["derive"] }
 sha2 = "0.10"
diff --git a/nomos-mix/message/src/sphinx/layered_cipher.rs b/nomos-mix/message/src/sphinx/layered_cipher.rs
index 03e3ed6f..5d4bbbbc 100644
--- a/nomos-mix/message/src/sphinx/layered_cipher.rs
+++ b/nomos-mix/message/src/sphinx/layered_cipher.rs
@@ -13,7 +13,7 @@ use sphinx_packet::{
     },
 };
 
-use super::{concat_bytes, parse_bytes};
+use super::parse_bytes;
 
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
@@ -118,13 +118,16 @@ impl<D: ConsistentLengthLayeredCipherData> ConsistentLengthLayeredCipher<D> {
         next_encrypted_data: Vec<u8>,
     ) -> Result<(Vec<u8>, HeaderIntegrityMac)> {
         // Concatenate the data with the encrypted subsequent layers and its MAC.
-        let total_data = concat_bytes(&[
-            &param.data.to_bytes(),
+        let data = param.data.to_bytes();
+        let total_data = itertools::chain!(
+            &data,
             next_mac.as_bytes(),
             // Truncate last bytes for the length-preserved decryption later.
             // They will be restored by a filler during the decryption process.
             &next_encrypted_data[..next_encrypted_data.len() - Self::single_layer_size()],
-        ]);
+        )
+        .copied()
+        .collect::<Vec<_>>();
 
         // Encrypt the concatenated bytes, and compute MAC.
         let mut encrypted = total_data;
@@ -160,8 +163,10 @@ impl<D: ConsistentLengthLayeredCipherData> ConsistentLengthLayeredCipher<D> {
         let random_bytes = random_bytes(self.total_size() - D::size() - fillers.len());
 
         // First, concat the data and the random bytes, and encrypt it.
-        let total_data_without_fillers =
-            concat_bytes(&[&last_param.data.to_bytes(), &random_bytes]);
+        let last_data = last_param.data.to_bytes();
+        let total_data_without_fillers = itertools::chain!(&last_data, &random_bytes)
+            .copied()
+            .collect::<Vec<_>>();
         let mut encrypted = total_data_without_fillers;
         self.apply_streamcipher(
             &mut encrypted,
diff --git a/nomos-mix/message/src/sphinx/mod.rs b/nomos-mix/message/src/sphinx/mod.rs
index f69b2f7b..c40c4a34 100644
--- a/nomos-mix/message/src/sphinx/mod.rs
+++ b/nomos-mix/message/src/sphinx/mod.rs
@@ -63,14 +63,6 @@ impl MixMessage for SphinxMessage {
     }
 }
 
-fn concat_bytes(bytes_list: &[&[u8]]) -> Vec<u8> {
-    let mut buf = Vec::with_capacity(bytes_list.iter().map(|bytes| bytes.len()).sum());
-    bytes_list
-        .iter()
-        .for_each(|bytes| buf.extend_from_slice(bytes));
-    buf
-}
-
 fn parse_bytes<'a>(data: &'a [u8], sizes: &[usize]) -> Result<Vec<&'a [u8]>, String> {
     let mut i = 0;
     sizes
diff --git a/nomos-mix/message/src/sphinx/packet.rs b/nomos-mix/message/src/sphinx/packet.rs
index 31b7a5bd..4a07fa0a 100644
--- a/nomos-mix/message/src/sphinx/packet.rs
+++ b/nomos-mix/message/src/sphinx/packet.rs
@@ -9,7 +9,7 @@ use sphinx_packet::{
 
 use crate::sphinx::ASYM_KEY_SIZE;
 
-use super::{concat_bytes, error::Error, parse_bytes, routing::EncryptedRoutingInformation};
+use super::{error::Error, parse_bytes, routing::EncryptedRoutingInformation};
 
 /// A packet that contains a header and a payload.
 /// The header and payload are encrypted for the selected recipients.
@@ -152,11 +152,15 @@ impl Packet {
     }
 
     pub fn to_bytes(&self) -> Vec<u8> {
-        concat_bytes(&[
-            &self.header.ephemeral_public_key.to_bytes(),
-            &self.header.encrypted_routing_info.to_bytes(),
+        let ephemeral_public_key = self.header.ephemeral_public_key.to_bytes();
+        let encrypted_routing_info = self.header.encrypted_routing_info.to_bytes();
+        itertools::chain!(
+            &ephemeral_public_key,
+            &encrypted_routing_info,
             &self.payload,
-        ])
+        )
+        .copied()
+        .collect()
     }
 
     pub fn from_bytes(data: &[u8], max_layers: usize) -> Result<Self, Error> {
diff --git a/nomos-mix/message/src/sphinx/routing.rs b/nomos-mix/message/src/sphinx/routing.rs
index 8d9f39a6..3963d111 100644
--- a/nomos-mix/message/src/sphinx/routing.rs
+++ b/nomos-mix/message/src/sphinx/routing.rs
@@ -8,7 +8,6 @@ use sphinx_packet::{
 };
 
 use super::{
-    concat_bytes,
     error::Error,
     layered_cipher::{
         ConsistentLengthLayeredCipher, ConsistentLengthLayeredCipherData, EncryptionParam, Key,
@@ -119,7 +118,9 @@ impl EncryptedRoutingInformation {
     }
 
     pub fn to_bytes(&self) -> Vec<u8> {
-        concat_bytes(&[self.mac.as_bytes(), &self.encrypted_routing_info])
+        itertools::chain!(self.mac.as_bytes(), &self.encrypted_routing_info)
+            .copied()
+            .collect()
     }
 
     pub fn from_bytes(data: &[u8], max_layers: usize) -> Result<Self, Error> {