adafruit · hathach · Jul 23, 2021
diff --git a/apps/self_update/CMakeLists.txt b/apps/self_update/CMakeLists.txt
@@ -1,3 +1,3 @@
-idf_component_register(SRCS self_update.c bootloader_bin.c
+idf_component_register(SRCS self_update.c tinyuf2_bin.c
                     INCLUDE_DIRS "."
                     REQUIRES boards)
diff --git a/ports/esp32s2/Makefile b/ports/esp32s2/Makefile
@@ -60,9 +60,13 @@ $(BUILD)/combined.bin: app
 SELF_BUILD = apps/self_update/$(BUILD)
 
 $(SELF_BUILD)/update-tinyuf2.bin: app
-	$(PYTHON3) $(TOP)/lib/uf2/utils/uf2conv.py --carray $(BUILD)/tinyuf2.bin -o $(TOP)/apps/self_update/bootloader_bin.c
+	$(PYTHON3) $(TOP)/lib/uf2/utils/uf2conv.py --carray $(BUILD)/tinyuf2.bin -o $(TOP)/apps/self_update/tinyuf2_bin.c
+	$(PYTHON3) $(TOP)/lib/uf2/utils/uf2conv.py --carray $(BUILD)/bootloader/bootloader.bin -o boards/boot2_bin.h
+	# rename variable in boot2_bin to avoid conflict
+	@sed -i 's/bindata/binboot2/g' boards/boot2_bin.h
 	idf.py -C apps/self_update/ -B$(SELF_BUILD) -DBOARD=$(BOARD) app
-	@rm $(TOP)/apps/self_update/bootloader_bin.c
+	@rm $(TOP)/apps/self_update/tinyuf2_bin.c
+	@rm boards/boot2_bin.h
 
 $(SELF_BUILD)/update-tinyuf2.uf2: $(SELF_BUILD)/update-tinyuf2.bin
 	$(PYTHON3) $(TOP)/lib/uf2/utils/uf2conv.py -f 0xbfdd4eee -b 0x0000 -c -o $@ $^

diff --git a/ports/esp32s2/apps/self_update/CMakeLists.txt b/ports/esp32s2/apps/self_update/CMakeLists.txt
@@ -25,5 +25,6 @@ include($ENV{IDF_PATH}/tools/cmake/project.cmake)
 set(SUPPORTED_TARGETS esp32s2)
 
 add_compile_definitions(TINYUF2_SELF_UPDATE)
+add_compile_definitions(CONFIG_SPI_FLASH_DANGEROUS_WRITE_ALLOWED)
 
 project(update-tinyuf2)
diff --git a/ports/esp32s2/boards/board_flash.c b/ports/esp32s2/boards/board_flash.c
@@ -127,27 +127,115 @@ void board_flash_write (uint32_t addr, void const *data, uint32_t len)
 //--------------------------------------------------------------------+
 
 #ifdef TINYUF2_SELF_UPDATE
-void board_self_update(const uint8_t * bootloader_bin, uint32_t bootloader_len)
+
+// boot2 binary converted by uf2conv.py
+#include "boot2_bin.h"
+
+static inline uint32_t uf2_min32(uint32_t x, uint32_t y)
+{
+  return (x < y) ? x : y;
+}
+
+// Erase and update boot stage2
+static void update_boot2(const uint8_t * data, uint32_t datalen)
+{
+  // boot2 is always at 0x1000
+  enum { BOOT2_ADDR = 0x1000 };
+
+  // max size of boot2 is 28KB
+  if ( datalen > 24*1024u ) return;
+
+  esp_flash_t * flash_chip = esp_flash_default_chip;
+  assert(flash_chip != NULL);
+
+  //------------- Verify if content matches -------------//
+  uint8_t* verify_buf = _fl_buf;
+  bool content_matches = true;
+
+  for(uint32_t count = 0; count < datalen; count += FLASH_CACHE_SIZE)
+  {
+    uint32_t const verify_len = uf2_min32(datalen - count, FLASH_CACHE_SIZE);
+    esp_flash_read(flash_chip, verify_buf, BOOT2_ADDR + count, verify_len);
+
+    if ( 0 != memcmp(data + count, verify_buf, verify_len) )
+    {
+      content_matches = false;
+      break;
+    }
+  }
+
+  PRINT_INT(content_matches);
+
+  // nothing to do
+  if (content_matches) return;
+
+  //------------- Erase & Flash -------------//
+  enum { SECTOR_SZ = 4096UL };
+
+  // make len aligned to 4K (round div)
+  uint32_t const erase_sz = (datalen + SECTOR_SZ - 1) / SECTOR_SZ;
+
+  // erase
+  esp_flash_erase_region(flash_chip, BOOT2_ADDR, erase_sz);
+
+  // flash
+  esp_flash_write(flash_chip, data, BOOT2_ADDR, datalen);
+}
+
+// Erase and write tinyuf2 partition
+static void update_tinyuf2(const uint8_t * data, uint32_t datalen)
 {
+  esp_partition_t const * part_uf2 = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_FACTORY, NULL);
+  assert(part_uf2 != NULL);
+
+  // Set UF2 as next boot regardless of flashing resule to prevent running this app again
+  esp_ota_set_boot_partition(part_uf2);
+
+  //------------- Verify if content matches -------------//
+  uint8_t* verify_buf = _fl_buf;
+  bool content_matches = true;
+
+  for(uint32_t count = 0; count < datalen; count += FLASH_CACHE_SIZE)
+  {
+    uint32_t const verify_len = uf2_min32(datalen - count, FLASH_CACHE_SIZE);
+    esp_partition_read(part_uf2, count, verify_buf, verify_len);
+
+    if ( 0 != memcmp(data + count, verify_buf, verify_len) )
+    {
+      content_matches = false;
+      break;
+    }
+  }
+
+  // nothing to do
+  if (!content_matches) return;
+
+  //------------- Erase & Flash -------------//
   enum { SECTOR_SZ = 4096UL };
-  esp_partition_t const * _part_uf2;
 
-  _part_uf2 = esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_FACTORY, NULL);
-  assert(_part_uf2 != NULL);
+  // make len aligned to 4K (round div)
+  uint32_t const erase_sz = (datalen + SECTOR_SZ - 1) / SECTOR_SZ;
+
+  // Erase partition
+  esp_partition_erase_range(part_uf2, 0, erase_sz);
 
-  // make len aligned to 4K
-  uint32_t erase_sz = (bootloader_len & ~(SECTOR_SZ-1));
-  if (bootloader_len & (SECTOR_SZ-1)) erase_sz += SECTOR_SZ;
+  // Write new data
+  esp_partition_write(part_uf2, 0, data, datalen);
 
-  // Erase old bootloader
-  esp_partition_erase_range(_part_uf2, 0, erase_sz);
+  // Set UF2 as next boot
+  esp_ota_set_boot_partition(part_uf2);
+}
+
+void board_self_update(const uint8_t * bootloader_bin, uint32_t bootloader_len)
+{
+  // Update TinyUF2 partition
+  update_tinyuf2(bootloader_bin, bootloader_len);
 
-  // Write new bootloader
-  esp_partition_write(_part_uf2, 0, bootloader_bin, bootloader_len);
+  // Update boot2 stage partition
+  update_boot2(binboot2, binboot2_len);
 
-  // Set UF2 as next boot and restart
-  esp_ota_set_boot_partition(_part_uf2);
+  // all done restart
   esp_restart();
 }
-#endif
 
+#endif
diff --git a/ports/test_ghostfat/boards.c b/ports/test_ghostfat/boards.c
@@ -24,43 +24,53 @@
 // Selected option #3.
 
 //------------- Flash -------------//
-uint32_t board_flash_size(void) { return CFG_UF2_FLASH_SIZE; }
+uint32_t board_flash_size (void)
+{
+  return CFG_UF2_FLASH_SIZE;
+}
 
 // not supported
-void     board_flash_write(uint32_t addr, void const *data, uint32_t len) {
-    (void)addr;
-    (void)data;
-    (void)len;
+void board_flash_write (uint32_t addr, void const *data, uint32_t len)
+{
+  (void) addr;
+  (void) data;
+  (void) len;
 }
 
 // not supported
-void     board_self_update(const uint8_t * bootloader_bin, uint32_t bootloader_len) {
-    (void)bootloader_bin;
-    (void)bootloader_len;
+void board_self_update (const uint8_t *bootloader_bin, uint32_t bootloader_len)
+{
+  (void) bootloader_bin;
+  (void) bootloader_len;
 }
 
 // not supported
-void     board_flash_flush(void) {}
+void board_flash_flush (void)
+{
+}
 
 //------------- Interesting part of flash support for this test -------------//
-void     board_flash_read (uint32_t addr, void* buffer, uint32_t len) {
-    if ((addr & 7) != 0) {
-        // TODO - need to copy part of the first eight bytes
-        exit(1); // failure exit
-        addr += 8 - (addr & 7);
-    }
+void board_flash_read (uint32_t addr, void *buffer, uint32_t len)
+{
+  if ( (addr & 7) != 0 )
+  {
+    // TODO - need to copy part of the first eight bytes
+    exit(1);    // failure exit
+    addr += 8 - (addr & 7);
+  }
 
-    // EMBED address in each 32 bits of the FLASH
-    uint32_t * dest = buffer;
-    size_t incBytes = sizeof(*dest);
-    uint32_t currentAddress = addr;
+  // EMBED address in each 32 bits of the FLASH
+  uint32_t *dest = buffer;
+  size_t incBytes = sizeof(*dest);
+  uint32_t currentAddress = addr;
 
-    while (len >= incBytes) {
-        memcpy(dest, &currentAddress, incBytes); // unaligned memory possible
+  while ( len >= incBytes )
+  {
+    memcpy(dest, &currentAddress, incBytes);    // unaligned memory possible
 
-        len -= incBytes;
-        dest++;
-        currentAddress += incBytes;
-    }
+    len -= incBytes;
+    dest++;
+    currentAddress += incBytes;
+  }
 }