Merge pull request #1114 from HathorNetwork/refactor/p2p-peer-use-peer-id

refactor(p2p): use PeerId type instead of str

Author: jansegre

hathor/builder/builder.py

@@ -47,7 +47,7 @@
     TransactionStorage,
 )
 from hathor.transaction.vertex_parser import VertexParser
-from hathor.util import Random, get_environment_info, not_none
+from hathor.util import Random, get_environment_info
 from hathor.verification.verification_service import VerificationService
 from hathor.verification.vertex_verifiers import VertexVerifiers
 from hathor.vertex_handler import VertexHandler
@@ -264,7 +264,7 @@ def build(self) -> BuildArtifacts:
             rng=self._rng,
             checkpoints=self._checkpoints,
             capabilities=self._capabilities,
-            environment_info=get_environment_info(self._cmdline, peer.id),
+            environment_info=get_environment_info(self._cmdline, str(peer.id)),
             bit_signaling_service=bit_signaling_service,
             verification_service=verification_service,
             cpu_mining_service=cpu_mining_service,
@@ -515,7 +515,7 @@ def _get_or_create_event_manager(self) -> EventManager:
             reactor = self._get_reactor()
             storage = self._get_or_create_event_storage()
             factory = EventWebsocketFactory(
-                peer_id=not_none(peer.id),
+                peer_id=str(peer.id),
                 network=settings.NETWORK_NAME,
                 reactor=reactor,
                 event_storage=storage,

hathor/builder/cli_builder.py

@@ -42,7 +42,7 @@
 from hathor.reactor import ReactorProtocol as Reactor
 from hathor.stratum import StratumFactory
 from hathor.transaction.vertex_parser import VertexParser
-from hathor.util import Random, not_none
+from hathor.util import Random
 from hathor.verification.verification_service import VerificationService
 from hathor.verification.vertex_verifiers import VertexVerifiers
 from hathor.vertex_handler import VertexHandler
@@ -225,7 +225,7 @@ def create_manager(self, reactor: Reactor) -> HathorManager:
 
         if self._args.x_enable_event_queue:
             self.event_ws_factory = EventWebsocketFactory(
-                peer_id=not_none(peer.id),
+                peer_id=str(peer.id),
                 network=network,
                 reactor=reactor,
                 event_storage=event_storage
@@ -354,7 +354,7 @@ def create_manager(self, reactor: Reactor) -> HathorManager:
             event_manager=event_manager,
             wallet=self.wallet,
             checkpoints=settings.CHECKPOINTS,
-            environment_info=get_environment_info(args=str(self._args), peer_id=peer.id),
+            environment_info=get_environment_info(args=str(self._args), peer_id=str(peer.id)),
             full_verification=full_verification,
             enable_event_queue=self._args.x_enable_event_queue,
             bit_signaling_service=bit_signaling_service,

hathor/manager.py

@@ -47,6 +47,7 @@
 from hathor.mining.cpu_mining_service import CpuMiningService
 from hathor.p2p.manager import ConnectionsManager
 from hathor.p2p.peer import Peer
+from hathor.p2p.peer_id import PeerId
 from hathor.profiler import get_cpu_profiler
 from hathor.pubsub import HathorEvents, PubSubManager
 from hathor.reactor import ReactorProtocol as Reactor
@@ -225,7 +226,7 @@ def __init__(
         self._full_verification = full_verification
 
         # List of whitelisted peers
-        self.peers_whitelist: list[str] = []
+        self.peers_whitelist: list[PeerId] = []
 
         # List of capabilities of the peer
         if capabilities is not None:
@@ -297,7 +298,7 @@ def start(self) -> None:
                 sys.exit(-1)
 
         if self._enable_event_queue:
-            self._event_manager.start(not_none(self.my_peer.id))
+            self._event_manager.start(str(not_none(self.my_peer.id)))
 
         self.state = self.NodeState.INITIALIZING
         self.pubsub.publish(HathorEvents.MANAGER_ON_START)
@@ -976,7 +977,7 @@ def on_new_tx(
     def has_sync_version_capability(self) -> bool:
         return self._settings.CAPABILITY_SYNC_VERSION in self.capabilities
 
-    def add_peer_to_whitelist(self, peer_id: str) -> None:
+    def add_peer_to_whitelist(self, peer_id: PeerId) -> None:
         if not self._settings.ENABLE_PEER_WHITELIST:
             return
 
@@ -985,7 +986,7 @@ def add_peer_to_whitelist(self, peer_id: str) -> None:
         else:
             self.peers_whitelist.append(peer_id)
 
-    def remove_peer_from_whitelist_and_disconnect(self, peer_id: str) -> None:
+    def remove_peer_from_whitelist_and_disconnect(self, peer_id: PeerId) -> None:
         if not self._settings.ENABLE_PEER_WHITELIST:
             return
 

hathor/metrics.py

@@ -253,7 +253,7 @@ def collect_peer_connection_metrics(self) -> None:
 
             metric = PeerConnectionMetrics(
                 connection_string=str(connection.entrypoint) if connection.entrypoint else "",
-                peer_id=connection.peer.id,
+                peer_id=str(connection.peer.id),
                 network=connection.network,
                 received_messages=connection.metrics.received_messages,
                 sent_messages=connection.metrics.sent_messages,

hathor/p2p/entrypoint.py

@@ -21,18 +21,14 @@
 from twisted.internet.interfaces import IStreamClientEndpoint
 from typing_extensions import Self
 
+from hathor.p2p.peer_id import PeerId
 from hathor.reactor import ReactorProtocol as Reactor
-from hathor.types import Hash
 
 
 class Protocol(Enum):
     TCP = 'tcp'
 
 
-class PeerId(Hash):
-    pass
-
-
 @dataclass(frozen=True, slots=True)
 class Entrypoint:
     """Endpoint description (returned from DNS query, or received from the p2p network) may contain a peer-id."""

hathor/p2p/manager.py

@@ -29,6 +29,7 @@
 from hathor.p2p.netfilter.factory import NetfilterFactory
 from hathor.p2p.peer import Peer
 from hathor.p2p.peer_discovery import PeerDiscovery
+from hathor.p2p.peer_id import PeerId
 from hathor.p2p.peer_storage import PeerStorage
 from hathor.p2p.protocol import HathorProtocol
 from hathor.p2p.rate_limiter import RateLimiter
@@ -51,11 +52,11 @@
 
 
 class _SyncRotateInfo(NamedTuple):
-    candidates: list[str]
-    old: set[str]
-    new: set[str]
-    to_disable: set[str]
-    to_enable: set[str]
+    candidates: list[PeerId]
+    old: set[PeerId]
+    new: set[PeerId]
+    to_disable: set[PeerId]
+    to_enable: set[PeerId]
 
 
 class _ConnectingPeer(NamedTuple):
@@ -79,7 +80,7 @@ class GlobalRateLimiter:
 
     manager: Optional['HathorManager']
     connections: set[HathorProtocol]
-    connected_peers: dict[str, HathorProtocol]
+    connected_peers: dict[PeerId, HathorProtocol]
     connecting_peers: dict[IStreamClientEndpoint, _ConnectingPeer]
     handshaking_peers: set[HathorProtocol]
     whitelist_only: bool
@@ -174,7 +175,7 @@ def __init__(
         self.lc_sync_update_interval: float = 5  # seconds
 
         # Peers that always have sync enabled.
-        self.always_enable_sync: set[str] = set()
+        self.always_enable_sync: set[PeerId] = set()
 
         # Timestamp of the last time sync was updated.
         self._last_sync_rotate: float = 0.
@@ -485,7 +486,7 @@ def iter_not_ready_endpoints(self) -> Iterable[Entrypoint]:
             else:
                 self.log.warn('handshaking protocol has empty connection string', protocol=protocol)
 
-    def is_peer_connected(self, peer_id: str) -> bool:
+    def is_peer_connected(self, peer_id: PeerId) -> bool:
         """
         :type peer_id: string (peer.id)
         """
@@ -729,7 +730,7 @@ def get_connection_to_drop(self, protocol: HathorProtocol) -> HathorProtocol:
         assert protocol.peer.id is not None
         assert protocol.my_peer.id is not None
         other_connection = self.connected_peers[protocol.peer.id]
-        if protocol.my_peer.id > protocol.peer.id:
+        if bytes(protocol.my_peer.id) > bytes(protocol.peer.id):
             # connection started by me is kept
             if not protocol.inbound:
                 # other connection is dropped
@@ -751,7 +752,7 @@ def drop_connection(self, protocol: HathorProtocol) -> None:
         self.log.debug('dropping connection', peer_id=protocol.peer.id, protocol=type(protocol).__name__)
         protocol.send_error_and_close_connection('Connection droped')
 
-    def drop_connection_by_peer_id(self, peer_id: str) -> None:
+    def drop_connection_by_peer_id(self, peer_id: PeerId) -> None:
         """ Drop a connection by peer id
         """
         protocol = self.connected_peers.get(peer_id)
@@ -765,9 +766,9 @@ def sync_update(self) -> None:
         except Exception:
             self.log.error('_sync_rotate_if_needed failed', exc_info=True)
 
-    def set_always_enable_sync(self, values: list[str]) -> None:
+    def set_always_enable_sync(self, values: list[PeerId]) -> None:
         """Set a new list of peers to always enable sync. This operation completely replaces the previous list."""
-        new: set[str] = set(values)
+        new: set[PeerId] = set(values)
 
         old = self.always_enable_sync
         if new == old:
@@ -792,14 +793,14 @@ def set_always_enable_sync(self, values: list[str]) -> None:
 
     def _calculate_sync_rotate(self) -> _SyncRotateInfo:
         """Calculate new sync rotation."""
-        current_enabled: set[str] = set()
+        current_enabled: set[PeerId] = set()
         for peer_id, conn in self.connected_peers.items():
             if conn.is_sync_enabled():
                 current_enabled.add(peer_id)
 
         candidates = list(self.connected_peers.keys())
         self.rng.shuffle(candidates)
-        selected_peers: set[str] = set(candidates[:self.MAX_ENABLED_SYNC])
+        selected_peers: set[PeerId] = set(candidates[:self.MAX_ENABLED_SYNC])
 
         to_disable = current_enabled - selected_peers
         to_enable = selected_peers - current_enabled

hathor/p2p/netfilter/matches.py

@@ -130,7 +130,7 @@ def match(self, context: 'NetfilterContext') -> bool:
         if context.protocol.peer is None:
             return False
 
-        if context.protocol.peer.id != self.peer_id:
+        if str(context.protocol.peer.id) != self.peer_id:
             return False
 
         return True

hathor/p2p/peer.py

@@ -32,6 +32,7 @@
 from hathor.conf.get_settings import get_global_settings
 from hathor.daa import DifficultyAdjustmentAlgorithm
 from hathor.p2p.entrypoint import Entrypoint
+from hathor.p2p.peer_id import PeerId
 from hathor.p2p.utils import discover_dns, generate_certificate
 from hathor.util import not_none
 
@@ -59,7 +60,7 @@ class Peer:
     Usually a peer will have only one entrypoint.
     """
 
-    id: Optional[str]
+    id: Optional[PeerId]
     entrypoints: list[Entrypoint]
     private_key: Optional[rsa.RSAPrivateKeyWithSerialization]
     public_key: Optional[rsa.RSAPublicKey]
@@ -135,15 +136,15 @@ def generate_keys(self, key_size: int = 2048) -> None:
         self.public_key = self.private_key.public_key()
         self.id = self.calculate_id()
 
-    def calculate_id(self) -> str:
+    def calculate_id(self) -> PeerId:
         """ Calculate and return the id based on the public key.
         """
         assert self.public_key is not None
         public_der = self.public_key.public_bytes(encoding=serialization.Encoding.DER,
                                                   format=serialization.PublicFormat.SubjectPublicKeyInfo)
         h1 = hashlib.sha256(public_der)
         h2 = hashlib.sha256(h1.digest())
-        return h2.hexdigest()
+        return PeerId(h2.digest())
 
     def get_public_key(self) -> str:
         """ Return the public key in DER encoding as an `str`.
@@ -189,7 +190,7 @@ def create_from_json(cls, data: dict[str, Any]) -> 'Peer':
         from a peer connection.
         """
         obj = cls(auto_generate_keys=False)
-        obj.id = data['id']
+        obj.id = PeerId(data['id'])
 
         if 'pubKey' in data:
             public_key_der = base64.b64decode(data['pubKey'])
@@ -252,7 +253,7 @@ def to_json(self, include_private_key: bool = False) -> dict[str, Any]:
                                                   format=serialization.PublicFormat.SubjectPublicKeyInfo)
         # This format is compatible with libp2p.
         result = {
-            'id': self.id,
+            'id': str(self.id),
             'pubKey': base64.b64encode(public_der).decode('utf-8'),
             'entrypoints': self.entrypoints_as_str(),
         }

hathor/p2p/peer_id.py

@@ -0,0 +1,19 @@
+# Copyright 2024 Hathor Labs
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from hathor.types import Hash
+
+
+class PeerId(Hash):
+    pass

hathor/p2p/peer_storage.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 
 from hathor.p2p.peer import Peer
+from hathor.p2p.peer_id import PeerId
 
 
-class PeerStorage(dict[str, Peer]):
+class PeerStorage(dict[PeerId, Peer]):
     """ PeerStorage is used to store all known peers in memory.
     It is a dict of peer objects, and peers can be retrieved by their `peer.id`.
     """

hathor/p2p/protocol.py

@@ -27,6 +27,7 @@
 from hathor.p2p.entrypoint import Entrypoint
 from hathor.p2p.messages import ProtocolMessages
 from hathor.p2p.peer import Peer
+from hathor.p2p.peer_id import PeerId
 from hathor.p2p.rate_limiter import RateLimiter
 from hathor.p2p.states import BaseState, HelloState, PeerIdState, ReadyState
 from hathor.p2p.sync_version import SyncVersion
@@ -192,7 +193,7 @@ def get_short_remote(self) -> str:
         assert self.transport is not None
         return format_address(self.transport.getPeer())
 
-    def get_peer_id(self) -> Optional[str]:
+    def get_peer_id(self) -> Optional[PeerId]:
         """Get peer id for logging."""
         if self.peer and self.peer.id:
             return self.peer.id
@@ -201,7 +202,7 @@ def get_peer_id(self) -> Optional[str]:
     def get_short_peer_id(self) -> Optional[str]:
         """Get short peer id for logging."""
         if self.peer and self.peer.id:
-            return self.peer.id[:7]
+            return str(self.peer.id)[:7]
         return None
 
     def get_logger_context(self) -> dict[str, Optional[str]]:

hathor/p2p/resources/status.py

@@ -64,7 +64,7 @@ def render_GET(self, request):
             status = {}
             status[conn.state.sync_agent.name] = conn.state.sync_agent.get_status()
             connected_peers.append({
-                'id': conn.peer.id,
+                'id': str(conn.peer.id),
                 'app_version': conn.app_version,
                 'current_time': now,
                 'uptime': now - conn.connection_time,
@@ -82,7 +82,7 @@ def render_GET(self, request):
         known_peers = []
         for peer in self.manager.connections.peer_storage.values():
             known_peers.append({
-                'id': peer.id,
+                'id': str(peer.id),
                 'entrypoints': peer.entrypoints_as_str(),
                 'last_seen': now - peer.last_seen,
                 'flags': [flag.value for flag in peer.flags],
@@ -102,7 +102,7 @@ def render_GET(self, request):
 
         data = {
             'server': {
-                'id': self.manager.connections.my_peer.id,
+                'id': str(self.manager.connections.my_peer.id),
                 'app_version': app,
                 'state': self.manager.state.value,
                 'network': self.manager.network,

hathor/p2p/states/peer_id.py

@@ -19,6 +19,7 @@
 from hathor.conf.settings import HathorSettings
 from hathor.p2p.messages import ProtocolMessages
 from hathor.p2p.peer import Peer
+from hathor.p2p.peer_id import PeerId
 from hathor.p2p.states.base import BaseState
 from hathor.util import json_dumps, json_loads
 
@@ -68,7 +69,7 @@ def send_peer_id(self) -> None:
         protocol = self.protocol
         my_peer = protocol.my_peer
         hello = {
-            'id': my_peer.id,
+            'id': str(my_peer.id),
             'pubKey': my_peer.get_public_key(),
             'entrypoints': my_peer.entrypoints_as_str(),
         }
@@ -139,7 +140,7 @@ async def handle_peer_id(self, payload: str) -> None:
 
         self.send_ready()
 
-    def _should_block_peer(self, peer_id: str) -> bool:
+    def _should_block_peer(self, peer_id: PeerId) -> bool:
         """ Determine if peer should not be allowed to connect.
 
         Currently this is only because the peer is not in a whitelist and whitelist blocking is active.