bootloader/test/test_handshake.py

#!/usr/bin/env python3
"""Host-tool activation handshake: bounded against a line that misbehaves.

`_handshake` drains the line after it sees a prompt, to absorb a real loader's
trailing bytes before it asks for the identity. That drain must be bounded: a
target that never falls quiet — a board stuck in a reset loop presents exactly
this, ~60 reboots/s of UART-reset garbage in which a stray 0x2b reads as a
prompt — otherwise spins the tool forever. Regression for that hang, plus a
control that a well-behaved loader still connects.

The handshake must also survive its own leftovers: after `--stay` the loader's
final prompt can still be in the USB pipeline when the next invocation opens
the port, and on a board wired to reset on open, that opening starts a fresh
activation window the stale prompt then betrays — the tool commits to an
identity read against a device that never heard its knock, and what it finally
collects is the application's banner. StaleDTRPort is that moment as a port.

Stdlib only, no device: host-tool logic, so it runs on every chip's preset
beside pureboot.planner.
"""
import importlib.util
import pathlib
import threading
import time

PB = pathlib.Path(__file__).resolve().parents[1] / "pureboot" / "pureboot.py"
_spec = importlib.util.spec_from_file_location("pureboot", PB)
pb = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(pb)

P = F = 0


def check(name, ok):
    global P, F
    P, F = P + (1 if ok else 0), F + (0 if ok else 1)
    print(f"  [{'PASS' if ok else 'FAIL'}] {name}")


class FloodPort:
    """A line that never falls quiet: read_available always returns bytes, and
    they contain a prompt. No identity ever completes."""

    def flush_input(self):
        pass

    def write(self, data):
        pass

    def read_available(self, wait):
        time.sleep(0.01)  # a real read waits; keep the busy loop off a core
        return b"+\x00\xff"

    def read_exact(self, count, timeout):
        raise pb.Error("no identity")


class LoaderPort:
    """A well-behaved pureboot 5: a prompt to the knock, then quiet, then the
    slim identity (version 5 + m328p signature) and a closing prompt."""

    def __init__(self):
        self.pending = b""
        self.exacts = 0

    def flush_input(self):
        self.pending = b""

    def write(self, data):
        if b"p" in data:
            self.pending = b"+"  # the prompt answers the knock, nothing else

    def read_available(self, wait):
        data, self.pending = self.pending, b""
        return data

    def read_exact(self, count, timeout):
        self.exacts += 1
        return b"\x05\x1e\x95\x0f" if self.exacts == 1 else b"+"  # identity, then prompt


class StaleDTRPort:
    """`--stay`, then a fresh invocation on a board that resets when its port
    opens. Three facts of that moment, all timed from the open: the previous
    session's final prompt is still in transit and lands only after the
    opening flush has already run; the reset holds the device off the line
    at first, eating anything written before it completes; and the fresh
    window is finite — once it expires the application boots and prints a
    banner whose bytes are what a pending identity read collects. A
    handshake that trusts the stale prompt spends the whole window waiting
    on a device that never heard its knock; one that drains the line first
    knocks into the real window and connects."""

    STALE_AT = 0.02   # the leftover prompt becomes visible (post-flush)
    READY_AT = 0.05   # reset complete, activation window opens
    WINDOW = 1.0      # window length; expiry boots the application

    def __init__(self):
        self.t0 = time.monotonic()
        # (visible-from, bytes): the line as a timed queue.
        self.queue = [(self.t0 + self.STALE_AT, b"+")]
        self.armed = False  # a 'p' heard inside the window arms 'b'
        self.booted = False

    def _boot_check(self):
        if not self.booted and time.monotonic() > self.t0 + self.READY_AT + self.WINDOW:
            self.booted = True
            self.queue.append((self.t0 + self.READY_AT + self.WINDOW,
                               b"W r libavr tempmon\r\n"))

    def _visible(self):
        self._boot_check()
        now = time.monotonic()
        return b"".join(d for t, d in self.queue if t <= now)

    def _consume(self, n):
        now = time.monotonic()
        left = []
        for t, d in self.queue:
            if t <= now and n:
                take = min(n, len(d))
                d = d[take:]
                n -= take
            if d:
                left.append((t, d))
        self.queue = left

    def flush_input(self):
        self._consume(len(self._visible()))

    def write(self, data):
        self._boot_check()
        now = time.monotonic()
        if now < self.t0 + self.READY_AT or self.booted:
            return  # still in reset, or the application owns the line
        if b"p" in data:
            self.armed = True
            self.queue.append((now + 0.01, b"+"))
        if b"b" in data and self.armed:
            # The slim identity (version 5 + m328p signature) and a prompt.
            self.queue.append((now + 0.01, b"\x05\x1e\x95\x0f+"))

    def read_available(self, wait):
        deadline = time.monotonic() + wait
        while True:
            data = self._visible()
            if data:
                self._consume(len(data))
                return data
            if time.monotonic() >= deadline:
                return b""
            time.sleep(0.005)

    def read_exact(self, count, timeout):
        deadline = time.monotonic() + timeout
        data = b""
        while len(data) < count:
            visible = self._visible()
            if visible:
                take = visible[:count - len(data)]
                self._consume(len(take))
                data += take
            elif time.monotonic() >= deadline:
                raise pb.Error(f"timeout: got {len(data)} of {count} bytes")
            else:
                time.sleep(0.005)
        return data


def terminates(port, wait, budget):
    """Run connect_autobaud in a thread; True if it returns/raises within
    `budget` seconds rather than hanging."""
    done = threading.Event()

    def run():
        try:
            pb.Loader(port).connect_autobaud(wait)
        except Exception:
            pass
        finally:
            done.set()

    threading.Thread(target=run, daemon=True).start()
    return done.wait(budget)


def main():
    # the hang: a flooding target must not spin the drain forever. With wait=0.5
    # the whole handshake has to give up well inside a few seconds.
    check("flooding target: handshake terminates, drain is bounded",
          terminates(FloodPort(), wait=0.5, budget=4.0))

    # the control: a real loader still connects and reads identity.
    info = pb.Loader(LoaderPort()).connect_autobaud(2.0)
    check("well-behaved loader still connects (version 5)", info.version == 5)

    # the stale prompt: a --stay leftover plus reset-on-open must not burn the
    # fresh window — the pre-knock drain absorbs it and the first real knock
    # lands inside the window.
    try:
        stale_ok = pb.Loader(StaleDTRPort()).connect(2.5).version == 5
    except pb.Error as failed:
        print(f"        ({failed})")
        stale_ok = False
    check("stale --stay prompt + reset-on-open: connects in the fresh window", stale_ok)

    print(f"\n  {P} passed, {F} failed")
    return 1 if F else 0


if __name__ == "__main__":
    raise SystemExit(main())