feat(write_flash): retry flashing if chip disconnects

esptool/cmds.py

@@ -13,6 +13,7 @@ import zlib
 import itertools
 
 from intelhex import IntelHex
+from serial import SerialException
 
 from .bin_image import ELFFile, ImageSegment, LoadFirmwareImage
 from .bin_image import (
@@ -579,12 +580,19 @@ def write_flash(esp, args):
         if compress:
             uncimage = image
             image = zlib.compress(uncimage, 9)
+        original_image = image  # Save the whole image in case retry is needed
+        # Try again if reconnect was successful
+        for attempt in range(1, esp.WRITE_FLASH_ATTEMPTS + 1):
+            try:
+                if compress:
                     # Decompress the compressed binary a block at a time,
                     # to dynamically calculate the timeout based on the real write size
                     decompress = zlib.decompressobj()
                     blocks = esp.flash_defl_begin(uncsize, len(image), address)
                 else:
-            blocks = esp.flash_begin(uncsize, address, begin_rom_encrypted=encrypted)
+                    blocks = esp.flash_begin(
+                        uncsize, address, begin_rom_encrypted=encrypted
+                    )
                 argfile.seek(0)  # in case we need it again
                 seq = 0
                 bytes_sent = 0  # bytes sent on wire
@@ -607,12 +615,12 @@ def write_flash(esp, args):
                         bytes_written += block_uncompressed
                         block_timeout = max(
                             DEFAULT_TIMEOUT,
-                    timeout_per_mb(ERASE_WRITE_TIMEOUT_PER_MB, block_uncompressed),
+                            timeout_per_mb(
+                                ERASE_WRITE_TIMEOUT_PER_MB, block_uncompressed
+                            ),
                         )
                         if not esp.IS_STUB:
-                    timeout = (
-                        block_timeout  # ROM code writes block to flash before ACKing
-                    )
+                            timeout = block_timeout  # ROM code writes block to flash before ACKing
                         esp.flash_defl_block(block, seq, timeout=timeout)
                         if esp.IS_STUB:
                             # Stub ACKs when block is received,
@@ -629,6 +637,32 @@ def write_flash(esp, args):
                     bytes_sent += len(block)
                     image = image[esp.FLASH_WRITE_SIZE :]
                     seq += 1
+                break
+            except SerialException:
+                if attempt == esp.WRITE_FLASH_ATTEMPTS or encrypted:
+                    # Already retried once or encrypted mode is disabled because of security reasons
+                    raise
+                print("\nLost connection, retrying...")
+                esp._port.close()
+                print("Waiting for the chip to reconnect", end="")
+                for _ in range(DEFAULT_CONNECT_ATTEMPTS):
+                    try:
+                        time.sleep(1)
+                        esp._port.open()
+                        print()  # Print new line which was suppressed by print(".")
+                        esp.connect()
+                        if esp.IS_STUB:
+                            # Hack to bypass the stub overwrite check
+                            esp.IS_STUB = False
+                            # Reflash stub because chip was reset
+                            esp = esp.run_stub()
+                        image = original_image
+                        break
+                    except SerialException:
+                        print(".", end="")
+                        sys.stdout.flush()
+                else:
+                    raise  # Reconnect limit reached
 
         if esp.IS_STUB:
             # Stub only writes each block to flash after 'ack'ing the receive,

esptool/loader.py

@@ -278,6 +278,9 @@ class ESPLoader(object):
     # Chip IDs that are no longer supported by esptool
     UNSUPPORTED_CHIPS = {6: "ESP32-S3(beta 3)"}
 
+    # Number of attempts to write flash data
+    WRITE_FLASH_ATTEMPTS = 2
+
     def __init__(self, port=DEFAULT_PORT, baud=ESP_ROM_BAUD, trace_enabled=False):
         """Base constructor for ESPLoader bootloader interaction