Mix: Packet header encryption

nomos-mix/message/Cargo.toml

@@ -4,6 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+rand_chacha = "0.3.1"
 serde = { version = "1.0", features = ["derive"] }
 sha2 = "0.10"
 sphinx-packet = "0.2"

nomos-mix/message/src/error.rs

@@ -4,10 +4,14 @@ pub enum Error {
     InvalidMixMessage,
     #[error("Payload is too large")]
     PayloadTooLarge,
-    #[error("Too many recipients")]
-    TooManyRecipients,
     #[error("Sphinx packet error: {0}")]
     SphinxPacketError(#[from] sphinx_packet::Error),
+    #[error("Invalid routing flag: {0}")]
+    InvalidRoutingFlag(u8),
+    #[error("Invalid routing length: {0}")]
+    InvalidEncryptedRoutingInfoLength(usize),
+    #[error("ConsistentLengthLayeredEncryptionError: {0}")]
+    ConsistentLengthLayeredEncryptionError(#[from] crate::layered_cipher::Error),
     #[error("Unwrapping a message is not allowed to this node")]
     /// e.g. the message cannot be unwrapped using the private key provided
     MsgUnwrapNotAllowed,

nomos-mix/message/src/layered_cipher.rs

@@ -0,0 +1,321 @@
+use rand_chacha::{
+    rand_core::{RngCore, SeedableRng},
+    ChaCha12Rng,
+};
+use sphinx_packet::{
+    constants::HEADER_INTEGRITY_MAC_SIZE,
+    crypto::STREAM_CIPHER_INIT_VECTOR,
+    header::{
+        keys::{HeaderIntegrityMacKey, StreamCipherKey},
+        mac::HeaderIntegrityMac,
+    },
+};
+
+use crate::{concat_bytes, parse_bytes};
+
+#[derive(thiserror::Error, Debug)]
+pub enum Error {
+    #[error("Invalid data length")]
+    InvalidDataLength,
+    #[error("Invalid params")]
+    InvalidParams,
+    #[error("Integrity MAC verification failed")]
+    IntegrityMacVerificationFailed,
+}
+
+type Result<T> = std::result::Result<T, Error>;
+
+/// A cipher to encrypt/decrypt a list of data of the same size using a list of keys.
+///
+/// The cipher performs the layered encryption.
+/// The following example shows the simplified output.
+/// - Input: [[data0, k0], [data1, k1]]
+/// - Output: encrypt(k0, [data0, encrypt(k1, [data1])])
+///
+/// The max number of layers is limited to the `max_layers` parameter.
+/// Even if the number of data and keys provided for encryption is smaller than `max_layers`,
+/// The cipher always produces the max-size output regardless of the number of data and keys provided,
+/// in order to ensure that all outputs generated by the cipher are the same size.
+///
+/// The cipher also provides the length-preserved decryption.
+/// Even if one layer of encryptions is decrypted, the length of decrypted data is
+/// the same as the length of the original data.
+/// For example:
+///   len(encrypt(k0, [data0, encrypt(k1, [data1])])) == len(encrypt(k1, [data1]))
+pub struct ConsistentLengthLayeredCipher {
+    /// All encrypted data produced by the cipher has the same size according to the `max_layers`.
+    pub max_layers: usize,
+    /// The size of data that will be included in each layer.
+    pub data_size: usize,
+}
+
+/// A parameter for one layer of encryption
+pub struct EncryptionParam {
+    /// A data to be included in the layer.
+    /// This must have the same size as `data_size` configured in the cipher.
+    pub data: Vec<u8>,
+    /// A [`Key`] to encrypt the layer that will include the [`Self::data`].
+    pub key: Key,
+}
+
+/// A set of keys to encrypt/decrypt a single layer.
+pub struct Key {
+    /// A 128-bit key for encryption/decryption
+    pub stream_cipher_key: StreamCipherKey,
+    /// A 128-bit key for computing/verifying integrity MAC
+    pub integrity_mac_key: HeaderIntegrityMacKey,
+}
+
+impl ConsistentLengthLayeredCipher {
+    pub fn new(max_layers: usize, data_size: usize) -> Self {
+        Self {
+            max_layers,
+            data_size,
+        }
+    }
+
+    /// The total size of fully encrypted output that includes all layers.
+    /// This size is determined by [`Self::data_size`] and [`Self::max_layers`].
+    pub fn total_size(&self) -> usize {
+        self.single_layer_size() * self.max_layers
+    }
+
+    /// The size of a single layer that contains a data and a MAC.
+    /// The MAC is used to verify integrity of the encrypted next layer.
+    fn single_layer_size(&self) -> usize {
+        self.data_size + HEADER_INTEGRITY_MAC_SIZE
+    }
+
+    /// Perform the layered encryption.
+    pub fn encrypt(&self, params: &[EncryptionParam]) -> Result<(Vec<u8>, HeaderIntegrityMac)> {
+        if params.is_empty() || params.len() > self.max_layers {
+            return Err(Error::InvalidParams);
+        }
+
+        params
+            .iter()
+            .take(params.len() - 1) // Exclude the last param that will be treated separately below.
+            .rev() // Data and keys must be used in reverse order to encrypt the inner-most layer first
+            .try_fold(self.build_last_layer(params)?, |(encrypted, mac), param| {
+                self.build_intermediate_layer(param, mac, encrypted)
+            })
+    }
+
+    /// Build an intermediate layer of encryption that wraps subsequent layers already encrypted.
+    /// The output has the same size as [`Self::total_size`],
+    /// regardless of how many subsequent layers that this layer wraps.
+    fn build_intermediate_layer(
+        &self,
+        param: &EncryptionParam,
+        next_mac: HeaderIntegrityMac,
+        next_encrypted_data: Vec<u8>,
+    ) -> Result<(Vec<u8>, HeaderIntegrityMac)> {
+        if param.data.len() != self.data_size {
+            return Err(Error::InvalidDataLength);
+        }
+
+        // Concatenate the data with the encrypted subsequent layers and its MAC.
+        let total_data = concat_bytes(&[
+            &param.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()],
+        ]);
+
+        // Encrypt the concatenated bytes, and compute MAC.
+        let mut encrypted = total_data;
+        self.apply_streamcipher(
+            &mut encrypted,
+            &param.key.stream_cipher_key,
+            StreamCipherOption::FromFront,
+        );
+        let mac = Self::compute_mac(&param.key.integrity_mac_key, &encrypted);
+
+        assert_eq!(encrypted.len(), self.total_size());
+        Ok((encrypted, mac))
+    }
+
+    /// Build the last layer of encryption.
+    /// The output has the same size as [`Self::total_size`] by using fillers,
+    /// even though it doesn't wrap any subsequent layer.
+    /// This is for the length-preserved decryption.
+    fn build_last_layer(
+        &self,
+        params: &[EncryptionParam],
+    ) -> Result<(Vec<u8>, HeaderIntegrityMac)> {
+        let last_param = params.last().ok_or(Error::InvalidParams)?;
+        if last_param.data.len() != self.data_size {
+            return Err(Error::InvalidDataLength);
+        }
+
+        // Build fillers that will be appended to the last data.
+        // The number of fillers must be the same as the number of intermediate layers
+        // that will be decrypted later.
+        let fillers = self.build_fillers(&params[..params.len() - 1]);
+        // Random bytes is used to fill the space between data and fillers.
+        let random_bytes = random_bytes(self.total_size() - self.data_size - fillers.len());
+
+        // First, concat the data and the random bytes, and encrypt it.
+        let total_data_without_fillers = concat_bytes(&[&last_param.data, &random_bytes]);
+        let mut encrypted = total_data_without_fillers;
+        self.apply_streamcipher(
+            &mut encrypted,
+            &last_param.key.stream_cipher_key,
+            StreamCipherOption::FromFront,
+        );
+
+        // Append fillers to the encrypted bytes, and compute MAC.
+        encrypted.extend(fillers);
+        let mac = Self::compute_mac(&last_param.key.integrity_mac_key, &encrypted);
+
+        assert_eq!(encrypted.len(), self.total_size());
+        Ok((encrypted, mac))
+    }
+
+    /// Build as many fillers as the number of keys provided.
+    /// Fillers are encrypted in accumulated manner by keys.
+    fn build_fillers(&self, params: &[EncryptionParam]) -> Vec<u8> {
+        params
+            .iter()
+            .map(|param| &param.key.stream_cipher_key)
+            .fold(Vec::new(), |mut fillers, key| {
+                fillers.extend(vec![0u8; self.single_layer_size()]);
+                self.apply_streamcipher(&mut fillers, key, StreamCipherOption::FromBack);
+                fillers
+            })
+    }
+
+    /// Unpack one layer of encryption by performing the length-preserved decryption.
+    pub fn unpack(
+        &self,
+        mac: &HeaderIntegrityMac,
+        encrypted_total_data: &[u8],
+        key: &Key,
+    ) -> Result<(Vec<u8>, HeaderIntegrityMac, Vec<u8>)> {
+        if encrypted_total_data.len() != self.total_size() {
+            return Err(Error::InvalidDataLength);
+        }
+        // If a wrong key is used, the decryption should fail.
+        if !mac.verify(key.integrity_mac_key, encrypted_total_data) {
+            return Err(Error::IntegrityMacVerificationFailed);
+        }
+
+        // Extend the encrypted data by the length of a single layer
+        // in order to restore the truncated part (a encrypted filler)
+        // by [`Self::build_intermediate_layer`] during the encryption process.
+        let total_data_with_zero_filler = encrypted_total_data
+            .iter()
+            .copied()
+            .chain(std::iter::repeat(0u8).take(self.single_layer_size()))
+            .collect::<Vec<_>>();
+
+        // Decrypt the extended data.
+        let mut decrypted = total_data_with_zero_filler;
+        self.apply_streamcipher(
+            &mut decrypted,
+            &key.stream_cipher_key,
+            StreamCipherOption::FromFront,
+        );
+
+        // Parse the decrypted data into 3 parts: data, MAC, and the next encrypted data.
+        let parsed = parse_bytes(
+            &decrypted,
+            &[self.data_size, HEADER_INTEGRITY_MAC_SIZE, self.total_size()],
+        )
+        .unwrap();
+        let data = parsed[0].to_vec();
+        let next_mac = HeaderIntegrityMac::from_bytes(parsed[1].try_into().unwrap());
+        let next_encrypted_data = parsed[2].to_vec();
+        Ok((data, next_mac, next_encrypted_data))
+    }
+
+    fn apply_streamcipher(&self, data: &mut [u8], key: &StreamCipherKey, opt: StreamCipherOption) {
+        let pseudorandom_bytes = sphinx_packet::crypto::generate_pseudorandom_bytes(
+            key,
+            &STREAM_CIPHER_INIT_VECTOR,
+            self.total_size() + self.single_layer_size(),
+        );
+        let pseudorandom_bytes = match opt {
+            StreamCipherOption::FromFront => &pseudorandom_bytes[..data.len()],
+            StreamCipherOption::FromBack => {
+                &pseudorandom_bytes[pseudorandom_bytes.len() - data.len()..]
+            }
+        };
+        Self::xor(data, pseudorandom_bytes)
+    }
+
+    // In-place XOR operation: b is applied to a.
+    fn xor(a: &mut [u8], b: &[u8]) {
+        assert_eq!(a.len(), b.len());
+        a.iter_mut().zip(b.iter()).for_each(|(x1, &x2)| *x1 ^= x2);
+    }
+
+    fn compute_mac(key: &HeaderIntegrityMacKey, data: &[u8]) -> HeaderIntegrityMac {
+        let mac = sphinx_packet::crypto::compute_keyed_hmac::<sha2::Sha256>(key, data).into_bytes();
+        assert!(mac.len() >= HEADER_INTEGRITY_MAC_SIZE);
+        HeaderIntegrityMac::from_bytes(
+            mac.into_iter()
+                .take(HEADER_INTEGRITY_MAC_SIZE)
+                .collect::<Vec<_>>()
+                .try_into()
+                .unwrap(),
+        )
+    }
+}
+
+fn random_bytes(size: usize) -> Vec<u8> {
+    let mut bytes = vec![0u8; size];
+    let mut rng = ChaCha12Rng::from_entropy();
+    rng.fill_bytes(&mut bytes);
+    bytes
+}
+
+enum StreamCipherOption {
+    FromFront,
+    FromBack,
+}
+
+#[cfg(test)]
+mod tests {
+    use sphinx_packet::{constants::INTEGRITY_MAC_KEY_SIZE, crypto::STREAM_CIPHER_KEY_SIZE};
+
+    use super::*;
+
+    #[test]
+    fn build_and_unpack() {
+        let data_size = 10;
+        let cipher = ConsistentLengthLayeredCipher::new(5, data_size);
+
+        let params = (0u8..3)
+            .map(|i| EncryptionParam {
+                data: vec![i; data_size],
+                key: Key {
+                    stream_cipher_key: [i * 10; STREAM_CIPHER_KEY_SIZE],
+                    integrity_mac_key: [i * 20; INTEGRITY_MAC_KEY_SIZE],
+                },
+            })
+            .collect::<Vec<_>>();
+
+        let (encrypted, mac) = cipher.encrypt(&params).unwrap();
+
+        let next_encrypted = encrypted.clone();
+        let (data, next_mac, next_encrypted) = cipher
+            .unpack(&mac, &next_encrypted, &params[0].key)
+            .unwrap();
+        assert_eq!(data, params[0].data);
+        assert_eq!(next_encrypted.len(), encrypted.len());
+
+        let (data, next_mac, next_encrypted) = cipher
+            .unpack(&next_mac, &next_encrypted, &params[1].key)
+            .unwrap();
+        assert_eq!(data, params[1].data);
+        assert_eq!(next_encrypted.len(), encrypted.len());
+
+        let (data, _, next_encrypted) = cipher
+            .unpack(&next_mac, &next_encrypted, &params[2].key)
+            .unwrap();
+        assert_eq!(data, params[2].data);
+        assert_eq!(next_encrypted.len(), encrypted.len());
+    }
+}

nomos-mix/message/src/lib.rs

@@ -1,5 +1,7 @@
 mod error;
+mod layered_cipher;
 pub mod packet;
+mod routing;
 
 pub use error::Error;
 
@@ -62,3 +64,26 @@ pub fn unwrap_message(message: &[u8]) -> Result<(Vec<u8>, bool), Error> {
 pub fn is_drop_message(message: &[u8]) -> bool {
     message == DROP_MESSAGE
 }
+
+pub(crate) 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
+}
+
+pub(crate) fn parse_bytes<'a>(data: &'a [u8], sizes: &[usize]) -> Result<Vec<&'a [u8]>, String> {
+    let mut i = 0;
+    sizes
+        .iter()
+        .map(|&size| {
+            if i + size > data.len() {
+                return Err("The sum of sizes exceeds the length of the input slice".to_string());
+            }
+            let slice = &data[i..i + size];
+            i += size;
+            Ok(slice)
+        })
+        .collect()
+}

nomos-mix/message/src/packet.rs

@@ -1,11 +1,17 @@
-use crate::Error;
-use serde::{Deserialize, Serialize};
-use sphinx_packet::constants::NODE_ADDRESS_LENGTH;
+use crate::{routing::EncryptedRoutingInformation, Error};
+use sphinx_packet::{
+    constants::NODE_ADDRESS_LENGTH,
+    header::{
+        keys::RoutingKeys,
+        routing::{FINAL_HOP, FORWARD_HOP},
+    },
+    payload::Payload,
+};
 
 /// A packet that contains a header and a payload.
 /// The header and payload are encrypted for the selected recipients.
 /// This packet can be serialized and sent over the network.
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug)]
 pub struct Packet {
     header: Header,
     // This crate doesn't limit the payload size.
@@ -14,29 +20,18 @@ pub struct Packet {
 }
 
 /// The packet header
-#[derive(Debug, Serialize, Deserialize)]
+#[derive(Debug)]
 struct Header {
     /// The ephemeral public key for a recipient to derive the shared secret
     /// which can be used to decrypt the header and payload.
     ephemeral_public_key: x25519_dalek::PublicKey,
-    // TODO: Length-preserved layered encryption on RoutingInfo
-    routing_info: RoutingInfo,
-}
-
-#[derive(Debug, Serialize, Deserialize)]
-struct RoutingInfo {
-    // TODO: Change this to `is_final_layer: bool`
-    // by implementing the length-preserved layered encryption.
-    // It's not good to expose the info that how many layers remain to the intermediate recipients.
-    remaining_layers: u8,
-    // TODO:: Add the following fields
-    // header_integrity_hamc
-    // additional data (e.g. incentivization)
+    encrypted_routing_info: EncryptedRoutingInformation,
 }
 
 impl Packet {
     pub fn build(
         recipient_pubkeys: &[x25519_dalek::PublicKey],
+        max_layers: usize,
         payload: &[u8],
         payload_size: usize,
     ) -> Result<Self, Error> {
@@ -44,6 +39,10 @@ impl Packet {
         let ephemeral_privkey = x25519_dalek::StaticSecret::random();
         let key_material = Self::derive_key_material(recipient_pubkeys, &ephemeral_privkey);
 
+        // Build the encrypted routing information.
+        let encrypted_routing_info =
+            EncryptedRoutingInformation::new(&key_material.routing_keys, max_layers)?;
+
         // Encrypt the payload for all recipients.
         let payload_keys = key_material
             .routing_keys
@@ -59,16 +58,13 @@ impl Packet {
         Ok(Packet {
             header: Header {
                 ephemeral_public_key: x25519_dalek::PublicKey::from(&ephemeral_privkey),
-                routing_info: RoutingInfo {
-                    remaining_layers: u8::try_from(recipient_pubkeys.len())
-                        .map_err(|_| Error::TooManyRecipients)?,
-                },
+                encrypted_routing_info,
             },
             payload: payload.into_bytes(),
         })
     }
 
-    fn derive_key_material(
+    pub(crate) fn derive_key_material(
         recipient_pubkeys: &[x25519_dalek::PublicKey],
         ephemeral_privkey: &x25519_dalek::StaticSecret,
     ) -> sphinx_packet::header::keys::KeyMaterial {
@@ -90,6 +86,7 @@ impl Packet {
     pub fn unpack(
         &self,
         private_key: &x25519_dalek::StaticSecret,
+        max_layers: usize,
     ) -> Result<UnpackedPacket, Error> {
         // Derive the routing keys for the recipient
         let routing_keys = sphinx_packet::header::SphinxHeader::compute_routing_keys(
@@ -101,25 +98,40 @@ impl Packet {
         let payload = sphinx_packet::payload::Payload::from_bytes(&self.payload)?;
         let payload = payload.unwrap(&routing_keys.payload_key)?;
 
-        // If this is the last layer of encryption, return the decrypted payload.
-        if self.header.routing_info.remaining_layers == 1 {
-            return Ok(UnpackedPacket::FullyUnpacked(payload.recover_plaintext()?));
+        // Unpack the routing information
+        let (routing_info, next_encrypted_routing_info) = self
+            .header
+            .encrypted_routing_info
+            .unpack(&routing_keys, max_layers)?;
+        match routing_info.flag {
+            FORWARD_HOP => Ok(UnpackedPacket::ToForward(self.build_next_packet(
+                &routing_keys,
+                next_encrypted_routing_info,
+                payload,
+            ))),
+            FINAL_HOP => Ok(UnpackedPacket::FullyUnpacked(payload.recover_plaintext()?)),
+            _ => return Err(Error::InvalidRoutingFlag(routing_info.flag)),
+        }
     }
 
+    fn build_next_packet(
+        &self,
+        routing_keys: &RoutingKeys,
+        next_encrypted_routing_info: EncryptedRoutingInformation,
+        payload: Payload,
+    ) -> Packet {
         // Derive the new ephemeral public key for the next recipient
         let next_ephemeral_pubkey = Self::derive_next_ephemeral_public_key(
             &self.header.ephemeral_public_key,
             &routing_keys.blinding_factor,
         );
-        Ok(UnpackedPacket::ToForward(Packet {
+        Packet {
             header: Header {
                 ephemeral_public_key: next_ephemeral_pubkey,
-                routing_info: RoutingInfo {
-                    remaining_layers: self.header.routing_info.remaining_layers - 1,
-                },
+                encrypted_routing_info: next_encrypted_routing_info,
             },
             payload: payload.into_bytes(),
-        }))
+        }
     }
 
     /// Derive the next ephemeral public key for the next recipient.
@@ -135,6 +147,38 @@ impl Packet {
         let new_shared_secret = blinding_factor.diffie_hellman(cur_ephemeral_pubkey);
         x25519_dalek::PublicKey::from(new_shared_secret.to_bytes())
     }
+
+    pub fn to_bytes(&self) -> Vec<u8> {
+        let mut bytes = Vec::new();
+        bytes.extend_from_slice(&self.header.ephemeral_public_key.to_bytes());
+        bytes.extend_from_slice(&self.header.encrypted_routing_info.to_bytes());
+        bytes.extend_from_slice(&self.payload);
+        bytes
+    }
+
+    pub fn from_bytes(data: &[u8], max_layers: usize) -> Result<Self, Error> {
+        let mut i = 0;
+        let public_key_bytes: [u8; 32] = data[i..i + 32].try_into().unwrap();
+        let ephemeral_public_key = x25519_dalek::PublicKey::from(public_key_bytes);
+        i += 32;
+
+        let encrypted_routing_info_size = EncryptedRoutingInformation::size(max_layers);
+        let encrypted_routing_info = EncryptedRoutingInformation::from_bytes(
+            &data[i..i + encrypted_routing_info_size],
+            max_layers,
+        )?;
+        i += encrypted_routing_info_size;
+
+        let payload = data[i..].to_vec();
+
+        Ok(Packet {
+            header: Header {
+                ephemeral_public_key,
+                encrypted_routing_info,
+            },
+            payload,
+        })
+    }
 }
 
 pub enum UnpackedPacket {
@@ -144,14 +188,12 @@ pub enum UnpackedPacket {
 
 #[cfg(test)]
 mod tests {
-    use nomos_core::wire;
-
     use super::*;
 
     #[test]
     fn unpack() {
         // Prepare keys of two recipients
-        let recipient_privkeys = (0..2)
+        let recipient_privkeys = (0..3)
             .map(|_| x25519_dalek::StaticSecret::random())
             .collect::<Vec<_>>();
         let recipient_pubkeys = recipient_privkeys
@@ -160,18 +202,26 @@ mod tests {
             .collect::<Vec<_>>();
 
         // Build a packet
+        let max_layers = 5;
         let payload = [10u8; 512];
-        let packet = Packet::build(&recipient_pubkeys, &payload, 1024).unwrap();
+        let packet = Packet::build(&recipient_pubkeys, max_layers, &payload, 1024).unwrap();
 
         // The 1st recipient unpacks the packet
-        let packet = match packet.unpack(&recipient_privkeys[0]).unwrap() {
+        let packet = match packet.unpack(&recipient_privkeys[0], max_layers).unwrap() {
             UnpackedPacket::ToForward(packet) => packet,
             UnpackedPacket::FullyUnpacked(_) => {
                 panic!("The unpacked packet should be the ToFoward type");
             }
         };
         // The 2nd recipient unpacks the packet
-        match packet.unpack(&recipient_privkeys[1]).unwrap() {
+        let packet = match packet.unpack(&recipient_privkeys[1], max_layers).unwrap() {
+            UnpackedPacket::ToForward(packet) => packet,
+            UnpackedPacket::FullyUnpacked(_) => {
+                panic!("The unpacked packet should be the ToFoward type");
+            }
+        };
+        // The last recipient unpacks the packet
+        match packet.unpack(&recipient_privkeys[2], max_layers).unwrap() {
             UnpackedPacket::ToForward(_) => {
                 panic!("The unpacked packet should be the FullyUnpacked type");
             }
@@ -185,34 +235,25 @@ mod tests {
     #[test]
     fn unpack_with_wrong_keys() {
         // Build a packet with two public keys
+        let max_layers = 5;
         let payload = [10u8; 512];
         let packet = Packet::build(
             &(0..2)
                 .map(|_| x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::random()))
                 .collect::<Vec<_>>(),
+            max_layers,
             &payload,
             1024,
         )
         .unwrap();
 
-        // The 1st recipient unpacks the packet with an wrong key
-        let packet = match packet
-            .unpack(&x25519_dalek::StaticSecret::random())
-            .unwrap()
-        {
-            UnpackedPacket::ToForward(packet) => packet,
-            UnpackedPacket::FullyUnpacked(_) => {
-                panic!("The unpacked packet should be the ToFoward type");
-            }
-        };
-        // The 2nd recipient unpacks the packet with an wrong key
         assert!(packet
-            .unpack(&x25519_dalek::StaticSecret::random())
+            .unpack(&x25519_dalek::StaticSecret::random(), max_layers)
             .is_err());
     }
 
     #[test]
-    fn consistent_size_serialization() {
+    fn consistent_size_after_unpack() {
         // Prepare keys of two recipients
         let recipient_privkeys = (0..2)
             .map(|_| x25519_dalek::StaticSecret::random())
@@ -223,28 +264,47 @@ mod tests {
             .collect::<Vec<_>>();
 
         // Build a packet
+        let max_layers = 5;
         let payload = [10u8; 512];
-        let packet = Packet::build(&recipient_pubkeys, &payload, 1024).unwrap();
+        let packet = Packet::build(&recipient_pubkeys, max_layers, &payload, 1024).unwrap();
 
         // Calculate the expected packet size
         let pubkey_size = 32;
-        let routing_info_size = 1;
-        let payload_length_enconding_size = 8;
         let payload_size = 1024;
         let packet_size =
-            pubkey_size + routing_info_size + payload_length_enconding_size + payload_size;
+            pubkey_size + EncryptedRoutingInformation::size(max_layers) + payload_size;
 
         // The serialized packet size must be the same as the expected size.
-        assert_eq!(wire::serialize(&packet).unwrap().len(), packet_size);
+        assert_eq!(packet.to_bytes().len(), packet_size);
 
         // The unpacked packet size must be the same as the original packet size.
-        match packet.unpack(&recipient_privkeys[0]).unwrap() {
+        match packet.unpack(&recipient_privkeys[0], max_layers).unwrap() {
             UnpackedPacket::ToForward(packet) => {
-                assert_eq!(wire::serialize(&packet).unwrap().len(), packet_size);
+                assert_eq!(packet.to_bytes().len(), packet_size);
             }
             UnpackedPacket::FullyUnpacked(_) => {
                 panic!("The unpacked packet should be the ToFoward type");
             }
         }
     }
+
+    #[test]
+    fn consistent_size_with_any_num_layers() {
+        let max_layers = 5;
+        let payload = [10u8; 512];
+
+        // Build a packet with 2 recipients
+        let recipient_pubkeys = (0..2)
+            .map(|_| x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::random()))
+            .collect::<Vec<_>>();
+        let packet1 = Packet::build(&recipient_pubkeys, max_layers, &payload, 1024).unwrap();
+
+        // Build a packet with 3 recipients
+        let recipient_pubkeys = (0..3)
+            .map(|_| x25519_dalek::PublicKey::from(&x25519_dalek::StaticSecret::random()))
+            .collect::<Vec<_>>();
+        let packet2 = Packet::build(&recipient_pubkeys, max_layers, &payload, 1024).unwrap();
+
+        assert_eq!(packet1.to_bytes().len(), packet2.to_bytes().len());
+    }
 }

nomos-mix/message/src/routing.rs

@@ -0,0 +1,139 @@
+use sphinx_packet::{
+    constants::HEADER_INTEGRITY_MAC_SIZE,
+    header::{
+        keys::RoutingKeys,
+        mac::HeaderIntegrityMac,
+        routing::{RoutingFlag, FINAL_HOP, FORWARD_HOP},
+    },
+};
+
+use crate::{
+    concat_bytes,
+    layered_cipher::{ConsistentLengthLayeredCipher, EncryptionParam, Key},
+    parse_bytes, Error,
+};
+
+/// A routing information that will be contained in a packet header
+/// in the encrypted format.
+pub struct RoutingInformation {
+    pub flag: RoutingFlag,
+    // Add additional fields here
+}
+
+impl RoutingInformation {
+    pub fn new(flag: RoutingFlag) -> Self {
+        Self { flag }
+    }
+
+    pub fn to_bytes(&self) -> Vec<u8> {
+        vec![self.flag]
+    }
+
+    pub fn from_bytes(data: &[u8]) -> Result<Self, Error> {
+        if data.len() != Self::size() {
+            return Err(Error::InvalidEncryptedRoutingInfoLength(data.len()));
+        }
+        Ok(Self { flag: data[0] })
+    }
+
+    pub fn size() -> usize {
+        std::mem::size_of::<RoutingFlag>()
+    }
+}
+
+/// Encrypted routing information that will be contained in a packet header.
+#[derive(Debug)]
+pub struct EncryptedRoutingInformation {
+    /// A MAC to verify the integrity of [`Self::encrypted_routing_info`].
+    mac: HeaderIntegrityMac,
+    /// The actual encrypted routing information produced by [`ConsistentLengthLayeredCipher`].
+    /// Its size should be the same as [`ConsistentLengthLayeredCipher::total_size`].
+    encrypted_routing_info: Vec<u8>,
+}
+
+impl EncryptedRoutingInformation {
+    /// Build all [`RoutingInformation`]s for the provides keys,
+    /// and encrypt them using [`ConsistentLengthLayeredCipher`].
+    pub fn new(routing_keys: &[RoutingKeys], max_layers: usize) -> Result<Self, Error> {
+        let cipher = Self::layered_cipher(max_layers);
+        let params = routing_keys
+            .iter()
+            .enumerate()
+            .map(|(i, k)| {
+                let flag = if i == routing_keys.len() - 1 {
+                    FINAL_HOP
+                } else {
+                    FORWARD_HOP
+                };
+                EncryptionParam {
+                    data: RoutingInformation::new(flag).to_bytes(),
+                    key: Self::layered_cipher_key(k),
+                }
+            })
+            .collect::<Vec<_>>();
+        let (encrypted, mac) = cipher.encrypt(&params)?;
+
+        Ok(Self {
+            mac,
+            encrypted_routing_info: encrypted,
+        })
+    }
+
+    /// Unpack one layer of encryptions using the key provided.
+    /// Returns the decrypted routing information
+    /// and the next [`EncryptedRoutingInformation`] to be unpacked further.
+    pub fn unpack(
+        &self,
+        routing_key: &RoutingKeys,
+        max_layers: usize,
+    ) -> Result<(RoutingInformation, Self), Error> {
+        let cipher = Self::layered_cipher(max_layers);
+        let (routing_info, next_mac, next_encrypted_routing_info) = cipher.unpack(
+            &self.mac,
+            &self.encrypted_routing_info,
+            &Self::layered_cipher_key(routing_key),
+        )?;
+        Ok((
+            RoutingInformation::from_bytes(&routing_info)?,
+            Self {
+                mac: next_mac,
+                encrypted_routing_info: next_encrypted_routing_info,
+            },
+        ))
+    }
+
+    fn layered_cipher(max_layers: usize) -> ConsistentLengthLayeredCipher {
+        ConsistentLengthLayeredCipher::new(max_layers, RoutingInformation::size())
+    }
+
+    fn layered_cipher_key(routing_key: &RoutingKeys) -> Key {
+        Key {
+            stream_cipher_key: routing_key.stream_cipher_key,
+            integrity_mac_key: routing_key.header_integrity_hmac_key,
+        }
+    }
+
+    pub fn to_bytes(&self) -> Vec<u8> {
+        concat_bytes(&[self.mac.as_bytes(), &self.encrypted_routing_info])
+    }
+
+    pub fn from_bytes(data: &[u8], max_layers: usize) -> Result<Self, Error> {
+        let parsed = parse_bytes(
+            data,
+            &[
+                HEADER_INTEGRITY_MAC_SIZE,
+                Self::layered_cipher(max_layers).total_size(),
+            ],
+        )
+        .map_err(|_| Error::InvalidEncryptedRoutingInfoLength(data.len()))?;
+        Ok(Self {
+            mac: HeaderIntegrityMac::from_bytes(parsed[0].try_into().unwrap()),
+            encrypted_routing_info: parsed[1].to_vec(),
+        })
+    }
+
+    pub fn size(max_layers: usize) -> usize {
+        let cipher = ConsistentLengthLayeredCipher::new(max_layers, RoutingInformation::size());
+        HEADER_INTEGRITY_MAC_SIZE + cipher.total_size()
+    }
+}