firmware flash --uf2: accept RP2350 images, and copy to unmounted WSL drives

cli/src/mpftp/uf2.py

@@ -28,6 +28,7 @@
 from __future__ import annotations
 
 import os
+import re
 import struct
 import subprocess
 import time
@@ -47,6 +48,8 @@
 UF2_FLAG_FAMILY_ID = 0x00002000
 UF2_FLAG_MD5_CHECKSUM = 0x00004000
 
+UF2_FAMILY_RP2XXX_ABSOLUTE = 0xE48BFF57
+
 # Family IDs from the upstream registry (microsoft/uf2 utils/uf2families.json).
 # Vendored rather than read from a sibling checkout so the engine stays
 # self-contained; unknown IDs degrade to hex, never to an error.
@@ -176,7 +179,13 @@ def parse_uf2(path: Path, max_blocks: int = 0) -> dict:
     payload_bytes = 0
     addr_lo: Optional[int] = None
     addr_hi: Optional[int] = None
-    declared_total: Optional[int] = None
+    # A UF2 file is a series of runs, each numbered from 0 with its own
+    # numBlocks. picotool's RP2350 images open with a short run in the
+    # absolute family, so each run is checked against its own header.
+    run_key: Optional[tuple] = None
+    run_declared = 0
+    run_count = 0
+    runs: list[tuple[int, int]] = []
     counted = 0
 
     with path.open("rb") as fh:
@@ -201,8 +210,20 @@ def parse_uf2(path: Path, max_blocks: int = 0) -> dict:
                 continue
             counted += 1
             payload_bytes += payload_len
-            if declared_total is None:
-                declared_total = num_blocks
+            if flags & UF2_FLAG_FAMILY_ID and family == UF2_FAMILY_RP2XXX_ABSOLUTE:
+                # picotool's RP2350-E10 workaround block declares numBlocks 2
+                # and carries only itself; the bootloader reads it as a hint
+                # and it must stay in the file, so it takes no part in the
+                # count of any run.
+                addr_lo = addr if addr_lo is None else min(addr_lo, addr)
+                addr_hi = addr + payload_len if addr_hi is None else max(addr_hi, addr + payload_len)
+                continue
+            key = (family if flags & UF2_FLAG_FAMILY_ID else None, num_blocks)
+            if run_key is None or key != run_key or blk_no == 0:
+                if run_key is not None:
+                    runs.append((run_declared, run_count))
+                run_key, run_declared, run_count = key, num_blocks, 0
+            run_count += 1
             if flags & UF2_FLAG_FAMILY_ID and family not in families:
                 families.append(family)
             addr_lo = addr if addr_lo is None else min(addr_lo, addr)
@@ -213,12 +234,17 @@ def parse_uf2(path: Path, max_blocks: int = 0) -> dict:
 
     if counted == 0:
         raise Uf2Error(f"{path.name}: no flashable blocks")
-    if not max_blocks and declared_total is not None and declared_total != counted:
-        # A bootloader waits for numBlocks blocks before committing, so a file
+    runs.append((run_declared, run_count))
+    if not max_blocks:
+        # A bootloader waits for numBlocks blocks before committing, so a run
         # that disagrees with its own header will hang rather than flash.
-        raise Uf2Error(
-            f"{path.name}: header declares {declared_total} blocks but the file has {counted}"
-        )
+        for declared, have in runs:
+            if declared != have:
+                raise Uf2Error(
+                    f"{path.name}: header declares {declared} blocks but the file has {have}"
+                    if len(runs) == 1 else
+                    f"{path.name}: a run of blocks declares {declared} but has {have}"
+                )
     if not families:
         # Legal, but it means the file targets whatever it is copied onto.
         warnings.append(
@@ -264,9 +290,28 @@ def parse_info_uf2(text: str) -> dict:
     return info
 
 
+def is_windows_drive(root: Path) -> bool:
+    """True for ``H:\\`` seen from a Linux python (WSL), where it is no path.
+
+    A removable drive that WSL never mounted can only be reached through
+    Windows, so reading, probing and copying go through PowerShell.
+    """
+    return os.name != "nt" and re.match(r"^[A-Za-z]:[\\/]*$", str(root)) is not None
+
+
+def _powershell(command: str, timeout: int = 60) -> subprocess.CompletedProcess:
+    return subprocess.run(
+        ["powershell.exe", "-NoProfile", "-Command", command],
+        capture_output=True, text=True, timeout=timeout,
+    )
+
+
 def read_volume_info(root: Path) -> dict:
     """Read INFO_UF2.TXT from a mounted bootloader volume, or {} if unreadable."""
     try:
+        if is_windows_drive(root):
+            done = _powershell(f"Get-Content -Raw '{str(root)[:2]}\\INFO_UF2.TXT'", 15)
+            return parse_info_uf2(done.stdout) if done.returncode == 0 else {}
         return parse_info_uf2((root / "INFO_UF2.TXT").read_text(errors="replace"))
     except Exception:
         return {}
@@ -347,10 +392,7 @@ def find_uf2_volumes(host: str) -> list[dict]:
     """
     found: list[dict] = []
     for root in _candidate_roots(host):
-        try:
-            if not (root / "INFO_UF2.TXT").is_file():
-                continue
-        except OSError:
+        if not volume_present(root):
             continue
         info = read_volume_info(root)
         found.append({
@@ -373,6 +415,9 @@ def looks_like_volume(dev: str) -> bool:
 def volume_present(root: Path) -> bool:
     """True while the bootloader volume is still mounted."""
     try:
+        if is_windows_drive(root):
+            done = _powershell(f"Test-Path '{str(root)[:2]}\\INFO_UF2.TXT'", 15)
+            return done.stdout.strip() == "True"
         return (root / "INFO_UF2.TXT").is_file()
     except OSError:
         # A vanishing removable volume can raise rather than return False.
@@ -412,6 +457,8 @@ def copy_uf2(src: Path, dest: Path) -> int:
     errors *after* the final write are expected and are left for the caller's
     disappearance check to adjudicate.
     """
+    if is_windows_drive(dest.parent):
+        return _copy_via_windows(src, dest)
     written = 0
     with src.open("rb") as fin, dest.open("wb") as fout:
         while True:
@@ -428,3 +475,22 @@ def copy_uf2(src: Path, dest: Path) -> int:
             # error by itself -- wait_for_volume_gone decides.
             pass
     return written
+
+
+def _copy_via_windows(src: Path, dest: Path) -> int:
+    """Copy onto a Windows drive that this Linux python cannot reach (WSL).
+
+    ``-ErrorAction Stop`` matters: ``Copy-Item`` otherwise fails without a
+    nonzero exit, which reads as success.
+    """
+    win_src = subprocess.check_output(["wslpath", "-w", str(src)], text=True).strip()
+    target = f"{str(dest.parent)[:2]}\\{dest.name}"
+    done = _powershell(
+        f"Copy-Item -LiteralPath '{win_src}' -Destination '{target}' -ErrorAction Stop",
+        120,
+    )
+    # The board may unmount before PowerShell sees the copy finish; the
+    # caller's disappearance check adjudicates that case.
+    if done.returncode != 0 and volume_present(dest.parent):
+        raise OSError(done.stderr.strip() or "Copy-Item failed")
+    return src.stat().st_size

cli/tests/test_uf2.py

@@ -6,6 +6,7 @@
 
 from __future__ import annotations
 
+import os
 import struct
 import tempfile
 import unittest
@@ -31,6 +32,26 @@ def _write_uf2(path: Path, blocks: list[bytes]) -> Path:
     return path
 
 
+class WindowsDriveTests(unittest.TestCase):
+    def test_drive_letter_is_unreachable_from_linux_python(self) -> None:
+        with mock.patch.object(uf2.os, "name", "posix"):
+            self.assertTrue(uf2.is_windows_drive(Path("H:\\")))
+            self.assertFalse(uf2.is_windows_drive(Path("/mnt/h")))
+
+    def test_copy_failure_on_such_a_drive_is_not_taken_for_a_reboot(self) -> None:
+        done = mock.Mock(returncode=1, stderr="no space", stdout="True")
+        with (
+            mock.patch.object(uf2.os, "name", "posix"),
+            mock.patch.object(uf2.subprocess, "check_output", return_value="H:\\fw.uf2"),
+            mock.patch.object(uf2.subprocess, "run", return_value=done),
+            tempfile.TemporaryDirectory() as d,
+        ):
+            src = Path(d) / "fw.uf2"
+            src.write_bytes(b"x" * 512)
+            with self.assertRaises(OSError):
+                uf2.copy_uf2(src, Path("H:\\") / "fw.uf2")
+
+
 class ParseUf2Tests(unittest.TestCase):
     def setUp(self) -> None:
         self.tmp = tempfile.TemporaryDirectory()
@@ -70,6 +91,31 @@ def test_rejects_block_count_disagreeing_with_header(self) -> None:
             uf2.parse_uf2(_write_uf2(self.dir / "short.uf2", blocks))
         self.assertIn("declares 9 blocks", str(cm.exception))
 
+    def test_multi_run_file_checks_each_run_against_its_own_header(self) -> None:
+        # RP2350 images from picotool: one absolute-family block that declares 2, then the image.
+        absolute = [_block(0x10FFFF00, b"\x01" * 256, 0, 2, family=0xE48BFF57)]
+        image = [_block(0x10000000 + i * 256, b"\xaa" * 256, i, 5, family=0xE48BFF59)
+                 for i in range(5)]
+        meta = uf2.parse_uf2(_write_uf2(self.dir / "rp2350.uf2", absolute + image))
+        self.assertEqual(meta["blocks"], 6)
+        self.assertEqual(meta["family_names"], ["RP2350_ARM_S"])
+        self.assertEqual(meta["warnings"], [])
+
+    def test_multi_run_file_with_a_short_image_run_is_rejected(self) -> None:
+        absolute = [_block(0x10FFFF00, b"\x01" * 256, 0, 2, family=0xE48BFF57)]
+        image = [_block(0x10000000 + i * 256, b"\xaa" * 256, i, 4, family=0xE48BFF59)
+                 for i in range(3)]
+        with self.assertRaises(uf2.Uf2Error):
+            uf2.parse_uf2(_write_uf2(self.dir / "bad.uf2", absolute + image))
+
+    def test_real_rp2350_image_parses(self) -> None:
+        # Point MPFTP_TEST_RP2350_UF2 at any picotool-made RP2350 image.
+        image = os.environ.get("MPFTP_TEST_RP2350_UF2")
+        if not image:
+            self.skipTest("MPFTP_TEST_RP2350_UF2 not set")
+        meta = uf2.parse_uf2(Path(image))
+        self.assertEqual(meta["family_names"], ["RP2350_ARM_S"])
+
     def test_not_main_flash_blocks_are_not_counted(self) -> None:
         blocks = [
             _block(0, b"\x01" * 16, 0, 1, flags=uf2.UF2_FLAG_NOT_MAIN_FLASH),