Remote_Wifi_Switch/main/ota.c

/*
 * ota.c
 *
 *  Created on: Feb 13, 2023
 *      Author: Sword
 */


/* OTA example
   This example code is in the Public Domain (or CC0 licensed, at your option.)
   Unless required by applicable law or agreed to in writing, this
   software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
   CONDITIONS OF ANY KIND, either express or implied.
 */
#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_system.h"
#include "esp_event.h"
#include "esp_ota_ops.h"
#include "esp_app_format.h"
#include "esp_log.h"
#include "esp_flash_partitions.h"
#include "esp_partition.h"
#include "nvs.h"
#include "nvs_flash.h"
#include "driver/gpio.h"
#include "ota.h"


#define BUFFSIZE 2048
#define HASH_LEN 32 /* SHA-256 digest length */
#define LOG_LOCAL_LEVEL ESP_LOG_INFO

static const char *TAG = "OTA";

/*an ota data write buffer ready to write to the flash*/
static char *ota_write_data;
static ota_bytes_status_t ota_bytes = OTA_BYTES_WRITTEN_WAIT;
static uint32_t ota_file_size = OTA_WITH_SEQUENTIAL_WRITES;

esp_ota_handle_t update_handle;
const esp_partition_t *update_partition = NULL;
const esp_partition_t *configured = NULL;
const esp_partition_t *running = NULL;

static uint8_t ota_needed = OTA_NOT_NEEDED;
static int data_read = -1;
int binary_file_length;
bool image_header_was_checked;

#define OTA_URL_SIZE 256


/******************************************************/
/********** Static Functions for OTA-Process **********/
/******************************************************/
static void print_sha256 (const uint8_t *image_hash, const char *label)
{
	char hash_print[HASH_LEN * 2 + 1];
	hash_print[HASH_LEN * 2] = 0;
	for (int i = 0; i < HASH_LEN; ++i) {
		sprintf(&hash_print[i * 2], "%02x", image_hash[i]);
	}
	ESP_LOGI(TAG, "%s: %s", label, hash_print);
}

static void ota_start_session(void)
{
	/* update handle : set by esp_ota_begin(), must be freed via esp_ota_end() */
	update_handle = 0;
	update_partition = NULL;

	/* Get partition info of currently configured boot app (the partition for current app)*/
	configured = esp_ota_get_boot_partition();
	/* Get partition info of currently running app (the partition for current app)*/
	running = esp_ota_get_running_partition();


	ESP_LOGI(TAG, "Starting OTA session");

	/* The result of esp_ota_get_boot_partition() is usually the same as esp_ota_get_running_partition().
	 * The two results are not equal if the configured boot partition does not contain a valid app (meaning that
	 * the running partition will be an app that the bootloader chose via fallback) */
	if (configured != running) {
		ESP_LOGW(TAG, "Configured OTA boot partition at offset 0x%08x, but running from offset 0x%08x",
				(unsigned int)configured->address, (unsigned int)running->address);
		ESP_LOGW(TAG, "(This can happen if either the OTA boot data or preferred boot image become corrupted somehow.)");
	}
	ESP_LOGI(TAG, "Running partition type %d subtype %d (offset 0x%08x)",
			running->type, running->subtype, (unsigned int)running->address);

	/* The following line returns the next OTA app partition which should be written with a new firmware*/
	/* So, update_partition is the partition where the new image will be written*/
	update_partition = esp_ota_get_next_update_partition(NULL);

	assert(update_partition != NULL);
	ESP_LOGI(TAG, "Writing to partition subtype %d at offset 0x%x",
			update_partition->subtype, (unsigned int)update_partition->address);

	/*deal with all receive packet*/
	image_header_was_checked = false;

}

static void ota_process_incoming_image(void)
{
	esp_err_t err;

	if ((data_read < 0) && (OTA_NEEDED == ota_needed)) {
		ESP_LOGE(TAG, "Error: No OTA data have been received yet");
	}
	else if ((data_read > 0) && (OTA_NEEDED == ota_needed)) {
		if (image_header_was_checked == false) {
			esp_app_desc_t new_app_info;
			if (data_read > sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t) + sizeof(esp_app_desc_t)) {
				// check current version with downloading
				/* 1- Check the firmware version of the new fetched-app*/
				memcpy(&new_app_info, &ota_write_data[sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t)], sizeof(esp_app_desc_t));
				ESP_LOGI(TAG, "New firmware version: %s", new_app_info.version);

				/* 2- Check the firmware version of the current running app*/
				esp_app_desc_t running_app_info;
				if (esp_ota_get_partition_description(running, &running_app_info) == ESP_OK) {
					ESP_LOGI(TAG, "Running firmware version: %s", running_app_info.version);
				}

				/* 3- Check the firmware version of the last disabled/aborted/invalid app (which is an app that )*/
				const esp_partition_t* last_invalid_app = esp_ota_get_last_invalid_partition();
				esp_app_desc_t invalid_app_info;
				if (esp_ota_get_partition_description(last_invalid_app, &invalid_app_info) == ESP_OK) {
					/*---> Stop OTA updating*/
					ESP_LOGI(TAG, "Last invalid firmware version: %s", invalid_app_info.version);
				}

				// check current version with last invalid partition
				/* 4- Compare the new fetched-image with the last invalid-image*/
				if (last_invalid_app != NULL) {
					if (memcmp(invalid_app_info.version, new_app_info.version, sizeof(new_app_info.version)) == 0) {
						/*---> Stop OTA updating*/
						ESP_LOGW(TAG, "New version is the same as invalid version.");
						ESP_LOGW(TAG, "Previously, there was an attempt to launch the firmware with %s version, but it failed.", invalid_app_info.version);
						ESP_LOGW(TAG, "The firmware has been rolled back to the previous version.");

						/* EXIT OTA task*/
						ota_bytes = OTA_BYTES_WRITTEN_FAILED;
						return;
					}
				}

				/* 5- Compare the new-fetched-image with the current running-image*/
				if (memcmp(new_app_info.version, running_app_info.version, sizeof(new_app_info.version)) == 0) {
					/*---> Stop OTA updating*/
					ESP_LOGW(TAG, "Current running version is the same as a new. We will not continue the update.");

					/* EXIT OTA task*/
					ota_bytes = OTA_BYTES_WRITTEN_FAILED;
					return;
				}

				/* 7- Start an OTA-update writing to the specified partition.
							  The specified partition is erased to the specified image size.
							  If image size is not yet known, pass OTA_SIZE_UNKNOWN which will cause the entire partition to be erased.*/
				err = esp_ota_begin(update_partition, ota_file_size, &update_handle);
				if (err != ESP_OK) {
					/*---> Stop OTA updating*/
					ESP_LOGE(TAG, "esp_ota_begin failed (%s)", esp_err_to_name(err));
					esp_ota_abort(update_handle);
					ESP_LOGI(TAG, "Aborting the update-data partition");

					/* EXIT OTA task*/
					ota_bytes = OTA_BYTES_WRITTEN_FAILED;
					return;
				}

				/* 6- Set the image_header_checked flag to be true*/
				image_header_was_checked = true;

				ESP_LOGI(TAG, "esp_ota_begin succeeded");
			} else {
				/*---> Stop OTA updating*/
				ESP_LOGE(TAG, "received package is not fit len");
				esp_ota_abort(update_handle);
				ESP_LOGI(TAG, "Aborting the update-data partition");

				/* EXIT OTA task*/
				ota_bytes = OTA_BYTES_WRITTEN_FAILED;
				return;
			}
		}
		/* 8- Write OTA update data to partition.
		 *    This function can be called multiple times as data is received during the OTA operation.
		 *    Data is written sequentially to the partition */
		err = esp_ota_write( update_handle, (const void *)ota_write_data, data_read);
		if (err != ESP_OK) {
			/*---> Stop OTA updating*/
			esp_ota_abort(update_handle);
			ESP_LOGI(TAG, "Aborting the update-data partition");

			/* EXIT OTA task*/
			ota_bytes = OTA_BYTES_WRITTEN_FAILED;
			return;
		}
		/* 9- Set the binary file length*/
		binary_file_length += data_read;
		/* Reset data_read variable to 0*/
		data_read = 0;
		ESP_LOGD(TAG, "Written bytes length %d", binary_file_length);
		ota_bytes = OTA_BYTES_WRITTEN_SUCCESS;
	}
	else if ((data_read == 0) && (OTA_FINISHED == ota_needed)) {
		ESP_LOGI(TAG, "OTA session finished");
		ota_bytes = OTA_BYTES_WRITTEN_SUCCESS;
		return;
	}
}

static void ota_end_session(void)
{
	esp_err_t err;
	if(OTA_BYTES_WRITTEN_SUCCESS == ota_bytes)
	{
		ESP_LOGI(TAG, "Total Write binary data length: %d", binary_file_length);
		/* 10- Checks if entire data in the response has been read without any error */
		if ( OTA_FINISHED != ota_needed ) {
			ESP_LOGI(TAG, "OTA session stopped for errors but didn't get finished");
			esp_ota_abort(update_handle);
			ESP_LOGI(TAG, "Aborting the update-data partition");
			/* EXIT OTA task*/
			return ;
		}

		/* 11- Finish OTA update and validate newly written app image*/
		err = esp_ota_end(update_handle);
		if (err != ESP_OK) {
			/*---> Stop OTA updating*/
			if (err == ESP_ERR_OTA_VALIDATE_FAILED) {
				ESP_LOGI(TAG, "Image validation failed, image is corrupted");
			} else {
				ESP_LOGI(TAG, "esp_ota_end failed (%s)!", esp_err_to_name(err));
			}
			/* EXIT OTA task*/
			return ;
		}

		/* 12- Set the new validated-written image to be the boot partition (Configure OTA data for a new boot partition)
		 *     On the next restart, ESP will boot from that new partition*/
		err = esp_ota_set_boot_partition(update_partition);
		if (err != ESP_OK) {
			/*Stop OTA update*/
			ESP_LOGI(TAG, "esp_ota_set_boot_partition failed (%s)!", esp_err_to_name(err));

			/* EXIT OTA task*/
			return ;
		}
		/* 13- Restart the device (ESP32)*/
		ESP_LOGI(TAG, "Prepare to restart system!");
		esp_restart();
		return ;
	}
	else if(OTA_BYTES_WRITTEN_FAILED == ota_bytes)
	{
		ESP_LOGI(TAG,"No Changing in boot_partition");
	}
}

static bool diagnostic(void)
{
	/* gpio_config_t io_conf;
    io_conf.intr_type    = GPIO_INTR_DISABLE;
    io_conf.mode         = GPIO_MODE_INPUT;
    io_conf.pin_bit_mask = (1ULL << CONFIG_EXAMPLE_GPIO_DIAGNOSTIC);
    io_conf.pull_down_en = GPIO_PULLDOWN_DISABLE;
    io_conf.pull_up_en   = GPIO_PULLUP_ENABLE;
    gpio_config(&io_conf);

    ESP_LOGI(TAG, "Diagnostics (5 sec)...");
    vTaskDelay(5000 / portTICK_PERIOD_MS);

    bool diagnostic_is_ok = gpio_get_level(CONFIG_EXAMPLE_GPIO_DIAGNOSTIC);

    gpio_reset_pin(CONFIG_EXAMPLE_GPIO_DIAGNOSTIC);*/
	return true;
}


void ota_set_needed(void)
{
	ota_needed = OTA_NEEDED;
}

void ota_set_not_needed(void)
{
	ota_needed = OTA_NOT_NEEDED;
}

void ota_set_finished(void)
{
	ota_needed = OTA_FINISHED;
}

/******************************************************/
/* Functions that gonna be used in comms state machine*/
/******************************************************/
esp_err_t ota_init(uint32_t fileSize)
{
	/*This function should be invoked in the ota_file_open_cb() in comms.c (instead of the OTA_Offline_Init())*/
	uint8_t sha_256[HASH_LEN] = { 0 };
	esp_partition_t partition;

	// get sha256 digest for the partition table
	partition.address   = ESP_PARTITION_TABLE_OFFSET;
	partition.size      = ESP_PARTITION_TABLE_MAX_LEN;
	partition.type      = ESP_PARTITION_TYPE_DATA;
	esp_partition_get_sha256(&partition, sha_256);
	print_sha256(sha_256, "SHA-256 for the partition table: ");

	// get sha256 digest for bootloader
	partition.address   = ESP_BOOTLOADER_OFFSET;
	partition.size      = ESP_PARTITION_TABLE_OFFSET;
	partition.type      = ESP_PARTITION_TYPE_APP;
	esp_partition_get_sha256(&partition, sha_256);
	print_sha256(sha_256, "SHA-256 for bootloader: ");

	// get sha256 digest for running partition
	esp_partition_get_sha256(esp_ota_get_running_partition(), sha_256);
	print_sha256(sha_256, "SHA-256 for current firmware: ");

	const esp_partition_t *running = esp_ota_get_running_partition();
	esp_ota_img_states_t ota_state;
	if (esp_ota_get_state_partition(running, &ota_state) == ESP_OK) {
		if (ota_state == ESP_OTA_IMG_PENDING_VERIFY) {
			// run diagnostic function ...
			bool diagnostic_is_ok = diagnostic();
			if (diagnostic_is_ok) {
				ESP_LOGI(TAG, "Diagnostics completed successfully! Continuing execution ...");
				esp_ota_mark_app_valid_cancel_rollback();
			} else {
				ESP_LOGE(TAG, "Diagnostics failed! Start rollback to the previous version ...");
				esp_ota_mark_app_invalid_rollback_and_reboot();

			}
		}
	}

	/*set the ota file size */
	if(0 == fileSize)
		fileSize = OTA_WITH_SEQUENTIAL_WRITES;
	else
		ota_file_size = fileSize;

	/* Initialize OTA session*/
	ota_start_session();

	/* */
	return ESP_OK;
}

esp_err_t ota_get_data_to_buffer(const char* buf, int buf_len)
{
	/* This function should be called in comms.c inside the callback function ota_file_data_cb() */
	/* This function should copy the bytes that we get from qfread command via uart*/
	/* handle of qfread cmd stores the incoming bytes from modem via uart in buf */
	/* After that we gonna use ota-buffer in ota_process_incoming_image() function to start ota-session*/

	esp_err_t retval = ESP_FAIL;
	int index = 0;

	/*Check if the number of bytes to be processed is greater than ota-buffer size*/
	/*if(BUFFSIZE < buf_len)
	{
		ESP_LOGI(TAG, "Number of bytes to be processed is greater than the size of OTA buffer");
		return retval;
	}*/

	/* Set the status of ota-bytes to be processed ---> OTA_BYTES_WRITTEN_WAIT*/
	ota_bytes = OTA_BYTES_WRITTEN_WAIT;

	/* Reset the number of bytes that have been read to 0 (because the initial value is -1)*/
	data_read = 0;

	/*Copy the incoming bytes (from uart) to ota-buffer */
	//for(index = 0; index < buf_len; index++)
	{
		ota_write_data = (char*)buf;
		data_read = buf_len;
	}

	/*Start OTA_Writing session*/
	ota_set_needed();

	/*OTA process and writing for the copied bytes*/
	ota_process_incoming_image();

	/* Check the status of the bytes after OTA-processing/OTA-writing */
	if(OTA_BYTES_WRITTEN_SUCCESS == ota_bytes)
		retval = ESP_OK;
	else if(OTA_BYTES_WRITTEN_FAILED == ota_bytes)
		retval = ESP_FAIL;

	return retval;
}

void ota_finish_processing(void)
{
	/*This function should be invoked in the COMMS_STATE_WAIT_MCU_UPDATE state in comms state machine
	 *  (instead of the OTA_terminate_connection() )*/

	/* Reset data_read variable to 0*/
	data_read = 0;

	/* Rise ota_finish flag*/
	ota_set_finished();

	/* Rise the bytes_success flag*/
	ota_process_incoming_image();

	/* End the ota_session by validating the image and then rebooting */
	ota_end_session();
}