ssh_stamp_esp32/
flash.rs

1// SPDX-FileCopyrightText: 2026 Roman Valls Guimera <brainstorm@nopcode.org>
2// SPDX-FileCopyrightText: 2026 Julio Beltran Ortega <jubeormk1@gmail.com>
3// SPDX-FileCopyrightText: 2026 pancake <pancake@nopcode.org>
4// SPDX-FileCopyrightText: 2026 Anthony Tambasco <anthony.tambasco@fastmail.com>
5//
6// SPDX-License-Identifier: GPL-3.0-or-later
7
8//! Flash storage and OTA implementation for ESP32 family
9//!
10//! Provides access to flash storage for configuration persistence and firmware updates.
11
12use embedded_storage::nor_flash::NorFlash;
13use esp_bootloader_esp_idf::ota::OtaImageState;
14use esp_bootloader_esp_idf::ota_updater::OtaUpdater;
15use esp_bootloader_esp_idf::partitions::PARTITION_TABLE_MAX_LEN;
16use esp_hal::peripherals::FLASH;
17use esp_storage::FlashStorage;
18use log::{debug, error};
19use once_cell::sync::OnceCell;
20use ssh_stamp_hal::{FlashError, HalError, OtaActions};
21use sunset_async::SunsetMutex;
22
23const FLASH_BUF_SIZE: usize = FlashStorage::SECTOR_SIZE as usize;
24
25/// Flash storage singleton
26static FLASH_STORAGE: OnceCell<SunsetMutex<FlashBuffer<'static>>> = OnceCell::new();
27
28/// Flash buffer holding both storage and read/write buffer
29#[derive(Debug)]
30pub struct FlashBuffer<'d> {
31    pub flash: FlashStorage<'d>,
32    pub buf: [u8; FLASH_BUF_SIZE],
33}
34
35impl<'d> FlashBuffer<'d> {
36    #[must_use]
37    pub fn new(flash: FlashStorage<'static>) -> Self {
38        Self {
39            flash,
40            buf: [0u8; FLASH_BUF_SIZE],
41        }
42    }
43
44    /// Get mutable references to both flash and buffer
45    pub fn split_ref_mut(&mut self) -> (&mut FlashStorage<'d>, &mut [u8]) {
46        (&mut self.flash, &mut self.buf)
47    }
48}
49
50/// Initialize flash storage
51pub fn init(flash: FLASH<'static>) {
52    let fl = FlashBuffer::new(FlashStorage::new(flash));
53
54    let Ok(()) = FLASH_STORAGE.set(SunsetMutex::new(fl)) else {
55        log::warn!("Flash storage already initialized");
56        return;
57    };
58}
59
60/// Get flash storage and buffer
61pub fn get_flash_n_buffer() -> Option<&'static SunsetMutex<FlashBuffer<'static>>> {
62    FLASH_STORAGE.get()
63}
64
65/// OTA writer for ESP32
66#[derive(Debug, Copy, Clone)]
67pub struct EspOtaWriter {}
68
69impl EspOtaWriter {
70    #[must_use]
71    pub fn new() -> Self {
72        EspOtaWriter {}
73    }
74
75    async fn next_ota_size() -> Result<u32, HalError> {
76        let Some(fb) = get_flash_n_buffer() else {
77            error!("Flash storage not initialized");
78            return Err(HalError::Flash(FlashError::InternalError));
79        };
80        let mut fb = fb.lock().await;
81
82        let (storage, _) = fb.split_ref_mut();
83        let mut buff_ota = [0u8; PARTITION_TABLE_MAX_LEN];
84
85        let mut ota = OtaUpdater::new(storage, &mut buff_ota)
86            .map_err(|_| HalError::Flash(FlashError::InternalError))?;
87        let (target_partition, _) = ota
88            .next_partition()
89            .map_err(|_| HalError::Flash(FlashError::InternalError))?;
90
91        u32::try_from(target_partition.partition_size())
92            .map_err(|_| HalError::Flash(FlashError::InternalError))
93    }
94
95    async fn write_to_target(offset: u32, data: &[u8]) -> Result<(), HalError> {
96        let Some(fb) = get_flash_n_buffer() else {
97            error!("Flash storage not initialized");
98            return Err(HalError::Flash(FlashError::InternalError));
99        };
100        let mut fb = fb.lock().await;
101
102        let (storage, _) = fb.split_ref_mut();
103        let mut buff_ota = [0u8; PARTITION_TABLE_MAX_LEN];
104
105        let mut ota = OtaUpdater::new(storage, &mut buff_ota)
106            .map_err(|_| HalError::Flash(FlashError::InternalError))?;
107        let (mut target_partition, part_type) = ota
108            .next_partition()
109            .map_err(|_| HalError::Flash(FlashError::InternalError))?;
110
111        debug!("Flashing image to {part_type:?}");
112        debug!(
113            "Writing data to target_partition at offset {}, with len {}",
114            offset,
115            data.len()
116        );
117
118        // `as_nor_flash` refuses a region whose partition is encrypted,
119        // rather than silently writing plaintext through it.
120        let mut target_partition = target_partition
121            .as_nor_flash()
122            .map_err(|_| HalError::Flash(FlashError::Write))?;
123        NorFlash::write(&mut target_partition, offset, data)
124            .map_err(|_| HalError::Flash(FlashError::Write))?;
125
126        Ok(())
127    }
128
129    async fn activate_next_ota_slot() -> Result<(), HalError> {
130        let Some(fb) = get_flash_n_buffer() else {
131            error!("Flash storage not initialized");
132            return Err(HalError::Flash(FlashError::InternalError));
133        };
134        let mut fb = fb.lock().await;
135
136        let (storage, _) = fb.split_ref_mut();
137        let mut buff_ota = [0u8; PARTITION_TABLE_MAX_LEN];
138
139        let mut ota = OtaUpdater::new(storage, &mut buff_ota)
140            .map_err(|_| HalError::Flash(FlashError::InternalError))?;
141
142        ota.activate_next_partition()
143            .map_err(|_| HalError::Flash(FlashError::Write))?;
144        ota.set_current_ota_state(OtaImageState::New)
145            .map_err(|_| HalError::Flash(FlashError::Write))?;
146
147        Ok(())
148    }
149}
150
151impl Default for EspOtaWriter {
152    fn default() -> Self {
153        Self::new()
154    }
155}
156
157impl OtaActions for EspOtaWriter {
158    async fn try_validating_current_ota_partition() -> Result<(), HalError> {
159        let Some(fb) = get_flash_n_buffer() else {
160            error!("Flash storage not initialized");
161            return Err(HalError::Flash(FlashError::InternalError));
162        };
163        let mut fb = fb.lock().await;
164
165        let (storage, _) = fb.split_ref_mut();
166        let mut buff_ota = [0u8; PARTITION_TABLE_MAX_LEN];
167
168        let mut ota = OtaUpdater::new(storage, &mut buff_ota)
169            .map_err(|_| HalError::Flash(FlashError::InternalError))?;
170        ota.selected_partition()
171            .map_err(|_| HalError::Flash(FlashError::InternalError))?;
172
173        debug!("current image state {:?}", ota.current_ota_state());
174
175        let state_result = ota.current_ota_state();
176        if let Ok(state) = state_result
177            && (state == esp_bootloader_esp_idf::ota::OtaImageState::New
178                || state == esp_bootloader_esp_idf::ota::OtaImageState::PendingVerify)
179        {
180            ota.set_current_ota_state(esp_bootloader_esp_idf::ota::OtaImageState::Valid)
181                .map_err(|_| HalError::Flash(FlashError::Write))?;
182            debug!("Changed state to VALID");
183        }
184
185        Ok(())
186    }
187
188    async fn get_ota_partition_size() -> Result<u32, HalError> {
189        Self::next_ota_size().await
190    }
191
192    async fn write_ota_data(&self, offset: u32, data: &[u8]) -> Result<(), HalError> {
193        Self::write_to_target(offset, data).await
194    }
195
196    async fn finalize_ota_update(&mut self) -> Result<(), HalError> {
197        Self::activate_next_ota_slot().await
198    }
199
200    fn reset_device(&self) -> ! {
201        esp_hal::system::software_reset()
202    }
203}
ssh_stamp_esp32/flash.rs

ssh_stamp_esp32/
flash.rs
