bootloader/test/test_planner.py

#!/usr/bin/env python3
"""Host-tool unit tests — the pure planning and policy logic, no simulator:
the flash-programming orders and their recovery properties, the reset-vector
surgery, the staging-slot composition, the mega boot-fuse decode, and the
update preflight's error/warning matrix (fuse combinations simavr cannot
model reach it here as synthetic bytes).

Usage: test_planner.py <tool_py>
"""

import sys


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


def expect_error(what, fn, *needles):
    try:
        fn()
    except Exception as error:
        for needle in needles:
            if needle not in str(error):
                fail(f"{what}: error lacks {needle!r}: {error}")
        return
    fail(f"{what}: no error raised")


def info_of(pb, base, page, patch, flash):
    raw = bytes((0x50, 0x42, 1, 0x1E, 0x93, 0x0B, page, base & 0xFF, base >> 8,
                 0, 2, 1 if patch else 0))
    info = pb.Info(raw)
    assert info.flash_size == flash
    return info


def rjmp_decode(word, at, flash_words):
    if word & 0xF000 != 0xC000:
        fail(f"not an rjmp: {word:#06x}")
    offset = word & 0x0FFF
    if offset >= 0x800:
        offset -= 0x1000
    return (at + 1 + offset) % flash_words


def main():
    import os
    sys.path.insert(0, os.path.dirname(os.path.abspath(sys.argv[1])))
    import pureboot as pb

    tiny = info_of(pb, 0x1E00, 64, True, 0x2000)
    mega = info_of(pb, 0x7E00, 128, False, 0x8000)

    # mega_boot: BOOTSZ words and the BOOTRST sense, DS40002061B §27.
    for bits, start in ((0b11, 0x7E00), (0b10, 0x7C00), (0b01, 0x7800), (0b00, 0x7000)):
        prog, at = pb.mega_boot((0xF8 | (bits << 1)) & ~1)
        if not prog or at != start:
            fail(f"mega_boot BOOTSZ={bits:02b} programmed: {prog} {at:#06x}")
        prog, at = pb.mega_boot(0xF8 | (bits << 1) | 1)
        if prog or at != start:
            fail(f"mega_boot BOOTSZ={bits:02b} unprogrammed: {prog} {at:#06x}")

    # Surgery: word 0 lands on the loader, the trampoline on the original
    # entry — checked with an independent decoder.
    app = bytes((0xC0 | 0x00, 0xC0)) + bytes((0x12,)) * 300  # rjmp .+0x00C0... entry word 0xC0C0
    entry = rjmp_decode(app[0] | (app[1] << 8), 0, tiny.flash_size // 2)
    pages = pb.plan_flash(app, tiny)
    word0 = pages[0][0] | (pages[0][1] << 8)
    if rjmp_decode(word0, 0, tiny.flash_size // 2) != tiny.base // 2:
        fail("surgery: patched word 0 misses the loader")
    tp = pages[tiny.base - 64]
    tramp = tp[62] | (tp[63] << 8)
    if rjmp_decode(tramp, (tiny.base - 2) // 2, tiny.flash_size // 2) != entry:
        fail("surgery: trampoline misses the original entry")
    expect_error("non-rjmp vector", lambda: pb.plan_flash(bytes((0x0C, 0x94)) + app[2:], tiny), "not an rjmp")
    looped = bytearray(app)
    word = pb.rjmp_to(0, tiny.base // 2, tiny.flash_size // 2)
    looped[0], looped[1] = word & 0xFF, word >> 8
    expect_error("read-back image", lambda: pb.plan_flash(bytes(looped), tiny), "read-back")
    expect_error("oversize image", lambda: pb.plan_flash(bytes(0x1DFF), tiny), "application flash ends")

    # Ordering: patched vector puts page 0 first and the trampoline second;
    # a boot section puts page 0 last. Blank pages drop only when erased.
    order = pb.covered(pages, tiny, skip_blank=False)
    if order[0] != 0 or order[1] != tiny.base - 64:
        fail(f"tiny order starts {order[:2]}, want page 0 then trampoline page")
    if sorted(order[2:]) != order[2:]:
        fail("tiny order tail not ascending")
    mega_pages = pb.plan_flash(bytes((0xFF,)) * 600, mega)
    morder = pb.covered(mega_pages, mega, skip_blank=False)
    if morder[-1] != 0 or sorted(morder[:-1]) != morder[:-1]:
        fail(f"mega order {morder}, want ascending with page 0 last")
    blanky = {0: pages[0], 64: bytes((0xFF,)) * 64, 128: pages[128], tiny.base - 64: tp}
    slim = pb.covered(blanky, tiny, skip_blank=True)
    if 64 in slim or 0 not in slim or tiny.base - 64 not in slim:
        fail(f"skip_blank order wrong: {slim}")

    # Staging content: the identical image plus the through-word on a
    # patched-vector chip; hard size clamps either way.
    image = bytes(range(256)) * 2  # 512 B — too big for a tiny slot
    expect_error("tiny staging size", lambda: pb.staging_content(image, tiny), "510")
    staged = pb.staging_content(image[:508], tiny)
    through = staged[510] | (staged[511] << 8)
    if rjmp_decode(through, (tiny.base - 2) // 2, tiny.flash_size // 2) != tiny.base // 2:
        fail("through-word misses the resident base")
    if pb.staging_content(image, mega) != image:
        fail("mega staging content should be the bare image")
    expect_error("mega staging size", lambda: pb.staging_content(image + b"!", mega), "512")

    # The embedded info block: found in a synthetic binary, absent in noise.
    binary = bytes((0xAA,)) * 10 + tiny.raw + bytes((0xBB,)) * 10
    found = pb.image_info(binary)
    if found is None or found.raw != tiny.raw:
        fail("image_info misses the embedded block")
    if pb.image_info(bytes((0xAA,)) * 40) is not None:
        fail("image_info invents a block")

    # Update preflight: the full fuse matrix, plus target mismatch.
    other = info_of(pb, 0x1E00, 32, True, 0x2000)
    expect_error("wrong-target image", lambda: pb.update_preflight(binary, other, None), "another target")
    expect_error("mega needs fuses", lambda: pb.update_preflight(bytes((0xAA,)) * 8 + mega.raw, mega, None),
                 "--assume-fuses")
    mega_image = bytes((0xAA,)) * 8 + mega.raw

    def fuses(high):
        return bytes((0xFF, 0xFF, 0xFF, high))

    expect_error("BOOTSZ 512 B", lambda: pb.update_preflight(mega_image, mega, fuses(0xFE)),
                 "cannot self-update", "BOOTSZ")
    notes = pb.update_preflight(mega_image, mega, fuses(0xFD))  # 1 KB, BOOTRST unprogrammed
    if not any("BOOTRST unprogrammed" in n for n in notes):
        fail(f"1K/unprogrammed notes: {notes}")
    notes = pb.update_preflight(mega_image, mega, fuses(0xFC))  # 1 KB, BOOTRST programmed
    if not any("staging slot" in n for n in notes):
        fail(f"1K/programmed notes: {notes}")
    notes = pb.update_preflight(mega_image, mega, fuses(0xFA))  # 2 KB, BOOTRST programmed
    if not any("application flash" in n for n in notes):
        fail(f"2K/programmed notes: {notes}")
    if pb.update_preflight(bytes((0xAA,)) * 8 + tiny.raw, tiny, None) != []:
        fail("tiny preflight should pass without fuses")

    # The walk-region refusal: BOOTRST aimed below the loader plus app data
    # in the walk span errors without --force; erased spans and unprogrammed
    # BOOTRST pass.
    deep = {0x7800: bytes((1,)) * 128}
    expect_error("walk region", lambda: pb.check_walk_region(deep, mega, fuses(0xFA), False), "--force")
    pb.check_walk_region(deep, mega, fuses(0xFA), True)
    pb.check_walk_region(deep, mega, fuses(0xFB), False)  # BOOTRST unprogrammed
    pb.check_walk_region({0x7800: bytes((0xFF,)) * 128}, mega, fuses(0xFA), False)
    pb.check_walk_region(deep, mega, None, False)  # fuses unknown: no check

    print("test_planner: all planner and policy checks pass")


if __name__ == "__main__":
    main()