esp_partition_table/
table.rs

1use crate::{Md5Data, PartitionBuffer, PartitionEntry, PartitionError, PartitionMd5};
2
3/// Partition table info
4#[derive(Clone, Copy, Debug, PartialEq, Eq)]
5pub struct PartitionTable {
6    /// Address of table
7    pub addr: u32,
8
9    /// Size of table
10    pub size: usize,
11}
12
13impl PartitionTable {
14    /// Address of partition table
15    pub const DEFAULT_ADDR: u32 = 0x8000;
16
17    /// Maximum size of partition table
18    pub const MAX_SIZE: usize = 0x1000;
19
20    /// Maxumum number of partition entries
21    pub const MAX_ENTRIES: usize = Self::MAX_SIZE / PartitionEntry::SIZE;
22}
23
24impl Default for PartitionTable {
25    fn default() -> Self {
26        Self::new(Self::DEFAULT_ADDR, Self::MAX_SIZE)
27    }
28}
29
30impl PartitionTable {
31    /// Instantiate partition table with specified address and size
32    pub fn new(addr: u32, size: usize) -> Self {
33        Self { addr, size }
34    }
35
36    /// Get maximum number of entries
37    pub fn max_entries(&self) -> usize {
38        self.size / PartitionEntry::SIZE
39    }
40}
41
42#[derive(Clone, Copy, Debug)]
43enum InternalState {
44    Init,
45    Proc,
46    Done,
47}
48
49/// Partition table reader state
50#[derive(Clone)]
51pub struct PartitionReaderState {
52    offset: u32,
53    end: u32,
54
55    #[cfg(feature = "md5")]
56    md5: Result<Md5Data, md5::Context>,
57
58    #[cfg(feature = "md5")]
59    calc_md5: bool,
60
61    stored_md5: Option<Md5Data>,
62
63    state: InternalState,
64}
65
66impl PartitionReaderState {
67    /// Instantiate reader state
68    ///
69    /// If `md5` feature isn't enabled `calc_md5` argument will be ignored.
70    pub fn new(offset: u32, length: usize, calc_md5: bool) -> Self {
71        #[cfg(not(feature = "md5"))]
72        let _ = calc_md5;
73
74        Self {
75            offset,
76            end: offset + length as u32,
77
78            #[cfg(feature = "md5")]
79            md5: Err(md5::Context::new()),
80
81            #[cfg(feature = "md5")]
82            calc_md5,
83
84            stored_md5: None,
85
86            state: InternalState::Proc,
87        }
88    }
89
90    /// Get current offset
91    pub fn offset(&self) -> u32 {
92        self.offset
93    }
94
95    /// Reader reached end of data
96    pub fn is_done(&self) -> bool {
97        matches!(self.state, InternalState::Done)
98    }
99
100    /// Get stored MD5 checksum
101    pub fn stored_md5(&self) -> Option<&Md5Data> {
102        self.stored_md5.as_ref()
103    }
104
105    /// Get computed MD5 checksum
106    pub fn actual_md5(&self) -> Option<&Md5Data> {
107        #[cfg(feature = "md5")]
108        {
109            self.md5.as_ref().ok()
110        }
111
112        #[cfg(not(feature = "md5"))]
113        {
114            None
115        }
116    }
117
118    /// Check partition table consistency
119    pub fn check_md5(&self) -> Option<bool> {
120        #[cfg(feature = "md5")]
121        if let (Some(stored_md5), Some(actual_md5)) = (self.stored_md5(), self.actual_md5()) {
122            Some(stored_md5 == actual_md5)
123        } else {
124            None
125        }
126
127        #[cfg(not(feature = "md5"))]
128        {
129            None
130        }
131    }
132
133    fn check(&mut self) -> Result<(), PartitionError> {
134        if self.offset >= self.end {
135            self.state = InternalState::Done;
136        }
137
138        if matches!(self.state, InternalState::Done) {
139            Err(PartitionError::NotEnoughData)
140        } else {
141            Ok(())
142        }
143    }
144
145    /// Read partition data from buffer
146    pub fn read(&mut self, buffer: &PartitionBuffer) -> Result<PartitionEntry, PartitionError> {
147        self.check()?;
148
149        let result = match *buffer
150            .split_first_chunk()
151            .ok_or(PartitionError::NotEnoughData)?
152            .0
153        {
154            PartitionEntry::MAGIC => {
155                #[cfg(feature = "md5")]
156                if self.calc_md5 {
157                    if let Err(ctx) = &mut self.md5 {
158                        ctx.consume(buffer);
159                    }
160                }
161
162                buffer.try_into()
163            }
164            PartitionMd5::MAGIC => match buffer.try_into() {
165                Ok(PartitionMd5 { data }) => {
166                    self.stored_md5 = Some(data);
167                    self.offset += PartitionEntry::SIZE as u32;
168                    Err(PartitionError::NotEnoughData)
169                }
170                Err(error) => Err(error),
171            },
172            [0xff, 0xff] => Err(PartitionError::NotEnoughData),
173            _ => Err(PartitionError::InvalidMagic),
174        };
175
176        if let Err(error) = &result {
177            if let PartitionError::NotEnoughData = error {
178                #[cfg(feature = "md5")]
179                if self.calc_md5 && self.md5.is_err() {
180                    self.md5 = Ok(self.md5.as_mut().unwrap_err().clone().compute().into());
181                }
182            }
183
184            self.state = InternalState::Done;
185        } else {
186            self.offset += PartitionEntry::SIZE as u32;
187        }
188
189        result
190    }
191}
192
193/// Partition table writer state
194pub struct PartitionWriterState {
195    offset: u32,
196    end: u32,
197
198    #[cfg(feature = "md5")]
199    md5: Result<Md5Data, md5::Context>,
200
201    #[cfg(feature = "md5")]
202    write_md5: bool,
203
204    state: InternalState,
205}
206
207impl PartitionWriterState {
208    /// Instantiate writer state
209    ///
210    /// If `md5` feature isn't enabled `write_md5` argument will be ignored.
211    pub fn new(offset: u32, length: usize, write_md5: bool) -> Self {
212        #[cfg(not(feature = "md5"))]
213        let _ = write_md5;
214
215        Self {
216            offset,
217            end: offset + length as u32,
218
219            #[cfg(feature = "md5")]
220            md5: Err(md5::Context::new()),
221
222            #[cfg(feature = "md5")]
223            write_md5,
224
225            state: InternalState::Init,
226        }
227    }
228
229    /// Get current offset
230    pub fn offset(&self) -> u32 {
231        self.offset
232    }
233
234    /// Writer reached end of data
235    pub fn is_done(&self) -> bool {
236        matches!(self.state, InternalState::Done)
237    }
238
239    /// Get computed MD5 checksum
240    pub fn actual_md5(&self) -> Option<&Md5Data> {
241        #[cfg(feature = "md5")]
242        {
243            self.md5.as_ref().ok()
244        }
245
246        #[cfg(not(feature = "md5"))]
247        {
248            None
249        }
250    }
251
252    fn check(&mut self) -> Result<(), PartitionError> {
253        if self.offset >= self.end {
254            self.state = InternalState::Done;
255        }
256
257        match self.state {
258            InternalState::Init => {
259                self.state = InternalState::Proc;
260                Ok(())
261            }
262            InternalState::Proc => {
263                self.offset += PartitionEntry::SIZE as u32;
264                Ok(())
265            }
266            InternalState::Done => Err(PartitionError::TooManyData),
267        }
268    }
269
270    /// Write partition data into buffer
271    ///
272    /// If `md5` feature is used and partition is None then MD5 checksum will be written.
273    pub fn write(
274        &mut self,
275        buffer: &mut PartitionBuffer,
276        partition: impl AsRef<PartitionEntry>,
277    ) -> Result<(), PartitionError> {
278        self.check()?;
279
280        partition.as_ref().to_bytes(buffer)?;
281
282        #[cfg(feature = "md5")]
283        if self.write_md5 {
284            if let Err(ctx) = &mut self.md5 {
285                ctx.consume(buffer);
286            }
287        }
288
289        Ok(())
290    }
291
292    /// Write partition MD5 into buffer
293    ///
294    /// If `md5` feature is used and partition is None then MD5 checksum will be written.
295    pub fn write_md5(&mut self, buffer: &mut PartitionBuffer) -> Result<(), PartitionError> {
296        self.check()?;
297
298        self.state = InternalState::Done;
299
300        #[cfg(not(feature = "md5"))]
301        let _ = buffer;
302
303        #[cfg(feature = "md5")]
304        if self.write_md5 && self.md5.is_err() {
305            let md5 = PartitionMd5::from(self.md5.as_mut().unwrap_err().clone().compute());
306            md5.to_bytes(buffer)?;
307            self.md5 = Ok(md5.into());
308        }
309
310        Ok(())
311    }
312
313    /// Finalize writer
314    pub fn finish(&mut self) {
315        self.state = InternalState::Done;
316    }
317}
318
319#[cfg(test)]
320mod test {
321    use crate::*;
322
323    #[test]
324    fn read_partitions() {
325        let table = include_bytes!("../tests/partitions.bin");
326        let data = &table[..];
327        let mut reader = PartitionReaderState::new(0, data.len(), true);
328
329        let (part, data) = data.split_first_chunk().unwrap();
330        let part = reader.read(part).unwrap();
331        assert_eq!(part.type_, PartitionType::Data(DataPartitionType::Nvs));
332        assert_eq!(part.offset, 36 << 10);
333        assert_eq!(part.size, 24 << 10);
334        assert_eq!(part.name(), "nvs");
335        assert!(!part.encrypted);
336
337        let (part, data) = data.split_first_chunk().unwrap();
338        let part = reader.read(part).unwrap();
339        assert_eq!(part.type_, PartitionType::Data(DataPartitionType::Phy));
340        assert_eq!(part.offset, 60 << 10);
341        assert_eq!(part.size, 4 << 10);
342        assert_eq!(part.name(), "phy_init");
343        assert!(!part.encrypted);
344
345        let (part, data) = data.split_first_chunk().unwrap();
346        let part = reader.read(part).unwrap();
347        assert_eq!(part.type_, PartitionType::App(AppPartitionType::Factory));
348        assert_eq!(part.offset, 64 << 10);
349        assert_eq!(part.size, 3 << 20);
350        assert_eq!(part.name(), "factory");
351        assert!(!part.encrypted);
352
353        let (part, data) = data.split_first_chunk().unwrap();
354        let part = reader.read(part).unwrap();
355        assert_eq!(part.type_, PartitionType::Data(DataPartitionType::CoreDump));
356        assert_eq!(part.offset, (64 << 10) + (3 << 20));
357        assert_eq!(part.size, 64 << 10);
358        assert_eq!(part.name(), "coredump");
359        assert!(!part.encrypted);
360
361        let (part, data) = data.split_first_chunk().unwrap();
362        let part = reader.read(part).unwrap();
363        assert_eq!(part.type_, PartitionType::Data(DataPartitionType::Nvs));
364        assert_eq!(part.offset, (128 << 10) + (3 << 20));
365        assert_eq!(part.size, 64 << 10);
366        assert_eq!(part.name(), "nvs_ext");
367        assert!(!part.encrypted);
368
369        let (part, data) = data.split_first_chunk().unwrap();
370        assert!(matches!(
371            reader.read(part).unwrap_err(),
372            PartitionError::NotEnoughData
373        ));
374        if let Some(md5) = reader.check_md5() {
375            assert!(md5);
376        }
377
378        let (part, _) = data.split_first_chunk().unwrap();
379        assert!(matches!(
380            reader.read(part).unwrap_err(),
381            PartitionError::NotEnoughData
382        ));
383    }
384
385    #[test]
386    fn read_partitions_ota() {
387        let table = include_bytes!("../tests/partitions-ota.bin");
388        let data = &table[..];
389        let mut reader = PartitionReaderState::new(0, data.len(), true);
390
391        let (part, data) = data.split_first_chunk().unwrap();
392        let part = reader.read(part).unwrap();
393        assert_eq!(part.type_, PartitionType::Data(DataPartitionType::Nvs));
394        assert_eq!(part.offset, 36 << 10);
395        assert_eq!(part.size, 16 << 10);
396        assert_eq!(part.name(), "nvs");
397        assert!(!part.encrypted);
398
399        let (part, data) = data.split_first_chunk().unwrap();
400        let part = reader.read(part).unwrap();
401        assert_eq!(part.type_, PartitionType::Data(DataPartitionType::Ota));
402        assert_eq!(part.offset, 52 << 10);
403        assert_eq!(part.size, 8 << 10);
404        assert_eq!(part.name(), "otadata");
405        assert!(!part.encrypted);
406
407        let (part, data) = data.split_first_chunk().unwrap();
408        let part = reader.read(part).unwrap();
409        assert_eq!(part.type_, PartitionType::Data(DataPartitionType::Phy));
410        assert_eq!(part.offset, 60 << 10);
411        assert_eq!(part.size, 4 << 10);
412        assert_eq!(part.name(), "phy_init");
413        assert!(!part.encrypted);
414
415        let (part, data) = data.split_first_chunk().unwrap();
416        let part = reader.read(part).unwrap();
417        assert_eq!(part.type_, PartitionType::App(AppPartitionType::Factory));
418        assert_eq!(part.offset, 64 << 10);
419        assert_eq!(part.size, 1 << 20);
420        assert_eq!(part.name(), "factory");
421        assert!(!part.encrypted);
422
423        let (part, data) = data.split_first_chunk().unwrap();
424        let part = reader.read(part).unwrap();
425        assert_eq!(part.type_, PartitionType::App(AppPartitionType::Ota(0)));
426        assert_eq!(part.offset, (64 << 10) + (1 << 20));
427        assert_eq!(part.size, 1 << 20);
428        assert_eq!(part.name(), "ota_0");
429        assert!(!part.encrypted);
430
431        let (part, data) = data.split_first_chunk().unwrap();
432        let part = reader.read(part).unwrap();
433        assert_eq!(part.type_, PartitionType::App(AppPartitionType::Ota(1)));
434        assert_eq!(part.offset, (64 << 10) + (2 << 20));
435        assert_eq!(part.size, 1 << 20);
436        assert_eq!(part.name(), "ota_1");
437        assert!(!part.encrypted);
438
439        let (part, data) = data.split_first_chunk().unwrap();
440        let part = reader.read(part).unwrap();
441        assert_eq!(part.type_, PartitionType::Data(DataPartitionType::CoreDump));
442        assert_eq!(part.offset, (64 << 10) + (3 << 20));
443        assert_eq!(part.size, 64 << 10);
444        assert_eq!(part.name(), "coredump");
445        assert!(!part.encrypted);
446
447        let (part, data) = data.split_first_chunk().unwrap();
448        let part = reader.read(part).unwrap();
449        assert_eq!(part.type_, PartitionType::Data(DataPartitionType::Nvs));
450        assert_eq!(part.offset, (128 << 10) + (3 << 20));
451        assert_eq!(part.size, 64 << 10);
452        assert_eq!(part.name(), "nvs_ext");
453        assert!(!part.encrypted);
454
455        let (part, data) = data.split_first_chunk().unwrap();
456        assert!(matches!(
457            reader.read(part).unwrap_err(),
458            PartitionError::NotEnoughData
459        ));
460        if let Some(md5) = reader.check_md5() {
461            assert!(md5);
462        }
463
464        let (part, _) = data.split_first_chunk().unwrap();
465        assert!(matches!(
466            reader.read(part).unwrap_err(),
467            PartitionError::NotEnoughData
468        ));
469    }
470
471    #[test]
472    fn write_partitions() {
473        let src_table = include_bytes!("../tests/partitions.bin");
474        let mut dst_table = [0u8; PartitionTable::MAX_SIZE];
475
476        let mut src_data = &src_table[..];
477        let mut dst_data = &mut dst_table[..];
478        let mut reader = PartitionReaderState::new(0, src_data.len(), true);
479        let mut writer = PartitionWriterState::new(0, dst_data.len(), true);
480
481        loop {
482            let (src_part, next_src_data) = src_data.split_first_chunk().unwrap();
483            src_data = next_src_data;
484            let part = match reader.read(src_part) {
485                Ok(part) => Some(part),
486                Err(PartitionError::NotEnoughData) => None,
487                Err(error) => panic!("{error:?}"),
488            };
489            let (dst_part, next_dst_data) = dst_data.split_first_chunk_mut().unwrap();
490            dst_data = next_dst_data;
491            if let Some(part) = part {
492                writer.write(dst_part, part).unwrap();
493            } else {
494                writer.write_md5(dst_part).unwrap();
495                break;
496            }
497        }
498
499        let len = src_table.len()
500            - src_data.len()
501            - if !cfg!(feature = "md5") {
502                PartitionEntry::SIZE
503            } else {
504                0
505            };
506
507        assert_eq!(&dst_table[..len], &src_table[..len]);
508    }
509
510    #[test]
511    fn write_partitions_ota() {
512        let src_table = include_bytes!("../tests/partitions-ota.bin");
513        let mut dst_table = [0u8; PartitionTable::MAX_SIZE];
514
515        let mut src_data = &src_table[..];
516        let mut dst_data = &mut dst_table[..];
517        let mut reader = PartitionReaderState::new(0, src_data.len(), true);
518        let mut writer = PartitionWriterState::new(0, dst_data.len(), true);
519
520        loop {
521            let (src_part, next_src_data) = src_data.split_first_chunk().unwrap();
522            src_data = next_src_data;
523            let part = match reader.read(src_part) {
524                Ok(part) => Some(part),
525                Err(PartitionError::NotEnoughData) => None,
526                Err(error) => panic!("{error:?}"),
527            };
528            let (dst_part, next_dst_data) = dst_data.split_first_chunk_mut().unwrap();
529            dst_data = next_dst_data;
530            if let Some(part) = part {
531                writer.write(dst_part, part).unwrap();
532            } else {
533                writer.write_md5(dst_part).unwrap();
534                break;
535            }
536        }
537
538        let len = src_table.len()
539            - src_data.len()
540            - if !cfg!(feature = "md5") {
541                PartitionEntry::SIZE
542            } else {
543                0
544            };
545
546        assert_eq!(&dst_table[..len], &src_table[..len]);
547    }
548}
esp_partition_table/table.rs

esp_partition_table/
table.rs
