amebazii/types/image/
ota.rs

1// Firmware/OTA Firmware Layout
2// ------------------------------------------------------------------------------------
3//
4// ┌────────────────────────────────────────┐  ┐
5// │        OTA Signature (32 bytes)        │  │
6// ├────────────────────────────────────────┤  │
7// │        Public Key 0 (32 bytes)         │  │
8// │                  ...                   │  │
9// │        Public Key 5 (32 bytes)         │  │
10// ├────────────────────────────────────────┤  │
11// │      SubImage 0 Header (96 bytes)      │  │
12// ├────────────────────────────────────────┤  │
13// │       SubImage 0 FST (96 bytes)        │  │ SubImage 0 Hash
14// ├────────────────────────────────────────┤  │
15// │ SubImage 0 Section 0 Header (96 bytes) │  │
16// ├────────────────────────────────────────┤  │
17// │       SubImage 0 Section 0 Data        │  │
18// └────────────────────────────────────────┘  │
19//                   .                         │
20//                   .                         │
21//                   .                         │
22// ┌────────────────────────────────────────┐  │
23// │ SubImage 0 Section N Header (96 bytes) │  │
24// ├────────────────────────────────────────┤  │
25// │       SubImage 0 Section N Data        │  │
26// ├────────────────────────────────────────┤  ┘
27// │       SubImage 0 Hash (32 bytes)       │
28// ├────────────────────────────────────────┤
29// │      SubImage 1 Header (96 bytes)      │
30// ├────────────────────────────────────────┤
31// │       SubImage 1 FST (96 bytes)        │
32// ├────────────────────────────────────────┤
33// │ SubImage 1 Section 0 Header (96 bytes) │
34// ├────────────────────────────────────────┤
35// │       SubImage 1 Section 0 Data        │
36// └────────────────────────────────────────┘
37//                     .
38
39use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
40use std::{
41    io::{self, Cursor, Write},
42    vec,
43};
44
45use crate::{
46    error::Error,
47    is_valid_data,
48    types::{
49        enums::HashAlgo,
50        from_stream,
51        fst::FST,
52        header::{ImageHeader, KeyBlock},
53        section::Section,
54        BinarySize, DataRefType, DataType, FromStream, ToStream,
55    },
56    util::{skip_aligned, write_fill},
57    write_aligned, write_data, write_padding,
58};
59
60use super::{AsImage, EncryptedOr};
61
62/// Represents a sub-image, including a header, FST (Firmware Security Table), sections, and hash for integrity verification.
63///
64/// This struct provides methods to manipulate sections, retrieve data, and manage the sub-image’s hash.
65#[derive(Debug)]
66pub struct SubImage {
67    /// The header of the sub-image containing general information about the sub-image.
68    pub header: ImageHeader,
69
70    // REVISIT: this struct does not cover the use-case of an encrypted sub-image!!
71    /// The Firmware Security Table (FST) associated with the sub-image.
72    pub fst: EncryptedOr<FST>,
73
74    /// The collection of sections in the sub-image.
75    sections: EncryptedOr<Vec<Section>>,
76
77    /// The hash of the sub-image used for integrity verification.
78    hash: [u8; 32],
79}
80
81impl Default for SubImage {
82    /// Creates a new `SubImage` with default values.
83    ///
84    /// The default `SubImage` is initialized as follows:
85    /// - The `header` is initialized with the default value of `ImageHeader`.
86    /// - The `fst` is initialized with the default value of `FST`.
87    /// - The `sections` is an empty vector.
88    /// - The `hash` is set to an array of 32 `0xFF` bytes (indicating an uninitialized or invalid hash).
89    fn default() -> Self {
90        SubImage {
91            header: ImageHeader::default(),
92            fst: EncryptedOr::Plain(FST::default()),
93            sections: EncryptedOr::Plain(Vec::new()),
94            hash: [0xFF; 32],
95        }
96    }
97}
98
99impl SubImage {
100    /// Returns a reference to the hash of the sub-image.
101    ///
102    /// # Returns:
103    /// - A reference to the 32-byte hash of the sub-image.
104    ///
105    pub fn get_hash(&self) -> &[u8; 32] {
106        &self.hash
107    }
108
109    /// Returns a reference to the sections in the sub-image.
110    ///
111    /// This method provides access to the sub-image's sections as an immutable slice.
112    ///
113    /// # Returns:
114    /// - A reference to the `Vec<Section>` representing the sections in the sub-image.
115    ///
116    pub fn get_sections(&self) -> &[Section] {
117        match &self.sections {
118            EncryptedOr::Plain(sections) => sections,
119            EncryptedOr::Encrypted(_) => panic!("SubImage is encrypted"),
120        }
121    }
122
123    /// Returns a mutable reference to the sections in the sub-image.
124    ///
125    /// This method provides access to the sub-image's sections as a mutable slice,
126    /// allowing for modification of the sections.
127    ///
128    /// # Returns:
129    /// - A mutable reference to the `Vec<Section>` representing the sections in the sub-image.
130    ///
131    pub fn get_sections_mut(&mut self) -> &mut [Section] {
132        match &mut self.sections {
133            EncryptedOr::Plain(sections) => sections,
134            EncryptedOr::Encrypted(_) => panic!("SubImage is encrypted"),
135        }
136    }
137
138    /// Adds a new section to the sub-image.
139    ///
140    /// This method appends the provided `section` to the list of sections in the sub-image.
141    ///
142    /// # Arguments:
143    /// - `section`: The section to add to the sub-image.
144    ///
145    pub fn add_section(&mut self, section: Section) {
146        match &mut self.sections {
147            EncryptedOr::Plain(sections) => sections.push(section),
148            EncryptedOr::Encrypted(_) => panic!("SubImage is encrypted"),
149        }
150    }
151
152    /// Removes the section at the specified index from the sub-image.
153    ///
154    /// This method removes the section at the given `index` from the list of sections.
155    /// If the index is out of bounds, the method will panic.
156    ///
157    /// # Arguments:
158    /// - `index`: The index of the section to remove.
159    ///
160    pub fn rem_section_at(&mut self, index: usize) {
161        match &mut self.sections {
162            EncryptedOr::Plain(sections) => sections.remove(index),
163            EncryptedOr::Encrypted(_) => panic!("SubImage is encrypted"),
164        };
165    }
166
167    /// Returns a reference to the section at the specified index, if it exists.
168    ///
169    /// This method retrieves the section at the specified index. If the index is out of bounds,
170    /// `None` is returned.
171    ///
172    /// # Arguments:
173    /// - `index`: The index of the section to retrieve.
174    ///
175    /// # Returns:
176    /// - `Option<&Section>`: `Some(section)` if the section exists, or `None` if the index is out of bounds.
177    ///
178    pub fn get_section(&self, index: usize) -> Option<&Section> {
179        match &self.sections {
180            EncryptedOr::Plain(sections) => sections.get(index),
181            EncryptedOr::Encrypted(_) => panic!("SubImage is encrypted"),
182        }
183    }
184
185    /// Returns a mutable reference to the section at the specified index, if it exists.
186    ///
187    /// This method retrieves the section at the specified index. If the index is out of bounds,
188    /// `None` is returned.
189    ///
190    /// # Arguments:
191    /// - `index`: The index of the section to retrieve.
192    ///
193    /// # Returns:
194    /// - `Option<&mut Section>`: `Some(section)` if the section exists, or `None` if the index is out of bounds.
195    ///
196    pub fn get_section_mut(&mut self, index: usize) -> Option<&mut Section> {
197        match &mut self.sections {
198            EncryptedOr::Plain(sections) => sections.get_mut(index),
199            EncryptedOr::Encrypted(_) => panic!("SubImage is encrypted"),
200        }
201    }
202
203    /// Reads the signature for this `SubImage` from a binary stream and computes its hash.
204    ///
205    /// This function reads the header and segment data of the `SubImage` from the given reader,
206    /// and then computes the hash of the data using the specified hashing algorithm (`algo`).
207    /// Optionally, a key can be provided for use by certain hashing algorithms (e.g., HMAC).
208    ///
209    /// # Arguments:
210    /// - `reader`: A mutable reference to a reader that implements both `io::Read` and `io::Seek`.
211    /// - `algo`: The hash algorithm to use for computing the signature (e.g., SHA-256, HMAC).
212    /// - `key`: An optional key for algorithms that require one (e.g., HMAC). If no key is needed, this can be `None`.
213    ///
214    /// # Returns:
215    /// - `Result<Vec<u8>, Error>`: Returns the computed signature as a `Vec<u8>` if successful,
216    ///   or an `Error` if there is an issue reading the data or computing the hash.
217    pub fn signature_from_stream<R>(
218        &self,
219        reader: &mut R,
220        algo: HashAlgo,
221        key: Option<&[u8]>,
222    ) -> Result<Vec<u8>, Error>
223    where
224        R: io::Read + io::Seek,
225    {
226        let mut buffer = vec![0x00; ImageHeader::binary_size() + self.header.segment_size as usize];
227        reader.read_exact(&mut buffer)?;
228        algo.compute_hash(&buffer, key)
229    }
230}
231
232impl AsImage for SubImage {
233    /// Set the signature for the SubImage.
234    ///
235    /// # Arguments:
236    /// - `signature`: A slice containing the signature to set.
237    fn set_signature(&mut self, signature: &[u8]) {
238        self.hash.copy_from_slice(signature);
239    }
240
241    /// Set the segment size for the SubImage.
242    ///
243    /// # Arguments:
244    /// - `size`: The size to set for the SubImage's segment.
245    fn set_segment_size(&mut self, size: u32) {
246        self.header.segment_size = size;
247        // REVISIT: by default, the partition size is 0
248        // if let EncryptedOr::Plain(fst) = &mut self.fst {
249        //     fst.partition_size = size - FST::binary_size() as u32;
250        // }
251    }
252
253    /// Build the segment size for the SubImage.
254    ///
255    /// # Returns:
256    /// The total segment size, calculated by adding the size of the `ImageHeader`, the `FST`,
257    /// and the aligned sizes of all the sections.
258    fn build_segment_size(&self) -> u32 {
259        // Segment size does not include the hash or image padding
260        FST::binary_size() as u32
261            + match &self.sections {
262                EncryptedOr::Plain(sections) => sections
263                    .iter()
264                    .map(Section::build_aligned_size)
265                    .sum::<u32>(),
266                EncryptedOr::Encrypted(sections_data) => sections_data.len() as u32,
267            }
268    }
269
270    /// Build the signature for the SubImage.
271    ///
272    /// This function generates a signature by calculating the hash of the image's content,
273    /// including the header, firmware security table (FST), and sections.
274    ///
275    /// # Arguments:
276    /// - `key`: A byte slice containing the key used to generate the signature.
277    ///
278    /// # Returns:
279    /// A `Result<Vec<u8>, crate::error::Error>` that contains:
280    /// - `Ok(Vec<u8>)`: The signature as a byte vector.
281    /// - `Err(Error)`: An error if signature calculation fails (e.g., unsupported hash algorithm).
282    fn build_signature(&self, key: Option<&[u8]>) -> Result<Vec<u8>, crate::error::Error> {
283        let hash_algo = match &self.fst {
284            // TODO: what should we use in case of an encrypted FST?
285            EncryptedOr::Encrypted(_) => Some(HashAlgo::Sha256),
286            EncryptedOr::Plain(fst) => fst.hash_algo,
287        };
288
289        let mut buffer = vec![0u8; ImageHeader::binary_size() + self.build_segment_size() as usize];
290        let mut writer = Cursor::new(&mut buffer);
291
292        // Write the header, FST, and sections to the buffer.
293        self.header.write_to(&mut writer)?;
294        self.fst.write_to(&mut writer)?;
295        self.sections.write_to(&mut writer)?;
296
297        // Compute the hash using the FST's hash algorithm.
298        match hash_algo {
299            Some(algo) => Ok(algo.compute_hash(&buffer, key)?.to_vec()),
300            None => Err(Error::NotImplemented(
301                "SubImage::build_signature".to_string(),
302            )),
303        }
304    }
305}
306
307impl FromStream for SubImage {
308    /// Reads a `SubImage` from a binary stream.
309    ///
310    /// # Arguments:
311    /// - `reader`: The stream from which the data will be read. This must implement `std::io::Read` and `std::io::Seek`.
312    fn read_from<R>(&mut self, reader: &mut R) -> Result<(), Error>
313    where
314        R: io::Read + io::Seek,
315    {
316        self.header.read_from(reader)?;
317        if self.header.is_encrypt {
318            self.fst = EncryptedOr::Encrypted(vec![0; FST::binary_size() as usize]);
319            self.fst.read_from(reader)?;
320
321            let mut sections =
322                vec![0; self.header.segment_size as usize - FST::binary_size() as usize];
323            reader.read_exact(&mut sections)?;
324            self.sections = EncryptedOr::Encrypted(sections);
325        } else {
326            self.fst.read_from(reader)?;
327
328            let mut sections = Vec::new();
329            loop {
330                let section: Section = from_stream(reader)?;
331                let has_next = section.header.has_next();
332
333                sections.push(section);
334                if !has_next {
335                    break;
336                }
337            }
338            self.sections = EncryptedOr::Plain(sections);
339        }
340        reader.read_exact(&mut self.hash)?;
341        skip_aligned(reader, if self.header.has_next() { 0x4000 } else { 0x40 })?;
342        Ok(())
343    }
344}
345
346impl ToStream for SubImage {
347    /// Writes a `SubImage` to a binary stream.
348    ///
349    /// # Arguments:
350    /// - `writer`: The stream to which the data will be written. This must implement `std::io::Write` and `std::io::Seek`.
351    ///
352    fn write_to<W>(&self, writer: &mut W) -> Result<(), Error>
353    where
354        W: io::Write + io::Seek,
355    {
356        self.header.write_to(writer)?;
357        self.fst.write_to(writer)?;
358        self.sections.write_to(writer)?;
359        writer.write_all(&self.hash)?;
360
361        let align = if self.header.has_next() { 0x4000 } else { 0x40 };
362        write_aligned!(writer, align, 0x87, optional);
363        Ok(())
364    }
365}
366
367/// Represents an OTA (Over-The-Air) image.
368///
369/// An `OTAImage` is used in the context of firmware updates, where the image consists of
370/// multiple subimages (representing different sections of the firmware), each of which
371/// may be signed and encrypted. The `keyblock` holds encrypted keys, `public_keys` are used for
372/// verifying signatures, and `checksum` ensures data integrity.
373///
374/// **Note**: The encryption and signature verification are currently are using the hash key
375/// specified in the partition table!
376#[derive(Debug)]
377pub struct OTAImage {
378    /// The key block containing cryptographic keys for encryption and signature verification.
379    pub keyblock: KeyBlock,
380
381    /// Public keys (up to 5) used for signature verification.
382    public_keys: [DataType<32>; 5],
383
384    /// A collection of subimages contained in the OTA image.
385    subimages: Vec<SubImage>,
386
387    /// A checksum value for verifying the integrity of the OTA image.
388    pub checksum: Option<u32>,
389}
390
391impl Default for OTAImage {
392    /// Creates a default `OTAImage` with an empty keyblock, no public keys, no subimages, and a checksum of -1.
393    fn default() -> Self {
394        OTAImage {
395            keyblock: KeyBlock::default(),
396            public_keys: [None; 5],
397            subimages: Vec::new(),
398            checksum: None,
399        }
400    }
401}
402
403impl OTAImage {
404    /// Returns a slice of the subimages contained in the OTA image.
405    ///
406    /// # Returns:
407    /// - A reference to a slice containing all the subimages.
408    pub fn get_subimages(&self) -> &[SubImage] {
409        return &self.subimages;
410    }
411
412    /// Returns a mutable slice of the subimages contained in the OTA image.
413    ///
414    /// # Returns:
415    /// - A mutable reference to a slice containing all the subimages.
416    pub fn get_subimages_mut(&mut self) -> &mut [SubImage] {
417        return &mut self.subimages;
418    }
419
420    /// Retrieves a specific subimage by its index.
421    ///
422    /// # Arguments:
423    /// - `index`: The index of the subimage to retrieve.
424    ///
425    /// # Returns:
426    /// - `Some(SubImage)` if the subimage exists at the given index, `None` otherwise.
427    pub fn get_subimage(&self, index: usize) -> Option<&SubImage> {
428        return self.subimages.get(index);
429    }
430
431    /// Retrieves a mutable reference to a specific subimage by its index.
432    ///
433    /// # Arguments:
434    /// - `index`: The index of the subimage to retrieve.
435    ///
436    /// # Returns:
437    /// - `Some(&mut SubImage)` if the subimage exists at the given index, `None` otherwise.
438    pub fn get_subimage_mut(&mut self, index: usize) -> Option<&mut SubImage> {
439        return self.subimages.get_mut(index);
440    }
441
442    /// Adds a new subimage to the OTA image.
443    ///
444    /// # Arguments:
445    /// - `subimage`: The `SubImage` to add to the OTA image.
446    pub fn add_subimage(&mut self, subimage: SubImage) {
447        self.subimages.push(subimage);
448    }
449
450    /// Removes a subimage from the OTA image at the specified index.
451    ///
452    /// # Arguments:
453    /// - `index`: The index of the subimage to remove.
454    pub fn rem_subimage_at(&mut self, index: usize) {
455        self.subimages.remove(index);
456    }
457
458    /// Returns the encryption public key from the keyblock, which is used for OTA signature verification.
459    ///
460    /// # Returns:
461    /// - A reference to the encryption public key (32 bytes) used for signature verification.
462    pub fn get_ota_signature(&self) -> &[u8; 32] {
463        return self.keyblock.get_enc_pubkey();
464    }
465
466    /// Retrieves a specific public key used for signature verification from the OTA image.
467    ///
468    /// # Arguments:
469    /// - `index`: The index (0-4) of the public key to retrieve.
470    ///
471    /// # Returns:
472    /// - A reference to the public key at the specified index, if it exists.
473    pub fn get_public_key(&self, index: u8) -> DataRefType<'_, 32> {
474        return self.public_keys[index as usize].as_ref();
475    }
476
477    pub fn get_public_keys(&self) -> &[DataType<32>; 5] {
478        &self.public_keys
479    }
480}
481
482// cryptographic ops
483impl OTAImage {
484    /// Builds the OTA image signature, which is the hash result of the first SubImage's header.
485    ///
486    /// # Arguments:
487    /// - `key`: An optional key to be used in the hash calculation (may be `None` if no key is provided).
488    ///
489    /// # Returns:
490    /// - `Result<Vec<u8>, crate::error::Error>`: The computed signature as a vector of bytes on success,
491    ///   or an error if the computation cannot be completed (e.g., `fst.hash_algo` is `None`).
492    pub fn build_ota_signature(&self, key: Option<&[u8]>) -> Result<Vec<u8>, crate::error::Error> {
493        let mut buffer = Vec::with_capacity(ImageHeader::binary_size());
494        // according to spec:
495        // OTA signature: The hash result of the 1st Image header “Sub FW Image 0 Header”
496        // which is:
497        // ├────────────────────────────────────────┤
498        // │      SubImage 0 Header (96 bytes)      │
499        // ├────────────────────────────────────────┤
500        //
501        // the key is the hash key from the partition table record for this image
502        if let Some(subimage) = self.get_subimage(0) {
503            if let EncryptedOr::Plain(fst) = &subimage.fst {
504                if let Some(algo) = &fst.hash_algo {
505                    let mut writer = Cursor::new(&mut buffer);
506                    subimage.header.write_to(&mut writer)?;
507                    return algo.compute_hash(&buffer, key);
508                } else {
509                    return Err(Error::NotImplemented(
510                        "OTAImage::build_ota_signature: subimage[0].fst.hash_algo is None"
511                            .to_string(),
512                    ));
513                }
514            }
515            return Err(Error::NotImplemented(
516                "OTAImage::build_ota_signature: subimage[0].fst is encrypted".to_string(),
517            ));
518        }
519
520        Err(Error::NotImplemented(
521            "OTAImage::build_ota_signature: subimage[0] not found".to_string(),
522        ))
523    }
524
525    pub fn build_first_signature(
526        &self,
527        key: Option<&[u8]>,
528    ) -> Result<Vec<u8>, crate::error::Error> {
529        // Hash: calculated with Encrypted FW image if image encryption is on
530        // The 1st sub-image
531        //  From OTA Signature to the last image section, including all padding bytes
532        let ota_signature = self.build_ota_signature(key)?;
533        if let Some(subimage) = self.get_subimage(0) {
534            if let EncryptedOr::Plain(fst) = &subimage.fst {
535                if let Some(algo) = &fst.hash_algo {
536                    // OTA signature + public keys
537                    let segment_size = subimage.build_segment_size() as usize;
538                    let mut buffer = vec![
539                        0u8;
540                        ImageHeader::binary_size()
541                            + segment_size as usize
542                            + 224
543                    ];
544                    let mut writer = Cursor::new(&mut buffer);
545
546                    writer.write_all(&ota_signature)?;
547                    writer.write_all(self.keyblock.get_hash_pubkey())?;
548                    for pubkey in &self.public_keys {
549                        match pubkey {
550                            Some(raw_key) => writer.write_all(raw_key)?,
551                            None => writer.write_all(&[0xFF; 32])?,
552                        };
553                    }
554                    // signature does not include the hash itself
555                    subimage.header.write_to(&mut writer)?;
556                    subimage.fst.write_to(&mut writer)?;
557                    subimage.sections.write_to(&mut writer)?;
558                    // build hash without existing subimage hash
559                    return algo.compute_hash(&buffer, key);
560                } else {
561                    return Err(Error::NotImplemented(
562                        "OTAImage::build_ota_signature: subimage[0].fst.hash_algo is None"
563                            .to_string(),
564                    ));
565                }
566            }
567            return Err(Error::NotImplemented(
568                "OTAImage::build_ota_signature: subimage[0].fst is encrypted".to_string(),
569            ));
570        }
571
572        Err(Error::NotImplemented(
573            "OTAImage::build_ota_signature: subimage[0] not found".to_string(),
574        ))
575    }
576
577    /// Reads the OTA signature from a stream and computes its hash using a specified algorithm.
578    ///
579    /// This function reads the `OTAImage` signature data from the provided reader, computes
580    /// its hash using the specified `HashAlgo`, and returns the computed signature.
581    ///
582    /// The function assumes that the data read corresponds to the "OTA signature" section of
583    /// the `OTAImage` format, which is typically the first part of the image.
584    ///
585    /// # Arguments:
586    /// - `reader`: A mutable reference to a reader that implements `io::Read` and `io::Seek`.
587    ///   This will be used to read the OTA signature data.
588    /// - `algo`: The hash algorithm to use for computing the signature (e.g., SHA-256).
589    /// - `key`: An optional key to be used by certain hash algorithms (e.g., for HMAC).
590    ///   If the algorithm does not require a key, this can be `None`.
591    ///
592    /// # Returns:
593    /// - `Result<Vec<u8>, crate::error::Error>`: Returns the computed signature as a `Vec<u8>`,
594    ///   or an error if there is an issue reading the data or computing the hash.
595    pub fn ota_signature_from_stream<R>(
596        reader: &mut R,
597        algo: HashAlgo,
598        key: Option<&[u8]>,
599    ) -> Result<Vec<u8>, crate::error::Error>
600    where
601        R: io::Read + io::Seek,
602    {
603        let mut buffer = vec![0x00; ImageHeader::binary_size()];
604        reader.read_exact(&mut buffer)?;
605        algo.compute_hash(&buffer, key)
606    }
607
608    /// Sets the OTA image signature, specifically the public encryption key in the keyblock.
609    ///
610    /// # Arguments:
611    /// - `signature`: The signature (encryption public key) to set, which will replace the existing public key.
612    pub fn set_ota_signature(&mut self, signature: &[u8]) {
613        self.keyblock
614            .get_enc_pubkey_mut()
615            .copy_from_slice(signature);
616    }
617
618    /// Computes the checksum for the OTA image by writing it to a buffer.
619    ///
620    /// This method serializes the current `OTAImage` object into a byte buffer and then calculates
621    /// the checksum for the serialized data. The checksum is returned as a 32-bit unsigned integer.
622    ///
623    /// # Errors
624    ///
625    /// This function returns an error if the writing process to the buffer fails.
626    pub fn build_checksum(&self) -> Result<u32, Error> {
627        let mut buffer = Vec::new();
628        let mut cursor = std::io::Cursor::new(&mut buffer);
629
630        self.write_to(&mut cursor)?;
631        Ok(OTAImage::checksum_from_buffer(&buffer))
632    }
633
634    /// Updates the checksum field of the OTA image.
635    ///
636    /// This method resets the `checksum` field to `None` and then calculates and sets the new checksum
637    /// by calling `build_checksum`. It ensures that the checksum field is always up-to-date.
638    ///
639    /// # Example
640    ///
641    /// ```
642    /// let mut ota_image = /* ... */;
643    /// if let Err(e) = ota_image.update_checksum() {
644    ///     eprintln!("Failed to update checksum: {}", e);
645    /// }
646    /// ```
647    pub fn update_checksum(&mut self) -> Result<(), Error> {
648        self.checksum = None;
649        self.checksum = Some(self.build_checksum()?);
650        Ok(())
651    }
652
653    /// Updates the OTA image signature using the provided public key.
654    ///
655    /// This method generates a new OTA signature and updates the `keyblock` field with the
656    /// signature. If a key is provided, it will be used in the signing process; otherwise,
657    /// the default behavior is applied.
658    ///
659    /// # Arguments
660    ///
661    /// * `key` - An optional reference to a byte slice (`&[u8]`) representing the public key.
662    ///
663    /// # Errors
664    ///
665    /// Returns an error if the signing process fails.
666    pub fn update_ota_signature(&mut self, key: Option<&[u8]>) -> Result<(), Error> {
667        let new_signature = OTAImage::build_ota_signature(self, key)?;
668        self.keyblock
669            .get_enc_pubkey_mut()
670            .copy_from_slice(&new_signature);
671        Ok(())
672    }
673
674    /// Calculates a checksum from a byte buffer by summing all the byte values and applying a bitmask.
675    ///
676    /// # Arguments:
677    /// - `buf`: The byte buffer to compute the checksum from.
678    ///
679    /// # Returns:
680    /// - `i32`: The computed checksum as a 32-bit signed integer.
681    pub fn checksum_from_buffer(buf: &[u8]) -> u32 {
682        buf.iter().map(|&byte| byte as u32).sum::<u32>()
683    }
684
685    /// Calculates a checksum from a stream by reading the content into a buffer and computing its checksum.
686    ///
687    /// # Arguments:
688    /// - `reader`: A reader that implements `io::Read + io::Seek` from which the content will be read.
689    ///
690    /// # Returns:
691    /// - `Result<i32, Error>`: The checksum computed from the stream as a 32-bit signed integer, or an error if the reading fails.
692    pub fn checksum_from_stream<R>(reader: &mut R) -> Result<u32, Error>
693    where
694        R: io::Read + io::Seek,
695    {
696        let mut buffer = Vec::new();
697        // we assume this reader is at pos 0
698        reader.read_to_end(&mut buffer)?;
699        Ok(OTAImage::checksum_from_buffer(&buffer[..&buffer.len() - 4]))
700    }
701}
702
703impl FromStream for OTAImage {
704    /// Reads an `OTAImage` from a binary stream.
705    ///
706    /// This function assumes that the provided reader is positioned correctly and
707    /// that the stream contains the expected data format for the `OTAImage` struct.
708    ///
709    /// # Arguments:
710    /// - `reader`: A mutable reference to a reader that implements both `io::Read` and `io::Seek`.
711    ///
712    /// # Returns:
713    /// - `Result<(), Error>`: Returns `Ok(())` if the data is read and parsed successfully, or an `Error`
714    ///   if something goes wrong (e.g., invalid format, stream read errors).
715    fn read_from<R>(&mut self, reader: &mut R) -> Result<(), Error>
716    where
717        R: io::Read + io::Seek,
718    {
719        self.keyblock.read_from(reader)?;
720
721        // Read 5 public keys, validate each, and store the valid ones.
722        for i in 0..5 {
723            let mut key = [0x00; 32];
724            reader.read_exact(&mut key)?;
725            if is_valid_data!(&key) {
726                self.public_keys[i] = Some(key);
727            }
728        }
729
730        loop {
731            let subimage: SubImage = from_stream(reader)?;
732            let has_next = subimage.header.has_next();
733            self.subimages.push(subimage);
734
735            // If there is no next subimage, break out of the loop.
736            if !has_next {
737                break;
738            }
739        }
740
741        let checksum = reader.read_u32::<LittleEndian>()?;
742        self.checksum = match checksum {
743            0xFFFF_FFFF | 0x1A1A_1A1A => None,
744            v => Some(v),
745        };
746        Ok(())
747    }
748}
749
750impl ToStream for OTAImage {
751    /// Writes an `OTAImage` to a binary stream.
752    ///
753    /// # Arguments:
754    /// - `writer`: A mutable reference to a writer that implements both `io::Write` and `io::Seek`.
755    fn write_to<W>(&self, writer: &mut W) -> Result<(), Error>
756    where
757        W: io::Write + io::Seek,
758    {
759        self.keyblock.write_to(writer)?;
760        for key in &self.public_keys {
761            write_data!(writer, key, 32);
762        }
763        for subimage in &self.subimages {
764            subimage.write_to(writer)?;
765        }
766        if let Some(checksum) = self.checksum {
767            writer.write_u32::<LittleEndian>(checksum)?;
768        }
769        Ok(())
770    }
771}