use partiton table from flash for download FS (#19694)

pio-tools/custom_target.py

@@ -91,63 +91,11 @@ def _parse_size(value):
         return int(value[:-1]) * base
     return value
 
-def _parse_partitions(env):
-    partitions_csv = env.subst("$PARTITIONS_TABLE_CSV")
-    if not isfile(partitions_csv):
-        sys.stderr.write("Could not find the file %s with partitions "
-                         "table.\n" % partitions_csv)
-        env.Exit(1)
-        return
-
-    result = []
-    next_offset = 0
-    with open(partitions_csv) as fp:
-        for line in fp.readlines():
-            line = line.strip()
-            if not line or line.startswith("#"):
-                continue
-            tokens = [t.strip() for t in line.split(",")]
-            if len(tokens) < 5:
-                continue
-            partition = {
-                "name": tokens[0],
-                "type": tokens[1],
-                "subtype": tokens[2],
-                "offset": tokens[3] or next_offset,
-                "size": tokens[4],
-                "flags": tokens[5] if len(tokens) > 5 else None
-            }
-            result.append(partition)
-            next_offset = (_parse_size(partition['offset']) +
-                           _parse_size(partition['size']))
-    return result
-
-def esp32_fetch_spiffs_size(env):
-    spiffs = None
-    for p in _parse_partitions(env):
-        if p['type'] == "data" and p['subtype'] == "spiffs":
-            spiffs = p
-    if not spiffs:
-        sys.stderr.write(
-            env.subst("Could not find the `spiffs` section in the partitions "
-                      "table $PARTITIONS_TABLE_CSV\n"))
-        env.Exit(1)
-        return
-    env["SPIFFS_START"] = _parse_size(spiffs['offset'])
-    env["SPIFFS_SIZE"] = _parse_size(spiffs['size'])
-    env["SPIFFS_PAGE"] = int("0x100", 16)
-    env["SPIFFS_BLOCK"] = int("0x1000", 16)
-
 ## FS helpers for ESP8266
 # copied from https://github.com/platformio/platform-espressif8266/blob/develop/builder/main.py
 # Copyright 2014-present PlatformIO <contact@platformio.org>
 # Licensed under the Apache License, Version 2.0 (the "License");
 
-def _get_board_f_flash(env):
-    frequency = env.subst("$BOARD_F_FLASH")
-    frequency = str(frequency).replace("L", "")
-    return int(int(frequency) / 1000000)
-
 def _parse_ld_sizes(ldscript_path):
     assert ldscript_path
     result = {}
@@ -181,12 +129,6 @@ def _parse_ld_sizes(ldscript_path):
                     match.group(2))
     return result
 
-def _get_flash_size(env):
-    ldsizes = _parse_ld_sizes(env.GetActualLDScript())
-    if ldsizes['flash_size'] < 1048576:
-        return "%dK" % (ldsizes['flash_size'] / 1024)
-    return "%dM" % (ldsizes['flash_size'] / 1048576)
-
 def esp8266_fetch_fs_size(env):
     ldsizes = _parse_ld_sizes(env.GetActualLDScript())
     for key in ldsizes:
@@ -225,13 +167,58 @@ def esp8266_get_esptoolpy_reset_flags(resetmethod):
     return ["--before", resets[0], "--after", resets[1]]
 
 ## Script interface functions
+def parse_partition_table(content):
+    entries = [e for e in content.split(b'\xaaP') if len(e) > 0]
+    print("Partition data:")
+    for entry in entries:
+        type = entry[1]
+        if type in [0x82,0x83]: # SPIFFS or LITTLEFS
+            offset = int.from_bytes(entry[2:5], byteorder='little', signed=False)
+            size = int.from_bytes(entry[6:9], byteorder='little', signed=False)
+            print("type:",hex(type))
+            print("address:",hex(offset))
+            print("size:",hex(size))
+            env["SPIFFS_START"] = offset
+            env["SPIFFS_SIZE"] = size
+            env["SPIFFS_PAGE"] = int("0x100", 16)
+            env["SPIFFS_BLOCK"] = int("0x1000", 16)
+
+def get_partition_table():
+    esptoolpy = join(platform.get_package_dir("tool-esptoolpy") or "", "esptool.py")
+    upload_port = join(env.get("UPLOAD_PORT", "none"))
+    if "none" in upload_port:
+        env.AutodetectUploadPort()
+        upload_port = join(env.get("UPLOAD_PORT", "none"))
+    fs_file = join(env["PROJECT_DIR"], "partition_table_from_flash.bin")
+    esptoolpy_flags = [
+            "--chip", mcu,
+            "--port", upload_port,
+            "--baud",  env.subst("$UPLOAD_SPEED"),
+            "--before", "default_reset",
+            "--after", "hard_reset",
+            "read_flash",
+            "0x8000",
+            "0x1000",
+            fs_file
+    ]
+    esptoolpy_cmd = [env["PYTHONEXE"], esptoolpy] + esptoolpy_flags
+    print("Executing flash download command to read partition table.")
+    print(esptoolpy_cmd)
+    try:
+        returncode = subprocess.call(esptoolpy_cmd, shell=False)
+    except subprocess.CalledProcessError as exc:
+        print("Downloading failed with " + str(exc))
+    with open(fs_file, mode="rb") as file:
+        content = file.read()
+        parse_partition_table(content)
 
 def get_fs_type_start_and_length():
     platform = env["PIOPLATFORM"]
     if platform == "espressif32":
-        print("Retrieving filesystem info for ESP32.")
+        print(f"Retrieving filesystem info for {mcu}.")
         print("Partition file: " + str(env.subst("$PARTITIONS_TABLE_CSV")))
-        esp32_fetch_spiffs_size(env)
+        # esp32_fetch_spiffs_size(env)
+        get_partition_table()
         return SPIFFSInfo(env["SPIFFS_START"], env["SPIFFS_SIZE"], env["SPIFFS_PAGE"], env["SPIFFS_BLOCK"])
     elif platform == "espressif8266":
         print("Retrieving filesystem info for ESP8266.")
@@ -256,12 +243,13 @@ def get_fs_type_start_and_length():
     pass
 
 def download_fs(fs_info: FSInfo):
+    print(fs_info)
     esptoolpy = join(platform.get_package_dir("tool-esptoolpy") or "", "esptool.py")
     upload_port = join(env.get("UPLOAD_PORT", "none"))
     if "none" in upload_port:
         env.AutodetectUploadPort()
         upload_port = join(env.get("UPLOAD_PORT", "none"))
-    fs_file = join(env["PROJECT_DIR"], f"downloaded_fs_{hex(fs_info.start)}_{hex(fs_info.length)}.bin")
+    fs_file = join(env.subst("$BUILD_DIR"), f"downloaded_fs_{hex(fs_info.start)}_{hex(fs_info.length)}.bin")
     esptoolpy_flags = [
             "--chip", mcu,
             "--port", upload_port,
@@ -278,7 +266,7 @@ def download_fs(fs_info: FSInfo):
     print(esptoolpy_cmd)
     try:
         returncode = subprocess.call(esptoolpy_cmd, shell=False)
-        print("Downloaded filesystem binary.")
+        # print("Launched download of filesystem binary.")
         return (True, fs_file)
     except subprocess.CalledProcessError as exc:
         print("Downloading failed with " + str(exc))
@@ -289,6 +277,9 @@ def unpack_fs(fs_info: FSInfo, downloaded_file: str):
     # control the unpack directory
     unpack_dir = env.GetProjectOption("custom_unpack_dir", "unpacked_fs")
     #unpack_dir = "unpacked_fs"
+    if not os.path.exists(downloaded_file):
+        print(f"ERROR: {downloaded_file} with filesystem not found, maybe download failed due to upload_speed setting being too high.")
+        assert(0)
     try:
         if os.path.exists(unpack_dir):
             shutil.rmtree(unpack_dir)
@@ -316,10 +307,11 @@ def display_fs(extracted_dir):
 def command_download_fs(*args, **kwargs):
     print("Entrypoint")
     #print(env.Dump())
+    get_partition_table()
     info = get_fs_type_start_and_length()
     print("Parsed FS info: " + str(info))
     download_ok, downloaded_file = download_fs(info)
-    print("Download was okay: " + str(download_ok) + ". File at: "+ str(downloaded_file))
+    # print("Download was okay: " + str(download_ok) + ". File at: "+ str(downloaded_file)) # this is wrong
     unpack_ok, unpacked_dir = unpack_fs(info, downloaded_file)
     if unpack_ok is True:
         display_fs(unpacked_dir)
@@ -336,7 +328,7 @@ def upload_factory(*args, **kwargs):
         esptoolpy_flags = [
                 "--chip", mcu,
                 "--port", upload_port,
-                "--baud", upload_speed,
+                "--baud", env.subst("$UPLOAD_SPEED"),
                 "write_flash",
                 "0x0",
                 target_firm