chore(radio): improve UF2 (#6644)

radio/src/drivers/uf2_ghostfat.cpp

@@ -185,8 +185,17 @@ static_assert(ARRAY_SIZE(indexFile) < 512);
 #define NUM_FILES (ARRAY_SIZE(info))
 #define NUM_DIRENTRIES (NUM_FILES + 1) // Code adds volume label as first root directory entry
 
+#ifndef BOOTLOADER_ADDRESS
+#define BOOTLOADER_ADDRESS FIRMWARE_ADDRESS
+#endif
 
-#define UF2_SIZE           (current_flash_size() * 2 + 512)
+#if BOOTLOADER_ADDRESS == FIRMWARE_ADDRESS
+#define REBOOT_BLOCK 0
+#else
+#define REBOOT_BLOCK 1
+#endif
+
+#define UF2_SIZE           (current_flash_size() * 2 + 512 * REBOOT_BLOCK)
 #define UF2_SECTORS        (UF2_SIZE / 512)
 #define UF2_FIRST_SECTOR   (NUM_FILES + 1) // WARNING -- code presumes each non-UF2 file content fits in single sector
 #define UF2_LAST_SECTOR    (UF2_FIRST_SECTOR + UF2_SECTORS - 1)
@@ -206,7 +215,6 @@ static_assert(ARRAY_SIZE(indexFile) < 512);
 
 static_assert(NUM_DIRENTRIES < DIRENTRIES_PER_SECTOR * ROOT_DIR_SECTORS);
 
-
 static FAT_BootBlock const BootBlock = {
     .JumpInstruction      = {0xeb, 0x3c, 0x90},
     .OEMInfo              = {'M','S','W','I','N','4','.','1'},
@@ -271,7 +279,7 @@ static uint32_t current_flash_size(void)
   // only need to compute once
   if ( result == 0 ) {
     firmware_description_t const *fw_desc =
-        (firmware_description_t const *)FIRMWARE_ADDRESS;
+        (firmware_description_t const *)APP_START_ADDRESS;
 
     if (is_firmware_valid(fw_desc)) {
       // round up to 256 bytes
@@ -374,26 +382,26 @@ void uf2_fat_read_block(uint32_t block_no, uint8_t *data)
             sectionIdx -= NUM_FILES - 1;
 
             uint32_t addr = 0;
-            if (sectionIdx >= UF2_MAX_FW_SIZE / 256) return;
+            if (sectionIdx >= UF2_MAX_FW_SIZE / 256 + REBOOT_BLOCK) return;
 
             UF2_Block *bl = (UF2_Block *)data;
             bl->magicStart0 = UF2_MAGIC_START0;
             bl->magicStart1 = UF2_MAGIC_START1;
             bl->magicEnd = UF2_MAGIC_END;
             bl->blockNo = sectionIdx;
-            bl->numBlocks = current_flash_size() / 256 + 1;
+            bl->numBlocks = current_flash_size() / 256 + REBOOT_BLOCK;
             bl->flags = 0;
             bl->reserved = 0;
 
-            if (sectionIdx == BOOTLOADER_SIZE / 256) {
+            if (REBOOT_BLOCK && (sectionIdx == BOOTLOADER_SIZE / 256)) {
                 writeUF2RebootBlock(bl);
                 writeUF2FirmwareVersion(bl);
             } else {
                 if (sectionIdx < BOOTLOADER_SIZE / 256) {
                     addr = BOOTLOADER_ADDRESS + sectionIdx * 256;
                 } else {
-                    sectionIdx -= BOOTLOADER_SIZE / 256 + 1;
-                    addr = FIRMWARE_ADDRESS + sectionIdx * 256;
+                    sectionIdx -= BOOTLOADER_SIZE / 256 + REBOOT_BLOCK;
+                    addr = APP_START_ADDRESS + sectionIdx * 256;
                 }
 
                 bl->targetAddr = addr;
@@ -439,7 +447,7 @@ int uf2_fat_write_block(uint32_t block_no, uint8_t *data)
         // to reset properly
     } else {
         uint32_t wr_block = bl->blockNo;
-        if (wr_block >= BOOTLOADER_SIZE / 256) wr_block--;
+        if (REBOOT_BLOCK && (wr_block >= BOOTLOADER_SIZE / 256)) wr_block--;
 
         uint32_t erase_sector = wr_block / (UF2_ERASE_BLOCK_SIZE / 256);
         uint32_t mask = 1 << (erase_sector & 0x1F);

radio/src/io/uf2.cpp

@@ -31,6 +31,10 @@
 #include "thirdparty/FatFs/ff.h"
 #include "thirdparty/uf2/uf2.h"
 
+#ifndef BOOTLOADER_ADDRESS
+#define BOOTLOADER_ADDRESS FIRMWARE_ADDRESS
+#endif
+
 #define UF2_BLOCK_SIZE 512
 
 bool isUF2Block(const void* block, uint32_t len)

radio/util/elf2uf2.py

@@ -4,7 +4,7 @@ import argparse
 import struct
 import sys
 
-from typing import BinaryIO, Literal, Optional
+from typing import BinaryIO, Literal, Optional, Iterator
 from elftools.elf.elffile import ELFFile
 
 DATA_BLOCK_SIZE = 256
@@ -74,7 +74,7 @@ class UF2Writer:
     def _write_block(
         self,
         addr: int,
-        payload: bytes,
+        payload: bytes | bytearray,
         nblocks: int,
         /,
         flags=0,
@@ -146,7 +146,7 @@ class ELFReader:
 
         return segments
 
-    def iter_blocks(self, chunks, *, block_size=DATA_BLOCK_SIZE):
+    def iter_blocks(self, chunks, *, block_size=DATA_BLOCK_SIZE) -> Iterator[bytearray]:
         buffer = bytearray()
         for i, size in chunks:
             # fill buffer first
@@ -235,21 +235,27 @@ def elf_to_uf2(args: argparse.Namespace) -> None:
 
     insert_block = None
     if args.reboot is not None:
+        insert = None
+        addr = None
+        start = None
+
         if args.reboot == "auto":
-            addr, _ = reader.get_section(".reboot_buffer")
-            start, bl_size = reader.get_section(".bootloader")
-            insert = start + bl_size
+            if len(reader.segments) > 1:
+                addr, _ = reader.get_section(".reboot_buffer")
+                start, bl_size = reader.get_section(".bootloader")
+                insert = start + bl_size
         else:
             insert, addr, start = args.reboot
 
-        print(f"Insert reboot block @{insert:08X} ({addr:08X}/{start:08X})")
-        insert_block = UF2InsertBlock(
-            insert,
-            addr,
-            b"BDFU",
-            extensions=[UF2Extension(UF2_REBOOT_TAG, start.to_bytes(4, "little"))],
-            flags=1,
-        )
+        if insert and addr and start:
+            print(f"Insert reboot block @{insert:08X} ({addr:08X}/{start:08X})")
+            insert_block = UF2InsertBlock(
+                insert,
+                addr,
+                b"BDFU",
+                extensions=[UF2Extension(UF2_REBOOT_TAG, start.to_bytes(4, "little"))],
+                flags=1,
+            )
 
     uf2 = UF2Writer(
         args.output, family=args.family, device=args.device, version=args.version