esp32 no_std code to write ota partitions managed by esp-idf bootloader

updater.rs

use defmt::*;
use bytemuck::{try_from_bytes, Pod, Zeroable};
use defmt::Debug2Format;
use embedded_storage::{ReadStorage, Storage};
use esp_storage::FlashStorage;
use hal::rom::crc::crc32_le;

pub struct Partition {
    pub offset: u32,
    pub size: usize,
}

impl Partition {
    pub const fn get_range(&self) -> core::ops::Range<u32> {
        self.offset..(self.offset + self.size as u32)
    }
}

// structs and impl part-wise copied from esp-idf
// e.g. https://github.com/espressif/esp-idf/blob/be06a6f5ffe36f9554cfc91fe2036e0fc85fea60/components/bootloader_support/include/esp_flash_partitions.h#L61-L66
#[repr(C)]
#[derive(Clone, Format)]
struct OtaSelectEntry {
    ota_seq: u32,
    _seq_label: [u8; 20],
    ota_state: OtaImgState,
    crc: u32, // Crc only of ota_seq
}

#[repr(u32)]
#[derive(PartialEq, Copy, Clone, Format)]
#[allow(dead_code)]
enum OtaImgState {
    ImgNew = 0,
    ImgPendingVerify = 1,
    ImgValid = 2,
    ImgInvalid = 3,
    ImgAborted = 4,
    ImgUndefined = 0xFFFFFFFF,
}

impl OtaSelectEntry {
    fn is_invalid(&self) -> bool {
        self.ota_seq == u32::MAX || self.ota_state == OtaImgState::ImgInvalid || self.ota_state == OtaImgState::ImgAborted
    }

    fn is_valid(&self) -> bool {
        let bytes: [u8; 4] = unsafe { core::mem::transmute(self.ota_seq) };
        !self.is_invalid() && self.crc == crc32_le(u32::MAX, bytes.as_slice())
    }
}

struct State {
    partition: u8,     // OTA_0 ..... OTA_(N-1)
    base_address: u32, // Start of partition OTA_i
    size: u32,         // Size of the entire update
    written: u32,      // Written bytes
    actual_crc: u32,
    expected_crc: u32,
}
      
      
/// Firmware header used by the ESP-IDF bootloader.
///
/// ## Header documentation:
/// * [Header](https://docs.espressif.com/projects/esptool/en/latest/esp32c3/advanced-topics/firmware-image-format.html#file-header)
/// * [Extended header](https://docs.espressif.com/projects/esptool/en/latest/esp32c3/advanced-topics/firmware-image-format.html#extended-file-header)
#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C, packed)]
#[doc(alias = "esp_image_header_t")]
struct ImageHeader {
    magic: u8,
    segment_count: u8,
    /// Flash read mode (esp_image_spi_mode_t)
    flash_mode: u8,
    /// ..4 bits are flash chip size (esp_image_flash_size_t)
    /// 4.. bits are flash frequency (esp_image_spi_freq_t)
    #[doc(alias = "spi_size")]
    #[doc(alias = "spi_speed")]
    flash_config: u8,
    entry: u32,

    // extended header part
    wp_pin: u8,
    clk_q_drv: u8,
    d_cs_drv: u8,
    gd_wp_drv: u8,
    chip_id: u16,
    min_rev: u8,
    /// Minimum chip revision supported by image, in format: major * 100 + minor
    min_chip_rev_full: u16,
    /// Maximal chip revision supported by image, in format: major * 100 + minor
    max_chip_rev_full: u16,
    reserved: [u8; 4],
    append_digest: u8,
}

#[derive(Debug, Clone, Copy, Pod, Zeroable)]
#[repr(C, packed)]
struct SegmentHeader {
    addr: u32,
    length: u32,
}
      
pub enum UpdaterError {
    CrcMissmatch,
    FlashWriteError,
    OtherError(&'static str),
}
      

pub struct UpdaterEsp32<const N: usize> {
    flash: FlashStorage,
    ota_data_partition_start: u32,
    ota_select_entries: [OtaSelectEntry; 2],
    ota_partitions: [Partition; N],
    update_state: Option<State>,
}

impl<const N: usize> UpdaterEsp32<N> {
    pub fn new(ota_data_partition: Partition, ota_partitions: [Partition; N]) -> Result<Self, ()> {
        assert_eq!(core::mem::size_of::<OtaSelectEntry>(), 32);
        assert_eq!(ota_data_partition.size, (2 * FlashStorage::SECTOR_SIZE) as usize);
        assert_eq!(ota_data_partition.offset % FlashStorage::SECTOR_SIZE, 0);
        assert!(ota_partitions.iter().all(|p| p.offset % FlashStorage::SECTOR_SIZE == 0));

        let mut flash = FlashStorage::new();
        let mut ota_select_0 = [0; 32];
        flash
            .read(ota_data_partition.offset, &mut ota_select_0)
            .map_err(|e| error!("Failed to read OTA select entry 0: {:?}", Debug2Format(&e)))?;
        let ota_select_0 = unsafe { &*(ota_select_0.as_ptr() as *const OtaSelectEntry) };
        let mut ota_select_1 = [0; 32];
        flash
            .read(ota_data_partition.offset + FlashStorage::SECTOR_SIZE, &mut ota_select_1)
            .map_err(|e| error!("Failed to read OTA select entry 1: {:?}", Debug2Format(&e)))?;
        let ota_select_1 = unsafe { &*(ota_select_1.as_ptr() as *const OtaSelectEntry) };

        Ok(Self {
            flash,
            ota_data_partition_start: ota_data_partition.offset,
            ota_select_entries: [ota_select_0.clone(), ota_select_1.clone()],
            ota_partitions,
            update_state: None,
        })
    }

    fn get_active_select_entry(&self) -> Result<usize, ()> {
        if self.ota_select_entries[0].is_valid() && self.ota_select_entries[1].is_valid() {
            if self.ota_select_entries[0].ota_seq > self.ota_select_entries[1].ota_seq {
                trace!(
                    "OTA 0 ({}) is newer than OTA 1 ({})",
                    self.ota_select_entries[0].ota_seq,
                    self.ota_select_entries[1].ota_seq
                );
                Ok(0)
            } else {
                trace!(
                    "OTA 1 ({}) is newer than OTA 0 ({})",
                    self.ota_select_entries[1].ota_seq,
                    self.ota_select_entries[0].ota_seq
                );
                Ok(1)
            }
        } else if self.ota_select_entries[0].is_valid() {
            trace!("OTA 0 ({}) is valid, OTA 1 is invalid", self.ota_select_entries[0].ota_seq);
            Ok(0)
        } else if self.ota_select_entries[1].is_valid() {
            trace!("OTA 1 ({}) is valid, OTA 0 is invalid", self.ota_select_entries[1].ota_seq);
            Ok(1)
        } else {
            error!("No valid OTA Entries found");
            Err(())
        }
    }

    pub fn get_selected_boot_partition(&mut self) -> u8 {
        if self.ota_select_entries[0].is_invalid() && self.ota_select_entries[1].is_invalid() {
            debug!("No valid OTA partitions found. Booting from ota 0");
            return 0;
        }

        trace!("OTA Select 0: {:?}", self.ota_select_entries[0]);
        trace!("OTA Select 1: {:?}", self.ota_select_entries[1]);

        let Ok(valid_select) = self.get_active_select_entry() else { return 0 }; // TODO: Check... Not sure if this panic can really happen
        let select = &self.ota_select_entries[valid_select];

        let seq = select.ota_seq - 1; // Raw OTA sequence number. May be more than # of OTA slots
        let boot_index = seq % N as u32; // Actual OTA partition selection
        trace!("Mapping OTA Sequence {} to partition {}", seq, boot_index);
        boot_index as u8
    }

    pub unsafe fn rewrite_ota_select(&mut self, boot_index: u8) -> Result<(), ()> {
        assert!(boot_index < N as u8);
        let boot_index = boot_index as u32;
        match self.get_active_select_entry() {
            Ok(active_entry) => {
                let last_seq = self.ota_select_entries[active_entry].ota_seq;
                let mut i = 0;
                while last_seq > (boot_index + 1) % N as u32 + i * N as u32 {
                    i += 1;
                }
                let inactive_entry = (active_entry + 1) % 2;
                let select = &mut self.ota_select_entries[inactive_entry];
                select.ota_seq = (boot_index + 1) % N as u32 + i * N as u32;
                select.ota_state = OtaImgState::ImgNew;
                let bytes: [u8; 4] = unsafe { core::mem::transmute(select.ota_seq) };
                select.crc = crc32_le(u32::MAX, bytes.as_slice());
                self.flash
                    .write(self.ota_data_partition_start + inactive_entry as u32 * FlashStorage::SECTOR_SIZE, unsafe {
                        core::slice::from_raw_parts(select as *const OtaSelectEntry as *const u8, 32)
                    })
                    .map_err(|e| error!("Failed to write OTA select entry: {:?}", Debug2Format(&e)))?;
                debug!("Wrote new OTA: partition index: {} Sequence: {} Select: {}", boot_index, select.ota_seq, inactive_entry);
            }
            Err(_) => {
                let select = &mut self.ota_select_entries[0];
                select.ota_seq = boot_index as u32 + 1;
                select.ota_state = OtaImgState::ImgNew;
                let bytes: [u8; 4] = unsafe { core::mem::transmute(select.ota_seq) };
                select.crc = crc32_le(u32::MAX, bytes.as_slice());
                self.flash
                    .write(self.ota_data_partition_start, unsafe {
                        core::slice::from_raw_parts(select as *const OtaSelectEntry as *const u8, 32)
                    })
                    .map_err(|e| error!("Failed to write OTA select entry: {:?}", Debug2Format(&e)))?;
                debug!("Wrote new OTA: partition index: {} Sequence: {} Select0", boot_index, select.ota_seq);
            }
        }
        Ok(())
    }

    pub fn get_next_update_partition(&mut self) -> u8 {
        // we currently just select the next partition which is not the current selected one.
        // Better: https://github.com/espressif/esp-idf/blob/be06a6f5ffe36f9554cfc91fe2036e0fc85fea60/components/app_update/esp_ota_ops.c#L574
        let next = (self.get_selected_boot_partition() + 1) % N as u8;
        trace!("Next OTA partition: {}", next);
        next
    }

    pub fn get_current_partition_checksum(&mut self) -> Result<[u8; 32], ()> {
        // The ESP-IDF bootloader uses a custom image format with a header and multiple segments.
        // At the end of all segments there is a sha256 checksum, which we're interested in
        // https://docs.espressif.com/projects/esp-idf/en/stable/esp32/api-reference/system/app_image_format.html
        let partition = &self.ota_partitions[self.get_selected_boot_partition() as usize];
        let mut buf = [0; core::mem::size_of::<ImageHeader>()];
        self.flash
            .read(partition.offset, &mut buf)
            .map_err(|e| error!("Failed to read partition image header: {:?}", Debug2Format(&e)))?;
        let image_header: ImageHeader = *try_from_bytes(&buf).map_err(|_| error!("Failed to parse image header"))?;
        if image_header.magic != 0xe9 {
            return Err(error!("Invalid image header magic: {:x}", image_header.magic));
        }

        if image_header.segment_count > 4 {
            return Err(error!("Invalid segment count: {}", image_header.segment_count));
        }

        let mut offset = partition.offset + core::mem::size_of::<ImageHeader>() as u32;
        for i in 0..image_header.segment_count {
            let mut buf = [0; core::mem::size_of::<SegmentHeader>()];
            self.flash
                .read(offset, &mut buf)
                .map_err(|e| error!("Failed to read segment header {}: {:?}", i, Debug2Format(&e)))?;
            let segment: SegmentHeader = *try_from_bytes(&buf).map_err(|_| error!("Failed to parse segment header"))?;
            trace!("Segment {}: {:?}", i, Debug2Format(&segment));
            offset += segment.length + core::mem::size_of::<SegmentHeader>() as u32;
        }

        // The following code tries to find the (single) checksum byte and the 32 byte hash at the end of the image.
        // TODO: try to understand the esp idf code to validate if this is correct:
        // https://github.com/espressif/esp-idf/blob/003f3bb5dc7c8af8b71926b7a0118cfc503cab11/components/bootloader_support/src/esp_image_format.c#L898
        let unpadded_length = offset;
        let length = unpadded_length + 1; // Add a byte for the checksum
        let length = (length + 15) & !15; // Pad to next full 16 byte block
        let checksum_length = (length - unpadded_length) as usize;
        let mut buf = [0; 16 + 32];
        self.flash
            .read(offset, &mut buf)
            .map_err(|e| error!("Failed to checksum and hash {:?}", Debug2Format(&e)))?;

        let checksum = buf[checksum_length - 1];
        let hash = &buf[checksum_length..checksum_length + 32];
        info!(
            "App Checksum: {:x} Hash: {}",
            checksum,
            defmt_or_log::Display2Format(&{
                // a hacky way to output the hash as hex string
                // log has {:x?}, but does not support arrays
                // defmt has {:x}
                use core::fmt::Write;
                use heapless::String;
                let mut s: String<64> = String::new();
                for b in hash {
                    write!(s, "{:02x}", b).unwrap();
                }
                s
            })
        );
        if !buf[0..checksum_length - 1].iter().all(|&b| b == 0) {
            info!("Buf {:?}", Debug2Format(&buf));
            info!("Unpadded length: {:x} Padded length: {:x} Checksum length: {}", unpadded_length, length, checksum_length);
            return Err(error!("Invalid alignment handling"));
        }
        Ok(hash.try_into().unwrap())
    }

    pub fn begin_update(&mut self, size: u32, expected_crc: u32) -> Result<(), UpdaterError> {
        let partition = self.get_next_update_partition() as usize;
        if size == 0 {
            error!("OTA size is 0");
            return Err(UpdaterError::OtherError("SIZE_ZERO"));
        }
        let space = self.ota_partitions[partition].size as u32;
        if size > space {
            error!("OTA size {} is larger than partition size {}", size, space);
            return Err(UpdaterError::OtherError("SIZE_TOO_LARGE_FOR_PARTITION"));
        }

        info!("Starting update of size {}", size);
        self.update_state = Some(State {
            partition: partition as u8,
            base_address: self.ota_partitions[partition].offset,
            size,
            actual_crc: !0xffffffff,
            written: 0,
            expected_crc,
        });

        Ok(())
    }

    pub fn write_update(&mut self, offset: u32, data: &[u8]) -> Result<(), UpdaterError> {
        let state = self.update_state.as_mut().ok_or_else(|| UpdaterError::OtherError("NO_UPDATE"))?;
        if offset != state.written {
            error!("Invalid offset for OTA data. Writes must be consecutive");
            return Err(UpdaterError::OtherError("INVALID_OFFSET"));
        }

        let last_page = state.written / FlashStorage::SECTOR_SIZE == state.size / FlashStorage::SECTOR_SIZE;
        if last_page {
            let expected = state.size - state.written;
            if data.len() as u32 != expected {
                error!("Last OTA data block is not the expected size. Expected: {}, got: {}", expected, data.len());
                return Err(UpdaterError::OtherError("LAST_BLOCK_INVALID_SIZE"));
            }
        } else {
            if data.len() as u32 != FlashStorage::SECTOR_SIZE {
                error!("OTA data block is not the expected size. Expected: {}, got: {}", FlashStorage::SECTOR_SIZE, data.len());
                return Err(UpdaterError::OtherError("BLOCK_INVALID_SIZE"));
            }
        }

        self.flash.write(state.base_address + state.written, data).map_err(|e| {
            error!("Failed to write OTA data: {:?}", Debug2Format(&e));
            UpdaterError::FlashWriteError
        })?;
        state.actual_crc = crc32_le(state.actual_crc, data);
        state.written += data.len() as u32;
        Ok(())
    }

    pub fn finalize_update(&mut self) -> Result<(), UpdaterError> {
        let state = self.update_state.as_ref().ok_or_else(|| UpdaterError::OtherError("NO_UPDATE"))?;
        if state.written != state.size {
            error!("OTA data is not complete. Expected: {}, got: {}", state.size, state.written);
            return Err(UpdaterError::OtherError("INCOMPLETE_UPDATE"));
        }

        if state.actual_crc != state.expected_crc {
            error!("CRC32 mismatch. Expected: {:x}, got: {:x}", state.expected_crc, state.actual_crc);
            return Err(UpdaterError::CrcMissmatch);
        }

        unsafe { self.rewrite_ota_select(state.partition) }.map_err(|_| UpdaterError::OtherError("REWRITE_OTA_SELECT"))?;

        self.update_state = None;
        info!("CRC ok. Update done");
        Ok(())
    }
}

Use begin_update, write_update and finalize_update as entrypoints.
You probably dont need get_current_partition_checksum

Sorry for the missing docs :)

No guarantees, whatsoever...