From d6520f60696fbf03c0a17180c83b058b6fafe483 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Wed, 15 Jul 2026 12:29:00 -0400 Subject: [PATCH 1/4] Harden hardware profile handshake validation --- macfleet/comm/transport.py | 55 ++++++++++++++++++++++++- tests/test_security/test_hw_exchange.py | 35 ++++++++++++++++ 2 files changed, 88 insertions(+), 2 deletions(-) diff --git a/macfleet/comm/transport.py b/macfleet/comm/transport.py index 4244d3f..be439b1 100644 --- a/macfleet/comm/transport.py +++ b/macfleet/comm/transport.py @@ -12,6 +12,7 @@ import asyncio import json import logging +import math import socket import struct import time @@ -115,6 +116,57 @@ def to_json_bytes(self) -> bytes: """Stable JSON encoding (sort_keys so HMAC is deterministic).""" return json.dumps(asdict(self), sort_keys=True).encode("utf-8") + @staticmethod + def _require_int(payload: dict[str, Any], field_name: str, *, min_value: int, max_value: int) -> None: + value = payload.get(field_name) + if isinstance(value, bool) or not isinstance(value, int): + raise HandshakeHwValidationError(f"HW field {field_name!r} must be an integer") + if value < min_value or value > max_value: + raise HandshakeHwValidationError( + f"HW field {field_name!r} must be between {min_value} and {max_value}" + ) + + @staticmethod + def _require_non_negative_float(payload: dict[str, Any], field_name: str) -> None: + value = payload.get(field_name) + if isinstance(value, bool) or not isinstance(value, (int, float)): + raise HandshakeHwValidationError(f"HW field {field_name!r} must be a number") + numeric = float(value) + if not math.isfinite(numeric) or numeric < 0.0: + raise HandshakeHwValidationError(f"HW field {field_name!r} must be finite and non-negative") + payload[field_name] = numeric + + @staticmethod + def _require_bool(payload: dict[str, Any], field_name: str) -> None: + if not isinstance(payload.get(field_name), bool): + raise HandshakeHwValidationError(f"HW field {field_name!r} must be a boolean") + + @classmethod + def _validate_wire_payload(cls, payload: dict[str, Any]) -> dict[str, Any]: + known = {f for f in cls.__dataclass_fields__} + filtered = {k: v for k, v in payload.items() if k in known} + + for field_name, min_value, max_value in ( + ("wire_version", 0, 255), + ("gpu_cores", 0, 1024), + ("data_port", 0, 65535), + ): + if field_name in filtered: + cls._require_int(filtered, field_name, min_value=min_value, max_value=max_value) + + for field_name in ("ram_gb", "memory_bandwidth_gbps"): + if field_name in filtered: + cls._require_non_negative_float(filtered, field_name) + + if "chip_name" in filtered and not isinstance(filtered["chip_name"], str): + raise HandshakeHwValidationError("HW field 'chip_name' must be a string") + + for field_name in ("has_ane", "mps_available", "mlx_available"): + if field_name in filtered: + cls._require_bool(filtered, field_name) + + return filtered + @classmethod def from_json_bytes(cls, data: bytes) -> "HardwareExchange": """Parse JSON bytes back into a HardwareExchange. @@ -127,8 +179,7 @@ def from_json_bytes(cls, data: bytes) -> "HardwareExchange": if not isinstance(payload, dict): raise HandshakeHwValidationError("HW payload not a JSON object") # Only accept known fields; ignore extras for forward compat - known = {f for f in cls.__dataclass_fields__} - filtered = {k: v for k, v in payload.items() if k in known} + filtered = cls._validate_wire_payload(payload) return cls(**filtered) except (ValueError, TypeError, UnicodeDecodeError) as e: raise HandshakeHwValidationError(f"HW payload deserialization failed: {e}") from e diff --git a/tests/test_security/test_hw_exchange.py b/tests/test_security/test_hw_exchange.py index ed65da2..7d47845 100644 --- a/tests/test_security/test_hw_exchange.py +++ b/tests/test_security/test_hw_exchange.py @@ -66,6 +66,41 @@ def test_from_json_bytes_rejects_invalid_utf8(self): with pytest.raises(HandshakeHwValidationError): HardwareExchange.from_json_bytes(b'\xff\xfe\xfd') + @pytest.mark.parametrize( + "payload", + [ + b'{"wire_version": "1"}', + b'{"gpu_cores": "8"}', + b'{"gpu_cores": true}', + b'{"ram_gb": "16"}', + b'{"memory_bandwidth_gbps": []}', + b'{"chip_name": 123}', + b'{"has_ane": "yes"}', + b'{"mps_available": 1}', + b'{"mlx_available": null}', + b'{"data_port": "50052"}', + ], + ) + def test_from_json_bytes_rejects_wrong_field_types(self, payload): + with pytest.raises(HandshakeHwValidationError): + HardwareExchange.from_json_bytes(payload) + + @pytest.mark.parametrize( + "payload", + [ + b'{"wire_version": -1}', + b'{"wire_version": 256}', + b'{"gpu_cores": -1}', + b'{"ram_gb": -0.01}', + b'{"memory_bandwidth_gbps": -0.01}', + b'{"data_port": -1}', + b'{"data_port": 65536}', + ], + ) + def test_from_json_bytes_rejects_out_of_range_fields(self, payload): + with pytest.raises(HandshakeHwValidationError): + HardwareExchange.from_json_bytes(payload) + class TestSignVerifyHwProfile: def _key(self) -> bytes: From 805dfd81bcb417681c358283b8197d98c7420a41 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Wed, 15 Jul 2026 12:30:07 -0400 Subject: [PATCH 2/4] Bound network topology advertisements --- macfleet/pool/topology.py | 25 +++++++++++++++++++++---- tests/test_pool/test_topology.py | 21 +++++++++++++++++++++ 2 files changed, 42 insertions(+), 4 deletions(-) diff --git a/macfleet/pool/topology.py b/macfleet/pool/topology.py index 9952ecc..6efed50 100644 --- a/macfleet/pool/topology.py +++ b/macfleet/pool/topology.py @@ -193,6 +193,20 @@ def serialize_network_links(links: Sequence[NetworkLink]) -> str: return json.dumps(payload, separators=(",", ":")) +def _optional_float(item: dict, short_key: str, long_key: str) -> float: + value = item.get(short_key) or item.get(long_key) or 0.0 + if isinstance(value, bool): + raise ValueError(f"{short_key} must be numeric") + return float(value) + + +def _optional_int(item: dict, short_key: str, long_key: str, default: int) -> int: + value = item.get(short_key) or item.get(long_key) or default + if isinstance(value, bool): + raise ValueError(f"{short_key} must be an integer") + return int(value) + + def deserialize_network_links(payload: str) -> tuple[NetworkLink, ...]: """Parse serialized links, ignoring malformed entries.""" try: @@ -218,16 +232,17 @@ def deserialize_network_links(payload: str) -> tuple[NetworkLink, ...]: or link_type == LinkType.LOOPBACK ): continue - bandwidth_mbps = float(item.get("b") or item.get("bandwidth_mbps") or 0.0) - latency_ms = float(item.get("l") or item.get("latency_ms") or 0.0) - loss_rate = float(item.get("r") or item.get("loss_rate") or 0.0) - mtu = int(item.get("m") or item.get("mtu") or 1500) + bandwidth_mbps = _optional_float(item, "b", "bandwidth_mbps") + latency_ms = _optional_float(item, "l", "latency_ms") + loss_rate = _optional_float(item, "r", "loss_rate") + mtu = _optional_int(item, "m", "mtu", 1500) if ( not all(math.isfinite(value) and value >= 0 for value in ( bandwidth_mbps, latency_ms, loss_rate, )) + or loss_rate > 1 or mtu < 1 ): continue @@ -242,6 +257,8 @@ def deserialize_network_links(payload: str) -> tuple[NetworkLink, ...]: mtu=mtu, ) ) + if len(links) >= _MAX_SERIALIZED_LINKS: + break except (KeyError, TypeError, ValueError): continue return tuple(links) diff --git a/tests/test_pool/test_topology.py b/tests/test_pool/test_topology.py index 8b19bbf..e4ace3b 100644 --- a/tests/test_pool/test_topology.py +++ b/tests/test_pool/test_topology.py @@ -153,3 +153,24 @@ def test_deserialize_network_links_rejects_non_finite_telemetry(): payload = '[{"i":"en0","t":"ethernet","ip":"192.168.1.2","b":NaN}]' assert deserialize_network_links(payload) == () + + +def test_deserialize_network_links_rejects_boolean_telemetry(): + payload = '[{"i":"en0","t":"ethernet","ip":"192.168.1.2","b":true}]' + + assert deserialize_network_links(payload) == () + + +def test_deserialize_network_links_rejects_invalid_loss_rate(): + payload = '[{"i":"en0","t":"wifi","ip":"192.168.1.2","r":1.5}]' + + assert deserialize_network_links(payload) == () + + +def test_deserialize_network_links_caps_link_count(): + payload = json.dumps([ + {"i": f"en{i}", "t": "wifi", "ip": f"192.168.1.{i + 2}"} + for i in range(12) + ]) + + assert len(deserialize_network_links(payload)) == 8 From b7ddf26129abb48f90f888a2031e9bac30c9d8c5 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Wed, 15 Jul 2026 12:31:14 -0400 Subject: [PATCH 3/4] Validate discovered node metadata --- macfleet/pool/discovery.py | 16 +++++++ tests/test_security/test_fleet_discovery.py | 48 ++++++++++++++++++++- 2 files changed, 63 insertions(+), 1 deletion(-) diff --git a/macfleet/pool/discovery.py b/macfleet/pool/discovery.py index ef63eaa..ca6e609 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 @@ -48,10 +49,25 @@ class DiscoveredNode: network_links: tuple[NetworkLink, ...] = () def __post_init__(self) -> None: + if isinstance(self.port, bool) or not isinstance(self.port, int) or not (0 < self.port < 65536): + raise ValueError("port must be between 1 and 65535") + if ( + isinstance(self.compute_score, bool) + or not isinstance(self.compute_score, (int, float)) + or not math.isfinite(float(self.compute_score)) + or self.compute_score < 0 + ): + raise ValueError("compute_score must be finite and non-negative") # Backward compat: 2.1.x peers don't advertise data_port. Fall back to # heartbeat_port + 1 — matches the 50051/50052 convention. if self.data_port == 0: self.data_port = self.port + 1 + if ( + isinstance(self.data_port, bool) + or not isinstance(self.data_port, int) + or not (0 < self.data_port < 65536) + ): + raise ValueError("data_port must be between 1 and 65535") @property def link_type_list(self) -> list[str]: diff --git a/tests/test_security/test_fleet_discovery.py b/tests/test_security/test_fleet_discovery.py index 08ef4f8..7418c41 100644 --- a/tests/test_security/test_fleet_discovery.py +++ b/tests/test_security/test_fleet_discovery.py @@ -2,9 +2,11 @@ from __future__ import annotations +import math + import pytest -from macfleet.pool.discovery import MACFLEET_SERVICE_TYPE, ServiceRegistry +from macfleet.pool.discovery import MACFLEET_SERVICE_TYPE, DiscoveredNode, ServiceRegistry from macfleet.security.auth import DEFAULT_SERVICE_TYPE, SecurityConfig @@ -201,3 +203,47 @@ def test_remove_service_strips_service_type_suffix(self): None, service_type, f"node-xyz.{service_type}", ) assert seen == ["node-xyz"] + + +class TestDiscoveredNodeValidation: + @pytest.mark.parametrize( + ("port", "data_port"), + [ + (0, 50052), + (65536, 50052), + (50051, -1), + (50051, 65536), + ], + ) + def test_rejects_invalid_ports(self, port, data_port): + with pytest.raises(ValueError, match="port"): + DiscoveredNode( + hostname="mac.local", + node_id="node-a", + ip_address="127.0.0.1", + port=port, + gpu_cores=0, + ram_gb=0, + chip_name="unknown", + link_types="", + pool_version="2.2.1", + compute_score=0.0, + data_port=data_port, + ) + + @pytest.mark.parametrize("compute_score", [math.nan, math.inf, -0.1]) + def test_rejects_invalid_compute_score(self, compute_score): + with pytest.raises(ValueError, match="compute_score"): + DiscoveredNode( + hostname="mac.local", + node_id="node-a", + ip_address="127.0.0.1", + port=50051, + gpu_cores=0, + ram_gb=0, + chip_name="unknown", + link_types="", + pool_version="2.2.1", + compute_score=compute_score, + data_port=50052, + ) From 9dedd5625ea50b8d375a5f4b0b2c289fb284aa41 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Wed, 15 Jul 2026 12:33:43 -0400 Subject: [PATCH 4/4] Preserve zero topology metrics --- macfleet/pool/topology.py | 14 ++++++++++++-- tests/test_pool/test_topology.py | 13 +++++++++++++ 2 files changed, 25 insertions(+), 2 deletions(-) diff --git a/macfleet/pool/topology.py b/macfleet/pool/topology.py index 6efed50..a82d1ef 100644 --- a/macfleet/pool/topology.py +++ b/macfleet/pool/topology.py @@ -194,14 +194,24 @@ def serialize_network_links(links: Sequence[NetworkLink]) -> str: def _optional_float(item: dict, short_key: str, long_key: str) -> float: - value = item.get(short_key) or item.get(long_key) or 0.0 + if short_key in item: + value = item[short_key] + elif long_key in item: + value = item[long_key] + else: + value = 0.0 if isinstance(value, bool): raise ValueError(f"{short_key} must be numeric") return float(value) def _optional_int(item: dict, short_key: str, long_key: str, default: int) -> int: - value = item.get(short_key) or item.get(long_key) or default + if short_key in item: + value = item[short_key] + elif long_key in item: + value = item[long_key] + else: + value = default if isinstance(value, bool): raise ValueError(f"{short_key} must be an integer") return int(value) diff --git a/tests/test_pool/test_topology.py b/tests/test_pool/test_topology.py index e4ace3b..539125f 100644 --- a/tests/test_pool/test_topology.py +++ b/tests/test_pool/test_topology.py @@ -167,6 +167,19 @@ def test_deserialize_network_links_rejects_invalid_loss_rate(): assert deserialize_network_links(payload) == () +def test_deserialize_network_links_preserves_explicit_zero_metrics(): + payload = ( + '[{"i":"en0","t":"ethernet","ip":"192.168.1.2",' + '"b":0,"bandwidth_mbps":true,"l":0,"latency_ms":true}]' + ) + + restored = deserialize_network_links(payload) + + assert len(restored) == 1 + assert restored[0].bandwidth_mbps == 0.0 + assert restored[0].latency_ms == 0.0 + + def test_deserialize_network_links_caps_link_count(): payload = json.dumps([ {"i": f"en{i}", "t": "wifi", "ip": f"192.168.1.{i + 2}"}