From fb1894fb8c6cfd93c73e6b580f5afbe95c84e931 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Fri, 8 Nov 2024 11:20:30 +0700
Subject: [PATCH] Exclude peer who forwarded a message when forwarding the
 message (#32)

Co-authored-by: Daniel Sanchez <3danimanimal@gmail.com>
---
 simlib/mixnet-sims/src/node/mix/mod.rs | 30 +++++++++++++++++---------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/simlib/mixnet-sims/src/node/mix/mod.rs b/simlib/mixnet-sims/src/node/mix/mod.rs
index 6a8fedc..3608253 100644
--- a/simlib/mixnet-sims/src/node/mix/mod.rs
+++ b/simlib/mixnet-sims/src/node/mix/mod.rs
@@ -9,6 +9,7 @@ use crossbeam::channel;
 use futures::Stream;
 use lottery::StakeLottery;
 use multiaddr::Multiaddr;
+use netrunner::network::NetworkMessage;
 use netrunner::node::{Node, NodeId};
 use netrunner::{
     network::{InMemoryNetworkInterface, NetworkInterface, PayloadSize},
@@ -211,23 +212,27 @@ impl MixNode {
         }
     }
 
-    fn forward(&mut self, message: MixMessage) {
+    fn forward(&mut self, message: MixMessage, exclude_node: Option<NodeId>) {
         if !self.message_cache.insert(Self::sha256(&message.0)) {
             return;
         }
-        for node_id in self.settings.connected_peers.iter() {
+        for node_id in self
+            .settings
+            .connected_peers
+            .iter()
+            .filter(|&id| Some(*id) != exclude_node)
+        {
             self.network_interface
                 .send_message(*node_id, message.clone())
         }
     }
 
-    fn receive(&mut self) -> Vec<MixMessage> {
+    fn receive(&mut self) -> Vec<NetworkMessage<MixMessage>> {
         self.network_interface
             .receive_messages()
             .into_iter()
-            .map(|msg| msg.into_payload())
             // Retain only messages that have not been seen before
-            .filter(|msg| self.message_cache.insert(Self::sha256(&msg.0)))
+            .filter(|msg| self.message_cache.insert(Self::sha256(&msg.payload().0)))
             .collect()
     }
 
@@ -277,11 +282,16 @@ impl Node for MixNode {
                 self.persistent_sender.send(message).unwrap();
             }
         }
-        let received_messages = self.receive();
-        for message in received_messages {
+        // TODO: Generate cover message with probability
+        for network_message in self.receive() {
             // println!(">>>>> Node {}, message: {message:?}", self.id);
-            self.forward(message.clone());
-            self.blend_sender.send(message.0).unwrap();
+            self.forward(
+                network_message.payload().clone(),
+                Some(network_message.from),
+            );
+            self.blend_sender
+                .send(network_message.into_payload().0)
+                .unwrap();
         }
 
         // Proceed message blend
@@ -306,7 +316,7 @@ impl Node for MixNode {
         if let Poll::Ready(Some(msg)) =
             pin!(&mut self.persistent_transmission_messages).poll_next(&mut cx)
         {
-            self.forward(MixMessage(msg));
+            self.forward(MixMessage(msg), None);
         }
 
         self.state.step_id += 1;