From dda05738afb55c4a77d8f5232d031b47839c3332 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Tue, 4 Aug 2026 03:20:37 -0400 Subject: [PATCH 1/3] Validate discovered peer metadata --- macfleet/pool/discovery.py | 68 ++++++++++++++++++--- tests/test_pool/test_discovery.py | 98 +++++++++++++++++++++++++++++++ 2 files changed, 159 insertions(+), 7 deletions(-) create mode 100644 tests/test_pool/test_discovery.py diff --git a/macfleet/pool/discovery.py b/macfleet/pool/discovery.py index ef63eaa..c2df8cb 100644 --- a/macfleet/pool/discovery.py +++ b/macfleet/pool/discovery.py @@ -6,6 +6,7 @@ """ import asyncio +import math import socket import threading import time @@ -22,6 +23,32 @@ MACFLEET_SERVICE_TYPE = DEFAULT_SERVICE_TYPE DEFAULT_TTL = 120 +MAX_DISCOVERY_GPU_CORES = 1024 +MAX_DISCOVERY_RAM_GB = 65_536 + + +def _parse_bounded_int( + value: bytes, + *, + minimum: int, + maximum: int, + field: str, +) -> int: + """Parse a bounded integer from an untrusted mDNS TXT value.""" + parsed = int(value.decode()) + if not minimum <= parsed <= maximum: + raise ValueError( + f"{field} must be between {minimum} and {maximum}, got {parsed}" + ) + return parsed + + +def _parse_nonnegative_finite_float(value: bytes, *, field: str) -> float: + """Parse a non-negative finite float from an untrusted mDNS TXT value.""" + parsed = float(value.decode()) + if not math.isfinite(parsed) or parsed < 0: + raise ValueError(f"{field} must be non-negative and finite") + return parsed @dataclass @@ -156,20 +183,47 @@ def _prop(key: bytes, default: bytes) -> bytes: return val if isinstance(val, bytes) else default node_id = _prop(b"node_id", b"").decode() or hostname - gpu_cores = int(_prop(b"gpu_cores", b"0").decode()) - ram_gb = int(_prop(b"ram_gb", b"0").decode()) + gpu_cores = _parse_bounded_int( + _prop(b"gpu_cores", b"0"), + minimum=0, + maximum=MAX_DISCOVERY_GPU_CORES, + field="gpu_cores", + ) + ram_gb = _parse_bounded_int( + _prop(b"ram_gb", b"0"), + minimum=0, + maximum=MAX_DISCOVERY_RAM_GB, + field="ram_gb", + ) chip_name = _prop(b"chip_name", b"unknown").decode() link_types = _prop(b"link_types", b"").decode() pool_version = _prop(b"pool_version", b"0.0.0").decode() - compute_score = float(_prop(b"compute_score", b"0").decode()) - # data_port advertised since v2.2; 2.1.x peers lack this. DiscoveredNode - # __post_init__ falls back to heartbeat port + 1 when 0. - data_port = int(_prop(b"data_port", b"0").decode()) + compute_score = _parse_nonnegative_finite_float( + _prop(b"compute_score", b"0"), field="compute_score" + ) + # data_port advertised since v2.2; 2.1.x peers lack this. Derive + # heartbeat port + 1 when absent, after checking for overflow. + data_port = _parse_bounded_int( + _prop(b"data_port", b"0"), + minimum=0, + maximum=65_535, + field="data_port", + ) network_links = deserialize_network_links( _prop(b"network_links", b"").decode() ) - if info.port is None: + if ( + isinstance(info.port, bool) + or not isinstance(info.port, int) + or not 1 <= info.port <= 65_535 + ): + return None + if data_port == 0: + if info.port == 65_535: + return None + data_port = info.port + 1 + if data_port == info.port: return None return DiscoveredNode( diff --git a/tests/test_pool/test_discovery.py b/tests/test_pool/test_discovery.py new file mode 100644 index 0000000..70d7aeb --- /dev/null +++ b/tests/test_pool/test_discovery.py @@ -0,0 +1,98 @@ +"""Validation tests for mDNS peer discovery metadata.""" + +from __future__ import annotations + +import socket + +import pytest +from zeroconf import ServiceInfo + +from macfleet.pool.discovery import PoolServiceListener + +SERVICE_TYPE = "_macfleet._tcp.local." + + +def _service_info( + *, + properties: dict[bytes, bytes] | None = None, + port: int = 50051, +) -> ServiceInfo: + base_properties = { + b"node_id": b"peer-1", + b"gpu_cores": b"20", + b"ram_gb": b"64", + b"chip_name": b"Apple M4 Ultra", + b"link_types": b"ethernet", + b"pool_version": b"2.2.1", + b"compute_score": b"512.0", + b"data_port": b"50052", + } + if properties: + base_properties.update(properties) + return ServiceInfo( + SERVICE_TYPE, + f"peer-1.{SERVICE_TYPE}", + addresses=[socket.inet_aton("192.0.2.10")], + port=port, + properties=base_properties, + server="peer-1.local.", + ) + + +class TestDiscoveryNumericValidation: + def test_valid_metadata_is_parsed(self): + node = PoolServiceListener()._parse_service_info(_service_info()) + + assert node is not None + assert node.node_id == "peer-1" + assert node.port == 50051 + assert node.data_port == 50052 + assert node.gpu_cores == 20 + assert node.ram_gb == 64 + assert node.compute_score == 512.0 + + @pytest.mark.parametrize( + ("key", "value"), + [ + (b"gpu_cores", b"-1"), + (b"gpu_cores", b"1025"), + (b"gpu_cores", b"1.5"), + (b"ram_gb", b"-1"), + (b"ram_gb", b"65537"), + (b"ram_gb", b"16.5"), + ], + ) + def test_invalid_hardware_counts_reject_peer(self, key, value): + info = _service_info(properties={key: value}) + + assert PoolServiceListener()._parse_service_info(info) is None + + @pytest.mark.parametrize("value", [b"-1", b"nan", b"inf", b"-inf"]) + def test_invalid_compute_score_rejects_peer(self, value): + info = _service_info(properties={b"compute_score": value}) + + assert PoolServiceListener()._parse_service_info(info) is None + + @pytest.mark.parametrize("value", [b"-1", b"65536", b"1.5"]) + def test_invalid_data_port_rejects_peer(self, value): + info = _service_info(properties={b"data_port": value}) + + assert PoolServiceListener()._parse_service_info(info) is None + + def test_same_control_and_data_port_rejects_peer(self): + info = _service_info(properties={b"data_port": b"50051"}) + + assert PoolServiceListener()._parse_service_info(info) is None + + def test_missing_data_port_uses_next_control_port(self): + info = _service_info(properties={b"data_port": b"0"}, port=9000) + + node = PoolServiceListener()._parse_service_info(info) + + assert node is not None + assert node.data_port == 9001 + + def test_missing_data_port_rejects_control_port_at_maximum(self): + info = _service_info(properties={b"data_port": b"0"}, port=65535) + + assert PoolServiceListener()._parse_service_info(info) is None From f73f60724b87b070e30d6d28f0f5e8624826b308 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Tue, 4 Aug 2026 03:22:20 -0400 Subject: [PATCH 2/3] Bound discovery identity fields --- macfleet/pool/discovery.py | 82 ++++++++++++++++++++++++++++--- tests/test_pool/test_discovery.py | 32 ++++++++++++ 2 files changed, 108 insertions(+), 6 deletions(-) diff --git a/macfleet/pool/discovery.py b/macfleet/pool/discovery.py index c2df8cb..6491040 100644 --- a/macfleet/pool/discovery.py +++ b/macfleet/pool/discovery.py @@ -19,12 +19,20 @@ import macfleet from macfleet.pool.network import NetworkLink from macfleet.pool.topology import deserialize_network_links, serialize_network_links -from macfleet.security.auth import DEFAULT_SERVICE_TYPE, SecurityConfig +from macfleet.security.auth import ( + DEFAULT_SERVICE_TYPE, + MAX_NODE_ID_BYTES, + SecurityConfig, +) MACFLEET_SERVICE_TYPE = DEFAULT_SERVICE_TYPE DEFAULT_TTL = 120 MAX_DISCOVERY_GPU_CORES = 1024 MAX_DISCOVERY_RAM_GB = 65_536 +MAX_DISCOVERY_HOSTNAME_BYTES = 253 +MAX_DISCOVERY_CHIP_NAME_BYTES = 128 +MAX_DISCOVERY_LINK_TYPES_BYTES = 128 +MAX_DISCOVERY_VERSION_BYTES = 64 def _parse_bounded_int( @@ -51,6 +59,39 @@ def _parse_nonnegative_finite_float(value: bytes, *, field: str) -> float: return parsed +def _validate_bounded_text( + value: str, + *, + maximum_bytes: int, + field: str, + allow_empty: bool = True, +) -> str: + """Require bounded printable text before retaining or displaying it.""" + if (not allow_empty and not value) or len(value.encode()) > maximum_bytes: + raise ValueError(f"{field} exceeds its size limit or is empty") + if any(not character.isprintable() for character in value): + raise ValueError(f"{field} contains control characters") + return value + + +def _decode_bounded_text( + value: bytes, + *, + maximum_bytes: int, + field: str, + allow_empty: bool = True, +) -> str: + """Decode an untrusted mDNS TXT value and validate display safety.""" + if len(value) > maximum_bytes: + raise ValueError(f"{field} exceeds its size limit") + return _validate_bounded_text( + value.decode(), + maximum_bytes=maximum_bytes, + field=field, + allow_empty=allow_empty, + ) + + @dataclass class DiscoveredNode: """A node discovered via mDNS. @@ -176,13 +217,28 @@ def _parse_service_info(self, info: ServiceInfo) -> Optional[DiscoveredNode]: return None props = info.properties - hostname = info.server.rstrip(".") + hostname = _validate_bounded_text( + info.server.rstrip("."), + maximum_bytes=MAX_DISCOVERY_HOSTNAME_BYTES, + field="hostname", + allow_empty=False, + ) def _prop(key: bytes, default: bytes) -> bytes: val = props.get(key) return val if isinstance(val, bytes) else default - node_id = _prop(b"node_id", b"").decode() or hostname + node_id = _decode_bounded_text( + _prop(b"node_id", b""), + maximum_bytes=MAX_NODE_ID_BYTES, + field="node_id", + ) or hostname + node_id = _validate_bounded_text( + node_id, + maximum_bytes=MAX_NODE_ID_BYTES, + field="node_id", + allow_empty=False, + ) gpu_cores = _parse_bounded_int( _prop(b"gpu_cores", b"0"), minimum=0, @@ -195,9 +251,23 @@ def _prop(key: bytes, default: bytes) -> bytes: maximum=MAX_DISCOVERY_RAM_GB, field="ram_gb", ) - chip_name = _prop(b"chip_name", b"unknown").decode() - link_types = _prop(b"link_types", b"").decode() - pool_version = _prop(b"pool_version", b"0.0.0").decode() + chip_name = _decode_bounded_text( + _prop(b"chip_name", b"unknown"), + maximum_bytes=MAX_DISCOVERY_CHIP_NAME_BYTES, + field="chip_name", + allow_empty=False, + ) + link_types = _decode_bounded_text( + _prop(b"link_types", b""), + maximum_bytes=MAX_DISCOVERY_LINK_TYPES_BYTES, + field="link_types", + ) + pool_version = _decode_bounded_text( + _prop(b"pool_version", b"0.0.0"), + maximum_bytes=MAX_DISCOVERY_VERSION_BYTES, + field="pool_version", + allow_empty=False, + ) compute_score = _parse_nonnegative_finite_float( _prop(b"compute_score", b"0"), field="compute_score" ) diff --git a/tests/test_pool/test_discovery.py b/tests/test_pool/test_discovery.py index 70d7aeb..6842474 100644 --- a/tests/test_pool/test_discovery.py +++ b/tests/test_pool/test_discovery.py @@ -96,3 +96,35 @@ def test_missing_data_port_rejects_control_port_at_maximum(self): info = _service_info(properties={b"data_port": b"0"}, port=65535) assert PoolServiceListener()._parse_service_info(info) is None + + +class TestDiscoveryTextValidation: + @pytest.mark.parametrize( + ("key", "value"), + [ + (b"node_id", b"n" * 129), + (b"node_id", b"peer\x1b[31m"), + (b"chip_name", b"c" * 129), + (b"chip_name", b"M4\nforged-output"), + (b"link_types", b"l" * 129), + (b"pool_version", b"v" * 65), + ], + ) + def test_unsafe_text_property_rejects_peer(self, key, value): + info = _service_info(properties={key: value}) + + assert PoolServiceListener()._parse_service_info(info) is None + + def test_empty_node_id_falls_back_to_hostname(self): + info = _service_info(properties={b"node_id": b""}) + + node = PoolServiceListener()._parse_service_info(info) + + assert node is not None + assert node.node_id == "peer-1.local" + + def test_unsafe_hostname_rejects_peer(self): + info = _service_info() + info.server = "peer\nforged.local." + + assert PoolServiceListener()._parse_service_info(info) is None From e7b7c9b86604b4c58c743ae7f319ff8b6e18907c Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Tue, 4 Aug 2026 03:26:00 -0400 Subject: [PATCH 3/3] Keep discovery advertisements wire-safe --- macfleet/pool/discovery.py | 122 +++++++++++++++++++++++++++--- tests/test_pool/test_discovery.py | 66 +++++++++++++++- 2 files changed, 178 insertions(+), 10 deletions(-) diff --git a/macfleet/pool/discovery.py b/macfleet/pool/discovery.py index 6491040..fa4c351 100644 --- a/macfleet/pool/discovery.py +++ b/macfleet/pool/discovery.py @@ -33,6 +33,8 @@ MAX_DISCOVERY_CHIP_NAME_BYTES = 128 MAX_DISCOVERY_LINK_TYPES_BYTES = 128 MAX_DISCOVERY_VERSION_BYTES = 64 +MAX_MDNS_INSTANCE_NAME_BYTES = 63 +MAX_MDNS_TXT_ITEM_BYTES = 255 def _parse_bounded_int( @@ -44,6 +46,22 @@ def _parse_bounded_int( ) -> int: """Parse a bounded integer from an untrusted mDNS TXT value.""" parsed = int(value.decode()) + return _validate_bounded_int( + parsed, minimum=minimum, maximum=maximum, field=field + ) + + +def _validate_bounded_int( + value: object, + *, + minimum: int, + maximum: int, + field: str, +) -> int: + """Require an integer value within protocol-safe bounds.""" + if isinstance(value, bool) or not isinstance(value, int): + raise ValueError(f"{field} must be an integer") + parsed = value if not minimum <= parsed <= maximum: raise ValueError( f"{field} must be between {minimum} and {maximum}, got {parsed}" @@ -54,6 +72,14 @@ def _parse_bounded_int( def _parse_nonnegative_finite_float(value: bytes, *, field: str) -> float: """Parse a non-negative finite float from an untrusted mDNS TXT value.""" parsed = float(value.decode()) + return _validate_nonnegative_finite_number(parsed, field=field) + + +def _validate_nonnegative_finite_number(value: object, *, field: str) -> float: + """Require a numeric value suitable for scoring and serialization.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise ValueError(f"{field} must be a number") + parsed = float(value) if not math.isfinite(parsed) or parsed < 0: raise ValueError(f"{field} must be non-negative and finite") return parsed @@ -92,6 +118,42 @@ def _decode_bounded_text( ) +def _encode_bounded_text( + value: str, + *, + maximum_bytes: int, + field: str, + allow_empty: bool = True, +) -> bytes: + """Validate local text before encoding it into an mDNS TXT record.""" + return _validate_bounded_text( + value, + maximum_bytes=maximum_bytes, + field=field, + allow_empty=allow_empty, + ).encode() + + +def _mdns_txt_item_fits(key: bytes, value: bytes) -> bool: + """Return whether ``key=value`` fits one DNS TXT character-string.""" + return len(key) + 1 + len(value) <= MAX_MDNS_TXT_ITEM_BYTES + + +def _serialize_network_links_for_mdns( + network_links: tuple[NetworkLink, ...], +) -> bytes: + """Keep as many links as fit in mDNS's 255-byte TXT item limit.""" + selected: list[NetworkLink] = [] + encoded = b"[]" + key = b"network_links" + for link in network_links: + candidate = serialize_network_links((*selected, link)).encode() + if candidate != encoded and _mdns_txt_item_fits(key, candidate): + selected.append(link) + encoded = candidate + return encoded + + @dataclass class DiscoveredNode: """A node discovered via mDNS. @@ -390,22 +452,64 @@ def _build_properties( sensitive — it's just the TCP port the transport listens on. Peers need it to initiate the authenticated handshake. """ + node_id_bytes = _encode_bounded_text( + node_id, + maximum_bytes=MAX_MDNS_INSTANCE_NAME_BYTES, + field="node_id", + allow_empty=False, + ) + pool_version_bytes = _encode_bounded_text( + macfleet.__version__, + maximum_bytes=MAX_DISCOVERY_VERSION_BYTES, + field="pool_version", + allow_empty=False, + ) + data_port = _validate_bounded_int( + data_port, minimum=1, maximum=65_535, field="data_port" + ) + data_port_bytes = str(data_port).encode() if self._security.is_secure: return { - b"node_id": node_id.encode(), - b"pool_version": macfleet.__version__.encode(), - b"data_port": str(data_port).encode(), + b"node_id": node_id_bytes, + b"pool_version": pool_version_bytes, + b"data_port": data_port_bytes, } + gpu_cores = _validate_bounded_int( + gpu_cores, + minimum=0, + maximum=MAX_DISCOVERY_GPU_CORES, + field="gpu_cores", + ) + ram_gb = _validate_bounded_int( + ram_gb, + minimum=0, + maximum=MAX_DISCOVERY_RAM_GB, + field="ram_gb", + ) + chip_name_bytes = _encode_bounded_text( + chip_name, + maximum_bytes=MAX_DISCOVERY_CHIP_NAME_BYTES, + field="chip_name", + allow_empty=False, + ) + link_types_bytes = _encode_bounded_text( + link_types, + maximum_bytes=MAX_DISCOVERY_LINK_TYPES_BYTES, + field="link_types", + ) + compute_score = _validate_nonnegative_finite_number( + compute_score, field="compute_score" + ) return { - b"node_id": node_id.encode(), + b"node_id": node_id_bytes, b"gpu_cores": str(gpu_cores).encode(), b"ram_gb": str(ram_gb).encode(), - b"chip_name": chip_name.encode(), - b"link_types": link_types.encode(), - b"network_links": serialize_network_links(network_links).encode(), - b"pool_version": macfleet.__version__.encode(), + b"chip_name": chip_name_bytes, + b"link_types": link_types_bytes, + b"network_links": _serialize_network_links_for_mdns(network_links), + b"pool_version": pool_version_bytes, b"compute_score": f"{compute_score:.1f}".encode(), - b"data_port": str(data_port).encode(), + b"data_port": data_port_bytes, } def register_node( diff --git a/tests/test_pool/test_discovery.py b/tests/test_pool/test_discovery.py index 6842474..adec727 100644 --- a/tests/test_pool/test_discovery.py +++ b/tests/test_pool/test_discovery.py @@ -7,7 +7,9 @@ import pytest from zeroconf import ServiceInfo -from macfleet.pool.discovery import PoolServiceListener +from macfleet.pool.discovery import PoolServiceListener, ServiceRegistry +from macfleet.pool.network import LinkType, NetworkLink +from macfleet.pool.topology import deserialize_network_links SERVICE_TYPE = "_macfleet._tcp.local." @@ -128,3 +130,65 @@ def test_unsafe_hostname_rejects_peer(self): info.server = "peer\nforged.local." assert PoolServiceListener()._parse_service_info(info) is None + + +class TestDiscoveryAdvertisementValidation: + def _properties(self, **overrides): + values = { + "node_id": "node-1", + "gpu_cores": 20, + "ram_gb": 64, + "chip_name": "Apple M4 Ultra", + "link_types": "ethernet", + "compute_score": 512.0, + "data_port": 50052, + } + values.update(overrides) + return ServiceRegistry()._build_properties(**values) + + @pytest.mark.parametrize( + "overrides", + [ + {"node_id": "n" * 64}, + {"gpu_cores": -1}, + {"gpu_cores": 1025}, + {"ram_gb": 65537}, + {"chip_name": "M4\nforged-output"}, + {"link_types": "l" * 129}, + {"compute_score": float("nan")}, + {"data_port": 0}, + {"data_port": 65536}, + ], + ) + def test_invalid_local_property_is_rejected(self, overrides): + with pytest.raises(ValueError): + self._properties(**overrides) + + def test_rich_topology_is_trimmed_to_mdns_txt_limit(self): + links = tuple( + NetworkLink( + interface=f"en{index}", + link_type=LinkType.ETHERNET, + ip_address=f"192.0.2.{index + 1}", + bandwidth_mbps=1000.0, + latency_ms=1.0, + loss_rate=0.01, + mtu=9000, + ) + for index in range(8) + ) + + properties = self._properties(network_links=links) + encoded_links = properties[b"network_links"] + + assert len(b"network_links=") + len(encoded_links) <= 255 + restored = deserialize_network_links(encoded_links.decode()) + assert 0 < len(restored) < len(links) + ServiceInfo( + SERVICE_TYPE, + f"node-1.{SERVICE_TYPE}", + addresses=[socket.inet_aton("192.0.2.100")], + port=50051, + properties=properties, + server="node-1.local.", + )