adjusted message handler to allow for dual packet sizes

common/client-core/src/client/real_messages_control/message_handler.rs

@@ -180,8 +180,6 @@ where
             config.average_packet_delay,
             config.average_ack_delay,
         )
-        .with_custom_primary_packet_size(config.primary_packet_size)
-        .with_custom_secondary_packet_size(config.secondary_packet_size)
         .with_mix_hops(config.num_mix_hops);
 
         MessageHandler {
@@ -225,6 +223,7 @@ where
     fn optimal_packet_size(&self, msg: &NymMessage) -> PacketSize {
         // if secondary packet was never set, then it's obvious we have to use the primary packet
         let Some(secondary_packet) = self.config.secondary_packet_size else {
+            trace!("only primary packet size is available");
             return self.config.primary_packet_size
         };
 
@@ -232,10 +231,13 @@ where
             msg.required_packets(self.config.primary_packet_size, self.config.num_mix_hops);
         let secondary_count = msg.required_packets(secondary_packet, self.config.num_mix_hops);
 
+        trace!("This message would require: {primary_count} primary packets or {secondary_count} secondary packets...");
         // if there would be no benefit in using the secondary packet - use the primary (duh)
         if primary_count <= secondary_count {
+            trace!("so choosing primary for this message");
             self.config.primary_packet_size
         } else {
+            trace!("so choosing secondary for this message");
             secondary_packet
         }
     }

common/client-core/src/client/transmission_buffer.rs

@@ -171,8 +171,6 @@ impl<T> TransmissionBuffer<T> {
         &mut self,
         n: usize,
     ) -> Option<Vec<(TransmissionLane, T)>> {
-        // let start = Instant::now();
-
         if self.buffer.is_empty() {
             return None;
         }
@@ -187,8 +185,6 @@ impl<T> TransmissionBuffer<T> {
             items.push(next)
         }
 
-        // todo!("time time taken");
-
         Some(items)
     }
 

common/nymsphinx/anonymous-replies/src/reply_surb.rs

@@ -172,10 +172,8 @@ impl ReplySurb {
     pub fn apply_surb<M: AsRef<[u8]>>(
         self,
         message: M,
-        packet_size: Option<PacketSize>,
+        packet_size: PacketSize,
     ) -> Result<(SphinxPacket, NymNodeRoutingAddress), ReplySurbError> {
-        let packet_size = packet_size.unwrap_or_default();
-
         let message_bytes = message.as_ref();
         if message_bytes.len() != packet_size.plaintext_size() {
             return Err(ReplySurbError::UnpaddedMessageError);

common/nymsphinx/chunking/src/fragment.rs

@@ -191,6 +191,14 @@ impl Fragment {
         })
     }
 
+    /// based on the size of the embedded data, determines which predefined `PacketSize`
+    /// was used for construction of this `Fragment`
+    pub fn serialized_size(&self) -> usize {
+        // TODO: optimisation: determine the size of the header without actually serializing it...
+        let header_size = self.header.to_bytes().len();
+        header_size + self.payload_size()
+    }
+
     /// Convert this `Fragment` into vector of bytes which can be put into a sphinx packet.
     pub fn into_bytes(self) -> Vec<u8> {
         self.header
@@ -434,8 +442,7 @@ impl FragmentHeader {
         let frag_id = self.id | (1 << 31);
         let frag_id_bytes = frag_id.to_be_bytes();
         let bytes_prefix_iter = frag_id_bytes
-            .iter()
-            .cloned()
+            .into_iter()
             .chain(std::iter::once(self.total_fragments))
             .chain(std::iter::once(self.current_fragment));
 

common/nymsphinx/chunking/src/lib.rs

@@ -5,6 +5,8 @@ use crate::fragment::{linked_fragment_payload_max_len, unlinked_fragment_payload
 pub use set::split_into_sets;
 use thiserror::Error;
 
+pub const MIN_PADDING_OVERHEAD: usize = 1;
+
 // Future consideration: currently in a lot of places, the payloads have randomised content
 // which is not a perfect testing strategy as it might not detect some edge cases I never would
 // have assumed could be possible. A better approach would be to research some Fuzz testing

common/nymsphinx/params/Cargo.toml

@@ -9,7 +9,7 @@ repository = { workspace = true }
 
 [dependencies]
 thiserror = "1.0.37"
-serde = { workspace = true }
+serde = { workspace = true, features = ["derive"] }
 
 nym-crypto = { path = "../../crypto", features = ["hashing", "symmetric"] }
 nym-sphinx-types = { path = "../types" }

common/nymsphinx/params/src/packet_sizes.rs

@@ -7,21 +7,25 @@ use nym_sphinx_types::PAYLOAD_OVERHEAD_SIZE;
 use serde::{Deserialize, Serialize};
 use std::cmp::Ordering;
 use std::convert::TryFrom;
+use std::fmt::{Debug, Display, Formatter};
 use std::str::FromStr;
 use thiserror::Error;
 
+// each sphinx packet contains mandatory header and payload padding + markers
+const PACKET_OVERHEAD: usize = HEADER_SIZE + PAYLOAD_OVERHEAD_SIZE;
+
 // it's up to the smart people to figure those values out : )
-const REGULAR_PACKET_SIZE: usize = HEADER_SIZE + PAYLOAD_OVERHEAD_SIZE + 2 * 1024;
+const REGULAR_PACKET_SIZE: usize = 2 * 1024 + PACKET_OVERHEAD;
 // TODO: even though we have 16B IV, is having just 5B (FRAG_ID_LEN) of the ID possibly insecure?
 
 // TODO: I'm not entirely sure if we can easily extract `<AckEncryptionAlgorithm as NewStreamCipher>::NonceSize`
 // into a const usize before relevant stuff is stabilised in rust...
 const ACK_IV_SIZE: usize = 16;
 
-const ACK_PACKET_SIZE: usize = HEADER_SIZE + PAYLOAD_OVERHEAD_SIZE + ACK_IV_SIZE + FRAG_ID_LEN;
-const EXTENDED_PACKET_SIZE_8: usize = HEADER_SIZE + PAYLOAD_OVERHEAD_SIZE + 8 * 1024;
-const EXTENDED_PACKET_SIZE_16: usize = HEADER_SIZE + PAYLOAD_OVERHEAD_SIZE + 16 * 1024;
-const EXTENDED_PACKET_SIZE_32: usize = HEADER_SIZE + PAYLOAD_OVERHEAD_SIZE + 32 * 1024;
+const ACK_PACKET_SIZE: usize = ACK_IV_SIZE + FRAG_ID_LEN + PACKET_OVERHEAD;
+const EXTENDED_PACKET_SIZE_8: usize = 8 * 1024 + PACKET_OVERHEAD;
+const EXTENDED_PACKET_SIZE_16: usize = 16 * 1024 + PACKET_OVERHEAD;
+const EXTENDED_PACKET_SIZE_32: usize = 32 * 1024 + PACKET_OVERHEAD;
 
 #[derive(Debug, Error)]
 pub enum InvalidPacketSize {
@@ -36,7 +40,7 @@ pub enum InvalidPacketSize {
 }
 
 #[repr(u8)]
-#[derive(Clone, Copy, Debug, PartialEq, Eq, Default, Serialize, Deserialize)]
+#[derive(Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
 pub enum PacketSize {
     // for example instant messaging use case
     #[default]
@@ -91,6 +95,28 @@ impl FromStr for PacketSize {
     }
 }
 
+impl Display for PacketSize {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        match self {
+            PacketSize::RegularPacket => write!(f, "regular"),
+            PacketSize::AckPacket => write!(f, "ack"),
+            PacketSize::ExtendedPacket32 => write!(f, "extended32"),
+            PacketSize::ExtendedPacket8 => write!(f, "extended8"),
+            PacketSize::ExtendedPacket16 => write!(f, "extended16"),
+        }
+    }
+}
+
+impl Debug for PacketSize {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        let name = self.to_string();
+        let size = self.size();
+        let plaintext = self.plaintext_size();
+
+        write!(f, "{name} ({size} bytes / {plaintext} plaintext)")
+    }
+}
+
 impl TryFrom<u8> for PacketSize {
     type Error = InvalidPacketSize;
 
@@ -107,7 +133,7 @@ impl TryFrom<u8> for PacketSize {
 }
 
 impl PacketSize {
-    pub fn size(self) -> usize {
+    pub const fn size(self) -> usize {
         match self {
             PacketSize::RegularPacket => REGULAR_PACKET_SIZE,
             PacketSize::AckPacket => ACK_PACKET_SIZE,
@@ -117,11 +143,11 @@ impl PacketSize {
         }
     }
 
-    pub fn plaintext_size(self) -> usize {
+    pub const fn plaintext_size(self) -> usize {
         self.size() - HEADER_SIZE - PAYLOAD_OVERHEAD_SIZE
     }
 
-    pub fn payload_size(self) -> usize {
+    pub const fn payload_size(self) -> usize {
         self.size() - HEADER_SIZE
     }
 
@@ -157,6 +183,11 @@ impl PacketSize {
             None
         }
     }
+
+    pub fn get_type_from_plaintext(plaintext_size: usize) -> Result<Self, InvalidPacketSize> {
+        let packet_size = plaintext_size + PACKET_OVERHEAD;
+        Self::get_type(packet_size)
+    }
 }
 
 #[cfg(test)]

common/nymsphinx/src/message.rs

@@ -16,6 +16,8 @@ use rand::Rng;
 use std::fmt::{Display, Formatter};
 use thiserror::Error;
 
+pub(crate) const ACK_OVERHEAD: usize = MAX_NODE_ADDRESS_UNPADDED_LEN + PacketSize::AckPacket.size();
+
 #[derive(Debug, Error)]
 pub enum NymMessageError {
     #[error("{received} is not a valid type tag for a NymMessage")]
@@ -172,24 +174,9 @@ impl NymMessage {
         message_type_size + inner_size
     }
 
-    /// Determines the number of required packets of the provided size for the split message.
-    pub fn required_packets(&self, packet_size: PacketSize, num_mix_hops: u8) -> usize {
-        // let size = self.into_bytes()
-        let plaintext_per_packet = self.available_plaintext_per_packet(packet_size);
-        let serialized_len = self.serialized_size(num_mix_hops);
-
-        // same logic as in `pad_to_full_packet_lengths` for the additional '1' being added.
-        let (num_fragments, _) =
-            chunking::number_of_required_fragments(serialized_len + 1, plaintext_per_packet);
-
-        num_fragments
-    }
-
-    /// Length of plaintext (from the sphinx point of view) data that is available per sphinx
+    /// Length of plaintext (from the **sphinx** point of view) data that is available per sphinx
     /// packet.
-    pub fn available_plaintext_per_packet(&self, packet_size: PacketSize) -> usize {
-        let ack_overhead = MAX_NODE_ADDRESS_UNPADDED_LEN + PacketSize::AckPacket.size();
-
+    pub fn available_sphinx_plaintext_per_packet(&self, packet_size: PacketSize) -> usize {
         let variant_overhead = match self {
             // each plain or repliable packet attaches an ephemeral public key so that the recipient
             // could perform diffie-hellman with its own keys followed by a kdf to re-derive
@@ -200,7 +187,31 @@ impl NymMessage {
             NymMessage::Reply(_) => ReplySurbKeyDigestAlgorithm::output_size(),
         };
 
-        packet_size.plaintext_size() - ack_overhead - variant_overhead
+        // each packet will contain an ack + variant specific data (as described above)
+        packet_size.plaintext_size() - ACK_OVERHEAD - variant_overhead
+    }
+
+    /// Length of the actual (from the **message** point of view) data that is available in each packet.
+    pub fn true_available_plaintext_per_packet(&self, packet_size: PacketSize) -> usize {
+        let sphinx_plaintext = self.available_sphinx_plaintext_per_packet(packet_size);
+        sphinx_plaintext - chunking::MIN_PADDING_OVERHEAD
+    }
+
+    /// Determines the number of required packets of the provided size for the split message.
+    pub fn required_packets(&self, packet_size: PacketSize, num_mix_hops: u8) -> usize {
+        let plaintext_per_packet = self.true_available_plaintext_per_packet(packet_size);
+        let serialized_len = self.serialized_size(num_mix_hops);
+
+        let (num_fragments, _) =
+            chunking::number_of_required_fragments(serialized_len, plaintext_per_packet);
+
+        // by chunking I mean that currently the fragments hold variable amount of plaintext in them (I wish I had time to rewrite it...)
+        log::trace!(
+            "this message will use {serialized_len} bytes of PLAINTEXT (This does not account for Ack or chunking overhead). \
+            With {packet_size:?} PacketSize ({plaintext_per_packet} of usable plaintext available) it will require {num_fragments} packet(s).",
+        );
+
+        num_fragments
     }
 
     /// Pads the message so that after it gets chunked, it will occupy exactly N sphinx packets.
@@ -210,10 +221,14 @@ impl NymMessage {
 
         let bytes = self.into_bytes();
 
-        // 1 is added as there will always have to be at least a single byte of padding (1) added
+        // 1 (chunking::MIN_PADDING_OVERHEAD) is added as there will always have to be at least a single byte of padding (1) added
         // to be able to later distinguish the actual padding from the underlying message
+        // TODO: this whole `MIN_PADDING_OVERHEAD` feels very awkward. it should somehow be included in
+        // `available_plaintext_per_packet`
+        let total_required_bytes = bytes.len() + chunking::MIN_PADDING_OVERHEAD;
+
         let (packets_used, space_left) =
-            chunking::number_of_required_fragments(bytes.len() + 1, plaintext_per_packet);
+            chunking::number_of_required_fragments(total_required_bytes, plaintext_per_packet);
 
         let wasted_space = space_left as f32 / (bytes.len() + 1 + space_left) as f32;
         log::trace!("Padding {self_display}: {} of raw plaintext bytes are required. They're going to be put into {packets_used} sphinx packets with {space_left} bytes of leftover space. {wasted_space}% of packet capacity is going to be wasted.", bytes.len() + 1);

common/nymsphinx/src/preparer/mod.rs

@@ -1,8 +1,10 @@
 // Copyright 2021-2022 - Nym Technologies SA <contact@nymtech.net>
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::message::NymMessage;
+use crate::message::{NymMessage, ACK_OVERHEAD};
 use crate::NymsphinxPayloadBuilder;
+use nym_crypto::asymmetric::encryption;
+use nym_crypto::Digest;
 use nym_sphinx_acknowledgements::surb_ack::SurbAck;
 use nym_sphinx_acknowledgements::AckKey;
 use nym_sphinx_addressing::clients::Recipient;
@@ -11,7 +13,7 @@ use nym_sphinx_anonymous_replies::reply_surb::ReplySurb;
 use nym_sphinx_chunking::fragment::{Fragment, FragmentIdentifier};
 use nym_sphinx_forwarding::packet::MixPacket;
 use nym_sphinx_params::packet_sizes::PacketSize;
-use nym_sphinx_params::DEFAULT_NUM_MIX_HOPS;
+use nym_sphinx_params::{ReplySurbKeyDigestAlgorithm, DEFAULT_NUM_MIX_HOPS};
 use nym_sphinx_types::builder::SphinxPacketBuilder;
 use nym_sphinx_types::{delays, Delay};
 use nym_topology::{NymTopology, NymTopologyError};
@@ -45,12 +47,6 @@ pub struct MessagePreparer<R> {
     /// Instance of a cryptographically secure random number generator.
     rng: R,
 
-    /// Size of the target [`SphinxPacket`] into which the underlying is going to get split.
-    primary_packet_size: PacketSize,
-
-    /// Alternative size of the target [`SphinxPacket`] into which the underlying is going to get split.
-    secondary_packet_size: Option<PacketSize>,
-
     /// Address of this client which also represent an address to which all acknowledgements
     /// and surb-based are going to be sent.
     sender_address: Recipient,
@@ -82,8 +78,6 @@ where
             average_packet_delay,
             average_ack_delay,
             num_mix_hops: DEFAULT_NUM_MIX_HOPS,
-            primary_packet_size: PacketSize::RegularPacket,
-            secondary_packet_size: None,
         }
     }
 
@@ -93,18 +87,6 @@ where
         self
     }
 
-    /// Allows setting non-default size of the sphinx packets sent out.
-    pub fn with_custom_primary_packet_size(mut self, packet_size: PacketSize) -> Self {
-        self.primary_packet_size = packet_size;
-        self
-    }
-
-    /// Allows setting non-default size of the sphinx packets sent out.
-    pub fn with_custom_secondary_packet_size(mut self, packet_size: Option<PacketSize>) -> Self {
-        self.secondary_packet_size = packet_size;
-        self
-    }
-
     /// Overwrites existing sender address with the provided value.
     pub fn set_sender_address(&mut self, sender_address: Recipient) {
         self.sender_address = sender_address;
@@ -146,6 +128,16 @@ where
         ack_key: &AckKey,
         reply_surb: ReplySurb,
     ) -> Result<PreparedFragment, NymTopologyError> {
+        // each reply attaches the digest of the encryption key so that the recipient could
+        // lookup correct key for decryption,
+        let reply_overhead = ReplySurbKeyDigestAlgorithm::output_size();
+        let expected_plaintext = fragment.serialized_size() + ACK_OVERHEAD + reply_overhead;
+
+        // the reason we're unwrapping (or rather 'expecting') here rather than handling the error
+        // more gracefully is that this error should never be reached as it implies incorrect chunking
+        let packet_size = PacketSize::get_type_from_plaintext(expected_plaintext)
+            .expect("the message has been incorrectly fragmented");
+
         // this is not going to be accurate by any means. but that's the best estimation we can do
         let expected_forward_delay = Delay::new_from_millis(
             (self.average_packet_delay.as_millis() * self.num_mix_hops as u128) as u64,
@@ -162,9 +154,8 @@ where
 
         // the unwrap here is fine as the failures can only originate from attempting to use invalid payload lengths
         // and we just very carefully constructed a (presumably) valid one
-        let (sphinx_packet, first_hop_address) = reply_surb
-            .apply_surb(packet_payload, Some(self.primary_packet_size))
-            .unwrap();
+        let (sphinx_packet, first_hop_address) =
+            reply_surb.apply_surb(packet_payload, packet_size).unwrap();
 
         Ok(PreparedFragment {
             // the round-trip delay is the sum of delays of all hops on the forward route as
@@ -200,6 +191,17 @@ where
         ack_key: &AckKey,
         packet_recipient: &Recipient,
     ) -> Result<PreparedFragment, NymTopologyError> {
+        // each plain or repliable packet (i.e. not a reply) attaches an ephemeral public key so that the recipient
+        // could perform diffie-hellman with its own keys followed by a kdf to re-derive
+        // the packet encryption key
+        let non_reply_overhead = encryption::PUBLIC_KEY_SIZE;
+        let expected_plaintext = fragment.serialized_size() + ACK_OVERHEAD + non_reply_overhead;
+
+        // the reason we're unwrapping (or rather 'expecting') here rather than handling the error
+        // more gracefully is that this error should never be reached as it implies incorrect chunking
+        let packet_size = PacketSize::get_type_from_plaintext(expected_plaintext)
+            .expect("the message has been incorrectly fragmented");
+
         let fragment_identifier = fragment.fragment_identifier();
 
         // create an ack
@@ -223,7 +225,7 @@ where
         // create the actual sphinx packet here. With valid route and correct payload size,
         // there's absolutely no reason for this call to fail.
         let sphinx_packet = SphinxPacketBuilder::new()
-            .with_payload_size(self.primary_packet_size.payload_size())
+            .with_payload_size(packet_size.payload_size())
             .build_packet(packet_payload, &route, &destination, &delays)
             .unwrap();
 
@@ -263,7 +265,7 @@ where
         message: NymMessage,
         packet_size: PacketSize,
     ) -> Vec<Fragment> {
-        let plaintext_per_packet = message.available_plaintext_per_packet(packet_size);
+        let plaintext_per_packet = message.available_sphinx_plaintext_per_packet(packet_size);
 
         message
             .pad_to_full_packet_lengths(plaintext_per_packet)

sdk/rust/nym-sdk/examples/simple.rs

@@ -18,4 +18,6 @@ async fn main() {
     client
         .on_messages(|msg| println!("Received: {}", String::from_utf8_lossy(&msg.message)))
         .await;
+
+    client.disconnect().await;
 }