esphome/tests/integration/test_host_ota.py

"""End-to-end OTA tests on the host platform.

Exercises the native OTA protocol against a real host binary, then asserts
pid is preserved across the post-OTA execv. A second OTA on the post-exec
instance covers the FD_CLOEXEC path.
"""

from __future__ import annotations

import asyncio
from collections.abc import Generator
from contextlib import contextmanager
import functools
from pathlib import Path
import socket

import pytest

from esphome import espota2

from .conftest import run_binary, wait_and_connect_api_client
from .const import LOCALHOST, PORT_POLL_INTERVAL, PORT_WAIT_TIMEOUT
from .types import CompileFunction, ConfigWriter

DEVICE_NAME = "host-ota-test"
API_KEY = "AAECAwQFBgcICQoLDA0ODxAREhMUFRYXGBkaGxwdHh8="


@contextmanager
def _reserve_port() -> Generator[tuple[int, socket.socket]]:
    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    s.bind(("", 0))
    try:
        yield s.getsockname()[1], s
    finally:
        s.close()


async def _wait_for_port(host: str, port: int, timeout: float) -> None:
    """Poll until a TCP port accepts connections, or raise TimeoutError."""
    loop = asyncio.get_running_loop()
    deadline = loop.time() + timeout
    while loop.time() < deadline:
        try:
            _, writer = await asyncio.open_connection(host, port)
        except (ConnectionRefusedError, OSError):
            await asyncio.sleep(PORT_POLL_INTERVAL)
            continue
        writer.close()
        await writer.wait_closed()
        return
    raise TimeoutError(f"Port {port} on {host} did not open within {timeout}s")


@pytest.mark.asyncio
async def test_host_ota_self_update(
    yaml_config: str,
    write_yaml_config: ConfigWriter,
    compile_esphome: CompileFunction,
    reserved_tcp_port: tuple[int, socket.socket],
) -> None:
    """Self-OTA: upload the running binary back to itself, expect re-exec."""
    api_port, api_socket = reserved_tcp_port
    with _reserve_port() as (ota_port, ota_socket):
        yaml_config = yaml_config.replace("__OTA_PORT__", str(ota_port))
        config_path = await write_yaml_config(yaml_config)
        binary_path = await compile_esphome(config_path)
        api_socket.close()
        ota_socket.close()

        loop = asyncio.get_running_loop()
        ota_staged = loop.create_future()
        rebooted = loop.create_future()

        def on_log(line: str) -> None:
            if not ota_staged.done() and "OTA staged at" in line:
                ota_staged.set_result(True)
            if not rebooted.done() and "Rebooting safely" in line:
                rebooted.set_result(True)

        async with run_binary(binary_path, line_callback=on_log) as (proc, _lines):
            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
            pid_before = proc.pid
            async with wait_and_connect_api_client(port=api_port) as client:
                info_before = await client.device_info()
                assert info_before.name == DEVICE_NAME

            # espota2 is blocking; run in executor.
            rc, _ = await loop.run_in_executor(
                None, espota2.run_ota, LOCALHOST, ota_port, None, binary_path
            )
            assert rc == 0, "espota2 reported failure"

            await asyncio.wait_for(ota_staged, timeout=10.0)
            await asyncio.wait_for(rebooted, timeout=10.0)
            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)

            # execv preserves pid; mismatch means external respawn.
            assert proc.returncode is None, "process exited instead of execing"
            assert proc.pid == pid_before

            async with wait_and_connect_api_client(port=api_port) as client:
                info_after = await client.device_info()
                assert info_after.name == DEVICE_NAME
                assert info_after.name == info_before.name

            # Second OTA: catches FD_CLOEXEC regressions (EADDRINUSE on rebind).
            rc, _ = await loop.run_in_executor(
                None, espota2.run_ota, LOCALHOST, ota_port, None, binary_path
            )
            assert rc == 0, "second OTA failed -- listener leaked across execv"
            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
            assert proc.pid == pid_before


@pytest.mark.asyncio
async def test_host_ota_encrypted(
    yaml_config: str,
    write_yaml_config: ConfigWriter,
    compile_esphome: CompileFunction,
    reserved_tcp_port: tuple[int, socket.socket],
) -> None:
    """Encrypted self-OTA succeeds; a plaintext upload to the same device fails."""
    pytest.importorskip("aioesphomeapi.noise")
    api_port, api_socket = reserved_tcp_port
    with _reserve_port() as (ota_port, ota_socket):
        yaml_config = yaml_config.replace("__OTA_PORT__", str(ota_port))
        config_path = await write_yaml_config(yaml_config)
        binary_path = await compile_esphome(config_path)
        api_socket.close()
        ota_socket.close()

        loop = asyncio.get_running_loop()
        rebooted = loop.create_future()

        def on_log(line: str) -> None:
            if not rebooted.done() and "Rebooting safely" in line:
                rebooted.set_result(True)

        async with run_binary(binary_path, line_callback=on_log) as (proc, _lines):
            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
            pid_before = proc.pid

            # A plaintext upload must be refused with the device unharmed
            rc, _ = await loop.run_in_executor(
                None, espota2.run_ota, LOCALHOST, ota_port, None, binary_path
            )
            assert rc == 1, "plaintext upload to an encrypted device must fail"
            await asyncio.sleep(0.5)
            assert proc.returncode is None, "process died on rejected plaintext OTA"

            # The encrypted upload goes through and the device re-execs
            rc, _ = await loop.run_in_executor(
                None,
                functools.partial(
                    espota2.run_ota,
                    LOCALHOST,
                    ota_port,
                    None,
                    binary_path,
                    noise_psk=API_KEY,
                ),
            )
            assert rc == 0, "encrypted OTA reported failure"
            await asyncio.wait_for(rebooted, timeout=10.0)
            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
            assert proc.returncode is None, "process exited instead of execing"
            assert proc.pid == pid_before


class _RebootCounter:
    """Counts safe reboots so a test can wait for the nth one."""

    def __init__(self) -> None:
        self._seen = asyncio.Event()
        self.count = 0

    def on_log(self, line: str) -> None:
        if "Rebooting safely" in line:
            self.count += 1
            self._seen.set()

    async def wait(self, count: int, timeout: float = 10.0) -> None:
        async with asyncio.timeout(timeout):
            while self.count < count:
                self._seen.clear()
                await self._seen.wait()


async def _run_ota(
    ota_port: int, password: str | None, binary_path: Path, noise_psk: str | None
) -> int:
    """espota2 is blocking; run it in the executor and return its exit code."""
    rc, _ = await asyncio.get_running_loop().run_in_executor(
        None,
        functools.partial(
            espota2.run_ota,
            LOCALHOST,
            ota_port,
            password,
            binary_path,
            noise_psk=noise_psk,
        ),
    )
    return rc


@pytest.mark.asyncio
async def test_host_ota_api_key_offer_with_password(
    yaml_config: str,
    write_yaml_config: ConfigWriter,
    compile_esphome: CompileFunction,
    reserved_tcp_port: tuple[int, socket.socket],
) -> None:
    """With only an api key the device offers encryption without requiring
    it: the password still guards plaintext uploads, and the key alone
    authenticates an encrypted one, which is the enablement path for
    `ota: encryption:`."""
    pytest.importorskip("aioesphomeapi.noise")
    api_port, api_socket = reserved_tcp_port
    with _reserve_port() as (ota_port, ota_socket):
        yaml_config = yaml_config.replace("__OTA_PORT__", str(ota_port))
        config_path = await write_yaml_config(yaml_config)
        binary_path = await compile_esphome(config_path)
        api_socket.close()
        ota_socket.close()

        reboots = _RebootCounter()
        async with run_binary(binary_path, line_callback=reboots.on_log) as (
            proc,
            lines,
        ):
            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
            pid_before = proc.pid

            rc = await _run_ota(ota_port, None, binary_path, None)
            assert rc == 1, "plaintext upload without the password must fail"
            await asyncio.sleep(0.5)
            assert proc.returncode is None, "process died on rejected upload"

            rc = await _run_ota(ota_port, "hunter2", binary_path, None)
            assert rc == 0, "plaintext upload with the password must succeed"
            await reboots.wait(1)
            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
            assert proc.pid == pid_before

            rc = await _run_ota(ota_port, None, binary_path, API_KEY)
            assert rc == 0, "encrypted upload with the api key must succeed"
            await reboots.wait(2)
            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
            assert proc.returncode is None, "process exited instead of execing"
            assert proc.pid == pid_before
            assert any("Encryption: offered" in line for line in lines)


@pytest.mark.asyncio
async def test_host_ota_rejects_garbage(
    yaml_config: str,
    write_yaml_config: ConfigWriter,
    compile_esphome: CompileFunction,
    reserved_tcp_port: tuple[int, socket.socket],
    integration_test_dir,
) -> None:
    """Bogus payload is rejected and the device keeps running."""
    api_port, api_socket = reserved_tcp_port
    with _reserve_port() as (ota_port, ota_socket):
        yaml_config = yaml_config.replace("__OTA_PORT__", str(ota_port))
        config_path = await write_yaml_config(yaml_config)
        binary_path = await compile_esphome(config_path)

        # 192 bytes that are neither ELF nor Mach-O.
        bogus_path = integration_test_dir / "bogus.bin"
        bogus_path.write_bytes(b"NOT-AN-EXECUTABLE-AT-ALL" * 8)

        api_socket.close()
        ota_socket.close()

        async with run_binary(binary_path) as (proc, _lines):
            await _wait_for_port(LOCALHOST, api_port, PORT_WAIT_TIMEOUT)
            pid_before = proc.pid

            loop = asyncio.get_running_loop()
            rc, _ = await loop.run_in_executor(
                None, espota2.run_ota, LOCALHOST, ota_port, None, bogus_path
            )
            assert rc == 1

            await asyncio.sleep(0.5)
            assert proc.returncode is None, "process died on rejected OTA"
            assert proc.pid == pid_before

            async with wait_and_connect_api_client(port=api_port) as client:
                info = await client.device_info()
                assert info.name == DEVICE_NAME