bootloader/test/pbtest.py

#!/usr/bin/env python3
"""End-to-end protocol test: drive the simavr device with the real host tool
over its pty through flash, EEPROM, fuse and hand-over scenarios, and
cross-check the tool's view against the simulator's ground-truth dumps.

Usage: pbtest.py <device_bin> <pureboot_elf> <mcu> <hz> <base_hex> <page>
                 <baud> <eeprom_size> <app_bin> <tool_py> <workdir> [link]

The optional link is the runner's -l spec (usart1, sw:B5,B1, ...), for a
loader built off the chip's natural serial default.
"""

import os
import re
import sys


def fail(message):
    print(f"FAIL: {message}")
    sys.exit(1)


def rjmp_decode(word, at, flash_words):
    """Where an rjmp word at word-address `at` lands - deliberately written
    against the instruction-set definition (12-bit signed offset), not with
    the host tool's encoder, so an encoding bug cannot verify itself."""
    if word & 0xF000 != 0xC000:
        fail(f"word at {at * 2:#06x} is {word:#06x}, not an rjmp")
    offset = word & 0x0FFF
    if offset >= 0x800:
        offset -= 0x1000
    return (at + 1 + offset) % flash_words


def main():
    args = sys.argv[1:]
    link = args.pop() if len(args) == 12 else None
    (device_bin, elf, mcu, hz, base_hex, page, baud, eeprom_size, app_bin, tool, workdir) = args
    base, page, baud, eeprom_size = int(base_hex, 0), int(page), int(baud), int(eeprom_size)
    sys.path.insert(0, os.path.dirname(os.path.abspath(tool)))
    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    import pbsim
    import pureboot as pb

    os.makedirs(workdir, exist_ok=True)
    ee_image = bytes(range(0xA0, 0xB0))
    ee_path = os.path.join(workdir, "ee.bin")
    open(ee_path, "wb").write(ee_image)
    dump = os.path.join(workdir, "flash_dump.bin")
    read_flash = os.path.join(workdir, "readback_flash.bin")
    read_eeprom = os.path.join(workdir, "readback_eeprom.bin")

    # The geometry the host will discover, for computing the expected image:
    # the boot-sectioned megas need no vector surgery (the tinies and the
    # boot-section-less m48s do), the large chips speak word addresses, and
    # the page byte is the wire's 0-means-256.
    mega = mcu.startswith("atmega")
    patch = not mega or mcu.startswith("atmega48")
    # Where SRAM begins: the x8 and x4 megas push it past their extended I/O
    # space, everything else starts right after the plain I/O registers. The
    # loader keeps no statics and its stack sits at RAMEND, so the first SRAM
    # byte is free for the data-space probe below.
    classic = mcu in ("atmega8", "atmega8a", "atmega16", "atmega16a", "atmega32", "atmega32a")
    ram_base = 0x0100 if mega and not classic else 0x0060
    word_flash = base + pb.SLOT > 0x10000
    wire_base = base // 2 if word_flash else base
    flags = (1 if patch else 0) | (2 if word_flash else 0)
    info = pb.Info(
        bytes([ord("P"), ord("B"), pb.NEWEST_LOADER, 0, 0, 0, page & 0xFF])
        + bytes([wire_base & 0xFF, wire_base >> 8, eeprom_size & 0xFF, eeprom_size >> 8])
        + bytes([flags])
    )

    # A shared-line link (RX == TX in the -l spec) makes the host read every
    # byte it sends back off the line; all sessions then discard the echo.
    one_wire = bool(link) and re.fullmatch(r"sw:([A-H][0-7]),\1(@[01])?", link) is not None
    extra = ("--one-wire",) if one_wire else ()

    device = pbsim.Device(device_bin, elf, mcu, hz, base_hex, page, baud, dump, link=link)
    try:
        # Session 1: knock from reset, identify, program everything, stay.
        out = pbsim.run_tool(tool, device.pty, baud, *extra, "--info", "--fuses", "--flash", app_bin,
                             "--eeprom", ee_path, "--stay")
        for needed in ("version", "signature", "fuses", "verify:", "stays"):
            if needed not in out:
                fail(f"session 1 output lacks {needed!r}")

        # Session 2: reconnect into the live session, verify, dump, exercise
        # the data space; hand over is deferred - the pty must be reopened for
        # the APP banner first.
        probe = "c0ffee"
        out = pbsim.run_tool(tool, device.pty, baud, *extra, "--verify-flash", app_bin, "--verify-eeprom", ee_path,
                             "--read-flash", read_flash, "--read-eeprom", read_eeprom,
                             "--poke", f"{ram_base:#x}:{probe}", "--peek", f"{ram_base:#x}:3", "--stay")
        if out.count("verify:") != 2:
            fail("session 2 did not verify both memories")
        # What went into SRAM must come back out of it: the data space is one
        # more selector on the same transfer as flash and EEPROM, so a wrong
        # selector decode would show up here and nowhere else.
        if probe not in out.replace(" ", ""):
            fail(f"data-space round trip at {ram_base:#x} did not read back {probe}\n{out}")

        eeprom_back = open(read_eeprom, "rb").read()
        if eeprom_back[: len(ee_image)] != ee_image:
            fail("EEPROM read-back mismatch")

        # The expected post-surgery flash, straight from the tool's planner.
        pages = pb.plan_flash(open(app_bin, "rb").read(), info)
        flash_back = open(read_flash, "rb").read()
        for address, data in pages.items():
            if flash_back[address : address + page] != data:
                fail(f"flash read-back mismatch in page {address:#06x}")

        # An external reset re-enters through the patched word 0 (tinies; the
        # runner resets them to address 0 like silicon) or BOOTRST (mega).
        # The loader must answer a fresh knock, and the 'J' hand-over must
        # land in the application, which banners on the same link.
        device.reset()
        port = pb.Port(device.pty, baud)
        if one_wire:
            port = pb.OneWirePort(port)
        try:
            loader = pb.Loader(port)
            live = loader.connect(15)
            # The loader built from this tree must report a version the tool
            # beside it speaks - a bump the tool was never told about is a
            # loader it would refuse to talk to. Not equality with the newest:
            # the tool now spans two loader generations, the fixed-baud one
            # here and the unified autobaud loader that follows it.
            if not pb.OLDEST_LOADER <= live.version <= pb.NEWEST_LOADER:
                fail(f"loader reports pureboot {live.version}, the tool speaks "
                     f"{pb.OLDEST_LOADER}..{pb.NEWEST_LOADER}")

            # A fill addressed inside a page rather than at its base must still
            # consume exactly one page and prompt. The loader's own slot is the
            # target and the payload is erased-state bytes, so the probe can
            # disturb neither the image nor the page buffer it leaves behind: a
            # fill only loads the buffer, and nothing commits it. Hand-built
            # rather than through write_page(), which would follow the fill
            # with its erase and write; the point here is that the fill alone
            # consumes exactly one page whatever the address's low bits say.
            # Hand-sealed too - a protocol probe that borrowed the tool's own
            # frame builder could not tell a wrong frame from a wrong loader.
            wire = base + 1
            head = bytes((pb.OP_FILL, pb.selector(pb.SP_FLASH, wire), wire & 0xFF,
                          (wire >> 8) & 0xFF, page & 0xFF))
            seal = pb.SEAL
            for byte in head:
                seal ^= byte
            port.write(head + bytes((seal,)))
            if port.read_exact(1, 5.0) != pb.PROMPT:
                fail("the loader refused a correctly sealed fill")
            port.write(b"\xff" * page)
            if port.read_exact(1, 5.0) != pb.PROMPT:
                fail("unaligned fill did not return to the prompt")

            # And the seal itself, red: one wrong bit in the address of that
            # same frame must be refused outright. The verdict has to arrive
            # *before* the page would have been sent - that ordering is what
            # keeps a refusal from turning into a desync - so the probe sends
            # no payload at all and expects the loader straight back at the
            # command level.
            broken = bytearray(head + bytes((seal,)))
            broken[2] ^= 0x01
            port.write(bytes(broken))
            if port.read_exact(1, 5.0) != pb.NAK:
                fail("a header with a broken seal was not refused")
            if port.read_exact(1, 5.0) != pb.PROMPT:
                fail("the loader did not re-prompt after refusing a broken seal")

            loader.run_application()
            banner = port.read_exact(3, 5.0)
            if banner != b"APP":
                fail(f"application banner was {banner!r}")
        finally:
            port.close()
    finally:
        device.stop()

    # Ground truth: the simulator's own memories, against the host's view.
    flash_true = open(dump, "rb").read()
    if flash_true[:base] != flash_back:
        fail("host flash read-back differs from the simulator's flash")
    if flash_true[base] == 0xFF and flash_true[base + 1] == 0xFF:
        fail("loader region looks erased in the ground-truth dump")

    # The surgery, decoded independently: the patched vector must land on the
    # loader, the trampoline on the application's own entry (patched-vector
    # chips only - a boot-sectioned mega's word 0 stays the application's).
    if patch:
        flash_words = (base + pb.SLOT) // 2
        app = open(app_bin, "rb").read()
        word0 = flash_true[0] | (flash_true[1] << 8)
        if rjmp_decode(word0, 0, flash_words) != base // 2:
            fail("patched reset vector does not land on the loader base")
        trampoline = flash_true[base - 2] | (flash_true[base - 1] << 8)
        original = app[0] | (app[1] << 8)
        if rjmp_decode(trampoline, (base - 2) // 2, flash_words) != rjmp_decode(original, 0, flash_words):
            fail("trampoline does not land on the application's own entry")
    ee_true_path = dump + ".eeprom"
    if os.path.exists(ee_true_path):
        ee_true = open(ee_true_path, "rb").read()
        if ee_true[: len(ee_image)] != ee_image:
            fail("ground-truth EEPROM does not match what was programmed")

    print("pbtest: all scenarios pass")


if __name__ == "__main__":
    main()