feat(erase_region): Enable erasing in ROM bootloader and SDM

Author: radimkarnis

esptool/cmds.py

@@ -1151,12 +1151,27 @@ def erase_flash(esp, args):
                 "please use with caution, otherwise it may brick your device!"
             )
     print("Erasing flash (this may take a while)...")
+    if esp.CHIP_NAME != "ESP8266" and not esp.IS_STUB:
+        print(
+            "Note: You can use the erase_region command in ROM bootloader "
+            "mode to erase a specific region."
+        )
     t = time.time()
     esp.erase_flash()
-    print("Chip erase completed successfully in %.1fs" % (time.time() - t))
+    print(f"Chip erase completed successfully in {time.time() - t:.1f} seconds.")
 
 
 def erase_region(esp, args):
+    if args.address % ESPLoader.FLASH_SECTOR_SIZE != 0:
+        raise FatalError(
+            "Offset to erase from must be a multiple "
+            f"of {ESPLoader.FLASH_SECTOR_SIZE}"
+        )
+    if args.size % ESPLoader.FLASH_SECTOR_SIZE != 0:
+        raise FatalError(
+            "Size of data to erase must be a multiple "
+            f"of {ESPLoader.FLASH_SECTOR_SIZE}"
+        )
     if not args.force and esp.CHIP_NAME != "ESP8266" and not esp.secure_download_mode:
         if esp.get_flash_encryption_enabled() or esp.get_secure_boot_enabled():
             raise FatalError(
@@ -1167,8 +1182,12 @@ def erase_region(esp, args):
             )
     print("Erasing region (may be slow depending on size)...")
     t = time.time()
-    esp.erase_region(args.address, args.size)
-    print("Erase completed successfully in %.1f seconds." % (time.time() - t))
+    if esp.CHIP_NAME != "ESP8266" and not esp.IS_STUB:
+        # flash_begin triggers a flash erase, enabling erasing in ROM and SDM
+        esp.flash_begin(args.size, args.address, logging=False)
+    else:
+        esp.erase_region(args.address, args.size)
+    print(f"Erase completed successfully in {time.time() - t:.1f} seconds.")
 
 
 def run(esp, args):

esptool/loader.py

@@ -891,7 +891,7 @@ def mem_finish(self, entrypoint=0):
                 raise
             pass
 
-    def flash_begin(self, size, offset, begin_rom_encrypted=False):
+    def flash_begin(self, size, offset, begin_rom_encrypted=False, logging=True):
         """
         Start downloading to Flash (performs an erase)
 
@@ -916,7 +916,7 @@ def flash_begin(self, size, offset, begin_rom_encrypted=False):
         self.check_command(
             "enter Flash download mode", self.ESP_FLASH_BEGIN, params, timeout=timeout
         )
-        if size != 0 and not self.IS_STUB:
+        if size != 0 and not self.IS_STUB and logging:
             print("Took %.2fs to erase flash block" % (time.time() - t))
         return num_blocks
 
@@ -1196,10 +1196,6 @@ def erase_flash(self):
 
     @stub_function_only
     def erase_region(self, offset, size):
-        if offset % self.FLASH_SECTOR_SIZE != 0:
-            raise FatalError("Offset to erase from must be a multiple of 4096")
-        if size % self.FLASH_SECTOR_SIZE != 0:
-            raise FatalError("Size of data to erase must be a multiple of 4096")
         timeout = timeout_per_mb(ERASE_REGION_TIMEOUT_PER_MB, size)
         self.check_command(
             "erase region",

test/test_esptool.py

@@ -1067,6 +1067,18 @@ def test_large_region_erase(self):
         # verifies that erasing a large region doesn't time out
         self.run_esptool("erase_region 0x0 0x100000")
 
+    @pytest.mark.skipif(arg_chip == "esp8266", reason="Not supported on ESP8266")
+    def test_region_erase_no_stub(self):
+        self.run_esptool("write_flash 0x10000 images/one_kb.bin")
+        self.run_esptool("write_flash 0x11000 images/sector.bin")
+        self.verify_readback(0x10000, 0x400, "images/one_kb.bin")
+        self.verify_readback(0x11000, 0x1000, "images/sector.bin")
+        # erase only the flash sector containing one_kb.bin
+        self.run_esptool("--no-stub erase_region 0x10000 0x1000")
+        self.verify_readback(0x11000, 0x1000, "images/sector.bin")
+        empty = self.readback(0x10000, 0x1000)
+        assert empty == b"\xFF" * 0x1000
+
 
 class TestSectorBoundaries(EsptoolTestCase):
     def test_end_sector(self):