OTA: Fix issues with encrypted OTA

- OTA source can write non-16-byte multiples of data
- Assumption that empty ota_data is 0xFFFFFFFF untrue when encrypted

components/app_update/esp_ota_ops.c

@@ -27,6 +27,7 @@
 #include "esp_spi_flash.h"
 #include "esp_image_format.h"
 #include "esp_secure_boot.h"
+#include "esp_flash_encrypt.h"
 #include "sdkconfig.h"
 
 #include "esp_ota_ops.h"
@@ -44,6 +45,10 @@ typedef struct ota_ops_entry_ {
     esp_partition_t part;
     uint32_t erased_size;
     uint32_t wrote_size;
+#ifdef CONFIG_FLASH_ENCRYPTION_ENABLED
+    uint8_t partial_bytes;
+    uint8_t partial_data[16];
+#endif
     LIST_ENTRY(ota_ops_entry_) entries;
 } ota_ops_entry_t;
 
@@ -126,7 +131,41 @@ esp_err_t esp_ota_write(esp_ota_handle_t handle, const void *data, size_t size)
                 return ESP_ERR_OTA_VALIDATE_FAILED;
             }
 
-            ret = esp_partition_write(&it->part, it->wrote_size, data, size);
+#ifdef CONFIG_FLASH_ENCRYPTION_ENABLED
+            if (esp_flash_encryption_enabled()) {
+                /* Can only write 16 byte blocks to flash, so need to cache anything else */
+                size_t copy_len;
+
+                /* check if we have partially written data from earlier */
+                if (it->partial_bytes != 0) {
+                    copy_len = OTA_MIN(16 - it->partial_bytes, size);
+                    memcpy(it->partial_data + it->partial_bytes, data_bytes, copy_len);
+                    it->partial_bytes += copy_len;
+                    if (it->partial_bytes != 16) {
+                        return ESP_OK; /* nothing to write yet, just filling buffer */
+                    }
+                    /* write 16 byte to partition */
+                    ret = esp_partition_write(&it->part, it->wrote_size, it->partial_data, 16);
+                    if (ret != ESP_OK) {
+                        return ret;
+                    }
+                    it->partial_bytes = 0;
+                    memset(it->partial_data, 0xFF, 16);
+                    it->wrote_size += 16;
+                    data_bytes += copy_len;
+                    size -= copy_len;
+                }
+
+                /* check if we need to save trailing data that we're about to write */
+                it->partial_bytes = size % 16;
+                if (it->partial_bytes != 0) {
+                    size -= it->partial_bytes;
+                    memcpy(it->partial_data, data_bytes + size, it->partial_bytes);
+                }
+            }
+#endif
+
+            ret = esp_partition_write(&it->part, it->wrote_size, data_bytes, size);
             if(ret == ESP_OK){
                 it->wrote_size += size;
             }
@@ -143,6 +182,8 @@ esp_err_t esp_ota_end(esp_ota_handle_t handle)
 {
     ota_ops_entry_t *it;
     size_t image_size;
+    esp_err_t __attribute__((unused)) ret;
+
     for (it = LIST_FIRST(&s_ota_ops_entries_head); it != NULL; it = LIST_NEXT(it, entries)) {
         if (it->handle == handle) {
             // an ota handle need to be ended after erased and wrote data in it
@@ -150,6 +191,18 @@ esp_err_t esp_ota_end(esp_ota_handle_t handle)
                 return ESP_ERR_INVALID_ARG;
             }
 
+#ifdef CONFIG_FLASH_ENCRYPTION_ENABLED
+            if (it->partial_bytes > 0 && esp_flash_encryption_enabled()) {
+                /* Write out last 16 bytes, if necessary */
+                ret = esp_partition_write(&it->part, it->wrote_size, it->partial_data, 16);
+                if (ret != ESP_OK) {
+                    return ret;
+                }
+                it->wrote_size += 16;
+                it->partial_bytes = 0;
+            }
+#endif
+
             if (esp_image_basic_verify(it->part.address, true, &image_size) != ESP_OK) {
                 return ESP_ERR_OTA_VALIDATE_FAILED;
             }
@@ -279,11 +332,9 @@ static esp_err_t esp_rewrite_ota_data(esp_partition_subtype_t subtype)
             }
             return rewrite_ota_seq((SUB_TYPE_ID(subtype) + 1) % ota_app_count + i * ota_app_count, 0, find_partition);
 
-        } else if (s_ota_select[0].ota_seq == 0xFFFFFFFF && s_ota_select[1].ota_seq == 0xFFFFFFFF) {
-            return rewrite_ota_seq(SUB_TYPE_ID(subtype) + 1, 0, find_partition);
-
         } else {
-            return ESP_ERR_OTA_SELECT_INFO_INVALID;
+            /* Both OTA slots are invalid, probably because unformatted... */
+            return rewrite_ota_seq(SUB_TYPE_ID(subtype) + 1, 0, find_partition);
         }
 
     } else {
@@ -385,7 +436,7 @@ const esp_partition_t *esp_ota_get_boot_partition(void)
                                         ESP_PARTITION_SUBTYPE_APP_OTA_MIN + (s_ota_select[1].ota_seq - 1) % ota_app_count, NULL);
 
     } else {
-        ESP_LOGE(TAG, "ota data invalid, no current app. Falling back to factory");
+        ESP_LOGE(TAG, "ota data invalid, no current app. Assuming factory");
         return esp_partition_find_first(ESP_PARTITION_TYPE_APP, ESP_PARTITION_SUBTYPE_APP_FACTORY, NULL);
     }
 }

components/bootloader_support/src/bootloader_flash.c

@@ -16,6 +16,7 @@
 #include <bootloader_flash.h>
 #include <esp_log.h>
 #include <esp_spi_flash.h> /* including in bootloader for error values */
+#include <esp_flash_encrypt.h>
 
 #ifndef BOOTLOADER_BUILD
 /* Normal app version maps to esp_spi_flash.h operations...
@@ -48,7 +49,11 @@ void bootloader_munmap(const void *mapping)
 
 esp_err_t bootloader_flash_read(size_t src, void *dest, size_t size, bool allow_decrypt)
 {
-    return spi_flash_read(src, dest, size);
+    if (allow_decrypt && esp_flash_encryption_enabled()) {
+        return spi_flash_read_encrypted(src, dest, size);
+    } else {
+        return spi_flash_read(src, dest, size);
+    }
 }
 
 esp_err_t bootloader_flash_write(size_t dest_addr, void *src, size_t size, bool write_encrypted)

components/spi_flash/partition.c

@@ -166,10 +166,14 @@ static esp_err_t load_partitions()
         item->info.type = it->type;
         item->info.subtype = it->subtype;
         item->info.encrypted = it->flags & PART_FLAG_ENCRYPTED;
-        if (esp_flash_encryption_enabled() && it->type == PART_TYPE_APP) {
-            /* All app partitions are encrypted if encryption is turned on */
+        if (esp_flash_encryption_enabled() && (
+                it->type == PART_TYPE_APP
+                || (it->type == PART_TYPE_DATA && it->subtype == PART_SUBTYPE_DATA_OTA))) {
+            /* If encryption is turned on, all app partitions and OTA data
+               are always encrypted */
             item->info.encrypted = true;
         }
+
         // it->label may not be zero-terminated
         strncpy(item->info.label, (const char*) it->label, sizeof(it->label));
         item->info.label[sizeof(it->label)] = 0;

examples/system/ota/main/ota.c

@@ -124,7 +124,7 @@ static bool read_past_http_header(char text[], int total_len, esp_ota_handle_t o
 
             esp_err_t err = esp_ota_write( out_handle, (const void *)ota_write_data, i_write_len);
             if (err != ESP_OK) {
-                ESP_LOGE(TAG, "Error: esp_ota_write failed! err=%x", err);
+                ESP_LOGE(TAG, "Error: esp_ota_write failed! err=0x%x", err);
                 return false;
             } else {
                 ESP_LOGI(TAG, "esp_ota_write header OK");
@@ -275,7 +275,7 @@ void main_task(void *pvParameter)
             memcpy(ota_write_data, text, buff_len);
             err = esp_ota_write( out_handle, (const void *)ota_write_data, buff_len);
             if (err != ESP_OK) {
-                ESP_LOGE(TAG, "Error: esp_ota_write failed! err=%x", err);
+                ESP_LOGE(TAG, "Error: esp_ota_write failed! err=0x%x", err);
                 task_fatal_error();
             }
             binary_file_length += buff_len;