feat(app): Add decryption bidirectional chunking

apps/inheritance_app/inheritance_decrypt_data.c

@@ -72,11 +72,16 @@
 #include "inheritance_api.h"
 #include "inheritance_context.h"
 #include "inheritance_priv.h"
+#include "memzero.h"
+#include "pb.h"
+#include "pb_decode.h"
+#include "pb_encode.h"
 #include "reconstruct_wallet_flow.h"
 #include "status_api.h"
 #include "ui_core_confirm.h"
 #include "ui_screens.h"
 #include "utils.h"
+#include "wallet.h"
 #include "wallet_list.h"
 
 /*****************************************************************************
@@ -86,7 +91,7 @@
 /*****************************************************************************
  * PRIVATE MACROS AND DEFINES
  *****************************************************************************/
-
+#define DECRYPTED_CHUNK_SIZE (2048)
 /*****************************************************************************
  * PRIVATE TYPEDEFS
  *****************************************************************************/
@@ -122,6 +127,48 @@ static bool validate_request_data(
  */
 STATIC bool inheritance_handle_initiate_query(inheritance_query_t *query);
 
+/**
+ * @brief Function responsible for decoding pb_encoded buffer to @ref
+ * inheritance_decrypt_data_with_pin_encrypted_data_structure_t
+ *
+ * @return true if decoding successful, false otherwise.
+ */
+static bool decode_inheritance_encrypted_data(
+    const uint8_t *data,
+    uint16_t data_size,
+    inheritance_decrypt_data_with_pin_encrypted_data_structure_t
+        *encrypted_data);
+
+/**
+ * @brief Retrieves pb_encoded encrypted data chunks from the host and creates
+ * input buffer for @ref decode_inheritance_encrypted_data().
+ *
+ * @return true if data is successfully retrieved and decoded, false otherwise.
+ */
+static bool inheritance_get_encrypted_data(inheritance_query_t *query);
+
+/**
+ * @brief Creates a pb encoded buffer of @ref
+ * inheritance_decrypt_data_with_pin_decrypted_data_structure_t to be sent to
+ * host
+ *
+ * @return true if encoding is successful, false otherwise.
+ */
+static bool get_pb_encoded_buffer(
+    const inheritance_decrypt_data_with_pin_decrypted_data_structure_t *result,
+    uint8_t *buffer,
+    uint16_t max_buffer_len,
+    size_t *bytes_written_out);
+
+/**
+ * @brief Sends input buffer to host in chunks
+ *
+ * @return true if chunking successful, false otherwise.
+ */
+static bool inheritance_send_in_chunks(inheritance_query_t *query,
+                                       const uint8_t *buffer,
+                                       const size_t buffer_len);
+
 /**
  * @brief Sends the decrypted data as a response.
  *
@@ -159,11 +206,12 @@ static bool decrypt_message_data(void);
 static bool decrypt_data(void);
 
 /**
- * @brief Displays decrypted data or performs related actions based on the data.
+ * @brief Displays decrypted data or performs related actions based on the
+ * data.
  *
- * This function processes the decrypted data and either shows it on the device
- * screen (if the tag indicates display-only) or prepares it for further
- * response.
+ * This function processes the decrypted data and either shows it on the
+ * device screen (if the tag indicates display-only) or prepares it for
+ * further response.
  *
  * @return true if the operation is successful, false otherwise.
  */
@@ -233,61 +281,211 @@ STATIC bool inheritance_handle_initiate_query(inheritance_query_t *query) {
 
   set_app_flow_status(INHERITANCE_DECRYPT_DATA_STATUS_USER_CONFIRMED);
 
-  decryption_context->request_pointer = &(query->decrypt.initiate);
+  memcpy(decryption_context->wallet_id,
+         query->decrypt.initiate.wallet_id,
+         WALLET_ID_SIZE);
+  inheritance_result_t result =
+      init_inheritance_result(INHERITANCE_RESULT_DECRYPT_TAG);
+  result.decrypt.which_response =
+      INHERITANCE_DECRYPT_DATA_WITH_PIN_RESPONSE_CONFIRMATION_TAG;
+  inheritance_send_result(&result);
+  delay_scr_init(ui_text_processing, DELAY_TIME);
   return true;
 }
 
-static bool send_decrypted_data(inheritance_query_t *query) {
-  inheritance_result_t result =
+static bool decode_inheritance_encrypted_data(
+    const uint8_t *data,
+    uint16_t data_size,
+    inheritance_decrypt_data_with_pin_encrypted_data_structure_t
+        *encrypted_data) {
+  if (NULL == data || NULL == encrypted_data || 0 == data_size) {
+    inheritance_send_error(ERROR_COMMON_ERROR_CORRUPT_DATA_TAG,
+                           ERROR_DATA_FLOW_DECODING_FAILED);
+    return false;
+  }
+
+  memzero(encrypted_data,
+          sizeof(inheritance_decrypt_data_with_pin_encrypted_data_structure_t));
+
+  pb_istream_t stream = pb_istream_from_buffer(data, data_size);
+
+  bool status = pb_decode(
+      &stream,
+      INHERITANCE_DECRYPT_DATA_WITH_PIN_ENCRYPTED_DATA_STRUCTURE_FIELDS,
+      encrypted_data);
+
+  if (false == status) {
+    inheritance_send_error(ERROR_COMMON_ERROR_CORRUPT_DATA_TAG,
+                           ERROR_DATA_FLOW_DECODING_FAILED);
+  }
+
+  return status;
+}
+
+static bool inheritance_get_encrypted_data(inheritance_query_t *query) {
+  uint8_t encoded_data[INHERITANCE_PACKET_MAX_SIZE] = {0};
+  inheritance_result_t response =
       init_inheritance_result(INHERITANCE_RESULT_DECRYPT_TAG);
-  result.decrypt.which_response =
-      INHERITANCE_DECRYPT_DATA_WITH_PIN_RESPONSE_MESSAGES_TAG;
-  if (!inheritance_get_query(query, INHERITANCE_QUERY_DECRYPT_TAG) ||
-      !check_which_request(
-          query, INHERITANCE_DECRYPT_DATA_WITH_PIN_REQUEST_INITIATE_TAG)) {
+  const inheritance_decrypt_data_with_pin_encrypted_data_t *encrypted_data =
+      &(query->decrypt.encrypted_data);
+  const common_chunk_payload_t *payload = &(encrypted_data->chunk_payload);
+  const common_chunk_payload_chunk_t *chunk = &(payload->chunk);
+  uint32_t total_size = 0;
+  uint32_t size = 0;
+  while (1) {
+    // req plain data chunk from host
+    if (!inheritance_get_query(query, INHERITANCE_QUERY_DECRYPT_TAG) ||
+        !check_which_request(
+            query,
+            INHERITANCE_DECRYPT_DATA_WITH_PIN_REQUEST_ENCRYPTED_DATA_TAG)) {
+      return false;
+    }
+    if (size == 0) {
+      total_size = chunk->size + payload->remaining_size;
+    }
+
+    if (false == query->decrypt.encrypted_data.has_chunk_payload ||
+        payload->chunk_index >= payload->total_chunks ||
+        size + chunk->size > total_size) {
+      inheritance_send_error(ERROR_COMMON_ERROR_CORRUPT_DATA_TAG,
+                             ERROR_DATA_FLOW_INVALID_DATA);
+      return false;
+    }
+
+    memcpy(encoded_data + size, chunk->bytes, chunk->size);
+    size += chunk->size;
+
+    // Send chunk ack to host
+    response.decrypt.which_response =
+        INHERITANCE_DECRYPT_DATA_WITH_PIN_RESPONSE_DATA_ACCEPTED_TAG;
+    response.decrypt.data_accepted.has_chunk_ack = true;
+    response.decrypt.data_accepted.chunk_ack.chunk_index = payload->chunk_index;
+    inheritance_send_result(&response);
+
+    // If no data remaining to be received from the host, then exit
+    if (0 == payload->remaining_size ||
+        payload->chunk_index + 1 == payload->total_chunks) {
+      break;
+    }
+  }
+  if (!decode_inheritance_encrypted_data(
+          encoded_data, total_size, &decryption_context->encrypted_data)) {
     return false;
   }
+  return true;
+}
 
-  memcpy(&result.decrypt.messages,
-         &decryption_context->response,
-         sizeof(inheritance_decrypt_data_with_pin_messages_response_t));
+static bool get_pb_encoded_buffer(
+    const inheritance_decrypt_data_with_pin_decrypted_data_structure_t *result,
+    uint8_t *buffer,
+    uint16_t max_buffer_len,
+    size_t *bytes_written_out) {
+  if (NULL == result || NULL == buffer || NULL == bytes_written_out) {
+    return false;
+  }
+  pb_ostream_t stream = pb_ostream_from_buffer(buffer, max_buffer_len);
 
-  inheritance_send_result(&result);
+  bool status = pb_encode(
+      &stream,
+      INHERITANCE_DECRYPT_DATA_WITH_PIN_DECRYPTED_DATA_STRUCTURE_FIELDS,
+      result);
+
+  if (true == status) {
+    *bytes_written_out = stream.bytes_written;
+  }
+
+  return status;
+}
+
+static bool inheritance_send_in_chunks(inheritance_query_t *query,
+                                       const uint8_t *buffer,
+                                       const size_t buffer_len) {
+  size_t total_count =
+      ((buffer_len + DECRYPTED_CHUNK_SIZE - 1) / DECRYPTED_CHUNK_SIZE);
+  size_t remaining_size = (size_t)buffer_len;
+  size_t offset = 0;
+  inheritance_result_t result =
+      init_inheritance_result(INHERITANCE_RESULT_DECRYPT_TAG);
+  result.decrypt.which_response =
+      INHERITANCE_DECRYPT_DATA_WITH_PIN_RESPONSE_DECRYPTED_DATA_TAG;
+  uint32_t *index = &result.decrypt.decrypted_data.chunk_payload.chunk_index;
+  result.decrypt.decrypted_data.chunk_payload.total_chunks = total_count;
+
+  for (*index = 0; *index < total_count; (*index)++) {
+    if (!inheritance_get_query(query, INHERITANCE_QUERY_DECRYPT_TAG) ||
+        !check_which_request(
+            query,
+            INHERITANCE_DECRYPT_DATA_WITH_PIN_REQUEST_DECRYPTED_DATA_REQUEST_TAG)) {
+      return false;
+    }
+    // chunk_payload validation checks
+    if (query->decrypt.decrypted_data_request.has_chunk_ack == false ||
+        query->decrypt.decrypted_data_request.chunk_ack.chunk_index != *index) {
+      return false;
+    }
+    size_t chunk_size = (remaining_size > DECRYPTED_CHUNK_SIZE)
+                            ? DECRYPTED_CHUNK_SIZE
+                            : remaining_size;
+    remaining_size -= chunk_size;
+    result.decrypt.decrypted_data.chunk_payload.remaining_size = remaining_size;
+    result.decrypt.decrypted_data.has_chunk_payload = true;
+    memcpy(result.decrypt.decrypted_data.chunk_payload.chunk.bytes,
+           buffer + offset,
+           chunk_size);
+    result.decrypt.decrypted_data.chunk_payload.chunk.size = chunk_size;
+    inheritance_send_result(&result);
+    offset += chunk_size;
+    if (remaining_size == 0) {
+      break;
+    }
+  }
+  return true;
+}
+
+static bool send_decrypted_data(inheritance_query_t *query) {
+  uint8_t
+      buffer[INHERITANCE_DECRYPT_DATA_WITH_PIN_DECRYPTED_DATA_STRUCTURE_SIZE] =
+          {0};
+  size_t bytes_encoded = 0;
+  if (!get_pb_encoded_buffer(&decryption_context->response_payload,
+                             buffer,
+                             sizeof(buffer),
+                             &bytes_encoded) ||
+      !inheritance_send_in_chunks(query, buffer, bytes_encoded)) {
+    // TODO: throw decryption failed error
+    return false;
+  }
   return true;
 }
 
 static bool decrypt_packet(void) {
-  decryption_context->packet_size =
-      decryption_context->request_pointer->encrypted_data.size;
-  memcpy(decryption_context->packet,
-         decryption_context->request_pointer->encrypted_data.bytes,
-         decryption_context->packet_size);
-  return session_aes_decrypt(decryption_context->packet,
-                             &decryption_context->packet_size) ==
+  return session_aes_decrypt(decryption_context->encrypted_data.data.bytes,
+                             &decryption_context->encrypted_data.data.size) ==
          SESSION_DECRYPT_PACKET_SUCCESS;
 }
 
 static bool deserialize_packet(void) {
   uint16_t packet_index = 0;
-  decryption_context->data_count = decryption_context->packet[packet_index++];
+  decryption_context->data_count =
+      decryption_context->encrypted_data.data.bytes[packet_index++];
   for (uint8_t index = 0; index < decryption_context->data_count; index++) {
     packet_index++;    ///< Skip tag
 
-    decryption_context->data[index].encrypted_data_size =
-        U16_READ_BE_ARRAY(&decryption_context->packet[packet_index]);
+    decryption_context->data[index].encrypted_data_size = U16_READ_BE_ARRAY(
+        &decryption_context->encrypted_data.data.bytes[packet_index]);
     packet_index += 2;    ///< Read length
 
     memcpy(decryption_context->data[index].encrypted_data,
-           &decryption_context->packet[packet_index],
+           &decryption_context->encrypted_data.data.bytes[packet_index],
            decryption_context->data[index].encrypted_data_size);
     packet_index += decryption_context->data[index].encrypted_data_size;
   }
 
-  return packet_index <= decryption_context->packet_size;
+  return packet_index <= decryption_context->encrypted_data.data.size;
 }
 
 static bool decrypt_message_data(void) {
-  return card_fetch_decrypt_data(decryption_context->request_pointer->wallet_id,
+  return card_fetch_decrypt_data(decryption_context->wallet_id,
                                  decryption_context->data,
                                  decryption_context->data_count) ==
          CARD_OPERATION_SUCCESS;
@@ -316,6 +514,7 @@ static bool decrypt_data(void) {
 
   } while (0);
   set_app_flow_status(INHERITANCE_DECRYPT_DATA_STATUS_MESSAGE_DECRYPTED);
+  delay_scr_init(ui_text_processing, DELAY_TIME);
   return status;
 }
 
@@ -335,14 +534,17 @@ static bool show_data(void) {
       }
     } else {
       uint16_t offset = 1;    // Skip tag
-      decryption_context->response.plain_data[response_count].message.size =
+      decryption_context->response_payload.decrypted_data[response_count]
+          .message.size =
           U16_READ_BE_ARRAY(decryption_context->data[i].plain_data + offset);
       offset += 2;    // Skip length
-      memcpy(
-          decryption_context->response.plain_data[response_count].message.bytes,
-          decryption_context->data[i].plain_data + offset,
-          decryption_context->response.plain_data[response_count].message.size);
-      decryption_context->response.plain_data_count = ++response_count;
+      memcpy(decryption_context->response_payload.decrypted_data[response_count]
+                 .message.bytes,
+             decryption_context->data[i].plain_data + offset,
+             decryption_context->response_payload.decrypted_data[response_count]
+                 .message.size);
+      decryption_context->response_payload.decrypted_data_count =
+          ++response_count;
     }
   }
 
@@ -359,8 +561,9 @@ void inheritance_decrypt_data(inheritance_query_t *query) {
       sizeof(inheritance_decryption_context_t));
   memzero(decryption_context, sizeof(inheritance_decryption_context_t));
 
-  if (inheritance_handle_initiate_query(query) && decrypt_data() &&
-      show_data() && send_decrypted_data(query)) {
+  if (inheritance_handle_initiate_query(query) &&
+      inheritance_get_encrypted_data(query) && decrypt_data() && show_data() &&
+      send_decrypted_data(query)) {
     delay_scr_init(ui_text_inheritance_decryption_flow_success, DELAY_TIME);
   } else {
     delay_scr_init(ui_text_inheritance_decryption_flow_failure, DELAY_TIME);

apps/inheritance_app/inheritance_encrypt_data.c

@@ -356,6 +356,7 @@ STATIC bool inheritance_encryption_handle_inititate_query(
   result.encrypt.which_response =
       INHERITANCE_ENCRYPT_DATA_WITH_PIN_RESPONSE_CONFIRMATION_TAG;
   inheritance_send_result(&result);
+  delay_scr_init(ui_text_processing, DELAY_TIME);
   return true;
 }
 
@@ -596,6 +597,7 @@ static bool encrypt_data(void) {
     }
   } while (0);
 
+  delay_scr_init(ui_text_processing, DELAY_TIME);
   return status;
 }
 
@@ -705,4 +707,4 @@ void inheritance_encrypt_data(inheritance_query_t *query) {
   memzero(encryption_context, sizeof(inheritance_encryption_context_t));
   free(encryption_context);
   encryption_context = NULL;
-}
+}