rockusb: read flash over USB, and reset a stale usbplug to MaskROM

src/defib/cli/app.py

@@ -695,17 +695,38 @@ def dump_flash_cmd(
     output_file: str = typer.Option("flash_dump.bin", "-o", "--output", help="Output binary file"),
     size: str = typer.Option("", "--size", help="Flash size (e.g., 8MB, 16MB) — auto-detect if empty"),
     output: str = typer.Option("human", "--output-mode", help="Output mode: human, json"),
+    chip: str = typer.Option("", "-c", "--chip", help="Chip model. Only needed for chips that dump over USB rather than a U-Boot console."),
+    partition: str = typer.Option("", "--partition", help="USB-recovery chips: dump one named partition instead of the whole flash"),
+    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)"),
+    power_cycle: bool = typer.Option(False, "--power-cycle", help="Auto power-cycle via the configured controller"),
 ) -> None:
-    """Dump flash contents via U-Boot serial console.
+    """Dump flash contents.
+
+    For most chips this drives a U-Boot serial console, so U-Boot must
+    already be running (connect first, or use after 'defib burn').
 
-    Requires U-Boot to be running on the device (connect to serial
-    console first, or use after 'defib burn').
+    Chips whose boot ROM only answers on USB (Rockchip) are read directly
+    over USB instead — pass -c along with --ddr/--usbplug, and -p is not
+    consulted. Use --partition to pull a single named region rather than the
+    whole device.
     """
     import asyncio
-    asyncio.run(_dump_flash_async(port, output_file, size, output))
+    asyncio.run(_dump_flash_async(
+        port, output_file, size, output, chip, partition,
+        ddr, usbplug, loader, wait, usb_path, power_cycle,
+    ))
 
 
-async def _dump_flash_async(port: str, output_file: str, size: str, output: str) -> None:
+async def _dump_flash_async(
+    port: str, output_file: str, size: str, output: str,
+    chip: str = "", partition: str = "",
+    ddr: str = "", usbplug: str = "", loader: str = "",
+    wait: float = 30.0, usb_path: str = "", power_cycle: bool = False,
+) -> None:
     import json as json_mod
 
     from rich.console import Console
@@ -715,6 +736,13 @@ async def _dump_flash_async(port: str, output_file: str, size: str, output: str)
 
     console = Console()
 
+    if chip and _recovery_mode_or_exit(chip, output) == "usb":
+        await _dump_flash_usb_async(
+            chip, output_file, partition, ddr, usbplug, loader,
+            wait, usb_path, power_cycle, output,
+        )
+        return
+
     # Parse size
     flash_size = None
     if size:
@@ -3718,7 +3746,12 @@ async def _open_usb_target(
     """
     from rich.console import Console
 
-    from defib.rockusb.device import DeviceMode, RockusbDevice, wait_for_device
+    from defib.rockusb.device import (
+        DeviceMode,
+        RockusbDevice,
+        RockusbUsbError,
+        wait_for_device,
+    )
     from defib.rockusb.recovery import RockchipRecovery
 
     console = Console()
@@ -3744,6 +3777,25 @@ async def _open_usb_target(
     device.open()
     recovery = RockchipRecovery(device)
 
+    # A usbplug left running by a previous invocation is not safe to reuse: it
+    # reliably wedges on the next command from a fresh process. Every tool in
+    # this space re-uploads the loader each run for that reason. So if we find
+    # one already in loader mode, send it back to MaskROM and start clean.
+    if found.mode is DeviceMode.LOADER:
+        if output == "human":
+            console.print("  Found a stale loader; returning it to MaskROM...")
+        try:
+            await recovery.return_to_maskrom(
+                usb_path=found.usb_path, recovery_ids=recovery_ids
+            )
+        except RockusbUsbError as e:
+            recovery.close()
+            raise RockusbUsbError(
+                f"{e}. A previous run left the loader wedged and it will not "
+                "reset — power-cycle the board (BOOT + replug) and retry."
+            ) from e
+        found = recovery.device_info
+
     if found.mode is DeviceMode.MASKROM:
         if output == "human":
             console.print("  Uploading DDR init and usbplug...")
@@ -3831,5 +3883,111 @@ async def _burn_usb_async(
         )
 
 
+async def _dump_flash_usb_async(
+    chip: str, output_file: str, partition: str,
+    ddr: str, usbplug: str, loader: str,
+    wait: float, usb_path: str, power_cycle: bool, output: str,
+) -> None:
+    """``dump-flash`` for chips whose boot ROM only answers on USB.
+
+    Reads through the usbplug's LBA space, which maps directly onto the
+    kernel's mtd partitions — a partition dumped here compares byte for byte
+    against one taken with ``cat /dev/mtdN`` on a running board.
+
+    Streams to disk: a full RV1106 is 255 MiB and there is no reason for it
+    to pass through memory twice.
+    """
+    import json as json_mod
+    from pathlib import Path
+
+    from rich.console import Console
+
+    from defib.profiles.loader import load_profile
+    from defib.rockusb.loader import LoaderFormatError
+    from defib.rockusb.protocol import SECTOR_SIZE, RockusbError
+
+    console = Console()
+    profile = load_profile(chip)
+
+    extent = None
+    if partition:
+        extent = profile.partitions.get(partition)
+        if extent is None:
+            known = ", ".join(sorted(profile.partitions)) or "(none declared)"
+            _usb_fail(output, f"{chip} declares no partition {partition!r}. Known: {known}")
+            return
+
+    recovery = None
+    written = 0
+    try:
+        blobs = _resolve_usb_loader(chip, ddr, usbplug, loader)
+        recovery = await _open_usb_target(
+            blobs, power_cycle, output, wait, "", usb_path,
+            profile.usb_recovery_ids,
+        )
+
+        info = await recovery.read_flash_info()
+        if output == "human":
+            console.print(f"  Flash: {info}")
+
+        if extent is not None:
+            start, sectors = extent.lba, extent.sectors
+        else:
+            start, sectors = 0, info.sectors
+
+        # Never read past what the device says it holds: the LBA commands
+        # stop answering there, and a dump that ends in a stall is worse than
+        # one that refuses to start.
+        if start + sectors > info.sectors:
+            _usb_fail(
+                output,
+                f"{partition or 'full dump'} runs to sector {start + sectors} "
+                f"but the flash reports only {info.sectors}",
+            )
+            return
+
+        target = Path(output_file)
+        if output == "human":
+            what = f"partition {partition!r}" if partition else "whole flash"
+            console.print(
+                f"  Dumping {what}: {sectors * SECTOR_SIZE / 1024 / 1024:.1f} MiB "
+                f"from lba {start} -> {target}"
+            )
+
+        # Stream to a temp file beside the destination and atomically replace
+        # it only once the whole transfer lands: a failure part-way must not
+        # truncate an existing backup or strand a partial image at the path.
+        tmp = target.with_name(target.name + ".partial")
+        ok = False
+        try:
+            with tmp.open("wb") as fh:
+                written = await recovery.dump_image(
+                    start, sectors, fh.write,
+                    on_progress=_usb_progress_printer(output),
+                )
+            tmp.replace(target)
+            ok = True
+        finally:
+            if not ok:
+                tmp.unlink(missing_ok=True)
+    except (RockusbError, LoaderFormatError, OSError, typer.BadParameter) as e:
+        _usb_fail(output, str(e))
+        return
+    finally:
+        if recovery is not None:
+            recovery.close()
+
+    if output == "json":
+        print(json_mod.dumps({
+            "event": "done", "success": True,
+            "file": output_file, "bytes": written,
+        }))
+    else:
+        console.print(
+            f"\n[green bold]Dumped {written / 1024 / 1024:.1f} MiB[/green bold] "
+            f"to {output_file}"
+        )
+
+
 def main() -> None:
     app()

src/defib/recovery/events.py

@@ -23,6 +23,7 @@ class Stage(str, Enum):
     # block writes it enables.
     USBPLUG = "usbplug"
     FLASH_WRITE = "flash_write"
+    FLASH_READ = "flash_read"
     COMPLETE = "complete"
 
 

src/defib/rockusb/__init__.py

@@ -30,21 +30,26 @@
 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 (
+    FLASH_INFO_LENGTH,
     SECTOR_SIZE,
     CommandStatus,
+    FlashInfo,
     Opcode,
     ResetSubcode,
     RockusbError,
     build_cbw,
     parse_csw,
+    parse_flash_info,
 )
 
 __all__ = [
     "CODE_471",
     "CODE_472",
+    "FLASH_INFO_LENGTH",
     "RK_RC4_KEY",
     "SECTOR_SIZE",
     "CommandStatus",
+    "FlashInfo",
     "LoaderBlobs",
     "LoaderFormatError",
     "Opcode",
@@ -53,6 +58,7 @@
     "build_cbw",
     "build_maskrom_chunks",
     "parse_csw",
+    "parse_flash_info",
     "parse_loader",
     "rc4",
     "rk_crc16",

src/defib/rockusb/protocol.py

@@ -12,6 +12,7 @@
 from __future__ import annotations
 
 import struct
+from dataclasses import dataclass
 from enum import IntEnum
 
 CBW_SIGNATURE = b"USBC"
@@ -72,6 +73,80 @@ def cdb_length(opcode: Opcode | int) -> int:
     )
 
 
+#: Bytes READ_FLASH_INFO returns.
+FLASH_INFO_LENGTH = 11
+
+
+@dataclass(frozen=True)
+class FlashInfo:
+    """What the usbplug reports about the flash behind it.
+
+    Sizes are the FTL's view, not the raw part: an RV1106 with a 256 MiB SPI
+    NAND reports 255.5 MiB, the remainder being the translation layer's own
+    reserve. That is the number a dump should trust — it is exactly the span
+    the LBA commands will answer for.
+    """
+
+    sectors: int
+    block_sectors: int
+    page_sectors: int
+    ecc_bits: int
+    access_time: int
+    manufacturer: int
+    flash_mask: int
+
+    @property
+    def size_bytes(self) -> int:
+        return self.sectors * SECTOR_SIZE
+
+    @property
+    def block_bytes(self) -> int:
+        """Erase block size. Matches the kernel's mtd ``erasesize``."""
+        return self.block_sectors * SECTOR_SIZE
+
+    @property
+    def page_bytes(self) -> int:
+        return self.page_sectors * SECTOR_SIZE
+
+    def __str__(self) -> str:
+        return (
+            f"{self.size_bytes / 1024 / 1024:.1f} MiB "
+            f"({self.sectors} sectors), "
+            f"block {self.block_bytes // 1024} KiB, page {self.page_bytes} B"
+        )
+
+
+def parse_flash_info(data: bytes) -> FlashInfo:
+    """Parse a READ_FLASH_INFO reply.
+
+    Raises:
+        RockusbError: on a short reply, or one claiming zero capacity — both
+            mean the usbplug never got the flash up, and dumping from it would
+            produce a convincing file full of nothing.
+    """
+    if len(data) < FLASH_INFO_LENGTH:
+        raise RockusbError(
+            f"short flash info: got {len(data)} bytes, want {FLASH_INFO_LENGTH}"
+        )
+    sectors, block, page, ecc, access, mfr, mask = struct.unpack_from(
+        "<IHBBBBB", data
+    )
+    if sectors == 0:
+        raise RockusbError(
+            "flash reports zero capacity — the loader is running but has not "
+            "brought the flash up"
+        )
+    return FlashInfo(
+        sectors=sectors,
+        block_sectors=block,
+        page_sectors=page,
+        ecc_bits=ecc,
+        access_time=access,
+        manufacturer=mfr,
+        flash_mask=mask,
+    )
+
+
 def residue_is_meaningful(opcode: Opcode | int) -> bool:
     """Whether this opcode's status wrapper reports a usable residue.