/*
 * SPDX-FileCopyrightText: 2015-2022 Espressif Systems (Shanghai) CO LTD
 *
 * SPDX-License-Identifier: Apache-2.0
 */
#include <string.h>
#include <sys/param.h>
#include <esp_cpu.h>
#include <bootloader_utility.h>
#include <bootloader_signature.h>
#include <esp_secure_boot.h>
#include <esp_fault.h>
#include <esp_log.h>
#include <esp_attr.h>
#include <spi_flash_mmap.h>
#include <bootloader_flash_priv.h>
#include <bootloader_random.h>
#include <bootloader_sha.h>
#include "bootloader_util.h"
#include "bootloader_common.h"
#include "esp_rom_sys.h"
#include "bootloader_memory_utils.h"
#include "soc/soc_caps.h"
#if CONFIG_IDF_TARGET_ESP32
#include "esp32/rom/secure_boot.h"
#elif CONFIG_IDF_TARGET_ESP32S2
#include "esp32s2/rom/secure_boot.h"
#elif CONFIG_IDF_TARGET_ESP32S3
#include "esp32s3/rom/secure_boot.h"
#elif CONFIG_IDF_TARGET_ESP32C3
#include "esp32c3/rom/secure_boot.h"
#elif CONFIG_IDF_TARGET_ESP32H2
#include "esp32h2/rom/secure_boot.h"
#elif CONFIG_IDF_TARGET_ESP32C2
#include "esp32c2/rom/rtc.h"
#include "esp32c2/rom/secure_boot.h"
#endif

/* Checking signatures as part of verifying images is necessary:
   - Always if secure boot is enabled
   - Differently in bootloader and/or app, depending on kconfig
*/
#ifdef BOOTLOADER_BUILD
#ifdef CONFIG_SECURE_SIGNED_ON_BOOT
#define SECURE_BOOT_CHECK_SIGNATURE 1
#else
#define SECURE_BOOT_CHECK_SIGNATURE 0
#endif
#else /* !BOOTLOADER_BUILD */
#ifdef CONFIG_SECURE_SIGNED_ON_UPDATE
#define SECURE_BOOT_CHECK_SIGNATURE 1
#else
#define SECURE_BOOT_CHECK_SIGNATURE 0
#endif
#endif

static const char *TAG = "esp_image";

#define HASH_LEN ESP_IMAGE_HASH_LEN

#define SIXTEEN_MB 0x1000000
#define ESP_ROM_CHECKSUM_INITIAL 0xEF

/* Headroom to ensure between stack SP (at time of checking) and data loaded from flash */
#define STACK_LOAD_HEADROOM 32768

#ifdef BOOTLOADER_BUILD
/* 64 bits of random data to obfuscate loaded RAM with, until verification is complete
   (Means loaded code isn't executable until after the secure boot check.)
*/
static uint32_t ram_obfs_value[2];

#endif

/* Return true if load_addr is an address the bootloader should load into */
static bool should_load(uint32_t load_addr);
/* Return true if load_addr is an address the bootloader should map via flash cache */
static bool should_map(uint32_t load_addr);

static esp_err_t process_segments(esp_image_metadata_t *data, bool silent, bool do_load, bootloader_sha256_handle_t sha_handle, uint32_t *checksum);
/* Load or verify a segment */
static esp_err_t process_segment(int index, uint32_t flash_addr, esp_image_segment_header_t *header, bool silent, bool do_load, bootloader_sha256_handle_t sha_handle, uint32_t *checksum);

/* split segment and verify if data_len is too long */
static esp_err_t process_segment_data(intptr_t load_addr, uint32_t data_addr, uint32_t data_len, bool do_load, bootloader_sha256_handle_t sha_handle, uint32_t *checksum);

/* Verify the main image header */
static esp_err_t verify_image_header(uint32_t src_addr, const esp_image_header_t *image, bool silent);

/* Verify a segment header */
static esp_err_t verify_segment_header(int index, const esp_image_segment_header_t *segment, uint32_t segment_data_offs, bool silent);

/* Log-and-fail macro for use in esp_image_load */
#define FAIL_LOAD(...) do {                         \
        if (!silent) {                              \
            ESP_LOGE(TAG, __VA_ARGS__);             \
        }                                           \
        goto err;                                   \
    }                                               \
    while(0)

#define CHECK_ERR(func) do {                        \
        if ((err = func) != ESP_OK) {               \
            goto err;                               \
        }                                           \
    }                                               \
    while(0)

static esp_err_t process_image_header(esp_image_metadata_t *data, uint32_t part_offset, bootloader_sha256_handle_t *sha_handle, bool do_verify, bool silent);
static esp_err_t process_appended_hash(esp_image_metadata_t *data, uint32_t part_len, bool do_verify, bool silent);
static esp_err_t process_checksum(bootloader_sha256_handle_t sha_handle, uint32_t checksum_word, esp_image_metadata_t *data, bool silent, bool skip_check_checksum);

static esp_err_t __attribute__((unused)) verify_secure_boot_signature(bootloader_sha256_handle_t sha_handle, esp_image_metadata_t *data, uint8_t *image_digest, uint8_t *verified_digest);
static esp_err_t __attribute__((unused)) verify_simple_hash(bootloader_sha256_handle_t sha_handle, esp_image_metadata_t *data);

static esp_err_t image_load(esp_image_load_mode_t mode, const esp_partition_pos_t *part, esp_image_metadata_t *data)
{
#ifdef BOOTLOADER_BUILD
    bool do_load   = (mode == ESP_IMAGE_LOAD) || (mode == ESP_IMAGE_LOAD_NO_VALIDATE);
    bool do_verify = (mode == ESP_IMAGE_LOAD) || (mode == ESP_IMAGE_VERIFY) || (mode == ESP_IMAGE_VERIFY_SILENT);
#else
    bool do_load   = false; // Can't load the image in app mode
    bool do_verify = true;  // In app mode is available only verify mode
#endif
    bool silent    = (mode == ESP_IMAGE_VERIFY_SILENT);
    esp_err_t err = ESP_OK;
    // checksum the image a word at a time. This shaves 30-40ms per MB of image size
    uint32_t checksum_word = ESP_ROM_CHECKSUM_INITIAL;
    uint32_t *checksum = (do_verify) ? &checksum_word : NULL;
    bootloader_sha256_handle_t sha_handle = NULL;
    bool verify_sha;
#if (SECURE_BOOT_CHECK_SIGNATURE == 1)
     /* used for anti-FI checks */
    uint8_t image_digest[HASH_LEN] = { [ 0 ... 31] = 0xEE };
    uint8_t verified_digest[HASH_LEN] = { [ 0 ... 31 ] = 0x01 };
#endif

    if (data == NULL || part == NULL) {
        return ESP_ERR_INVALID_ARG;
    }

#if CONFIG_SECURE_BOOT_V2_ENABLED
    // For Secure Boot V2, we do verify signature on bootloader which includes the SHA calculation.
    verify_sha = do_verify;
#else // Secure boot not enabled
    // For secure boot V1 on ESP32, we don't calculate SHA or verify signature on bootloaders.
    // (For non-secure boot, we don't verify any SHA-256 hash appended to the bootloader because
    // esptool.py may have rewritten the header - rely on esptool.py having verified the bootloader at flashing time, instead.)
    verify_sha = (part->offset != ESP_BOOTLOADER_OFFSET) && do_verify;
#endif

    if (part->size > SIXTEEN_MB) {
        err = ESP_ERR_INVALID_ARG;
        FAIL_LOAD("partition size 0x%x invalid, larger than 16MB", part->size);
    }

    bootloader_sha256_handle_t *p_sha_handle = &sha_handle;
    CHECK_ERR(process_image_header(data, part->offset, (verify_sha) ? p_sha_handle : NULL, do_verify, silent));
    CHECK_ERR(process_segments(data, silent, do_load, sha_handle, checksum));
    bool skip_check_checksum = !do_verify || esp_cpu_dbgr_is_attached();
    CHECK_ERR(process_checksum(sha_handle, checksum_word, data, silent, skip_check_checksum));
    CHECK_ERR(process_appended_hash(data, part->size, do_verify, silent));
    if (verify_sha) {
#if (SECURE_BOOT_CHECK_SIGNATURE == 1)
        // secure boot images have a signature appended
#if defined(BOOTLOADER_BUILD) && !defined(CONFIG_SECURE_BOOT)
        // If secure boot is not enabled in hardware, then
        // skip the signature check in bootloader when the debugger is attached.
        // This is done to allow for breakpoints in Flash.
        bool do_verify_sig = !esp_cpu_dbgr_is_attached();
#else // CONFIG_SECURE_BOOT
        bool do_verify_sig = true;
#endif // end checking for JTAG
        if (do_verify_sig) {
            err = verify_secure_boot_signature(sha_handle, data, image_digest, verified_digest);
            sha_handle = NULL; // verify_secure_boot_signature finishes sha_handle
        }
#else // SECURE_BOOT_CHECK_SIGNATURE
        // No secure boot, but SHA-256 can be appended for basic corruption detection
        if (sha_handle != NULL && !esp_cpu_dbgr_is_attached()) {
            err = verify_simple_hash(sha_handle, data);
            sha_handle = NULL; // calling verify_simple_hash finishes sha_handle
        }
#endif // SECURE_BOOT_CHECK_SIGNATURE
    } // verify_sha

    // bootloader may still have a sha256 digest handle open
    if (sha_handle != NULL) {
        bootloader_sha256_finish(sha_handle, NULL);
        sha_handle = NULL;
    }

    if (err != ESP_OK) {
        goto err;
    }

#ifdef BOOTLOADER_BUILD

#if (SECURE_BOOT_CHECK_SIGNATURE == 1)
    /* If signature was checked in bootloader build, verified_digest should equal image_digest

       This is to detect any fault injection that caused signature verification to not complete normally.

       Any attack which bypasses this check should be of limited use as the RAM contents are still obfuscated, therefore we do the check
       immediately before we deobfuscate.

       Note: the conditions for making this check are the same as for setting verify_sha above, but on ESP32 SB V1 we move the test for
       "only verify signature in bootloader" into the macro so it's tested multiple times.
     */
#if CONFIG_SECURE_BOOT_V2_ENABLED
    ESP_FAULT_ASSERT(!esp_secure_boot_enabled() || memcmp(image_digest, verified_digest, HASH_LEN) == 0);
#else // Secure Boot V1 on ESP32, only verify signatures for apps not bootloaders
    ESP_FAULT_ASSERT(data->start_addr == ESP_BOOTLOADER_OFFSET || memcmp(image_digest, verified_digest, HASH_LEN) == 0);
#endif

#endif // SECURE_BOOT_CHECK_SIGNATURE

    // Deobfuscate RAM
    if (do_load && ram_obfs_value[0] != 0 && ram_obfs_value[1] != 0) {
        for (int i = 0; i < data->image.segment_count; i++) {
            uint32_t load_addr = data->segments[i].load_addr;
            if (should_load(load_addr)) {
                uint32_t *loaded = (uint32_t *)load_addr;
                for (size_t j = 0; j < data->segments[i].data_len / sizeof(uint32_t); j++) {
                    loaded[j] ^= (j & 1) ? ram_obfs_value[0] : ram_obfs_value[1];
                }
            }
        }
    }
#endif // BOOTLOADER_BUILD

    // Success!
    return ESP_OK;

err:
    if (err == ESP_OK) {
        err = ESP_ERR_IMAGE_INVALID;
    }
    if (sha_handle != NULL) {
        // Need to finish the hash process to free the handle
        bootloader_sha256_finish(sha_handle, NULL);
    }
    // Prevent invalid/incomplete data leaking out
    bzero(data, sizeof(esp_image_metadata_t));
    return err;
}

esp_err_t bootloader_load_image(const esp_partition_pos_t *part, esp_image_metadata_t *data)
{
#if !defined(BOOTLOADER_BUILD)
    return ESP_FAIL;
#else
    esp_image_load_mode_t mode = ESP_IMAGE_LOAD;

#if !defined(CONFIG_SECURE_BOOT)
    /* Skip validation under particular configurations */
#if CONFIG_BOOTLOADER_SKIP_VALIDATE_ALWAYS
    mode = ESP_IMAGE_LOAD_NO_VALIDATE;
#elif CONFIG_BOOTLOADER_SKIP_VALIDATE_ON_POWER_ON
    if (esp_rom_get_reset_reason(0) == RESET_REASON_CHIP_POWER_ON
#if SOC_EFUSE_HAS_EFUSE_RST_BUG
        || esp_rom_get_reset_reason(0) == RESET_REASON_CORE_EFUSE_CRC
#endif
        ) {
        mode = ESP_IMAGE_LOAD_NO_VALIDATE;
    }
#endif // CONFIG_BOOTLOADER_SKIP_...
#endif // CONFIG_SECURE_BOOT

 return image_load(mode, part, data);
#endif // BOOTLOADER_BUILD
}

esp_err_t bootloader_load_image_no_verify(const esp_partition_pos_t *part, esp_image_metadata_t *data)
{
#ifdef BOOTLOADER_BUILD
    return image_load(ESP_IMAGE_LOAD_NO_VALIDATE, part, data);
#else
    return ESP_FAIL;
#endif
}

esp_err_t esp_image_verify(esp_image_load_mode_t mode, const esp_partition_pos_t *part, esp_image_metadata_t *data)
{
    return image_load(mode, part, data);
}

esp_err_t esp_image_get_metadata(const esp_partition_pos_t *part, esp_image_metadata_t *metadata)
{
    esp_err_t err;
    if (metadata == NULL || part == NULL || part->size > SIXTEEN_MB) {
        return ESP_ERR_INVALID_ARG;
    }

    bool silent = true;
    bool do_verify = false;
    bool do_load = false;
    CHECK_ERR(process_image_header(metadata, part->offset, NULL, do_verify, silent));
    CHECK_ERR(process_segments(metadata, silent, do_load, NULL, NULL));
    bool skip_check_checksum = true;
    CHECK_ERR(process_checksum(NULL, 0, metadata, silent, skip_check_checksum));
    CHECK_ERR(process_appended_hash(metadata, part->size, true, silent));
    return ESP_OK;
err:
    return err;
}

static esp_err_t verify_image_header(uint32_t src_addr, const esp_image_header_t *image, bool silent)
{
    esp_err_t err = ESP_OK;

    ESP_LOGD(TAG, "image header: 0x%02x 0x%02x 0x%02x 0x%02x %08x",
                image->magic,
                image->segment_count,
                image->spi_mode,
                image->spi_size,
                image->entry_addr);

    if (image->magic != ESP_IMAGE_HEADER_MAGIC) {
        FAIL_LOAD("image at 0x%x has invalid magic byte (nothing flashed here?)", src_addr);
    }

    // Checking the chip revision header *will* print a bunch of other info
    // regardless of silent setting as this may be important, but don't bother checking it
    // if it looks like the app partition is erased or otherwise garbage
    CHECK_ERR(bootloader_common_check_chip_validity(image, ESP_IMAGE_APPLICATION));

    if (image->segment_count > ESP_IMAGE_MAX_SEGMENTS) {
        FAIL_LOAD("image at 0x%x segment count %d exceeds max %d", src_addr, image->segment_count, ESP_IMAGE_MAX_SEGMENTS);
    }
    return err;
err:
    if (err == ESP_OK) {
        err = ESP_ERR_IMAGE_INVALID;
    }
    return err;
}

#ifdef BOOTLOADER_BUILD
/* Check the region load_addr - load_end doesn't overlap any memory used by the bootloader, registers, or other invalid memory
 */
static bool verify_load_addresses(int segment_index, intptr_t load_addr, intptr_t load_end, bool print_error, bool no_recurse)
{
    /* Addresses of static data and the "loader" section of bootloader IRAM, all defined in ld script */
    const char *reason = NULL;
    extern int _dram_start, _dram_end, _loader_text_start, _loader_text_end;
    void *load_addr_p = (void *)load_addr;
    void *load_inclusive_end_p = (void *)load_end - 0x1;
    void *load_exclusive_end_p = (void *)load_end;

    if (load_end == load_addr) {
        return true; // zero-length segments are fine
    }
    assert(load_end > load_addr); // data_len<16MB is checked in verify_segment_header() which is called before this, so this should always be true

    if (esp_ptr_in_dram(load_addr_p) && esp_ptr_in_dram(load_inclusive_end_p)) { /* Writing to DRAM */
        /* Check if we're clobbering the stack */
        intptr_t sp = (intptr_t)esp_cpu_get_sp();
        if (bootloader_util_regions_overlap(sp - STACK_LOAD_HEADROOM, SOC_ROM_STACK_START,
                                           load_addr, load_end)) {
            reason = "overlaps bootloader stack";
            goto invalid;
        }

        /* Check if we're clobbering static data

           (_dram_start.._dram_end includes bss, data, rodata sections in DRAM)
         */
        if (bootloader_util_regions_overlap((intptr_t)&_dram_start, (intptr_t)&_dram_end, load_addr, load_end)) {
            reason = "overlaps bootloader data";
            goto invalid;
        }

        /* LAST DRAM CHECK (recursive): for D/IRAM, check the equivalent IRAM addresses if needed

           Allow for the possibility that even though both pointers are IRAM, only part of the region is in a D/IRAM
           section. In which case we recurse to check the part which falls in D/IRAM.

           Note: We start with SOC_DIRAM_DRAM_LOW/HIGH and convert that address to IRAM to account for any reversing of word order
           (chip-specific).
         */
        if (!no_recurse && bootloader_util_regions_overlap(SOC_DIRAM_DRAM_LOW, SOC_DIRAM_DRAM_HIGH, load_addr, load_end)) {
            intptr_t iram_load_addr, iram_load_end;

            if (esp_ptr_in_diram_dram(load_addr_p)) {
                iram_load_addr = (intptr_t)esp_ptr_diram_dram_to_iram(load_addr_p);
            } else {
                iram_load_addr = (intptr_t)esp_ptr_diram_dram_to_iram((void *)SOC_DIRAM_DRAM_LOW);
            }

            if (esp_ptr_in_diram_dram(load_inclusive_end_p)) {
                iram_load_end = (intptr_t)esp_ptr_diram_dram_to_iram(load_exclusive_end_p);
            } else {
                iram_load_end = (intptr_t)esp_ptr_diram_dram_to_iram((void *)SOC_DIRAM_DRAM_HIGH);
            }

            if (iram_load_end < iram_load_addr) {
                return verify_load_addresses(segment_index, iram_load_end, iram_load_addr, print_error, true);
            } else {
                return verify_load_addresses(segment_index, iram_load_addr, iram_load_end, print_error, true);
            }
        }
    }
    else if (esp_ptr_in_iram(load_addr_p) && esp_ptr_in_iram(load_inclusive_end_p)) { /* Writing to IRAM */
        /* Check for overlap of 'loader' section of IRAM */
        if (bootloader_util_regions_overlap((intptr_t)&_loader_text_start, (intptr_t)&_loader_text_end,
                                            load_addr, load_end)) {
            reason = "overlaps loader IRAM";
            goto invalid;
        }

        /* LAST IRAM CHECK (recursive): for D/IRAM, check the equivalent DRAM address if needed

           Allow for the possibility that even though both pointers are IRAM, only part of the region is in a D/IRAM
           section. In which case we recurse to check the part which falls in D/IRAM.
           Note: We start with SOC_DIRAM_IRAM_LOW/HIGH and convert that address to DRAM to account for any reversing of word order
           (chip-specific).
         */
        if (!no_recurse && bootloader_util_regions_overlap(SOC_DIRAM_IRAM_LOW, SOC_DIRAM_IRAM_HIGH, load_addr, load_end)) {
            intptr_t dram_load_addr, dram_load_end;

            if (esp_ptr_in_diram_iram(load_addr_p)) {
                dram_load_addr = (intptr_t)esp_ptr_diram_iram_to_dram(load_addr_p);
            } else {
                dram_load_addr = (intptr_t)esp_ptr_diram_iram_to_dram((void *)SOC_DIRAM_IRAM_LOW);
            }

            if (esp_ptr_in_diram_iram(load_inclusive_end_p)) {
                dram_load_end = (intptr_t)esp_ptr_diram_iram_to_dram(load_exclusive_end_p);
            } else {
                dram_load_end = (intptr_t)esp_ptr_diram_iram_to_dram((void *)SOC_DIRAM_IRAM_HIGH);
            }

            if (dram_load_end < dram_load_addr) {
                return verify_load_addresses(segment_index, dram_load_end, dram_load_addr, print_error, true);
            } else {
                return verify_load_addresses(segment_index, dram_load_addr, dram_load_end, print_error, true);
            }
        }
    /* Sections entirely in RTC memory won't overlap with a vanilla bootloader but are valid load addresses, thus skipping them from the check */
    }
#if SOC_RTC_FAST_MEM_SUPPORTED
    else if (esp_ptr_in_rtc_iram_fast(load_addr_p) && esp_ptr_in_rtc_iram_fast(load_inclusive_end_p)){
        return true;
    } else if (esp_ptr_in_rtc_dram_fast(load_addr_p) && esp_ptr_in_rtc_dram_fast(load_inclusive_end_p)){
        return true;
    }
#endif

#if SOC_RTC_SLOW_MEM_SUPPORTED
    else if (esp_ptr_in_rtc_slow(load_addr_p) && esp_ptr_in_rtc_slow(load_inclusive_end_p)) {
        return true;
    }
#endif

    else { /* Not a DRAM or an IRAM or RTC Fast IRAM, RTC Fast DRAM or RTC Slow address */
        reason = "bad load address range";
        goto invalid;
    }
    return true;

 invalid:
    if (print_error) {
        ESP_LOGE(TAG, "Segment %d 0x%08x-0x%08x invalid: %s", segment_index, load_addr, load_end, reason);
    }
    return false;
}
#endif // BOOTLOADER_BUILD

static esp_err_t process_image_header(esp_image_metadata_t *data, uint32_t part_offset, bootloader_sha256_handle_t *sha_handle, bool do_verify, bool silent)
{
    esp_err_t err;
    bzero(data, sizeof(esp_image_metadata_t));
    data->start_addr = part_offset;

    ESP_LOGD(TAG, "reading image header @ 0x%x", data->start_addr);
    CHECK_ERR(bootloader_flash_read(data->start_addr, &data->image, sizeof(esp_image_header_t), true));

    if (do_verify) {
        // Calculate SHA-256 of image if secure boot is on, or if image has a hash appended
        if (SECURE_BOOT_CHECK_SIGNATURE || data->image.hash_appended) {
            if (sha_handle != NULL) {
                *sha_handle = bootloader_sha256_start();
                if (*sha_handle == NULL) {
                    return ESP_ERR_NO_MEM;
                }
                bootloader_sha256_data(*sha_handle, &data->image, sizeof(esp_image_header_t));
            }
        }
        CHECK_ERR(verify_image_header(data->start_addr, &data->image, silent));
    }
    data->image_len = sizeof(esp_image_header_t);
    return ESP_OK;
err:
    return err;
}

static esp_err_t process_segments(esp_image_metadata_t *data, bool silent, bool do_load, bootloader_sha256_handle_t sha_handle, uint32_t *checksum)
{
    esp_err_t err = ESP_OK;
    uint32_t start_segments = data->start_addr + data->image_len;
    uint32_t next_addr = start_segments;
    for (int i = 0; i < data->image.segment_count; i++) {
        esp_image_segment_header_t *header = &data->segments[i];
        ESP_LOGV(TAG, "loading segment header %d at offset 0x%x", i, next_addr);
        CHECK_ERR(process_segment(i, next_addr, header, silent, do_load, sha_handle, checksum));
        next_addr += sizeof(esp_image_segment_header_t);
        data->segment_data[i] = next_addr;
        next_addr += header->data_len;
    }
    // Segments all loaded, verify length
    uint32_t end_addr = next_addr;
    if (end_addr < data->start_addr) {
        FAIL_LOAD("image offset has wrapped");
    }

    data->image_len += end_addr - start_segments;
    ESP_LOGV(TAG, "image start 0x%08x end of last section 0x%08x", data->start_addr, end_addr);
    return err;
err:
    if (err == ESP_OK) {
        err = ESP_ERR_IMAGE_INVALID;
    }
    return err;
}

static esp_err_t process_segment(int index, uint32_t flash_addr, esp_image_segment_header_t *header, bool silent, bool do_load, bootloader_sha256_handle_t sha_handle, uint32_t *checksum)
{
    esp_err_t err;

    /* read segment header */
    err = bootloader_flash_read(flash_addr, header, sizeof(esp_image_segment_header_t), true);
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "bootloader_flash_read failed at 0x%08x", flash_addr);
        return err;
    }
    if (sha_handle != NULL) {
        bootloader_sha256_data(sha_handle, header, sizeof(esp_image_segment_header_t));
    }

    intptr_t load_addr = header->load_addr;
    uint32_t data_len = header->data_len;
    uint32_t data_addr = flash_addr + sizeof(esp_image_segment_header_t);

    ESP_LOGV(TAG, "segment data length 0x%x data starts 0x%x", data_len, data_addr);

    CHECK_ERR(verify_segment_header(index, header, data_addr, silent));

    if (data_len % 4 != 0) {
        FAIL_LOAD("unaligned segment length 0x%x", data_len);
    }

    bool is_mapping = should_map(load_addr);
    do_load = do_load && should_load(load_addr);

    if (!silent) {
        ESP_LOGI(TAG, "segment %d: paddr=%08x vaddr=%08x size=%05xh (%6d) %s",
                 index, data_addr, load_addr,
                 data_len, data_len,
                 (do_load) ? "load" : (is_mapping) ? "map" : "");
    }


#ifdef BOOTLOADER_BUILD
    /* Before loading segment, check it doesn't clobber bootloader RAM. */
    if (do_load && data_len > 0) {
        if (!verify_load_addresses(index, load_addr, load_addr + data_len, true, false)) {
            return ESP_ERR_IMAGE_INVALID;
        }
    }
#endif // BOOTLOADER_BUILD

    uint32_t free_page_count = bootloader_mmap_get_free_pages();
    ESP_LOGD(TAG, "free data page_count 0x%08x", free_page_count);

    uint32_t data_len_remain = data_len;
    while (data_len_remain > 0) {
#if (SECURE_BOOT_CHECK_SIGNATURE == 1) && defined(BOOTLOADER_BUILD)
        /* Double check the address verification done above */
        ESP_FAULT_ASSERT(!do_load || verify_load_addresses(0, load_addr, load_addr + data_len_remain, false, false));
#endif
        uint32_t offset_page = ((data_addr & MMAP_ALIGNED_MASK) != 0) ? 1 : 0;
        /* Data we could map in case we are not aligned to PAGE boundary is one page size lesser. */
        data_len = MIN(data_len_remain, ((free_page_count - offset_page) * SPI_FLASH_MMU_PAGE_SIZE));
        CHECK_ERR(process_segment_data(load_addr, data_addr, data_len, do_load, sha_handle, checksum));
        data_addr += data_len;
        data_len_remain -= data_len;
    }

    return ESP_OK;

err:
    if (err == ESP_OK) {
        err = ESP_ERR_IMAGE_INVALID;
    }

    return err;
}

static esp_err_t process_segment_data(intptr_t load_addr, uint32_t data_addr, uint32_t data_len, bool do_load, bootloader_sha256_handle_t sha_handle, uint32_t *checksum)
{
    // If we are not loading, and the checksum is empty, skip processing this
    // segment for data
    if (!do_load && checksum == NULL) {
        ESP_LOGD(TAG, "skipping checksum for segment");
        return ESP_OK;
    }

    const uint32_t *data = (const uint32_t *)bootloader_mmap(data_addr, data_len);
    if (!data) {
        ESP_LOGE(TAG, "bootloader_mmap(0x%x, 0x%x) failed",
                 data_addr, data_len);
        return ESP_FAIL;
    }

    if (checksum == NULL && sha_handle == NULL) {
        memcpy((void *)load_addr, data, data_len);
        bootloader_munmap(data);
        return ESP_OK;
    }

#ifdef BOOTLOADER_BUILD
    // Set up the obfuscation value to use for loading
    while (ram_obfs_value[0] == 0 || ram_obfs_value[1] == 0) {
        bootloader_fill_random(ram_obfs_value, sizeof(ram_obfs_value));
#if CONFIG_IDF_ENV_FPGA
        /* FPGA doesn't always emulate the RNG */
        ram_obfs_value[0] ^= 0x33;
        ram_obfs_value[1] ^= 0x66;
#endif
    }
    uint32_t *dest = (uint32_t *)load_addr;
#endif

    const uint32_t *src = data;

    for (size_t i = 0; i < data_len; i += 4) {
        int w_i = i / 4; // Word index
        uint32_t w = src[w_i];
        if (checksum != NULL) {
            *checksum ^= w;
        }
#ifdef BOOTLOADER_BUILD
        if (do_load) {
            dest[w_i] = w ^ ((w_i & 1) ? ram_obfs_value[0] : ram_obfs_value[1]);
        }
#endif
        // SHA_CHUNK determined experimentally as the optimum size
        // to call bootloader_sha256_data() with. This is a bit
        // counter-intuitive, but it's ~3ms better than using the
        // SHA256 block size.
        const size_t SHA_CHUNK = 1024;
        if (sha_handle != NULL && i % SHA_CHUNK == 0) {
            bootloader_sha256_data(sha_handle, &src[w_i],
                                   MIN(SHA_CHUNK, data_len - i));
        }
    }

    bootloader_munmap(data);

    return ESP_OK;
}

static esp_err_t verify_segment_header(int index, const esp_image_segment_header_t *segment, uint32_t segment_data_offs, bool silent)
{
    if ((segment->data_len & 3) != 0
            || segment->data_len >= SIXTEEN_MB) {
        if (!silent) {
            ESP_LOGE(TAG, "invalid segment length 0x%x", segment->data_len);
        }
        return ESP_ERR_IMAGE_INVALID;
    }

    uint32_t load_addr = segment->load_addr;
    bool map_segment = should_map(load_addr);

    /* Check that flash cache mapped segment aligns correctly from flash to its mapped address,
       relative to the 64KB page mapping size.
    */
    ESP_LOGV(TAG, "segment %d map_segment %d segment_data_offs 0x%x load_addr 0x%x",
             index, map_segment, segment_data_offs, load_addr);
    if (map_segment
            && ((segment_data_offs % SPI_FLASH_MMU_PAGE_SIZE) != (load_addr % SPI_FLASH_MMU_PAGE_SIZE))) {
        if (!silent) {
            ESP_LOGE(TAG, "Segment %d load address 0x%08x, doesn't match data 0x%08x",
                     index, load_addr, segment_data_offs);
        }
        return ESP_ERR_IMAGE_INVALID;
    }

    return ESP_OK;
}

static bool should_map(uint32_t load_addr)
{
    return (load_addr >= SOC_IROM_LOW && load_addr < SOC_IROM_HIGH)
           || (load_addr >= SOC_DROM_LOW && load_addr < SOC_DROM_HIGH);
}

static bool should_load(uint32_t load_addr)
{
    /* Reload the RTC memory segments whenever a non-deepsleep reset
       is occurring */
    bool load_rtc_memory = esp_rom_get_reset_reason(0) != RESET_REASON_CORE_DEEP_SLEEP;

    if (should_map(load_addr)) {
        return false;
    }

    if (load_addr < 0x10000000) {
        // Reserved for non-loaded addresses.
        // Current reserved values are
        // 0x0 (padding block)
        // 0x4 (unused, but reserved for an MD5 block)
        return false;
    }

    if (!load_rtc_memory) {
#if SOC_RTC_FAST_MEM_SUPPORTED
        if (load_addr >= SOC_RTC_IRAM_LOW && load_addr < SOC_RTC_IRAM_HIGH) {
            ESP_LOGD(TAG, "Skipping RTC fast memory segment at 0x%08x", load_addr);
            return false;
        }
        if (load_addr >= SOC_RTC_DRAM_LOW && load_addr < SOC_RTC_DRAM_HIGH) {
            ESP_LOGD(TAG, "Skipping RTC fast memory segment at 0x%08x", load_addr);
            return false;
        }
#endif

#if SOC_RTC_SLOW_MEM_SUPPORTED
        if (load_addr >= SOC_RTC_DATA_LOW && load_addr < SOC_RTC_DATA_HIGH) {
            ESP_LOGD(TAG, "Skipping RTC slow memory segment at 0x%08x", load_addr);
            return false;
        }
#endif
    }

    return true;
}

esp_err_t esp_image_verify_bootloader(uint32_t *length)
{
    esp_image_metadata_t data;
    esp_err_t err = esp_image_verify_bootloader_data(&data);
    if (length != NULL) {
        *length = (err == ESP_OK) ? data.image_len : 0;
    }
    return err;
}

esp_err_t esp_image_verify_bootloader_data(esp_image_metadata_t *data)
{
    if (data == NULL) {
        return ESP_ERR_INVALID_ARG;
    }
    const esp_partition_pos_t bootloader_part = {
        .offset = ESP_BOOTLOADER_OFFSET,
        .size = ESP_PARTITION_TABLE_OFFSET - ESP_BOOTLOADER_OFFSET,
    };
    return esp_image_verify(ESP_IMAGE_VERIFY,
                            &bootloader_part,
                            data);
}

static esp_err_t process_appended_hash(esp_image_metadata_t *data, uint32_t part_len, bool do_verify, bool silent)
{
    esp_err_t err = ESP_OK;
    if (data->image.hash_appended) {
        // Account for the hash in the total image length
        if (do_verify) {
            CHECK_ERR(bootloader_flash_read(data->start_addr + data->image_len, &data->image_digest, HASH_LEN, true));
        }
        data->image_len += HASH_LEN;
    }

    if (data->image_len > part_len) {
        FAIL_LOAD("Image length %d doesn't fit in partition length %d", data->image_len, part_len);
    }
    return err;
err:
    if (err == ESP_OK) {
        err = ESP_ERR_IMAGE_INVALID;
    }

    return err;
}

static esp_err_t process_checksum(bootloader_sha256_handle_t sha_handle, uint32_t checksum_word, esp_image_metadata_t *data, bool silent, bool skip_check_checksum)
{
    esp_err_t err = ESP_OK;
    uint32_t unpadded_length = data->image_len;
    uint32_t length = unpadded_length + 1; // Add a byte for the checksum
    length = (length + 15) & ~15; // Pad to next full 16 byte block
    length = length - unpadded_length;

    // Verify checksum
    WORD_ALIGNED_ATTR uint8_t buf[16] = {0};
    if (!skip_check_checksum || sha_handle != NULL) {
        CHECK_ERR(bootloader_flash_read(data->start_addr + unpadded_length, buf, length, true));
    }
    uint8_t read_checksum = buf[length - 1];
    uint8_t calc_checksum = (checksum_word >> 24) ^ (checksum_word >> 16) ^ (checksum_word >> 8) ^ (checksum_word >> 0);
    if (!skip_check_checksum && calc_checksum != read_checksum) {
        FAIL_LOAD("Checksum failed. Calculated 0x%x read 0x%x", calc_checksum, read_checksum);
    }
    if (sha_handle != NULL) {
        bootloader_sha256_data(sha_handle, buf, length);
    }
    data->image_len += length;

    return err;
err:
    if (err == ESP_OK) {
        err = ESP_ERR_IMAGE_INVALID;
    }

    return err;
}

static esp_err_t verify_secure_boot_signature(bootloader_sha256_handle_t sha_handle, esp_image_metadata_t *data, uint8_t *image_digest, uint8_t *verified_digest)
{
#if (SECURE_BOOT_CHECK_SIGNATURE == 1)
    uint32_t end = data->start_addr + data->image_len;

    ESP_LOGI(TAG, "Verifying image signature...");

    // For secure boot, we calculate the signature hash over the whole file, which includes any "simple" hash
    // appended to the image for corruption detection
    if (data->image.hash_appended) {
        const void *simple_hash = bootloader_mmap(end - HASH_LEN, HASH_LEN);
        bootloader_sha256_data(sha_handle, simple_hash, HASH_LEN);
        bootloader_munmap(simple_hash);
    }

#if CONFIG_SECURE_BOOT_V2_ENABLED
    // End of the image needs to be padded all the way to a 4KB boundary, after the simple hash
    // (for apps they are usually already padded due to --secure-pad-v2, only a problem if this option was not used.)
    uint32_t padded_end = (end + FLASH_SECTOR_SIZE - 1) & ~(FLASH_SECTOR_SIZE-1);
    if (padded_end > end) {
        const void *padding = bootloader_mmap(end, padded_end - end);
        bootloader_sha256_data(sha_handle, padding, padded_end - end);
        bootloader_munmap(padding);
        end = padded_end;
    }
#endif

    bootloader_sha256_finish(sha_handle, image_digest);

    // Log the hash for debugging
    bootloader_debug_buffer(image_digest, HASH_LEN, "Calculated secure boot hash");

    // Use hash to verify signature block
    esp_err_t err = ESP_ERR_IMAGE_INVALID;
#if CONFIG_SECURE_BOOT || CONFIG_SECURE_SIGNED_APPS_NO_SECURE_BOOT
    const void *sig_block;
    ESP_FAULT_ASSERT(memcmp(image_digest, verified_digest, HASH_LEN) != 0); /* sanity check that these values start differently */
#if defined(CONFIG_SECURE_SIGNED_APPS_ECDSA_SCHEME)
    sig_block = bootloader_mmap(data->start_addr + data->image_len, sizeof(esp_secure_boot_sig_block_t));
    err = esp_secure_boot_verify_ecdsa_signature_block(sig_block, image_digest, verified_digest);
#else
    sig_block = bootloader_mmap(end, sizeof(ets_secure_boot_signature_t));
    err = esp_secure_boot_verify_sbv2_signature_block(sig_block, image_digest, verified_digest);
#endif
    bootloader_munmap(sig_block);
#endif // CONFIG_SECURE_BOOT || CONFIG_SECURE_SIGNED_APPS_NO_SECURE_BOOT
    if (err != ESP_OK) {
        ESP_LOGE(TAG, "Secure boot signature verification failed");

        // Go back and check if the simple hash matches or not (we're off the fast path so we can re-hash the whole image now)
        ESP_LOGI(TAG, "Calculating simple hash to check for corruption...");
        const void *whole_image = bootloader_mmap(data->start_addr, data->image_len - HASH_LEN);
        if (whole_image != NULL) {
            sha_handle = bootloader_sha256_start();
            bootloader_sha256_data(sha_handle, whole_image, data->image_len - HASH_LEN);
            bootloader_munmap(whole_image);
            if (verify_simple_hash(sha_handle, data) != ESP_OK) {
                ESP_LOGW(TAG, "image corrupted on flash");
            } else {
                ESP_LOGW(TAG, "image valid, signature bad");
            }
        }
        return ESP_ERR_IMAGE_INVALID;
    }

#if CONFIG_SECURE_SIGNED_APPS_RSA_SCHEME || CONFIG_SECURE_SIGNED_APPS_ECDSA_V2_SCHEME
    // Adjust image length result to include the appended signature
    data->image_len = end - data->start_addr + sizeof(ets_secure_boot_signature_t);
#endif

#endif // SECURE_BOOT_CHECK_SIGNATURE
    return ESP_OK;
}

static esp_err_t verify_simple_hash(bootloader_sha256_handle_t sha_handle, esp_image_metadata_t *data)
{
    uint8_t image_hash[HASH_LEN] = { 0 };
    bootloader_sha256_finish(sha_handle, image_hash);

    // Log the hash for debugging
    bootloader_debug_buffer(image_hash, HASH_LEN, "Calculated hash");

    // Simple hash for verification only
    if (memcmp(data->image_digest, image_hash, HASH_LEN) != 0) {
        ESP_LOGE(TAG, "Image hash failed - image is corrupt");
        bootloader_debug_buffer(data->image_digest, HASH_LEN, "Expected hash");
#ifdef CONFIG_IDF_ENV_FPGA
        ESP_LOGW(TAG, "Ignoring invalid SHA-256 as running on FPGA");
        return ESP_OK;
#endif
        return ESP_ERR_IMAGE_INVALID;
    }

    return ESP_OK;
}

int esp_image_get_flash_size(esp_image_flash_size_t app_flash_size)
{
    switch (app_flash_size) {
    case ESP_IMAGE_FLASH_SIZE_1MB:
        return 1 * 1024 * 1024;
    case ESP_IMAGE_FLASH_SIZE_2MB:
        return 2 * 1024 * 1024;
    case ESP_IMAGE_FLASH_SIZE_4MB:
        return 4 * 1024 * 1024;
    case ESP_IMAGE_FLASH_SIZE_8MB:
        return 8 * 1024 * 1024;
    case ESP_IMAGE_FLASH_SIZE_16MB:
        return 16 * 1024 * 1024;
    case ESP_IMAGE_FLASH_SIZE_32MB:
        return 32 * 1024 * 1024;
    case ESP_IMAGE_FLASH_SIZE_64MB:
        return 64 * 1024 * 1024;
    case ESP_IMAGE_FLASH_SIZE_128MB:
        return 128 * 1024 * 1024;
    default:
        return 0;
    }
}
