feat(core/bootloader): make firmware flashing more safe

core/embed/bootloader/messages.c

@@ -349,13 +349,17 @@ void process_msg_Ping(uint8_t iface_num, uint32_t msg_size, uint8_t *buf) {
   MSG_SEND(Success);
 }
 
-static uint32_t firmware_remaining, firmware_block, chunk_requested;
+static uint32_t firmware_remaining;
+static uint32_t firmware_block;
+static uint32_t chunk_requested;
+static uint32_t erase_offset;
 
 void process_msg_FirmwareErase(uint8_t iface_num, uint32_t msg_size,
                                uint8_t *buf) {
   firmware_remaining = 0;
   firmware_block = 0;
   chunk_requested = 0;
+  erase_offset = 0;
 
   MSG_RECV_INIT(FirmwareErase);
   MSG_RECV(FirmwareErase);
@@ -411,8 +415,8 @@ static bool _read_payload(pb_istream_t *stream, const pb_field_t *field,
     // update loader but skip first block
     if (firmware_block > 0) {
       ui_screen_install_progress_upload(
-          250 + 750 * (firmware_block * IMAGE_CHUNK_SIZE + chunk_written) /
-                    (firmware_block * IMAGE_CHUNK_SIZE + firmware_remaining));
+          1000 * (firmware_block * IMAGE_CHUNK_SIZE + chunk_written) /
+          (firmware_block * IMAGE_CHUNK_SIZE + firmware_remaining));
     }
     // read data
     if (!pb_read(
@@ -660,8 +664,6 @@ int process_msg_FirmwareUpload(uint8_t iface_num, uint32_t msg_size,
         ensure(flash_area_erase_bulk(STORAGE_AREAS, STORAGE_AREAS_COUNT, NULL),
                NULL);
       }
-      ensure(flash_area_erase(&FIRMWARE_AREA, ui_screen_install_progress_erase),
-             NULL);
 
       headers_offset = IMAGE_HEADER_SIZE + vhdr.hdrlen;
       read_offset = IMAGE_INIT_CHUNK_SIZE;
@@ -713,22 +715,60 @@ int process_msg_FirmwareUpload(uint8_t iface_num, uint32_t msg_size,
     return UPLOAD_ERR_INVALID_CHUNK_HASH;
   }
 
-  ensure(flash_unlock_write(), NULL);
+  // buffer with the received data
+  const uint32_t *quadword_ptr = (const uint32_t *)CHUNK_BUFFER_PTR;
+  // number of received bytes
+  uint32_t bytes_remaining = chunk_size;
+  // offset into the FIRMWARE_AREA part of the flash 
+  uint32_t write_offset = firmware_block * IMAGE_CHUNK_SIZE;
 
-  const uint32_t *const src = (const uint32_t *const)CHUNK_BUFFER_PTR;
+  while (bytes_remaining > 0) {
+    // erase flash before writing
+    uint32_t bytes_erased = 0; 
+    
+    if (write_offset >= erase_offset) {
+      // erase the next flash section
+      ensure(
+          flash_area_erase_partial(&FIRMWARE_AREA, erase_offset, &bytes_erased),
+          NULL);
+      erase_offset += bytes_erased;
+    } else {
+      // some erased space left from the previous round => use it
+      bytes_erased = erase_offset - write_offset;
+    }
+  
+    // write the received data
+    uint32_t bytes_to_write = MIN(bytes_erased, bytes_remaining);
+    uint32_t write_end = write_offset + bytes_to_write;
 
-  for (int i = 0; i < chunk_size / (sizeof(uint32_t) * 4); i++) {
-    ensure(flash_area_write_quadword(
-               &FIRMWARE_AREA,
-               firmware_block * IMAGE_CHUNK_SIZE + i * 4 * sizeof(uint32_t),
-               &src[4 * i]),
-           NULL);
+    ensure(flash_unlock_write(), NULL);
+    while (write_offset < write_end) {
+      // write a quad word (16 bytes) to the flash
+      ensure(
+          flash_area_write_quadword(&FIRMWARE_AREA, write_offset, quadword_ptr),
+          NULL);
+      write_offset += 4 * sizeof(uint32_t);
+      quadword_ptr += 4;
+    }
+    ensure(flash_lock_write(), NULL);
+
+    bytes_remaining -= bytes_to_write;
   }
 
-  ensure(flash_lock_write(), NULL);
+  firmware_remaining -= chunk_requested;
+
+  if (firmware_remaining == 0) {
+    // erase the rest (unused part) of the FIRMWARE_AREA
+    uint32_t bytes_erased = 0;
+    do {
+      ensure(
+          flash_area_erase_partial(&FIRMWARE_AREA, erase_offset, &bytes_erased),
+          NULL);
+      erase_offset += bytes_erased;
+    } while (bytes_erased > 0);
+  }
 
   headers_offset = 0;
-  firmware_remaining -= chunk_requested;
   firmware_block++;
   firmware_upload_chunk_retry = FIRMWARE_UPLOAD_CHUNK_RETRY_COUNT;
 

core/embed/trezorhal/flash.h

@@ -57,6 +57,14 @@ secbool flash_write_word(uint16_t sector, uint32_t offset, uint32_t data);
 
 uint32_t flash_wait_and_clear_status_flags(void);
 
+// Erases the single sector in the designated flash area
+// The 'offset' parameter must indicate the relative sector offset within the
+// flash area If 'offset' is outside the bounds of the flash area,
+// 'bytes_erased' is set to 0 otherwise, 'bytes_erased' is set to the size of
+// the erased sector
+secbool flash_area_erase_partial(const flash_area_t *area, uint32_t offset,
+                                 uint32_t *bytes_erased);
+
 secbool __wur flash_otp_read(uint8_t block, uint8_t offset, uint8_t *data,
                              uint8_t datalen);
 secbool __wur flash_otp_write(uint8_t block, uint8_t offset,

core/embed/trezorhal/stm32f4/flash.c

@@ -167,6 +167,50 @@ secbool flash_area_erase_bulk(const flash_area_t *area, int count,
   return sectrue;
 }
 
+secbool flash_area_erase_partial(const flash_area_t *area, uint32_t offset,
+                                 uint32_t *bytes_erased) {
+  uint32_t sector_offset = 0;
+  *bytes_erased = 0;
+
+  for (int s = 0; s < area->num_subareas; s++) {
+    for (int i = 0; i < area->subarea[s].num_sectors; i++) {
+      uint32_t sector_index = area->subarea[s].first_sector + i;
+      uint32_t sector_size = FLASH_SECTOR_TABLE[sector_index + 1] -
+                             FLASH_SECTOR_TABLE[sector_index];
+
+      if (offset == sector_offset) {
+        ensure(flash_unlock_write(), NULL);
+
+        FLASH_EraseInitTypeDef erase_init = {
+            .TypeErase = FLASH_TYPEERASE_SECTORS,
+            .VoltageRange = FLASH_VOLTAGE_RANGE_3,
+            .Sector = sector_index,
+            .NbSectors = 1};
+
+        uint32_t sector_error;
+
+        if (HAL_FLASHEx_Erase(&erase_init, &sector_error) != HAL_OK) {
+          ensure(flash_lock_write(), NULL);
+          return secfalse;
+        }
+
+        ensure(flash_lock_write(), NULL);
+
+        *bytes_erased = sector_size;
+        return sectrue;
+      }
+
+      sector_offset += sector_size;
+    }
+  }
+
+  if (offset == sector_offset) {
+    return sectrue;
+  }
+
+  return secfalse;
+}
+
 secbool flash_write_byte(uint16_t sector, uint32_t offset, uint8_t data) {
   uint32_t address = (uint32_t)flash_get_address(sector, offset, 1);
   if (address == 0) {

core/embed/trezorhal/unix/flash.c

@@ -190,6 +190,36 @@ secbool flash_area_erase_bulk(const flash_area_t *area, int count,
   return sectrue;
 }
 
+secbool flash_area_erase_partial(const flash_area_t *area, uint32_t offset,
+                                 uint32_t *bytes_erased) {
+  uint32_t sector_offset = 0;
+  *bytes_erased = 0;
+
+  for (int s = 0; s < area->num_subareas; s++) {
+    for (int i = 0; i < area->subarea[s].num_sectors; i++) {
+      uint32_t sector_index = area->subarea[s].first_sector + i;
+      uint32_t sector_size = FLASH_SECTOR_TABLE[sector_index + 1] -
+                             FLASH_SECTOR_TABLE[sector_index];
+
+      if (offset == sector_offset) {
+        uint8_t *flash =
+            (uint8_t *)flash_get_address(sector_index, 0, sector_size);
+        memset(flash, 0xFF, sector_size);
+        *bytes_erased = sector_size;
+        return sectrue;
+      }
+
+      sector_offset += sector_size;
+    }
+  }
+
+  if (offset == sector_offset) {
+    return sectrue;
+  }
+
+  return secfalse;
+}
+
 secbool flash_write_byte(uint16_t sector, uint32_t offset, uint8_t data) {
   uint8_t *flash = (uint8_t *)flash_get_address(sector, offset, 1);
   if (!flash) {