ota: initial refactor of ota processing

diff --git a/main/process/ota.c b/main/process/ota.c
index eb166d7..0be6b3c 100644
--- a/main/process/ota.c
+++ b/main/process/ota.c
@@ -16,7 +16,6 @@
 #include <esp_efuse.h>
 #include <esp_ota_ops.h>
 
-#include <deflate.h>
 #include <mbedtls/sha256.h>
 #include <wally_core.h>
 
@@ -76,113 +75,45 @@ static int uncompressed_stream_writer(void* ctx, uint8_t* uncompressed, size_t l
 void ota_process(void* process_ptr)
 {
     JADE_LOGI("Starting: %d", xPortGetFreeHeapSize());
-
     jade_process_t* process = process_ptr;
-    bool uploading = false;
-    bool ota_end_called = false;
-
-    // We expect a current message to be present
-    ASSERT_CURRENT_MESSAGE(process, "ota");
-    GET_MSG_PARAMS(process);
-
-    const jade_msg_source_t ota_source = process->ctx.source;
-    if (keychain_has_pin()) {
-        // NOTE: ota from internal source is allowed (eg. QR codes or USB storage)
-        JADE_ASSERT(ota_source == (jade_msg_source_t)keychain_get_userdata() || ota_source == SOURCE_INTERNAL);
-        JADE_ASSERT(!keychain_has_temporary());
-    }
-
-    size_t firmwaresize = 0;
-    size_t compressedsize = 0;
-    if (!rpc_get_sizet("fwsize", &params, &firmwaresize) || !rpc_get_sizet("cmpsize", &params, &compressedsize)
-        || firmwaresize <= compressedsize) {
-        jade_process_reject_message(process, CBOR_RPC_BAD_PARAMETERS, "Bad filesize parameters");
-        goto cleanup;
-    }
 
-    // Optional field indicating preference for rich reply data
-    bool extended_replies = false;
-    rpc_get_boolean("extended_replies", &params, &extended_replies);
-
-    // Can accept either uploaded file data hash (legacy) or hash of the full/final firmware image (preferred)
-    uint8_t expected_hash[SHA256_LEN];
-    char* expected_hash_hexstr = NULL;
-    hash_type_t hash_type;
-    if (rpc_get_n_bytes("fwhash", &params, sizeof(expected_hash), expected_hash)) {
-        hash_type = HASHTYPE_FULLFWDATA;
-    } else if (rpc_get_n_bytes("cmphash", &params, sizeof(expected_hash), expected_hash)) {
-        hash_type = HASHTYPE_FILEDATA;
-    } else {
-        jade_process_reject_message(process, CBOR_RPC_BAD_PARAMETERS, "Cannot extract valid fw hash value");
-        goto cleanup;
-    }
-    JADE_WALLY_VERIFY(wally_hex_from_bytes(expected_hash, sizeof(expected_hash), &expected_hash_hexstr));
-    jade_process_wally_free_string_on_exit(process, expected_hash_hexstr);
-
-    // We will show a progress bar once the user has confirmed and the upload in progress
-    // Initially just show a message screen.
-    const char* message[] = { "Preparing for firmware", "", "update" };
-    display_message_activity(message, 3);
-    vTaskDelay(100 / portTICK_PERIOD_MS); // sleep a little bit to redraw screen
-
-    struct deflate_ctx* dctx = JADE_MALLOC_PREFER_SPIRAM(sizeof(struct deflate_ctx));
-    jade_process_free_on_exit(process, dctx);
-
-    jade_ota_ctx_t joctx = {
-        .sha_ctx = {},
-        .progress_bar = {0},
-        .hash_type = hash_type,
-        .dctx = dctx,
-        .id = { 0 },
-        .uncompressedsize = firmwaresize,
-        .remaining_uncompressed = firmwaresize,
-        .ota_return_status = OTA_ERR_SETUP,
-        .expected_source = ota_source,
-        .remaining_compressed = compressedsize,
-        .compressedsize = compressedsize,
-        .ota_handle = 0,
-        .firmwaresize = firmwaresize,
-        .expected_hash_hexstr = expected_hash_hexstr,
-        .expected_hash = expected_hash,
-        .extended_replies = extended_replies,
-        .validated_confirmed = false,
-    };
-
-    if (!ota_init(&joctx)) {
-        jade_process_reject_message(process, CBOR_RPC_INTERNAL_ERROR, "Failed to initialize OTA");
-        goto cleanup;
+    const bool is_delta = false;
+    jade_ota_ctx_t* joctx = ota_init(process, is_delta);
+    if (!joctx) {
+        return; // The message has already been rejected
     }
 
-    const int dret = deflate_init_write_compressed(dctx, compressedsize, uncompressed_stream_writer, &joctx);
-    JADE_ASSERT(!dret);
+    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
     jade_process_reply_to_message_ok(process);
-    uploading = true;
+    bool uploading = true;
 
-    joctx.ota_return_status = OTA_SUCCESS;
-    while (joctx.remaining_compressed) {
-        jade_process_get_in_message(&joctx, &handle_in_bin_data, true);
+    joctx->ota_return_status = OTA_SUCCESS;
+    while (joctx->remaining_compressed) {
+        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);
+        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);
+    JADE_ASSERT(joctx->validated_confirmed);
 
     // Uploading complete
     uploading = false;
 
     // Bail-out if the fw uncompressed to an unexpected size
-    if (joctx.remaining_uncompressed != 0) {
-        JADE_LOGE("Expected uncompressed size: %u, got %u", firmwaresize, firmwaresize - joctx.remaining_uncompressed);
-        joctx.ota_return_status = OTA_ERR_DECOMPRESS;
+    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;
     }
-    if (joctx.fwwritten != firmwaresize) {
-        JADE_LOGE("Expected amountof firmware written: %u, expected %u", joctx.fwwritten, firmwaresize);
-        joctx.ota_return_status = OTA_ERR_DECOMPRESS;
+    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;
     }
 
     // Expect a complete/request for status
@@ -195,29 +126,27 @@ void ota_process(void* process_ptr)
 
     // If all good with the upload do all final checks and then finalise the ota
     // and set the new boot partition, etc.
-    if (joctx.ota_return_status == OTA_SUCCESS) {
-        joctx.ota_return_status = post_ota_check(&joctx, &ota_end_called);
+    if (joctx->ota_return_status == OTA_SUCCESS) {
+        joctx->ota_return_status = post_ota_check(joctx);
     }
 
     // Send final message reply with final status
-    if (joctx.ota_return_status != OTA_SUCCESS) {
+    if (joctx->ota_return_status != OTA_SUCCESS) {
         uint8_t buf[256];
-        const char* error = ota_get_status_text(joctx.ota_return_status);
+        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;
     }
 
-    JADE_ASSERT(ota_end_called);
     jade_process_reply_to_message_ok(process);
     JADE_LOGI("Success");
 
 cleanup:
-    ota_free(&joctx);
 
     // If ota has been successful show message and reboot.
     // If error, show error-message and await user acknowledgement.
-    if (joctx.ota_return_status == OTA_SUCCESS) {
+    if (joctx->ota_return_status == OTA_SUCCESS) {
         JADE_LOGW("OTA successful - rebooting");
 
         const char* message[] = { "Upgrade successful!" };
@@ -226,27 +155,22 @@ cleanup:
         vTaskDelay(2500 / portTICK_PERIOD_MS);
         esp_restart();
     } else {
-        JADE_LOGE("OTA error %u: %s", joctx.ota_return_status, ota_get_status_text(joctx.ota_return_status));
-        if (joctx.validated_confirmed && !ota_end_called) {
-            // ota_begin has been called, cleanup
-            const esp_err_t err = esp_ota_abort(joctx.ota_handle);
-            JADE_ASSERT(err == ESP_OK);
-        }
+        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);
+        const char* status_text = ota_get_status_text(joctx->ota_return_status);
         if (uploading) {
-            JADE_ASSERT(joctx.id[0] != '\0');
+            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;
+                = 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), ota_source);
+            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) {
+        if (joctx->ota_return_status != OTA_ERR_SETUP && joctx->ota_return_status != OTA_ERR_USERDECLINED) {
             await_error_activity(&status_text, 1);
         }
     }
diff --git a/main/process/ota_delta.c b/main/process/ota_delta.c
index b10a510..a56ff08 100644
--- a/main/process/ota_delta.c
+++ b/main/process/ota_delta.c
@@ -21,7 +21,6 @@
 
 #include "process_utils.h"
 #include <bspatch.h>
-#include <deflate.h>
 
 // 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.
@@ -60,7 +59,7 @@ static int patch_stream_reader(const struct bspatch_stream* stream, void* buffer
     }
 
     // Return any non-zero error code from the read routine
-    const int ret = read_uncompressed(joctx->dctx, buffer, length);
+    const int ret = read_uncompressed(&joctx->dctx, buffer, length);
     if (ret) {
         HANDLE_NEW_ERROR(joctx, ret);
     }
@@ -146,123 +145,48 @@ void ota_delta_process(void* process_ptr)
 {
     JADE_LOGI("Starting: %d", xPortGetFreeHeapSize());
     jade_process_t* process = process_ptr;
-    bool uploading = false;
-    bool ota_end_called = false;
-    bool ota_begin_called = false;
-
-    // We expect a current message to be present
-    ASSERT_CURRENT_MESSAGE(process, "ota_delta");
-    GET_MSG_PARAMS(process);
-
-    const jade_msg_source_t ota_source = process->ctx.source;
-    if (keychain_has_pin()) {
-        // NOTE: ota from internal source is allowed (eg. QR codes or USB storage)
-        JADE_ASSERT(ota_source == (jade_msg_source_t)keychain_get_userdata() || ota_source == SOURCE_INTERNAL);
-        JADE_ASSERT(!keychain_has_temporary());
-    }
-
-    size_t firmwaresize = 0;
-    size_t compressedsize = 0;
-    size_t uncompressedpatchsize = 0;
 
-    if (!rpc_get_sizet("fwsize", &params, &firmwaresize) || !rpc_get_sizet("cmpsize", &params, &compressedsize)
-        || !rpc_get_sizet("patchsize", &params, &uncompressedpatchsize) || firmwaresize <= compressedsize
-        || uncompressedpatchsize <= compressedsize) {
-        jade_process_reject_message(process, CBOR_RPC_BAD_PARAMETERS, "Bad filesize parameters");
-        goto cleanup;
+    const bool is_delta = true;
+    jade_ota_ctx_t* joctx = ota_init(process, is_delta);
+    if (!joctx) {
+        return; // The message has already been rejected
     }
 
-    // Optional field indicating preference for rich reply data
-    bool extended_replies = false;
-    rpc_get_boolean("extended_replies", &params, &extended_replies);
-
-    // Can accept either uploaded file data hash (legacy) or hash of the full/final firmware image (preferred)
-    uint8_t expected_hash[SHA256_LEN];
-    char* expected_hash_hexstr = NULL;
-    hash_type_t hash_type;
-    if (rpc_get_n_bytes("fwhash", &params, sizeof(expected_hash), expected_hash)) {
-        hash_type = HASHTYPE_FULLFWDATA;
-    } else if (rpc_get_n_bytes("cmphash", &params, sizeof(expected_hash), expected_hash)) {
-        hash_type = HASHTYPE_FILEDATA;
-    } else {
-        jade_process_reject_message(process, CBOR_RPC_BAD_PARAMETERS, "Cannot extract valid fw hash value");
-        goto cleanup;
-    }
-    JADE_WALLY_VERIFY(wally_hex_from_bytes(expected_hash, sizeof(expected_hash), &expected_hash_hexstr));
-    jade_process_wally_free_string_on_exit(process, expected_hash_hexstr);
-
-    // We will show a progress bar once the user has confirmed and the upload in progress
-    // Initially just show a message screen.
-    const char* message[] = { "Preparing for firmware", "", "update" };
-    display_message_activity(message, 3);
-    vTaskDelay(100 / portTICK_PERIOD_MS); // sleep a little bit to redraw screen
-
-    struct deflate_ctx* dctx = JADE_MALLOC_PREFER_SPIRAM(sizeof(struct deflate_ctx));
-    jade_process_free_on_exit(process, dctx);
-
-    jade_ota_ctx_t joctx = {
-        .sha_ctx = {},
-        .progress_bar = {0},
-        .hash_type = hash_type,
-        .ota_handle = 0,
-        .dctx = dctx,
-        .id = { 0 },
-        .uncompressedsize = uncompressedpatchsize,
-        .remaining_uncompressed = uncompressedpatchsize,
-        .ota_return_status = OTA_ERR_SETUP,
-        .expected_source = ota_source,
-        .remaining_compressed = compressedsize,
-        .firmwaresize = firmwaresize,
-        .compressedsize = compressedsize,
-        .expected_hash_hexstr = expected_hash_hexstr,
-        .expected_hash = expected_hash,
-        .extended_replies = extended_replies,
-        .validated_confirmed = false,
-    };
-
-    int ret = deflate_init_read_uncompressed(dctx, compressedsize, compressed_stream_reader, &joctx);
+    int ret = deflate_init_read_uncompressed(&joctx->dctx, joctx->compressedsize, compressed_stream_reader, joctx);
     JADE_ASSERT(!ret);
 
-    if (!ota_init(&joctx)) {
-        jade_process_reject_message(process, CBOR_RPC_INTERNAL_ERROR, "Failed to initialize OTA");
-        goto cleanup;
-    }
-
-    ota_begin_called = true;
-
     // Send the ok response, which implies now we will get ota_data messages
     jade_process_reply_to_message_ok(process);
-    uploading = true;
+    bool uploading = true;
 
     struct bspatch_stream_n destination_firmware_stream_writer;
     // new partition
     destination_firmware_stream_writer.write = &ota_stream_writer;
-    destination_firmware_stream_writer.opaque = &joctx;
+    destination_firmware_stream_writer.opaque = joctx;
     // patch
     struct bspatch_stream stream;
     stream.read = &patch_stream_reader;
-    stream.opaque = &joctx;
+    stream.opaque = joctx;
     // old partition / base
     struct bspatch_stream_i basestream;
     basestream.read = &base_firmware_stream_reader;
-    basestream.opaque = &joctx;
+    basestream.opaque = joctx;
 
-    joctx.ota_return_status = OTA_SUCCESS;
+    joctx->ota_return_status = OTA_SUCCESS;
     ret = bspatch(
-        &basestream, joctx.running_partition->size, &destination_firmware_stream_writer, firmwaresize, &stream);
+        &basestream, joctx->running_partition->size, &destination_firmware_stream_writer, joctx->firmwaresize, &stream);
 
     if (ret != OTA_SUCCESS) {
         JADE_LOGE("Error applying patch: %d", ret);
-        joctx.ota_return_status = ret < 0 ? OTA_ERR_PATCH : ret;
+        joctx->ota_return_status = ret < 0 ? OTA_ERR_PATCH : ret;
         goto cleanup;
     }
-    JADE_ASSERT(ota_begin_called);
 
     // Uploading complete
     uploading = false;
 
-    if (joctx.fwwritten != firmwaresize) {
-        joctx.ota_return_status = OTA_ERR_PATCH;
+    if (joctx->fwwritten != joctx->firmwaresize) {
+        joctx->ota_return_status = OTA_ERR_PATCH;
     }
 
     // Expect a complete/request for status
@@ -275,29 +199,27 @@ void ota_delta_process(void* process_ptr)
 
     // If all good with the upload do all final checks and then finalise the ota
     // and set the new boot partition, etc.
-    if (joctx.ota_return_status == OTA_SUCCESS) {
-        joctx.ota_return_status = post_ota_check(&joctx, &ota_end_called);
+    if (joctx->ota_return_status == OTA_SUCCESS) {
+        joctx->ota_return_status = post_ota_check(joctx);
     }
 
     // Send final message reply with final status
-    if (joctx.ota_return_status != OTA_SUCCESS) {
+    if (joctx->ota_return_status != OTA_SUCCESS) {
         uint8_t buf[256];
-        const char* error = ota_get_status_text(joctx.ota_return_status);
+        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;
     }
 
-    JADE_ASSERT(ota_end_called);
     jade_process_reply_to_message_ok(process);
     JADE_LOGI("Success");
 
 cleanup:
-    ota_free(&joctx);
 
     // If ota has been successful show message and reboot.
     // If error, show error-message and await user acknowledgement.
-    if (joctx.ota_return_status == OTA_SUCCESS) {
+    if (joctx->ota_return_status == OTA_SUCCESS) {
         JADE_LOGW("OTA successful - rebooting");
 
         const char* message[] = { "Upgrade successful!" };
@@ -306,32 +228,27 @@ cleanup:
         vTaskDelay(2500 / portTICK_PERIOD_MS);
         esp_restart();
     } else {
-        JADE_LOGE("OTA error %u: %s", joctx.ota_return_status, ota_get_status_text(joctx.ota_return_status));
-        if (ota_begin_called && !ota_end_called) {
-            // ota_begin has been called, cleanup
-            const esp_err_t err = esp_ota_abort(joctx.ota_handle);
-            JADE_ASSERT(err == ESP_OK);
-        }
+        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);
+        const char* status_text = ota_get_status_text(joctx->ota_return_status);
         if (uploading) {
-            if (joctx.id[0] == '\0') {
+            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");
+                strcpy(joctx->id, "00");
             }
             const int error_code
-                = joctx.ota_return_status == OTA_ERR_USERDECLINED ? CBOR_RPC_USER_CANCELLED : CBOR_RPC_INTERNAL_ERROR;
+                = 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), ota_source);
+            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) {
+        if (joctx->ota_return_status != OTA_ERR_SETUP && joctx->ota_return_status != OTA_ERR_USERDECLINED) {
             await_error_activity(&status_text, 1);
         }
     }
diff --git a/main/process/ota_util.c b/main/process/ota_util.c
index 777eeb0..f13b57a 100644
--- a/main/process/ota_util.c
+++ b/main/process/ota_util.c
@@ -7,7 +7,6 @@
 #include "ota_defines.h"
 
 #include <ctype.h>
-#include <deflate.h>
 #include <esp_efuse.h>
 #include <sodium/utils.h>
 #include <string.h>
@@ -98,7 +97,7 @@ void handle_in_bin_data(void* ctx, uint8_t* data, const size_t rawsize)
     // For the time being, send ok *after* the processing/decompression steps.
 
     // Return any non-zero error code from the decompress routine
-    const int ret = joctx->dctx->write_compressed(joctx->dctx, (uint8_t* const)inbound_buf, written);
+    const int ret = joctx->dctx.write_compressed(&joctx->dctx, (uint8_t* const)inbound_buf, written);
     if (ret) {
         joctx->ota_return_status = ret < 0 ? OTA_ERR_DECOMPRESS : ret;
         return;
@@ -127,51 +126,129 @@ void handle_in_bin_data(void* ctx, uint8_t* data, const size_t rawsize)
     joctx->id[0] = '\0';
 }
 
-bool ota_init(jade_ota_ctx_t* joctx)
+static void ota_free(void* ctx)
 {
-    JADE_ASSERT(joctx);
+    jade_ota_ctx_t* joctx = (jade_ota_ctx_t*)ctx;
+    if (joctx) {
+        mbedtls_sha256_free(&joctx->sha_ctx);
+        wally_free_string(joctx->expected_hash_hexstr); // Ignore return value
+        if (joctx->ota_handle) {
+            // ota was started but not cleanly shutdown: abort it
+            const esp_err_t err = esp_ota_abort(joctx->ota_handle);
+            JADE_ASSERT(err == ESP_OK);
+            joctx->ota_handle = 0;
+        }
+    }
+}
 
-    mbedtls_sha256_init(&joctx->sha_ctx);
-    joctx->running_partition = esp_ota_get_running_partition();
-    JADE_ASSERT(joctx->running_partition);
-    JADE_LOGI("Running partition ptr: %p", joctx->running_partition);
+jade_ota_ctx_t* ota_init(jade_process_t* process, const bool is_delta)
+{
+    JADE_ASSERT(process);
+    jade_ota_ctx_t* joctx = NULL;
+    const char* errmsg = NULL;
+    int errcode = CBOR_RPC_BAD_PARAMETERS;
+
+    // We expect a current message to be present
+    ASSERT_CURRENT_MESSAGE(process, is_delta ? "ota_delta" : "ota");
+    GET_MSG_PARAMS(process);
+
+    const jade_msg_source_t ota_source = process->ctx.source;
+    if (keychain_has_pin()) {
+        // NOTE: ota from internal source is allowed (eg. QR codes or USB storage)
+        JADE_ASSERT(ota_source == (jade_msg_source_t)keychain_get_userdata() || ota_source == SOURCE_INTERNAL);
+        JADE_ASSERT(!keychain_has_temporary());
+    }
 
-    // Check partition
-    joctx->update_partition = esp_ota_get_next_update_partition(NULL);
-    JADE_LOGI("Update partition: %p", joctx->update_partition);
+    size_t firmwaresize = 0;
+    size_t compressedsize = 0;
+    size_t uncompressedpatchsize = 0;
 
-    if (joctx->update_partition == NULL) {
-        JADE_LOGE("Failed to get next update partition");
-        return false;
+    if (!rpc_get_sizet("fwsize", &params, &firmwaresize) || !rpc_get_sizet("cmpsize", &params, &compressedsize)
+        || firmwaresize <= compressedsize) {
+        errmsg = "Bad filesize parameters";
+        goto cleanup;
+    }
+    if (is_delta
+        && (!rpc_get_sizet("patchsize", &params, &uncompressedpatchsize) || uncompressedpatchsize <= compressedsize)) {
+        errmsg = "Bad delta filesize parameters";
+        goto cleanup;
     }
 
-    if (joctx->update_partition == joctx->running_partition) {
-        JADE_LOGE("Cannot OTA on running partition: %p", joctx->running_partition);
-        return false;
+    // Optional field indicating preference for rich reply data
+    bool extended_replies = false;
+    rpc_get_boolean("extended_replies", &params, &extended_replies);
+
+    // Can accept either uploaded file data hash (legacy) or hash of the full/final firmware image (preferred)
+    uint8_t expected_hash[SHA256_LEN];
+    hash_type_t hash_type;
+    if (rpc_get_n_bytes("fwhash", &params, sizeof(expected_hash), expected_hash)) {
+        hash_type = HASHTYPE_FULLFWDATA;
+    } else if (rpc_get_n_bytes("cmphash", &params, sizeof(expected_hash), expected_hash)) {
+        hash_type = HASHTYPE_FILEDATA;
+    } else {
+        errmsg = "Cannot extract valid fw hash value";
+        goto cleanup;
     }
+    errcode = CBOR_RPC_INTERNAL_ERROR; // From this point all errors are internal
 
+    // We will show a progress bar once the user has confirmed and the upload in progress
+    // Initially just show a message screen.
+    const char* message[] = { "Preparing for firmware", "", "update" };
+    display_message_activity(message, 3);
+    vTaskDelay(100 / portTICK_PERIOD_MS); // sleep a little bit to redraw screen
+
+    joctx = JADE_CALLOC_PREFER_SPIRAM(1, sizeof(jade_ota_ctx_t));
+    jade_process_call_on_exit(process, ota_free, joctx);
+
+    mbedtls_sha256_init(&joctx->sha_ctx);
     JADE_ZERO_VERIFY(mbedtls_sha256_starts(&joctx->sha_ctx, 0));
+    joctx->hash_type = hash_type;
+    JADE_STATIC_ASSERT(sizeof(joctx->expected_hash) == sizeof(expected_hash));
+    memcpy(joctx->expected_hash, expected_hash, sizeof(expected_hash));
+    JADE_WALLY_VERIFY(wally_hex_from_bytes(expected_hash, sizeof(expected_hash), &joctx->expected_hash_hexstr));
+    joctx->ota_return_status = OTA_ERR_SETUP;
+    joctx->expected_source = ota_source;
+    joctx->compressedsize = compressedsize;
+    joctx->remaining_compressed = joctx->compressedsize;
+    joctx->uncompressedsize = is_delta ? uncompressedpatchsize : firmwaresize;
+    joctx->remaining_uncompressed = joctx->uncompressedsize;
+    joctx->firmwaresize = firmwaresize;
+    joctx->fwwritten = 0;
+    joctx->extended_replies = extended_replies;
+    joctx->validated_confirmed = false;
 
-    const esp_err_t err = esp_ota_begin(joctx->update_partition, joctx->firmwaresize, &joctx->ota_handle);
-    if (err != ESP_OK) {
-        JADE_LOGE("Failed to begin ota, error: %d", err);
-        return false;
-    }
+    joctx->running_partition = esp_ota_get_running_partition();
+    joctx->update_partition = esp_ota_get_next_update_partition(NULL);
 
-    return true;
-}
+    // Check partitions
+    if (joctx->running_partition == NULL) {
+        errmsg = "Failed to get running partition";
+        goto cleanup;
+    } else if (joctx->update_partition == NULL) {
+        errmsg = "Failed to get next update partition";
+        goto cleanup;
+    } else if (joctx->update_partition == joctx->running_partition) {
+        errmsg = "Cannot OTA on running partition";
+        goto cleanup;
+    } else {
+        const esp_err_t err = esp_ota_begin(joctx->update_partition, joctx->firmwaresize, &joctx->ota_handle);
+        if (err != ESP_OK) {
+            errmsg = "Failed to begin ota";
+            goto cleanup;
+        }
+    }
 
-void ota_free(jade_ota_ctx_t* joctx)
-{
-    if (joctx) {
-        mbedtls_sha256_free(&joctx->sha_ctx);
+cleanup:
+    if (errmsg) {
+        JADE_LOGE("%s", errmsg);
+        jade_process_reject_message(process, errcode, errmsg);
     }
+    return joctx;
 }
 
-ota_status_t post_ota_check(jade_ota_ctx_t* joctx, bool* ota_end_called)
+ota_status_t post_ota_check(jade_ota_ctx_t* joctx)
 {
     JADE_ASSERT(joctx);
-    JADE_ASSERT(ota_end_called);
 
     // Ensure no cached error - if so return it now
     if (joctx->ota_return_status != OTA_SUCCESS) {
@@ -205,7 +282,7 @@ ota_status_t post_ota_check(jade_ota_ctx_t* joctx, bool* ota_end_called)
 
     // All good, finalise the ota and set the partition to boot
     esp_err_t err = esp_ota_end(joctx->ota_handle);
-    *ota_end_called = true;
+    joctx->ota_handle = 0;
 
     if (err != ESP_OK) {
         JADE_LOGE("esp_ota_end() returned %d", err);
diff --git a/main/process/ota_util.h b/main/process/ota_util.h
index 8234335..0e33a9c 100644
--- a/main/process/ota_util.h
+++ b/main/process/ota_util.h
@@ -4,6 +4,7 @@
 #include "../process.h"
 #include "../ui.h"
 #include "../utils/cbor_rpc.h"
+#include <deflate.h>
 #include <esp_app_format.h>
 #include <esp_ota_ops.h>
 #include <esp_partition.h>
@@ -52,21 +53,21 @@ typedef enum {
 typedef struct {
     // Context used to compute (compressed) firmware hash - ie. file as uploaded
     mbedtls_sha256_context sha_ctx;
+    struct deflate_ctx dctx;
     progress_bar_t progress_bar;
     hash_type_t hash_type;
     char id[MAXLEN_ID + 1];
-    const uint8_t* expected_hash;
-    const char* expected_hash_hexstr;
+    uint8_t expected_hash[32];
+    char* expected_hash_hexstr;
     const esp_partition_t* running_partition;
     const esp_partition_t* update_partition;
     esp_ota_handle_t ota_handle;
     ota_status_t ota_return_status;
-    struct deflate_ctx* dctx;
-    const jade_msg_source_t expected_source;
-    size_t remaining_uncompressed;
+    jade_msg_source_t expected_source;
+    size_t compressedsize;
     size_t remaining_compressed;
     size_t uncompressedsize;
-    size_t compressedsize;
+    size_t remaining_uncompressed;
     size_t firmwaresize;
     size_t fwwritten;
     bool extended_replies;
@@ -75,9 +76,8 @@ typedef struct {
 
 void handle_in_bin_data(void* ctx, uint8_t* data, size_t rawsize);
 
-bool ota_init(jade_ota_ctx_t* joctx);
-void ota_free(jade_ota_ctx_t* joctx);
-ota_status_t post_ota_check(jade_ota_ctx_t* joctx, bool* ota_end_called);
+jade_ota_ctx_t* ota_init(jade_process_t* process, bool is_delta);
+ota_status_t post_ota_check(jade_ota_ctx_t* joctx);
 ota_status_t ota_user_validation(jade_ota_ctx_t* joctx, const uint8_t* uncompressed);
 const char* ota_get_status_text(ota_status_t status);