OTA_Handler.h

Go to the documentation of this file.

#ifndef OTA_Handler_h
#define OTA_Handler_h
 
// Local include.
#include "Configuration.h"
 
// Local include.
#include "Callback_Watchdog.h"
#include "HashGenerator.h"
#include "OTA_Update_Callback.h"
#include "OTA_Failure_Response.h"
#include "Helper.h"
 
// Library includes.
#include <string.h>
 
 
// Firmware data keys.
char constexpr FW_STATE_DOWNLOADING[] = "DOWNLOADING";
char constexpr FW_STATE_DOWNLOADED[] = "DOWNLOADED";
char constexpr FW_STATE_UPDATING[] = "UPDATING";
char constexpr FW_STATE_FAILED[] = "FAILED";
char constexpr FW_STATE_UPDATED[] = "UPDATED";
 
// Log messages.
char constexpr OTA_CB_IS_NULL[] = "OTA update callback is NULL, has it been deleted";
char constexpr UNABLE_TO_REQUEST_CHUNCKS[] = "Unable to request firmware chunk";
char constexpr RECEIVED_UNEXPECTED_CHUNK[] = "Received chunk (%u), not the same as requested chunk (%u)";
char constexpr RECEIVED_UNEXPECTED_CHUNK_SIZE[] = "Received chunk size (%u), not the same as expected chunk size (%u)";
char constexpr ERROR_UPDATE_BEGIN[] = "Failed to initalize flash updater, ensure that the partition scheme has two app sections";
char constexpr ERROR_UPDATE_WRITE[] = "Only wrote (%u) bytes of binary data instead of expected (%u)";
char constexpr ERROR_UPDATE_END[] = "Error during flash updater not all bytes written";
char constexpr CHECKSUM_VERIFICATION_FAILED[] = "Calculated checksum (%s), not the same as expected checksum (%s)";
char constexpr FW_UPDATE_ABORTED[] = "Firmware update aborted";
char constexpr CHUNK_REQUEST_TIMED_OUT[] = "Failed to receive requested chunk (%u) in (%llu) us. Internet connection might have been lost";
#if THINGSBOARD_ENABLE_DEBUG
char constexpr FW_CHUNK[] = "Receive chunk (%u), with size (%u) bytes";
char constexpr HASH_EXPECTED[] = "Expected checksum: (%s)";
char constexpr CHECKSUM_VERIFICATION_SUCCESS[] = "Checksum is the same as expected";
char constexpr FW_UPDATE_SUCCESS[] = "Update success";
#endif // THINGSBOARD_ENABLE_DEBUG
 
 
template <typename Logger>
class OTA_Handler {
  public:
    OTA_Handler(Callback<bool, size_t const &, size_t const &>::function publish_callback, Callback<bool, char const * const, char const * const>::function send_fw_state_callback, Callback<bool>::function finish_callback)
      : m_fw_callback(nullptr)
      , m_publish_callback(publish_callback)
      , m_send_fw_state_callback(send_fw_state_callback)
      , m_finish_callback(finish_callback)
      , m_fw_size(0U)
      , m_fw_checksum()
      , m_fw_checksum_algorithm()
      , m_hash()
      , m_total_chunks(0U)
      , m_requested_chunks(0U)
      , m_retries(0U)
    {
#if !THINGSBOARD_ENABLE_STL
        m_subscribedInstance = nullptr;
#endif // !THINGSBOARD_ENABLE_STL
    }
 
    void Start_Firmware_Update(OTA_Update_Callback & fw_callback, size_t const & fw_size, char const * fw_checksum, mbedtls_md_type_t const & fw_checksum_algorithm) {
        m_fw_callback = &fw_callback;
        m_fw_size = fw_size;
        m_total_chunks = (m_fw_size / m_fw_callback->Get_Chunk_Size()) + 1U;
        (void)strncpy(m_fw_checksum, fw_checksum, sizeof(m_fw_checksum));
        m_fw_checksum_algorithm = fw_checksum_algorithm;
        auto & request_timeout = m_fw_callback->Get_Request_Timeout();
#if THINGSBOARD_ENABLE_STL
        request_timeout.Set_Timeout_Callback(std::bind(&OTA_Handler::Handle_Request_Timeout, this));
#else
        request_timeout.Set_Timeout_Callback(OTA_Handler::Static_Handle_Request_Timeout);
#endif // THINGSBOARD_ENABLE_STL
        Request_First_Firmware_Packet();
        (void)m_send_fw_state_callback.Call_Callback(FW_STATE_DOWNLOADING, "");
    }
 
    void Stop_Firmware_Update()  {
        auto & request_timeout = m_fw_callback->Get_Request_Timeout();
        request_timeout.Stop_Timeout_Timer();
        auto fw_updater = m_fw_callback->Get_Updater();
        fw_updater->reset();
        Logger::printfln(FW_UPDATE_ABORTED);
        Handle_Failure(OTA_Failure_Response::RETRY_NOTHING, FW_UPDATE_ABORTED);
        m_fw_callback = nullptr;
    }
 
    void Process_Firmware_Packet(size_t const & current_chunk, uint8_t * payload, size_t const & total_bytes)  {
        if (current_chunk != m_requested_chunks) {
            Logger::printfln(RECEIVED_UNEXPECTED_CHUNK, current_chunk, m_requested_chunks);
            return;
        }
        size_t expected_chunk_size = 0U;
        if (!Received_Valid_Chunk_Size(total_bytes, expected_chunk_size)) {
            Logger::printfln(RECEIVED_UNEXPECTED_CHUNK_SIZE, expected_chunk_size, total_bytes);
            return;
        }
 
        auto & request_timeout = m_fw_callback->Get_Request_Timeout();
        request_timeout.Stop_Timeout_Timer();
    #if THINGSBOARD_ENABLE_DEBUG
        Logger::printfln(FW_CHUNK, current_chunk, total_bytes);
    #endif // THINGSBOARD_ENABLE_DEBUG
 
        auto fw_updater = m_fw_callback->Get_Updater();
        if (current_chunk == 0U && !fw_updater->begin(m_fw_size)) {
            Logger::printfln(ERROR_UPDATE_BEGIN);
            return Handle_Failure(OTA_Failure_Response::RETRY_UPDATE, ERROR_UPDATE_BEGIN);
        }
 
        auto const written_bytes = fw_updater->write(payload, total_bytes);
        if (written_bytes != total_bytes) {
            char message[Helper::Calculate_Print_Size(ERROR_UPDATE_WRITE, written_bytes, total_bytes)] = {};
            (void)snprintf(message, sizeof(message), ERROR_UPDATE_WRITE, written_bytes, total_bytes);
            Logger::printfln(message);
            return Handle_Failure(OTA_Failure_Response::RETRY_UPDATE, message);
        }
 
        // Update hash value only if writing with updater implementation was a success, result is ignored,
        // because it can only fail if the input parameters are invalid
        (void)m_hash.update(payload, total_bytes);
 
        m_requested_chunks = current_chunk + 1;
        m_fw_callback->Call_Progress_Callback(m_requested_chunks, m_total_chunks);
 
        // Ensure to check if the update was cancelled during the progress callback,
        // if it was the callback variable was reset and there is no need to request the next firmware packet
        if (m_fw_callback == nullptr) {
            Logger::printfln(OTA_CB_IS_NULL);
            return Handle_Failure(OTA_Failure_Response::RETRY_NOTHING, OTA_CB_IS_NULL);
        }
 
        Reset_Retries();
        Request_Next_Firmware_Packet();
    }
 
#if !THINGSBOARD_USE_ESP_TIMER
    void update() {
        auto & request_timeout = m_fw_callback->Get_Request_Timeout();
        request_timeout.Update_Timeout_Timer();
    }
#endif // !THINGSBOARD_USE_ESP_TIMER
 
  private:
    void Reset_Retries() {
        m_retries = m_fw_callback->Get_Chunk_Retries();
    }
 
    bool Received_Valid_Chunk_Size(size_t const & received_chunk_size, size_t & expected_chunk_size) {
        bool const is_last_chunk = m_requested_chunks + 1 >= m_total_chunks;
        if (is_last_chunk) {
            auto const last_chunk_expected_size = m_fw_size % m_fw_callback->Get_Chunk_Size();
            expected_chunk_size = last_chunk_expected_size;
            return received_chunk_size == last_chunk_expected_size;
        }
        expected_chunk_size = m_fw_callback->Get_Chunk_Size();
        return received_chunk_size == m_fw_callback->Get_Chunk_Size();
    }
 
    void Request_First_Firmware_Packet()  {
        m_requested_chunks = 0U;
        Reset_Retries();
        // Hash start result is ignored, because it can only fail if the input parameters are invalid
        (void)m_hash.start(m_fw_checksum_algorithm);
        auto & request_timeout = m_fw_callback->Get_Request_Timeout();
        request_timeout.Stop_Timeout_Timer();
        auto fw_updater = m_fw_callback->Get_Updater();
        fw_updater->reset();
        Request_Next_Firmware_Packet();
    }
 
    void Request_Next_Firmware_Packet()  {
        // Check if we have already requested and handled the last remaining chunk
        if (m_requested_chunks >= m_total_chunks) {
            Finish_Firmware_Update();   
            return;
        }
 
        if (!m_publish_callback.Call_Callback(m_fw_callback->Get_Request_ID(), m_requested_chunks)) {
            Logger::printfln(UNABLE_TO_REQUEST_CHUNCKS);
        }
 
        // Request timeout gets started no matter if publishing previous request was successful or not in hopes,
        // that after the given timeout the callback calls this method again and can then publish the request successfully.
        // This works because the request fails most of the time, because the internet connection might have been temporarily disconnected.
        // Therefore waiting a while and then retrying, means we might be reconnected again
        auto & request_timeout = m_fw_callback->Get_Request_Timeout();
        request_timeout.Start_Timeout_Timer();
    }
 
    void Finish_Firmware_Update()  {
        (void)m_send_fw_state_callback.Call_Callback(FW_STATE_DOWNLOADED, "");
        auto const calculated_checksum = m_hash.finish();
 
        if (strncmp(m_fw_checksum, calculated_checksum.hash, strlen(m_fw_checksum)) != 0) {
            char message[Helper::Calculate_Print_Size(CHECKSUM_VERIFICATION_FAILED, calculated_checksum.hash, m_fw_checksum)] = {};
            (void)snprintf(message, sizeof(message), CHECKSUM_VERIFICATION_FAILED, calculated_checksum.hash, m_fw_checksum);
            Logger::printfln(message);
            return Handle_Failure(OTA_Failure_Response::RETRY_UPDATE, message);
        }
 
    #if THINGSBOARD_ENABLE_DEBUG
        Logger::printfln(CHECKSUM_VERIFICATION_SUCCESS);
    #endif // THINGSBOARD_ENABLE_DEBUG
 
        auto fw_updater = m_fw_callback->Get_Updater();
        if (!fw_updater->end()) {
            Logger::printfln(ERROR_UPDATE_END);
            return Handle_Failure(OTA_Failure_Response::RETRY_UPDATE, ERROR_UPDATE_END);
        }
 
    #if THINGSBOARD_ENABLE_DEBUG
        Logger::printfln(FW_UPDATE_SUCCESS);
    #endif // THINGSBOARD_ENABLE_DEBUG
 
        (void)m_send_fw_state_callback.Call_Callback(FW_STATE_UPDATING, "");
        m_fw_callback->Call_Callback(true);
        (void)m_finish_callback.Call_Callback();
    }
 
    void Handle_Failure(OTA_Failure_Response failure_response, char const * error_message)  {
        if (m_retries <= 0) {
            Abort_Firmware_Update(error_message);
            return;
        }
 
        // Decrease the amount of retries of downloads for the current chunk,
        // reset as soon as the next chunk has been received and handled successfully
        m_retries--;
 
        switch (failure_response) {
            case OTA_Failure_Response::RETRY_CHUNK:
                Request_Next_Firmware_Packet();
                break;
            case OTA_Failure_Response::RETRY_UPDATE:
                Request_First_Firmware_Packet();
                break;
            case OTA_Failure_Response::RETRY_NOTHING:
                Abort_Firmware_Update(error_message);
                break;
        }
    }
 
    void Abort_Firmware_Update(char const * error_message) {
        (void)m_send_fw_state_callback.Call_Callback(FW_STATE_FAILED, error_message);
        if (m_fw_callback != nullptr) {
            m_fw_callback->Call_Callback(false);
        }
        (void)m_finish_callback.Call_Callback();
    }
 
    void Handle_Request_Timeout()  {
        auto const & request_timeout = m_fw_callback->Get_Request_Timeout();
        uint64_t const & timeout = request_timeout.Get_Timeout();
        char message[Helper::Calculate_Print_Size(CHUNK_REQUEST_TIMED_OUT, m_requested_chunks, timeout)] = {};
        (void)snprintf(message, sizeof(message), CHUNK_REQUEST_TIMED_OUT, m_requested_chunks, timeout);
        Logger::printfln(message);
        Handle_Failure(OTA_Failure_Response::RETRY_CHUNK, message);
    }
 
#if !THINGSBOARD_ENABLE_STL
    static bool Static_Handle_Request_Timeout() {
        if (m_subscribedInstance == nullptr) {
            return false;
        }
        return m_subscribedInstance->Handle_Request_Timeout();
    }
 
    static OTA_Handler *m_subscribedInstance;
#endif // !THINGSBOARD_ENABLE_STL
 
    OTA_Update_Callback                                    *m_fw_callback = {};                            // Callback method that contains configuration information, about the over the air update
    Callback<bool, size_t const &, size_t const &>         m_publish_callback = {};                  // Callback that is used to request the firmware chunk of the firmware binary with the given chunk number
    Callback<bool, char const * const, char const * const> m_send_fw_state_callback = {};            // Callback that is used to send information about the current state of the over the air update
    Callback<bool>                                         m_finish_callback = {};                   // Callback that is called once the update has been finished and the user should be informed of the failure or success of the over the air update
    size_t                                                 m_fw_size = {};                           // Total size of the firmware binary we will receive. Allows for a binary size of up to theoretically 4 GB
    char                                                   m_fw_checksum[MAX_STRING_HASH_SIZE] = {}; // Checksum of the complete firmware binary, should be the same as the actually written data in the end
    mbedtls_md_type_t                                      m_fw_checksum_algorithm = {};             // Algorithm type used to hash the firmware binary
    HashGenerator                                          m_hash = {};                              // Class instance that allows to generate a hash from received firmware binary data
    size_t                                                 m_total_chunks = {};                      // Total amount of chunks that need to be received to get the complete firmware binary
    size_t                                                 m_requested_chunks = {};                  // Amount of successfully requested and received firmware binary chunks
    uint8_t                                                m_retries = {};                           // Amount of request retries we attempt for each chunk, increasing makes the connection more stable
};
 
#if !THINGSBOARD_ENABLE_STL
OTA_Handler *OTA_Handler::m_subscribedInstance = nullptr;
#endif
 
#endif // OTA_Handler_h