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, signature=(0x1E, 0x93, 0x0B), word_flash=False):
    scale = 2 if word_flash else 1
    wire_base = base // scale
    flags = (1 if patch else 0) | (2 if word_flash else 0)
    raw = bytes((0x50, 0x42, 1, *signature, page & 0xFF, wire_base & 0xFF, wire_base >> 8,
                 0, 2, flags))
    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, signature=(0x1E, 0x95, 0x0F))

    # mega_boot: BOOTSZ words and the BOOTRST sense per chip — the fuse byte
    # index (HIGH everywhere but the m168A's EXTENDED) and the per-family
    # ladders (Atmel-2486/2466/2503/8271/42719). Synthetic 'F' replies: only
    # the boot byte carries meaning.
    cases = (
        ((0x1E, 0x93, 0x07), 0x2000, 3, {0b11: 0x1F00, 0b10: 0x1E00, 0b01: 0x1C00, 0b00: 0x1800}),  # m8
        ((0x1E, 0x94, 0x03), 0x4000, 3, {0b11: 0x3F00, 0b10: 0x3E00, 0b01: 0x3C00, 0b00: 0x3800}),  # m16
        ((0x1E, 0x95, 0x02), 0x8000, 3, {0b11: 0x7E00, 0b10: 0x7C00, 0b01: 0x7800, 0b00: 0x7000}),  # m32
        ((0x1E, 0x94, 0x06), 0x4000, 2, {0b11: 0x3F00, 0b10: 0x3E00, 0b01: 0x3C00, 0b00: 0x3800}),  # m168A
        ((0x1E, 0x95, 0x0F), 0x8000, 3, {0b11: 0x7E00, 0b10: 0x7C00, 0b01: 0x7800, 0b00: 0x7000}),  # m328P
        ((0x1E, 0x97, 0x05), 0x20000, 3, {0b11: 0x1FC00, 0b10: 0x1F800, 0b01: 0x1F000, 0b00: 0x1E000}),  # 1284P
    )
    for signature, flash, which, ladder in cases:
        chip = info_of(pb, flash - 512, 128 if flash < 0x20000 else 0, False, flash,
                       signature=signature, word_flash=flash > 0x10000)
        for bits, start in ladder.items():
            fuses = bytearray((0xFF, 0xFF, 0xFF, 0xFF))
            fuses[which] = (0xF8 | (bits << 1)) & ~1
            prog, at = pb.mega_boot(chip, bytes(fuses))
            if not prog or at != start:
                fail(f"mega_boot {signature[1]:02x}{signature[2]:02x} BOOTSZ={bits:02b} programmed: {prog} {at:#07x}")
            fuses[which] |= 1
            prog, at = pb.mega_boot(chip, bytes(fuses))
            if prog or at != start:
                fail(f"mega_boot {signature[1]:02x}{signature[2]:02b} unprogrammed: {prog} {at:#07x}")

    # Word-addressed info decode: the 1284P's base/page ride the wire scaled.
    big = info_of(pb, 0x1FE00, 0, False, 0x20000, signature=(0x1E, 0x97, 0x05), word_flash=True)
    if big.page != 256 or big.base != 0x1FE00 or big.stage != 0x1FC00:
        fail(f"word-addressed info decode: page {big.page}, base {big.base:#x}, stage {big.stage:#x}")

    # 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()