From 063efcf5c3d219c878e644e4d8f5a35cc5388097 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Thu, 23 May 2024 16:34:39 +0900
Subject: [PATCH] gossip wip

---
 mixnet/v2/sim/node.py       |  4 +--
 mixnet/v2/sim/p2p.py        | 63 +++++++++++++++++++++++++++++--------
 mixnet/v2/sim/simulation.py |  4 +--
 3 files changed, 54 insertions(+), 17 deletions(-)

diff --git a/mixnet/v2/sim/node.py b/mixnet/v2/sim/node.py
index 14a2851..ab41352 100644
--- a/mixnet/v2/sim/node.py
+++ b/mixnet/v2/sim/node.py
@@ -10,14 +10,14 @@ from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X
 
 from config import Config
 from sphinx import SphinxPacket, Attachment
-from p2p import P2p
+from p2p import P2P
 
 
 class Node:
     INCENTIVE_TX_SIZE = 512
     PADDING_SEPARATOR = b'\x01'
 
-    def __init__(self, id: int, env: simpy.Environment, p2p: P2p, config: Config):
+    def __init__(self, id: int, env: simpy.Environment, p2p: P2P, config: Config):
         self.id = id
         self.env = env
         self.p2p = p2p
diff --git a/mixnet/v2/sim/p2p.py b/mixnet/v2/sim/p2p.py
index 980d18a..c197284 100644
--- a/mixnet/v2/sim/p2p.py
+++ b/mixnet/v2/sim/p2p.py
@@ -1,5 +1,9 @@
 from __future__ import annotations
+
+import hashlib
 import random
+from abc import ABC, abstractmethod
+from collections import defaultdict
 from typing import TYPE_CHECKING
 
 import simpy
@@ -13,7 +17,7 @@ if TYPE_CHECKING:
     from node import Node
 
 
-class P2p:
+class P2P(ABC):
     def __init__(self, env: simpy.Environment, config: Config):
         self.env = env
         self.config = config
@@ -29,29 +33,62 @@ class P2p:
 
     # This should accept only bytes in practice,
     # but we accept SphinxPacket as well because we don't implement Sphinx deserialization.
+    @abstractmethod
     def broadcast(self, sender: "Node", msg: SphinxPacket | bytes):
-        self.log("Broadcasting a msg: %d bytes" % len(msg))
-
         # Adversary
         self.adversary.inspect_message_size(msg)
         self.adversary.observe_outgoing_message(sender)
-
         # Yield 0 to ensure that the broadcast is done in the same time step.
         # Without any yield, SimPy complains that the broadcast func is not a generator.
         yield self.env.timeout(0)
 
-        # TODO: gossipsub or something similar
+    @abstractmethod
+    def send(self, msg: SphinxPacket | bytes, receiver: "Node"):
+        # simulate network latency
+        yield self.env.timeout(random.uniform(0, self.config.p2p.max_network_latency))
+        # Measurement and adversary
+        self.measurement.measure_ingress(receiver, msg)
+        self.adversary.observe_incoming_message(receiver)
+
+    def log(self, msg):
+        print("P2P at %g: %s" % (self.env.now, msg))
+
+
+class NaiveBroadcastP2P(P2P):
+    def __init__(self, env: simpy.Environment, config: Config):
+        super().__init__(env, config)
+
+    # This should accept only bytes in practice,
+    # 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("Broadcasting a msg: %d bytes" % len(msg))
         for node in self.nodes:
             self.measurement.measure_egress(sender, msg)
             self.env.process(self.send(msg, node))
 
-    def send(self, msg: SphinxPacket | bytes, node: "Node"):
-        # simulate network latency
-        yield self.env.timeout(random.uniform(0, self.config.p2p.max_network_latency))
+    def send(self, msg: SphinxPacket | bytes, receiver: "Node"):
+        yield from super().send(msg, receiver)
+        self.env.process(receiver.receive_message(msg))
 
-        self.measurement.measure_ingress(node, msg)
-        self.adversary.observe_incoming_message(node)
-        self.env.process(node.receive_message(msg))
 
-    def log(self, msg):
-        print("P2P at %g: %s" % (self.env.now, msg))
+class GossipP2P(P2P):
+    def __init__(self, env: simpy.Environment, config: Config):
+        super().__init__(env, config)
+        self.topology = defaultdict(set)
+        self.message_cache = defaultdict(set)
+
+    def broadcast(self, sender: "Node", msg: SphinxPacket | bytes):
+        yield from super().broadcast(sender, msg)
+        self.log("Gossiping a msg: %d bytes" % len(msg))
+        for receiver in self.topology[sender]:
+            self.measurement.measure_egress(sender, msg)
+            self.env.process(self.send(msg, receiver))
+
+    def send(self, msg: SphinxPacket | bytes, receiver: "Node"):
+        yield from super().send(msg, receiver)
+        # receive the msg only if it hasn't been received before
+        msg_hash = hashlib.sha256(bytes(msg)).digest()
+        if msg_hash not in self.message_cache[receiver]:
+            self.message_cache[receiver].add(msg_hash)
+            self.env.process(receiver.receive_message(msg))
diff --git a/mixnet/v2/sim/simulation.py b/mixnet/v2/sim/simulation.py
index 3600668..b236807 100644
--- a/mixnet/v2/sim/simulation.py
+++ b/mixnet/v2/sim/simulation.py
@@ -4,7 +4,7 @@ import simpy
 
 from config import Config
 from node import Node
-from p2p import P2p
+from p2p import NaiveBroadcastP2P
 
 
 class Simulation:
@@ -12,7 +12,7 @@ class Simulation:
         random.seed()
         self.config = config
         self.env = simpy.Environment()
-        self.p2p = P2p(self.env, config)
+        self.p2p = NaiveBroadcastP2P(self.env, config)
         self.nodes = [Node(i, self.env, self.p2p, config) for i in range(config.mixnet.num_nodes)]
         self.p2p.add_node(self.nodes)