From 5a0412ff8361d0d8879b637b45243e60d2b872a0 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Fri, 8 Nov 2024 16:00:25 +0700
Subject: [PATCH] Measure both data and cover message counts (#40)

---
 simlib/mixnet-sims/src/node/mix/mod.rs   | 29 ++++++++++++++++--------
 simlib/mixnet-sims/src/node/mix/state.rs |  4 ++--
 2 files changed, 21 insertions(+), 12 deletions(-)

diff --git a/simlib/mixnet-sims/src/node/mix/mod.rs b/simlib/mixnet-sims/src/node/mix/mod.rs
index 7f1a710..eaf74b8 100644
--- a/simlib/mixnet-sims/src/node/mix/mod.rs
+++ b/simlib/mixnet-sims/src/node/mix/mod.rs
@@ -193,8 +193,8 @@ impl MixNode {
             state: MixnodeState {
                 node_id: id,
                 step_id: 0,
-                data_messages_generated: HashMap::new(),
-                data_messages_fully_unwrapped: HashMap::new(),
+                messages_generated: HashMap::new(),
+                messages_fully_unwrapped: HashMap::new(),
             },
             data_msg_lottery_update_time_sender,
             data_msg_lottery_interval,
@@ -279,11 +279,12 @@ impl Node for MixNode {
         let waker = futures::task::noop_waker();
         let mut cx = futures::task::Context::from_waker(&waker);
 
+        // Generate a data message probabilistically
         if let Poll::Ready(Some(_)) = pin!(&mut self.data_msg_lottery_interval).poll_next(&mut cx) {
             if self.data_msg_lottery.run() {
                 let payload = Payload::new();
                 self.state
-                    .data_messages_generated
+                    .messages_generated
                     .insert(payload.id(), self.state.step_id);
                 let message = self
                     .crypto_processor
@@ -292,9 +293,9 @@ impl Node for MixNode {
                 self.persistent_sender.send(message).unwrap();
             }
         }
-        // TODO: Generate cover message with probability
+
+        // Handle incoming messages
         for network_message in self.receive() {
-            // println!(">>>>> Node {}, message: {message:?}", self.id);
             self.forward(
                 network_message.payload().clone(),
                 Some(network_message.from),
@@ -313,14 +314,23 @@ impl Node for MixNode {
                 MixOutgoingMessage::FullyUnwrapped(payload) => {
                     let payload = Payload::load(payload);
                     self.state
-                        .data_messages_fully_unwrapped
+                        .messages_fully_unwrapped
                         .insert(payload.id(), self.state.step_id);
                     //TODO: create a tracing event
                 }
             }
         }
-        if let Poll::Ready(Some(msg)) = pin!(&mut self.cover_traffic).poll_next(&mut cx) {
-            let message = self.crypto_processor.wrap_message(&msg).unwrap();
+
+        // Generate a cover message probabilistically
+        if let Poll::Ready(Some(_)) = pin!(&mut self.cover_traffic).poll_next(&mut cx) {
+            let payload = Payload::new();
+            self.state
+                .messages_generated
+                .insert(payload.id(), self.state.step_id);
+            let message = self
+                .crypto_processor
+                .wrap_message(payload.as_bytes())
+                .unwrap();
             self.persistent_sender.send(message).unwrap();
         }
 
@@ -338,8 +348,7 @@ impl Node for MixNode {
         match ward {
             WardCondition::Max(_) => false,
             WardCondition::Sum(condition) => {
-                *condition.step_result.borrow_mut() +=
-                    self.state.data_messages_fully_unwrapped.len();
+                *condition.step_result.borrow_mut() += self.state.messages_fully_unwrapped.len();
                 false
             }
         }
diff --git a/simlib/mixnet-sims/src/node/mix/state.rs b/simlib/mixnet-sims/src/node/mix/state.rs
index 4f450b6..6ec9d27 100644
--- a/simlib/mixnet-sims/src/node/mix/state.rs
+++ b/simlib/mixnet-sims/src/node/mix/state.rs
@@ -16,8 +16,8 @@ pub struct MixnodeState {
     #[serde(serialize_with = "serialize_node_id_as_index")]
     pub node_id: NodeId,
     pub step_id: usize,
-    pub data_messages_generated: HashMap<PayloadId, usize>,
-    pub data_messages_fully_unwrapped: HashMap<PayloadId, usize>,
+    pub messages_generated: HashMap<PayloadId, usize>,
+    pub messages_fully_unwrapped: HashMap<PayloadId, usize>,
 }
 
 #[derive(Serialize)]