opentitanlib/image/
image.rs

1// Copyright lowRISC contributors (OpenTitan project).
2// Licensed under the Apache License, Version 2.0, see LICENSE for details.
3// SPDX-License-Identifier: Apache-2.0
4
5use std::borrow::Cow;
6use std::collections::HashSet;
7use std::convert::TryInto;
8use std::fs::File;
9use std::io::{Read, Write};
10use std::mem::{align_of, offset_of, size_of};
11use std::path::{Path, PathBuf};
12
13use anyhow::{Context, Result, bail, ensure};
14use ml_dsa::{MlDsa87, VerifyingKey};
15use sphincsplus::{SphincsPlus, SpxDomain, SpxPublicKey};
16use thiserror::Error;
17use zerocopy::FromBytes;
18
19use crate::crypto::ecdsa::{EcdsaPublicKey, EcdsaRawPublicKey, EcdsaRawSignature};
20use crate::crypto::mldsa::mldsa_verify;
21use crate::crypto::sha256::Sha256Digest;
22use crate::image::manifest::{
23    CHIP_MANIFEST_VERSION_MAJOR2, CHIP_ROM_EXT_IDENTIFIER, CHIP_ROM_EXT_SIZE_MAX,
24    MANIFEST_EXT_ID_MLDSA87_KEY, MANIFEST_EXT_ID_MLDSA87_SIGNATURE, MANIFEST_EXT_ID_SPX_KEY,
25    MANIFEST_EXT_ID_SPX_SIGNATURE, MLDSA87_IMAGE_CONTEXT, MLDSA87_PUBLIC_KEY_BYTES,
26    MLDSA87_SIGNATURE_BYTES, Manifest, ManifestExtHeader, ManifestKind, SigverifySpxSignature,
27};
28use crate::image::manifest_def::{ManifestSigverifyBuffer, ManifestSpec};
29use crate::image::manifest_ext::{ManifestExtEntry, ManifestExtEntrySpec};
30use crate::util::file::{FromReader, ToWriter};
31use crate::util::parse_int::ParseInt;
32
33#[derive(Debug, Error)]
34pub enum ImageError {
35    #[error("Incomplete read: expected to read {0} bytes but read {1} bytes")]
36    IncompleteRead(usize, usize),
37    #[error("Failed to parse image manifest.")]
38    Parse,
39    #[error("Extension data overflows flash image.")]
40    ExtensionOverflow,
41    #[error("Extension table index is out of bounds.")]
42    BadExtensionTableIndex,
43    #[error("Extension ID 0x{0:x} not in manifest table.")]
44    NoExtensionTableEntry(u32),
45    #[error("Extension 0x{0:x} is not aligned to word boundary.")]
46    BadExtensionAlignment(u32),
47    #[error("Invalid placement of signed extension 0x{0:x}.")]
48    MisplacedSignedExtension(u32),
49    #[error("Invalid manifest major version: {0}. ECDSA support requires major version {1}.")]
50    InvalidManifestVersionforEcdsa(u16, u16),
51    #[error("Invalid signed region end 0x{0:x}.")]
52    BadSignedRegionEnd(u32),
53}
54
55pub enum MainSignatureParams {
56    Ecdsa(EcdsaRawPublicKey, EcdsaRawSignature),
57}
58
59pub struct SpxSignatureParams {
60    key: SpxPublicKey,
61    signature: [u8; 7856],
62}
63
64/// The ML-DSA-87 key and signature extensions.
65pub struct Mldsa87SignatureParams {
66    key: VerifyingKey<MlDsa87>,
67    signature: [u8; MLDSA87_SIGNATURE_BYTES],
68}
69
70/// The manifest always holds an ECDSA-P256 public key and signature whereas the
71/// SPX+ and ML-DSA-87 keys and signatures are optional manifest extensions.
72pub struct SigverifyParams {
73    pub main_sig_params: MainSignatureParams,
74    pub spx_sig_params: Option<SpxSignatureParams>,
75    pub mldsa87_sig_params: Option<Mldsa87SignatureParams>,
76    pub spx_hash_reversal_bug: bool,
77}
78
79impl SigverifyParams {
80    pub fn new(
81        main_sig_params: MainSignatureParams,
82        spx_sig_params: Option<SpxSignatureParams>,
83        mldsa87_sig_params: Option<Mldsa87SignatureParams>,
84    ) -> Self {
85        SigverifyParams {
86            main_sig_params,
87            spx_sig_params,
88            mldsa87_sig_params,
89            spx_hash_reversal_bug: false,
90        }
91    }
92    pub fn with_hash_reversal_bug(mut self, bug: bool) -> Self {
93        self.spx_hash_reversal_bug = bug;
94        self
95    }
96
97    // Verify the main image signature.
98    pub fn verify(&self, digest: &Sha256Digest) -> Result<()> {
99        match &self.main_sig_params {
100            MainSignatureParams::Ecdsa(key, sig) => {
101                let ecdsa_key: EcdsaPublicKey = key.try_into()?;
102                ecdsa_key.verify(digest, sig)?;
103            }
104        }
105        Ok(())
106    }
107
108    // Verify the optional SPX+ signature.
109    pub fn spx_verify(&self, b: &[u8], domain: SpxDomain) -> Result<()> {
110        if let Some(spx) = &self.spx_sig_params {
111            let msg = match domain {
112                SpxDomain::PreHashedSha256 => Cow::from(if self.spx_hash_reversal_bug {
113                    Sha256Digest::hash(b).to_vec_rev()
114                } else {
115                    Sha256Digest::hash(b).to_vec()
116                }),
117                _ => Cow::from(b),
118            };
119            spx.key.verify(domain, &spx.signature, &msg)?;
120        } else {
121            bail!("No SPX signature found");
122        }
123
124        Ok(())
125    }
126
127    // Verify the optional ML-DSA-87 signature (pure mode).
128    pub fn mldsa87_verify(&self, b: &[u8]) -> Result<()> {
129        let mldsa = self
130            .mldsa87_sig_params
131            .as_ref()
132            .context("No ML-DSA-87 signature found")?;
133        mldsa_verify(&mldsa.key, MLDSA87_IMAGE_CONTEXT, &mldsa.signature, b)
134    }
135}
136
137/// A buffer with the same alignment as `Manifest` for storing image data.
138#[repr(C)]
139#[derive(Debug)]
140pub struct ImageData {
141    pub bytes: [u8; Image::MAX_SIZE],
142    _align: [Manifest; 0],
143}
144
145impl Default for ImageData {
146    fn default() -> Self {
147        ImageData {
148            bytes: [0xFF; Image::MAX_SIZE],
149            _align: [],
150        }
151    }
152}
153
154#[derive(Debug, Default)]
155pub struct Image {
156    // TODO(cfrantz): We should use Box::new_uninit to create this, as the
157    // scheme of allowing Box::default to initialize this creates a copy on
158    // the stack and then copies it into the Box.  Unfortunately, Box::new_uninit
159    // is nightly-only (rust-lang#63291).
160    //
161    // For now, I've increased the thread stack size for opentitanlib_test to
162    // avoid the test overflowing its stack while initializing this item.
163    data: Box<ImageData>,
164    pub size: usize,
165}
166
167#[derive(Debug)]
168pub struct SubImage<'a> {
169    pub kind: ManifestKind,
170    pub offset: usize,
171    pub manifest: &'a Manifest,
172    pub data: &'a [u8],
173}
174
175#[derive(Debug)]
176pub enum ImageChunk {
177    Concat(PathBuf),
178    Offset(PathBuf, usize),
179}
180
181#[derive(Debug, Default)]
182pub struct ImageAssembler {
183    pub size: usize,
184    pub mirrored: bool,
185    pub chunks: Vec<ImageChunk>,
186}
187
188impl FromReader for Image {
189    /// Reads in an `Image`.
190    fn from_reader(mut r: impl Read) -> Result<Self> {
191        let mut image = Image::default();
192        image.size = r.read(&mut image.data.bytes)?;
193        Ok(image)
194    }
195}
196
197impl ToWriter for Image {
198    /// Writes out the `Image`.
199    fn to_writer(&self, w: &mut impl Write) -> Result<()> {
200        w.write_all(&self.data.bytes[..self.size])?;
201        Ok(())
202    }
203}
204
205impl Image {
206    pub const MAX_SIZE: usize = 1024 * 1024;
207
208    /// Perform a sanity check to ensure that the image comes with manifest.
209    ///
210    /// Note that passing the sanity check doesn't guarantee that the manifest is valid.
211    pub fn manifest_sanity_check(&self) -> Result<()> {
212        let manifest = self.borrow_manifest()?;
213        let len = self.data.bytes.len() as u32;
214
215        ensure!(manifest.signed_region_end <= len);
216        ensure!(manifest.length <= len);
217        ensure!(manifest.code_start < len);
218        ensure!(manifest.code_end < len);
219        ensure!(manifest.entry_point < len);
220        ensure!(manifest.extensions.entries.iter().all(|x| x.offset < len));
221
222        // We allow the ROM_EXT to be larger than maximum slot size to support provisioning flows
223        // and future ROM_EXT expansion if needed.
224        if (manifest.identifier == CHIP_ROM_EXT_IDENTIFIER)
225            && (manifest.length > CHIP_ROM_EXT_SIZE_MAX)
226        {
227            log::warn!("ROM_EXT is larger than 64k. Link offsets may need recalculating.");
228        }
229
230        Ok(())
231    }
232
233    // Retrieve the SPX+ signature from the image, if present.
234    fn get_spx_signature(&self) -> Result<Option<SpxSignatureParams>> {
235        let ext_tab = self.borrow_manifest()?.extensions.entries;
236        let key_o = ext_tab
237            .iter()
238            .find(|e| e.identifier == MANIFEST_EXT_ID_SPX_KEY);
239        let sig_o = ext_tab
240            .iter()
241            .find(|e| e.identifier == MANIFEST_EXT_ID_SPX_SIGNATURE);
242
243        match (key_o, sig_o) {
244            (Some(key_e), Some(sig_e)) => {
245                const KEY_SIZE: usize = 32; // SPX+ public key size.
246                const SIG_SIZE: usize = std::mem::size_of::<SigverifySpxSignature>();
247
248                let mut key_bytes = [0u8; KEY_SIZE];
249                let mut signature = [0u8; SIG_SIZE];
250                let k_ofs = (key_e.offset + 8) as usize;
251                let s_ofs = (sig_e.offset + 8) as usize;
252
253                key_bytes.copy_from_slice(&self.data.bytes[k_ofs..k_ofs + KEY_SIZE]);
254                signature.copy_from_slice(&self.data.bytes[s_ofs..s_ofs + SIG_SIZE]);
255
256                let key = SpxPublicKey::from_bytes(SphincsPlus::Sha2128sSimple, &key_bytes)?;
257
258                Ok(Some(SpxSignatureParams { key, signature }))
259            }
260            (_, _) => Ok(None),
261        }
262    }
263
264    // Retrieve the ML-DSA-87 signature from the image, if present.
265    fn get_mldsa87_signature(&self) -> Result<Option<Mldsa87SignatureParams>> {
266        let ext_tab = self.borrow_manifest()?.extensions.entries;
267        let ext_data = |id: u32, len: usize| {
268            ext_tab
269                .iter()
270                .find(|e| e.identifier == id)
271                .map(|e| {
272                    let ofs = e.offset as usize + size_of::<ManifestExtHeader>();
273                    self.data
274                        .bytes
275                        .get(ofs..ofs + len)
276                        .ok_or(ImageError::ExtensionOverflow)
277                })
278                .transpose()
279        };
280
281        match (
282            ext_data(MANIFEST_EXT_ID_MLDSA87_KEY, MLDSA87_PUBLIC_KEY_BYTES)?,
283            ext_data(MANIFEST_EXT_ID_MLDSA87_SIGNATURE, MLDSA87_SIGNATURE_BYTES)?,
284        ) {
285            (Some(key), Some(signature)) => Ok(Some(Mldsa87SignatureParams {
286                key: VerifyingKey::decode(&key.try_into()?),
287                signature: signature.try_into()?,
288            })),
289            (_, _) => Ok(None),
290        }
291    }
292
293    pub fn get_sigverify_params_from_manifest(&self) -> Result<SigverifyParams> {
294        let manifest = self.borrow_manifest()?;
295        let manifest_def: ManifestSpec = manifest.try_into()?;
296
297        let spx_sig_params = self.get_spx_signature()?;
298        let mldsa87_sig_params = self.get_mldsa87_signature()?;
299
300        let pub_key = manifest_def
301            .pub_key()
302            .ok_or(ImageError::Parse)?
303            .to_le_bytes();
304
305        let signature = manifest_def
306            .signature()
307            .ok_or(ImageError::Parse)?
308            .to_le_bytes();
309
310        let ecdsa_pub_key = EcdsaRawPublicKey::read(&mut std::io::Cursor::new(pub_key))?;
311        let ecdsa_sig = EcdsaRawSignature::read(&mut std::io::Cursor::new(signature))?;
312
313        Ok(SigverifyParams::new(
314            MainSignatureParams::Ecdsa(ecdsa_pub_key, ecdsa_sig),
315            spx_sig_params,
316            mldsa87_sig_params,
317        ))
318    }
319
320    /// Overwrites all fields in the image's manifest that are defined in `other`.
321    pub fn overwrite_manifest(&mut self, other: ManifestSpec) -> Result<()> {
322        let manifest = self.borrow_manifest_mut()?;
323        let mut manifest_def: ManifestSpec = (&*manifest).try_into()?;
324        manifest_def.overwrite_fields(other);
325        *manifest = manifest_def.try_into()?;
326        Ok(())
327    }
328
329    /// Adds an extension to the signed region of this `Image`.
330    ///
331    /// This will take all the signed extensions in `spec` and append them to the image.
332    /// This should be called before adding any unsigned extensions to ensure all extensions that
333    /// are a part of the signature exist within the contiguous signed region of the image.
334    pub fn add_signed_manifest_extensions(&mut self, spec: &[ManifestExtEntrySpec]) -> Result<()> {
335        spec.iter()
336            .filter(|entry_spec| entry_spec.is_signed())
337            .try_for_each(|entry_spec| {
338                self.add_manifest_extension(ManifestExtEntry::from_spec(entry_spec)?)
339            })
340    }
341
342    /// Adds an extension to the unsigned region of this `Image`.
343    ///
344    /// This will take all the unsigned extensions in `spec` and append them to the image.
345    /// This should only be called once all signed extensions have been added.
346    pub fn add_unsigned_manifest_extensions(
347        &mut self,
348        spec: &[ManifestExtEntrySpec],
349    ) -> Result<()> {
350        spec.iter()
351            .filter(|entry_spec| !entry_spec.is_signed())
352            .try_for_each(|entry_spec| {
353                self.add_manifest_extension(ManifestExtEntry::from_spec(entry_spec)?)
354            })
355    }
356
357    /// Adds an extension to the end of this `Image`.
358    pub fn add_manifest_extension(&mut self, entry: ManifestExtEntry) -> Result<()> {
359        let manifest = self.borrow_manifest()?;
360
361        // A copy of the extension table since we can't borrow as mutable.
362        let mut ext_table = manifest.extensions.entries;
363
364        let entry_id = entry.header().identifier;
365
366        // Update the offset in the extension table.
367        let ext_table_entry = ext_table
368            .iter_mut()
369            .find(|e| e.identifier == entry_id)
370            .ok_or(ImageError::NoExtensionTableEntry(entry_id))?;
371
372        // If the extension already exists, overwrite it, else append it to the end of the
373        // image.
374        let offset = if ext_table_entry.offset != 0 {
375            ext_table_entry.offset
376        } else {
377            let alignment = entry.alignment() as usize;
378            let header_size = size_of::<ManifestExtHeader>();
379            // Align the offset so that the first data element inside the extension (following the
380            // manifest header) starts on a boundary aligned to `alignment` (typically a 2KB page boundary).
381            let aligned_size = if alignment > header_size {
382                ((self.size + header_size + alignment - 1) & !(alignment - 1)) - header_size
383            } else {
384                (self.size + alignment - 1) & !(alignment - 1)
385            };
386            ensure!(
387                aligned_size <= Image::MAX_SIZE,
388                ImageError::ExtensionOverflow
389            );
390            self.size = aligned_size;
391            ensure!(
392                self.size.is_multiple_of(align_of::<u32>()),
393                ImageError::BadExtensionAlignment(entry_id)
394            );
395            self.size.try_into()?
396        };
397        ext_table_entry.offset = offset;
398
399        // Write the extension to the end of the image.
400        let ext_bytes = entry.to_vec();
401        let end_index = offset
402            .checked_add(ext_bytes.len().try_into()?)
403            .ok_or(ImageError::ExtensionOverflow)?;
404        let extension_slice = self
405            .data
406            .bytes
407            .get_mut(offset as usize..end_index as usize)
408            .ok_or(ImageError::ExtensionOverflow)?;
409        extension_slice.copy_from_slice(ext_bytes.as_slice());
410        self.size = std::cmp::max(end_index as usize, self.size);
411
412        let manifest = self.borrow_manifest_mut()?;
413        manifest.extensions.entries = ext_table;
414
415        Ok(())
416    }
417
418    /// Allocates space for a manifest extension and sets the offset in the extension table.
419    ///
420    /// This function is similar to `add_manifest_extension`, but doesn't populate any data for the
421    /// extension. This is necessary to properly set the `length` field of the manifest and the
422    /// `offset` field of the manifest extension entry for signing.
423    pub fn allocate_manifest_extension(&mut self, id: u32, len: usize) -> Result<()> {
424        let aligned_size = (self.size + 3) & !3;
425        ensure!(
426            aligned_size + len <= Image::MAX_SIZE,
427            ImageError::ExtensionOverflow
428        );
429        self.size = aligned_size;
430        let offset = self.size as u32;
431        self.borrow_manifest_mut()?
432            .extensions
433            .entries
434            .iter_mut()
435            .find(|e| e.identifier == id)
436            .ok_or(ImageError::NoExtensionTableEntry(id))?
437            .offset = offset;
438
439        self.size = self
440            .size
441            .checked_add(len)
442            .ok_or(ImageError::ExtensionOverflow)?;
443
444        Ok(())
445    }
446
447    /// Clears any entry in the manifest extension table that doesn't have an offset.
448    ///
449    /// This allows a manifest definition with a populated extension table to be used even when
450    /// extensions aren't provided.
451    pub fn drop_null_extensions(&mut self) -> Result<()> {
452        let manifest = self.borrow_manifest_mut()?;
453
454        manifest.extensions.entries.iter_mut().for_each(|e| {
455            if e.offset == 0 {
456                e.identifier = 0;
457            }
458        });
459
460        Ok(())
461    }
462
463    /// Updates the signature field in the `Manifest` with the provided ECDSA signature.
464    pub fn update_ecdsa_signature(&mut self, signature: EcdsaRawSignature) -> Result<()> {
465        let manifest = self.borrow_manifest_mut()?;
466
467        // TODO(moidx): Remove check once we have migrated away from RSA keys
468        // and start using a key type field in the manifest.
469        //
470        // Note(cfrantz): I have disabled this error return because we want to test
471        // manifests with bad version numbers.  I've replaced the error return with
472        // a log message.  As stated by moidx@, we'll remove this once we've
473        // completely migrated away from RSA keys.
474        // ensure!(
475        //     manifest.manifest_version.major == CHIP_MANIFEST_VERSION_MAJOR2,
476        //     ImageError::InvalidManifestVersionforEcdsa(
477        //         manifest.manifest_version.major,
478        //         CHIP_MANIFEST_VERSION_MAJOR2
479        //     )
480        // );
481        if manifest.manifest_version.major != CHIP_MANIFEST_VERSION_MAJOR2 {
482            log::error!(
483                "Invalid manifest version for ECDSA: {:?}",
484                manifest.manifest_version
485            );
486        }
487
488        let mut manifest_def: ManifestSpec = (&*manifest).try_into()?;
489
490        let signature_bytes = signature
491            .r
492            .iter()
493            .chain(signature.s.iter())
494            .copied()
495            .collect::<Vec<u8>>();
496
497        manifest_def.update_signature(ManifestSigverifyBuffer::from_le_bytes(signature_bytes)?);
498        *manifest = manifest_def.try_into()?;
499        Ok(())
500    }
501
502    /// Updates the pub_key field in the `Manifest` with the ECDSA public key.
503    pub fn update_ecdsa_public_key(&mut self, ecdsa_public_key: EcdsaRawPublicKey) -> Result<()> {
504        let manifest = self.borrow_manifest_mut()?;
505        let mut manifest_def: ManifestSpec = (&*manifest).try_into()?;
506
507        let key_bytes = ecdsa_public_key
508            .x
509            .iter()
510            .chain(ecdsa_public_key.y.iter())
511            .copied()
512            .collect::<Vec<u8>>();
513
514        manifest_def.update_pub_key(ManifestSigverifyBuffer::from_le_bytes(key_bytes)?);
515        *manifest = manifest_def.try_into()?;
516        Ok(())
517    }
518
519    pub fn subimages(&self) -> Result<Vec<SubImage<'_>>> {
520        let mut result = Vec::new();
521        let mut offset = 0;
522        while offset < self.size {
523            let m = &self.data.bytes[offset..offset + size_of::<Manifest>()];
524            let manifest = Manifest::ref_from_bytes(m).map_err(|_| ImageError::Parse)?;
525            let kind = ManifestKind(manifest.identifier);
526            let mut size = 1;
527            if kind.is_known_value() {
528                size = manifest.length as usize;
529                result.push(SubImage {
530                    kind,
531                    offset,
532                    manifest,
533                    data: &self.data.bytes[offset..offset + size],
534                });
535            }
536            // Round up to next 64K offset.
537            offset += (size + 65535) & !65535;
538        }
539        Ok(result)
540    }
541
542    pub fn bytes(&self) -> &[u8] {
543        &self.data.bytes[..self.size]
544    }
545
546    pub fn borrow_manifest(&self) -> Result<&Manifest> {
547        let manifest_slice = &self.data.bytes[0..size_of::<Manifest>()];
548        let manifest = Manifest::ref_from_bytes(manifest_slice).map_err(|_| ImageError::Parse)?;
549        Ok(manifest)
550    }
551
552    pub fn borrow_manifest_mut(&mut self) -> Result<&mut Manifest> {
553        let manifest_slice = &mut self.data.bytes[0..size_of::<Manifest>()];
554        let manifest = Manifest::mut_from_bytes(manifest_slice).map_err(|_| ImageError::Parse)?;
555        Ok(manifest)
556    }
557
558    /// Updates the length field in the `Manifest`.
559    pub fn update_length(&mut self) -> Result<usize> {
560        self.borrow_manifest_mut()?.length = self.size as u32;
561        Ok(self.size)
562    }
563
564    /// Sets the `signed_region_end` field to the end of the last signed extension.
565    ///
566    /// The end of the signed region is computed as the offset of the first extension that is not
567    /// in `signed_ids` or the size of the image if there is no such extension.
568    pub fn update_signed_region(&mut self, signed_ids: &HashSet<u32>) -> Result<()> {
569        let image_size = self.size as u32;
570        let mut first_unsigned_ext = 0u32;
571        let manifest = self.borrow_manifest_mut()?;
572
573        let mut ext_table = manifest.extensions.entries;
574        ext_table.sort_by(|a, b| a.offset.cmp(&b.offset));
575        for e in ext_table {
576            // Ignore any extensions that haven't been added to the image.
577            if e.offset == 0 {
578                continue;
579            }
580            if signed_ids.contains(&e.identifier) {
581                // Since extensions are sorted by offset, if we've already seen an unsigned
582                // extension then the signed region is discontinuous.
583                ensure!(
584                    first_unsigned_ext == 0,
585                    ImageError::MisplacedSignedExtension(e.identifier)
586                );
587            } else if first_unsigned_ext == 0 {
588                first_unsigned_ext = e.offset;
589            }
590        }
591
592        manifest.signed_region_end = if first_unsigned_ext == 0 {
593            image_size
594        } else {
595            first_unsigned_ext
596        };
597
598        Ok(())
599    }
600
601    /// Operates on the signed region of the image.
602    pub fn map_signed_region<F, R>(&self, f: F) -> Result<R>
603    where
604        F: FnOnce(&[u8]) -> R,
605    {
606        let end = self.borrow_manifest()?.signed_region_end;
607        let region = self
608            .data
609            .bytes
610            .get(offset_of!(Manifest, usage_constraints)..end as usize)
611            .ok_or(ImageError::BadSignedRegionEnd(end))?;
612        Ok(f(region))
613    }
614
615    /// Compute the SHA256 digest for the signed portion of the `Image`.
616    pub fn compute_digest(&self) -> Result<Sha256Digest> {
617        self.map_signed_region(|v| Sha256Digest::hash(v))
618    }
619}
620
621impl SubImage<'_> {
622    /// Operates on the signed region of the image.
623    pub fn map_signed_region<F, R>(&self, f: F) -> Result<R>
624    where
625        F: FnOnce(&[u8]) -> R,
626    {
627        let end = self.manifest.signed_region_end;
628        let region = self
629            .data
630            .get(offset_of!(Manifest, usage_constraints)..end as usize)
631            .ok_or(ImageError::BadSignedRegionEnd(end))?;
632        Ok(f(region))
633    }
634
635    /// Compute the SHA256 digest for the signed portion of the `Image`.
636    pub fn compute_digest(&self) -> Result<Sha256Digest> {
637        self.map_signed_region(|v| Sha256Digest::hash(v))
638    }
639}
640
641impl ImageAssembler {
642    /// Creates an `ImageAssembler` with a given `size` and mirroring parameters.
643    pub fn with_params(size: usize, mirrored: bool) -> Self {
644        ImageAssembler {
645            size,
646            mirrored,
647            ..Default::default()
648        }
649    }
650
651    /// Creates an `ImageAssembler` with default parameters for OpenTitan: a 1MiB image which is mirrored.
652    pub fn new() -> Self {
653        Self::with_params(0x100000, true)
654    }
655
656    /// Parse a list of strings into chunks to be assembled.
657    /// Each string may be a filename or a filename@offset describing where in the assembled image the contents of the file should appear.
658    /// The offset is an integer expressed in any of the bases accepted by [`ParseInt`].
659    pub fn parse(&mut self, chunks: &[impl AsRef<str>]) -> Result<()> {
660        for chunk in chunks {
661            if let Some((file, offset)) = chunk.as_ref().split_once('@') {
662                self.chunks.push(ImageChunk::Offset(
663                    PathBuf::from(file),
664                    usize::from_str(offset)?,
665                ));
666            } else {
667                self.chunks
668                    .push(ImageChunk::Concat(PathBuf::from(chunk.as_ref())));
669            }
670        }
671        Ok(())
672    }
673
674    // Read a file into a buffer.  Ensure the entire file is read.
675    fn read(path: &Path, buf: &mut [u8]) -> Result<usize> {
676        let mut file = File::open(path)?;
677        let len = file.metadata()?.len() as usize;
678        ensure!(len <= buf.len(), ImageError::IncompleteRead(len, buf.len()));
679        if buf[..len].iter().any(|&c| c != 0xff) {
680            bail!("{path:?} Overlapped during image assembly");
681        }
682        let n = file.read(buf)?;
683        ensure!(len == n, ImageError::IncompleteRead(len, n));
684        Ok(n)
685    }
686
687    /// Assemble the image according to the parameters and parsed chunk specifications.
688    pub fn assemble(&self) -> Result<Vec<u8>> {
689        let size = if self.mirrored {
690            self.size / 2
691        } else {
692            self.size
693        };
694        let mut image = vec![0xff; size];
695        let mut pos = 0;
696        for chunk in &self.chunks {
697            match chunk {
698                ImageChunk::Concat(path) => {
699                    let n = Self::read(path, &mut image[pos..])?;
700                    pos += n;
701                }
702                ImageChunk::Offset(path, offset) => {
703                    let n = Self::read(path, &mut image[*offset..])?;
704                    pos = offset + n;
705                }
706            }
707        }
708        if self.mirrored {
709            image.extend_from_within(..size);
710        }
711        Ok(image)
712    }
713}
714
715#[cfg(test)]
716mod tests {
717    use super::*;
718    use crate::util::testdata;
719
720    #[test]
721    fn test_assemble_concat() -> Result<()> {
722        // Test image assembly by concatenation.
723        let mut image = ImageAssembler::with_params(16, false);
724        image.parse(&[
725            testdata("image/hello.txt").to_str().unwrap(),
726            testdata("image/world.txt").to_str().unwrap(),
727        ])?;
728        let data = image.assemble()?;
729        assert_eq!(data, b"HelloWorld\xff\xff\xff\xff\xff\xff");
730        Ok(())
731    }
732
733    #[test]
734    fn test_assemble_offset() -> Result<()> {
735        // Test image assembly by explicit offsets.
736        let mut image = ImageAssembler::with_params(16, false);
737        image.parse(&[
738            testdata("image/hello.txt@0").to_str().unwrap(),
739            testdata("image/world.txt@0x8").to_str().unwrap(),
740        ])?;
741        let data = image.assemble()?;
742        assert_eq!(data, b"Hello\xff\xff\xffWorld\xff\xff\xff");
743        Ok(())
744    }
745
746    #[test]
747    fn test_assemble_mirrored() -> Result<()> {
748        // Test image assembly with mirroring.
749        let mut image = ImageAssembler::with_params(20, true);
750        image.parse(&[
751            testdata("image/hello.txt").to_str().unwrap(),
752            testdata("image/world.txt").to_str().unwrap(),
753        ])?;
754        let data = image.assemble()?;
755        assert_eq!(data, b"HelloWorldHelloWorld");
756        Ok(())
757    }
758
759    #[test]
760    fn test_assemble_mirrored_offset_error() -> Result<()> {
761        // Test image assembly where one of the source files isn't read completely.
762        let mut image = ImageAssembler::with_params(16, true);
763        image.parse(&[
764            testdata("image/hello.txt@0").to_str().unwrap(),
765            testdata("image/world.txt@0x5").to_str().unwrap(),
766        ])?;
767        let err = image.assemble().unwrap_err();
768        assert_eq!(
769            err.to_string(),
770            "Incomplete read: expected to read 5 bytes but read 3 bytes"
771        );
772        Ok(())
773    }
774
775    #[test]
776    fn test_load_image() {
777        // Read and write back image.
778        let image = Image::read_from_file(&testdata("image/test_image.bin")).unwrap();
779        image
780            .write_to_file(&testdata("image/test_image_out.bin"))
781            .unwrap();
782
783        // Ensure the result is identical to the original.
784        let (mut orig_bytes, mut res_bytes) = (Vec::<u8>::new(), Vec::<u8>::new());
785        File::open(testdata("image/test_image.bin"))
786            .unwrap()
787            .read_to_end(&mut orig_bytes)
788            .unwrap();
789        File::open(testdata("image/test_image_out.bin"))
790            .unwrap()
791            .read_to_end(&mut res_bytes)
792            .unwrap();
793        assert_eq!(orig_bytes, res_bytes);
794    }
795}