ota: share error handling code

diff --git a/main/process/ota.c b/main/process/ota.c
index cb46132..85e1f6f 100644
--- a/main/process/ota.c
+++ b/main/process/ota.c
@@ -86,81 +86,16 @@ void ota_process(void* process_ptr)
     int ret = deflate_init_write_compressed(&joctx->dctx, joctx->compressedsize, uncompressed_stream_writer, joctx);
     JADE_ASSERT(!ret);
 
-    // Send the ok response, which implies now we will get ota_data messages
+    // Send the ok response. The caller should then send ota_data messages
     jade_process_reply_to_message_ok(process);
-    bool uploading = true;
 
+    // Uncompress, verify and write data from incoming ota_data messages
     joctx->ota_return_status = OTA_SUCCESS;
-    while (joctx->remaining_compressed) {
+    while (joctx->remaining_compressed && joctx->ota_return_status == OTA_SUCCESS) {
         jade_process_get_in_message(joctx, &handle_in_bin_data, true);
-
-        // NOTE: the ota_return_status can be set via ptr in joctx
-        if (joctx->ota_return_status != OTA_SUCCESS) {
-            JADE_LOGE("Error on ota_data message: %d", joctx->ota_return_status);
-            goto cleanup;
-        }
-    }
-    JADE_ASSERT(joctx->validated_confirmed);
-
-    // Bail-out if the fw uncompressed to an unexpected size
-    if (joctx->remaining_uncompressed != 0) {
-        JADE_LOGE("Expected uncompressed size: %u, got %u", joctx->firmwaresize,
-            joctx->firmwaresize - joctx->remaining_uncompressed);
-        joctx->ota_return_status = OTA_ERR_DECOMPRESS;
-        goto cleanup;
-    }
-    if (joctx->fwwritten != joctx->firmwaresize) {
-        JADE_LOGE("Expected amountof firmware written: %u, expected %u", joctx->fwwritten, joctx->firmwaresize);
-        joctx->ota_return_status = OTA_ERR_DECOMPRESS;
-        goto cleanup;
-    }
-
-    // Uploading complete
-    uploading = false;
-
-    // Expect a complete/request for status
-    jade_process_load_in_message(process, true);
-    if (!IS_CURRENT_MESSAGE(process, "ota_complete")) {
-        // Protocol error
-        jade_process_reject_message(process, CBOR_RPC_PROTOCOL_ERROR, "Unexpected message, expecting 'ota_complete'");
-        return;
-    }
-
-    // If all good with the upload do all final checks and then finalise the ota
-    // and set the new boot partition, etc.
-    ota_finalize(process, joctx);
-
-    // Send final message reply with final status
-    if (joctx->ota_return_status != OTA_SUCCESS) {
-        uint8_t buf[256];
-        const char* error = ota_get_status_text(joctx->ota_return_status);
-        jade_process_reject_message_ex(process->ctx, CBOR_RPC_INTERNAL_ERROR, "Error completing OTA",
-            (const uint8_t*)error, strlen(error), buf, sizeof(buf));
-        goto cleanup;
     }
 
-cleanup:
-
-    // Show error-message and await user acknowledgement.
-    if (joctx->ota_return_status != OTA_SUCCESS) {
-        JADE_LOGE("OTA error %u: %s", joctx->ota_return_status, ota_get_status_text(joctx->ota_return_status));
-
-        // If we get here and we have not finished loading the data, send an error message
-        const char* status_text = ota_get_status_text(joctx->ota_return_status);
-        if (uploading) {
-            JADE_ASSERT(joctx->id[0] != '\0');
-            const int error_code
-                = joctx->ota_return_status == OTA_ERR_USERDECLINED ? CBOR_RPC_USER_CANCELLED : CBOR_RPC_INTERNAL_ERROR;
-
-            uint8_t buf[256];
-            jade_process_reject_message_with_id(joctx->id, error_code, "Error uploading OTA data",
-                (const uint8_t*)status_text, strlen(status_text), buf, sizeof(buf), joctx->expected_source);
-        }
-
-        // If the error is not 'did not start' or 'user declined', show an error screen
-        if (joctx->ota_return_status != OTA_ERR_SETUP && joctx->ota_return_status != OTA_ERR_USERDECLINED) {
-            await_error_activity(&status_text, 1);
-        }
-    }
+    // Finalise the ota and reboot, or send an error and return
+    ota_finalize(process, joctx, is_delta);
 }
 #endif // AMALGAMATED_BUILD
diff --git a/main/process/ota_delta.c b/main/process/ota_delta.c
index 83c961e..392f60e 100644
--- a/main/process/ota_delta.c
+++ b/main/process/ota_delta.c
@@ -157,9 +157,8 @@ void ota_delta_process(void* process_ptr)
     int ret = deflate_init_read_uncompressed(&joctx->dctx, joctx->compressedsize, compressed_stream_reader, joctx);
     JADE_ASSERT(!ret);
 
-    // Send the ok response, which implies now we will get ota_data messages
+    // Send the ok response. The caller should then send ota_data messages
     jade_process_reply_to_message_ok(process);
-    bool uploading = true;
 
     struct bspatch_stream_n destination_firmware_stream_writer;
     // new partition
@@ -174,6 +173,7 @@ void ota_delta_process(void* process_ptr)
     basestream.read = &base_firmware_stream_reader;
     basestream.opaque = joctx;
 
+    // Uncompress, verify and apply the patch from incoming ota_data messages
     joctx->ota_return_status = OTA_SUCCESS;
     ret = bspatch(
         &basestream, joctx->running_partition->size, &destination_firmware_stream_writer, joctx->firmwaresize, &stream);
@@ -181,65 +181,9 @@ void ota_delta_process(void* process_ptr)
     if (ret != OTA_SUCCESS) {
         JADE_LOGE("Error applying patch: %d", ret);
         joctx->ota_return_status = ret < 0 ? OTA_ERR_PATCH : ret;
-        goto cleanup;
     }
 
-    if (joctx->fwwritten != joctx->firmwaresize) {
-        joctx->ota_return_status = OTA_ERR_PATCH;
-        goto cleanup;
-    }
-
-    // Uploading complete
-    uploading = false;
-
-    // Expect a complete/request for status
-    jade_process_load_in_message(process, true);
-    if (!IS_CURRENT_MESSAGE(process, "ota_complete")) {
-        // Protocol error
-        jade_process_reject_message(process, CBOR_RPC_PROTOCOL_ERROR, "Unexpected message, expecting 'ota_complete'");
-        return;
-    }
-
-    // If all good with the upload do all final checks and then finalise the ota
-    // and set the new boot partition, etc.
-    ota_finalize(process, joctx);
-
-    // Send final message reply with final status
-    if (joctx->ota_return_status != OTA_SUCCESS) {
-        uint8_t buf[256];
-        const char* error = ota_get_status_text(joctx->ota_return_status);
-        jade_process_reject_message_ex(process->ctx, CBOR_RPC_INTERNAL_ERROR, "Error completing OTA delta",
-            (const uint8_t*)error, strlen(error), buf, sizeof(buf));
-        goto cleanup;
-    }
-
-cleanup:
-
-    // Show error-message and await user acknowledgement.
-    if (joctx->ota_return_status != OTA_SUCCESS) {
-        JADE_LOGE("OTA error %u: %s", joctx->ota_return_status, ota_get_status_text(joctx->ota_return_status));
-
-        // If we get here and we have not finished loading the data, send an error message
-        const char* status_text = ota_get_status_text(joctx->ota_return_status);
-        if (uploading) {
-            if (joctx->id[0] == '\0') {
-                // This should not happen under normal circumstances, but it could occur if the delta
-                // uploaded is not appropriate for the base/running firmware (or perhaps is corrupted).
-                // In that case bspatch() can fail unexpectedly - default the id.
-                strcpy(joctx->id, "00");
-            }
-            const int error_code
-                = joctx->ota_return_status == OTA_ERR_USERDECLINED ? CBOR_RPC_USER_CANCELLED : CBOR_RPC_INTERNAL_ERROR;
-
-            uint8_t buf[256];
-            jade_process_reject_message_with_id(joctx->id, error_code, "Error uploading OTA delta data",
-                (const uint8_t*)status_text, strlen(status_text), buf, sizeof(buf), joctx->expected_source);
-        }
-
-        // If the error is not 'did not start' or 'user declined', show an error screen
-        if (joctx->ota_return_status != OTA_ERR_SETUP && joctx->ota_return_status != OTA_ERR_USERDECLINED) {
-            await_error_activity(&status_text, 1);
-        }
-    }
+    // Finalise the ota and reboot, or send an error and return
+    ota_finalize(process, joctx, is_delta);
 }
 #endif // AMALGAMATED_BUILD
diff --git a/main/process/ota_util.c b/main/process/ota_util.c
index 9701030..ac56cbd 100644
--- a/main/process/ota_util.c
+++ b/main/process/ota_util.c
@@ -20,6 +20,44 @@ extern esp_app_desc_t running_app_info;
 const __attribute__((section(".rodata_custom_desc"))) esp_custom_app_desc_t custom_app_desc
     = { .version = 1, .board_type = JADE_OTA_BOARD_TYPE, .features = JADE_OTA_FEATURES, .config = JADE_OTA_CONFIG };
 
+static const char* ota_get_status_text(const ota_status_t status)
+{
+    switch (status) {
+    case OTA_SUCCESS:
+        return "OK";
+    case OTA_ERR_SETUP:
+        return "OTA_ERR_SETUP";
+    case OTA_ERR_INIT:
+        return "OTA_ERR_INIT";
+    case OTA_ERR_BADPARTITION:
+        return "OTA_ERR_BADPARTITION";
+    case OTA_ERR_DECOMPRESS:
+        return "OTA_ERR_DECOMPRESS";
+    case OTA_ERR_WRITE:
+        return "OTA_ERR_WRITE";
+    case OTA_ERR_FINISH:
+        return "OTA_ERR_FINISH";
+    case OTA_ERR_SETPARTITION:
+        return "OTA_ERR_SETPARTITION";
+    case OTA_ERR_BADDATA:
+        return "OTA_ERR_BADDATA";
+    case OTA_ERR_NODOWNGRADE:
+        return "OTA_ERR_NODOWNGRADE";
+    case OTA_ERR_INVALIDFW:
+        return "OTA_ERR_INVALIDFW";
+    case OTA_ERR_USERDECLINED:
+        return "OTA_ERR_USERDECLINED";
+    case OTA_ERR_BADHASH:
+        return "OTA_ERR_BADHASH";
+    case OTA_ERR_PATCH:
+        return "OTA_ERR_PATCH";
+    case OTA_ERR_PROTOCOL:
+        return "OTA_ERR_PROTOCOL";
+    default:
+        return "OTA_ERR_UNKNOWN";
+    }
+}
+
 static void reply_ok(const void* ctx, CborEncoder* container)
 {
     JADE_ASSERT(ctx);
@@ -246,22 +284,35 @@ cleanup:
     return joctx;
 }
 
-void ota_finalize(jade_process_t* process, jade_ota_ctx_t* joctx)
+void ota_finalize(jade_process_t* process, jade_ota_ctx_t* joctx, const bool is_delta)
 {
     JADE_ASSERT(joctx);
 
-    // Ensure no cached error - if so return it now
     if (joctx->ota_return_status != OTA_SUCCESS) {
-        return;
+        goto error; // An error has already occured, return it
+    }
+
+    // To reach this far without error, the user must have confirmed
+    JADE_ASSERT(joctx->validated_confirmed);
+
+    // Expect an ota_complete message
+    jade_process_load_in_message(process, true);
+    if (!IS_CURRENT_MESSAGE(process, "ota_complete")) {
+        joctx->ota_return_status = OTA_ERR_PROTOCOL; // Protocol error
+        goto error;
     }
 
+    if (joctx->fwwritten != joctx->firmwaresize) {
+        JADE_LOGE("OTA checks failed: written: %u/%u", joctx->fwwritten, joctx->firmwaresize);
+        joctx->ota_return_status = is_delta ? OTA_ERR_PATCH : OTA_ERR_DECOMPRESS;
+        goto error;
+    }
     if (joctx->remaining_compressed || joctx->remaining_uncompressed || !joctx->compressedsize
         || !joctx->uncompressedsize) {
-        JADE_LOGE("OTA checks failed: uncompressed size: %u, compressed size: %u, remaining compressed %u, remaining "
-                  "uncompressed %u",
-            joctx->uncompressedsize, joctx->compressedsize, joctx->remaining_compressed, joctx->remaining_uncompressed);
+        JADE_LOGE("OTA checks failed: uncompressed: %u/%u, compressed: %u/%u", joctx->uncompressedsize,
+            joctx->remaining_uncompressed, joctx->compressedsize, joctx->remaining_compressed);
         joctx->ota_return_status = OTA_ERR_INIT;
-        return;
+        goto error;
     }
 
     // Verify calculated compressed file hash matches expected
@@ -279,7 +330,7 @@ void ota_finalize(jade_process_t* process, jade_ota_ctx_t* joctx)
         JADE_WALLY_VERIFY(wally_free_string(calc_hash_hexstr));
 
         joctx->ota_return_status = OTA_ERR_BADHASH;
-        return;
+        goto error;
     }
 
     // All good, finalise the ota and set the partition to boot
@@ -289,14 +340,14 @@ void ota_finalize(jade_process_t* process, jade_ota_ctx_t* joctx)
     if (err != ESP_OK) {
         JADE_LOGE("esp_ota_end() returned %d", err);
         joctx->ota_return_status = OTA_ERR_FINISH;
-        return;
+        goto error;
     }
 
     err = esp_ota_set_boot_partition(joctx->update_partition);
     if (err != ESP_OK) {
         JADE_LOGE("esp_ota_set_boot_partition() returned %d", err);
         joctx->ota_return_status = OTA_ERR_SETPARTITION;
-        return;
+        goto error;
     }
 
     // OTA completed without errors. send an ok and reboot
@@ -308,7 +359,37 @@ void ota_finalize(jade_process_t* process, jade_ota_ctx_t* joctx)
     display_message_activity(message, 1);
 
     vTaskDelay(2500 / portTICK_PERIOD_MS);
-    esp_restart();
+    esp_restart(); // Does not return
+    return; // Unreachable
+
+error:
+    // We have an error, send an error response.
+    const char* status_text = ota_get_status_text(joctx->ota_return_status);
+    JADE_LOGE("OTA error: %s", status_text);
+
+    int errcode = CBOR_RPC_INTERNAL_ERROR;
+    if (joctx->ota_return_status == OTA_ERR_USERDECLINED) {
+        errcode = CBOR_RPC_USER_CANCELLED;
+    } else if (joctx->ota_return_status == OTA_ERR_PROTOCOL) {
+        errcode = CBOR_RPC_PROTOCOL_ERROR;
+    }
+
+    uint8_t buf[256];
+    if (joctx->id[0] != '\0') {
+        // Send error response to the ota_data message we were processing.
+        jade_process_reject_message_with_id(joctx->id, errcode, "Error uploading OTA data", (const uint8_t*)status_text,
+            strlen(status_text), buf, sizeof(buf), joctx->expected_source);
+    } else {
+        // Send error response to the ota_complete message.
+        // If we didn't get an ota_complete, sets the reply id as "00".
+        jade_process_reject_message_ex(process->ctx, errcode, "Error completing OTA", (const uint8_t*)status_text,
+            strlen(status_text), buf, sizeof(buf));
+    }
+
+    // If the error is not 'did not start' or 'user declined', show an error screen
+    if (joctx->ota_return_status != OTA_ERR_SETUP && joctx->ota_return_status != OTA_ERR_USERDECLINED) {
+        await_error_activity(&status_text, 1);
+    }
 }
 
 // NOTE: 'dest' is assumed to be at least as long as 'strlen(src)'
@@ -402,39 +483,4 @@ void ota_user_validate(jade_ota_ctx_t* joctx, const uint8_t* uncompressed)
     joctx->validated_confirmed = true;
 }
 
-const char* ota_get_status_text(const ota_status_t status)
-{
-    switch (status) {
-    case OTA_SUCCESS:
-        return "OK";
-    case OTA_ERR_SETUP:
-        return "OTA_ERR_SETUP";
-    case OTA_ERR_INIT:
-        return "OTA_ERR_INIT";
-    case OTA_ERR_BADPARTITION:
-        return "OTA_ERR_BADPARTITION";
-    case OTA_ERR_DECOMPRESS:
-        return "OTA_ERR_DECOMPRESS";
-    case OTA_ERR_WRITE:
-        return "OTA_ERR_WRITE";
-    case OTA_ERR_FINISH:
-        return "OTA_ERR_FINISH";
-    case OTA_ERR_SETPARTITION:
-        return "OTA_ERR_SETPARTITION";
-    case OTA_ERR_BADDATA:
-        return "OTA_ERR_BADDATA";
-    case OTA_ERR_NODOWNGRADE:
-        return "OTA_ERR_NODOWNGRADE";
-    case OTA_ERR_INVALIDFW:
-        return "OTA_ERR_INVALIDFW";
-    case OTA_ERR_USERDECLINED:
-        return "OTA_ERR_USERDECLINED";
-    case OTA_ERR_BADHASH:
-        return "OTA_ERR_BADHASH";
-    case OTA_ERR_PATCH:
-        return "OTA_ERR_PATCH";
-    default:
-        return "OTA_ERR_UNKNOWN";
-    }
-}
 #endif // AMALGAMATED_BUILD
diff --git a/main/process/ota_util.h b/main/process/ota_util.h
index 0c1fe99..ff34b87 100644
--- a/main/process/ota_util.h
+++ b/main/process/ota_util.h
@@ -48,6 +48,7 @@ typedef enum {
     OTA_ERR_USERDECLINED,
     OTA_ERR_BADHASH,
     OTA_ERR_PATCH,
+    OTA_ERR_PROTOCOL,
 } ota_status_t;
 
 typedef struct {
@@ -78,7 +79,6 @@ void handle_in_bin_data(void* ctx, uint8_t* data, size_t rawsize);
 
 jade_ota_ctx_t* ota_init(jade_process_t* process, bool is_delta);
 void ota_user_validate(jade_ota_ctx_t* joctx, const uint8_t* uncompressed);
-void ota_finalize(jade_process_t* process, jade_ota_ctx_t* joctx);
-const char* ota_get_status_text(ota_status_t status);
+void ota_finalize(jade_process_t* process, jade_ota_ctx_t* joctx, bool is_delta);
 
 #endif /* JADE_OTA_UTIL_H_ */