From a0ddaf5106a67f10168cfbc8cfedf169f4010098 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Sat, 25 May 2024 16:05:25 +0900
Subject: [PATCH] refactor: p2p measuring

---
 mixnet/v2/sim/p2p.py | 53 +++++++++++++++-----------------------------
 1 file changed, 18 insertions(+), 35 deletions(-)

diff --git a/mixnet/v2/sim/p2p.py b/mixnet/v2/sim/p2p.py
index cc3310e..8817497 100644
--- a/mixnet/v2/sim/p2p.py
+++ b/mixnet/v2/sim/p2p.py
@@ -39,11 +39,23 @@ class P2P(ABC):
         # Without any yield, SimPy complains that the broadcast func is not a generator.
         yield self.env.timeout(0)
 
-    @abstractmethod
-    def receive(self, msg: SphinxPacket | bytes, sender: "Node", receiver: "Node"):
+    def send(self, msg: SphinxPacket | bytes, sender: "Node", receiver: "Node", is_first_of_msg: bool):
+        if is_first_of_msg:
+            self.adversary.inspect_message_size(msg)
+            self.adversary.observe_sending_node(sender)
+        self.measurement.measure_egress(sender, msg)
+
         # simulate network latency
         yield self.env.timeout(random.uniform(0, self.config.p2p.max_network_latency))
 
+        self.measurement.measure_ingress(receiver, msg)
+        self.adversary.observe_receiving_node(receiver)
+        self.receive(msg, sender, receiver)
+
+    @abstractmethod
+    def receive(self, msg: SphinxPacket | bytes, sender: "Node", receiver: "Node"):
+        pass
+
     def log(self, msg):
         print(f"t={self.env.now:.3f}: P2P: {msg}")
 
@@ -57,27 +69,12 @@ class NaiveBroadcastP2P(P2P):
     # but we accept SphinxPacket as well because we don't implement Sphinx deserialization.
     def broadcast(self, sender: "Node", msg: SphinxPacket | bytes):
         yield from super().broadcast(sender, msg)
+
         self.log(f"Node:{sender.id}: Broadcasting a msg: {len(msg)} bytes")
-
-        cnt = 0
-        for receiver in self.nodes:
-            self.measurement.measure_egress(sender, msg)
-            self.env.process(self.receive(msg, sender, receiver))
-            cnt += 1
-
-        if cnt > 0:
-            # Adversary
-            self.adversary.inspect_message_size(msg)
-            self.adversary.observe_sending_node(sender)
+        for i, receiver in enumerate(self.nodes):
+            self.env.process(self.send(msg, sender, receiver, i == 0))
 
     def receive(self, msg: SphinxPacket | bytes, sender: "Node", receiver: "Node"):
-        yield from super().receive(msg, sender, receiver)
-
-        # Measurement
-        self.measurement.measure_ingress(receiver, msg)
-        # Adversary
-        self.adversary.observe_receiving_node(receiver)
-
         self.env.process(receiver.receive_message(msg))
 
 
@@ -122,24 +119,10 @@ class GossipP2P(P2P):
             # Don't gossip the message if it was received from the node who is going to be the receiver,
             # which means that the node already knows the message.
             if receiver != self.message_cache[sender][msg_hash]:
-                self.measurement.measure_egress(sender, msg)
-                self.env.process(self.receive(msg, sender, receiver))
+                self.env.process(self.send(msg, sender, receiver, cnt == 0))
                 cnt += 1
 
-        if cnt > 0:
-            # Adversary
-            self.adversary.inspect_message_size(msg)
-            self.adversary.observe_sending_node(sender)
-
     def receive(self, msg: SphinxPacket | bytes, sender: "Node", receiver: "Node"):
-        yield from super().receive(msg, sender, receiver)
-
-        # Measure ingress regardless of whether the message has been received before,
-        # because the node doesn't know if the message is a duplicate before receiving it from the network.
-        self.measurement.measure_ingress(receiver, msg)
-        # Adversary
-        self.adversary.observe_receiving_node(receiver)
-
         # Receive/gossip the msg only if it hasn't been received before. If not, just ignore the msg.
         # i.e. each message is received/gossiped at most once by each node.
         msg_hash = hashlib.sha256(bytes(msg)).digest()