From 833668ac62bc172fc13dac1e235e02d01dea91c5 Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Sun, 2 Aug 2026 13:06:14 -0400 Subject: [PATCH 1/2] Validate network probe inputs --- macfleet/pool/network.py | 17 +++++++++++++++++ tests/test_pool/test_network.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/macfleet/pool/network.py b/macfleet/pool/network.py index 3f18855..94320d4 100644 --- a/macfleet/pool/network.py +++ b/macfleet/pool/network.py @@ -7,6 +7,7 @@ from __future__ import annotations import asyncio +import math import socket import subprocess import time @@ -15,6 +16,17 @@ from typing import Optional +def _validate_positive_finite(name: str, value: float) -> float: + if ( + isinstance(value, bool) + or not isinstance(value, (int, float)) + or not math.isfinite(float(value)) + or value <= 0 + ): + raise ValueError(f"{name} must be a positive finite number") + return float(value) + + class LinkType(Enum): """Network link types, ordered by typical bandwidth.""" WIFI = "wifi" @@ -249,6 +261,7 @@ def _parse_ifconfig() -> list[NetworkLink]: async def measure_latency(host: str, port: int, timeout: float = 5.0) -> float: """Measure TCP round-trip latency to a peer in milliseconds.""" + timeout = _validate_positive_finite("timeout", timeout) try: start = time.monotonic() _, writer = await asyncio.wait_for( @@ -278,7 +291,11 @@ async def measure_bandwidth( confirms the kernel send buffer received the bytes, not the wire. Requires a listening peer that accepts and discards data. """ + payload_mb = _validate_positive_finite("payload_mb", payload_mb) + timeout = _validate_positive_finite("timeout", timeout) payload_bytes = int(payload_mb * 1024 * 1024) + if payload_bytes == 0: + raise ValueError("payload_mb must represent at least one byte") chunk_size = 65536 try: diff --git a/tests/test_pool/test_network.py b/tests/test_pool/test_network.py index c4654cb..7219390 100644 --- a/tests/test_pool/test_network.py +++ b/tests/test_pool/test_network.py @@ -1,5 +1,7 @@ """Tests for network detection and link scoring.""" +import pytest + from macfleet.pool import network from macfleet.pool.network import ( LinkType, @@ -96,3 +98,29 @@ def test_get_topology(self): topo = get_network_topology() assert isinstance(topo, NetworkTopology) assert topo.hostname != "" + + +class TestNetworkProbes: + @pytest.mark.parametrize( + "timeout", + [0, -1, True, float("nan"), float("inf"), "soon"], + ) + async def test_latency_rejects_invalid_timeout(self, timeout): + with pytest.raises(ValueError, match="timeout"): + await network.measure_latency("127.0.0.1", 9, timeout=timeout) + + @pytest.mark.parametrize( + "timeout", + [0, -1, True, float("nan"), float("inf"), "soon"], + ) + async def test_bandwidth_rejects_invalid_timeout(self, timeout): + with pytest.raises(ValueError, match="timeout"): + await network.measure_bandwidth("127.0.0.1", 9, timeout=timeout) + + @pytest.mark.parametrize( + "payload_mb", + [0, -1, True, float("nan"), float("inf"), "large"], + ) + async def test_bandwidth_rejects_invalid_payload_size(self, payload_mb): + with pytest.raises(ValueError, match="payload_mb"): + await network.measure_bandwidth("127.0.0.1", 9, payload_mb=payload_mb) From d58ff24ebeaaee32e2cb1af32a6b05450d52d34c Mon Sep 17 00:00:00 2001 From: Vikranth Reddimasu Date: Sun, 2 Aug 2026 13:08:23 -0400 Subject: [PATCH 2/2] Bound network probe lifecycles --- macfleet/pool/network.py | 97 +++++++++++----------- tests/test_pool/test_network.py | 143 ++++++++++++++++++++++++++++++++ 2 files changed, 192 insertions(+), 48 deletions(-) diff --git a/macfleet/pool/network.py b/macfleet/pool/network.py index 94320d4..f92d9c0 100644 --- a/macfleet/pool/network.py +++ b/macfleet/pool/network.py @@ -260,20 +260,26 @@ def _parse_ifconfig() -> list[NetworkLink]: async def measure_latency(host: str, port: int, timeout: float = 5.0) -> float: - """Measure TCP round-trip latency to a peer in milliseconds.""" + """Measure TCP connection latency to a peer in milliseconds. + + ``timeout`` bounds the complete probe, including connection cleanup. + Returns -1 when the peer cannot complete the probe in time. + """ timeout = _validate_positive_finite("timeout", timeout) + writer = None try: - start = time.monotonic() - _, writer = await asyncio.wait_for( - asyncio.open_connection(host, port), - timeout=timeout, - ) - elapsed = (time.monotonic() - start) * 1000.0 - writer.close() - await writer.wait_closed() - return elapsed + async with asyncio.timeout(timeout): + start = time.monotonic() + _, writer = await asyncio.open_connection(host, port) + elapsed = (time.monotonic() - start) * 1000.0 + writer.close() + await writer.wait_closed() + return elapsed except (OSError, asyncio.TimeoutError): return -1.0 + finally: + if writer is not None: + writer.close() async def measure_bandwidth( @@ -284,12 +290,11 @@ async def measure_bandwidth( ) -> float: """Measure TCP bandwidth to a peer in Mbps. - Sends a payload, waits for the receiver to close the connection - (signaling it has drained the bytes), and divides total elapsed - time by payload size. The wait_closed() bound is the closest the - stdlib affords to a real ACK — without it, writer.drain() only - confirms the kernel send buffer received the bytes, not the wire. - Requires a listening peer that accepts and discards data. + Sends a payload, then waits for the receiver to close the connection + after draining the bytes. ``timeout`` bounds the complete probe: connect, + transfer, receiver acknowledgement, and connection cleanup. Returns 0 + when any phase fails or the receiver does not acknowledge completion. + Requires a listening peer that reads until EOF and then closes. """ payload_mb = _validate_positive_finite("payload_mb", payload_mb) timeout = _validate_positive_finite("timeout", timeout) @@ -297,43 +302,39 @@ async def measure_bandwidth( if payload_bytes == 0: raise ValueError("payload_mb must represent at least one byte") chunk_size = 65536 + writer = None try: - reader, writer = await asyncio.wait_for( - asyncio.open_connection(host, port), - timeout=timeout, - ) - - data = b"\x00" * chunk_size - start = time.monotonic() - sent = 0 - while sent < payload_bytes: - to_send = min(chunk_size, payload_bytes - sent) - writer.write(data[:to_send]) - await writer.drain() - sent += to_send - - # Half-close the writer side so the peer sees EOF, then wait - # for the connection to fully close — this is our acknowledgement - # signal that everything we sent has actually been read on the - # other end (or the receiver hung up). Either way, the elapsed - # time more closely reflects wire bandwidth than the local - # send-buffer drain time alone. - if writer.can_write_eof(): + async with asyncio.timeout(timeout): + reader, writer = await asyncio.open_connection(host, port) + + data = b"\x00" * chunk_size + start = time.monotonic() + sent = 0 + while sent < payload_bytes: + to_send = min(chunk_size, payload_bytes - sent) + writer.write(data[:to_send]) + await writer.drain() + sent += to_send + + # Half-close so the peer sees EOF. A full close from the peer is + # the acknowledgement that its handler finished draining data. + if not writer.can_write_eof(): + return 0.0 writer.write_eof() - try: - await asyncio.wait_for(reader.read(), timeout=timeout) - except asyncio.TimeoutError: - pass - elapsed = time.monotonic() - start - writer.close() - await writer.wait_closed() - - if elapsed > 0: - return (payload_bytes * 8) / (elapsed * 1_000_000) # Mbps - return 0.0 + await reader.read() + elapsed = time.monotonic() - start + writer.close() + await writer.wait_closed() + + if elapsed > 0: + return (payload_bytes * 8) / (elapsed * 1_000_000) # Mbps + return 0.0 except (OSError, asyncio.TimeoutError): return 0.0 + finally: + if writer is not None: + writer.close() def get_network_topology() -> NetworkTopology: diff --git a/tests/test_pool/test_network.py b/tests/test_pool/test_network.py index 7219390..ec0f917 100644 --- a/tests/test_pool/test_network.py +++ b/tests/test_pool/test_network.py @@ -1,5 +1,8 @@ """Tests for network detection and link scoring.""" +import asyncio +import time + import pytest from macfleet.pool import network @@ -124,3 +127,143 @@ async def test_bandwidth_rejects_invalid_timeout(self, timeout): async def test_bandwidth_rejects_invalid_payload_size(self, payload_mb): with pytest.raises(ValueError, match="payload_mb"): await network.measure_bandwidth("127.0.0.1", 9, payload_mb=payload_mb) + + async def test_latency_timeout_bounds_connection_close(self, monkeypatch): + writer = _ProbeWriter(wait_closed_delay=None) + + async def open_connection(host, port): + return _ProbeReader(), writer + + monkeypatch.setattr(network.asyncio, "open_connection", open_connection) + + started = time.monotonic() + result = await asyncio.wait_for( + network.measure_latency("peer", 5000, timeout=0.01), + timeout=0.2, + ) + + assert result == -1.0 + assert time.monotonic() - started < 0.15 + assert writer.closed + + async def test_bandwidth_requires_receiver_acknowledgement(self, monkeypatch): + writer = _ProbeWriter() + + async def open_connection(host, port): + return _ProbeReader(read_delay=None), writer + + monkeypatch.setattr(network.asyncio, "open_connection", open_connection) + + result = await network.measure_bandwidth( + "peer", + 5000, + payload_mb=0.001, + timeout=0.01, + ) + + assert result == 0.0 + assert writer.closed + + async def test_bandwidth_timeout_is_end_to_end(self, monkeypatch): + writer = _ProbeWriter() + + async def open_connection(host, port): + await asyncio.sleep(0.03) + return _ProbeReader(read_delay=0.03), writer + + monkeypatch.setattr(network.asyncio, "open_connection", open_connection) + + started = time.monotonic() + result = await network.measure_bandwidth( + "peer", + 5000, + payload_mb=0.001, + timeout=0.05, + ) + + assert result == 0.0 + assert time.monotonic() - started < 0.09 + assert writer.closed + + async def test_bandwidth_timeout_bounds_stalled_drain(self, monkeypatch): + writer = _ProbeWriter(drain_delay=None) + + async def open_connection(host, port): + return _ProbeReader(), writer + + monkeypatch.setattr(network.asyncio, "open_connection", open_connection) + + result = await asyncio.wait_for( + network.measure_bandwidth( + "peer", + 5000, + payload_mb=0.001, + timeout=0.01, + ), + timeout=0.2, + ) + + assert result == 0.0 + assert writer.closed + + async def test_bandwidth_reports_measurement_after_acknowledgement(self, monkeypatch): + writer = _ProbeWriter() + + async def open_connection(host, port): + return _ProbeReader(), writer + + monkeypatch.setattr(network.asyncio, "open_connection", open_connection) + + result = await network.measure_bandwidth( + "peer", + 5000, + payload_mb=0.001, + timeout=0.1, + ) + + assert result > 0.0 + assert writer.bytes_written == int(0.001 * 1024 * 1024) + assert writer.eof_written + assert writer.closed + + +class _ProbeReader: + def __init__(self, read_delay=0.0): + self.read_delay = read_delay + + async def read(self): + if self.read_delay is None: + await asyncio.Future() + await asyncio.sleep(self.read_delay) + return b"" + + +class _ProbeWriter: + def __init__(self, drain_delay=0.0, wait_closed_delay=0.0): + self.drain_delay = drain_delay + self.wait_closed_delay = wait_closed_delay + self.bytes_written = 0 + self.eof_written = False + self.closed = False + + def write(self, data): + self.bytes_written += len(data) + + async def drain(self): + if self.drain_delay is None: + await asyncio.Future() + await asyncio.sleep(self.drain_delay) + + def can_write_eof(self): + return True + + def write_eof(self): + self.eof_written = True + + def close(self): + self.closed = True + + async def wait_closed(self): + if self.wait_closed_delay is None: + await asyncio.Future() + await asyncio.sleep(self.wait_closed_delay)