espemu: eFuses, flash encryption, esp32c5, and a working hard reset (CII-251)

pytest-embedded-espemu/README.md

@@ -2,7 +2,7 @@
 
 pytest-embedded service for running tests on [esp-emu](https://github.com/espressif/esp-emulator), Espressif's lightweight emulator for ESP RISC-V series SoCs, instead of real targets.
 
-Supported targets: `esp32c3`, `esp32c6`, `esp32h2`, `esp32p4`, `esp32s31`.
+Supported targets: `esp32c3`, `esp32c5`, `esp32c6`, `esp32h2`, `esp32p4`, `esp32s31`.
 
 #### Usage
 

pytest-embedded-espemu/pytest_embedded_espemu/__init__.py

@@ -5,15 +5,18 @@
 from pytest_embedded.utils import lazy_load
 
 DEFAULT_IMAGE_FN = 'espemu_image.bin'
+ENCRYPTED_IMAGE_FN = f'encrypted_{DEFAULT_IMAGE_FN}'
 
 from .dut import EspEmuDut  # noqa
 from .espemu import EspEmu  # noqa
+from .serial import EspEmuSerial  # noqa
 
 __getattr__ = lazy_load(
     importlib.import_module(__name__),
     {
         'EspEmu': EspEmu,
         'EspEmuDut': EspEmuDut,
+        'EspEmuSerial': EspEmuSerial,
     },
     {
         'EspEmuApp': '.app',  # requires idf
@@ -23,6 +26,7 @@
 
 __all__ = [
     'DEFAULT_IMAGE_FN',
+    'ENCRYPTED_IMAGE_FN',
     'EspEmu',
     'EspEmuApp',
     'EspEmuDut',

pytest-embedded-espemu/pytest_embedded_espemu/app.py

@@ -6,7 +6,7 @@
 from pytest_embedded.log import MessageQueue, live_print_call
 from pytest_embedded_idf.app import IdfApp
 
-from . import DEFAULT_IMAGE_FN
+from . import DEFAULT_IMAGE_FN, ENCRYPTED_IMAGE_FN
 
 
 class EspEmuApp(IdfApp):
@@ -22,6 +22,8 @@ def __init__(
         msg_queue: MessageQueue,
         espemu_image_path: str | None = None,
         skip_regenerate_image: bool | None = False,
+        encrypt: bool | None = False,
+        keyfile: str | None = None,
         **kwargs,
     ):
         self._q = msg_queue
@@ -30,9 +32,39 @@ def __init__(
 
         self.image_path = espemu_image_path or os.path.join(self.binary_path, DEFAULT_IMAGE_FN)
         self.skip_regenerate_image = skip_regenerate_image
+        self.encrypt = encrypt
+        self.keyfile = keyfile
+
+        if self.encrypt:
+            self.encrypted_image_path = os.path.join(self.binary_path, ENCRYPTED_IMAGE_FN)
 
         self.create_image()
 
+    def _write_encrypted_image(self, seek: int = 0) -> None:
+        """
+        Encrypt the merged image with the flash encryption key, as the ROM would
+        have written it.
+        """
+        with contextlib.redirect_stdout(self._q):
+            live_print_call(
+                [
+                    sys.executable,
+                    '-m',
+                    'espsecure',
+                    'encrypt-flash-data',
+                    # every target the emulator supports encrypts flash with
+                    # XTS-AES, not the scheme the original esp32 uses
+                    '--aes-xts',
+                    '--keyfile',
+                    self.keyfile,
+                    '--output',
+                    self.encrypted_image_path,
+                    '--address',
+                    str(seek),
+                    self.image_path,
+                ],
+            )
+
     def create_image(self) -> None:
         """
         Create the image, if it doesn't exist.
@@ -65,3 +97,8 @@ def create_image(self) -> None:
                 ],
                 cwd=self.binary_path,
             )
+
+        if self.encrypt:
+            if self.keyfile is None or not os.path.exists(self.keyfile):
+                raise ValueError("Flash Encryption key file doesn't exist")
+            self._write_encrypted_image()

pytest-embedded-espemu/pytest_embedded_espemu/dut.py

@@ -3,6 +3,7 @@
 from pytest_embedded.dut import Dut
 
 from .espemu import EspEmu
+from .serial import EspEmuSerial
 
 
 class EspEmuDut(Dut):
@@ -27,6 +28,11 @@ def __init__(
         # target is known here as well; keep it optional for other app classes.
         self.target: str | None = getattr(self.app, 'target', None)
 
+        # `IdfSerial` is what ESP-IDF tests call for device state — a reset,
+        # an erase, an eFuse burn. Give them something that answers, so a test
+        # that needs one reports which operation it needed.
+        self.serial = EspEmuSerial(espemu)
+
         self._hard_reset_func = self.espemu._hard_reset
 
     def write(self, s: AnyStr) -> None:

pytest-embedded-espemu/pytest_embedded_espemu/espemu.py

@@ -1,4 +1,9 @@
+import logging
+import os
 import shlex
+import socket
+import subprocess
+import time
 import typing as t
 
 from pytest_embedded.log import DuplicateStdoutPopen
@@ -20,14 +25,25 @@ class EspEmu(DuplicateStdoutPopen):
 
     ESPEMU_PROG_PATH = 'esp-emu'
 
-    SUPPORTED_TARGETS: t.ClassVar[tuple] = ('esp32c3', 'esp32c6', 'esp32h2', 'esp32p4', 'esp32s31')
+    SUPPORTED_TARGETS: t.ClassVar[tuple] = ('esp32c3', 'esp32c5', 'esp32c6', 'esp32h2', 'esp32p4', 'esp32s31')
+
+    # esp-emu reads and writes a QEMU compatible eFuse image of this size. A
+    # blank one means "nothing burned", which is the state a test starts from
+    # unless it is given an image of its own.
+    EFUSE_IMAGE_SIZE = 336
+
+    # Boot to the ROM loader instead of the firmware, so esptool and espefuse
+    # can drive the emulator over the UART socket. There are no modem control
+    # lines on a socket, so the strap has to be set at launch.
+    DOWNLOAD_MODE_STRAP = '0x02'
 
     def __init__(
         self,
         espemu_image_path: str | None = None,
         espemu_prog_path: str | None = None,
         espemu_cli_args: str | None = None,
         espemu_extra_args: str | None = None,
+        espemu_efuse_path: str | None = None,
         app: t.Optional['EspEmuApp'] = None,
         **kwargs,
     ):
@@ -37,6 +53,7 @@ def __init__(
             espemu_prog_path: esp-emu program path
             espemu_cli_args: esp-emu CLI arguments
             espemu_extra_args: esp-emu CLI extra arguments, will be appended to `espemu_cli_args`
+            espemu_efuse_path: eFuse image the emulator reads at start and writes back on exit
             app: `EspEmuApp` instance, used to detect the target chip
         """
         self.app = app
@@ -55,22 +72,152 @@ def __init__(
 
         espemu_prog_path = espemu_prog_path or self.ESPEMU_PROG_PATH
 
+        self.espemu_prog_path = espemu_prog_path
+        self.image_path = image_path
+        self.target = target
+        self.efuse_path = espemu_efuse_path
+
+        efuse_args = []
+        if self.efuse_path:
+            self._create_efuse_image(self.efuse_path)
+            logging.debug('The eFuse image will be saved to: %s', self.efuse_path)
+            efuse_args = ['--efuse', self.efuse_path]
+        # A chip reset has no in-band form — on hardware it is a DTR/RTS toggle
+        # into EN, and pyserial's `socket://` handler ignores modem control
+        # lines — so it goes over esp-emu's control channel. Older binaries do
+        # not have the flag, and passing an unknown one makes them exit, so ask
+        # first and stay resettable-or-not accordingly.
+        self.control_port: int | None = None
+        control_args = []
+        if self._supports_control_channel(espemu_prog_path):
+            self.control_port = self._free_port()
+            control_args = ['--control-tcp', f'127.0.0.1:{self.control_port}']
+
         cmd = [
             espemu_prog_path,
             '--chip',
             target,
             '--firmware',
             image_path,
+            *efuse_args,
+            *control_args,
             *shlex.split(espemu_cli_args or ''),
             *shlex.split(espemu_extra_args or ''),
         ]
 
         super().__init__(cmd=cmd, **kwargs)
 
+    @classmethod
+    def _create_efuse_image(cls, path: str) -> None:
+        """Create a blank eFuse image, keeping one that already exists."""
+        if os.path.exists(path):
+            return
+
+        with open(path, 'wb') as f:
+            f.write(b'\x00' * cls.EFUSE_IMAGE_SIZE)
+
+    def execute_efuse_command(self, command: str) -> None:
+        """
+        Run an espefuse command against the emulator.
+
+        A second emulator instance is started in download mode with its UART on
+        a socket, since the running one is booted into the firmware and a socket
+        carries no reset lines. The instance writes the eFuse image back on
+        exit, so the burned bits are there the next time the emulator starts.
+
+        Args:
+            command: espefuse command line, e.g. "burn-custom-mac 00:11:22:33:44:55"
+        """
+        import espefuse
+
+        if not self.efuse_path:
+            raise ValueError('No eFuse image set. Please use --espemu-efuse-path')
+
+        available_port = self._free_port()
+
+        child = subprocess.Popen(
+            [
+                self.espemu_prog_path,
+                '--chip',
+                self.target,
+                '--firmware',
+                self.image_path,
+                '--efuse',
+                self.efuse_path,
+                '--strap-mode',
+                self.DOWNLOAD_MODE_STRAP,
+                '--uart-tcp',
+                f'127.0.0.1:{available_port}',
+            ],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+        )
+        try:
+            self._wait_for_port(available_port)
+            args = [arg for arg in shlex.split(command) if arg != '--do-not-confirm']
+            espefuse.main(
+                [
+                    '--port',
+                    f'socket://127.0.0.1:{available_port}',
+                    '--chip',
+                    self.target,
+                    '--do-not-confirm',
+                    *args,
+                ]
+            )
+        finally:
+            child.terminate()
+            child.wait(timeout=10)
+
+    @staticmethod
+    def _wait_for_port(port: int, timeout: float = 30) -> None:
+        """Wait until the emulator accepts connections on its UART socket."""
+        deadline = time.time() + timeout
+        while time.time() < deadline:
+            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+                s.settimeout(0.5)
+                if s.connect_ex(('127.0.0.1', port)) == 0:
+                    return
+            time.sleep(0.1)
+
+        raise TimeoutError(f'esp-emu did not open its UART socket on port {port} within {timeout}s')
+
+    @staticmethod
+    def _free_port() -> int:
+        """An unused local port for one of the emulator's sockets."""
+        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+            s.bind(('127.0.0.1', 0))
+            return int(s.getsockname()[1])
+
+    @classmethod
+    def _supports_control_channel(cls, prog_path: str) -> bool:
+        """Whether this esp-emu build accepts `--control-tcp`."""
+        try:
+            out = subprocess.run([prog_path, '--help'], capture_output=True, text=True, timeout=15)
+        except (OSError, subprocess.SubprocessError):
+            return False
+        return '--control-tcp' in (out.stdout + out.stderr)
+
+    def control_command(self, command: str, timeout: float = 30) -> str:
+        """Send one command to the emulator's control channel and return its reply."""
+        if self.control_port is None:
+            raise NotImplementedError(
+                'this esp-emu build has no --control-tcp channel; '
+                'a newer emulator is needed for device-state operations'
+            )
+        with socket.create_connection(('127.0.0.1', self.control_port), timeout=timeout) as sock:
+            sock.sendall(f'{command}\n'.encode())
+            with sock.makefile('rb') as reader:
+                reply = reader.readline().decode().strip()
+        if reply.startswith('err:'):
+            raise RuntimeError(f'esp-emu rejected {command!r}: {reply}')
+        return reply
+
     def _hard_reset(self):
         """
-        esp-emu has no reset API. Raising `NotImplementedError` makes
-        `IdfUnityDutMixin` fall back to re-triggering the Unity test menu
-        with a newline instead of resetting the target.
+        Reset the emulated chip, the way a DTR/RTS toggle does on hardware.
+
+        The emulator process keeps running, so the dut's output stream, its
+        expect history and the log all continue across the reset.
         """
-        raise NotImplementedError('esp-emu does not support resetting; relaunch the emulator instead')
+        self.control_command('reset')

pytest-embedded-espemu/pytest_embedded_espemu/serial.py

@@ -0,0 +1,47 @@
+import logging
+import typing as t
+
+if t.TYPE_CHECKING:
+    from .espemu import EspEmu
+
+
+class EspEmuSerial:
+    """
+    The device-state half of an emulated dut.
+
+    ESP-IDF tests reach the chip through two channels: the console stream
+    (``dut.expect``) and a control channel (``dut.serial``) that resets it,
+    erases flash and burns eFuses. The second one is esptool-backed on
+    hardware, and an emulator has no serial port behind it, so the operations
+    that are expressible go over esp-emu's control channel instead.
+
+    Only ``hard_reset`` is implemented so far. The rest of ``IdfSerial`` needs
+    the emulator in download mode with its UART on a socket, so esptool can
+    drive it; until then those methods raise, and the test reports what it
+    needed rather than an ``AttributeError`` on a missing attribute.
+    """
+
+    def __init__(self, espemu: 'EspEmu') -> None:
+        self.espemu = espemu
+
+    @property
+    def port(self) -> str:
+        """The control channel's address, in the place a port name would be."""
+        if self.espemu.control_port is None:
+            return 'espemu'
+        return f'espemu://127.0.0.1:{self.espemu.control_port}'
+
+    def hard_reset(self) -> None:
+        """Reset the chip, the way a DTR/RTS toggle does on hardware."""
+        logging.debug('hard resetting the emulated chip')
+        self.espemu._hard_reset()
+
+    def close(self) -> None:
+        """Nothing to close: the control channel is opened per command."""
+
+    def __getattr__(self, name: str) -> t.Any:
+        raise NotImplementedError(
+            f'esp-emu cannot do dut.serial.{name}() yet. Only hard_reset is '
+            f'implemented; flash and eFuse operations need the emulator in '
+            f'download mode with its UART on a socket.'
+        )