Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
75 changes: 72 additions & 3 deletions hathor/p2p/netfilter/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,15 +17,84 @@
from hathor.p2p.netfilter.rule import NetfilterRule
from hathor.p2p.netfilter.targets import NetfilterReject

# Global mapping to track peer_id -> rule UUID for blacklist management
_peer_id_to_rule_uuid: dict[str, str] = {}

def add_peer_id_blacklist(peer_id_blacklist: list[str]) -> None:
""" Add a list of peer ids to a blacklist using netfilter reject

def add_blacklist_peers(peer_ids: str | list[str]) -> list[str]:
"""Add peer(s) to the blacklist.

Args:
peer_ids: A single peer_id string or a list of peer_id strings

Returns:
List of peer_ids that were successfully added (not already blacklisted)
"""
if isinstance(peer_ids, str):
peer_ids = [peer_ids]

post_peerid = get_table('filter').get_chain('post_peerid')
added_peers: list[str] = []

for peer_id in peer_id_blacklist:
for peer_id in peer_ids:
if not peer_id:
continue

# Skip if already blacklisted
if peer_id in _peer_id_to_rule_uuid:
continue

match = NetfilterMatchPeerId(peer_id)
rule = NetfilterRule(match, NetfilterReject())
post_peerid.add_rule(rule)
_peer_id_to_rule_uuid[peer_id] = rule.uuid
added_peers.append(peer_id)

return added_peers


def remove_blacklist_peers(peer_ids: str | list[str]) -> list[str]:
"""Remove peer(s) from the blacklist.

Args:
peer_ids: A single peer_id string or a list of peer_id strings

Returns:
List of peer_ids that were successfully removed
"""
if isinstance(peer_ids, str):
peer_ids = [peer_ids]

post_peerid = get_table('filter').get_chain('post_peerid')
removed_peers: list[str] = []

for peer_id in peer_ids:
if not peer_id:
continue

rule_uuid = _peer_id_to_rule_uuid.get(peer_id)
if rule_uuid is None:
continue

if post_peerid.delete_rule(rule_uuid):
del _peer_id_to_rule_uuid[peer_id]
removed_peers.append(peer_id)

return removed_peers


def list_blacklist_peers() -> list[str]:
"""List all currently blacklisted peer_ids.

Returns:
List of blacklisted peer_id strings
"""
return list(_peer_id_to_rule_uuid.keys())


def add_peer_id_blacklist(peer_id_blacklist: list[str]) -> None:
"""Add a list of peer ids to a blacklist using netfilter reject.

This is a legacy function that wraps add_blacklist_peers for backward compatibility.
"""
add_blacklist_peers(peer_id_blacklist)
56 changes: 56 additions & 0 deletions hathor/sysctl/p2p/manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
import os

from hathor.p2p.manager import ConnectionsManager
from hathor.p2p.netfilter.utils import add_blacklist_peers, list_blacklist_peers, remove_blacklist_peers
from hathor.p2p.peer_id import PeerId
from hathor.p2p.sync_version import SyncVersion
from hathor.p2p.utils import discover_hostname
Expand Down Expand Up @@ -122,6 +123,21 @@ def __init__(self, connections: ConnectionsManager) -> None:
None,
self.reload_entrypoints_and_connections,
)
self.register(
'blacklist.add_peers',
None,
self.set_blacklist_add_peers,
)
self.register(
'blacklist.remove_peers',
None,
self.set_blacklist_remove_peers,
)
self.register(
'blacklist.list_peers',
self.get_blacklist_list_peers,
None,
)

def set_force_sync_rotate(self) -> None:
"""Force a sync rotate."""
Expand Down Expand Up @@ -269,3 +285,43 @@ def refresh_auto_hostname(self) -> None:
def reload_entrypoints_and_connections(self) -> None:
"""Kill all connections and reload entrypoints from the peer config file."""
self.connections.reload_entrypoints_and_connections()

@signal_handler_safe
def set_blacklist_add_peers(self, peer_ids: str | list[str]) -> None:
"""Add peer(s) to the blacklist. Accepts a single peer-id string or a list of peer-ids."""
# Validate peer IDs
peer_id_list = [peer_ids] if isinstance(peer_ids, str) else peer_ids
try:
for peer_id in peer_id_list:
if peer_id: # Skip empty strings
PeerId(peer_id) # Validate format
except ValueError as e:
raise SysctlException(f'Invalid peer-id format: {e}')

added_peers = add_blacklist_peers(peer_ids)
if added_peers:
self.log.info('Added peers to blacklist', peer_ids=added_peers)
else:
self.log.info('No new peers added to blacklist (already blacklisted or empty)')

@signal_handler_safe
def set_blacklist_remove_peers(self, peer_ids: str | list[str]) -> None:
"""Remove peer(s) from the blacklist. Accepts a single peer-id string or a list of peer-ids."""
# Validate peer IDs
peer_id_list = [peer_ids] if isinstance(peer_ids, str) else peer_ids
try:
for peer_id in peer_id_list:
if peer_id: # Skip empty strings
PeerId(peer_id) # Validate format
except ValueError as e:
raise SysctlException(f'Invalid peer-id format: {e}')

removed_peers = remove_blacklist_peers(peer_ids)
if removed_peers:
self.log.info('Removed peers from blacklist', peer_ids=removed_peers)
else:
self.log.info('No peers removed from blacklist (not found or empty)')

def get_blacklist_list_peers(self) -> list[str]:
"""List all currently blacklisted peer_ids."""
return list_blacklist_peers()
141 changes: 140 additions & 1 deletion hathor_tests/p2p/netfilter/test_utils.py
Original file line number Diff line number Diff line change
@@ -1,9 +1,23 @@
from hathor.p2p.netfilter import get_table
from hathor.p2p.netfilter.utils import add_peer_id_blacklist
from hathor.p2p.netfilter.utils import (
add_blacklist_peers,
add_peer_id_blacklist,
list_blacklist_peers,
remove_blacklist_peers,
)
from hathor_tests import unittest


class NetfilterUtilsTest(unittest.TestCase):
def setUp(self) -> None:
"""Clean up rules and tracking before each test."""
super().setUp()
post_peerid = get_table('filter').get_chain('post_peerid')
post_peerid.rules = []
# Clear the global tracking dictionary
from hathor.p2p.netfilter import utils
utils._peer_id_to_rule_uuid.clear()

def test_peer_id_blacklist(self) -> None:
post_peerid = get_table('filter').get_chain('post_peerid')

Expand All @@ -24,3 +38,128 @@ def test_peer_id_blacklist(self) -> None:
self.assertEqual(data['match']['type'], 'NetfilterMatchPeerId')
self.assertIn(data['match']['match_params']['peer_id'], blacklist)
self.assertEqual(data['target']['type'], 'NetfilterReject')

def test_add_blacklist_peers_with_list(self) -> None:
"""Test adding multiple peers with a list."""
post_peerid = get_table('filter').get_chain('post_peerid')

# Initially empty
self.assertEqual(len(post_peerid.rules), 0)
self.assertEqual(list_blacklist_peers(), [])

# Add peers
peer_ids = ['peer1', 'peer2', 'peer3']
added = add_blacklist_peers(peer_ids)

# All peers should be added
self.assertEqual(sorted(added), sorted(peer_ids))
self.assertEqual(len(post_peerid.rules), 3)
self.assertEqual(sorted(list_blacklist_peers()), sorted(peer_ids))

def test_add_blacklist_peers_with_string(self) -> None:
"""Test adding a single peer with a string."""
post_peerid = get_table('filter').get_chain('post_peerid')

# Add single peer
peer_id = 'single_peer'
added = add_blacklist_peers(peer_id)

self.assertEqual(added, [peer_id])
self.assertEqual(len(post_peerid.rules), 1)
self.assertEqual(list_blacklist_peers(), [peer_id])

def test_add_blacklist_peers_skip_duplicates(self) -> None:
"""Test that adding duplicate peers is skipped."""
post_peerid = get_table('filter').get_chain('post_peerid')

# Add peers first time
peer_ids = ['peer1', 'peer2']
added1 = add_blacklist_peers(peer_ids)
self.assertEqual(sorted(added1), sorted(peer_ids))
self.assertEqual(len(post_peerid.rules), 2)

# Try to add same peers again
added2 = add_blacklist_peers(peer_ids)
self.assertEqual(added2, []) # Nothing added
self.assertEqual(len(post_peerid.rules), 2) # Still 2 rules

# Add mix of new and existing
added3 = add_blacklist_peers(['peer1', 'peer3'])
self.assertEqual(added3, ['peer3']) # Only new peer added
self.assertEqual(len(post_peerid.rules), 3)

def test_add_blacklist_peers_skip_empty(self) -> None:
"""Test that empty strings are skipped."""
peer_ids = ['peer1', '', 'peer2', '']
added = add_blacklist_peers(peer_ids)

self.assertEqual(sorted(added), ['peer1', 'peer2'])
self.assertEqual(sorted(list_blacklist_peers()), ['peer1', 'peer2'])

def test_remove_blacklist_peers_with_list(self) -> None:
"""Test removing multiple peers with a list."""
# Add peers first
peer_ids = ['peer1', 'peer2', 'peer3']
add_blacklist_peers(peer_ids)
self.assertEqual(sorted(list_blacklist_peers()), sorted(peer_ids))

# Remove some peers
to_remove = ['peer1', 'peer3']
removed = remove_blacklist_peers(to_remove)

self.assertEqual(sorted(removed), sorted(to_remove))
self.assertEqual(list_blacklist_peers(), ['peer2'])

def test_remove_blacklist_peers_with_string(self) -> None:
"""Test removing a single peer with a string."""
# Add peers first
add_blacklist_peers(['peer1', 'peer2'])

# Remove one peer
removed = remove_blacklist_peers('peer1')

self.assertEqual(removed, ['peer1'])
self.assertEqual(list_blacklist_peers(), ['peer2'])

def test_remove_blacklist_peers_nonexistent(self) -> None:
"""Test removing peers that don't exist."""
# Add one peer
add_blacklist_peers('peer1')

# Try to remove nonexistent peers
removed = remove_blacklist_peers(['peer2', 'peer3'])

self.assertEqual(removed, [])
self.assertEqual(list_blacklist_peers(), ['peer1'])

# Remove mix of existing and nonexistent
removed2 = remove_blacklist_peers(['peer1', 'peer2'])
self.assertEqual(removed2, ['peer1'])
self.assertEqual(list_blacklist_peers(), [])

def test_remove_blacklist_peers_skip_empty(self) -> None:
"""Test that empty strings are skipped during removal."""
add_blacklist_peers(['peer1', 'peer2'])

removed = remove_blacklist_peers(['peer1', '', 'peer2'])

self.assertEqual(sorted(removed), ['peer1', 'peer2'])
self.assertEqual(list_blacklist_peers(), [])

def test_list_blacklist_peers(self) -> None:
"""Test listing blacklisted peers."""
# Initially empty
self.assertEqual(list_blacklist_peers(), [])

# Add some peers
peer_ids = ['peer1', 'peer2', 'peer3']
add_blacklist_peers(peer_ids)
self.assertEqual(sorted(list_blacklist_peers()), sorted(peer_ids))

# Remove one
remove_blacklist_peers('peer2')
self.assertEqual(sorted(list_blacklist_peers()), ['peer1', 'peer3'])

# Remove all
remove_blacklist_peers(['peer1', 'peer3'])
self.assertEqual(list_blacklist_peers(), [])
Loading
Loading