esp32: Improve support for OTA updates.

This commit adds several small items to improve the support for OTA
updates on an esp32:

- a partition table for 4MB flash modules that has two OTA partitions ready
  to go to do updates
- a GENERIC_OTA board that uses that partition table and that enables
  automatic roll-back in the bootloader
- a new esp32.Partition.mark_app_valid_cancel_rollback() class-method to
  signal that the boot is successful and should not be rolled back at the
  next reset
- an automated test for doing an OTA update
- documentation updates

docs/library/esp32.rst

@@ -65,7 +65,8 @@ Functions
 Flash partitions
 ----------------
 
-This class gives access to the partitions in the device's flash memory.
+This class gives access to the partitions in the device's flash memory and includes
+methods to enable over-the-air (OTA) updates.
 
 .. class:: Partition(id)
 
@@ -75,7 +76,8 @@ This class gives access to the partitions in the device's flash memory.
 .. classmethod:: Partition.find(type=TYPE_APP, subtype=0xff, label=None)
 
     Find a partition specified by *type*, *subtype* and *label*.  Returns a
-    (possibly empty) list of Partition objects.
+    (possibly empty) list of Partition objects. Note: ``subtype=0xff`` matches any subtype
+    and ``label=None`` matches any label.
 
 .. method:: Partition.info()
 
@@ -98,6 +100,19 @@ This class gives access to the partitions in the device's flash memory.
 .. method:: Partition.get_next_update()
 
     Gets the next update partition after this one, and returns a new Partition object.
+    Typical usage is ``Partition(Partition.RUNNING).get_next_update()``
+    which returns the next partition to update given the current running one.
+
+.. classmethod:: Partition.mark_app_valid_cancel_rollback()
+
+    Signals that the current boot is considered successful.
+    Calling ``mark_app_valid_cancel_rollback`` is required on the first boot of a new
+    partition to avoid an automatic rollback at the next boot.
+    This uses the ESP-IDF "app rollback" feature with "CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE"
+    and  an ``OSError(-261)`` is raised if called on firmware that doesn't have the
+    feature enabled.
+    It is OK to call ``mark_app_valid_cancel_rollback`` on every boot and it is not
+    necessary when booting firmare that was loaded using esptool.
 
 Constants
 ~~~~~~~~~
@@ -105,12 +120,16 @@ Constants
 .. data:: Partition.BOOT
           Partition.RUNNING
 
-    Used in the `Partition` constructor to fetch various partitions.
+    Used in the `Partition` constructor to fetch various partitions: ``BOOT`` is the
+    partition that will be booted at the next reset and ``RUNNING`` is the currently
+    running partition.
 
 .. data:: Partition.TYPE_APP
           Partition.TYPE_DATA
 
-    Used in `Partition.find` to specify the partition type.
+    Used in `Partition.find` to specify the partition type: ``APP`` is for bootable
+    firmware partitions (typically labelled ``factory``, ``ota_0``, ``ota_1``), and
+    ``DATA`` is for other partitions, e.g. ``nvs``, ``otadata``, ``phy_init``, ``vfs``.
 
 .. data:: HEAP_DATA
           HEAP_EXEC

ports/esp32/boards/GENERIC_OTA/mpconfigboard.h

@@ -0,0 +1,2 @@
+#define MICROPY_HW_BOARD_NAME "4MB/OTA module"
+#define MICROPY_HW_MCU_NAME "ESP32"

ports/esp32/boards/GENERIC_OTA/mpconfigboard.mk

@@ -0,0 +1,4 @@
+SDKCONFIG += boards/sdkconfig.base
+SDKCONFIG += boards/GENERIC_OTA/sdkconfig.board
+
+PART_SRC = partitions-ota.csv

ports/esp32/boards/GENERIC_OTA/sdkconfig.board

@@ -0,0 +1,4 @@
+CONFIG_BOOTLOADER_APP_ROLLBACK_ENABLE=y
+
+# ESP-IDF v3:
+CONFIG_APP_ROLLBACK_ENABLE=y

ports/esp32/esp32_partition.c

@@ -209,6 +209,15 @@ STATIC mp_obj_t esp32_partition_get_next_update(mp_obj_t self_in) {
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_1(esp32_partition_get_next_update_obj, esp32_partition_get_next_update);
 
+STATIC mp_obj_t esp32_partition_mark_app_valid_cancel_rollback(mp_obj_t cls_in) {
+    check_esp_err(esp_ota_mark_app_valid_cancel_rollback());
+    return mp_const_none;
+}
+STATIC MP_DEFINE_CONST_FUN_OBJ_1(esp32_partition_mark_app_valid_cancel_rollback_fun_obj,
+    esp32_partition_mark_app_valid_cancel_rollback);
+STATIC MP_DEFINE_CONST_CLASSMETHOD_OBJ(esp32_partition_mark_app_valid_cancel_rollback_obj,
+    MP_ROM_PTR(&esp32_partition_mark_app_valid_cancel_rollback_fun_obj));
+
 STATIC const mp_rom_map_elem_t esp32_partition_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_find), MP_ROM_PTR(&esp32_partition_find_obj) },
 
@@ -218,6 +227,7 @@ STATIC const mp_rom_map_elem_t esp32_partition_locals_dict_table[] = {
     { MP_ROM_QSTR(MP_QSTR_ioctl), MP_ROM_PTR(&esp32_partition_ioctl_obj) },
 
     { MP_ROM_QSTR(MP_QSTR_set_boot), MP_ROM_PTR(&esp32_partition_set_boot_obj) },
+    { MP_ROM_QSTR(MP_QSTR_mark_app_valid_cancel_rollback), MP_ROM_PTR(&esp32_partition_mark_app_valid_cancel_rollback_obj) },
     { MP_ROM_QSTR(MP_QSTR_get_next_update), MP_ROM_PTR(&esp32_partition_get_next_update_obj) },
 
     { MP_ROM_QSTR(MP_QSTR_BOOT), MP_ROM_INT(ESP32_PARTITION_BOOT) },

ports/esp32/partitions-ota.csv

@@ -0,0 +1,9 @@
+# Partition table for MicroPython with OTA support using 4MB flash
+# Name,   Type, SubType, Offset,  Size, Flags
+# Note: if you change the phy_init or app partition offset, make sure to change the offset in Kconfig.projbuild
+nvs,      data, nvs,     0x9000,   0x4000,
+otadata,  data, ota,     0xd000,   0x2000,
+phy_init, data, phy,     0xf000,   0x1000,
+ota_0,    app,  ota_0,   0x10000,  0x180000,
+ota_1,    app,  ota_1,   0x190000, 0x180000,
+vfs,      data, fat,     0x310000, 0x0f0000,

tests/esp32/partition_ota.py

@@ -0,0 +1,117 @@
+# Test ESP32 OTA updates, including automatic roll-back.
+# Running this test requires firmware with an OTA Partition, such as the GENERIC_OTA "board".
+# This test also requires patience as it copies the boot partition into the other OTA slot.
+
+import machine
+from esp32 import Partition
+
+# start by checking that the running partition table has OTA partitions, 'cause if
+# it doesn't there's nothing we can test
+cur = Partition(Partition.RUNNING)
+cur_name = cur.info()[4]
+if not cur_name.startswith("ota_"):
+    print("SKIP")
+    raise SystemExit
+
+DEBUG = True
+
+
+def log(*args):
+    if DEBUG:
+        print(*args)
+
+
+# replace boot.py with the test code that will run on each reboot
+import uos
+
+try:
+    uos.rename("boot.py", "boot-orig.py")
+except:
+    pass
+with open("boot.py", "w") as f:
+    f.write("DEBUG=" + str(DEBUG))
+    f.write(
+        """
+import machine
+from esp32 import Partition
+cur = Partition(Partition.RUNNING)
+cur_name = cur.info()[4]
+
+def log(*args):
+    if DEBUG: print(*args)
+
+from step import STEP, EXPECT
+log("Running partition: " + cur_name + " STEP=" + str(STEP) + " EXPECT=" + EXPECT)
+if cur_name != EXPECT:
+    print("\\x04FAILED: step " + str(STEP) + " expected " + EXPECT + " got " + cur_name + "\\x04")
+
+if STEP == 0:
+    log("Not confirming boot ok and resetting back into first")
+    nxt = cur.get_next_update()
+    with open("step.py", "w") as f:
+        f.write("STEP=1\\nEXPECT=\\"" + nxt.info()[4] + "\\"\\n")
+    machine.reset()
+elif STEP == 1:
+    log("Booting into second partition again")
+    nxt = cur.get_next_update()
+    nxt.set_boot()
+    with open("step.py", "w") as f:
+        f.write("STEP=2\\nEXPECT=\\"" + nxt.info()[4] + "\\"\\n")
+    machine.reset()
+elif STEP == 2:
+    log("Confirming boot ok and rebooting into same partition")
+    Partition.mark_app_valid_cancel_rollback()
+    with open("step.py", "w") as f:
+        f.write("STEP=3\\nEXPECT=\\"" + cur_name + "\\"\\n")
+    machine.reset()
+elif STEP == 3:
+    log("Booting into original partition")
+    nxt = cur.get_next_update()
+    nxt.set_boot()
+    with open("step.py", "w") as f:
+        f.write("STEP=4\\nEXPECT=\\"" + nxt.info()[4] + "\\"\\n")
+    machine.reset()
+elif STEP == 4:
+    log("Confirming boot ok and DONE!")
+    Partition.mark_app_valid_cancel_rollback()
+    import uos
+    uos.remove("step.py")
+    uos.remove("boot.py")
+    uos.rename("boot-orig.py", "boot.py")
+    print("\\nSUCCESS!\\n\\x04\\x04")
+
+"""
+    )
+
+
+def copy_partition(src, dest):
+    log("Partition copy: {} --> {}".format(src.info(), dest.info()))
+    sz = src.info()[3]
+    if dest.info()[3] != sz:
+        raise ValueError("Sizes don't match: {} vs {}".format(sz, dest.info()[3]))
+    addr = 0
+    blk = bytearray(4096)
+    while addr < sz:
+        if sz - addr < 4096:
+            blk = blk[: sz - addr]
+        if addr & 0xFFFF == 0:
+            # need to show progress to run-tests else it times out
+            print("   ... 0x{:06x}".format(addr))
+        src.readblocks(addr >> 12, blk)
+        dest.writeblocks(addr >> 12, blk)
+        addr += len(blk)
+
+
+# get things started by copying the current partition into the next slot and rebooting
+print("Copying current to next partition")
+nxt = cur.get_next_update()
+copy_partition(cur, nxt)
+print("Partition copied, booting into it")
+nxt.set_boot()
+
+# the step.py file is used to keep track of state across reboots
+# EXPECT is the name of the partition we expect to reboot into
+with open("step.py", "w") as f:
+    f.write('STEP=0\nEXPECT="' + nxt.info()[4] + '"\n')
+
+machine.reset()

tests/esp32/partition_ota.py.exp

@@ -0,0 +1,15 @@
+Copying current to next partition
+########
+Partition copied, booting into it
+########
+Not confirming boot ok and resetting back into first
+########
+Booting into second partition again
+########
+Confirming boot ok and rebooting into same partition
+########
+Booting into original partition
+########
+Confirming boot ok and DONE!
+
+SUCCESS!

tests/run-tests

@@ -56,6 +56,7 @@ def run_micropython(pyb, args, test_file, is_special=False):
     special_tests = (
         'micropython/meminfo.py', 'basics/bytes_compare3.py',
         'basics/builtin_help.py', 'thread/thread_exc2.py',
+        'esp32/partition_ota.py',
     )
     had_crash = False
     if pyb is None: