diff --git a/src/defib/cli/app.py b/src/defib/cli/app.py index 2c2c4cf..b1cafbc 100644 --- a/src/defib/cli/app.py +++ b/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() diff --git a/src/defib/recovery/events.py b/src/defib/recovery/events.py index f427acd..f8e4617 100644 --- a/src/defib/recovery/events.py +++ b/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" diff --git a/src/defib/rockusb/__init__.py b/src/defib/rockusb/__init__.py index d779849..6c4e79b 100644 --- a/src/defib/rockusb/__init__.py +++ b/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", diff --git a/src/defib/rockusb/protocol.py b/src/defib/rockusb/protocol.py index 29759dc..053f547 100644 --- a/src/defib/rockusb/protocol.py +++ b/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( + " bool: """Whether this opcode's status wrapper reports a usable residue. diff --git a/src/defib/rockusb/recovery.py b/src/defib/rockusb/recovery.py index cc061b6..54efc6b 100644 --- a/src/defib/rockusb/recovery.py +++ b/src/defib/rockusb/recovery.py @@ -24,15 +24,19 @@ from defib.recovery.events import ProgressEvent, Stage from defib.rockusb.device import ( DeviceMode, + FoundDevice, 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 ( + FLASH_INFO_LENGTH, SECTOR_SIZE, + FlashInfo, Opcode, ResetSubcode, + parse_flash_info, split_lba_transfers, ) @@ -188,12 +192,88 @@ async def read_image(self, start_lba: int, sectors: int) -> bytes: ) return bytes(out) + async def read_flash_info(self) -> FlashInfo: + """Capacity and geometry as the FTL sees it.""" + return parse_flash_info( + await asyncio.to_thread( + self._device.command, + Opcode.READ_FLASH_INFO, + read_length=FLASH_INFO_LENGTH, + ) + ) + + async def dump_image( + self, + start_lba: int, + sectors: int, + sink: Callable[[bytes], None], + on_progress: Callable[[ProgressEvent], None] | None = None, + ) -> int: + """Read ``sectors`` sectors, handing each chunk straight to ``sink``. + + Streams rather than returning bytes: a full RV1106 dump is 255 MiB, + and holding that in memory to write it out again helps nobody. + + Returns the number of bytes read. + """ + done = 0 + total = sectors * SECTOR_SIZE + for lba, count in split_lba_transfers(start_lba, sectors): + chunk = await asyncio.to_thread( + self._device.command, + Opcode.READ_LBA, + address=lba, + count=count, + read_length=count * SECTOR_SIZE, + ) + sink(chunk) + done += len(chunk) + _emit(on_progress, ProgressEvent(Stage.FLASH_READ, done, total, f"lba {lba}")) + return done + 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 ) + @property + def device_info(self) -> FoundDevice: + """The currently-held device's :class:`FoundDevice`.""" + return self._device._found + + async def return_to_maskrom( + self, + usb_path: str | None = None, + recovery_ids: Sequence[int] | None = None, + timeout: float = 15.0, + ) -> RockusbDevice: + """Send a running usbplug back to MaskROM and re-acquire it there. + + Reusing a usbplug a previous process left behind wedges the next + command, so a fresh run resets to MaskROM and re-uploads. Returns the + new MaskROM-mode device handle; the old one is closed. + + Raises: + RockusbUsbError: if the loader will not accept the reset (already + wedged) or MaskROM never reappears. + """ + if self._device.mode is DeviceMode.MASKROM: + return self._device + + await self.reset(ResetSubcode.MASKROM) + self._device.close() + await asyncio.sleep(USBPLUG_SETTLE) + + found = await wait_for_device( + timeout=timeout, mode=DeviceMode.MASKROM, + usb_path=usb_path, recovery_ids=recovery_ids, + ) + device = RockusbDevice(found) + device.open() + self._device = device + return device + def close(self) -> None: """Release the underlying device. diff --git a/tests/test_rockusb_device.py b/tests/test_rockusb_device.py index 9070333..ea72c78 100644 --- a/tests/test_rockusb_device.py +++ b/tests/test_rockusb_device.py @@ -588,3 +588,99 @@ def test_block_read_write_still_declare_the_data_phase(self): for op in (Opcode.WRITE_LBA, Opcode.READ_LBA): cbw = build_cbw(tag=1, opcode=op, address=0, count=2) assert self._transfer_length(cbw) == 2 * SECTOR_SIZE, op + + +class TestReturnToMaskrom: + """A usbplug a previous process left running wedges the next command, so a + fresh run resets it to MaskROM and re-uploads. This is the reset step. + """ + + class _Dev: + def __init__(self, mode, after=None): + from defib.rockusb.device import FoundDevice + self._mode = mode + self._found = FoundDevice( + mode=mode, bus=1, address=1, product_id=0x110C, + handle=None, port_numbers=(1,), + ) + self.closed = False + + @property + def mode(self): + return self._mode + + def close(self): + self.closed = True + + async def test_noop_when_already_maskrom(self, monkeypatch): + from defib.rockusb.device import DeviceMode + from defib.rockusb.recovery import RockchipRecovery + + dev = self._Dev(DeviceMode.MASKROM) + r = RockchipRecovery(dev) + out = await r.return_to_maskrom() + assert out is dev + assert not dev.closed + + async def test_device_info_exposes_the_handle(self): + from defib.rockusb.device import DeviceMode + from defib.rockusb.recovery import RockchipRecovery + + dev = self._Dev(DeviceMode.LOADER) + r = RockchipRecovery(dev) + assert r.device_info.mode is DeviceMode.LOADER + + async def test_reset_and_reacquire(self, monkeypatch): + """A healthy loader resets, the old handle is closed, and a fresh + MaskROM handle comes back.""" + import defib.rockusb.recovery as mod + from defib.rockusb.device import DeviceMode, FoundDevice + from defib.rockusb.recovery import RockchipRecovery + + loader = self._Dev(DeviceMode.LOADER) + r = RockchipRecovery(loader) + + reset_calls = [] + + async def fake_reset(subcode): + reset_calls.append(subcode) + + monkeypatch.setattr(r, "reset", fake_reset) + monkeypatch.setattr(mod.asyncio, "sleep", _no_sleep) + + maskrom_found = FoundDevice( + mode=DeviceMode.MASKROM, bus=1, address=2, product_id=0x110C, + handle=object(), port_numbers=(1,), + ) + + async def fake_wait(**kwargs): + assert kwargs["mode"] is DeviceMode.MASKROM + return maskrom_found + + opened = [] + monkeypatch.setattr(mod, "wait_for_device", fake_wait) + monkeypatch.setattr( + mod, "RockusbDevice", + lambda found: _FakeReopened(found, opened), + ) + + from defib.rockusb.protocol import ResetSubcode + + await r.return_to_maskrom() + assert reset_calls == [ResetSubcode.MASKROM] + assert loader.closed + assert opened == ["opened"] + + +async def _no_sleep(_seconds): + return None + + +class _FakeReopened: + def __init__(self, found, log): + self._found = found + self._log = log + self.mode = found.mode + + def open(self): + self._log.append("opened") diff --git a/tests/test_rockusb_flashinfo.py b/tests/test_rockusb_flashinfo.py new file mode 100644 index 0000000..07f0ac3 --- /dev/null +++ b/tests/test_rockusb_flashinfo.py @@ -0,0 +1,146 @@ +"""Tests for READ_FLASH_INFO parsing. + +The reply is what a dump trusts for "how much is there", so a wrong decode +either truncates the image or runs off the end into stalled reads. +""" + +import struct + +import pytest + +from defib.rockusb.protocol import ( + FLASH_INFO_LENGTH, + FlashInfo, + RockusbError, + parse_flash_info, +) + +# Captured from a Luckfox Pico Max (RV1106G3) over rockusb. +REAL = bytes.fromhex("00fc070000010400280001") + + +class TestParseRealDevice: + def test_decodes_the_captured_reply(self): + info = parse_flash_info(REAL) + assert info.sectors == 523264 + assert info.block_sectors == 256 + assert info.page_sectors == 4 + + def test_capacity_is_just_under_the_part_size(self): + """256 MiB of SPI NAND, less the FTL's own reserve.""" + info = parse_flash_info(REAL) + assert 255.0 < info.size_bytes / 1024 / 1024 < 256.0 + + def test_block_size_matches_the_kernels_erasesize(self): + """/proc/mtd on the same board reports erasesize 0x20000.""" + assert parse_flash_info(REAL).block_bytes == 0x20000 + + def test_page_size_is_2k(self): + assert parse_flash_info(REAL).page_bytes == 2048 + + def test_str_is_readable(self): + text = str(parse_flash_info(REAL)) + assert "255.5 MiB" in text + assert "block 128 KiB" in text + + +class TestRejections: + def test_short_reply_rejected(self): + with pytest.raises(RockusbError, match="short flash info"): + parse_flash_info(REAL[:-1]) + + def test_zero_capacity_rejected(self): + """A loader that is running but never brought flash up would + otherwise yield a convincing file full of nothing.""" + bad = struct.pack("