change: make command independent on status byte length

This commit adds passing of expected response data length to check_command,
which all command independent on status byte length. This enables the
--after no-reset-no-sync option for all chips when stub is running and
unifies read_reg verification.

Author: Dzarda7

esptool/loader.py

@@ -291,9 +291,6 @@ class ESPLoader(object):
     IROM_MAP_START = 0x40200000
     IROM_MAP_END = 0x40300000
 
-    # The number of bytes in the UART response that signify command status
-    STATUS_BYTES_LENGTH = 4
-
     # Bootloader flashing offset
     BOOTLOADER_FLASH_OFFSET = 0x0
 
@@ -518,36 +515,54 @@ def command(
         raise FatalError("Response doesn't match request.")
 
     def check_command(
-        self, op_description, op=None, data=b"", chk=0, timeout=DEFAULT_TIMEOUT
+        self,
+        op_description,
+        op=None,
+        data=b"",
+        chk=0,
+        resp_data_len=0,
+        timeout=DEFAULT_TIMEOUT,
     ):
         """
         Execute a command with 'command', check the result code and throw an appropriate
         FatalError if it fails.
 
         Returns the "result" of a successful command.
         """
-        val, data = self.command(op, data, chk, timeout=timeout)
 
-        # things are a bit weird here, bear with us
+        # The status bytes are first two bytes after the expected response data,
+        # ROM bootloaders of the chips (except ESP8266) return another 2 bytes,
+        # which are reserved for future use, these are ignored here.
+        STATUS_BYTES_LENGTH = 2
 
-        # the status bytes are the last 2/4 bytes in the data (depending on chip)
-        if len(data) < self.STATUS_BYTES_LENGTH:
-            raise FatalError(
-                f"Failed to {op_description}. "
-                f"Only got {len(data)} byte status response."
-            )
-        status_bytes = data[-self.STATUS_BYTES_LENGTH :]
-        # only care if the first one is non-zero. If it is, the second byte is a reason.
-        if byte(status_bytes, 0) != 0:
+        # Execute the command and get the result
+        val, data = self.command(op, data, chk, timeout=timeout)
+
+        # Check if we have enough data,
+        # including the expected response data and status bytes
+        if len(data) < resp_data_len + STATUS_BYTES_LENGTH:
+            # If we don't have enough data, check the first 2 bytes as status
+            status_bytes = data[0:2]
+            # Only care if the first byte is non-zero.
+            # If it is, the second byte is a reason.
+            if status_bytes[0] != 0:
+                raise FatalError.WithResult(f"Failed to {op_description}", status_bytes)
+            else:
+                raise FatalError(
+                    f"Failed to {op_description}. "
+                    f"Only got {len(data)} byte status response."
+                )
+        # The status bytes are positioned after the expected response data
+        # (first two bytes after resp_data_len are the status bytes)
+        status_bytes = data[resp_data_len : resp_data_len + STATUS_BYTES_LENGTH]
+        # Only care if the first byte is non-zero.
+        # If it is, the second byte is a reason.
+        if status_bytes[0] != 0:
             raise FatalError.WithResult(f"Failed to {op_description}", status_bytes)
 
-        # if we had more data than just the status bytes, return it as the result
-        # (this is used by the md5sum command, maybe other commands?)
-        if len(data) > self.STATUS_BYTES_LENGTH:
-            return data[: -self.STATUS_BYTES_LENGTH]
+        if resp_data_len > 0:
+            return data[:resp_data_len]
         else:
-            # otherwise, just return the 'val' field which comes from the reply header
-            # (this is used by read_reg)
             return val
 
     def flush_input(self):
@@ -856,17 +871,10 @@ def _post_connect(self):
 
     def read_reg(self, addr, timeout=DEFAULT_TIMEOUT):
         """Read memory address in target"""
-        # we don't call check_command here because read_reg() function is called
-        # when detecting chip type, and the way we check for success
-        # (STATUS_BYTES_LENGTH) is different for different chip types (!)
-        val, data = self.command(
-            self.ESP_CMDS["READ_REG"], struct.pack("<I", addr), timeout=timeout
+        command = struct.pack("<I", addr)
+        return self.check_command(
+            "read target memory", self.ESP_CMDS["READ_REG"], command, timeout=timeout
         )
-        if byte(data, 0) != 0:
-            raise FatalError.WithResult(
-                "Failed to read register address %08x" % addr, data
-            )
-        return val
 
     def write_reg(self, addr, value, mask=0xFFFFFFFF, delay_us=0, delay_after_us=0):
         """Write to memory address in target"""
@@ -1099,11 +1107,25 @@ def parse_security_flags(flags_value):
                 parsed_flags[flag_name] = (flags_value & flag_mask) != 0
             return parsed_flags
 
-        res = self.check_command(
-            "get security info", self.ESP_CMDS["GET_SECURITY_INFO"], b""
-        )
-        esp32s2 = True if len(res) == 12 else False
-        res = struct.unpack("<IBBBBBBBB" if esp32s2 else "<IBBBBBBBBII", res)
+        # Try with 20 bytes first (most chips), fall back to 12 bytes (ESP32-S2)
+        try:
+            res = self.check_command(
+                "get security info",
+                self.ESP_CMDS["GET_SECURITY_INFO"],
+                b"",
+                resp_data_len=20,
+            )
+            res = struct.unpack("<IBBBBBBBBII", res)
+            esp32s2 = False
+        except FatalError:
+            res = self.check_command(
+                "get security info",
+                self.ESP_CMDS["GET_SECURITY_INFO"],
+                b"",
+                resp_data_len=12,
+            )
+            res = struct.unpack("<IBBBBBBBB", res)
+            esp32s2 = True
 
         security_info = {
             "flags": res[0],
@@ -1347,20 +1369,21 @@ def flash_defl_finish(self, reboot=False):
     def flash_md5sum(self, addr, size):
         # the MD5 command returns additional bytes in the standard
         # command reply slot
+        RESP_DATA_LEN = 32
+        RESP_DATA_LEN_STUB = 16
         timeout = timeout_per_mb(MD5_TIMEOUT_PER_MB, size)
         res = self.check_command(
             "calculate md5sum",
             self.ESP_CMDS["SPI_FLASH_MD5"],
             struct.pack("<IIII", addr, size, 0, 0),
+            resp_data_len=RESP_DATA_LEN_STUB if self.IS_STUB else RESP_DATA_LEN,
             timeout=timeout,
         )
 
-        if len(res) == 32:
+        if not self.IS_STUB:
             return res.decode("utf-8")  # already hex formatted
-        elif len(res) == 16:
-            return hexify(res).lower()
         else:
-            raise FatalError("MD5Sum command returned unexpected result: %r" % res)
+            return hexify(res).lower()
 
     @stub_and_esp32_function_only
     def change_baud(self, baud):
@@ -1756,7 +1779,6 @@ class StubMixin:
     """
 
     FLASH_WRITE_SIZE = 0x4000  # Default value, can be overridden
-    STATUS_BYTES_LENGTH = 2
     IS_STUB = True
 
     def __init__(self, rom_loader):

esptool/targets/esp32.py

@@ -25,9 +25,6 @@ class ESP32ROM(ESPLoader):
     DROM_MAP_START = 0x3F400000
     DROM_MAP_END = 0x3F800000
 
-    # ESP32 uses a 4 byte status reply
-    STATUS_BYTES_LENGTH = 4
-
     SPI_REG_BASE = 0x3FF42000
     SPI_USR_OFFS = 0x1C
     SPI_USR1_OFFS = 0x20
@@ -393,6 +390,7 @@ def read_flash_slow(self, offset, length, progress_fn):
                     "read flash block",
                     self.ESP_CMDS["READ_FLASH_SLOW"],
                     struct.pack("<II", offset + len(data), block_len),
+                    resp_data_len=BLOCK_LEN,
                 )
             except FatalError:
                 log.note("Consider specifying the flash size argument.")

esptool/targets/esp8266.py

@@ -29,8 +29,6 @@ class ESP8266ROM(ESPLoader):
 
     UART_CLKDIV_REG = 0x60000014
 
-    STATUS_BYTES_LENGTH = 2
-
     XTAL_CLK_DIVIDER = 2
 
     FLASH_SIZES = {

test/test_esptool.py

@@ -1075,19 +1075,14 @@ def test_stub_reuse_with_synchronization(self):
         )  # do sync before (without reset it talks to the flasher stub)
         assert "Manufacturer:" in res
 
-    @pytest.mark.skipif(arg_chip != "esp8266", reason="ESP8266 only")
     def test_stub_reuse_without_synchronization(self):
         """
         Keep the flasher stub running and reuse it the next time
         without synchronization.
-
-        Synchronization is necessary for chips where the ROM bootloader has different
-        status length in comparison to the flasher stub.
-        Therefore, this is ESP8266 only test.
         """
         res = self.run_esptool("--after no-reset-stub flash-id")
         assert "Manufacturer:" in res
-        res = self.run_esptool("--before no-reset-no-sync flash-id")
+        res = self.run_esptool("--before no-reset-no-sync flash-id", preload=False)
         assert "Manufacturer:" in res