Use proper size for drop message

nomos-mix/message/src/sphinx/layered_cipher.rs

@@ -82,9 +82,9 @@ impl<D: ConsistentLengthLayeredCipherData> ConsistentLengthLayeredCipher<D> {
     }
 
     /// The total size of fully encrypted output that includes all layers.
-    /// This size is determined by [`D::size`] and [`Self::max_layers`].
+    /// This size is determined by [`D::size`] and [`max_layers`].
-    pub fn total_size(&self) -> usize {
+    pub const fn total_size(max_layers: usize) -> usize {
-        Self::SINGLE_LAYER_SIZE * self.max_layers
+        Self::SINGLE_LAYER_SIZE * max_layers
     }
 
     /// The size of a single layer that contains a data and a MAC.
@@ -136,7 +136,7 @@ impl<D: ConsistentLengthLayeredCipherData> ConsistentLengthLayeredCipher<D> {
         );
         let mac = Self::compute_mac(&param.key.integrity_mac_key, &encrypted);
 
-        assert_eq!(encrypted.len(), self.total_size());
+        assert_eq!(encrypted.len(), Self::total_size(self.max_layers));
         Ok((encrypted, mac))
     }
 
@@ -158,7 +158,8 @@ impl<D: ConsistentLengthLayeredCipherData> ConsistentLengthLayeredCipher<D> {
         // because there is no next encrypted layer.
         // Instead, random bytes are used to fill the space between data and fillers.
         // The size of random bytes depends on the [`self.max_layers`].
-        let random_bytes = random_bytes(self.total_size() - D::SIZE - fillers.len());
+        let random_bytes =
+            random_bytes(Self::total_size(self.max_layers) - D::SIZE - fillers.len());
 
         // First, concat the data and the random bytes, and encrypt it.
         let last_data = last_param.data.to_bytes();
@@ -176,7 +177,7 @@ impl<D: ConsistentLengthLayeredCipherData> ConsistentLengthLayeredCipher<D> {
         encrypted.extend(fillers);
         let mac = Self::compute_mac(&last_param.key.integrity_mac_key, &encrypted);
 
-        assert_eq!(encrypted.len(), self.total_size());
+        assert_eq!(encrypted.len(), Self::total_size(self.max_layers));
         Ok((encrypted, mac))
     }
 
@@ -205,7 +206,7 @@ impl<D: ConsistentLengthLayeredCipherData> ConsistentLengthLayeredCipher<D> {
         encrypted_total_data: &[u8],
         key: &Key,
     ) -> Result<(Vec<u8>, HeaderIntegrityMac, Vec<u8>)> {
-        if encrypted_total_data.len() != self.total_size() {
+        if encrypted_total_data.len() != Self::total_size(self.max_layers) {
             return Err(Error::InvalidCipherTextLength);
         }
         // If a wrong key is used, the decryption should fail.
@@ -233,7 +234,11 @@ impl<D: ConsistentLengthLayeredCipherData> ConsistentLengthLayeredCipher<D> {
         // Parse the decrypted data into 3 parts: data, MAC, and the next encrypted data.
         let parsed = parse_bytes(
             &decrypted,
-            &[D::SIZE, HEADER_INTEGRITY_MAC_SIZE, self.total_size()],
+            &[
+                D::SIZE,
+                HEADER_INTEGRITY_MAC_SIZE,
+                Self::total_size(self.max_layers),
+            ],
         )
         .unwrap();
         let data = parsed[0].to_vec();
@@ -246,7 +251,7 @@ impl<D: ConsistentLengthLayeredCipherData> ConsistentLengthLayeredCipher<D> {
         let pseudorandom_bytes = sphinx_packet::crypto::generate_pseudorandom_bytes(
             key,
             &STREAM_CIPHER_INIT_VECTOR,
-            self.total_size() + Self::SINGLE_LAYER_SIZE,
+            Self::total_size(self.max_layers) + Self::SINGLE_LAYER_SIZE,
         );
         let pseudorandom_bytes = match opt {
             StreamCipherOption::FromFront => &pseudorandom_bytes[..data.len()],

nomos-mix/message/src/sphinx/mod.rs

@@ -20,9 +20,7 @@ impl MixMessage for SphinxMessage {
     type PrivateKey = [u8; ASYM_KEY_SIZE];
     type Error = Error;
 
-    // TODO: Remove DROP_MESSAGE. Currently, an arbitrary size (2048) is used,
+    const DROP_MESSAGE: &'static [u8] = &[0; Packet::size(MAX_LAYERS, PADDED_PAYLOAD_SIZE)];
-    // but we've decided to remove drop messages from the spec.
-    const DROP_MESSAGE: &'static [u8] = &[0; 2048];
 
     fn build_message(
         payload: &[u8],

nomos-mix/message/src/sphinx/packet.rs

@@ -189,6 +189,10 @@ impl Packet {
             payload: parsed[2].to_vec(),
         })
     }
+
+    pub const fn size(max_layers: usize, max_payload_size: usize) -> usize {
+        ASYM_KEY_SIZE + EncryptedRoutingInformation::size(max_layers) + max_payload_size
+    }
 }
 
 pub enum UnpackedPacket {
@@ -281,8 +285,7 @@ mod tests {
             Packet::build(&recipient_pubkeys, max_layers, &payload, max_payload_size).unwrap();
 
         // Calculate the expected packet size
-        let packet_size =
+        let packet_size = Packet::size(max_layers, max_payload_size);
-            ASYM_KEY_SIZE + EncryptedRoutingInformation::size(max_layers) + max_payload_size;
 
         // The serialized packet size must be the same as the expected size.
         assert_eq!(packet.to_bytes().len(), packet_size);

nomos-mix/message/src/sphinx/routing.rs

@@ -53,11 +53,13 @@ pub struct EncryptedRoutingInformation {
     pub encrypted_routing_info: Vec<u8>,
 }
 
+type LayeredCipher = ConsistentLengthLayeredCipher<RoutingInformation>;
+
 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 cipher = LayeredCipher::new(max_layers);
         let params = routing_keys
             .iter()
             .enumerate()
@@ -89,7 +91,7 @@ impl EncryptedRoutingInformation {
         routing_key: &RoutingKeys,
         max_layers: usize,
     ) -> Result<(RoutingInformation, Self), Error> {
-        let cipher = Self::layered_cipher(max_layers);
+        let cipher = LayeredCipher::new(max_layers);
         let (routing_info, next_mac, next_encrypted_routing_info) = cipher.unpack(
             &self.mac,
             &self.encrypted_routing_info,
@@ -104,10 +106,6 @@ impl EncryptedRoutingInformation {
         ))
     }
 
-    fn layered_cipher(max_layers: usize) -> ConsistentLengthLayeredCipher<RoutingInformation> {
-        ConsistentLengthLayeredCipher::<RoutingInformation>::new(max_layers)
-    }
-
     fn layered_cipher_key(routing_key: &RoutingKeys) -> Key {
         Key {
             stream_cipher_key: routing_key.stream_cipher_key,
@@ -126,7 +124,7 @@ impl EncryptedRoutingInformation {
             data,
             &[
                 HEADER_INTEGRITY_MAC_SIZE,
-                Self::layered_cipher(max_layers).total_size(),
+                LayeredCipher::total_size(max_layers),
             ],
         )
         .map_err(|_| Error::InvalidEncryptedRoutingInfoLength(data.len()))?;
@@ -136,7 +134,7 @@ impl EncryptedRoutingInformation {
         })
     }
 
-    pub fn size(max_layers: usize) -> usize {
+    pub const fn size(max_layers: usize) -> usize {
-        HEADER_INTEGRITY_MAC_SIZE + Self::layered_cipher(max_layers).total_size()
+        HEADER_INTEGRITY_MAC_SIZE + LayeredCipher::total_size(max_layers)
     }
 }