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