amebazii/types/image/pt.rs
1use std::io::{self, Cursor};
2
3use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
4use serde::{Deserialize, Serialize};
5
6use crate::{
7 error::Error,
8 is_valid_data, read_padding, read_valid_data,
9 types::{
10 enums::{KeyExportOp, PartitionType},
11 from_stream,
12 header::{ImageHeader, KeyBlock},
13 BinarySize, DataRefType, DataType, FromStream, ToStream,
14 },
15 util::{hmac_sha256, write_fill},
16 write_data, write_padding,
17};
18
19use super::{AsImage, EncryptedOr};
20
21/// Represents the configuration of a hardware trap.
22///
23/// A `TrapConfig` structure holds information related to a trap configuration, including
24/// the pin, port, level, and validity of the trap. The fields are packed into a 16-bit
25/// integer with specific bitwise encoding. The layout of the 16-bit value is as follows:
26///
27/// ```text
28/// Layout (16-bit integer):
29/// 0 8 15
30/// 0 1 2 3 4 5 6 7 0 1 2 3 4 5 6 7
31/// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
32/// | pin | port|l| |v|
33/// +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
34/// ```
35///
36/// Where:
37/// - `v` (bit 15) represents the validity of the trap configuration (`1` for valid, `0` for invalid).
38/// - `l` (bit 8) represents the level of the trap (`0` or `1`).
39/// - `port` (bits 5 to 7) represents the port number (3 bits, value range from 0 to 7).
40/// - `pin` (bits 0 to 4) represents the pin number (5 bits, value range from 0 to 31).
41///
42/// # Type Implementation:
43/// - `TrapConfig` can be created from a 16-bit integer using the `From<u16>` trait, and
44/// it can be converted back to a 16-bit integer using the `Into<u16>` trait.
45///
46/// The packing and unpacking logic allows easy conversion between the `TrapConfig` struct
47/// and a single 16-bit integer, which is useful for hardware register manipulation.
48#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
49pub struct TrapConfig {
50 /// Whether the trap configuration is valid (1 bit).
51 #[serde(default)]
52 pub valid: bool,
53
54 /// The level of the trap (1 bit, 0 or 1).
55 #[serde(default)]
56 pub level: u8,
57
58 /// The port number (3 bits, value range: 0-7).
59 #[serde(default)]
60 pub port: u8,
61
62 /// The pin number (5 bits, value range: 0-31).
63 #[serde(default)]
64 pub pin: u8,
65}
66
67impl Default for TrapConfig {
68 /// Returns a default `TrapConfig` with all fields set to 0 or `false`.
69 fn default() -> Self {
70 TrapConfig {
71 valid: false,
72 level: 0,
73 port: 0,
74 pin: 0,
75 }
76 }
77}
78
79impl From<u16> for TrapConfig {
80 /// Converts a 16-bit integer to a `TrapConfig` by unpacking the respective bits
81 /// according to the layout described above.
82 ///
83 /// # Arguments:
84 /// - `value`: The 16-bit integer to convert into a `TrapConfig`.
85 ///
86 /// # Returns:
87 /// A `TrapConfig` where each field is populated based on the respective bits
88 /// in the 16-bit integer.
89 fn from(value: u16) -> Self {
90 TrapConfig {
91 valid: (value >> 15) & 0x1 != 0,
92 level: ((value >> 8) & 0x1) as u8,
93 port: ((value >> 5) & 7) as u8,
94 pin: (value & 0x1F) as u8,
95 }
96 }
97}
98
99impl Into<u16> for TrapConfig {
100 /// Converts a `TrapConfig` back into a 16-bit integer by packing the fields into
101 /// their respective bit positions.
102 ///
103 /// # Returns:
104 /// A 16-bit integer representing the packed values of the `TrapConfig` fields.
105 fn into(self) -> u16 {
106 ((self.valid as u16) << 15)
107 | ((self.level as u16) << 8)
108 | ((self.port as u16) << 5)
109 | self.pin as u16
110 }
111}
112
113/// Represents a firmware partition record.
114///
115/// A `Record` struct encapsulates information about a partition within a firmware image.
116/// This record stores a hash key for the partition, which will be later used to verify
117/// the signatures of the partition contents.
118///
119/// # Layout (64 bytes):
120/// ```text
121/// +---------------+---+---+---+---+---+---+---+----------+--------------+----+----+----+----+----+----+
122/// | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 |
123/// +========+===============+===+===+===+===+===+===+===+==========+==============+====+====+====+====+====+====+
124/// | 0x00 | start_addr: u32 | length: u32 | type: u8 | dbg_skip: u8 | |
125/// +--------+---------------------------+---------------+----------+--------------+-----------------------------+
126/// | 0x10 | key_valid: u8 | |
127/// +--------+---------------+-----------------------------------------------------------------------------------+
128/// | 0x20 | hash_key: bytes[32] |
129/// +--------+---------------------------------------------------------------------------------------------------+
130/// ```
131#[derive(Debug)]
132pub struct Record {
133 /// The starting address of the partition in the firmware image (4 bytes).
134 pub start_addr: u32,
135
136 /// The length of the partition in bytes (4 bytes).
137 pub length: u32,
138
139 /// The partition type (1 byte). This is an enum (`PartitionType`).
140 pub part_type: PartitionType,
141
142 /// A flag that indicates whether debugging should be skipped for this partition (1 byte).
143 /// `true` means skip debugging, `false` means do not skip.
144 pub dbg_skip: bool,
145
146 /// A 32-byte hash key associated with this partition. By default, it's invalid
147 hash_key: DataType<32>,
148}
149
150impl BinarySize for Record {
151 /// Returns the size of the `Record` structure in bytes (64 bytes).
152 #[inline]
153 fn binary_size() -> usize {
154 0x40
155 }
156}
157
158impl Default for Record {
159 /// Returns a default `Record` with zeroed and invalid fields.
160 ///
161 /// The default values are as follows:
162 /// - `start_addr`: `0`
163 /// - `length`: `0`
164 /// - `part_type`: `PartitionType::PartTab` (default type)
165 /// - `dbg_skip`: `false`
166 /// - `hash_key`: `None` (invalid hash key)
167 fn default() -> Self {
168 Record {
169 start_addr: 0,
170 length: 0,
171 part_type: PartitionType::PartTab, // Default to PartitionTab type
172 dbg_skip: false,
173 hash_key: None, // Invalid hash key by default
174 }
175 }
176}
177
178impl Record {
179 // ------------------------------------------------------------------------------------
180 // instance methods
181 // ------------------------------------------------------------------------------------
182
183 /// Checks whether the `hash_key` is valid.
184 ///
185 /// # Returns:
186 /// - `true` if the `hash_key` is valid (non-`None` and passes the validation check).
187 /// - `false` otherwise (e.g., if the key is `None` or invalid).
188 pub fn hash_key_valid(&self) -> bool {
189 match &self.hash_key {
190 None => false,
191 Some(key) => is_valid_data!(key), // checks key validity using the macro
192 }
193 }
194
195 /// Returns a reference to the `hash_key`.
196 pub fn get_hash_key(&self) -> DataRefType<'_, 32> {
197 self.hash_key.as_ref()
198 }
199
200 /// Sets the `hash_key` to a new value.
201 ///
202 /// # Arguments:
203 /// - `key`: A reference to a 32-byte array slice that represents the new `hash_key`.
204 /// If `None` is passed, the `hash_key` will be cleared.
205 pub fn set_hash_key(&mut self, key: DataType<32>) {
206 self.hash_key = key;
207 }
208}
209
210impl FromStream for Record {
211 /// Parses a `Record` from a binary stream.
212 ///
213 /// # Arguments:
214 /// - `reader`: A mutable reference to a reader that implements `std::io::Read` and `std::io::Seek`.
215 ///
216 /// # Returns:
217 /// - `Ok(())` if the record was successfully parsed.
218 /// - `Err(Error)` if there was an issue reading from the stream.
219 fn read_from<R>(&mut self, reader: &mut R) -> Result<(), Error>
220 where
221 R: io::Read + io::Seek,
222 {
223 self.start_addr = reader.read_u32::<LittleEndian>()?;
224 self.length = reader.read_u32::<LittleEndian>()?;
225
226 self.part_type = PartitionType::try_from(reader.read_u8()?)?;
227 self.dbg_skip = reader.read_u8()? != 0;
228
229 // Skip 6 bytes of padding.
230 read_padding!(reader, 6);
231
232 // Check if the hash_key is valid (using a specific flag).
233 if reader.read_u8()? & 0x1 != 0 {
234 read_padding!(reader, 15);
235 read_valid_data!(self.hash_key, 32, reader);
236 } else {
237 // Skip 47 bytes if the hash_key is not valid
238 read_padding!(reader, 47);
239 }
240 Ok(())
241 }
242}
243
244impl ToStream for Record {
245 /// Writes a `Record` to a binary stream.
246 ///
247 /// # Arguments:
248 /// - `writer`: A mutable reference to a writer that implements `std::io::Write`.
249 fn write_to<W>(&self, writer: &mut W) -> Result<(), Error>
250 where
251 W: io::Write,
252 {
253 writer.write_u32::<LittleEndian>(self.start_addr)?;
254 writer.write_u32::<LittleEndian>(self.length)?;
255
256 writer.write_u8(self.part_type as u8)?;
257 writer.write_u8(self.dbg_skip as u8)?;
258
259 write_padding!(writer, 6);
260 writer.write_u8(self.hash_key_valid() as u8)?;
261 write_padding!(writer, 15);
262 write_data!(writer, self.hash_key, 32);
263 Ok(())
264 }
265}
266
267/// =====================================================================================
268/// Partition Table (PartTab)
269/// =====================================================================================
270///
271/// The `PartTab` struct represents a partition table for the flash, containing various
272/// metadata and configuration related to the partitioning, as well as firmware-specific
273/// data such as the state of firmware updates, trap configurations, and key export
274/// operations.
275///
276/// # Layout:
277/// ```text
278/// +-----------------+------------------+---+----------+---------+-------------+-------------+---+---+---+-------+-------+-------+------+----+----------------+
279/// | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 | 8 | 9 | 10 | 11 | 12 | 13 | 14 | 15 |
280/// +========+=================+==================+===+==========+=========+=============+=============+===+===+===+=======+=======+=======+======+====+================+
281/// | 0x00 | rma_w_state: u8 | rma_ov_state: u8 | | eFWV: u8 | num: u8 | fw1_idx: u8 | fw2_idx: u8 | | ota_trap: u16 | mp_trap: u16 | | key_exp_op: u8 |
282/// +--------+-----------------+------------------+---+----------+---------+-------------+-------------+-----------+---------------+--------------+----+----------------+
283/// | 0x10 | user_len: u32 | user_ext: bytes[12] |
284/// +--------+---------------------------------------------------+------------------------------------------------------------------------------------------------------+
285/// | 0x20 | records: Record * num |
286/// +--------+----------------------------------------------------------------------------------------------------------------------------------------------------------+
287/// | 0x30 | |
288/// +--------+----------------------------------------------------------------------------------------------------------------------------------------------------------+
289/// | 0x40 | |
290/// +--------+----------------------------------------------------------------------------------------------------------------------------------------------------------+
291/// | 0x50 | |
292/// +--------+----------------------------------------------------------------------------------------------------------------------------------------------------------+
293/// | 0x60 | user_bin: bytes[user_len] |
294/// +--------+----------------------------------------------------------------------------------------------------------------------------------------------------------+
295/// ```
296#[derive(Debug)]
297#[allow(non_snake_case)]
298pub struct PartTab {
299 pub rma_w_state: u8,
300 pub rma_ov_state: u8,
301 pub eFWV: u8,
302 pub fw1_idx: u8,
303 pub fw2_idx: u8,
304 pub ota_trap: TrapConfig,
305 pub mp_trap: TrapConfig,
306 pub key_exp_op: KeyExportOp,
307 user_ext: [u8; 12],
308 records: Vec<Record>,
309 user_bin: Vec<u8>,
310}
311
312impl Default for PartTab {
313 fn default() -> Self {
314 PartTab {
315 rma_w_state: 0xFF,
316 rma_ov_state: 0xFF,
317 eFWV: 0,
318 fw1_idx: 0,
319 fw2_idx: 0,
320 ota_trap: TrapConfig::default(),
321 mp_trap: TrapConfig::default(),
322 key_exp_op: KeyExportOp::None,
323 user_ext: [0xFF; 12],
324 records: Vec::new(),
325 user_bin: Vec::new(),
326 }
327 }
328}
329
330impl PartTab {
331 /// Returns the records in the partition table.
332 ///
333 /// # Returns:
334 /// - A slice of `Record` structs.
335 pub fn get_records(&self) -> &[Record] {
336 return &self.records;
337 }
338
339 /// Returns the user binary data.
340 ///
341 /// This method provides access to the raw user binary data in the partition table.
342 ///
343 /// # Returns:
344 /// - A slice of bytes representing the user binary data.
345 pub fn get_user_bin(&self) -> &[u8] {
346 return &self.user_bin;
347 }
348
349 /// Returns the user extension data (12 bytes).
350 ///
351 /// This method provides access to the 12-byte user extension field, which can be used
352 /// for storing additional metadata related to the partition table.
353 ///
354 /// # Returns:
355 /// - A reference to the 12-byte array representing the user extension data.
356 pub fn get_user_ext(&self) -> &[u8] {
357 return &self.user_ext;
358 }
359
360 /// Sets the user binary data in the partition table.
361 ///
362 /// # Arguments:
363 /// - `user_bin`: A slice of bytes representing the new user binary data to store.
364 pub fn set_user_bin(&mut self, user_bin: &[u8]) {
365 self.user_bin.extend_from_slice(user_bin);
366 }
367
368 /// Sets the user extension data in the partition table.
369 ///
370 /// # Arguments:
371 /// - `user_ext`: A slice of bytes representing the new user extension data (12 bytes).
372 pub fn set_user_ext(&mut self, user_ext: &[u8]) {
373 self.user_ext.copy_from_slice(user_ext);
374 }
375
376 /// Adds a new partition record to the partition table.
377 ///
378 /// # Arguments:
379 /// - `record`: The `Record` struct that defines the new partition entry to add.
380 pub fn add_record(&mut self, record: Record) {
381 self.records.push(record);
382 }
383
384 /// Creates and returns a new record for a specific partition type.
385 ///
386 /// # Arguments:
387 /// - `part_type`: The `PartitionType` enum value representing the type of the new record.
388 ///
389 /// # Returns:
390 /// - A mutable reference to the newly created record.
391 pub fn new_record(&mut self, part_type: PartitionType) -> &mut Record {
392 let mut r = Record::default();
393 r.part_type = part_type;
394 self.records.push(r);
395 return self.records.last_mut().unwrap();
396 }
397
398 /// Returns the record for a specific partition type.
399 ///
400 /// # Arguments:
401 /// - `part_type`: The `PartitionType` enum value representing the partition type.
402 ///
403 /// # Returns:
404 /// - `Some(&Record)` if the record for the given partition type is found.
405 /// - `None` if the record for the given partition type is not found.
406 pub fn get_record(&self, part_type: PartitionType) -> Option<&Record> {
407 return self.records.iter().find(|r| r.part_type == part_type);
408 }
409
410 /// Returns the record for a specific partition type.
411 ///
412 /// # Arguments:
413 /// - `part_type`: The `PartitionType` enum value representing the partition type.
414 ///
415 /// # Returns:
416 /// - `Some(&mut Record)` if the record for the given partition type is found.
417 /// - `None` if the record for the given partition type is not found.
418 pub fn get_record_mut(&mut self, part_type: PartitionType) -> Option<&mut Record> {
419 return self.records.iter_mut().find(|r| r.part_type == part_type);
420 }
421
422 /// Removes a partition record from the partition table.
423 ///
424 /// # Arguments:
425 /// - `part_type`: The `PartitionType` enum value representing the partition type to remove.
426 pub fn rem_record(&mut self, part_type: PartitionType) {
427 self.records.retain(|r| r.part_type != part_type);
428 }
429
430 /// Checks if a record exists for a specific partition type.
431 ///
432 /// # Arguments:
433 /// - `part_type`: The `PartitionType` enum value representing the partition type to check.
434 ///
435 /// # Returns:
436 /// - `true` if a record exists for the given partition type.
437 /// - `false` if no record exists for the given partition type.
438 pub fn has_record(&self, part_type: PartitionType) -> bool {
439 self.get_record(part_type).is_some()
440 }
441}
442
443impl FromStream for PartTab {
444 /// Parses a `PartTab` from a binary stream.
445 ///
446 /// # Returns:
447 /// - `Ok(())` if the `PartTab` was successfully parsed.
448 /// - `Err(Error)` if there was an issue reading from the stream.
449 ///
450 fn read_from<R>(&mut self, reader: &mut R) -> Result<(), Error>
451 where
452 R: io::Read + io::Seek,
453 {
454 self.rma_w_state = reader.read_u8()?;
455 self.rma_ov_state = reader.read_u8()?;
456 self.eFWV = reader.read_u8()?;
457 read_padding!(reader, 1);
458
459 let num = reader.read_u8()? as u32;
460 self.fw1_idx = reader.read_u8()?;
461 self.fw2_idx = reader.read_u8()?;
462 read_padding!(reader, 3);
463
464 self.ota_trap = TrapConfig::from(reader.read_u16::<LittleEndian>()?);
465 self.mp_trap = TrapConfig::from(reader.read_u16::<LittleEndian>()?);
466
467 // Skip the byte set to 0xFF manually in generate_pt_table()
468 read_padding!(reader, 1);
469 self.key_exp_op = KeyExportOp::try_from(reader.read_u8()?)?;
470
471 let mut user_len = reader.read_u32::<LittleEndian>()?;
472 reader.read_exact(&mut self.user_ext)?;
473
474 // Read the partition records (num + 1, including boot record).
475 for _ in 0..=num {
476 self.records.push(from_stream(reader)?);
477 }
478
479 // See #1 for details. Even though we parse the user data here,
480 // we don't verify its length. This causes issues with malformed
481 // partition tables.
482 if user_len > 0x100 {
483 // REVISIT: this length seems to be correct as a fallback mechanism
484 // as it is the maximum length supported by the image tool.
485 user_len = 0x100; // == 256
486 }
487
488 self.user_bin = vec![0xFF; user_len as usize];
489 reader.read_exact(&mut self.user_bin)?;
490 Ok(())
491 }
492}
493
494impl ToStream for PartTab {
495 /// Writes a `PartTab` to a binary stream.
496 ///
497 /// # Arguments:
498 /// - `writer`: A mutable reference to a writer that implements `std::io::Write`.
499 fn write_to<W>(&self, writer: &mut W) -> Result<(), Error>
500 where
501 W: io::Write + io::Seek,
502 {
503 writer.write_u8(self.rma_w_state)?;
504 writer.write_u8(self.rma_ov_state)?;
505 writer.write_u8(self.eFWV)?;
506 writer.write_u8(0x000)?;
507
508 if self.records.is_empty() {
509 return Err(Error::InvalidState("Empty partition table".to_string()));
510 }
511 writer.write_u8((self.records.len() - 1) as u8)?;
512 writer.write_u8(self.fw1_idx)?;
513 writer.write_u8(self.fw2_idx)?;
514 write_padding!(writer, 3);
515
516 writer.write_u16::<LittleEndian>(self.ota_trap.into())?;
517 writer.write_u16::<LittleEndian>(self.mp_trap.into())?;
518
519 // Skip the byte set to 0xFF manually in generate_pt_table()
520 write_padding!(writer, 1);
521 writer.write_u8(self.key_exp_op as u8)?;
522
523 writer.write_u32::<LittleEndian>(self.user_bin.len() as u32)?;
524 writer.write_all(&self.user_ext)?;
525
526 // Write the partition records (num + 1, including boot record).
527 for record in &self.records {
528 record.write_to(writer)?;
529 }
530
531 // Write the user binary data.
532 writer.write_all(&self.user_bin)?;
533 Ok(())
534 }
535}
536
537/// =====================================================================================
538/// Partition Table Image (PartitionTableImage)
539/// =====================================================================================
540///
541/// The `PartitionTableImage` struct represents a complete partition table image, including
542/// the key block, header, partition table, and a hash value. It provides methods to read
543/// the image from a stream, retrieve the hash, and generate its signature.
544#[derive(Debug)]
545pub struct PartitionTableImage {
546 pub keyblock: KeyBlock,
547 pub header: ImageHeader,
548 pub pt: EncryptedOr<PartTab>,
549 hash: [u8; 32],
550}
551
552impl FromStream for PartitionTableImage {
553 /// Parses a `PartitionTableImage` from a binary stream.
554 ///
555 /// # Arguments:
556 /// - `reader`: A mutable reference to a reader that implements `std::io::Read` and `std::io::Seek`.
557 ///
558 /// # Returns:
559 /// - `Ok(())` if the partition table image was successfully parsed.
560 /// - `Err(Error)` if there was an issue reading from the stream.
561 fn read_from<R>(&mut self, reader: &mut R) -> Result<(), crate::error::Error>
562 where
563 R: std::io::Read + std::io::Seek,
564 {
565 // Read the components of the partition table image
566 self.keyblock
567 .read_from(reader)
568 .map_err(|e| Error::MalformedKeyblock(e.to_string()))?;
569
570 self.header
571 .read_from(reader)
572 .map_err(|e| Error::MalformedImageHeader(e.to_string()))?;
573
574 // Save the current position to determine the expected size later
575 let start_pos = reader.stream_position()?;
576 if self.header.is_encrypt {
577 self.pt = EncryptedOr::Encrypted(vec![0x00; self.header.segment_size as usize]);
578 self.pt
579 .read_from(reader)
580 .map_err(|e| Error::MalfromedPartTab(e.to_string()))?;
581 } else {
582 self.pt = EncryptedOr::Plain(
583 from_stream(reader).map_err(|e| Error::MalfromedPartTab(e.to_string()))?,
584 );
585 }
586 let current_pos = reader.stream_position()?;
587 let target_pos = start_pos + self.header.segment_size as u64;
588
589 // If the stream is behind of the expected position, seek back
590 if current_pos < target_pos {
591 reader.seek(io::SeekFrom::Current((target_pos - current_pos) as i64))?;
592 }
593 reader.read_exact(&mut self.hash)?;
594 Ok(())
595 }
596}
597
598impl PartitionTableImage {
599 /// Returns a reference to the 32-byte hash value of the partition table image.
600 ///
601 /// # Returns:
602 /// - A reference to the 32-byte hash array.
603 pub fn get_hash(&self) -> &[u8; 32] {
604 return &self.hash;
605 }
606
607 /// Creates a signature for the partition table image using HMAC-SHA256.
608 ///
609 /// This function generates a signature by using the HMAC (Hash-based Message Authentication
610 /// Code) algorithm with SHA-256. It takes a key and the partition table image data (excluding
611 /// the `hash` field) as inputs, and returns the resulting signature as a vector of bytes.
612 ///
613 /// # Arguments:
614 /// - `reader`: A mutable reference to a reader that implements `std::io::Read` and `std::io::Seek`.
615 /// - `key`: The key to be used in the HMAC algorithm, which should be a byte slice.
616 ///
617 /// # Returns:
618 /// - `Ok(Vec<u8>)` containing the cryptographic signature.
619 /// - `Err(Error)` if an error occurs during reading or signature generation.
620 ///
621 /// # Example:
622 /// ```rust
623 /// let signature = pt_image.create_signature(&mut reader, &key).unwrap();
624 /// ```
625 pub fn create_signature<R>(&self, reader: &mut R, key: &[u8]) -> Result<Vec<u8>, Error>
626 where
627 R: io::Read + io::Seek,
628 {
629 // Buffer for reading the data, size calculated based on header and segment size
630 let mut buffer =
631 vec![0xFF; 64 + ImageHeader::binary_size() + self.header.segment_size as usize];
632
633 reader.read_exact(&mut buffer)?;
634 return Ok(hmac_sha256(key, &buffer)?.to_vec());
635 }
636}
637
638impl AsImage for PartitionTableImage {
639 /// Computes the segment size for the partition table image.
640 ///
641 /// The segment size includes the sizes of the keyblock, header, partition table records,
642 /// and the user binary data.
643 ///
644 /// # Returns:
645 /// - `u32`: The computed segment size.
646 fn build_segment_size(&self) -> u32 {
647 // segment size is partition table static size + partition table records + user data length
648 let new_size = match &self.pt {
649 EncryptedOr::Encrypted(data) => data.len() as u32,
650 EncryptedOr::Plain(data) => {
651 (0x20 + ((data.records.len() + 1) * Record::binary_size()) + data.user_bin.len())
652 as u32
653 }
654 };
655
656 // align size to 0x20
657 return new_size + (0x20 - (new_size % 0x20));
658 }
659
660 /// Computes the signature for the partition table image.
661 ///
662 /// This method generates the HMAC SHA-256 signature for the image using the provided key.
663 ///
664 /// # Arguments:
665 /// - `key`: The key used to compute the HMAC SHA-256 signature.
666 ///
667 /// # Returns:
668 /// - `Result<Vec<u8>, crate::error::Error>`: The computed signature as a vector of bytes.
669 fn build_signature(&self, key: Option<&[u8]>) -> Result<Vec<u8>, Error> {
670 let mut buffer =
671 vec![0xFF; 64 + ImageHeader::binary_size() + self.build_segment_size() as usize];
672 let mut writer = Cursor::new(&mut buffer);
673
674 self.keyblock.write_to(&mut writer)?;
675 self.header.write_to(&mut writer)?;
676 self.pt.write_to(&mut writer)?;
677 Ok(hmac_sha256(key.unwrap(), &buffer)?.to_vec())
678 }
679
680 /// Sets the signature for the partition table image.
681 ///
682 /// This method sets the signature in the image, typically after it has been calculated.
683 ///
684 /// # Arguments:
685 /// - `signature`: The signature to set in the image.
686 fn set_signature(&mut self, signature: &[u8]) {
687 self.hash.copy_from_slice(signature);
688 }
689
690 /// Sets the segment size for the partition table image.
691 ///
692 /// This method allows setting the segment size manually.
693 ///
694 /// # Arguments:
695 /// - `size`: The segment size to set.
696 fn set_segment_size(&mut self, size: u32) {
697 self.header.segment_size = size;
698 }
699}
700
701impl Default for PartitionTableImage {
702 /// Returns a default `PartitionTableImage` with default values for all fields.
703 ///
704 /// The `keyblock`, `header`, and `pt` are initialized with their respective defaults,
705 /// and the `hash` field is set to an array of `0xFF` bytes (representing an uninitialized hash).
706 ///
707 /// # Returns:
708 /// - A `PartitionTableImage` with all fields set to their default values.
709 ///
710 /// # Example:
711 /// ```rust
712 /// let default_pt_image = PartitionTableImage::default();
713 /// ```
714 fn default() -> Self {
715 PartitionTableImage {
716 keyblock: KeyBlock::default(),
717 header: ImageHeader::default(),
718 pt: EncryptedOr::Plain(PartTab::default()),
719 hash: [0xFF; 32],
720 }
721 }
722}
723
724impl ToStream for PartitionTableImage {
725 /// Writes a `PartitionTableImage` to a binary stream.
726 ///
727 /// Note that this method does not check for valid segment size or hash values, and the
728 /// padding is applied automatically as part of the partition table write process.
729 ///
730 /// # Arguments:
731 /// - `writer`: A mutable reference to a writer that implements the `std::io::Write` and
732 /// `std::io::Seek` traits.
733 ///
734 /// # Returns:
735 /// - `Ok(())` if the write operation is successful.
736 /// - `Err(Error)` if there is an error during the write operation.
737 fn write_to<W>(&self, writer: &mut W) -> Result<(), Error>
738 where
739 W: io::Write + io::Seek,
740 {
741 // Write the keyblock and header to the stream
742 self.keyblock.write_to(writer)?;
743 self.header.write_to(writer)?;
744
745 // Create a buffer to hold the partition table (with padding applied)
746 let mut pt_buffer = vec![0xFF; self.header.segment_size as usize];
747 let mut pt_writer = Cursor::new(&mut pt_buffer);
748 self.pt.write_to(&mut pt_writer)?;
749
750 writer.write_all(&pt_buffer)?;
751 writer.write_all(&self.hash)?;
752 Ok(())
753 }
754}