From 425a3e2684a5c58c9575ff0ae27674c5ad39b034 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Mon, 10 Jun 2024 20:51:45 +0900
Subject: [PATCH] fix observing final msg

---
 mixnet/v2/sim/adversary.py | 17 +++++++++--------
 mixnet/v2/sim/p2p.py       |  4 +++-
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/mixnet/v2/sim/adversary.py b/mixnet/v2/sim/adversary.py
index 3029b35..9e2591f 100644
--- a/mixnet/v2/sim/adversary.py
+++ b/mixnet/v2/sim/adversary.py
@@ -33,14 +33,9 @@ class Adversary:
     def inspect_message_size(self, msg: SphinxPacket | bytes):
         self.message_sizes.append(len(msg))
 
-    def observe_receiving_node(self, sender: "Node", receiver: "Node", msg: SphinxPacket | bytes):
-        cur_window = len(self.msg_pools_per_window) - 1
-        self.msg_pools_per_window[cur_window][receiver].append(self.env.now)
-        self.msgs_received_per_window[cur_window][receiver].add(sender)
-
-        origin_id = receiver.inspect_message(msg)
-        if origin_id is not None:
-            self.final_msgs_received[receiver][cur_window].append((sender, origin_id))
+    def observe_receiving_node(self, sender: "Node", receiver: "Node"):
+        self.msg_pools_per_window[-1][receiver].append(self.env.now)
+        self.msgs_received_per_window[-1][receiver].add(sender)
         # if node not in self.node_states[self.env.now]:
         #     self.node_states[self.env.now][node] = NodeState.RECEIVING
 
@@ -53,6 +48,12 @@ class Adversary:
             self.senders_around_interval.update({sender})
         # self.node_states[self.env.now][node] = NodeState.SENDING
 
+    def observe_if_final_msg(self, sender: "Node", receiver: "Node", msg: SphinxPacket | bytes):
+        origin_id = receiver.inspect_message(msg)
+        if origin_id is not None:
+            cur_window = len(self.msgs_received_per_window) - 1
+            self.final_msgs_received[receiver][cur_window].append((sender, origin_id))
+
     def is_around_message_interval(self, time: SimTime):
         return time % self.config.mixnet.message_interval <= self.config.mixnet.max_message_prep_time
 
diff --git a/mixnet/v2/sim/p2p.py b/mixnet/v2/sim/p2p.py
index fc6158c..48d175e 100644
--- a/mixnet/v2/sim/p2p.py
+++ b/mixnet/v2/sim/p2p.py
@@ -52,7 +52,7 @@ class P2P(ABC):
             yield self.env.timeout(self.config.p2p.random_network_latency())
 
             self.measurement.measure_ingress(receiver, msg)
-            self.adversary.observe_receiving_node(sender, receiver, msg)
+            self.adversary.observe_receiving_node(sender, receiver)
         self.receive(msg, hops_traveled + 1, sender, receiver)
 
     @abstractmethod
@@ -80,6 +80,7 @@ class NaiveBroadcastP2P(P2P):
     def receive(self, msg: SphinxPacket | bytes, hops_traveled: int, sender: "Node", receiver: "Node"):
         msg_hash = hashlib.sha256(bytes(msg)).digest()
         self.measurement.update_message_hops(msg_hash, hops_traveled)
+        self.adversary.observe_if_final_msg(sender, receiver, msg)
         self.env.process(receiver.receive_message(msg))
 
 
@@ -134,6 +135,7 @@ class GossipP2P(P2P):
         if msg_hash not in self.message_cache[receiver]:
             self.message_cache[receiver][msg_hash] = sender
             self.measurement.update_message_hops(msg_hash, hops_traveled)
+            self.adversary.observe_if_final_msg(sender, receiver, msg)
 
             # Receive and gossip
             self.env.process(receiver.receive_message(msg))