اضافه شدن قابلیت تشخیص حافظه

scripts/serial_command.py

@@ -82,16 +82,37 @@ class UBoot:
         return output
 
 
+LAYOUTS = {
+    16: {"uboot": (0, 0x100000), "dtb": (0x100000, 0x10000),
+         "kernel": (0x110000, 0x500000), "rootfs": (0x610000, 0x900000)},
+    32: {"uboot": (0, 0x200000), "dtb": (0x200000, 0x110000),
+         "kernel": (0x310000, 0x500000), "rootfs": (0x810000, 0x17f0000)},
+}
+
+
+def detect_flash(cli):
+    output = cli.command("sf probe 0:0", r"^SF: Detected .+", timeout=20)
+    capacities = re.findall(r"^SF: Detected .*\btotal (\d+) MiB\s*$", output, re.MULTILINE)
+    if len(capacities) != 1 or int(capacities[0]) not in LAYOUTS:
+        raise RuntimeError("Unsupported or ambiguous flash capacity; expected 16 or 32 MiB")
+    capacity = int(capacities[0])
+    print(f"\nDetected {capacity} MiB ({capacity * 8} Mbit); selecting matching flash layout.")
+    return capacity
+
+
 def transfer(cli, firmware, *, dtb=False, kernel=False, rootfs=False, rootfs_size=None):
+    capacity = detect_flash(cli)
+    part = "rootfs" if rootfs else "kernel" if kernel else "dtb" if dtb else "uboot"
+    offset, limit = LAYOUTS[capacity][part]
     address = 0x41000000 if (dtb or rootfs) else 0x42000000
-    offset = 0x310000 if kernel else (0x200000 if dtb else 0)
-    limit = 0x500000 if kernel else (0x110000 if dtb else 0x200000)
-    dump_size = 64 if kernel else 32
-    filename = "zImage" if kernel else ("sun8i-v3s-licheepi-zero-dock.dtb" if dtb else "u-boot-sunxi-with-spl.bin")
-    if rootfs:
-        if rootfs_size is None or rootfs_size < 64:
-            raise RuntimeError("Actual RootFS partition size is required (at least 64 bytes)")
-        offset, limit, dump_size, filename = 0x810000, rootfs_size, 64, "rootfs.squashfs"
+    dump_size = 64 if (kernel or rootfs) else 32
+    filename = {"uboot": "u-boot-sunxi-with-spl.bin", "dtb": "sun8i-v3s-licheepi-zero-dock.dtb",
+                "kernel": "zImage", "rootfs": "rootfs.squashfs"}[part]
+    if rootfs_size is not None and rootfs:
+        if not 64 <= rootfs_size <= limit:
+            raise RuntimeError("RootFS size override exceeds the selected partition or is invalid")
+        limit = rootfs_size
+    print(f"Selected {part}: flash offset {offset:#x}, maximum size {limit:#x}")
     image = firmware.read_bytes()
     if not dump_size <= len(image) <= limit:
         raise RuntimeError(f"Firmware size must be between {dump_size} and {limit:#x} bytes")
@@ -122,7 +143,6 @@ def transfer(cli, firmware, *, dtb=False, kernel=False, rootfs=False, rootfs_siz
     if dumped != image[:dump_size]:
         raise RuntimeError(f"First {dump_size} bytes in RAM differ from the local firmware file")
     print(f"\nTransfer verified: size and first {dump_size} bytes match.")
-    cli.command("sf probe 0:0", r"^SF: Detected .+", timeout=20)
     cli.command(
         f"sf write {address:#x} {offset:#x} ${{filesize}}",
         rf"^SF: {len(image)} bytes @ {offset:#x} Written: OK\s*$", timeout=120,
@@ -142,10 +162,11 @@ def transfer(cli, firmware, *, dtb=False, kernel=False, rootfs=False, rootfs_siz
 
 
 def erase_flash(cli):
-    cli.command("sf probe 0:0", r"^SF: Detected .+total 32 MiB\s*$", timeout=20)
+    capacity = detect_flash(cli)
+    length = capacity * 1024 * 1024
     cli.command(
-        "sf erase 0x0 0x2000000",
-        r"^SF: 33554432 bytes @ 0x0 Erased: OK\s*$", timeout=600,
+        f"sf erase 0x0 {length:#x}",
+        rf"^SF: {length} bytes @ 0x0 Erased: OK\s*$", timeout=600,
     )
     cli.command(
         "sf read 0x42800000 0x0 0x100",
@@ -184,8 +205,6 @@ def main():
     parser.add_argument("--rootfs-size", type=lambda value: int(value, 0), help="Actual RootFS partition size in bytes (decimal or 0x-prefixed hex)")
     parser.add_argument("--confirm-squashfs-mtdblock3", action="store_true", help="Confirm actual RootFS is SquashFS on mtdblock3 and set bootargs")
     args = parser.parse_args()
-    if args.mode == "transfer-rootfs" and (args.rootfs_size is None or args.rootfs_size < 64):
-        parser.error("--rootfs-size must specify the actual RootFS partition size (at least 64 bytes)")
     if args.firmware is None:
         filename = "rootfs.squashfs" if args.mode == "transfer-rootfs" else "zImage" if args.mode == "transfer-kernel" else ("sun8i-v3s-licheepi-zero-dock.dtb" if args.mode == "transfer-dtb" else "u-boot-sunxi-with-spl.bin")
         args.firmware = Path(__file__).resolve().parent.parent / "firmware" / filename