ota: move valid/confirmed flag and fw-written value into ota ctx struct

Remove wrapping context structs and pass main ota ctx into all relevant
ota functions and callbacks, so that all interesting ota progress data
is consolidated into the one structure.

Author: Jamie C. Driver

main/process/ota.c

@@ -21,54 +21,52 @@
 
 #include "process_utils.h"
 
-typedef struct {
-    bool* prevalidated;
-    jade_ota_ctx_t* joctx;
-} ota_deflate_ctx_t;
-
 /* this is called by the deflate library when it has uncompressed data to write */
 static int uncompressed_stream_writer(void* ctx, uint8_t* uncompressed, size_t length)
 {
     JADE_ASSERT(ctx);
     JADE_ASSERT(uncompressed);
     JADE_ASSERT(length);
 
-    ota_deflate_ctx_t* octx = (ota_deflate_ctx_t*)ctx;
-    JADE_ASSERT(octx->joctx);
+    jade_ota_ctx_t* joctx = (jade_ota_ctx_t*)ctx;
 
-    if (!*octx->prevalidated && length >= CUSTOM_HEADER_MIN_WRITE) {
-        const enum ota_status res = ota_user_validation(octx->joctx, uncompressed);
+    if (!*joctx->validated_confirmed && length >= CUSTOM_HEADER_MIN_WRITE) {
+        const enum ota_status res = ota_user_validation(joctx, uncompressed);
         if (res != SUCCESS) {
             JADE_LOGE("ota_user_validation() error, %u", res);
-            *octx->joctx->ota_return_status = res;
+            *joctx->ota_return_status = res;
             return res;
         }
-        *octx->prevalidated = true;
+        *joctx->validated_confirmed = true;
     }
 
-    const esp_err_t res = esp_ota_write(*octx->joctx->ota_handle, (const void*)uncompressed, length);
+    const esp_err_t res = esp_ota_write(*joctx->ota_handle, (const void*)uncompressed, length);
     if (res != ESP_OK) {
         JADE_LOGE("ota_write() error: %u", res);
-        *octx->joctx->ota_return_status = ERROR_WRITE;
+        *joctx->ota_return_status = ERROR_WRITE;
         return DEFLATE_ERROR;
     }
 
-    if (octx->joctx->hash_type == HASHTYPE_FULLFWDATA) {
+    if (joctx->hash_type == HASHTYPE_FULLFWDATA) {
         // Add written to hash calculation
-        JADE_ZERO_VERIFY(mbedtls_sha256_update(octx->joctx->sha_ctx, uncompressed, length));
+        JADE_ZERO_VERIFY(mbedtls_sha256_update(joctx->sha_ctx, uncompressed, length));
     }
 
-    *octx->joctx->remaining_uncompressed -= length;
-    const size_t written = octx->joctx->uncompressedsize - *octx->joctx->remaining_uncompressed;
+    *joctx->remaining_uncompressed -= length;
+    joctx->fwwritten += length;
+
+    // For a full ota, the amount of fw data uncompressed should always be equal to the
+    // amount of new firmware we have written, as it should be the same thing.
+    JADE_ASSERT(joctx->uncompressedsize - *joctx->remaining_uncompressed == joctx->fwwritten);
 
-    if (written > CUSTOM_HEADER_MIN_WRITE && !*octx->prevalidated) {
+    if (joctx->fwwritten > CUSTOM_HEADER_MIN_WRITE && !*joctx->validated_confirmed) {
         return DEFLATE_ERROR;
     }
 
     /* Update the progress bar once the user has confirmed and upload is in progress */
-    if (*octx->prevalidated) {
-        JADE_ASSERT(octx->joctx->progress_bar.progress_bar);
-        update_progress_bar(&octx->joctx->progress_bar, octx->joctx->uncompressedsize, written);
+    if (*joctx->validated_confirmed) {
+        JADE_ASSERT(joctx->progress_bar.progress_bar);
+        update_progress_bar(&joctx->progress_bar, joctx->uncompressedsize, joctx->fwwritten);
     }
 
     return DEFLATE_OK;
@@ -81,7 +79,7 @@ void ota_process(void* process_ptr)
     jade_process_t* process = process_ptr;
     bool uploading = false;
     enum ota_status ota_return_status = ERROR_OTA_SETUP;
-    bool prevalidated = false;
+    bool validated_confirmed = false;
     bool ota_end_called = false;
 
     char id[MAXLEN_ID + 1];
@@ -134,10 +132,6 @@ void ota_process(void* process_ptr)
     struct deflate_ctx* dctx = JADE_MALLOC_PREFER_SPIRAM(sizeof(struct deflate_ctx));
     jade_process_free_on_exit(process, dctx);
 
-    ota_deflate_ctx_t octx = {
-        .prevalidated = &prevalidated,
-    };
-
     size_t remaining_uncompressed = firmwaresize;
 
     jade_ota_ctx_t joctx = {
@@ -146,6 +140,7 @@ void ota_process(void* process_ptr)
         .hash_type = hash_type,
         .dctx = dctx,
         .id = id,
+        .validated_confirmed = &validated_confirmed,
         .uncompressedsize = firmwaresize,
         .remaining_uncompressed = &remaining_uncompressed,
         .ota_return_status = &ota_return_status,
@@ -158,14 +153,12 @@ void ota_process(void* process_ptr)
         .expected_hash = expected_hash,
     };
 
-    octx.joctx = &joctx;
-
     if (!ota_init(&joctx)) {
         jade_process_reject_message(process, CBOR_RPC_INTERNAL_ERROR, "Failed to initialize OTA", NULL);
         goto cleanup;
     }
 
-    const int dret = deflate_init_write_compressed(dctx, compressedsize, uncompressed_stream_writer, &octx);
+    const int dret = deflate_init_write_compressed(dctx, compressedsize, uncompressed_stream_writer, &joctx);
     JADE_ASSERT(!dret);
 
     // Send the ok response, which implies now we will get ota_data messages
@@ -182,7 +175,7 @@ void ota_process(void* process_ptr)
             goto cleanup;
         }
     }
-    JADE_ASSERT(prevalidated);
+    JADE_ASSERT(validated_confirmed);
 
     // Uploading complete
     uploading = false;
@@ -192,6 +185,10 @@ void ota_process(void* process_ptr)
         JADE_LOGE("Expected uncompressed size: %u, got %u", firmwaresize, firmwaresize - remaining_uncompressed);
         ota_return_status = ERROR_DECOMPRESS;
     }
+    if (joctx.fwwritten != firmwaresize) {
+        JADE_LOGE("Expected amountof firmware written: %u, expected %u", joctx.fwwritten, firmwaresize);
+        ota_return_status = ERROR_DECOMPRESS;
+    }
 
     // Expect a complete/request for status
     jade_process_load_in_message(process, true);
@@ -234,7 +231,7 @@ void ota_process(void* process_ptr)
         esp_restart();
     } else {
         JADE_LOGE("OTA error %u: %s", ota_return_status, MESSAGES[ota_return_status]);
-        if (prevalidated && !ota_end_called) {
+        if (validated_confirmed && !ota_end_called) {
             // ota_begin has been called, cleanup
             const esp_err_t err = esp_ota_abort(ota_handle);
             JADE_ASSERT(err == ESP_OK);

main/process/ota_delta.c

@@ -22,14 +22,6 @@
 #include <bspatch.h>
 #include <deflate.h>
 
-typedef struct {
-    jade_ota_ctx_t* joctx;
-    char* id;
-    struct deflate_ctx* const dctx;
-    size_t written;
-    bool header_validated;
-} bsdiff_ctx_t;
-
 // Error reply in ota_delta is complicated by the fact that we reply 'ok' when we push the received patch data
 // into the decompressor, but we carry on copying the base firmware and inflating/applying patch data.
 // If an error occurs at this stage - we have no message id to reply to so must cache the error until we do receive
@@ -59,36 +51,36 @@ typedef struct {
 // NOTE: uses macros above so may return error immediately, or may just cache it for later return
 static int patch_stream_reader(const struct bspatch_stream* stream, void* buffer, int length)
 {
-    bsdiff_ctx_t* bctx = (bsdiff_ctx_t*)stream->opaque;
-    JADE_ASSERT(bctx);
+    jade_ota_ctx_t* joctx = (jade_ota_ctx_t*)stream->opaque;
+    JADE_ASSERT(joctx);
 
     if (length <= 0) {
-        HANDLE_NEW_ERROR(bctx->joctx, ERROR_PATCH);
+        HANDLE_NEW_ERROR(joctx, ERROR_PATCH);
     }
 
     // Return any non-zero error code from the read routine
-    const int ret = read_uncompressed(bctx->dctx, buffer, length);
+    const int ret = read_uncompressed(joctx->dctx, buffer, length);
     if (ret) {
-        HANDLE_NEW_ERROR(bctx->joctx, ret);
+        HANDLE_NEW_ERROR(joctx, ret);
     }
 
-    *bctx->joctx->remaining_uncompressed -= length;
+    *joctx->remaining_uncompressed -= length;
 
     return SUCCESS;
 }
 
 // NOTE: uses macros above so may return error immediately, or may just cache it for later return
 static int base_firmware_stream_reader(const struct bspatch_stream_i* stream, void* buffer, int pos, int length)
 {
-    bsdiff_ctx_t* bctx = (bsdiff_ctx_t*)stream->opaque;
-    JADE_ASSERT(bctx);
+    jade_ota_ctx_t* joctx = (jade_ota_ctx_t*)stream->opaque;
+    JADE_ASSERT(joctx);
 
     // If currently in error, return immediately without reading anything
-    HANDLE_ANY_CACHED_ERROR(bctx->joctx);
+    HANDLE_ANY_CACHED_ERROR(joctx);
 
-    if (length <= 0 || pos + length >= bctx->joctx->running_partition->size
-        || esp_partition_read(bctx->joctx->running_partition, pos, buffer, length) != ESP_OK) {
-        HANDLE_NEW_ERROR(bctx->joctx, ERROR_PATCH);
+    if (length <= 0 || pos + length >= joctx->running_partition->size
+        || esp_partition_read(joctx->running_partition, pos, buffer, length) != ESP_OK) {
+        HANDLE_NEW_ERROR(joctx, ERROR_PATCH);
     }
 
     return SUCCESS;
@@ -97,37 +89,41 @@ static int base_firmware_stream_reader(const struct bspatch_stream_i* stream, vo
 // NOTE: uses macros above so may return error immediately, or may just cache it for later return
 static int ota_stream_writer(const struct bspatch_stream_n* stream, const void* buffer, int length)
 {
-    bsdiff_ctx_t* bctx = (bsdiff_ctx_t*)stream->opaque;
-    JADE_ASSERT(bctx);
+    jade_ota_ctx_t* joctx = (jade_ota_ctx_t*)stream->opaque;
+    JADE_ASSERT(joctx);
 
     // If currently in error, return immediately without writing anything
-    HANDLE_ANY_CACHED_ERROR(bctx->joctx);
+    HANDLE_ANY_CACHED_ERROR(joctx);
 
-    if (length <= 0 || esp_ota_write(*bctx->joctx->ota_handle, buffer, length) != ESP_OK) {
-        HANDLE_NEW_ERROR(bctx->joctx, ERROR_PATCH);
+    if (length <= 0 || esp_ota_write(*joctx->ota_handle, buffer, length) != ESP_OK) {
+        HANDLE_NEW_ERROR(joctx, ERROR_PATCH);
     }
 
-    if (!bctx->header_validated && length >= CUSTOM_HEADER_MIN_WRITE) {
-        const enum ota_status validation = ota_user_validation(bctx->joctx, (uint8_t*)buffer);
+    if (!*joctx->validated_confirmed && length >= CUSTOM_HEADER_MIN_WRITE) {
+        const enum ota_status validation = ota_user_validation(joctx, (uint8_t*)buffer);
         if (validation != SUCCESS) {
-            HANDLE_NEW_ERROR(bctx->joctx, validation);
+            HANDLE_NEW_ERROR(joctx, validation);
         }
-        bctx->header_validated = true;
+        *joctx->validated_confirmed = true;
     }
 
-    if (bctx->joctx->hash_type == HASHTYPE_FULLFWDATA) {
+    if (joctx->hash_type == HASHTYPE_FULLFWDATA) {
         // Add written to hash calculation
-        JADE_ZERO_VERIFY(mbedtls_sha256_update(bctx->joctx->sha_ctx, buffer, length));
+        JADE_ZERO_VERIFY(mbedtls_sha256_update(joctx->sha_ctx, buffer, length));
     }
 
-    bctx->written += length;
+    joctx->fwwritten += length;
+
+    // For a patch, the amount of patch data uncompressed should always be more than the
+    // amount of new firmware we have written, because of additional patch meta-data.
+    JADE_ASSERT(joctx->uncompressedsize - *joctx->remaining_uncompressed > joctx->fwwritten);
 
-    if (bctx->written > CUSTOM_HEADER_MIN_WRITE && !bctx->header_validated) {
-        HANDLE_NEW_ERROR(bctx->joctx, ERROR_PATCH);
+    if (joctx->fwwritten > CUSTOM_HEADER_MIN_WRITE && !*joctx->validated_confirmed) {
+        HANDLE_NEW_ERROR(joctx, ERROR_PATCH);
     }
 
-    if (bctx->header_validated) {
-        update_progress_bar(&bctx->joctx->progress_bar, bctx->joctx->firmwaresize, bctx->written);
+    if (*joctx->validated_confirmed) {
+        update_progress_bar(&joctx->progress_bar, joctx->firmwaresize, joctx->fwwritten);
     }
 
     return SUCCESS;
@@ -137,13 +133,12 @@ static int compressed_stream_reader(void* ctx)
 {
     JADE_ASSERT(ctx);
 
-    bsdiff_ctx_t* bctx = (bsdiff_ctx_t*)ctx;
-    JADE_ASSERT(bctx->joctx);
+    jade_ota_ctx_t* joctx = (jade_ota_ctx_t*)ctx;
 
     // NOTE: the ota_return_status can be set via ptr in joctx
     // Return any error here as it can be returned to the caller in the message reply
-    jade_process_get_in_message(bctx->joctx, &handle_in_bin_data, true);
-    return *bctx->joctx->ota_return_status;
+    jade_process_get_in_message(joctx, &handle_in_bin_data, true);
+    return *joctx->ota_return_status;
 }
 
 void ota_delta_process(void* process_ptr)
@@ -208,10 +203,7 @@ void ota_delta_process(void* process_ptr)
     struct deflate_ctx* dctx = JADE_MALLOC_PREFER_SPIRAM(sizeof(struct deflate_ctx));
     jade_process_free_on_exit(process, dctx);
 
-    bsdiff_ctx_t bctx = {
-        .dctx = dctx,
-    };
-
+    bool validated_confirmed = false;
     size_t remaining_uncompressed = uncompressedpatchsize;
 
     jade_ota_ctx_t joctx = {
@@ -221,6 +213,7 @@ void ota_delta_process(void* process_ptr)
         .ota_handle = &ota_handle,
         .dctx = dctx,
         .id = id,
+        .validated_confirmed = &validated_confirmed,
         .uncompressedsize = uncompressedpatchsize,
         .remaining_uncompressed = &remaining_uncompressed,
         .ota_return_status = &ota_return_status,
@@ -232,9 +225,7 @@ void ota_delta_process(void* process_ptr)
         .expected_hash = expected_hash,
     };
 
-    bctx.joctx = &joctx;
-
-    int ret = deflate_init_read_uncompressed(dctx, compressedsize, compressed_stream_reader, &bctx);
+    int ret = deflate_init_read_uncompressed(dctx, compressedsize, compressed_stream_reader, &joctx);
     JADE_ASSERT(!ret);
 
     if (!ota_init(&joctx)) {
@@ -251,15 +242,15 @@ void ota_delta_process(void* process_ptr)
     struct bspatch_stream_n destination_firmware_stream_writer;
     // new partition
     destination_firmware_stream_writer.write = &ota_stream_writer;
-    destination_firmware_stream_writer.opaque = &bctx;
+    destination_firmware_stream_writer.opaque = &joctx;
     // patch
     struct bspatch_stream stream;
     stream.read = &patch_stream_reader;
-    stream.opaque = &bctx;
+    stream.opaque = &joctx;
     // old partition / base
     struct bspatch_stream_i basestream;
     basestream.read = &base_firmware_stream_reader;
-    basestream.opaque = &bctx;
+    basestream.opaque = &joctx;
 
     ota_return_status = SUCCESS;
     ret = bspatch(
@@ -275,7 +266,7 @@ void ota_delta_process(void* process_ptr)
     // Uploading complete
     uploading = false;
 
-    if (bctx.written != firmwaresize) {
+    if (joctx.fwwritten != firmwaresize) {
         ota_return_status = ERROR_PATCH;
     }
 

main/process/ota_util.h

@@ -47,11 +47,13 @@ typedef struct {
     enum ota_status* ota_return_status;
     struct deflate_ctx* dctx;
     const jade_msg_source_t* expected_source;
+    bool* validated_confirmed;
     size_t* const remaining_uncompressed;
     size_t remaining_compressed;
     size_t uncompressedsize;
     size_t compressedsize;
     size_t firmwaresize;
+    size_t fwwritten;
 } jade_ota_ctx_t;
 
 enum ota_status {