1use digest::Digest;
2#[cfg(target_os = "none")]
3use embassy_embedded_hal::flash::partition::BlockingPartition;
4#[cfg(target_os = "none")]
5use embassy_sync::blocking_mutex::raw::NoopRawMutex;
6use embedded_storage::nor_flash::NorFlash;
7
8use super::FirmwareUpdaterConfig;
9use crate::{FirmwareUpdaterError, State, BOOT_MAGIC, DFU_DETACH_MAGIC, STATE_ERASE_VALUE, SWAP_MAGIC};
10
11pub struct BlockingFirmwareUpdater<'d, DFU: NorFlash, STATE: NorFlash> {
14    dfu: DFU,
15    state: BlockingFirmwareState<'d, STATE>,
16    last_erased_dfu_sector_index: Option<usize>,
17}
18
19#[cfg(target_os = "none")]
20impl<'a, DFU: NorFlash, STATE: NorFlash>
21    FirmwareUpdaterConfig<BlockingPartition<'a, NoopRawMutex, DFU>, BlockingPartition<'a, NoopRawMutex, STATE>>
22{
23    pub fn from_linkerfile_blocking(
55        dfu_flash: &'a embassy_sync::blocking_mutex::Mutex<NoopRawMutex, core::cell::RefCell<DFU>>,
56        state_flash: &'a embassy_sync::blocking_mutex::Mutex<NoopRawMutex, core::cell::RefCell<STATE>>,
57    ) -> Self {
58        extern "C" {
59            static __bootloader_state_start: u32;
60            static __bootloader_state_end: u32;
61            static __bootloader_dfu_start: u32;
62            static __bootloader_dfu_end: u32;
63        }
64
65        let dfu = unsafe {
66            let start = &__bootloader_dfu_start as *const u32 as u32;
67            let end = &__bootloader_dfu_end as *const u32 as u32;
68            trace!("DFU: 0x{:x} - 0x{:x}", start, end);
69
70            BlockingPartition::new(dfu_flash, start, end - start)
71        };
72        let state = unsafe {
73            let start = &__bootloader_state_start as *const u32 as u32;
74            let end = &__bootloader_state_end as *const u32 as u32;
75            trace!("STATE: 0x{:x} - 0x{:x}", start, end);
76
77            BlockingPartition::new(state_flash, start, end - start)
78        };
79
80        Self { dfu, state }
81    }
82}
83
84impl<'d, DFU: NorFlash, STATE: NorFlash> BlockingFirmwareUpdater<'d, DFU, STATE> {
85    pub fn new(config: FirmwareUpdaterConfig<DFU, STATE>, aligned: &'d mut [u8]) -> Self {
92        Self {
93            dfu: config.dfu,
94            state: BlockingFirmwareState::new(config.state, aligned),
95            last_erased_dfu_sector_index: None,
96        }
97    }
98
99    pub fn get_state(&mut self) -> Result<State, FirmwareUpdaterError> {
105        self.state.get_state()
106    }
107
108    #[cfg(feature = "_verify")]
120    pub fn verify_and_mark_updated(
121        &mut self,
122        _public_key: &[u8; 32],
123        _signature: &[u8; 64],
124        _update_len: u32,
125    ) -> Result<(), FirmwareUpdaterError> {
126        assert!(_update_len <= self.dfu.capacity() as u32);
127
128        self.state.verify_booted()?;
129
130        #[cfg(feature = "ed25519-dalek")]
131        {
132            use ed25519_dalek::{Signature, SignatureError, Verifier, VerifyingKey};
133
134            use crate::digest_adapters::ed25519_dalek::Sha512;
135
136            let into_signature_error = |e: SignatureError| FirmwareUpdaterError::Signature(e.into());
137
138            let public_key = VerifyingKey::from_bytes(_public_key).map_err(into_signature_error)?;
139            let signature = Signature::from_bytes(_signature);
140
141            let mut message = [0; 64];
142            let mut chunk_buf = [0; 2];
143            self.hash::<Sha512>(_update_len, &mut chunk_buf, &mut message)?;
144
145            public_key.verify(&message, &signature).map_err(into_signature_error)?;
146            return self.state.mark_updated();
147        }
148        #[cfg(feature = "ed25519-salty")]
149        {
150            use salty::{PublicKey, Signature};
151
152            use crate::digest_adapters::salty::Sha512;
153
154            fn into_signature_error<E>(_: E) -> FirmwareUpdaterError {
155                FirmwareUpdaterError::Signature(signature::Error::default())
156            }
157
158            let public_key = PublicKey::try_from(_public_key).map_err(into_signature_error)?;
159            let signature = Signature::try_from(_signature).map_err(into_signature_error)?;
160
161            let mut message = [0; 64];
162            let mut chunk_buf = [0; 2];
163            self.hash::<Sha512>(_update_len, &mut chunk_buf, &mut message)?;
164
165            let r = public_key.verify(&message, &signature);
166            trace!(
167                "Verifying with public key {}, signature {} and message {} yields ok: {}",
168                public_key.to_bytes(),
169                signature.to_bytes(),
170                message,
171                r.is_ok()
172            );
173            r.map_err(into_signature_error)?;
174            return self.state.mark_updated();
175        }
176        #[cfg(not(any(feature = "ed25519-dalek", feature = "ed25519-salty")))]
177        {
178            Err(FirmwareUpdaterError::Signature(signature::Error::new()))
179        }
180    }
181
182    pub fn hash<D: Digest>(
184        &mut self,
185        update_len: u32,
186        chunk_buf: &mut [u8],
187        output: &mut [u8],
188    ) -> Result<(), FirmwareUpdaterError> {
189        let mut digest = D::new();
190        for offset in (0..update_len).step_by(chunk_buf.len()) {
191            self.dfu.read(offset, chunk_buf)?;
192            let len = core::cmp::min((update_len - offset) as usize, chunk_buf.len());
193            digest.update(&chunk_buf[..len]);
194        }
195        output.copy_from_slice(digest.finalize().as_slice());
196        Ok(())
197    }
198
199    pub fn read_dfu(&mut self, offset: u32, buf: &mut [u8]) -> Result<(), FirmwareUpdaterError> {
206        self.dfu.read(offset, buf)?;
207        Ok(())
208    }
209
210    #[cfg(not(feature = "_verify"))]
212    pub fn mark_updated(&mut self) -> Result<(), FirmwareUpdaterError> {
213        self.state.mark_updated()
214    }
215
216    pub fn mark_dfu(&mut self) -> Result<(), FirmwareUpdaterError> {
218        self.state.verify_booted()?;
219        self.state.mark_dfu()
220    }
221
222    pub fn mark_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
224        self.state.mark_booted()
225    }
226
227    pub fn write_firmware(&mut self, offset: usize, data: &[u8]) -> Result<(), FirmwareUpdaterError> {
253        self.state.verify_booted()?;
255
256        let mut remaining_data = data;
258        let mut offset = offset;
259
260        while !remaining_data.is_empty() {
262            let current_sector = offset / DFU::ERASE_SIZE;
264            let sector_start = current_sector * DFU::ERASE_SIZE;
265            let sector_end = sector_start + DFU::ERASE_SIZE;
266            let need_erase = self
268                .last_erased_dfu_sector_index
269                .map_or(true, |last_erased_sector| current_sector != last_erased_sector);
270
271            if need_erase {
273                self.dfu.erase(sector_start as u32, sector_end as u32)?;
274                self.last_erased_dfu_sector_index = Some(current_sector);
275            }
276
277            let write_size = core::cmp::min(remaining_data.len(), sector_end - offset);
279            let (data_chunk, rest) = remaining_data.split_at(write_size);
281
282            self.dfu.write(offset as u32, data_chunk)?;
284
285            remaining_data = rest;
287            offset += write_size;
288        }
289
290        Ok(())
291    }
292
293    pub fn prepare_update(&mut self) -> Result<&mut DFU, FirmwareUpdaterError> {
299        self.state.verify_booted()?;
300        self.dfu.erase(0, self.dfu.capacity() as u32)?;
301
302        Ok(&mut self.dfu)
303    }
304}
305
306pub struct BlockingFirmwareState<'d, STATE> {
310    state: STATE,
311    aligned: &'d mut [u8],
312}
313
314impl<'d, STATE: NorFlash> BlockingFirmwareState<'d, STATE> {
315    pub fn from_config<DFU: NorFlash>(config: FirmwareUpdaterConfig<DFU, STATE>, aligned: &'d mut [u8]) -> Self {
322        Self::new(config.state, aligned)
323    }
324
325    pub fn new(state: STATE, aligned: &'d mut [u8]) -> Self {
332        assert_eq!(aligned.len(), STATE::WRITE_SIZE);
333        Self { state, aligned }
334    }
335
336    fn verify_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
338        let state = self.get_state()?;
339        if state == State::Boot || state == State::DfuDetach || state == State::Revert {
340            Ok(())
341        } else {
342            Err(FirmwareUpdaterError::BadState)
343        }
344    }
345
346    pub fn get_state(&mut self) -> Result<State, FirmwareUpdaterError> {
352        self.state.read(0, &mut self.aligned)?;
353        Ok(State::from(&self.aligned))
354    }
355
356    pub fn mark_updated(&mut self) -> Result<(), FirmwareUpdaterError> {
358        self.set_magic(SWAP_MAGIC)
359    }
360
361    pub fn mark_dfu(&mut self) -> Result<(), FirmwareUpdaterError> {
363        self.set_magic(DFU_DETACH_MAGIC)
364    }
365
366    pub fn mark_booted(&mut self) -> Result<(), FirmwareUpdaterError> {
368        self.set_magic(BOOT_MAGIC)
369    }
370
371    fn set_magic(&mut self, magic: u8) -> Result<(), FirmwareUpdaterError> {
372        self.state.read(0, &mut self.aligned)?;
373
374        if self.aligned.iter().any(|&b| b != magic) {
375            self.state.read(STATE::WRITE_SIZE as u32, &mut self.aligned)?;
377
378            if self.aligned.iter().any(|&b| b != STATE_ERASE_VALUE) {
379                } else {
381                self.aligned.fill(!STATE_ERASE_VALUE);
383                self.state.write(STATE::WRITE_SIZE as u32, &self.aligned)?;
384            }
385
386            self.state.erase(0, self.state.capacity() as u32)?;
388
389            self.aligned.fill(magic);
391            self.state.write(0, &self.aligned)?;
392        }
393        Ok(())
394    }
395}
396
397#[cfg(test)]
398mod tests {
399    use core::cell::RefCell;
400
401    use embassy_embedded_hal::flash::partition::BlockingPartition;
402    use embassy_sync::blocking_mutex::raw::NoopRawMutex;
403    use embassy_sync::blocking_mutex::Mutex;
404    use sha1::{Digest, Sha1};
405
406    use super::*;
407    use crate::mem_flash::MemFlash;
408
409    #[test]
410    fn can_verify_sha1() {
411        let flash = Mutex::<NoopRawMutex, _>::new(RefCell::new(MemFlash::<131072, 4096, 8>::default()));
412        let state = BlockingPartition::new(&flash, 0, 4096);
413        let dfu = BlockingPartition::new(&flash, 65536, 65536);
414        let mut aligned = [0; 8];
415
416        let update = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66];
417        let mut to_write = [0; 4096];
418        to_write[..7].copy_from_slice(update.as_slice());
419
420        let mut updater = BlockingFirmwareUpdater::new(FirmwareUpdaterConfig { dfu, state }, &mut aligned);
421        updater.write_firmware(0, to_write.as_slice()).unwrap();
422        let mut chunk_buf = [0; 2];
423        let mut hash = [0; 20];
424        updater
425            .hash::<Sha1>(update.len() as u32, &mut chunk_buf, &mut hash)
426            .unwrap();
427
428        assert_eq!(Sha1::digest(update).as_slice(), hash);
429    }
430
431    #[test]
432    fn can_verify_sha1_sector_bigger_than_chunk() {
433        let flash = Mutex::<NoopRawMutex, _>::new(RefCell::new(MemFlash::<131072, 4096, 8>::default()));
434        let state = BlockingPartition::new(&flash, 0, 4096);
435        let dfu = BlockingPartition::new(&flash, 65536, 65536);
436        let mut aligned = [0; 8];
437
438        let update = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66];
439        let mut to_write = [0; 4096];
440        to_write[..7].copy_from_slice(update.as_slice());
441
442        let mut updater = BlockingFirmwareUpdater::new(FirmwareUpdaterConfig { dfu, state }, &mut aligned);
443        let mut offset = 0;
444        for chunk in to_write.chunks(1024) {
445            updater.write_firmware(offset, chunk).unwrap();
446            offset += chunk.len();
447        }
448        let mut chunk_buf = [0; 2];
449        let mut hash = [0; 20];
450        updater
451            .hash::<Sha1>(update.len() as u32, &mut chunk_buf, &mut hash)
452            .unwrap();
453
454        assert_eq!(Sha1::digest(update).as_slice(), hash);
455    }
456
457    #[test]
458    fn can_verify_sha1_sector_smaller_than_chunk() {
459        let flash = Mutex::<NoopRawMutex, _>::new(RefCell::new(MemFlash::<131072, 1024, 8>::default()));
460        let state = BlockingPartition::new(&flash, 0, 4096);
461        let dfu = BlockingPartition::new(&flash, 65536, 65536);
462        let mut aligned = [0; 8];
463
464        let update = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66];
465        let mut to_write = [0; 4096];
466        to_write[..7].copy_from_slice(update.as_slice());
467
468        let mut updater = BlockingFirmwareUpdater::new(FirmwareUpdaterConfig { dfu, state }, &mut aligned);
469        let mut offset = 0;
470        for chunk in to_write.chunks(2048) {
471            updater.write_firmware(offset, chunk).unwrap();
472            offset += chunk.len();
473        }
474        let mut chunk_buf = [0; 2];
475        let mut hash = [0; 20];
476        updater
477            .hash::<Sha1>(update.len() as u32, &mut chunk_buf, &mut hash)
478            .unwrap();
479
480        assert_eq!(Sha1::digest(update).as_slice(), hash);
481    }
482
483    #[test]
484    fn can_verify_sha1_cross_sector_boundary() {
485        let flash = Mutex::<NoopRawMutex, _>::new(RefCell::new(MemFlash::<131072, 1024, 8>::default()));
486        let state = BlockingPartition::new(&flash, 0, 4096);
487        let dfu = BlockingPartition::new(&flash, 65536, 65536);
488        let mut aligned = [0; 8];
489
490        let update = [0x00, 0x11, 0x22, 0x33, 0x44, 0x55, 0x66];
491        let mut to_write = [0; 4096];
492        to_write[..7].copy_from_slice(update.as_slice());
493
494        let mut updater = BlockingFirmwareUpdater::new(FirmwareUpdaterConfig { dfu, state }, &mut aligned);
495        let mut offset = 0;
496        for chunk in to_write.chunks(896) {
497            updater.write_firmware(offset, chunk).unwrap();
498            offset += chunk.len();
499        }
500        let mut chunk_buf = [0; 2];
501        let mut hash = [0; 20];
502        updater
503            .hash::<Sha1>(update.len() as u32, &mut chunk_buf, &mut hash)
504            .unwrap();
505
506        assert_eq!(Sha1::digest(update).as_slice(), hash);
507    }
508}