diff --git a/hathor/builder/builder.py b/hathor/builder/builder.py index 8160bae17..cffa2c08e 100644 --- a/hathor/builder/builder.py +++ b/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, diff --git a/hathor/builder/cli_builder.py b/hathor/builder/cli_builder.py index 54e7b8fa4..f13b88dad 100644 --- a/hathor/builder/cli_builder.py +++ b/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, diff --git a/hathor/manager.py b/hathor/manager.py index 0218fd663..28eb7930d 100644 --- a/hathor/manager.py +++ b/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 diff --git a/hathor/metrics.py b/hathor/metrics.py index 92b73d3ad..046c8c54d 100644 --- a/hathor/metrics.py +++ b/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, diff --git a/hathor/p2p/entrypoint.py b/hathor/p2p/entrypoint.py index 3340f784d..23ead1199 100644 --- a/hathor/p2p/entrypoint.py +++ b/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.""" diff --git a/hathor/p2p/manager.py b/hathor/p2p/manager.py index 4cf4e240e..f70cd1424 100644 --- a/hathor/p2p/manager.py +++ b/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 diff --git a/hathor/p2p/netfilter/matches.py b/hathor/p2p/netfilter/matches.py index d686ac7aa..d7eb714bd 100644 --- a/hathor/p2p/netfilter/matches.py +++ b/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 diff --git a/hathor/p2p/peer.py b/hathor/p2p/peer.py index 9d40ec7b6..bb48e46fd 100644 --- a/hathor/p2p/peer.py +++ b/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,7 +136,7 @@ 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 @@ -143,7 +144,7 @@ def calculate_id(self) -> str: 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(), } diff --git a/hathor/p2p/peer_id.py b/hathor/p2p/peer_id.py new file mode 100644 index 000000000..9a3c5eaca --- /dev/null +++ b/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 diff --git a/hathor/p2p/peer_storage.py b/hathor/p2p/peer_storage.py index 6efb8a549..b6be116b7 100644 --- a/hathor/p2p/peer_storage.py +++ b/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`. """ diff --git a/hathor/p2p/protocol.py b/hathor/p2p/protocol.py index eeec521f2..5a77bdacb 100644 --- a/hathor/p2p/protocol.py +++ b/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]]: diff --git a/hathor/p2p/resources/status.py b/hathor/p2p/resources/status.py index 225665930..bd099d3e1 100644 --- a/hathor/p2p/resources/status.py +++ b/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, diff --git a/hathor/p2p/states/peer_id.py b/hathor/p2p/states/peer_id.py index 439ccd4b0..7534487b3 100644 --- a/hathor/p2p/states/peer_id.py +++ b/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. diff --git a/hathor/p2p/states/ready.py b/hathor/p2p/states/ready.py index c04d97b4c..1d011cf00 100644 --- a/hathor/p2p/states/ready.py +++ b/hathor/p2p/states/ready.py @@ -165,7 +165,7 @@ def send_peers(self, peer_list: Iterable['Peer']) -> None: for peer in peer_list: if peer.entrypoints: data.append({ - 'id': peer.id, + 'id': str(peer.id), 'entrypoints': peer.entrypoints_as_str(), }) self.send_message(ProtocolMessages.PEERS, json_dumps(data)) diff --git a/hathor/p2p/utils.py b/hathor/p2p/utils.py index 7141030d1..c0a25f3d8 100644 --- a/hathor/p2p/utils.py +++ b/hathor/p2p/utils.py @@ -31,6 +31,7 @@ from hathor.indexes.height_index import HeightInfo from hathor.p2p.entrypoint import Entrypoint from hathor.p2p.peer_discovery import DNSPeerDiscovery +from hathor.p2p.peer_id import PeerId from hathor.transaction.genesis import get_representation_for_all_genesis @@ -142,7 +143,7 @@ def parse_file(text: str, *, header: Optional[str] = None) -> list[str]: return list(nonblank_lines) -def parse_whitelist(text: str, *, header: Optional[str] = None) -> set[str]: +def parse_whitelist(text: str, *, header: Optional[str] = None) -> set[PeerId]: """ Parses the list of whitelist peer ids Example: @@ -161,12 +162,7 @@ def parse_whitelist(text: str, *, header: Optional[str] = None) -> set[str]: """ lines = parse_file(text, header=header) - peerids = {line.split()[0] for line in lines} - for peerid in peerids: - bpeerid = bytes.fromhex(peerid) - if len(bpeerid) != 32: - raise ValueError('invalid peerid size') - return peerids + return {PeerId(line.split()[0]) for line in lines} def format_address(addr: IAddress) -> str: diff --git a/hathor/stratum/stratum.py b/hathor/stratum/stratum.py index 92111a431..e3658d457 100644 --- a/hathor/stratum/stratum.py +++ b/hathor/stratum/stratum.py @@ -670,7 +670,7 @@ def create_job_tx(self, jobid: UUID) -> BaseTransaction: assert self.miner_id is not None # Only get first 32 bytes of peer_id because block data is limited to 100 bytes - data = '{}-{}-{}'.format(peer_id[:32], self.miner_id.hex, jobid.hex).encode() + data = '{}-{}-{}'.format(str(peer_id)[:32], self.miner_id.hex, jobid.hex).encode() data = data[:self._settings.BLOCK_DATA_MAX_SIZE] block = self.manager.generate_mining_block(data=data, address=self.miner_address, merge_mined=self.merged_mining) diff --git a/hathor/sysctl/p2p/manager.py b/hathor/sysctl/p2p/manager.py index e821039bd..9f9856a42 100644 --- a/hathor/sysctl/p2p/manager.py +++ b/hathor/sysctl/p2p/manager.py @@ -15,6 +15,7 @@ import os from hathor.p2p.manager import ConnectionsManager +from hathor.p2p.peer_id import PeerId from hathor.p2p.sync_version import SyncVersion from hathor.p2p.utils import discover_hostname from hathor.sysctl.exception import SysctlException @@ -161,11 +162,11 @@ def set_lc_sync_update_interval(self, value: float) -> None: def get_always_enable_sync(self) -> list[str]: """Return the list of sync-always-enabled peers.""" - return list(self.connections.always_enable_sync) + return list(map(str, self.connections.always_enable_sync)) def set_always_enable_sync(self, values: list[str]) -> None: """Change the list of sync-always-enabled peers.""" - self.connections.set_always_enable_sync(values) + self.connections.set_always_enable_sync(list(map(PeerId, values))) def set_always_enable_sync_readtxt(self, file_path: str) -> None: """Update the list of sync-always-enabled peers from a file.""" @@ -174,7 +175,7 @@ def set_always_enable_sync_readtxt(self, file_path: str) -> None: values: list[str] with open(file_path, 'r') as fp: values = parse_text(fp.read()) - self.connections.set_always_enable_sync(values) + self.connections.set_always_enable_sync(list(map(PeerId, values))) def get_max_enabled_sync(self) -> int: """Return the maximum number of peers running sync simultaneously.""" @@ -230,7 +231,11 @@ def set_kill_connection(self, peer_id: str, force: bool = False) -> None: self.connections.disconnect_all_peers(force=force) return - conn = self.connections.connected_peers.get(peer_id, None) + try: + peer_id_obj = PeerId(peer_id) + except ValueError: + raise SysctlException('invalid peer-id') + conn = self.connections.connected_peers.get(peer_id_obj, None) if conn is None: self.log.warn('Killing connection', peer_id=peer_id) raise SysctlException('peer-id is not connected') diff --git a/tests/event/event_simulation_tester.py b/tests/event/event_simulation_tester.py index 5541c6397..1383a4fd8 100644 --- a/tests/event/event_simulation_tester.py +++ b/tests/event/event_simulation_tester.py @@ -40,8 +40,7 @@ def _create_artifacts(self) -> None: .enable_event_queue() artifacts = self.simulator.create_artifacts(builder) - assert peer.id is not None - self.peer_id: str = peer.id + self.peer_id: str = str(peer.id) self.manager = artifacts.manager self.manager.allow_mining_without_peers() self.settings = artifacts.settings diff --git a/tests/p2p/test_peer_id.py b/tests/p2p/test_peer_id.py index b947f8337..f870f824a 100644 --- a/tests/p2p/test_peer_id.py +++ b/tests/p2p/test_peer_id.py @@ -8,6 +8,7 @@ from hathor.p2p.entrypoint import Entrypoint from hathor.p2p.peer import InvalidPeerIdException, Peer +from hathor.p2p.peer_id import PeerId from hathor.p2p.peer_storage import PeerStorage from hathor.util import not_none from tests import unittest @@ -17,7 +18,7 @@ class PeerIdTest(unittest.TestCase): def test_invalid_id(self) -> None: p1 = Peer() - p1.id = not_none(p1.id)[::-1] + p1.id = PeerId(str(not_none(p1.id))[::-1]) self.assertRaises(InvalidPeerIdException, p1.validate) def test_invalid_public_key(self) -> None: diff --git a/tests/resources/p2p/test_status.py b/tests/resources/p2p/test_status.py index 7ab42ae68..5b156326f 100644 --- a/tests/resources/p2p/test_status.py +++ b/tests/resources/p2p/test_status.py @@ -89,10 +89,10 @@ def test_get_with_one_peer(self): self.assertGreater(server_data['uptime'], 0) self.assertEqual(len(known_peers), 1) - self.assertEqual(known_peers[0]['id'], self.manager2.my_peer.id) + self.assertEqual(known_peers[0]['id'], str(self.manager2.my_peer.id)) self.assertEqual(len(connections['connected_peers']), 1) - self.assertEqual(connections['connected_peers'][0]['id'], self.manager2.my_peer.id) + self.assertEqual(connections['connected_peers'][0]['id'], str(self.manager2.my_peer.id)) @inlineCallbacks def test_connecting_peers(self): diff --git a/tests/sysctl/test_p2p.py b/tests/sysctl/test_p2p.py index 6a3980c5e..ec0366888 100644 --- a/tests/sysctl/test_p2p.py +++ b/tests/sysctl/test_p2p.py @@ -2,6 +2,7 @@ import tempfile from unittest.mock import MagicMock +from hathor.p2p.peer_id import PeerId from hathor.sysctl import ConnectionsManagerSysctl from hathor.sysctl.exception import SysctlException from tests import unittest @@ -100,9 +101,12 @@ def test_always_enable_sync(self): connections = manager.connections sysctl = ConnectionsManagerSysctl(connections) - sysctl.unsafe_set('always_enable_sync', ['peer-1', 'peer-2']) - self.assertEqual(connections.always_enable_sync, {'peer-1', 'peer-2'}) - self.assertEqual(set(sysctl.get('always_enable_sync')), {'peer-1', 'peer-2'}) + peer_id_1 = '0e2bd0d8cd1fb6d040801c32ec27e8986ce85eb8810b6c878dcad15bce3b5b1e' + peer_id_2 = '2ff0d2c80c50f724de79f132a2f8cae576c64b57ea531d400577adf7db3e7c15' + + sysctl.unsafe_set('always_enable_sync', [peer_id_1, peer_id_2]) + self.assertEqual(connections.always_enable_sync, {PeerId(peer_id_1), PeerId(peer_id_2)}) + self.assertEqual(set(sysctl.get('always_enable_sync')), {peer_id_1, peer_id_2}) sysctl.unsafe_set('always_enable_sync', []) self.assertEqual(connections.always_enable_sync, set()) @@ -110,8 +114,8 @@ def test_always_enable_sync(self): with tempfile.TemporaryDirectory() as dir_path: content = [ - 'peer-id-1', - 'peer-id-2', + peer_id_1, + peer_id_2, ] file_path = os.path.join(dir_path, 'a.txt') @@ -120,7 +124,7 @@ def test_always_enable_sync(self): fp.close() sysctl.unsafe_set('always_enable_sync.readtxt', file_path) - self.assertEqual(connections.always_enable_sync, set(content)) + self.assertEqual(connections.always_enable_sync, {PeerId(peer_id_1), PeerId(peer_id_2)}) self.assertEqual(set(sysctl.get('always_enable_sync')), set(content)) def test_available_sync_versions(self): @@ -166,9 +170,9 @@ def test_kill_one_connection(self): p2p_manager = manager.connections sysctl = ConnectionsManagerSysctl(p2p_manager) - peer_id = 'my-peer-id' + peer_id = '0e2bd0d8cd1fb6d040801c32ec27e8986ce85eb8810b6c878dcad15bce3b5b1e' conn = MagicMock() - p2p_manager.connected_peers[peer_id] = conn + p2p_manager.connected_peers[PeerId(peer_id)] = conn self.assertEqual(conn.disconnect.call_count, 0) sysctl.unsafe_set('kill_connection', peer_id) self.assertEqual(conn.disconnect.call_count, 1)