From 2065060cf3977f22df0591f8a1ece935d08b0891 Mon Sep 17 00:00:00 2001
From: Youngjoon Lee <5462944+youngjoon-lee@users.noreply.github.com>
Date: Fri, 5 Jul 2024 17:19:02 +0900
Subject: [PATCH] handle PeeringDegreeReached properly

---
 mixnet/node.py           | 17 +++++++++++++++--
 mixnet/sim/simulation.py |  7 +++++--
 2 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/mixnet/node.py b/mixnet/node.py
index b24134e..c5e1933 100644
--- a/mixnet/node.py
+++ b/mixnet/node.py
@@ -71,6 +71,12 @@ class Node:
         inbound_conn: SimplexConnection,
         outbound_conn: SimplexConnection,
     ):
+        if (
+            not self.mixgossip_channel.can_accept_conn()
+            or not peer.mixgossip_channel.can_accept_conn()
+        ):
+            raise PeeringDegreeReached
+
         self.mixgossip_channel.add_conn(
             DuplexConnection(
                 inbound_conn,
@@ -122,10 +128,13 @@ class MixGossipChannel:
         # https://docs.python.org/3/library/asyncio-task.html#asyncio.create_task
         self.tasks = set()
 
+    def can_accept_conn(self) -> bool:
+        return len(self.conns) < self.peering_degree
+
     def add_conn(self, conn: DuplexConnection):
-        if len(self.conns) >= self.peering_degree:
+        if not self.can_accept_conn():
             # For simplicity of the spec, reject the connection if the peering degree is reached.
-            raise ValueError("The peering degree is reached.")
+            raise PeeringDegreeReached()
 
         self.conns.append(conn)
         task = self.framework.spawn(self.__process_inbound_conn(conn))
@@ -223,3 +232,7 @@ def parse_msg(data: bytes) -> tuple[MsgType, bytes]:
 
 def build_noise_packet() -> bytes:
     return build_msg(MsgType.NOISE, bytes(DEFAULT_PAYLOAD_SIZE))
+
+
+class PeeringDegreeReached(Exception):
+    pass
diff --git a/mixnet/sim/simulation.py b/mixnet/sim/simulation.py
index b5520b8..2d873fc 100644
--- a/mixnet/sim/simulation.py
+++ b/mixnet/sim/simulation.py
@@ -3,7 +3,7 @@ import usim
 import mixnet.framework.usim as usimfw
 from mixnet.config import GlobalConfig, MixMembership, NodeInfo
 from mixnet.framework.framework import Framework
-from mixnet.node import Node
+from mixnet.node import Node, PeeringDegreeReached
 from mixnet.sim.config import Config
 from mixnet.sim.connection import MeteredRemoteSimplexConnection
 from mixnet.sim.stats import ConnectionStats
@@ -54,7 +54,10 @@ class Simulation:
                 self.create_conn(),
             )
             peer = nodes[(i + 1) % len(nodes)]
-            node.connect(peer, inbound_conn, outbound_conn)
+            try:
+                node.connect(peer, inbound_conn, outbound_conn)
+            except PeeringDegreeReached:
+                continue
             conn_stats.register(node, inbound_conn, outbound_conn)
             conn_stats.register(peer, outbound_conn, inbound_conn)