diff --git a/pyproject.toml b/pyproject.toml index 19630c6..0bf9c80 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ [project.optional-dependencies] tui = ["textual>=0.40"] web = ["fastapi>=0.100", "uvicorn>=0.20", "jinja2>=3.1"] +rockchip = ["pyusb>=1.2"] dev = [ "pytest>=7", "pytest-asyncio>=0.21", @@ -67,6 +68,12 @@ strict = true module = "defib.tui.*" disallow_subclassing_any = false +# pyusb ships no stubs and is an optional extra, so it is absent from the +# type-check environment entirely. +[[tool.mypy.overrides]] +module = ["usb", "usb.*"] +ignore_missing_imports = true + [tool.pytest.ini_options] asyncio_mode = "auto" testpaths = ["tests"] diff --git a/src/defib/cli/app.py b/src/defib/cli/app.py index 93d7f9e..2c2c4cf 100644 --- a/src/defib/cli/app.py +++ b/src/defib/cli/app.py @@ -24,19 +24,30 @@ def burn( poe_port_override: str = typer.Option("", "--poe-port", help="Explicit MikroTik ether port (e.g. ether3) — overrides comment-based auto-discovery. Requires --power-cycle."), output: str = typer.Option("human", "--output", help="Output mode: human, json, quiet"), debug: bool = typer.Option(False, "-d", "--debug", help="Enable debug logging"), + ddr: str = typer.Option("", "--ddr", help="USB-recovery chips: DDR-init blob (rkbin rv1106_ddr_*.bin)"), + usbplug: str = typer.Option("", "--usbplug", help="USB-recovery chips: usbplug blob (rkbin rv1106_usbplug_*.bin)"), + loader: str = typer.Option("", "--loader", help="USB-recovery chips: RKBOOT container (MiniLoaderAll.bin), instead of --ddr/--usbplug"), + wait: float = typer.Option(30.0, "--wait", help="USB-recovery chips: seconds to wait for the board to enumerate"), + usb_path: str = typer.Option("", "--usb-path", help="USB-recovery chips: pin to one physical port path (e.g. 1-4.2); required when several Rockchip boards are attached"), ) -> None: - """Recover a device by uploading firmware via UART serial. + """Wake a dead device so its flash can be written. - If no firmware file is specified with -f, automatically downloads - the appropriate U-Boot from OpenIPC releases. + For most chips this uploads U-Boot into RAM over UART; if no file is + given with -f the right one is downloaded from OpenIPC. + + Chips whose boot ROM only answers on USB (Rockchip) instead take + --ddr/--usbplug (or --loader) and ignore -p: an erased flash drops them + into MaskROM on its own at power-up, so --power-cycle is all it takes. """ import asyncio - asyncio.run(_burn_async(chip, file, port, send_break, terminal, power_cycle, poe_port_override, output, debug)) + asyncio.run(_burn_async(chip, file, port, send_break, terminal, power_cycle, poe_port_override, output, debug, ddr, usbplug, loader, wait, usb_path)) async def _burn_async( chip: str, file: str, port: str, send_break: bool, terminal: bool, power_cycle: bool, poe_port_override: str, output: str, debug: bool, + ddr: str = "", usbplug: str = "", loader: str = "", wait: float = 30.0, + usb_path: str = "", ) -> None: import json as json_mod import logging @@ -54,6 +65,13 @@ async def _burn_async( else: logging.basicConfig(level=logging.INFO) + if _recovery_mode_or_exit(chip, output) == "usb": + await _burn_usb_async( + chip, ddr, usbplug, loader, power_cycle, output, wait, + poe_port_override, usb_path, + ) + return + # Resolve firmware: local file or auto-download from OpenIPC firmware_path = file if not firmware_path: @@ -2212,6 +2230,19 @@ async def _install_async( else: logging.basicConfig(level=logging.INFO) + if _recovery_mode_or_exit(chip, output) == "usb": + # Writing flash over USB is not landed yet: this chip's images have + # nowhere correct to go until the UBI-vs-raw rootfs question is + # settled, and the write path has never run against hardware. Wake + # the board with `burn` and flash it by other means meanwhile. + # Routed through _usb_fail so --output json still gets an event. + _usb_fail( + output, + f"{chip} recovers over USB, and `install` does not support that " + "yet — use `defib burn` to bring the board up.", + ) + return + if nand: layout = _NAND_LAYOUT flash_cmd = "nand" @@ -3611,5 +3642,194 @@ async def _replace_in_tftp(name: str, data: bytes) -> None: print(json_mod.dumps({"event": "done", "success": True, "partitions": len(partitions)})) +def _recovery_mode_or_exit(chip: str, output: str) -> str: + """Resolve how ``chip`` is recovered, reporting profile problems cleanly. + + ``recovery_mode`` deliberately lets a bad variant or malformed profile + propagate rather than defaulting to UART, so this turns that into a + message rather than a traceback. + """ + from defib.profiles.loader import recovery_mode + + try: + return recovery_mode(chip) + except ValueError as e: + _usb_fail(output, str(e)) + raise AssertionError("unreachable") from e # pragma: no cover + + +def _usb_progress_printer(output: str) -> Any: + """Progress callback that emits JSON lines, or nothing in other modes.""" + import json as json_mod + + from defib.recovery.events import ProgressEvent + + def on_progress(event: ProgressEvent) -> None: + if output == "json": + print(json_mod.dumps({ + "event": "progress", "stage": event.stage.value, + "sent": event.bytes_sent, "total": event.bytes_total, + "percent": round(event.percent, 1), + }), flush=True) + + return on_progress + + +def _resolve_usb_loader(chip: str, ddr: str, usbplug: str, loader: str) -> Any: + """Build the MaskROM loader blobs for a USB-recovery chip. + + ``--loader`` takes an RKBOOT container; ``--ddr``/``--usbplug`` take the + two bare images. Both forms exist because the blobs Rockchip publishes for + RV1106 carry no container header — which is exactly what ``rkdeveloptool + db`` refuses to load. + + Neither is bundled: they are vendor binaries, so the profile records only + the filenames it expects and the user supplies them. + """ + from pathlib import Path + + from defib.profiles.loader import load_profile + from defib.rockusb.loader import parse_loader, raw_blobs + + if loader: + return parse_loader(Path(loader).read_bytes()) + if ddr and usbplug: + return raw_blobs(Path(ddr).read_bytes(), Path(usbplug).read_bytes()) + + profile = load_profile(chip) + expected = " and ".join( + n for n in (profile.loader_ddr, profile.loader_usbplug) if n + ) + raise typer.BadParameter( + f"{chip} recovers over USB and needs its vendor loader blobs: pass " + f"--ddr and --usbplug (expected {expected}), or --loader for an " + "RKBOOT container." + ) + + +async def _open_usb_target( + blobs: Any, power_cycle: bool, output: str, wait: float, + poe_port_override: str = "", usb_path: str = "", + recovery_ids: Any = None, +) -> Any: + """Power-cycle if asked, catch the board, and get its usbplug running. + + Returns a ``RockchipRecovery`` ready to touch flash. + """ + from rich.console import Console + + from defib.rockusb.device import DeviceMode, RockusbDevice, wait_for_device + from defib.rockusb.recovery import RockchipRecovery + + console = Console() + + if power_cycle: + from defib.power.factory import power_controller_from_env + + controller = power_controller_from_env() + if output == "human": + console.print(f" Power-cycling via {controller.name()}...") + try: + await controller.power_cycle(poe_port_override) + finally: + await controller.close() + + found = await wait_for_device( + timeout=wait, usb_path=usb_path or None, recovery_ids=recovery_ids + ) + if output == "human": + console.print(f" Found {found}") + + device = RockusbDevice(found) + device.open() + recovery = RockchipRecovery(device) + + if found.mode is DeviceMode.MASKROM: + if output == "human": + console.print(" Uploading DDR init and usbplug...") + # Pin re-enumeration to the port this board is plugged into, so a + # second board appearing mid-upload cannot be adopted instead. + await recovery.download_boot( + blobs, + on_progress=_usb_progress_printer(output), + usb_path=found.usb_path, + recovery_ids=recovery_ids, + ) + + return recovery + + +def _usb_fail(output: str, message: str) -> None: + """Report a USB-recovery failure and exit non-zero.""" + import json as json_mod + + from rich.console import Console + from rich.markup import escape + + if output == "json": + print(json_mod.dumps({"event": "error", "message": message})) + else: + # Escaped: messages carry file paths and the "defib[rockchip]" install + # hint, either of which Rich would read as markup. + Console().print(f"[red]{escape(message)}[/red]") + raise typer.Exit(1) + + +async def _burn_usb_async( + chip: str, ddr: str, usbplug: str, loader: str, + power_cycle: bool, output: str, wait: float, + poe_port_override: str = "", usb_path: str = "", +) -> None: + """``burn`` for chips whose boot ROM only answers on USB. + + The UART path uploads U-Boot into RAM; the USB equivalent is getting the + usbplug running, after which flash is writable. It stops there — writing + images is ``install``'s job. + """ + import json as json_mod + + from rich.console import Console + + from defib.profiles.loader import load_profile + from defib.rockusb.loader import LoaderFormatError + from defib.rockusb.protocol import RockusbError + + console = Console() + recovery_ids = load_profile(chip).usb_recovery_ids + + # Loader resolution sits inside the handler: an unreadable file, a + # malformed container or missing arguments would otherwise escape as + # Typer text or a traceback, and in JSON mode that means no error event + # at all. + recovery = None + try: + blobs = _resolve_usb_loader(chip, ddr, usbplug, loader) + recovery = await _open_usb_target( + blobs, power_cycle, output, wait, poe_port_override, usb_path, + recovery_ids, + ) + flash_id = await recovery.read_flash_id() + except (RockusbError, LoaderFormatError, OSError, typer.BadParameter) as e: + _usb_fail(output, str(e)) + return + finally: + # Release the interface either way. A claimed handle outliving the + # command leaves the next attempt unable to open the board, which is + # indistinguishable from a device that has stopped responding. + if recovery is not None: + recovery.close() + + if output == "json": + print(json_mod.dumps({ + "event": "done", "success": True, "flash_id": flash_id.hex(), + })) + elif output != "quiet": + console.print(f" Flash ID: {flash_id.hex()}") + console.print( + "\n[green bold]Device is awake[/green bold] and in loader mode, " + "ready for a flashing tool." + ) + + def main() -> None: app() diff --git a/src/defib/flashdump.py b/src/defib/flashdump.py index 7bae16e..fbeac96 100644 --- a/src/defib/flashdump.py +++ b/src/defib/flashdump.py @@ -95,7 +95,10 @@ def get_ram_staging_addr(chip: str) -> int: try: from defib.profiles.loader import load_profile profile = load_profile(chip_lower) - uboot_addr = int(profile.addresses[2], 16) + # The property raises for USB-recovery chips, which have no load + # addresses at all — the except below then falls through to the + # prefix lookup, which is the right answer for them. + uboot_addr = profile.uboot_address # Derive RAM base from U-Boot address (aligned to 0x40000000 boundary) ram_base = uboot_addr & 0xF0000000 return ram_base + RAM_STAGING_OFFSET diff --git a/src/defib/profiles/data/rv1106.json b/src/defib/profiles/data/rv1106.json new file mode 100644 index 0000000..5cc7b32 --- /dev/null +++ b/src/defib/profiles/data/rv1106.json @@ -0,0 +1,16 @@ +{ + "name": "rv1106", + "RECOVERY": "usb", + "USB_RECOVERY_IDS": [4364], + "LOADER_DDR": "rv1106_ddr_924MHz_v1.15.bin", + "LOADER_USBPLUG": "rv1106_usbplug_v1.09.bin", + "PARTITIONS": { + "env": {"lba": 0, "sectors": 512}, + "idblock": {"lba": 512, "sectors": 512}, + "uboot": {"lba": 1024, "sectors": 1024}, + "boot": {"lba": 2048, "sectors": 8192}, + "oem": {"lba": 10240, "sectors": 61440}, + "userdata": {"lba": 71680, "sectors": 20480}, + "rootfs": {"lba": 92160, "sectors": 430080} + } +} diff --git a/src/defib/profiles/loader.py b/src/defib/profiles/loader.py index eb90ba1..748fe70 100644 --- a/src/defib/profiles/loader.py +++ b/src/defib/profiles/loader.py @@ -107,6 +107,25 @@ def load_profile(chip_name: str, profiles_dir: Path | None = None) -> SoCProfile raise ValueError(f"Alias chain too deep for chip: {chip_name}") +def recovery_mode(chip_name: str, profiles_dir: Path | None = None) -> str: + """How ``chip_name`` is reached when it is dead: ``"uart"`` or ``"usb"``. + + Defaults to ``"uart"`` only when the chip has no profile at all, which + covers the V500 and CV6xx families whose chip lists live in their + protocol modules rather than in JSON. + + Anything else — an unknown variant, an alias loop, malformed JSON, a + profile that fails validation — propagates. Treating those as UART would + quietly route a USB-recovery chip into the serial workflow and report the + wrong problem: ``rv1106:typo`` should say the variant is unknown, not + fail later trying to open a serial port. + """ + try: + return load_profile(chip_name, profiles_dir).recovery + except FileNotFoundError: + return "uart" + + def list_variants(chip_name: str, profiles_dir: Path | None = None) -> list[str]: """Return the list of board variants declared for ``chip_name``. diff --git a/src/defib/profiles/schema.py b/src/defib/profiles/schema.py index 8fff775..5240cf6 100644 --- a/src/defib/profiles/schema.py +++ b/src/defib/profiles/schema.py @@ -2,16 +2,50 @@ from __future__ import annotations -from pydantic import BaseModel, Field, PrivateAttr +from typing import Literal + +from pydantic import BaseModel, Field, PrivateAttr, model_validator + + +class FlashPartition(BaseModel): + """Where a partition lives, in 512-byte sectors. + + The size matters as much as the offset: without it an oversized image + would be written straight through into whatever follows. + """ + + lba: int = Field(description="Starting sector") + sectors: int = Field(description="Length in sectors") + + @property + def end_lba(self) -> int: + """First sector past this partition.""" + return self.lba + self.sectors + + @property + def size_bytes(self) -> int: + return self.sectors * 512 class SoCProfile(BaseModel): - """A SoC configuration profile for the standard HiSilicon protocol. + """A SoC configuration profile. - Fields match the JSON profile format used by the original burn tool. + Field names match the JSON profile format used by the original burn tool, + which was HiSilicon-only. ``RECOVERY`` widens that: chips whose boot ROM + has no UART download path at all — Rockchip's, which is USB-only — declare + ``"RECOVERY": "usb"`` and supply none of the DDR/SPL bytecode below. """ name: str = Field(description="Internal chip name") + recovery: Literal["uart", "usb"] = Field( + default="uart", alias="RECOVERY", + description=( + "How a dead board is reached. 'uart' drives the boot ROM over " + "serial with the bytecode in this profile. 'usb' means the boot " + "ROM only answers on USB, so the fields below do not apply and " + "the loader blobs are named by LOADER_* instead." + ), + ) prestep0: list[int] | None = Field( default=None, alias="PRESTEP0", description="Pre-DDR init bytecode (sent before DDRSTEP0)", @@ -20,17 +54,19 @@ class SoCProfile(BaseModel): default=None, alias="PRESTEP1", description="DDR training verification bytecode (sent after DDRSTEP0)", ) - ddrstep0: list[int] = Field(alias="DDRSTEP0", description="DDR initialization bytecode") - addresses: list[str] = Field( - alias="ADDRESS", + ddrstep0: list[int] | None = Field( + default=None, alias="DDRSTEP0", description="DDR initialization bytecode" + ) + addresses: list[str] | None = Field( + default=None, alias="ADDRESS", description="Load addresses: [ddr_step, spl, uboot]", ) - file_lengths: list[str] = Field( - alias="FILELEN", + file_lengths: list[str] | None = Field( + default=None, alias="FILELEN", description="Size limits: [ddr_step_max, spl_max]", ) - step_lengths: list[str] = Field( - alias="STEPLEN", + step_lengths: list[str] | None = Field( + default=None, alias="STEPLEN", description="Step frame sizes: [ddr_step, spl]", ) sram_limit: str | None = Field( @@ -56,10 +92,76 @@ class SoCProfile(BaseModel): "`spl_data` property." ), ) + usb_recovery_ids: list[int] = Field( + default_factory=list, alias="USB_RECOVERY_IDS", + description=( + "USB recovery only. Product ids the SoC presents while it is " + "actually in a recovery mode. Needed because a running board can " + "expose an unrelated gadget under the same vendor id — a Luckfox " + "booted into Linux shows 2207:0019 (RNDIS+ADB), which must never " + "be mistaken for something waiting to be flashed." + ), + ) + loader_ddr: str | None = Field( + default=None, alias="LOADER_DDR", + description=( + "USB recovery only. Filename of the DDR-init blob the boot ROM " + "expects first (rkbin's rv1106_ddr_*.bin). Not bundled — it is a " + "vendor binary the user supplies." + ), + ) + loader_usbplug: str | None = Field( + default=None, alias="LOADER_USBPLUG", + description=( + "USB recovery only. Filename of the second-stage blob that takes " + "over USB and exposes flash (rkbin's rv1106_usbplug_*.bin)." + ), + ) + partitions: dict[str, FlashPartition] = Field( + default_factory=dict, alias="PARTITIONS", + description=( + "USB recovery only. Partition name to its extent, used to place " + "firmware images without the caller computing offsets and to " + "reject images that would not fit." + ), + ) + # Bytes of the SPL_BLOB, populated by the loader (not in JSON, not # validated by pydantic). None if `spl_blob` is unset. _spl_data: bytes | None = PrivateAttr(default=None) + @model_validator(mode="after") + def _check_recovery_fields(self) -> SoCProfile: + """Each recovery family must carry its own required fields. + + Without this the UART fields could quietly go missing on a UART chip + and only surface as a confusing ``NoneType`` deep inside a burn. + """ + if self.recovery == "uart": + missing = [ + alias + for alias, value in ( + ("DDRSTEP0", self.ddrstep0), + ("ADDRESS", self.addresses), + ("FILELEN", self.file_lengths), + ("STEPLEN", self.step_lengths), + ) + if value is None + ] + if missing: + raise ValueError( + f"profile '{self.name}' uses UART recovery but is missing " + f"{', '.join(missing)}" + ) + return self + + def _require_uart(self, field: str) -> None: + if self.recovery != "uart": + raise ValueError( + f"'{self.name}' recovers over {self.recovery}, not UART — " + f"{field} is not defined for it" + ) + @property def spl_data(self) -> bytes | None: """Pre-built SPL bytes if the profile declares an `SPL_BLOB`.""" @@ -67,18 +169,26 @@ def spl_data(self) -> bytes | None: @property def ddr_step_address(self) -> int: + self._require_uart("ddr_step_address") + assert self.addresses is not None return int(self.addresses[0], 16) @property def spl_address(self) -> int: + self._require_uart("spl_address") + assert self.addresses is not None return int(self.addresses[1], 16) @property def uboot_address(self) -> int: + self._require_uart("uboot_address") + assert self.addresses is not None return int(self.addresses[2], 16) @property def spl_max_size(self) -> int: + self._require_uart("spl_max_size") + assert self.file_lengths is not None return int(self.file_lengths[1], 16) @property @@ -89,6 +199,8 @@ def spl_sram_limit(self) -> int | None: @property def ddr_step_data(self) -> bytes: + self._require_uart("ddr_step_data") + assert self.ddrstep0 is not None return bytes(self.ddrstep0) @property diff --git a/src/defib/recovery/events.py b/src/defib/recovery/events.py index 8fd30f8..f427acd 100644 --- a/src/defib/recovery/events.py +++ b/src/defib/recovery/events.py @@ -19,6 +19,10 @@ class Stage(str, Enum): AUX_AREA = "aux_area" BOOT_IMAGE = "boot_image" BOARD_ID = "board_id" + # Rockchip USB recovery: the usbplug upload that follows DDR_INIT, and the + # block writes it enables. + USBPLUG = "usbplug" + FLASH_WRITE = "flash_write" COMPLETE = "complete" diff --git a/src/defib/rockusb/__init__.py b/src/defib/rockusb/__init__.py new file mode 100644 index 0000000..d779849 --- /dev/null +++ b/src/defib/rockusb/__init__.py @@ -0,0 +1,59 @@ +"""Rockchip MaskROM / rockusb recovery over USB. + +This is a separate subsystem from :mod:`defib.protocol`, and deliberately so. +Every protocol in that package is a UART boot-ROM dialect layered on the +byte-stream :class:`~defib.transport.base.Transport`. Rockchip's boot ROM has +no UART download path at all — recovery is USB-only, and splits into two +stages that are not byte streams: + +1. **MaskROM** — vendor control transfers (``bRequest=0x0C``) push a DDR init + blob and then a "usbplug" blob into SRAM. See :mod:`.maskrom`. +2. **rockusb** — the usbplug then speaks a USB Mass-Storage-shaped bulk + protocol (CBW/CSW) with Rockchip opcodes. See :mod:`.protocol`. + +Because neither stage fits ``Transport``, none of this registers with the +``defib.protocols`` entry-point group. + +The payoff for a test rig: when SPI NAND holds no valid IDB — an erased or +half-written flash — the boot ROM falls into MaskROM *by itself* at power-up, +with no button press and no strap. Power-cycling the board is enough to make +it recoverable, which is what makes unattended recovery possible at all. + +Protocol details were derived from the MIT-licensed xboot/xrock and the +BSD-2-licensed rkflashtool. rkdeveloptool is GPL-2 and was **not** used as a +source for this implementation. +""" + +from __future__ import annotations + +from defib.rockusb.codec import RK_RC4_KEY, rc4, rk_crc16 +from defib.rockusb.loader import LoaderBlobs, LoaderFormatError, parse_loader +from defib.rockusb.maskrom import CODE_471, CODE_472, build_maskrom_chunks +from defib.rockusb.protocol import ( + SECTOR_SIZE, + CommandStatus, + Opcode, + ResetSubcode, + RockusbError, + build_cbw, + parse_csw, +) + +__all__ = [ + "CODE_471", + "CODE_472", + "RK_RC4_KEY", + "SECTOR_SIZE", + "CommandStatus", + "LoaderBlobs", + "LoaderFormatError", + "Opcode", + "ResetSubcode", + "RockusbError", + "build_cbw", + "build_maskrom_chunks", + "parse_csw", + "parse_loader", + "rc4", + "rk_crc16", +] diff --git a/src/defib/rockusb/codec.py b/src/defib/rockusb/codec.py new file mode 100644 index 0000000..50417ad --- /dev/null +++ b/src/defib/rockusb/codec.py @@ -0,0 +1,58 @@ +"""Wire codecs for the Rockchip MaskROM stage: CRC-16 and RC4. + +Both are needed only by :mod:`defib.rockusb.maskrom` — the rockusb bulk stage +carries no checksum of its own (USB already provides one) and is never +encrypted. +""" + +from __future__ import annotations + +from defib.protocol.crc import CRC_TABLE + +# Rockchip's boot ROM fixes this key for the MaskROM code path. Published in +# rkflashtool's README.maskrom as the argument to ``openssl rc4 -K``, and +# identical to the key constant in xboot/xrock. +RK_RC4_KEY = bytes( + [124, 78, 3, 4, 85, 5, 9, 7, 45, 44, 123, 56, 23, 13, 23, 17] +) + +# CRC-16/CCITT seed used by the MaskROM loader. Note this is *not* the same +# variant as :func:`defib.protocol.crc.calc_crc`, which HiSilicon seeds at 0 +# and finalises with two zero bytes. Only the polynomial table is shared. +CRC16_INIT = 0xFFFF + + +def rk_crc16(data: bytes | bytearray, crc: int = CRC16_INIT) -> int: + """CRC-16/CCITT as the Rockchip boot ROM computes it. + + Poly 0x1021, seeded 0xFFFF, MSB-first, no final XOR — matching + ``rkcrc16()`` in rkflashtool's BSD-2 ``rkcrc.h``:: + + crc = (crc << 8) ^ crc16table[(crc >> 8) ^ *buf++]; + """ + for byte in data: + crc = ((crc << 8) & 0xFFFF) ^ CRC_TABLE[((crc >> 8) ^ byte) & 0xFF] + return crc & 0xFFFF + + +def rc4(data: bytes | bytearray, key: bytes = RK_RC4_KEY) -> bytes: + """Plain RC4. Self-inverse, so this both encrypts and decrypts. + + Only used when a loader's header does *not* set the "RC4 disabled" flag. + Newer parts — RV1106 among them — ship loaders built with RC4 off, so this + path is normally dead for that SoC. + """ + s = list(range(256)) + j = 0 + for i in range(256): + j = (j + s[i] + key[i % len(key)]) & 0xFF + s[i], s[j] = s[j], s[i] + + out = bytearray(len(data)) + i = j = 0 + for n, byte in enumerate(data): + i = (i + 1) & 0xFF + j = (j + s[i]) & 0xFF + s[i], s[j] = s[j], s[i] + out[n] = byte ^ s[(s[i] + s[j]) & 0xFF] + return bytes(out) diff --git a/src/defib/rockusb/device.py b/src/defib/rockusb/device.py new file mode 100644 index 0000000..df1a3e3 --- /dev/null +++ b/src/defib/rockusb/device.py @@ -0,0 +1,578 @@ +"""libusb plumbing for the two Rockchip USB stages. + +Kept deliberately thin: everything that can be decided without a device on the +bus lives in :mod:`.maskrom`, :mod:`.protocol` and :mod:`.loader`. What is +left here is enumeration, endpoint discovery and the blocking transfers, all +of which need real hardware. + +``pyusb`` is an optional dependency — install the ``rockchip`` extra. It is +imported lazily so that a defib install without it keeps working for every +UART-based SoC. +""" + +from __future__ import annotations + +import asyncio +import logging +import secrets +from collections.abc import Sequence +from dataclasses import dataclass +from enum import Enum +from typing import Any + +from defib.rockusb.protocol import ( + CSW_LENGTH, + CommandStatus, + Opcode, + RockusbError, + build_cbw, + parse_csw, + residue_is_meaningful, +) + +logger = logging.getLogger(__name__) + +ROCKCHIP_VID = 0x2207 + +# The rockusb interface advertises itself as vendor-specific. MaskROM exposes +# a barer descriptor, so interface matching is best-effort and we fall back to +# scanning every interface for a bulk pair. +ROCKUSB_CLASS = 0xFF +ROCKUSB_SUBCLASS = 0x06 +ROCKUSB_PROTOCOL = 0x05 + + +class DeviceMode(str, Enum): + """Which of the two stages the device is currently answering.""" + + MASKROM = "maskrom" + LOADER = "loader" + + +class RockusbUsbError(RockusbError): + """USB-level failure: no device, cannot claim, transfer error.""" + + +def _require_usb() -> Any: + try: + import usb.core # noqa: PLC0415 + import usb.util # noqa: PLC0415 + except ImportError as e: # pragma: no cover - depends on install extras + raise RockusbUsbError( + "pyusb is required for Rockchip USB recovery — " + "install it with: pip install 'defib[rockchip]'" + ) from e + return usb + + +@dataclass +class FoundDevice: + """A Rockchip device on the bus, and which stage it is in.""" + + mode: DeviceMode + bus: int + address: int + product_id: int + handle: Any # usb.core.Device + port_numbers: tuple[int, ...] = () + + @property + def usb_path(self) -> str: + """Physical topology path, e.g. ``1-4.2``. + + This is the only stable identity a Rockchip board has across the + MaskROM-to-loader transition: the USB address is reassigned when the + usbplug re-enumerates, and every board shares one VID:PID. The port + path is where it is plugged in, so it survives. + """ + if not self.port_numbers: + return f"{self.bus}-?" + return f"{self.bus}-{'.'.join(str(p) for p in self.port_numbers)}" + + def __str__(self) -> str: + return ( + f"{self.mode.value} device {ROCKCHIP_VID:04x}:{self.product_id:04x} " + f"at {self.usb_path} (bus {self.bus} addr {self.address})" + ) + + +def _classify(dev: Any) -> DeviceMode: + """MaskROM or loader? + + Both stages enumerate under the same VID:PID — on RV1106, ``2207:110c`` — + so the product id cannot separate them. + + xrock uses the low bit of ``bcdUSB`` for this. Measured on an RV1106 that + does not hold: the boot ROM and the running usbplug both report 0x0200, + so that test calls the usbplug MaskROM and the caller waits out a + re-enumeration that already happened. + + What does differ is the string descriptors. The boot ROM ships a bare + descriptor with no strings at all; the usbplug names itself:: + + MaskROM iManufacturer=0 iProduct=0 + usbplug iManufacturer=1 "RockChip" iProduct=2 "USB-MSC" + + Only the index is read here, which lives in the device descriptor — no + string fetch, so this stays cheap and cannot fail on a device that + refuses string reads. + """ + named = bool(dev.iProduct) or bool(dev.iManufacturer) + return DeviceMode.LOADER if named else DeviceMode.MASKROM + + +def find_devices( + usb_path: str | None = None, + recovery_ids: Sequence[int] | None = None, +) -> list[FoundDevice]: + """Every Rockchip device in a recovery mode, optionally pinned to a port. + + ``recovery_ids`` restricts the search to product ids that mean "waiting to + be flashed". Without it the vendor id alone is far too broad: a Luckfox + booted into Linux presents 2207:0019, an RNDIS+ADB gadget that shares the + vendor id and — because its bcdUSB is 0x0200 — even looks like MaskROM to + the bcdUSB test. Uploading a loader into a healthy running board is not a + mistake worth leaving reachable. + """ + usb = _require_usb() + + out: list[FoundDevice] = [] + for dev in usb.core.find(idVendor=ROCKCHIP_VID, find_all=True): + if recovery_ids and dev.idProduct not in recovery_ids: + logger.debug( + "ignoring %04x:%04x — not a recovery product id", + ROCKCHIP_VID, dev.idProduct, + ) + continue + found = FoundDevice( + mode=_classify(dev), + bus=dev.bus, + address=dev.address, + product_id=dev.idProduct, + handle=dev, + port_numbers=tuple(getattr(dev, "port_numbers", None) or ()), + ) + if usb_path is None or found.usb_path == usb_path: + out.append(found) + return out + + +def find_device( + usb_path: str | None = None, + recovery_ids: Sequence[int] | None = None, +) -> FoundDevice | None: + """A single Rockchip device in a recovery mode, or None. + + Raises: + RockusbUsbError: if more than one is present and no ``usb_path`` + narrows it down. Picking one arbitrarily would mean flashing + whichever board happened to enumerate first, which on a rack is + how you write firmware to the wrong device. + """ + devices = find_devices(usb_path, recovery_ids) + if not devices: + return None + if len(devices) > 1: + paths = ", ".join(sorted(d.usb_path for d in devices)) + raise RockusbUsbError( + f"{len(devices)} Rockchip devices present ({paths}) — " + "narrow it down with --usb-path" + ) + return devices[0] + + +async def wait_for_device( + timeout: float = 30.0, + mode: DeviceMode | None = None, + poll_interval: float = 0.25, + usb_path: str | None = None, + recovery_ids: Sequence[int] | None = None, +) -> FoundDevice: + """Poll until a matching device appears. + + Used both for the initial "power-cycle an erased board and catch it in + MaskROM" step and for the re-enumeration that follows the usbplug upload. + Pass ``usb_path`` for the latter: the board keeps its port path across + re-enumeration, so pinning it stops the session hopping to a different + board that happened to appear meanwhile. + + Args: + timeout: seconds to keep looking. + mode: require this stage; None accepts either. + poll_interval: seconds between scans. + usb_path: restrict to this physical port path (e.g. ``1-4.2``). + + Raises: + RockusbUsbError: if nothing matching shows up in time, or if several + candidates are present and ``usb_path`` does not disambiguate. + """ + loop = asyncio.get_event_loop() + deadline = loop.time() + timeout + seen: str | None = None + + while loop.time() < deadline: + found = await asyncio.to_thread(find_device, usb_path, recovery_ids) + if found is not None: + if mode is None or found.mode is mode: + return found + seen = str(found) + await asyncio.sleep(poll_interval) + + want = f" in {mode.value} mode" if mode else "" + where = f" at {usb_path}" if usb_path else "" + if seen: + raise RockusbUsbError( + f"no Rockchip device{want}{where} after {timeout:.0f}s — " + f"saw {seen} instead" + ) + raise RockusbUsbError( + f"no Rockchip device{want}{where} after {timeout:.0f}s. " + "Power-cycle the board; an erased flash enters MaskROM on its own." + ) + + +class RockusbDevice: + """An opened Rockchip device: control transfers and the bulk command loop. + + Call :meth:`open` before use and :meth:`close` after, or use it as a + context manager. + """ + + def __init__(self, found: FoundDevice, timeout_ms: int = 5000) -> None: + self._found = found + self._dev = found.handle + self._timeout_ms = timeout_ms + self._interface: Any = None + self._ep_in: Any = None + self._ep_out: Any = None + self._detached_interface: int | None = None + + @property + def mode(self) -> DeviceMode: + return self._found.mode + + def __enter__(self) -> RockusbDevice: + self.open() + return self + + def __exit__(self, *exc: object) -> None: + self.close() + + def open(self) -> None: + """Claim the interface and locate the bulk endpoints. + + Endpoints are read from the descriptor rather than hardcoded; the + conventional 0x02/0x81 pair is not guaranteed across usbplug builds. + """ + usb = _require_usb() + dev = self._dev + + # Only configure a device that is not configured already. Issuing + # SET_CONFIGURATION to a running usbplug resets the endpoint data + # toggles underneath it, after which the next read comes back + # [Errno 75] Overflow — the device answers out of step with what the + # host asked for. Costly to diagnose, trivial to avoid. + try: + cfg = dev.get_active_configuration() + except usb.core.USBError: + try: + dev.set_configuration() + except usb.core.USBError as e: + raise RockusbUsbError( + f"cannot configure {self._found}: {e}" + ) from e + cfg = dev.get_active_configuration() + for intf in cfg: + ep_out = usb.util.find_descriptor( + intf, + custom_match=lambda e: ( + usb.util.endpoint_direction(e.bEndpointAddress) + == usb.util.ENDPOINT_OUT + and usb.util.endpoint_type(e.bmAttributes) + == usb.util.ENDPOINT_TYPE_BULK + ), + ) + ep_in = usb.util.find_descriptor( + intf, + custom_match=lambda e: ( + usb.util.endpoint_direction(e.bEndpointAddress) + == usb.util.ENDPOINT_IN + and usb.util.endpoint_type(e.bmAttributes) + == usb.util.ENDPOINT_TYPE_BULK + ), + ) + if ep_in is not None and ep_out is not None: + self._interface, self._ep_in, self._ep_out = intf, ep_in, ep_out + break + + if self._interface is None: + # MaskROM legitimately has no bulk pair — control transfers only. + if self._found.mode is DeviceMode.MASKROM: + logger.debug("%s: no bulk endpoints, MaskROM control-only", self._found) + return + raise RockusbUsbError( + f"{self._found}: no bulk IN/OUT endpoint pair found" + ) + + # Detach from the interface actually being claimed, not interface 0 — + # the bulk pair is not guaranteed to live on the first one, and + # detaching the wrong interface leaves the right one bound and the + # claim below failing. + number = self._interface.bInterfaceNumber + try: + if dev.is_kernel_driver_active(number): + dev.detach_kernel_driver(number) + self._detached_interface = number + except (NotImplementedError, usb.core.USBError): + # Not all backends/platforms implement this; only Linux binds a + # kernel driver here in the first place. + pass + + try: + usb.util.claim_interface(dev, number) + self._clear_stalls() + self._drain_stale_input() + except usb.core.USBError as e: + # Give the kernel driver back before bailing out. Leaving it + # detached would strand the interface for whatever owned it, + # turning a failed attempt into a lasting one. + self._reattach_kernel_driver() + raise RockusbUsbError( + f"cannot claim interface {number} on {self._found}: {e} " + "(need a udev rule for 2207:* or root)" + ) from e + + def _reattach_kernel_driver(self) -> None: + """Undo a detach, if we did one. Best-effort by nature.""" + if self._detached_interface is None: + return + try: + self._dev.attach_kernel_driver(self._detached_interface) + except Exception: # pragma: no cover - platform dependent + logger.debug( + "could not reattach kernel driver to interface %d", + self._detached_interface, exc_info=True, + ) + finally: + self._detached_interface = None + + def close(self) -> None: + usb = _require_usb() + try: + if self._interface is not None: + usb.util.release_interface(self._dev, self._interface.bInterfaceNumber) + usb.util.dispose_resources(self._dev) + except Exception: # pragma: no cover - teardown is best-effort + logger.debug("cleanup failed for %s", self._found, exc_info=True) + self._reattach_kernel_driver() + + # -- MaskROM stage ---------------------------------------------------- + + def control_write(self, code: int, payload: bytes) -> int: + """One MaskROM code-upload control transfer. + + Returns the number of bytes the host actually pushed. + + Raises: + RockusbUsbError: on a transfer error, or a short write. A + truncated loader upload surfaces much later as a + re-enumeration timeout that looks like a dead board, so it is + caught where it happens. + """ + try: + written = int(self._dev.ctrl_transfer( + bmRequestType=0x40, + bRequest=0x0C, + wValue=0x0000, + wIndex=code, + data_or_wLength=payload, + timeout=self._timeout_ms, + )) + if written != len(payload): + raise RockusbUsbError( + f"MaskROM control transfer short (code {code:#06x}): " + f"wrote {written} of {len(payload)} bytes" + ) + return written + except RockusbUsbError: + raise + except Exception as e: + raise RockusbUsbError( + f"MaskROM control transfer failed (code {code:#06x}, " + f"{len(payload)} bytes): {e}" + ) from e + + # -- rockusb bulk stage ----------------------------------------------- + + def _clear_stalls(self) -> None: + """Clear a halt condition left on either endpoint. + + The other half of inheriting a device from an abandoned attempt: a + transfer that failed part-way can leave an endpoint halted, and every + subsequent transfer on it then times out — including the very first + command wrapper, which makes the device look dead rather than merely + out of step. + + Best-effort. A device that has genuinely gone away will fail here and + that is fine; the caller finds out on the next real transfer. + """ + for ep in (self._ep_out, self._ep_in): + if ep is None: + continue + try: + ep.clear_halt() + except Exception: + logger.debug( + "could not clear halt on endpoint %#04x", + ep.bEndpointAddress, exc_info=True, + ) + + def _drain_stale_input(self) -> None: + """Discard anything the device still has queued from a previous run. + + A recovery tool is routinely pointed at a device some earlier attempt + abandoned mid-transaction, leaving an unread status wrapper in the + pipe. The next command then reads that stale wrapper as its data + phase — a 13-byte CSW landing in a 5-byte buffer surfaces as + ``[Errno 75] Overflow`` and every command after it desynchronises. + + The serial transports already open by reading until the line goes + quiet; this is the same idea one layer down. + """ + if self._ep_in is None: + return + mps = getattr(self._ep_in, "wMaxPacketSize", 512) or 512 + for _ in range(8): + try: + stale = self._ep_in.read(mps, 50) + except Exception: + return # nothing waiting: a timeout here is the good outcome + if not stale: + return + logger.debug("discarded %d stale bytes from a previous session", len(stale)) + + def _read_bulk(self, length: int) -> bytes: + """Read ``length`` bytes from the bulk IN endpoint. + + The request is rounded up to the endpoint's max packet size and the + result trimmed. A bulk IN transfer whose buffer is not a multiple of + that size fails with ``[Errno 75] Overflow`` the moment the device + answers with a full packet — which the usbplug does for short replies + like READ_FLASH_ID's five bytes. + """ + mps = getattr(self._ep_in, "wMaxPacketSize", 512) or 512 + rounded = ((length + mps - 1) // mps) * mps + data = bytes(self._ep_in.read(rounded, self._timeout_ms)) + return data[:length] + + def _write_bulk(self, data: bytes, what: str, opcode: Opcode | int) -> None: + """Push ``data`` out the bulk endpoint, insisting all of it lands. + + pyusb returns how many bytes it managed; a short count means the + device took less than we handed it, which for a payload write is a + partial flash write wearing the costume of a successful one. + """ + try: + written = int(self._ep_out.write(data, self._timeout_ms)) + except Exception as e: + raise RockusbUsbError( + f"rockusb {what} write failed (opcode {int(opcode):#04x}): {e}" + ) from e + if written != len(data): + raise RockusbUsbError( + f"rockusb {what} write short (opcode {int(opcode):#04x}): " + f"wrote {written} of {len(data)} bytes" + ) + + def _require_bulk(self) -> None: + if self._ep_in is None or self._ep_out is None: + raise RockusbUsbError( + f"{self._found}: bulk commands need the usbplug running — " + "device is still in MaskROM, send the loader first" + ) + + def command( + self, + opcode: Opcode | int, + *, + subcode: int = 0, + address: int = 0, + count: int = 0, + data_out: bytes | None = None, + read_length: int = 0, + tolerate_disconnect: bool = False, + ) -> bytes: + """Run one CBW / optional data phase / CSW exchange. + + Returns any data read during an IN transfer, otherwise ``b""``. + + Raises: + RockusbUsbError: on a transfer error, a non-zero status, or a + short transfer. A device is entitled to move less than it was + asked to and still report OK, reporting the shortfall in the + status wrapper's residue; treating that as success is how a + partial write gets announced as a finished install. + """ + self._require_bulk() + tag = secrets.randbits(32) + direction_in = read_length > 0 + transfer_length = read_length if direction_in else len(data_out or b"") + + cbw = build_cbw( + tag, + opcode, + subcode=subcode, + address=address, + count=count, + transfer_length=transfer_length, + direction_in=direction_in, + ) + + # The command itself must always get out intact. Only the status + # wrapper may legitimately go missing, and only when the caller is + # expecting the device to drop off the bus — see `tolerate_disconnect`. + self._write_bulk(cbw, "command wrapper", opcode) + try: + payload = b"" + if direction_in: + payload = self._read_bulk(read_length) + elif data_out: + self._write_bulk(data_out, "payload", opcode) + except RockusbUsbError: + raise + except Exception as e: + raise RockusbUsbError( + f"rockusb data phase failed (opcode {int(opcode):#04x}): {e}" + ) from e + + try: + csw = bytes(self._ep_in.read(CSW_LENGTH, self._timeout_ms)) + except Exception as e: + if tolerate_disconnect: + logger.debug( + "no status wrapper for opcode %#04x — device already gone: %s", + int(opcode), e, + ) + return b"" + raise RockusbUsbError( + f"rockusb status read failed (opcode {int(opcode):#04x}): {e}" + ) from e + + _, residue, status = parse_csw(csw, expected_tag=tag) + if status != CommandStatus.OK: + raise RockusbUsbError( + f"rockusb command {int(opcode):#04x} failed " + f"(status {status}, residue {residue})" + ) + if residue and residue_is_meaningful(opcode): + raise RockusbUsbError( + f"rockusb command {int(opcode):#04x} moved " + f"{transfer_length - residue} of {transfer_length} bytes " + f"(residue {residue})" + ) + if direction_in and len(payload) != read_length: + raise RockusbUsbError( + f"rockusb command {int(opcode):#04x} returned " + f"{len(payload)} of {read_length} bytes" + ) + return payload diff --git a/src/defib/rockusb/loader.py b/src/defib/rockusb/loader.py new file mode 100644 index 0000000..64e0959 --- /dev/null +++ b/src/defib/rockusb/loader.py @@ -0,0 +1,175 @@ +"""Parse a Rockchip loader (``MiniLoaderAll.bin``) into its two SRAM blobs. + +The container is an ``RKBOOT`` archive: a fixed header, then tables of entries +pointing at blobs elsewhere in the same file. We only care about the 471 +(DDR init) and 472 (usbplug) tables. + +Not every loader Rockchip publishes has this container. The standalone +``rv1106_ddr_*.bin`` / ``rv1106_usbplug_*.bin`` files in rkbin are bare images +with no ``RKBOOT`` header at all, which is exactly why ``rkdeveloptool db`` +rejects them (rockchip-linux/rkdeveloptool#105). Those are handled by +:func:`raw_blobs`, which skips the container entirely — the same escape hatch +``xrock maskrom `` uses. +""" + +from __future__ import annotations + +import struct +from dataclasses import dataclass, field + +TAG_BOOT = b"BOOT" +TAG_LOADER = b"LDR " + +ENTRY_471 = 1 +ENTRY_472 = 2 +ENTRY_LOADER = 4 + +# dwDataOffset, dwDataSize, dwDataDelay — the last 12 bytes of every entry. +_ENTRY_TRAILER = 12 +_ENTRY_NAME_BYTES = 40 # WCHAR szName[20], UTF-16LE + + +class LoaderFormatError(Exception): + """The file is not a loader we know how to read.""" + + +@dataclass(frozen=True) +class LoaderEntry: + """One blob referenced by the loader's entry table.""" + + name: str + data: bytes + delay_ms: int + + +@dataclass(frozen=True) +class LoaderBlobs: + """What the MaskROM stage needs out of a loader file.""" + + ddr: list[LoaderEntry] = field(default_factory=list) + usbplug: list[LoaderEntry] = field(default_factory=list) + rc4_disabled: bool = True + + @property + def use_rc4(self) -> bool: + """Whether blobs must be RC4'd before upload.""" + return not self.rc4_disabled + + +def raw_blobs(ddr: bytes, usbplug: bytes, *, use_rc4: bool = False) -> LoaderBlobs: + """Wrap two bare images as if they had come out of a container. + + For RV1106 this is the normal path, since rkbin ships its DDR and usbplug + images headerless. + + Raises: + LoaderFormatError: if either image is empty — uploading nothing would + look like a board that never came back. + """ + empty = [n for n, b in (("ddr", ddr), ("usbplug", usbplug)) if not b] + if empty: + raise LoaderFormatError(f"{' and '.join(empty)} image is empty") + return LoaderBlobs( + ddr=[LoaderEntry(name="ddr", data=ddr, delay_ms=0)], + usbplug=[LoaderEntry(name="usbplug", data=usbplug, delay_ms=0)], + rc4_disabled=not use_rc4, + ) + + +def _parse_entries( + data: bytes, offset: int, entry_size: int, count: int +) -> list[LoaderEntry]: + """Read ``count`` entries of ``entry_size`` bytes starting at ``offset``. + + The three fields we need sit at the *end* of each entry, so they are read + backwards from the entry stride rather than forwards from its start. That + sidesteps the one genuinely ambiguous field in the struct — ``emType`` is a + C enum, so whether it occupies 1 byte or 4 depends on how the producing + toolchain packed it, and getting it wrong would silently shift every + subsequent offset. + """ + if entry_size < _ENTRY_TRAILER + _ENTRY_NAME_BYTES: + raise LoaderFormatError( + f"entry size {entry_size} too small to hold a name plus offsets" + ) + + entries: list[LoaderEntry] = [] + for i in range(count): + base = offset + i * entry_size + if base + entry_size > len(data): + raise LoaderFormatError( + f"entry {i} at {base:#x} runs past end of file ({len(data)} bytes)" + ) + + name_at = base + entry_size - _ENTRY_TRAILER - _ENTRY_NAME_BYTES + name = ( + data[name_at : name_at + _ENTRY_NAME_BYTES] + .decode("utf-16-le", errors="replace") + .rstrip("\x00") + ) + data_offset, data_size, delay_ms = struct.unpack_from( + " len(data): + raise LoaderFormatError( + f"entry {name!r} blob at {data_offset:#x}+{data_size} " + f"runs past end of file ({len(data)} bytes)" + ) + entries.append( + LoaderEntry( + name=name, + data=data[data_offset : data_offset + data_size], + delay_ms=delay_ms, + ) + ) + return entries + + +def parse_loader(data: bytes) -> LoaderBlobs: + """Parse an ``RKBOOT`` container. + + Raises: + LoaderFormatError: if the magic is wrong or the tables do not fit. A + headerless rkbin image lands here — use :func:`raw_blobs` for those. + """ + if len(data) < 64: + raise LoaderFormatError(f"file too short to be a loader ({len(data)} bytes)") + if data[:4] not in (TAG_BOOT, TAG_LOADER): + raise LoaderFormatError( + f"expected {TAG_BOOT!r} or {TAG_LOADER!r} magic, got {data[:4]!r} — " + "a bare rkbin image? those have no container, pass them to raw_blobs()" + ) + + # Header layout up to the entry tables: + # uiTag[4] usSize[2] dwVersion[4] dwMergeVersion[4] + # STRUCT_RKTIME[7] emSupportChip[4] + # then the three (count, offset, size) triples. + table_at = 4 + 2 + 4 + 4 + 7 + 4 + + n471, off471, size471 = struct.unpack_from("device, device recipient) + bRequest = 0x0C + wValue = 0x0000 + wIndex = 0x0471 | 0x0472 + +``0x0471`` carries the DDR init blob, which runs in SRAM, brings up DRAM and +returns to the boot ROM. ``0x0472`` then carries the "usbplug" blob, which +takes over the USB device and starts answering the bulk protocol in +:mod:`defib.rockusb.protocol`. + +This module deliberately holds no USB code — :func:`build_maskrom_chunks` +turns a blob into the exact sequence of control-transfer payloads, so the +framing (which is where the fiddly parts live) is unit-testable with no +hardware and no libusb. +""" + +from __future__ import annotations + +import struct + +from defib.rockusb.codec import rc4, rk_crc16 + +CODE_471 = 0x0471 +CODE_472 = 0x0472 + +# Boot ROM accepts at most this much per control transfer. +CHUNK_SIZE = 4096 + + +def build_maskrom_chunks( + blob: bytes, + *, + use_rc4: bool = False, +) -> list[bytes]: + """Frame ``blob`` into the control-transfer payloads the boot ROM expects. + + The blob is optionally RC4'd, gets a big-endian CRC-16 appended, and is + then split into 4096-byte chunks. Two size-dependent quirks, both copied + from the reference implementations, are folded in here: + + * ``len(blob) % 4096 == 4095`` — a zero byte is appended *before* the CRC + is computed. Without it the two CRC bytes would straddle a chunk + boundary, which the boot ROM does not accept. + * ``len(blob) % 4096 == 4094`` — appending the CRC makes the payload an + exact multiple of the chunk size, leaving no short packet to signal the + end of the transfer. A trailing one-byte chunk is emitted to terminate + it. + + Args: + blob: raw DDR-init or usbplug image. + use_rc4: encrypt before checksumming. Loaders whose header sets the + "RC4 disabled" flag — which includes every RV1106 loader Rockchip + ships — must leave this off. + + Returns: + Payloads to send, in order, each as one control transfer. + """ + payload = bytearray(blob) + + remainder = len(payload) % CHUNK_SIZE + needs_terminator = remainder == CHUNK_SIZE - 2 + if remainder == CHUNK_SIZE - 1: + payload.append(0x00) + + if use_rc4: + # Whole-buffer, matching xrock. rkdeveloptool is widely read as + # encrypting per 4096-byte block instead; the two agree only when the + # blob is under one chunk. Untested here because RV1106 ships RC4-off + # loaders — verify against hardware before trusting this on a part + # that actually needs encryption. + payload = bytearray(rc4(bytes(payload))) + + payload += struct.pack(">H", rk_crc16(payload)) + + chunks = [ + bytes(payload[i : i + CHUNK_SIZE]) + for i in range(0, len(payload), CHUNK_SIZE) + ] + if needs_terminator: + chunks.append(b"\x00") + return chunks diff --git a/src/defib/rockusb/protocol.py b/src/defib/rockusb/protocol.py new file mode 100644 index 0000000..ea434d5 --- /dev/null +++ b/src/defib/rockusb/protocol.py @@ -0,0 +1,220 @@ +"""rockusb bulk stage: CBW/CSW framing and Rockchip opcodes. + +Once the usbplug is running the device speaks a dialect of USB Mass Storage +Bulk-Only Transport: a 31-byte command block wrapper out, an optional data +phase, then a 13-byte status wrapper back. The command descriptor block is +Rockchip's own, and — unlike the wrappers around it — is big-endian. + +As with :mod:`defib.rockusb.maskrom`, this module is pure framing so it can be +tested without libusb. +""" + +from __future__ import annotations + +import struct +from enum import IntEnum + +CBW_SIGNATURE = b"USBC" +CSW_SIGNATURE = b"USBS" + +CBW_LENGTH = 31 +CSW_LENGTH = 13 + +#: The usbplug presents flash as 512-byte logical blocks regardless of the +#: underlying medium; on SPI NAND it runs Rockchip's FTL underneath, so bad +#: blocks, wear levelling and ECC are all handled device-side. +SECTOR_SIZE = 512 + +#: Conservative ceiling for one READ_LBA/WRITE_LBA. The count field is 16-bit +#: so the protocol permits far more, but the usbplug's own buffering is +#: undocumented and the reference tools stay at or below this. +MAX_SECTORS_PER_TRANSFER = 512 + +DIRECTION_IN = 0x80 +DIRECTION_OUT = 0x00 + + +class Opcode(IntEnum): + """Rockchip CDB operation codes (only the ones we actually issue).""" + + TEST_UNIT_READY = 0x00 + READ_FLASH_ID = 0x01 + ERASE_NORMAL = 0x06 + READ_LBA = 0x14 + WRITE_LBA = 0x15 + READ_FLASH_INFO = 0x1A + READ_CHIP_INFO = 0x1B + ERASE_LBA = 0x25 + READ_CAPABILITY = 0xAA + RESET_DEVICE = 0xFF + + +#: Declared CDB length, which is *not* the 16 bytes the field occupies on the +#: wire. Commands carrying an address declare 10; the rest declare 6. Sending +#: 16 makes the usbplug ignore the wrapper outright — the command never lands +#: and the host waits out its timeout with no error to explain it. +CDB_LENGTH_SHORT = 6 +CDB_LENGTH_ADDRESSED = 10 + +_ADDRESSED_OPCODES = frozenset({ + Opcode.READ_LBA, + Opcode.WRITE_LBA, + Opcode.ERASE_LBA, +}) + + +def cdb_length(opcode: Opcode | int) -> int: + """How long this opcode's command block claims to be.""" + return ( + CDB_LENGTH_ADDRESSED + if opcode in _ADDRESSED_OPCODES + else CDB_LENGTH_SHORT + ) + + +def residue_is_meaningful(opcode: Opcode | int) -> bool: + """Whether this opcode's status wrapper reports a usable residue. + + Mass Storage says residue is a little-endian count of bytes *not* + transferred, and on the LBA path an RV1106 usbplug honours that — a + full-sector read comes back with residue 0. + + Everything else reports nonsense. Measured:: + + TEST_UNIT_READY transfer 0 residue 0x06000000 + READ_FLASH_ID transfer 5 residue 0x05000000 + READ_CAPABILITY transfer 8 residue 0x08000000 + READ_FLASH_INFO transfer 11 residue 0x0B000000 + READ_LBA transfer 512 residue 0 + + Those are the transfer lengths written big-endian, i.e. "none of it + arrived" — while the data plainly did arrive and is correct. xrock never + looks at residue at all, which is presumably why nobody noticed. + + So the check is kept exactly where it earns its keep: the block path, the + one place a short transfer means a partially written flash. + """ + return opcode in _ADDRESSED_OPCODES + + +class ResetSubcode(IntEnum): + """Sub-selector for :attr:`Opcode.RESET_DEVICE`.""" + + NORMAL = 0 + RESET_MSC = 1 + POWEROFF = 2 + MASKROM = 3 + DISCONNECT = 4 + + +class CommandStatus(IntEnum): + OK = 0 + FAILED = 1 + + +class RockusbError(Exception): + """A rockusb command failed, or the reply did not frame correctly.""" + + +def build_cbw( + tag: int, + opcode: Opcode | int, + *, + subcode: int = 0, + address: int = 0, + count: int = 0, + transfer_length: int = 0, + direction_in: bool = False, +) -> bytes: + """Build the 31-byte command block wrapper. + + Mind the mixed endianness: the wrapper's ``tag`` and ``length`` are + little-endian, but ``address`` and ``count`` inside the CDB are + big-endian. + + Args: + tag: caller-chosen id, echoed back in the status wrapper. + opcode: Rockchip operation code. + subcode: CDB byte 1 — the reset selector, or the read/write method. + address: starting LBA, for the block opcodes. + count: sector count, for the block opcodes. + transfer_length: bytes in the data phase. Defaults to + ``count * SECTOR_SIZE`` when a count is given. + direction_in: True when the data phase flows device to host. + """ + if transfer_length == 0 and count: + transfer_length = count * SECTOR_SIZE + + cdb = struct.pack( + ">BB I B H 7x", + int(opcode), + subcode, + address, + 0, + count, + ) + assert len(cdb) == 16, f"CDB must be 16 bytes, got {len(cdb)}" + + return ( + CBW_SIGNATURE + + struct.pack(" tuple[int, int, int]: + """Parse the 13-byte command status wrapper. + + Returns: + ``(tag, residue, status)``. + + Raises: + RockusbError: on a short read, a bad signature, or a tag that does not + match ``expected_tag``. A mismatched tag means replies have got out + of step with commands, which is not something to paper over. + """ + if len(data) < CSW_LENGTH: + raise RockusbError( + f"short status wrapper: got {len(data)} bytes, want {CSW_LENGTH}" + ) + if data[:4] != CSW_SIGNATURE: + raise RockusbError( + f"bad status signature {data[:4]!r}, want {CSW_SIGNATURE!r}" + ) + + tag, residue, status = struct.unpack_from(" list[tuple[int, int]]: + """Split a block range into ``(lba, count)`` pairs of at most ``max_sectors``.""" + if total_sectors < 0: + raise ValueError(f"negative sector count: {total_sectors}") + if max_sectors <= 0: + raise ValueError(f"max_sectors must be positive, got {max_sectors}") + + out: list[tuple[int, int]] = [] + lba = start_lba + remaining = total_sectors + while remaining > 0: + count = min(remaining, max_sectors) + out.append((lba, count)) + lba += count + remaining -= count + return out diff --git a/src/defib/rockusb/recovery.py b/src/defib/rockusb/recovery.py new file mode 100644 index 0000000..cc061b6 --- /dev/null +++ b/src/defib/rockusb/recovery.py @@ -0,0 +1,224 @@ +"""End-to-end Rockchip USB recovery. + +The flow this exists to serve: a board whose SPI NAND was erased or +half-written. Its boot ROM finds no valid IDB, gives up on flash and falls +into MaskROM by itself at power-up — no button, no strap, no UART. So a rig +that can only cut power can still bring the board back: + + power_cycle() -> wait_for_device(MASKROM) -> download_boot() + -> write_image(...) -> reset() + +:meth:`RockchipRecovery.download_boot` is the part that turns a MaskROM device +into one that can actually touch flash; everything after it is ordinary block +writes, because the usbplug runs Rockchip's FTL and presents SPI NAND as flat +512-byte sectors with bad blocks and ECC already handled. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import Sequence +from typing import Callable + +from defib.recovery.events import ProgressEvent, Stage +from defib.rockusb.device import ( + DeviceMode, + RockusbDevice, + wait_for_device, +) +from defib.rockusb.loader import LoaderBlobs +from defib.rockusb.maskrom import CODE_471, CODE_472, build_maskrom_chunks +from defib.rockusb.protocol import ( + SECTOR_SIZE, + Opcode, + ResetSubcode, + split_lba_transfers, +) + +logger = logging.getLogger(__name__) + +#: The usbplug needs a moment after the last 472 chunk before it drops off the +#: bus and comes back as a loader-mode device. +USBPLUG_SETTLE = 1.0 + + +def _emit(cb: Callable[[ProgressEvent], None] | None, event: ProgressEvent) -> None: + if cb is not None: + cb(event) + + +class RockchipRecovery: + """Drive one board from MaskROM through to flashed and rebooted.""" + + def __init__(self, device: RockusbDevice) -> None: + self._device = device + + # -- stage 1: get the usbplug running --------------------------------- + + async def download_boot( + self, + blobs: LoaderBlobs, + on_progress: Callable[[ProgressEvent], None] | None = None, + reenumerate_timeout: float = 15.0, + usb_path: str | None = None, + recovery_ids: Sequence[int] | None = None, + ) -> RockusbDevice: + """Upload DDR init then usbplug, and wait for the device to come back. + + Returns the *new* opened device — the old handle is stale once the + usbplug re-enumerates, so callers must use the returned one. + + Pass ``usb_path`` to require that the device coming back is the same + one that went away. Every Rockchip board shares a VID:PID and gets a + fresh USB address on re-enumeration, so without it a second board + appearing mid-upload could be adopted by this session and flashed by + mistake. + """ + if self._device.mode is not DeviceMode.MASKROM: + logger.info("device already past MaskROM; skipping loader upload") + return self._device + + for entry in blobs.ddr: + await self._upload( + CODE_471, entry.data, blobs.use_rc4, Stage.DDR_INIT, entry.name, on_progress + ) + if entry.delay_ms: + await asyncio.sleep(entry.delay_ms / 1000.0) + + for entry in blobs.usbplug: + await self._upload( + CODE_472, entry.data, blobs.use_rc4, Stage.USBPLUG, entry.name, on_progress + ) + if entry.delay_ms: + await asyncio.sleep(entry.delay_ms / 1000.0) + + self._device.close() + await asyncio.sleep(USBPLUG_SETTLE) + + found = await wait_for_device( + timeout=reenumerate_timeout, mode=DeviceMode.LOADER, + usb_path=usb_path, recovery_ids=recovery_ids, + ) + device = RockusbDevice(found) + device.open() + self._device = device + _emit( + on_progress, + ProgressEvent(Stage.USBPLUG, 1, 1, f"usbplug running: {found}"), + ) + return device + + async def _upload( + self, + code: int, + blob: bytes, + use_rc4: bool, + stage: Stage, + name: str, + on_progress: Callable[[ProgressEvent], None] | None, + ) -> None: + chunks = build_maskrom_chunks(blob, use_rc4=use_rc4) + # Total is the framed size, not the blob's: framing appends a CRC and + # sometimes a terminator packet, so measuring real bytes sent against + # the raw length reports more than 100%. + total = sum(len(c) for c in chunks) + sent = 0 + for chunk in chunks: + # Count what the wire took, not what we handed it. control_write + # already rejects a short transfer; crediting the full chunk here + # anyway would make the progress bar lie about it. + sent += await asyncio.to_thread(self._device.control_write, code, chunk) + _emit( + on_progress, + ProgressEvent(stage, sent, total, f"{name} -> {code:#06x}"), + ) + logger.debug( + "uploaded %s (%d bytes in %d chunks) to %#06x", name, len(blob), len(chunks), code + ) + + # -- stage 2: touch flash --------------------------------------------- + + async def write_image( + self, + start_lba: int, + data: bytes, + on_progress: Callable[[ProgressEvent], None] | None = None, + ) -> None: + """Write ``data`` to flash starting at ``start_lba``. + + Short trailing data is zero-padded up to a sector; the usbplug has no + concept of a partial block. + """ + if len(data) % SECTOR_SIZE: + data = data + bytes(SECTOR_SIZE - (len(data) % SECTOR_SIZE)) + + total = len(data) // SECTOR_SIZE + written = 0 + for lba, count in split_lba_transfers(start_lba, total): + offset = (lba - start_lba) * SECTOR_SIZE + await asyncio.to_thread( + self._device.command, + Opcode.WRITE_LBA, + address=lba, + count=count, + data_out=data[offset : offset + count * SECTOR_SIZE], + ) + written += count + _emit( + on_progress, + ProgressEvent( + Stage.FLASH_WRITE, + written * SECTOR_SIZE, + total * SECTOR_SIZE, + f"lba {lba}", + ), + ) + + async def read_image(self, start_lba: int, sectors: int) -> bytes: + """Read ``sectors`` sectors back, for verification.""" + out = bytearray() + for lba, count in split_lba_transfers(start_lba, sectors): + out += await asyncio.to_thread( + self._device.command, + Opcode.READ_LBA, + address=lba, + count=count, + read_length=count * SECTOR_SIZE, + ) + return bytes(out) + + async def read_flash_id(self) -> bytes: + """Flash ID bytes — a cheap "is the usbplug really alive" probe.""" + return await asyncio.to_thread( + self._device.command, Opcode.READ_FLASH_ID, read_length=5 + ) + + def close(self) -> None: + """Release the underlying device. + + Whoever drove the recovery owns the handle; leaving it claimed makes + the next attempt fail to open the board, which looks exactly like + hardware that has stopped answering. + """ + self._device.close() + + async def reset(self, subcode: ResetSubcode = ResetSubcode.NORMAL) -> None: + """Reset the device. + + ``ResetSubcode.MASKROM`` comes back in MaskROM rather than booting, + which is how you chain several flash operations without needing the + board's power cut in between. + + A device is entitled to drop off the bus before acknowledging its own + reset, so a missing status wrapper is tolerated — but only that. + Failing to send the command at all, or getting back an explicit + failure, means the board never reset and must not be reported as + though it did. + """ + await asyncio.to_thread( + self._device.command, + Opcode.RESET_DEVICE, + subcode=int(subcode), + tolerate_disconnect=True, + ) diff --git a/tests/test_profiles_usb_recovery.py b/tests/test_profiles_usb_recovery.py new file mode 100644 index 0000000..2fe2f70 --- /dev/null +++ b/tests/test_profiles_usb_recovery.py @@ -0,0 +1,157 @@ +"""Tests for the USB-recovery discriminator on SoC profiles. + +Chips whose boot ROM has no UART download path (Rockchip) declare +``"RECOVERY": "usb"`` and carry none of the DDR/SPL bytecode. These tests pin +down that the two families stay distinguishable and that neither can silently +be missing what it needs. +""" + +import json +from pathlib import Path + +import pytest +from pydantic import ValidationError + +from defib.profiles.loader import load_profile, recovery_mode +from defib.profiles.schema import FlashPartition, SoCProfile + +PROFILES_DIR = Path(__file__).parent.parent / "src" / "defib" / "profiles" / "data" + + +class TestRecoveryDiscriminator: + def test_uart_is_the_default(self): + """Every pre-existing profile predates the field and must stay UART.""" + assert load_profile("hi3516cv300", PROFILES_DIR).recovery == "uart" + + def test_rv1106_is_usb(self): + assert load_profile("rv1106", PROFILES_DIR).recovery == "usb" + + def test_recovery_mode_helper(self): + assert recovery_mode("hi3516cv300", PROFILES_DIR) == "uart" + assert recovery_mode("rv1106", PROFILES_DIR) == "usb" + + def test_unknown_chip_defaults_to_uart(self): + """V500/CV6xx chips have no JSON profile, so absence must not be + mistaken for USB recovery.""" + assert recovery_mode("hi3516dv500", PROFILES_DIR) == "uart" + assert recovery_mode("no-such-chip-at-all", PROFILES_DIR) == "uart" + + +class TestValidation: + def test_uart_profile_missing_bytecode_is_rejected(self): + with pytest.raises(ValidationError, match="DDRSTEP0"): + SoCProfile.model_validate({"name": "bogus"}) + + def test_uart_error_names_every_missing_field(self): + with pytest.raises(ValidationError) as excinfo: + SoCProfile.model_validate({"name": "bogus", "DDRSTEP0": [1]}) + message = str(excinfo.value) + for alias in ("ADDRESS", "FILELEN", "STEPLEN"): + assert alias in message + + def test_usb_profile_needs_no_bytecode(self): + profile = SoCProfile.model_validate({"name": "x", "RECOVERY": "usb"}) + assert profile.recovery == "usb" + assert profile.ddrstep0 is None + + def test_uart_properties_raise_on_usb_profile(self): + profile = SoCProfile.model_validate({"name": "x", "RECOVERY": "usb"}) + for attr in ( + "ddr_step_address", "spl_address", "uboot_address", + "spl_max_size", "ddr_step_data", + ): + with pytest.raises(ValueError, match="recovers over usb"): + getattr(profile, attr) + + def test_unknown_recovery_value_rejected(self): + with pytest.raises(ValidationError): + SoCProfile.model_validate({"name": "x", "RECOVERY": "jtag"}) + + +class TestRv1106Profile: + def test_declares_its_loader_blobs(self): + profile = load_profile("rv1106", PROFILES_DIR) + assert profile.loader_ddr == "rv1106_ddr_924MHz_v1.15.bin" + assert profile.loader_usbplug == "rv1106_usbplug_v1.09.bin" + + def test_partition_lbas_match_the_vendor_byte_layout(self): + """Read off a real Luckfox Pico Max, not the docs. + + Its U-Boot env and /proc/mtd both give:: + + spi-nand0:256K(env),256K@256K(idblock),512K(uboot),4M(boot), + 30M(oem),10M(userdata),210M(rootfs) + + Note rootfs is 210M. Luckfox's published layout says 80M, which is + what this profile shipped with until hardware contradicted it. + """ + partitions = load_profile("rv1106", PROFILES_DIR).partitions + K, M = 1024, 1024 * 1024 + expected = { + "env": (0, 256 * K), + "idblock": (256 * K, 256 * K), + "uboot": (512 * K, 512 * K), + "boot": (1 * M, 4 * M), + "oem": (5 * M, 30 * M), + "userdata": (35 * M, 10 * M), + "rootfs": (45 * M, 210 * M), + } + actual = { + name: (p.lba * 512, p.size_bytes) for name, p in partitions.items() + } + assert actual == expected + + def test_partitions_tile_without_gaps_or_overlap(self): + """A gap or an overlap here would show up as one image silently + landing inside another.""" + partitions = load_profile("rv1106", PROFILES_DIR).partitions + extents = sorted(partitions.values(), key=lambda p: p.lba) + for lower, upper in zip(extents, extents[1:]): + assert lower.end_lba == upper.lba + + def test_idblock_stays_at_its_fixed_offset(self): + """The boot ROM looks for the IDB at 0x40000; moving it bricks the + board in a way no button recovers.""" + partitions = load_profile("rv1106", PROFILES_DIR).partitions + assert partitions["idblock"].lba * 512 == 0x40000 + + def test_profile_json_is_minimal(self): + """A USB profile carrying UART bytecode would mean someone copied the + wrong template.""" + data = json.loads((PROFILES_DIR / "rv1106.json").read_text()) + assert not {"DDRSTEP0", "PRESTEP0", "ADDRESS", "FILELEN"} & set(data) + + +class TestFlashPartition: + def test_end_lba(self): + assert FlashPartition(lba=100, sectors=50).end_lba == 150 + + def test_size_bytes(self): + assert FlashPartition(lba=0, sectors=8192).size_bytes == 4 * 1024 * 1024 + + +class TestRecoveryModeErrorHandling: + """recovery_mode() must default to UART only for a chip with no profile. + + Swallowing every ValueError would route a mistyped USB chip into the + serial workflow and then report the wrong problem entirely — a failure to + open a serial port, rather than "that variant does not exist". + """ + + def test_missing_profile_defaults_to_uart(self): + assert recovery_mode("no-such-chip-at-all", PROFILES_DIR) == "uart" + + def test_unknown_variant_propagates(self): + with pytest.raises(ValueError, match="Unknown variant"): + recovery_mode("rv1106:typo", PROFILES_DIR) + + def test_unknown_variant_on_uart_chip_also_propagates(self): + with pytest.raises(ValueError, match="Unknown variant"): + recovery_mode("hi3516cv300:nope", PROFILES_DIR) + + def test_declared_variant_still_resolves(self): + from defib.profiles.loader import list_variants + + for chip in ("hi3516cv300", "rv1106"): + for variant in list_variants(chip, PROFILES_DIR): + assert recovery_mode(f"{chip}:{variant}", PROFILES_DIR) in ("uart", "usb") diff --git a/tests/test_rockusb_codec.py b/tests/test_rockusb_codec.py new file mode 100644 index 0000000..b2bb399 --- /dev/null +++ b/tests/test_rockusb_codec.py @@ -0,0 +1,95 @@ +"""Tests for the Rockchip MaskROM wire codecs (CRC-16 and RC4).""" + +import struct + +import pytest + +from defib.rockusb.codec import CRC16_INIT, RK_RC4_KEY, rc4, rk_crc16 + + +def _crc16_bitwise(data: bytes, crc: int = CRC16_INIT) -> int: + """Independent bit-by-bit CRC-CCITT, to check the table-driven version. + + Poly 0x1021, MSB-first, no final XOR. + """ + for byte in data: + crc ^= byte << 8 + for _ in range(8): + crc = ((crc << 1) ^ 0x1021) if (crc & 0x8000) else (crc << 1) + crc &= 0xFFFF + return crc + + +class TestRkCrc16: + @pytest.mark.parametrize( + "data", + [ + b"", + b"\x00", + b"\xff", + b"123456789", + bytes(range(256)), + b"\x00" * 4096, + bytes(range(256)) * 17, + ], + ) + def test_matches_bitwise_reference(self, data): + assert rk_crc16(data) == _crc16_bitwise(data) + + def test_check_vector(self): + # CRC-16/CCITT-FALSE over b"123456789" is the standard 0x29B1 check + # value; this pins the seed and bit order, not just self-consistency. + assert rk_crc16(b"123456789") == 0x29B1 + + def test_empty_returns_seed(self): + assert rk_crc16(b"") == CRC16_INIT + + def test_in_range(self): + assert 0 <= rk_crc16(bytes(range(256))) <= 0xFFFF + + def test_differs_from_hisilicon_variant(self): + """Guard against anyone 'simplifying' this to reuse calc_crc(). + + The HiSilicon helper seeds at 0 and pads with two zero bytes; feeding a + Rockchip blob through it would produce a loader the boot ROM rejects. + """ + from defib.protocol.crc import calc_crc + + data = b"\xde\xad\xbe\xef" + assert rk_crc16(data) != calc_crc(data) + + def test_bytearray_accepted(self): + assert rk_crc16(bytearray(b"abc")) == rk_crc16(b"abc") + + +class TestRc4: + def test_known_vector(self): + # RFC 6229-style vector: key "Key", plaintext "Plaintext". + assert rc4(b"Plaintext", b"Key").hex() == "bbf316e8d940af0ad3" + + def test_second_known_vector(self): + assert rc4(b"pedia", b"Wiki").hex() == "1021bf0420" + + def test_self_inverse(self): + plain = bytes(range(256)) * 3 + assert rc4(rc4(plain)) == plain + + def test_default_key_is_the_rockchip_constant(self): + assert len(RK_RC4_KEY) == 16 + assert RK_RC4_KEY.hex() == "7c4e0304550509072d2c7b38170d1711" + + def test_length_preserved(self): + assert len(rc4(b"\x00" * 1000)) == 1000 + + def test_empty(self): + assert rc4(b"") == b"" + + +class TestCrcAppendOrder: + def test_appended_big_endian(self): + """The boot ROM wants the high byte first.""" + payload = b"\x01\x02\x03" + crc = rk_crc16(payload) + packed = struct.pack(">H", crc) + assert packed[0] == (crc >> 8) & 0xFF + assert packed[1] == crc & 0xFF diff --git a/tests/test_rockusb_device.py b/tests/test_rockusb_device.py new file mode 100644 index 0000000..cd9e596 --- /dev/null +++ b/tests/test_rockusb_device.py @@ -0,0 +1,559 @@ +"""Tests for Rockchip USB device selection and transport behaviour. + +Most of what is pinned here was learnt from a Luckfox Pico Max on the bench: +which device on the bus is actually a recovery target, which stage it is in, +and the several ways an inherited or half-finished transaction makes a +healthy board look dead. +""" + +import pytest + +from defib.rockusb.device import DeviceMode, FoundDevice +from defib.rockusb.protocol import CSW_SIGNATURE, Opcode + + +class TestFoundDevice: + def _dev(self, bus=1, ports=(4, 2), mode=DeviceMode.MASKROM): + return FoundDevice( + mode=mode, bus=bus, address=7, product_id=0x110C, + handle=None, port_numbers=ports, + ) + + def test_usb_path_is_the_topology_path(self): + assert self._dev().usb_path == "1-4.2" + + def test_usb_path_survives_readdressing(self): + """Address changes when the usbplug re-enumerates; the port path is + what identifies the same physical board across that.""" + before = self._dev() + after = FoundDevice( + mode=DeviceMode.LOADER, bus=1, address=9, product_id=0x110C, + handle=None, port_numbers=(4, 2), + ) + assert before.usb_path == after.usb_path + assert before.address != after.address + + def test_distinct_ports_are_distinct_paths(self): + assert self._dev(ports=(4, 2)).usb_path != self._dev(ports=(4, 3)).usb_path + + def test_missing_port_numbers_degrade_visibly(self): + assert self._dev(ports=()).usb_path == "1-?" + + def test_str_mentions_the_path(self): + assert "1-4.2" in str(self._dev()) + + +class TestRecoveryIdFilter: + """A running Luckfox presents 2207:0019 — an RNDIS+ADB gadget sharing the + vendor id, whose bcdUSB of 0x0200 makes it look like MaskROM to the + bcdUSB test. Found on real hardware: without a product-id filter, defib + would have uploaded a loader into a healthy booted board. + """ + + class _FakeUsbDevice: + def __init__(self, pid, bcd=0x0200, strings=False): + self.idProduct = pid + self.bcdUSB = bcd + # MaskROM ships no string descriptors; anything past it names + # itself. Measured on an RV1106: bcdUSB is 0x0200 for both, so + # these indices are the only thing separating the two stages. + self.iManufacturer = 1 if strings else 0 + self.iProduct = 2 if strings else 0 + self.bus = 7 + self.address = 3 + self.port_numbers = (1,) + + def _patch(self, monkeypatch, devices): + import defib.rockusb.device as mod + + class FakeCore: + @staticmethod + def find(**kwargs): + return list(devices) + + monkeypatch.setattr(mod, "_require_usb", lambda: type("U", (), {"core": FakeCore})) + + def test_runtime_adb_gadget_ignored(self, monkeypatch): + from defib.rockusb.device import find_devices + + self._patch(monkeypatch, [self._FakeUsbDevice(0x0019)]) + assert find_devices(recovery_ids=[0x110C]) == [] + + def test_recovery_device_still_found(self, monkeypatch): + from defib.rockusb.device import find_devices + + self._patch(monkeypatch, [self._FakeUsbDevice(0x110C)]) + found = find_devices(recovery_ids=[0x110C]) + assert len(found) == 1 + assert found[0].product_id == 0x110C + + def test_recovery_device_picked_out_of_a_mixed_bus(self, monkeypatch): + from defib.rockusb.device import find_devices + + self._patch( + monkeypatch, + [self._FakeUsbDevice(0x0019), self._FakeUsbDevice(0x110C)], + ) + found = find_devices(recovery_ids=[0x110C]) + assert [d.product_id for d in found] == [0x110C] + + def test_no_ids_means_no_filter(self, monkeypatch): + """Back-compat: callers that pass nothing keep the old broad search.""" + from defib.rockusb.device import find_devices + + self._patch(monkeypatch, [self._FakeUsbDevice(0x0019)]) + assert len(find_devices()) == 1 + + def test_rv1106_profile_declares_the_recovery_id(self): + from defib.profiles.loader import load_profile + + assert load_profile("rv1106").usb_recovery_ids == [0x110C] + + +class TestModeClassification: + """Measured on an RV1106: both stages report bcdUSB 0x0200, so xrock's + low-bit test calls a running usbplug MaskROM. Only the string descriptors + tell them apart. + """ + + def _dev(self, *, strings): + return TestRecoveryIdFilter._FakeUsbDevice(0x110C, strings=strings) + + def test_bare_descriptor_is_maskrom(self): + from defib.rockusb.device import DeviceMode, _classify + + assert _classify(self._dev(strings=False)) is DeviceMode.MASKROM + + def test_named_device_is_loader(self): + from defib.rockusb.device import DeviceMode, _classify + + assert _classify(self._dev(strings=True)) is DeviceMode.LOADER + + def test_bcdusb_is_not_consulted(self): + """Both real stages report 0x0200; keying on it regresses the bug.""" + from defib.rockusb.device import DeviceMode, _classify + + maskrom = TestRecoveryIdFilter._FakeUsbDevice(0x110C, bcd=0x0200, strings=False) + loader = TestRecoveryIdFilter._FakeUsbDevice(0x110C, bcd=0x0200, strings=True) + assert _classify(maskrom) is DeviceMode.MASKROM + assert _classify(loader) is DeviceMode.LOADER + + def test_manufacturer_alone_is_enough(self): + from defib.rockusb.device import DeviceMode, _classify + + dev = TestRecoveryIdFilter._FakeUsbDevice(0x110C, strings=False) + dev.iManufacturer = 1 + assert _classify(dev) is DeviceMode.LOADER + + +class TestResidueScope: + """Measured on an RV1106 usbplug: residue is only honoured on the LBA + path. Everything else reports its transfer length written big-endian — + i.e. "none of it arrived" — while the data plainly did arrive. + """ + + def test_lba_opcodes_are_checked(self): + from defib.rockusb.protocol import Opcode, residue_is_meaningful + + for op in (Opcode.READ_LBA, Opcode.WRITE_LBA, Opcode.ERASE_LBA): + assert residue_is_meaningful(op) + + def test_other_opcodes_are_not(self): + from defib.rockusb.protocol import Opcode, residue_is_meaningful + + for op in ( + Opcode.TEST_UNIT_READY, Opcode.READ_FLASH_ID, + Opcode.READ_CAPABILITY, Opcode.READ_CHIP_INFO, + Opcode.READ_FLASH_INFO, Opcode.RESET_DEVICE, + ): + assert not residue_is_meaningful(op) + + def test_observed_garbage_residues_would_be_ignored(self): + """The literal values seen on the bench, each the transfer length + byte-swapped. Enforcing residue here would break every probe.""" + from defib.rockusb.protocol import Opcode, residue_is_meaningful + + observed = { + Opcode.TEST_UNIT_READY: 0x06000000, + Opcode.READ_FLASH_ID: 0x05000000, + Opcode.READ_CAPABILITY: 0x08000000, + Opcode.READ_FLASH_INFO: 0x0B000000, + } + for op, residue in observed.items(): + assert residue != 0 and not residue_is_meaningful(op) + + +class TestStaleInputDrain: + """A recovery tool is routinely pointed at a device an earlier attempt + abandoned mid-transaction. The unread status wrapper it left behind gets + read as the next command's data phase — 13 bytes into a 5-byte buffer is + [Errno 75] Overflow, and everything after it desynchronises. + """ + + class _Ep: + wMaxPacketSize = 512 + + def __init__(self, queued): + self.queued = list(queued) + self.reads = 0 + + def read(self, size, timeout=None): + self.reads += 1 + if not self.queued: + raise TimeoutError("nothing waiting") + return self.queued.pop(0) + + def _device(self, queued): + from defib.rockusb.device import RockusbDevice + + found = FoundDevice( + mode=DeviceMode.LOADER, bus=1, address=1, product_id=0x110C, + handle=None, port_numbers=(1,), + ) + d = RockusbDevice(found) + d._ep_in = self._Ep(queued) + return d + + def test_stale_bytes_are_discarded(self): + d = self._device([b"USBS" + b"\x00" * 9]) + d._drain_stale_input() + assert d._ep_in.queued == [] + + def test_drain_stops_at_the_first_timeout(self): + d = self._device([]) + d._drain_stale_input() + assert d._ep_in.reads == 1 + + def test_drain_is_bounded(self): + """A device stuck streaming must not hang the open.""" + d = self._device([b"x" * 512] * 100) + d._drain_stale_input() + assert d._ep_in.reads <= 8 + + def test_no_endpoint_is_harmless(self): + """MaskROM is reached before endpoints are known.""" + d = self._device([]) + d._ep_in = None + d._drain_stale_input() + + +class TestBulkReadRounding: + """Bulk IN buffers must be a multiple of the endpoint's max packet size.""" + + class _Ep: + wMaxPacketSize = 512 + + def __init__(self): + self.requested = None + + def read(self, size, timeout=None): + self.requested = size + return b"\x5a" * size + + def _device(self): + from defib.rockusb.device import RockusbDevice + + found = FoundDevice( + mode=DeviceMode.LOADER, bus=1, address=1, product_id=0x110C, + handle=None, port_numbers=(1,), + ) + d = RockusbDevice(found) + d._ep_in = self._Ep() + return d + + def test_short_read_is_rounded_up(self): + d = self._device() + out = d._read_bulk(5) + assert d._ep_in.requested == 512 + assert len(out) == 5 + + def test_exact_multiple_is_unchanged(self): + d = self._device() + d._read_bulk(1024) + assert d._ep_in.requested == 1024 + + def test_partial_sector_rounds_to_next_packet(self): + d = self._device() + d._read_bulk(513) + assert d._ep_in.requested == 1024 + + +class TestResidueHandling: + """A device may move less than asked and still report status OK.""" + + class _FakeEndpoint: + def __init__(self, replies=None): + self.written = [] + self._replies = list(replies or []) + + def write(self, data, timeout=None): + self.written.append(bytes(data)) + return len(data) + + def read(self, size, timeout=None): + return self._replies.pop(0) + + def _device(self, residue: int, status: int = 0): + from defib.rockusb.device import RockusbDevice + + found = FoundDevice( + mode=DeviceMode.LOADER, bus=1, address=1, product_id=0x110C, + handle=None, port_numbers=(1,), + ) + device = RockusbDevice(found) + device._ep_out = self._FakeEndpoint() + # CSW tag is echoed from the CBW the device just received. + device._ep_in = self._FakeEndpoint() + + import struct + + def read(size, timeout=None): + cbw = device._ep_out.written[0] + tag = struct.unpack_from(" len(blob) + + def test_progress_never_exceeds_total(self): + from defib.rockusb.maskrom import build_maskrom_chunks + + for size in (1, 3, 4094, 4095, 4096, 10000): + chunks = build_maskrom_chunks(b"\xa5" * size) + total = sum(len(c) for c in chunks) + sent = 0 + for c in chunks: + sent += len(c) + assert sent <= total + assert sent == total diff --git a/tests/test_rockusb_loader.py b/tests/test_rockusb_loader.py new file mode 100644 index 0000000..df2942a --- /dev/null +++ b/tests/test_rockusb_loader.py @@ -0,0 +1,176 @@ +"""Tests for RKBOOT loader container parsing. + +Entries are read backwards from their stride because ``emType`` is a C enum of +ambiguous width. The parametrised stride tests below are what make that +approach worth having. +""" + +import struct + +import pytest + +from defib.rockusb.loader import ( + LoaderFormatError, + parse_loader, + raw_blobs, +) + +# uiTag[4] usSize[2] dwVersion[4] dwMergeVersion[4] RKTIME[7] emSupportChip[4] +TABLE_AT = 25 + + +def _build_loader( + ddr: bytes = b"\xdd" * 32, + usbplug: bytes = b"\xbb" * 64, + *, + entry_size: int = 57, + rc4_flag: int = 1, + ddr_delay: int = 0, + magic: bytes = b"BOOT", +) -> bytes: + """Synthesise a minimal but structurally honest RKBOOT container.""" + header_len = TABLE_AT + 6 * 3 + 2 + 57 + off471 = header_len + off472 = off471 + entry_size + blob_at = off472 + entry_size + + out = bytearray(blob_at) + out[0:4] = magic + struct.pack_into(" None: + out[base] = entry_size + name_at = base + entry_size - 12 - 40 + out[name_at : name_at + 40] = name.encode("utf-16-le").ljust(40, b"\x00") + struct.pack_into(" bytes: + data = bytearray(_build_loader()) + data[TABLE_AT] = n471 + data[TABLE_AT + 6] = n472 + return bytes(data) + + def test_no_ddr_entries_rejected(self): + with pytest.raises(LoaderFormatError, match="DDR init"): + parse_loader(self._with_counts(0, 1)) + + def test_no_usbplug_entries_rejected(self): + with pytest.raises(LoaderFormatError, match="usbplug"): + parse_loader(self._with_counts(1, 0)) + + def test_neither_stage_names_both(self): + with pytest.raises(LoaderFormatError) as excinfo: + parse_loader(self._with_counts(0, 0)) + assert "DDR init" in str(excinfo.value) + assert "usbplug" in str(excinfo.value) + + +class TestRawBlobs: + def test_wraps_bare_images(self): + blobs = raw_blobs(b"\x01" * 16, b"\x02" * 32) + assert blobs.ddr[0].data == b"\x01" * 16 + assert blobs.usbplug[0].data == b"\x02" * 32 + + def test_defaults_to_rc4_off(self): + assert raw_blobs(b"a", b"b").use_rc4 is False + + def test_rc4_can_be_forced_on(self): + assert raw_blobs(b"a", b"b", use_rc4=True).use_rc4 is True + + def test_empty_ddr_rejected(self): + with pytest.raises(LoaderFormatError, match="ddr"): + raw_blobs(b"", b"b") + + def test_empty_usbplug_rejected(self): + with pytest.raises(LoaderFormatError, match="usbplug"): + raw_blobs(b"a", b"") diff --git a/tests/test_rockusb_maskrom.py b/tests/test_rockusb_maskrom.py new file mode 100644 index 0000000..4a38b58 --- /dev/null +++ b/tests/test_rockusb_maskrom.py @@ -0,0 +1,108 @@ +"""Tests for MaskROM control-transfer framing. + +The size quirks around the 4096-byte chunk boundary are the whole reason this +framing lives in a pure function — they are easy to get wrong and impossible +to notice without hardware, because a mis-framed loader simply never comes +back. +""" + +import struct + +import pytest + +from defib.rockusb.codec import rc4, rk_crc16 +from defib.rockusb.maskrom import CHUNK_SIZE, build_maskrom_chunks + + +def _rejoin(chunks: list[bytes]) -> bytes: + return b"".join(chunks) + + +class TestChunking: + @pytest.mark.parametrize("size", [1, 100, 4095, 4096, 4097, 8192, 100_000]) + def test_no_chunk_exceeds_limit(self, size): + chunks = build_maskrom_chunks(b"\xa5" * size) + assert all(len(c) <= CHUNK_SIZE for c in chunks) + + def test_small_blob_is_one_chunk(self): + chunks = build_maskrom_chunks(b"\x01\x02\x03") + assert len(chunks) == 1 + assert len(chunks[0]) == 5 # 3 payload + 2 CRC + + def test_crc_appended_big_endian(self): + blob = b"\xde\xad\xbe\xef" + chunks = build_maskrom_chunks(blob) + assert _rejoin(chunks) == blob + struct.pack(">H", rk_crc16(blob)) + + def test_exact_chunk_multiple_spills_crc(self): + blob = b"\x11" * CHUNK_SIZE + chunks = build_maskrom_chunks(blob) + assert len(chunks) == 2 + assert len(chunks[0]) == CHUNK_SIZE + assert len(chunks[1]) == 2 # just the CRC + + +class TestBoundaryQuirks: + def test_4095_pads_before_crc(self): + """A 4095-byte blob gains a zero byte so the CRC stays contiguous.""" + blob = b"\x5a" * (CHUNK_SIZE - 1) + chunks = build_maskrom_chunks(blob) + + padded = blob + b"\x00" + assert _rejoin(chunks) == padded + struct.pack(">H", rk_crc16(padded)) + # Padding matters: the CRC must cover it, not the unpadded blob. + assert rk_crc16(padded) != rk_crc16(blob) + assert len(chunks[0]) == CHUNK_SIZE + assert len(chunks[1]) == 2 + + def test_4094_appends_terminator_packet(self): + """CRC lands the payload on an exact multiple, so a short packet is + needed to close the transfer.""" + blob = b"\x5a" * (CHUNK_SIZE - 2) + chunks = build_maskrom_chunks(blob) + + assert len(chunks) == 2 + assert len(chunks[0]) == CHUNK_SIZE + assert chunks[-1] == b"\x00" + assert chunks[0] == blob + struct.pack(">H", rk_crc16(blob)) + + def test_4094_plus_full_chunk_also_terminated(self): + blob = b"\x5a" * (CHUNK_SIZE + CHUNK_SIZE - 2) + chunks = build_maskrom_chunks(blob) + assert chunks[-1] == b"\x00" + assert sum(len(c) for c in chunks[:-1]) % CHUNK_SIZE == 0 + + @pytest.mark.parametrize("size", [4093, 4096, 4097]) + def test_no_terminator_when_not_needed(self, size): + chunks = build_maskrom_chunks(b"\x5a" * size) + assert chunks[-1] != b"\x00" or len(chunks[-1]) > 1 + + +class TestRc4Path: + def test_rc4_off_by_default(self): + """RV1106 loaders are built RC4-off; defaulting the other way would + silently corrupt every upload.""" + blob = b"\x01" * 64 + assert _rejoin(build_maskrom_chunks(blob)).startswith(blob) + + def test_rc4_on_encrypts_then_checksums(self): + blob = b"\x01" * 64 + chunks = build_maskrom_chunks(blob, use_rc4=True) + + encrypted = rc4(blob) + assert _rejoin(chunks) == encrypted + struct.pack(">H", rk_crc16(encrypted)) + + def test_rc4_changes_output(self): + blob = b"\x01" * 64 + assert build_maskrom_chunks(blob) != build_maskrom_chunks(blob, use_rc4=True) + + +class TestDeterminism: + def test_stable_across_calls(self): + blob = bytes(range(256)) * 40 + assert build_maskrom_chunks(blob) == build_maskrom_chunks(blob) + + def test_input_not_mutated(self): + blob = bytearray(b"\x01\x02\x03") + build_maskrom_chunks(bytes(blob)) + assert blob == bytearray(b"\x01\x02\x03") diff --git a/tests/test_rockusb_protocol.py b/tests/test_rockusb_protocol.py new file mode 100644 index 0000000..00cc939 --- /dev/null +++ b/tests/test_rockusb_protocol.py @@ -0,0 +1,168 @@ +"""Tests for rockusb CBW/CSW framing. + +The wrapper is little-endian but the command block inside it is big-endian. +Most of these tests exist to pin that down. +""" + +import struct + +import pytest + +from defib.rockusb.protocol import ( + CBW_LENGTH, + CBW_SIGNATURE, + CSW_LENGTH, + CSW_SIGNATURE, + DIRECTION_IN, + DIRECTION_OUT, + MAX_SECTORS_PER_TRANSFER, + SECTOR_SIZE, + Opcode, + ResetSubcode, + RockusbError, + build_cbw, + parse_csw, + split_lba_transfers, +) + + +def _csw(tag: int, residue: int = 0, status: int = 0, signature: bytes = CSW_SIGNATURE): + return signature + struct.pack("