Overide number of mix hops separately per packet (#4205)

* Try passing mix_hops all the way down the call chain

* Set zero mix_hops manually

* fix

* also set zero mix hops in tun listener

* fix

* mix hops for surbs

* Another case covered

* Remove mix_hops config from Traffic

* clippy

* Add comment about why we added new functions

* Update comment

* Add surb_mix_hops to wasm config

* Remove temporary added mix_hops = 0

* Remove another temporary added mix_hops = 0

* Add comment about the limitation of num_mix_hops

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

@@ -28,6 +28,7 @@ pub enum InputMessage {
         recipient: Recipient,
         data: Vec<u8>,
         lane: TransmissionLane,
+        mix_hops: Option<u8>,
     },
 
     /// Creates a message used for a duplex anonymous communication where the recipient
@@ -43,6 +44,7 @@ pub enum InputMessage {
         data: Vec<u8>,
         reply_surbs: u32,
         lane: TransmissionLane,
+        mix_hops: Option<u8>,
     },
 
     /// Attempt to use our internally received and stored `ReplySurb` to send the message back
@@ -92,6 +94,29 @@ impl InputMessage {
             recipient,
             data,
             lane,
+            mix_hops: None,
+        };
+        if let Some(packet_type) = packet_type {
+            InputMessage::new_wrapper(message, packet_type)
+        } else {
+            message
+        }
+    }
+
+    // IMHO `new_regular` should take `mix_hops: Option<u8>` as an argument instead of creating
+    // this function, but that would potentially break backwards compatibility with the current API
+    pub fn new_regular_with_custom_hops(
+        recipient: Recipient,
+        data: Vec<u8>,
+        lane: TransmissionLane,
+        packet_type: Option<PacketType>,
+        mix_hops: u8,
+    ) -> Self {
+        let message = InputMessage::Regular {
+            recipient,
+            data,
+            lane,
+            mix_hops: Some(mix_hops),
         };
         if let Some(packet_type) = packet_type {
             InputMessage::new_wrapper(message, packet_type)
@@ -112,6 +137,31 @@ impl InputMessage {
             data,
             reply_surbs,
             lane,
+            mix_hops: None,
+        };
+        if let Some(packet_type) = packet_type {
+            InputMessage::new_wrapper(message, packet_type)
+        } else {
+            message
+        }
+    }
+
+    // IMHO `new_anonymous` should take `mix_hops: Option<u8>` as an argument instead of creating
+    // this function, but that would potentially break backwards compatibility with the current API
+    pub fn new_anonymous_with_custom_hops(
+        recipient: Recipient,
+        data: Vec<u8>,
+        reply_surbs: u32,
+        lane: TransmissionLane,
+        packet_type: Option<PacketType>,
+        mix_hops: u8,
+    ) -> Self {
+        let message = InputMessage::Anonymous {
+            recipient,
+            data,
+            reply_surbs,
+            lane,
+            mix_hops: Some(mix_hops),
         };
         if let Some(packet_type) = packet_type {
             InputMessage::new_wrapper(message, packet_type)

common/client-core/src/client/real_messages_control/acknowledgement_control/input_message_listener.rs

@@ -73,10 +73,11 @@ where
         content: Vec<u8>,
         lane: TransmissionLane,
         packet_type: PacketType,
+        mix_hops: Option<u8>,
     ) {
         if let Err(err) = self
             .message_handler
-            .try_send_plain_message(recipient, content, lane, packet_type)
+            .try_send_plain_message(recipient, content, lane, packet_type, mix_hops)
             .await
         {
             warn!("failed to send a plain message - {err}")
@@ -90,10 +91,18 @@ where
         reply_surbs: u32,
         lane: TransmissionLane,
         packet_type: PacketType,
+        mix_hops: Option<u8>,
     ) {
         if let Err(err) = self
             .message_handler
-            .try_send_message_with_reply_surbs(recipient, content, reply_surbs, lane, packet_type)
+            .try_send_message_with_reply_surbs(
+                recipient,
+                content,
+                reply_surbs,
+                lane,
+                packet_type,
+                mix_hops,
+            )
             .await
         {
             warn!("failed to send a repliable message - {err}")
@@ -106,8 +115,9 @@ where
                 recipient,
                 data,
                 lane,
+                mix_hops,
             } => {
-                self.handle_plain_message(recipient, data, lane, PacketType::Mix)
+                self.handle_plain_message(recipient, data, lane, PacketType::Mix, mix_hops)
                     .await
             }
             InputMessage::Anonymous {
@@ -115,9 +125,17 @@ where
                 data,
                 reply_surbs,
                 lane,
+                mix_hops,
             } => {
-                self.handle_repliable_message(recipient, data, reply_surbs, lane, PacketType::Mix)
-                    .await
+                self.handle_repliable_message(
+                    recipient,
+                    data,
+                    reply_surbs,
+                    lane,
+                    PacketType::Mix,
+                    mix_hops,
+                )
+                .await
             }
             InputMessage::Reply {
                 recipient_tag,
@@ -135,8 +153,9 @@ where
                     recipient,
                     data,
                     lane,
+                    mix_hops,
                 } => {
-                    self.handle_plain_message(recipient, data, lane, packet_type)
+                    self.handle_plain_message(recipient, data, lane, packet_type, mix_hops)
                         .await
                 }
                 InputMessage::Anonymous {
@@ -144,9 +163,17 @@ where
                     data,
                     reply_surbs,
                     lane,
+                    mix_hops,
                 } => {
-                    self.handle_repliable_message(recipient, data, reply_surbs, lane, packet_type)
-                        .await
+                    self.handle_repliable_message(
+                        recipient,
+                        data,
+                        reply_surbs,
+                        lane,
+                        packet_type,
+                        mix_hops,
+                    )
+                    .await
                 }
                 InputMessage::Reply {
                     recipient_tag,

common/client-core/src/client/real_messages_control/acknowledgement_control/mod.rs

@@ -69,6 +69,7 @@ pub(crate) struct PendingAcknowledgement {
     message_chunk: Fragment,
     delay: SphinxDelay,
     destination: PacketDestination,
+    mix_hops: Option<u8>,
 }
 
 impl PendingAcknowledgement {
@@ -77,11 +78,13 @@ impl PendingAcknowledgement {
         message_chunk: Fragment,
         delay: SphinxDelay,
         recipient: Recipient,
+        mix_hops: Option<u8>,
     ) -> Self {
         PendingAcknowledgement {
             message_chunk,
             delay,
             destination: PacketDestination::KnownRecipient(recipient.into()),
+            mix_hops,
         }
     }
 
@@ -98,6 +101,9 @@ impl PendingAcknowledgement {
                 recipient_tag,
                 extra_surb_request,
             },
+            // Messages sent using SURBs are using the number of mix hops set by the recipient when
+            // they provided the SURBs, so it doesn't make sense to include it here.
+            mix_hops: None,
         }
     }
 

common/client-core/src/client/real_messages_control/acknowledgement_control/retransmission_request_listener.rs

@@ -49,12 +49,18 @@ where
         packet_recipient: Recipient,
         chunk_data: Fragment,
         packet_type: PacketType,
+        mix_hops: Option<u8>,
     ) -> Result<PreparedFragment, PreparationError> {
         debug!("retransmitting normal packet...");
 
         // TODO: Figure out retransmission packet type signaling
         self.message_handler
-            .try_prepare_single_chunk_for_sending(packet_recipient, chunk_data, packet_type)
+            .try_prepare_single_chunk_for_sending(
+                packet_recipient,
+                chunk_data,
+                packet_type,
+                mix_hops,
+            )
             .await
     }
 
@@ -89,6 +95,7 @@ where
                     **recipient,
                     timed_out_ack.message_chunk.clone(),
                     packet_type,
+                    timed_out_ack.mix_hops,
                 )
                 .await
             }

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

@@ -418,9 +418,10 @@ where
         message: Vec<u8>,
         lane: TransmissionLane,
         packet_type: PacketType,
+        mix_hops: Option<u8>,
     ) -> Result<(), PreparationError> {
         let message = NymMessage::new_plain(message);
-        self.try_split_and_send_non_reply_message(message, recipient, lane, packet_type)
+        self.try_split_and_send_non_reply_message(message, recipient, lane, packet_type, mix_hops)
             .await
     }
 
@@ -430,6 +431,7 @@ where
         recipient: Recipient,
         lane: TransmissionLane,
         packet_type: PacketType,
+        mix_hops: Option<u8>,
     ) -> Result<(), PreparationError> {
         debug!("Sending non-reply message with packet type {packet_type}");
         // TODO: I really dislike existence of this assertion, it implies code has to be re-organised
@@ -461,6 +463,7 @@ where
                 &self.config.ack_key,
                 &recipient,
                 packet_type,
+                mix_hops,
             )?;
 
             let real_message = RealMessage::new(
@@ -468,7 +471,8 @@ where
                 Some(fragment.fragment_identifier()),
             );
             let delay = prepared_fragment.total_delay;
-            let pending_ack = PendingAcknowledgement::new_known(fragment, delay, recipient);
+            let pending_ack =
+                PendingAcknowledgement::new_known(fragment, delay, recipient, mix_hops);
 
             real_messages.push(real_message);
             pending_acks.push(pending_ack);
@@ -485,6 +489,7 @@ where
         recipient: Recipient,
         amount: u32,
         packet_type: PacketType,
+        mix_hops: Option<u8>,
     ) -> Result<(), PreparationError> {
         debug!("Sending additional reply SURBs with packet type {packet_type}");
         let sender_tag = self.get_or_create_sender_tag(&recipient);
@@ -501,6 +506,7 @@ where
             recipient,
             TransmissionLane::AdditionalReplySurbs,
             packet_type,
+            mix_hops,
         )
         .await?;
 
@@ -517,6 +523,7 @@ where
         num_reply_surbs: u32,
         lane: TransmissionLane,
         packet_type: PacketType,
+        mix_hops: Option<u8>,
     ) -> Result<(), SurbWrappedPreparationError> {
         debug!("Sending message with reply SURBs with packet type {packet_type}");
         let sender_tag = self.get_or_create_sender_tag(&recipient);
@@ -527,7 +534,7 @@ where
         let message =
             NymMessage::new_repliable(RepliableMessage::new_data(message, sender_tag, reply_surbs));
 
-        self.try_split_and_send_non_reply_message(message, recipient, lane, packet_type)
+        self.try_split_and_send_non_reply_message(message, recipient, lane, packet_type, mix_hops)
             .await?;
 
         log::trace!("storing {} reply keys", reply_keys.len());
@@ -541,6 +548,7 @@ where
         recipient: Recipient,
         chunk: Fragment,
         packet_type: PacketType,
+        mix_hops: Option<u8>,
     ) -> Result<PreparedFragment, PreparationError> {
         debug!("Sending single chunk with packet type {packet_type}");
         let topology_permit = self.topology_access.get_read_permit().await;
@@ -554,6 +562,7 @@ where
                 &self.config.ack_key,
                 &recipient,
                 packet_type,
+                mix_hops,
             )
             .unwrap();
 

common/client-core/src/client/replies/reply_controller/mod.rs

@@ -516,6 +516,7 @@ where
                     recipient,
                     to_send,
                     nym_sphinx::params::PacketType::Mix,
+                    self.config.reply_surbs.surb_mix_hops,
                 )
                 .await
             {

common/client-core/src/config/mod.rs

@@ -607,6 +607,10 @@ pub struct ReplySurbs {
     /// This is going to be superseded by key rotation once implemented.
     #[serde(with = "humantime_serde")]
     pub maximum_reply_key_age: Duration,
+
+    /// Specifies the number of mixnet hops the packet should go through. If not specified, then
+    /// the default value is used.
+    pub surb_mix_hops: Option<u8>,
 }
 
 impl Default for ReplySurbs {
@@ -622,6 +626,7 @@ impl Default for ReplySurbs {
             maximum_reply_surb_drop_waiting_period: DEFAULT_MAXIMUM_REPLY_SURB_DROP_WAITING_PERIOD,
             maximum_reply_surb_age: DEFAULT_MAXIMUM_REPLY_SURB_AGE,
             maximum_reply_key_age: DEFAULT_MAXIMUM_REPLY_KEY_AGE,
+            surb_mix_hops: None,
         }
     }
 }

common/client-core/src/config/old_config_v1_1_30.rs

@@ -155,6 +155,7 @@ impl From<ConfigV1_1_30> for Config {
                         .maximum_reply_surb_drop_waiting_period,
                     maximum_reply_surb_age: value.debug.reply_surbs.maximum_reply_surb_age,
                     maximum_reply_key_age: value.debug.reply_surbs.maximum_reply_key_age,
+                    surb_mix_hops: None,
                 },
             },
         }

common/node-tester-utils/src/tester.rs

@@ -258,6 +258,7 @@ where
             &address,
             &address,
             PacketType::Mix,
+            None,
         )?)
     }
 

common/nymsphinx/src/preparer/mod.rs

@@ -181,6 +181,7 @@ pub trait FragmentPreparer {
     /// - compute vk_b = g^x || v_b
     /// - compute sphinx_plaintext = SURB_ACK || g^x || v_b
     /// - compute sphinx_packet = Sphinx(recipient, sphinx_plaintext)
+    #[allow(clippy::too_many_arguments)]
     fn prepare_chunk_for_sending(
         &mut self,
         fragment: Fragment,
@@ -189,6 +190,7 @@ pub trait FragmentPreparer {
         packet_sender: &Recipient,
         packet_recipient: &Recipient,
         packet_type: PacketType,
+        mix_hops: Option<u8>,
     ) -> 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
@@ -226,7 +228,8 @@ pub trait FragmentPreparer {
         };
 
         // generate pseudorandom route for the packet
-        let hops = self.num_mix_hops();
+        let hops = mix_hops.unwrap_or(self.num_mix_hops());
+        log::trace!("Preparing chunk for sending with {} mix hops", hops);
         let route =
             topology.random_route_to_gateway(self.rng(), hops, packet_recipient.gateway())?;
         let destination = packet_recipient.as_sphinx_destination();
@@ -389,6 +392,7 @@ where
         ack_key: &AckKey,
         packet_recipient: &Recipient,
         packet_type: PacketType,
+        mix_hops: Option<u8>,
     ) -> Result<PreparedFragment, NymTopologyError> {
         let sender = self.sender_address;
 
@@ -400,6 +404,7 @@ where
             &sender,
             packet_recipient,
             packet_type,
+            mix_hops,
         )
     }
 

common/nymsphinx/src/receiver.rs

@@ -169,6 +169,10 @@ pub struct SphinxMessageReceiver {
 
 impl SphinxMessageReceiver {
     /// Allows setting non-default number of expected mix hops in the network.
+    // IMPORTANT NOTE: this is among others used to deserialize SURBs. Meaning that this is a
+    // global setting and currently always set to the default value. The implication is that it is
+    // not currently possible to have different number of hops for different SURB messages. So,
+    // don't try to use <3 mix hops for SURBs until this is refactored.
     #[must_use]
     pub fn with_mix_hops(mut self, hops: u8) -> Self {
         self.num_mix_hops = hops;

common/wasm/client-core/src/config/mod.rs

@@ -434,6 +434,10 @@ pub struct ReplySurbsWasm {
     /// Defines maximum amount of time given reply key is going to be valid for.
     /// This is going to be superseded by key rotation once implemented.
     pub maximum_reply_key_age_ms: u32,
+
+    /// Defines how many mix nodes the reply surb should go through.
+    /// If not set, the default value is going to be used.
+    pub surb_mix_hops: Option<u8>,
 }
 
 impl Default for ReplySurbsWasm {
@@ -463,6 +467,7 @@ impl From<ReplySurbsWasm> for ConfigReplySurbs {
             maximum_reply_key_age: Duration::from_millis(
                 reply_surbs.maximum_reply_key_age_ms as u64,
             ),
+            surb_mix_hops: reply_surbs.surb_mix_hops,
         }
     }
 }
@@ -484,6 +489,7 @@ impl From<ConfigReplySurbs> for ReplySurbsWasm {
                 .as_millis() as u32,
             maximum_reply_surb_age_ms: reply_surbs.maximum_reply_surb_age.as_millis() as u32,
             maximum_reply_key_age_ms: reply_surbs.maximum_reply_key_age.as_millis() as u32,
+            surb_mix_hops: reply_surbs.surb_mix_hops,
         }
     }
 }

common/wasm/client-core/src/config/override.rs

@@ -303,6 +303,9 @@ pub struct ReplySurbsWasmOverride {
     /// This is going to be superseded by key rotation once implemented.
     #[tsify(optional)]
     pub maximum_reply_key_age_ms: Option<u32>,
+
+    #[tsify(optional)]
+    pub surb_mix_hops: Option<u8>,
 }
 
 impl From<ReplySurbsWasmOverride> for ReplySurbsWasm {
@@ -337,6 +340,7 @@ impl From<ReplySurbsWasmOverride> for ReplySurbsWasm {
             maximum_reply_key_age_ms: value
                 .maximum_reply_key_age_ms
                 .unwrap_or(def.maximum_reply_key_age_ms),
+            surb_mix_hops: value.surb_mix_hops,
         }
     }
 }

service-providers/network-requester/src/reply.rs

@@ -189,6 +189,7 @@ impl MixnetAddress {
                     recipient: *recipient,
                     data: message,
                     lane: TransmissionLane::ConnectionId(connection_id),
+                    mix_hops: None,
                 }),
                 packet_type,
             },