amebazii/types/image/
mod.rs

1use super::{FromStream, ToStream};
2use crate::error::Error;
3use std::io;
4
5pub mod boot;
6pub use boot::BootImage;
7
8pub mod ota;
9pub use ota::{OTAImage, SubImage};
10
11pub mod pt;
12pub use pt::{PartTab, PartitionTableImage, Record, TrapConfig};
13
14pub type RawImage = Vec<u8>;
15
16/// A generic enum representing either encrypted or plain data.
17///
18/// The `EncryptedOr` enum is used to differentiate between encrypted data and unencrypted (plain) data.
19/// It allows to store either type of data in the same structure while providing methods to safely access
20/// or mutate the contents, depending on whether the data is encrypted or not.
21#[derive(Debug)]
22pub enum EncryptedOr<T> {
23    /// Contains encrypted data as a vector of bytes.
24    Encrypted(Vec<u8>),
25
26    /// Contains plain, unencrypted data of type `T`.
27    Plain(T),
28}
29
30impl<T> AsRef<T> for EncryptedOr<T> {
31    /// Returns a reference to the plain data if available.
32    ///
33    /// # Panics
34    /// Panics if the data is encrypted, as the method expects plain data.
35    fn as_ref(&self) -> &T {
36        match self {
37            EncryptedOr::Encrypted(_) => {
38                panic!("Cannot get reference to encrypted data when plain is expected")
39            }
40            EncryptedOr::Plain(t) => t,
41        }
42    }
43}
44
45impl<T> AsRef<[u8]> for EncryptedOr<T> {
46    /// Returns a reference to the encrypted data if available.
47    ///
48    /// # Panics
49    /// Panics if the data is plain, as the method expects encrypted data.
50    fn as_ref(&self) -> &[u8] {
51        match self {
52            EncryptedOr::Encrypted(v) => v,
53            EncryptedOr::Plain(_) => {
54                panic!("Cannot get reference to plain data when encrypted is expected")
55            }
56        }
57    }
58}
59
60impl<T> EncryptedOr<T> {
61    /// Returns `true` if the data is encrypted.
62    ///
63    /// This method allows checking if the current instance of `EncryptedOr` contains encrypted data.
64    pub fn is_encrypted(&self) -> bool {
65        match self {
66            EncryptedOr::Encrypted(_) => true,
67            EncryptedOr::Plain(_) => false,
68        }
69    }
70
71    /// Returns `true` if the data is plain.
72    ///
73    /// This method allows checking if the current instance of `EncryptedOr` contains plain (unencrypted) data.
74    pub fn is_plain(&self) -> bool {
75        match self {
76            EncryptedOr::Encrypted(_) => false,
77            EncryptedOr::Plain(_) => true,
78        }
79    }
80
81    pub fn unwrap(self) -> T {
82        match self {
83            EncryptedOr::Encrypted(_) => panic!("Attempted to unwrap encrypted data"),
84            EncryptedOr::Plain(t) => t,
85        }
86    }
87}
88
89impl<T> AsMut<T> for EncryptedOr<T> {
90    /// Returns a mutable reference to the plain data if available.
91    ///
92    /// # Panics
93    /// Panics if the data is encrypted, as the method expects plain data.
94    fn as_mut(&mut self) -> &mut T {
95        match self {
96            EncryptedOr::Encrypted(_) => {
97                panic!("Cannot get mutable reference to encrypted data when plain is expected")
98            }
99            EncryptedOr::Plain(t) => t,
100        }
101    }
102}
103
104impl<T> AsMut<[u8]> for EncryptedOr<T> {
105    /// Returns a mutable reference to the encrypted data if available.
106    ///
107    /// # Panics
108    /// Panics if the data is plain, as the method expects encrypted data.
109    fn as_mut(&mut self) -> &mut [u8] {
110        match self {
111            EncryptedOr::Encrypted(v) => v,
112            EncryptedOr::Plain(_) => {
113                panic!("Cannot get mutable reference to encrypted data when plain is expected")
114            }
115        }
116    }
117}
118
119impl<T: ToStream> ToStream for EncryptedOr<T> {
120    /// Writes the data to a stream, either encrypted or plain.
121    ///
122    /// This method is used to serialize the data into a stream. If the data is encrypted, it writes the encrypted byte vector,
123    /// otherwise it serializes the plain data of type `T`.
124    fn write_to<W>(&self, writer: &mut W) -> Result<(), Error>
125    where
126        W: io::Write + io::Seek,
127    {
128        match self {
129            EncryptedOr::Encrypted(v) => writer.write_all(v)?,
130            EncryptedOr::Plain(t) => return t.write_to(writer),
131        }
132        Ok(())
133    }
134}
135
136impl<T: ToStream> ToStream for EncryptedOr<Vec<T>> {
137    fn write_to<W>(&self, writer: &mut W) -> Result<(), Error>
138    where
139        W: io::Write + io::Seek,
140    {
141        match self {
142            EncryptedOr::Encrypted(v) => writer.write_all(v)?,
143            EncryptedOr::Plain(t) => {
144                for item in t {
145                    item.write_to(writer)?;
146                }
147            }
148        }
149        Ok(())
150    }
151}
152
153impl<T: FromStream> FromStream for EncryptedOr<T> {
154    /// Reads the data from a stream, either encrypted or plain.
155    ///
156    /// This method deserializes the data from a stream. If the data is encrypted, it reads the encrypted byte vector,
157    /// otherwise it reads the plain data of type `T`.
158    fn read_from<R>(&mut self, reader: &mut R) -> Result<(), Error>
159    where
160        R: io::Read + io::Seek,
161    {
162        match self {
163            EncryptedOr::Encrypted(v) => reader.read_exact(v)?,
164            EncryptedOr::Plain(t) => return t.read_from(reader),
165        }
166        Ok(())
167    }
168}
169
170/// A trait that provides common functionality for image-like objects,
171/// such as computing and setting the segment size and signature.
172pub trait AsImage {
173    /// Computes the segment size for the image.
174    ///
175    /// The segment size typically represents the total size of the image segment,
176    /// including all of its components (e.g., header, records, user data, etc.).
177    ///
178    /// # Returns:
179    /// - `u32` representing the segment size.
180    ///
181    /// # Example:
182    /// ```rust
183    /// let segment_size = image.build_segment_size();
184    /// ```
185    fn build_segment_size(&self) -> u32;
186
187    /// Sets the segment size for the image.
188    ///
189    /// This method allows setting the segment size, typically after calculating
190    /// it or modifying the image in some way.
191    ///
192    /// # Arguments:
193    /// - `size`: The new segment size to set.
194    ///
195    /// # Example:
196    /// ```rust
197    /// image.set_segment_size(1024);
198    /// ```
199    fn set_segment_size(&mut self, size: u32);
200
201    /// Computes the signature for the image using the provided key.
202    ///
203    /// The signature is usually a hash or HMAC generated from the image data and
204    /// a secret key, often used for verification or authentication purposes.
205    ///
206    /// # Arguments:
207    /// - `key`: The key used to compute the signature.
208    ///
209    /// # Returns:
210    /// - `Result<Vec<u8>, crate::error::Error>`: The signature as a `Vec<u8>`, or an error.
211    ///
212    /// # Example:
213    /// ```rust
214    /// let signature = image.build_signature(&key);
215    /// ```
216    fn build_signature(&self, key: Option<&[u8]>) -> Result<Vec<u8>, crate::error::Error>;
217
218    /// Sets the signature for the image.
219    ///
220    /// This method allows setting the signature after computing it or for some
221    /// validation process.
222    ///
223    /// # Arguments:
224    /// - `signature`: The computed signature to set.
225    ///
226    /// # Example:
227    /// ```rust
228    /// image.set_signature(&signature);
229    /// ```
230    fn set_signature(&mut self, signature: &[u8]);
231}
232
233/// Builds the signature for a given image.
234///
235/// This function uses the `build_signature` method from the `AsImage` trait to generate
236/// the signature for the image using the provided key.
237///
238/// # Arguments:
239/// - `image`: The image-like object that implements `AsImage`.
240/// - `key`: The key used to compute the signature.
241///
242/// # Returns:
243/// - `Result<Vec<u8>, crate::error::Error>`: The computed signature.
244pub fn build_default_signature<I>(
245    image: &I,
246    key: Option<&[u8]>,
247) -> Result<Vec<u8>, crate::error::Error>
248where
249    I: AsImage,
250{
251    image.build_signature(key)
252}
253
254/// Sets the signature for a given image.
255///
256/// This function computes the signature using `build_default_signature` and then sets
257/// the signature for the image using `set_signature`.
258///
259/// # Arguments:
260/// - `image`: The image-like object that implements `AsImage`.
261/// - `key`: The key used to compute the signature.
262///
263/// # Returns:
264/// - `Result<(), crate::error::Error>`: An empty result on success, or an error.
265pub fn set_default_signature<I>(
266    image: &mut I,
267    key: Option<&[u8]>,
268) -> Result<(), crate::error::Error>
269where
270    I: AsImage,
271{
272    let signature = build_default_signature(image, key)?;
273    image.set_signature(&signature);
274    Ok(())
275}
276
277/// Builds the segment size for a given image.
278///
279/// # Arguments:
280/// - `image`: The image-like object that implements `AsImage`.
281///
282/// # Returns:
283/// - `u32`: The computed segment size.
284pub fn build_segment_size<I>(image: &I) -> u32
285where
286    I: AsImage,
287{
288    image.build_segment_size()
289}
290
291/// Sets the segment size for a given image.
292///
293/// # Arguments:
294/// - `image`: The image-like object that implements `AsImage`.
295///
296/// # Returns:
297/// - `()`: An empty result on success.
298pub fn set_default_segment_size<I>(image: &mut I)
299where
300    I: AsImage,
301{
302    image.set_segment_size(image.build_segment_size())
303}
amebazii/types/image/mod.rs

amebazii/types/image/
mod.rs
