Sanitize peer-supplied mDNS labels in discover CLI (#1689)

Author: bdraco

aioesphomeapi/_frame_helper/base.pxd

@@ -1,6 +1,7 @@
 
 import cython
 
+from .._sanitize cimport safe_label_str
 from ..connection cimport APIConnection
 
 
@@ -10,8 +11,6 @@ cdef int _MAX_NAME_LEN
 cdef int _MAX_MAC_LEN
 cdef int _MAX_EXPLANATION_LEN
 
-cpdef str safe_label_str(str raw, int limit)
-
 cdef class APIFrameHelper:
 
     cdef object _loop

aioesphomeapi/_frame_helper/base.py

@@ -6,6 +6,7 @@
 import logging
 from typing import TYPE_CHECKING, cast
 
+from .._sanitize import MAX_EXPLANATION_LEN, MAX_MAC_LEN, MAX_NAME_LEN, safe_label_str
 from ..core import SocketClosedAPIError
 
 if TYPE_CHECKING:
@@ -25,25 +26,20 @@
 _bytes = bytes
 
 
-# Caps match the firmware's actual wire-format limits:
-#   - name: ESPHOME_DEVICE_NAME_MAX_LEN = 31 (validate_hostname in core/config.py)
-#   - mac: MAC_ADDRESS_BUFFER_SIZE - 1 = 12 (lowercase hex, no separator)
-#   - explanation: 32-byte handshake-reject buffer minus the 1-byte failure code
-# A small extra margin on each lets benign forward-compat tweaks (e.g. firmware
-# bumping the max name length by a few chars) through without breaking clients.
-# MAX_* are the Python-importable forms (used by tests + connection.py); _MAX_*
-# are the cdef int aliases declared in base.pxd for hot-path C comparisons.
-MAX_NAME_LEN = 32
-MAX_MAC_LEN = 16
-MAX_EXPLANATION_LEN = 64
+# _MAX_* are the cdef int aliases declared in base.pxd for hot-path C
+# comparisons; the Python-importable MAX_* names live in aioesphomeapi._sanitize
+# and are re-exported here so existing `from .base import MAX_NAME_LEN` callers
+# keep working unchanged.
 _MAX_NAME_LEN = MAX_NAME_LEN
 _MAX_MAC_LEN = MAX_MAC_LEN
 _MAX_EXPLANATION_LEN = MAX_EXPLANATION_LEN
 
-
-def safe_label_str(raw: str, limit: _int) -> str:
-    """Strip non-printables and length-cap a peer-supplied label for log output."""
-    return "".join(filter(str.isprintable, raw))[:limit]
+__all__ = (
+    "MAX_EXPLANATION_LEN",
+    "MAX_MAC_LEN",
+    "MAX_NAME_LEN",
+    "safe_label_str",
+)
 
 
 class APIFrameHelper:

aioesphomeapi/_frame_helper/noise.pxd

@@ -1,12 +1,12 @@
 import cython
 
+from .._sanitize cimport safe_label_str
 from ..connection cimport APIConnection
 from .base cimport (
     APIFrameHelper,
     _MAX_EXPLANATION_LEN,
     _MAX_MAC_LEN,
     _MAX_NAME_LEN,
-    safe_label_str,
 )
 from .noise_encryption cimport EncryptCipher, DecryptCipher
 from .packets cimport make_noise_packets

aioesphomeapi/_frame_helper/noise.py

@@ -8,6 +8,7 @@
 from cryptography.exceptions import InvalidTag
 from noise.connection import NoiseConnection
 
+from .._sanitize import MAX_EXPLANATION_LEN, MAX_MAC_LEN, MAX_NAME_LEN, safe_label_str
 from ..core import (
     APIConnectionError,
     BadMACAddressAPIError,
@@ -19,14 +20,7 @@
     InvalidEncryptionKeyAPIError,
     ProtocolAPIError,
 )
-from .base import (
-    _LOGGER,
-    MAX_EXPLANATION_LEN,
-    MAX_MAC_LEN,
-    MAX_NAME_LEN,
-    APIFrameHelper,
-    safe_label_str,
-)
+from .base import _LOGGER, APIFrameHelper
 from .noise_encryption import ESPHOME_NOISE_BACKEND, DecryptCipher, EncryptCipher
 from .packets import make_noise_packets
 
@@ -51,9 +45,9 @@
 
 int_ = int
 
-# Cython resolves _MAX_* and safe_label_str via cimport from .base
-# (noise.pxd); these assignments are the pure-Python (SKIP_CYTHON=1) fallback
-# so callers below have a name to resolve.
+# Cython resolves _MAX_* via cimport from .base and safe_label_str via cimport
+# from .._sanitize (noise.pxd); these assignments are the pure-Python
+# (SKIP_CYTHON=1) fallback so callers below have a name to resolve.
 _MAX_NAME_LEN = MAX_NAME_LEN
 _MAX_MAC_LEN = MAX_MAC_LEN
 _MAX_EXPLANATION_LEN = MAX_EXPLANATION_LEN

aioesphomeapi/_sanitize.pxd

@@ -0,0 +1,2 @@
+
+cpdef str safe_label_str(str raw, int limit)

aioesphomeapi/_sanitize.py

@@ -0,0 +1,38 @@
+"""Sanitize peer-supplied labels before logging or printing them.
+
+Shared by the noise hello / handshake-reject paths and the discover CLI;
+strips non-printable characters and length-caps the result so a hostile
+remote can't inject ANSI escapes, newlines or oversized values into the
+operator's terminal or logs.
+"""
+
+from __future__ import annotations
+
+# Caps match the firmware's actual wire-format limits:
+#   - name: ESPHOME_DEVICE_NAME_MAX_LEN = 31 (validate_hostname in core/config.py)
+#   - mac: MAC_ADDRESS_BUFFER_SIZE - 1 = 12 (lowercase hex, no separator)
+#   - explanation: 32-byte handshake-reject buffer minus the 1-byte failure code
+# A small extra margin on each lets benign forward-compat tweaks (e.g. firmware
+# bumping the max name length by a few chars) through without breaking clients.
+MAX_NAME_LEN = 32
+MAX_MAC_LEN = 16
+MAX_EXPLANATION_LEN = 64
+
+__all__ = (
+    "MAX_EXPLANATION_LEN",
+    "MAX_MAC_LEN",
+    "MAX_NAME_LEN",
+    "safe_label_str",
+)
+
+
+# Alias so the `limit` annotation below isn't interpreted as a C-int type
+# declaration by Cython — the .pxd already declares `int limit` and a bare
+# `int` annotation would clash with it ("Function signature does not match
+# previous declaration"). Mirrors the same workaround in _frame_helper/base.py.
+_int = int
+
+
+def safe_label_str(raw: str, limit: _int) -> str:
+    """Strip non-printables and length-cap a peer-supplied label for log output."""
+    return "".join(filter(str.isprintable, raw))[:limit]

aioesphomeapi/discover.py

@@ -6,23 +6,44 @@
 import asyncio
 import contextlib
 import logging
+import re
 import sys
 
 from zeroconf import IPVersion, ServiceStateChange, Zeroconf
 from zeroconf.asyncio import AsyncServiceBrowser, AsyncServiceInfo, AsyncZeroconf
 
+from ._sanitize import safe_label_str
+
 FORMAT = "{: <7}|{: <32}|{: <15}|{: <12}|{: <16}|{: <10}|{: <32}"
 COLUMN_NAMES = ("Status", "Name", "Address", "MAC", "Version", "Platform", "Board")
 UNKNOWN = "unknown"
 
-
-def decode_bytes_or_unknown(data: str | bytes | None) -> str:
-    """Decode bytes or return unknown."""
+# Per-column display caps for peer-supplied mDNS labels, derived from the
+# FORMAT widths so a hostile broadcaster can't widen a column by stuffing a
+# long value; deriving them from FORMAT keeps the caps in lock-step if the
+# table layout is ever retuned.
+_COLUMN_WIDTHS = tuple(int(w) for w in re.findall(r"<\s*(\d+)", FORMAT))
+assert len(_COLUMN_WIDTHS) == len(COLUMN_NAMES), (
+    "FORMAT width count must match COLUMN_NAMES; update one and the other together"
+)
+_MAX_NAME_DISPLAY = _COLUMN_WIDTHS[COLUMN_NAMES.index("Name")]
+_MAX_MAC_DISPLAY = _COLUMN_WIDTHS[COLUMN_NAMES.index("MAC")]
+_MAX_VERSION_DISPLAY = _COLUMN_WIDTHS[COLUMN_NAMES.index("Version")]
+_MAX_PLATFORM_DISPLAY = _COLUMN_WIDTHS[COLUMN_NAMES.index("Platform")]
+_MAX_BOARD_DISPLAY = _COLUMN_WIDTHS[COLUMN_NAMES.index("Board")]
+
+
+def decode_mdns_label_or_unknown(
+    data: str | bytes | None, limit: int = _MAX_NAME_DISPLAY
+) -> str:
+    """Decode peer-supplied mDNS bytes, strip non-printables, length-cap."""
     if data is None:
         return UNKNOWN
     if isinstance(data, bytes):
-        return data.decode()
-    return data
+        # A device on the LAN can broadcast arbitrary bytes; use "replace" so
+        # a malformed UTF-8 payload doesn't raise out of the zeroconf callback.
+        data = data.decode("utf-8", "replace")
+    return safe_label_str(data, limit)
 
 
 def async_service_update(
@@ -32,15 +53,22 @@ def async_service_update(
     state_change: ServiceStateChange,
 ) -> None:
     """Service state changed."""
-    short_name = name.partition(".")[0]
+    # The mDNS service name is peer-controlled — sanitize before printing so
+    # a hostile broadcaster can't inject ANSI escapes / newlines / null bytes
+    # into the terminal.
+    short_name = safe_label_str(name.partition(".")[0], _MAX_NAME_DISPLAY)
     state = "OFFLINE" if state_change is ServiceStateChange.Removed else "ONLINE"
     info = AsyncServiceInfo(service_type, name)
     info.load_from_cache(zeroconf)
     properties = info.properties
-    mac = decode_bytes_or_unknown(properties.get(b"mac"))
-    version = decode_bytes_or_unknown(properties.get(b"version"))
-    platform = decode_bytes_or_unknown(properties.get(b"platform"))
-    board = decode_bytes_or_unknown(properties.get(b"board"))
+    mac = decode_mdns_label_or_unknown(properties.get(b"mac"), _MAX_MAC_DISPLAY)
+    version = decode_mdns_label_or_unknown(
+        properties.get(b"version"), _MAX_VERSION_DISPLAY
+    )
+    platform = decode_mdns_label_or_unknown(
+        properties.get(b"platform"), _MAX_PLATFORM_DISPLAY
+    )
+    board = decode_mdns_label_or_unknown(properties.get(b"board"), _MAX_BOARD_DISPLAY)
     address = ""
     if addresses := info.ip_addresses_by_version(IPVersion.V4Only):
         address = str(addresses[0])

setup.py

@@ -15,6 +15,7 @@
     from distutils.core import Extension
 
 TO_CYTHONIZE = [
+    "aioesphomeapi/_sanitize.py",
     "aioesphomeapi/client_base.py",
     "aioesphomeapi/connection.py",
     "aioesphomeapi/_frame_helper/base.py",

tests/test_discover.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+import re
+
+from aioesphomeapi.discover import (
+    _MAX_BOARD_DISPLAY,
+    _MAX_MAC_DISPLAY,
+    _MAX_NAME_DISPLAY,
+    _MAX_PLATFORM_DISPLAY,
+    _MAX_VERSION_DISPLAY,
+    COLUMN_NAMES,
+    FORMAT,
+    UNKNOWN,
+    decode_mdns_label_or_unknown,
+)
+
+
+def test_decode_mdns_label_or_unknown_none() -> None:
+    assert decode_mdns_label_or_unknown(None) == UNKNOWN
+
+
+def test_decode_mdns_label_or_unknown_str_passthrough() -> None:
+    assert decode_mdns_label_or_unknown("esp32-board") == "esp32-board"
+
+
+def test_decode_mdns_label_or_unknown_bytes_utf8() -> None:
+    assert decode_mdns_label_or_unknown(b"esp32-board") == "esp32-board"
+
+
+def test_decode_mdns_label_or_unknown_invalid_utf8_replaces() -> None:
+    # Hostile mDNS broadcaster sends non-UTF-8 bytes; result is the U+FFFD
+    # replacement character (one per invalid byte), never raises
+    # UnicodeDecodeError. Pinning the actual output, not just the type, keeps
+    # a future refactor from silently switching to UNKNOWN or empty string.
+    assert decode_mdns_label_or_unknown(b"\xff\xfe") == "\ufffd\ufffd"
+
+
+def test_decode_mdns_label_or_unknown_strips_control_chars() -> None:
+    # Strip the ESC byte that activates ANSI sequences, plus newline / CR /
+    # null / tab / etc. The trailing "[2J" is harmless printable text once
+    # the leading ESC is gone, so a hostile broadcaster can no longer clear
+    # the user's terminal from a discovery scan.
+    assert decode_mdns_label_or_unknown(b"\x1b[2Jvers\n1.0") == "[2Jvers1.0"
+    assert decode_mdns_label_or_unknown(b"line1\r\nline2") == "line1line2"
+    assert decode_mdns_label_or_unknown(b"col\tumn") == "column"
+
+
+def test_decode_mdns_label_or_unknown_strips_null_byte() -> None:
+    assert decode_mdns_label_or_unknown(b"esp\x0032") == "esp32"
+
+
+def test_decode_mdns_label_or_unknown_caps_length() -> None:
+    assert decode_mdns_label_or_unknown(b"x" * 200, limit=10) == "x" * 10
+
+
+def test_decode_mdns_label_or_unknown_default_limit_caps_long_str() -> None:
+    # Default cap is the Name column width from FORMAT.
+    assert len(decode_mdns_label_or_unknown("a" * 100)) == _MAX_NAME_DISPLAY
+
+
+def test_decode_mdns_label_or_unknown_unicode_printable_survives() -> None:
+    # safe_label_str uses str.isprintable so non-ASCII printable chars stay.
+    assert decode_mdns_label_or_unknown("café") == "café"
+
+
+def test_per_column_caps_match_format_widths() -> None:
+    # The per-column caps must equal the FORMAT widths so a peer-controlled
+    # value can never widen a column past its slot. If FORMAT changes and
+    # this assertion fires, update the cap derivation in discover.py — do
+    # not just bump the expected values.
+    widths = tuple(int(w) for w in re.findall(r"<\s*(\d+)", FORMAT))
+    assert widths[COLUMN_NAMES.index("Name")] == _MAX_NAME_DISPLAY
+    assert widths[COLUMN_NAMES.index("MAC")] == _MAX_MAC_DISPLAY
+    assert widths[COLUMN_NAMES.index("Version")] == _MAX_VERSION_DISPLAY
+    assert widths[COLUMN_NAMES.index("Platform")] == _MAX_PLATFORM_DISPLAY
+    assert widths[COLUMN_NAMES.index("Board")] == _MAX_BOARD_DISPLAY