disk/
composite.rs

1// Copyright 2019 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5use std::cmp::max;
6use std::cmp::min;
7use std::collections::HashSet;
8use std::convert::TryInto;
9use std::fs::File;
10use std::fs::OpenOptions;
11use std::io;
12use std::io::ErrorKind;
13use std::io::Read;
14use std::io::Seek;
15use std::io::SeekFrom;
16use std::io::Write;
17use std::ops::Range;
18use std::path::Path;
19use std::path::PathBuf;
20use std::sync::atomic::AtomicBool;
21use std::sync::atomic::Ordering;
22use std::sync::Arc;
23
24use async_trait::async_trait;
25use base::AsRawDescriptors;
26use base::FileAllocate;
27use base::FileReadWriteAtVolatile;
28use base::FileSetLen;
29use base::RawDescriptor;
30use base::VolatileSlice;
31use crc32fast::Hasher;
32use cros_async::BackingMemory;
33use cros_async::Executor;
34use cros_async::IoOptions;
35use cros_async::MemRegionIter;
36use protobuf::Message;
37use protos::cdisk_spec;
38use protos::cdisk_spec::ComponentDisk;
39use protos::cdisk_spec::CompositeDisk;
40use protos::cdisk_spec::ReadWriteCapability;
41use remain::sorted;
42use thiserror::Error;
43use uuid::Uuid;
44
45use crate::gpt;
46use crate::gpt::write_gpt_header;
47use crate::gpt::write_protective_mbr;
48use crate::gpt::GptPartitionEntry;
49use crate::gpt::GPT_BEGINNING_SIZE;
50use crate::gpt::GPT_END_SIZE;
51use crate::gpt::GPT_HEADER_SIZE;
52use crate::gpt::GPT_NUM_PARTITIONS;
53use crate::gpt::GPT_PARTITION_ENTRY_SIZE;
54use crate::gpt::SECTOR_SIZE;
55use crate::open_disk_file;
56use crate::AsyncDisk;
57use crate::DiskFile;
58use crate::DiskFileParams;
59use crate::DiskGetLen;
60use crate::ImageType;
61use crate::ToAsyncDisk;
62
63/// The amount of padding needed between the last partition entry and the first partition, to align
64/// the partition appropriately. The two sectors are for the MBR and the GPT header.
65const PARTITION_ALIGNMENT_SIZE: usize = GPT_BEGINNING_SIZE as usize
66    - 2 * SECTOR_SIZE as usize
67    - GPT_NUM_PARTITIONS as usize * GPT_PARTITION_ENTRY_SIZE as usize;
68const HEADER_PADDING_LENGTH: usize = SECTOR_SIZE as usize - GPT_HEADER_SIZE as usize;
69// Keep all partitions 4k aligned for performance.
70const PARTITION_SIZE_SHIFT: u8 = 12;
71
72// From https://en.wikipedia.org/wiki/GUID_Partition_Table#Partition_type_GUIDs.
73const LINUX_FILESYSTEM_GUID: Uuid = Uuid::from_u128(0x0FC63DAF_8483_4772_8E79_3D69D8477DE4);
74const EFI_SYSTEM_PARTITION_GUID: Uuid = Uuid::from_u128(0xC12A7328_F81F_11D2_BA4B_00A0C93EC93B);
75
76#[sorted]
77#[derive(Error, Debug)]
78pub enum Error {
79    #[error("failed to use underlying disk: \"{0}\"")]
80    DiskError(Box<crate::Error>),
81    #[error("duplicate GPT partition label \"{0}\"")]
82    DuplicatePartitionLabel(String),
83    #[error("failed to write GPT header: \"{0}\"")]
84    GptError(gpt::Error),
85    #[error("invalid magic header for composite disk format")]
86    InvalidMagicHeader,
87    #[error("invalid partition path {0:?}")]
88    InvalidPath(PathBuf),
89    #[error("failed to parse specification proto: \"{0}\"")]
90    InvalidProto(protobuf::Error),
91    #[error("invalid specification: \"{0}\"")]
92    InvalidSpecification(String),
93    #[error("no image files for partition {0:?}")]
94    NoImageFiles(PartitionInfo),
95    #[error("failed to open component file \"{1}\": \"{0}\"")]
96    OpenFile(io::Error, String),
97    #[error("failed to read specification: \"{0}\"")]
98    ReadSpecificationError(io::Error),
99    #[error("Read-write partition {0:?} size is not a multiple of {multiple}.", multiple = 1 << PARTITION_SIZE_SHIFT)]
100    UnalignedReadWrite(PartitionInfo),
101    #[error("unknown version {0} in specification")]
102    UnknownVersion(u64),
103    #[error("unsupported component disk type \"{0:?}\"")]
104    UnsupportedComponent(ImageType),
105    #[error("failed to write composite disk header: \"{0}\"")]
106    WriteHeader(io::Error),
107    #[error("failed to write specification proto: \"{0}\"")]
108    WriteProto(protobuf::Error),
109    #[error("failed to write zero filler: \"{0}\"")]
110    WriteZeroFiller(io::Error),
111}
112
113impl From<gpt::Error> for Error {
114    fn from(e: gpt::Error) -> Self {
115        Self::GptError(e)
116    }
117}
118
119pub type Result<T> = std::result::Result<T, Error>;
120
121#[derive(Debug)]
122struct ComponentDiskPart {
123    file: Box<dyn DiskFile>,
124    offset: u64, // Location in the block device visible to the guest
125    length: u64,
126    file_offset: u64, // Location within the host file
127    // Whether there have been any writes since the last fsync or fdatasync.
128    needs_flush: AtomicBool,
129}
130
131impl ComponentDiskPart {
132    fn range(&self) -> Range<u64> {
133        self.offset..(self.offset + self.length)
134    }
135}
136
137/// Represents a composite virtual disk made out of multiple component files. This is described on
138/// disk by a protocol buffer file that lists out the component file locations and their offsets
139/// and lengths on the virtual disk. The spaces covered by the component disks must be contiguous
140/// and not overlapping.
141#[derive(Debug)]
142pub struct CompositeDiskFile {
143    component_disks: Vec<ComponentDiskPart>,
144    // We keep the root composite file open so that the file lock is not dropped.
145    _disk_spec_file: File,
146}
147
148// TODO(b/271381851): implement `try_clone`. It allows virtio-blk to run multiple workers.
149impl DiskFile for CompositeDiskFile {}
150
151fn ranges_overlap(a: &Range<u64>, b: &Range<u64>) -> bool {
152    range_intersection(a, b).is_some()
153}
154
155fn range_intersection(a: &Range<u64>, b: &Range<u64>) -> Option<Range<u64>> {
156    let r = Range {
157        start: max(a.start, b.start),
158        end: min(a.end, b.end),
159    };
160    if r.is_empty() {
161        None
162    } else {
163        Some(r)
164    }
165}
166
167/// The version of the composite disk format supported by this implementation.
168const COMPOSITE_DISK_VERSION: u64 = 2;
169
170/// A magic string placed at the beginning of a composite disk file to identify it.
171pub const CDISK_MAGIC: &str = "composite_disk\x1d";
172
173impl CompositeDiskFile {
174    fn new(mut disks: Vec<ComponentDiskPart>, disk_spec_file: File) -> Result<CompositeDiskFile> {
175        disks.sort_by(|d1, d2| d1.offset.cmp(&d2.offset));
176        for s in disks.windows(2) {
177            if s[0].offset == s[1].offset {
178                return Err(Error::InvalidSpecification(format!(
179                    "Two disks at offset {}",
180                    s[0].offset
181                )));
182            }
183        }
184        Ok(CompositeDiskFile {
185            component_disks: disks,
186            _disk_spec_file: disk_spec_file,
187        })
188    }
189
190    /// Set up a composite disk by reading the specification from a file. The file must consist of
191    /// the CDISK_MAGIC string followed by one binary instance of the CompositeDisk protocol
192    /// buffer. Returns an error if it could not read the file or if the specification was invalid.
193    pub fn from_file(mut file: File, params: DiskFileParams) -> Result<CompositeDiskFile> {
194        file.seek(SeekFrom::Start(0))
195            .map_err(Error::ReadSpecificationError)?;
196        let mut magic_space = [0u8; CDISK_MAGIC.len()];
197        file.read_exact(&mut magic_space[..])
198            .map_err(Error::ReadSpecificationError)?;
199        if magic_space != CDISK_MAGIC.as_bytes() {
200            return Err(Error::InvalidMagicHeader);
201        }
202        let proto: cdisk_spec::CompositeDisk =
203            Message::parse_from_reader(&mut file).map_err(Error::InvalidProto)?;
204        if proto.version > COMPOSITE_DISK_VERSION {
205            return Err(Error::UnknownVersion(proto.version));
206        }
207        let mut disks: Vec<ComponentDiskPart> = proto
208            .component_disks
209            .iter()
210            .map(|disk| {
211                let writable = !params.is_read_only
212                    && disk.read_write_capability
213                        == cdisk_spec::ReadWriteCapability::READ_WRITE.into();
214                let component_path = PathBuf::from(&disk.file_path);
215                let path = if component_path.is_relative() || proto.version > 1 {
216                    params.path.parent().unwrap().join(component_path)
217                } else {
218                    component_path
219                };
220
221                // Note that a read-only parts of a composite disk should NOT be marked sparse,
222                // as the action of marking them sparse is a write. This may seem a little hacky,
223                // and it is; however:
224                //    (a)  there is not a good way to pass sparseness parameters per composite disk
225                //         part (the proto does not have fields for it).
226                //    (b)  this override of sorts always matches the correct user intent.
227                Ok(ComponentDiskPart {
228                    file: open_disk_file(DiskFileParams {
229                        path: path.to_owned(),
230                        is_read_only: !writable,
231                        is_sparse_file: params.is_sparse_file && writable,
232                        // TODO: Should pass `params.is_overlapped` through here. Needs testing.
233                        // is_overlapped: false,
234                        is_direct: params.is_direct,
235                        lock: params.lock,
236                        depth: params.depth + 1,
237                        ..Default::default()
238                    })
239                    .map_err(|e| Error::DiskError(Box::new(e)))?,
240                    offset: disk.offset,
241                    length: 0, // Assigned later
242                    file_offset: disk.file_offset,
243                    needs_flush: AtomicBool::new(false),
244                })
245            })
246            .collect::<Result<Vec<ComponentDiskPart>>>()?;
247        disks.sort_by(|d1, d2| d1.offset.cmp(&d2.offset));
248        for i in 0..(disks.len() - 1) {
249            let length = disks[i + 1].offset - disks[i].offset;
250            if length == 0 {
251                let text = format!("Two disks at offset {}", disks[i].offset);
252                return Err(Error::InvalidSpecification(text));
253            }
254            if let Some(disk) = disks.get_mut(i) {
255                disk.length = length;
256            } else {
257                let text = format!("Unable to set disk length {length}");
258                return Err(Error::InvalidSpecification(text));
259            }
260        }
261        if let Some(last_disk) = disks.last_mut() {
262            if proto.length <= last_disk.offset {
263                let text = format!(
264                    "Full size of disk doesn't match last offset. {} <= {}",
265                    proto.length, last_disk.offset
266                );
267                return Err(Error::InvalidSpecification(text));
268            }
269            last_disk.length = proto.length - last_disk.offset;
270        } else {
271            let text = format!("Unable to set last disk length to end at {}", proto.length);
272            return Err(Error::InvalidSpecification(text));
273        }
274
275        CompositeDiskFile::new(disks, file)
276    }
277
278    fn length(&self) -> u64 {
279        if let Some(disk) = self.component_disks.last() {
280            disk.offset + disk.length
281        } else {
282            0
283        }
284    }
285
286    fn disk_at_offset(&self, offset: u64) -> io::Result<&ComponentDiskPart> {
287        self.component_disks
288            .iter()
289            .find(|disk| disk.range().contains(&offset))
290            .ok_or_else(|| {
291                io::Error::new(
292                    ErrorKind::InvalidData,
293                    format!("no disk at offset {offset}"),
294                )
295            })
296    }
297}
298
299impl DiskGetLen for CompositeDiskFile {
300    fn get_len(&self) -> io::Result<u64> {
301        Ok(self.length())
302    }
303}
304
305impl FileSetLen for CompositeDiskFile {
306    fn set_len(&self, _len: u64) -> io::Result<()> {
307        Err(io::Error::other("unsupported operation"))
308    }
309}
310
311// Implements Read and Write targeting volatile storage for composite disks.
312//
313// Note that reads and writes will return early if crossing component disk boundaries.
314// This is allowed by the read and write specifications, which only say read and write
315// have to return how many bytes were actually read or written. Use read_exact_volatile
316// or write_all_volatile to make sure all bytes are received/transmitted.
317//
318// If one of the component disks does a partial read or write, that also gets passed
319// transparently to the parent.
320impl FileReadWriteAtVolatile for CompositeDiskFile {
321    fn read_at_volatile(&self, slice: VolatileSlice, offset: u64) -> io::Result<usize> {
322        let cursor_location = offset;
323        let disk = self.disk_at_offset(cursor_location)?;
324        let subslice = if cursor_location + slice.size() as u64 > disk.offset + disk.length {
325            let new_size = disk.offset + disk.length - cursor_location;
326            slice
327                .sub_slice(0, new_size as usize)
328                .map_err(|e| io::Error::new(ErrorKind::InvalidData, e.to_string()))?
329        } else {
330            slice
331        };
332        disk.file
333            .read_at_volatile(subslice, cursor_location - disk.offset + disk.file_offset)
334    }
335    fn write_at_volatile(&self, slice: VolatileSlice, offset: u64) -> io::Result<usize> {
336        let cursor_location = offset;
337        let disk = self.disk_at_offset(cursor_location)?;
338        let subslice = if cursor_location + slice.size() as u64 > disk.offset + disk.length {
339            let new_size = disk.offset + disk.length - cursor_location;
340            slice
341                .sub_slice(0, new_size as usize)
342                .map_err(|e| io::Error::new(ErrorKind::InvalidData, e.to_string()))?
343        } else {
344            slice
345        };
346
347        let bytes = disk
348            .file
349            .write_at_volatile(subslice, cursor_location - disk.offset + disk.file_offset)?;
350        disk.needs_flush.store(true, Ordering::SeqCst);
351        Ok(bytes)
352    }
353}
354
355impl AsRawDescriptors for CompositeDiskFile {
356    fn as_raw_descriptors(&self) -> Vec<RawDescriptor> {
357        self.component_disks
358            .iter()
359            .flat_map(|d| d.file.as_raw_descriptors())
360            .collect()
361    }
362}
363
364struct AsyncComponentDiskPart {
365    file: Box<dyn AsyncDisk>,
366    offset: u64, // Location in the block device visible to the guest
367    length: u64,
368    file_offset: u64, // Location within the host file
369    needs_flush: AtomicBool,
370}
371
372pub struct AsyncCompositeDiskFile {
373    component_disks: Vec<AsyncComponentDiskPart>,
374}
375
376impl DiskGetLen for AsyncCompositeDiskFile {
377    fn get_len(&self) -> io::Result<u64> {
378        Ok(self.length())
379    }
380}
381
382impl FileSetLen for AsyncCompositeDiskFile {
383    fn set_len(&self, _len: u64) -> io::Result<()> {
384        Err(io::Error::other("unsupported operation"))
385    }
386}
387
388impl FileAllocate for AsyncCompositeDiskFile {
389    fn allocate(&self, offset: u64, length: u64) -> io::Result<()> {
390        let range = offset..(offset + length);
391        let disks = self
392            .component_disks
393            .iter()
394            .filter(|disk| ranges_overlap(&disk.range(), &range));
395        for disk in disks {
396            if let Some(intersection) = range_intersection(&range, &disk.range()) {
397                disk.file.allocate(
398                    intersection.start - disk.offset + disk.file_offset,
399                    intersection.end - intersection.start,
400                )?;
401                disk.needs_flush.store(true, Ordering::SeqCst);
402            }
403        }
404        Ok(())
405    }
406}
407
408impl ToAsyncDisk for CompositeDiskFile {
409    fn to_async_disk(self: Box<Self>, ex: &Executor) -> crate::Result<Box<dyn AsyncDisk>> {
410        Ok(Box::new(AsyncCompositeDiskFile {
411            component_disks: self
412                .component_disks
413                .into_iter()
414                .map(|disk| -> crate::Result<_> {
415                    Ok(AsyncComponentDiskPart {
416                        file: disk.file.to_async_disk(ex)?,
417                        offset: disk.offset,
418                        length: disk.length,
419                        file_offset: disk.file_offset,
420                        needs_flush: disk.needs_flush,
421                    })
422                })
423                .collect::<crate::Result<Vec<_>>>()?,
424        }))
425    }
426}
427
428impl AsyncComponentDiskPart {
429    fn range(&self) -> Range<u64> {
430        self.offset..(self.offset + self.length)
431    }
432
433    fn set_needs_flush(&self) {
434        self.needs_flush.store(true, Ordering::SeqCst);
435    }
436}
437
438impl AsyncCompositeDiskFile {
439    fn length(&self) -> u64 {
440        if let Some(disk) = self.component_disks.last() {
441            disk.offset + disk.length
442        } else {
443            0
444        }
445    }
446
447    fn disk_at_offset(&self, offset: u64) -> io::Result<&AsyncComponentDiskPart> {
448        self.component_disks
449            .iter()
450            .find(|disk| disk.range().contains(&offset))
451            .ok_or_else(|| {
452                io::Error::new(
453                    ErrorKind::InvalidData,
454                    format!("no disk at offset {offset}"),
455                )
456            })
457    }
458
459    fn disks_in_range<'a>(&'a self, range: &Range<u64>) -> Vec<&'a AsyncComponentDiskPart> {
460        self.component_disks
461            .iter()
462            .filter(|disk| ranges_overlap(&disk.range(), range))
463            .collect()
464    }
465}
466
467#[async_trait(?Send)]
468impl AsyncDisk for AsyncCompositeDiskFile {
469    async fn flush(&self) -> crate::Result<()> {
470        futures::future::try_join_all(self.component_disks.iter().map(|c| c.file.flush())).await?;
471        Ok(())
472    }
473
474    async fn fsync(&self) -> crate::Result<()> {
475        // NOTE: The fsync implementation isn't really async, so no point in adding concurrency
476        // here unless we introduce a blocking threadpool.
477        for disk in self.component_disks.iter() {
478            if disk.needs_flush.fetch_and(false, Ordering::SeqCst) {
479                if let Err(e) = disk.file.fsync().await {
480                    disk.set_needs_flush();
481                    return Err(e);
482                }
483            }
484        }
485        Ok(())
486    }
487
488    async fn fdatasync(&self) -> crate::Result<()> {
489        // NOTE: The fdatasync implementation isn't really async, so no point in adding concurrency
490        // here unless we introduce a blocking threadpool.
491        for disk in self.component_disks.iter() {
492            if disk.needs_flush.fetch_and(false, Ordering::SeqCst) {
493                if let Err(e) = disk.file.fdatasync().await {
494                    disk.set_needs_flush();
495                    return Err(e);
496                }
497            }
498        }
499        Ok(())
500    }
501
502    async fn read_to_mem<'a>(
503        &'a self,
504        file_offset: u64,
505        mem: Arc<dyn BackingMemory + Send + Sync>,
506        mem_offsets: MemRegionIter<'a>,
507        options: IoOptions,
508    ) -> crate::Result<usize> {
509        let disk = self
510            .disk_at_offset(file_offset)
511            .map_err(crate::Error::ReadingData)?;
512        let remaining_disk = disk.offset + disk.length - file_offset;
513        disk.file
514            .read_to_mem(
515                file_offset - disk.offset + disk.file_offset,
516                mem,
517                mem_offsets.take_bytes(remaining_disk.try_into().unwrap()),
518                options,
519            )
520            .await
521    }
522
523    async fn write_from_mem<'a>(
524        &'a self,
525        file_offset: u64,
526        mem: Arc<dyn BackingMemory + Send + Sync>,
527        mem_offsets: MemRegionIter<'a>,
528        options: IoOptions,
529    ) -> crate::Result<usize> {
530        let disk = self
531            .disk_at_offset(file_offset)
532            .map_err(crate::Error::ReadingData)?;
533        let remaining_disk = disk.offset + disk.length - file_offset;
534        let n = disk
535            .file
536            .write_from_mem(
537                file_offset - disk.offset + disk.file_offset,
538                mem,
539                mem_offsets.take_bytes(remaining_disk.try_into().unwrap()),
540                options,
541            )
542            .await?;
543        disk.set_needs_flush();
544        Ok(n)
545    }
546
547    async fn punch_hole(&self, file_offset: u64, length: u64) -> crate::Result<()> {
548        let range = file_offset..(file_offset + length);
549        let disks = self.disks_in_range(&range);
550        for disk in disks {
551            if let Some(intersection) = range_intersection(&range, &disk.range()) {
552                disk.file
553                    .punch_hole(
554                        intersection.start - disk.offset + disk.file_offset,
555                        intersection.end - intersection.start,
556                    )
557                    .await?;
558                disk.set_needs_flush();
559            }
560        }
561        Ok(())
562    }
563
564    async fn write_zeroes_at(&self, file_offset: u64, length: u64) -> crate::Result<()> {
565        let range = file_offset..(file_offset + length);
566        let disks = self.disks_in_range(&range);
567        for disk in disks {
568            if let Some(intersection) = range_intersection(&range, &disk.range()) {
569                disk.file
570                    .write_zeroes_at(
571                        intersection.start - disk.offset + disk.file_offset,
572                        intersection.end - intersection.start,
573                    )
574                    .await?;
575                disk.set_needs_flush();
576            }
577        }
578        Ok(())
579    }
580}
581
582/// Information about a partition to create.
583#[derive(Clone, Debug, Eq, PartialEq)]
584pub struct PartitionInfo {
585    pub label: String,
586    pub path: PathBuf,
587    pub partition_type: ImagePartitionType,
588    pub writable: bool,
589    pub size: u64,
590    pub part_guid: Option<Uuid>,
591}
592
593impl PartitionInfo {
594    fn aligned_size(&self) -> u64 {
595        self.size.next_multiple_of(1 << PARTITION_SIZE_SHIFT)
596    }
597}
598
599/// The type of partition.
600#[derive(Copy, Clone, Debug, Eq, PartialEq)]
601pub enum ImagePartitionType {
602    LinuxFilesystem,
603    EfiSystemPartition,
604}
605
606impl ImagePartitionType {
607    fn guid(self) -> Uuid {
608        match self {
609            Self::LinuxFilesystem => LINUX_FILESYSTEM_GUID,
610            Self::EfiSystemPartition => EFI_SYSTEM_PARTITION_GUID,
611        }
612    }
613}
614
615/// Write protective MBR and primary GPT table.
616fn write_beginning(
617    file: &mut impl Write,
618    disk_guid: Uuid,
619    partitions: &[u8],
620    partition_entries_crc32: u32,
621    secondary_table_offset: u64,
622    disk_size: u64,
623) -> Result<()> {
624    // Write the protective MBR to the first sector.
625    write_protective_mbr(file, disk_size)?;
626
627    // Write the GPT header, and pad out to the end of the sector.
628    write_gpt_header(
629        file,
630        disk_guid,
631        partition_entries_crc32,
632        secondary_table_offset,
633        false,
634    )?;
635    file.write_all(&[0; HEADER_PADDING_LENGTH])
636        .map_err(Error::WriteHeader)?;
637
638    // Write partition entries, including unused ones.
639    file.write_all(partitions).map_err(Error::WriteHeader)?;
640
641    // Write zeroes to align the first partition appropriately.
642    file.write_all(&[0; PARTITION_ALIGNMENT_SIZE])
643        .map_err(Error::WriteHeader)?;
644
645    Ok(())
646}
647
648/// Write secondary GPT table.
649fn write_end(
650    file: &mut impl Write,
651    disk_guid: Uuid,
652    partitions: &[u8],
653    partition_entries_crc32: u32,
654    secondary_table_offset: u64,
655) -> Result<()> {
656    // Write partition entries, including unused ones.
657    file.write_all(partitions).map_err(Error::WriteHeader)?;
658
659    // Write the GPT header, and pad out to the end of the sector.
660    write_gpt_header(
661        file,
662        disk_guid,
663        partition_entries_crc32,
664        secondary_table_offset,
665        true,
666    )?;
667    file.write_all(&[0; HEADER_PADDING_LENGTH])
668        .map_err(Error::WriteHeader)?;
669
670    Ok(())
671}
672
673/// Create the `GptPartitionEntry` for the given partition.
674fn create_gpt_entry(partition: &PartitionInfo, offset: u64) -> GptPartitionEntry {
675    let mut partition_name: Vec<u16> = partition.label.encode_utf16().collect();
676    partition_name.resize(36, 0);
677
678    GptPartitionEntry {
679        partition_type_guid: partition.partition_type.guid(),
680        unique_partition_guid: partition.part_guid.unwrap_or(Uuid::new_v4()),
681        first_lba: offset / SECTOR_SIZE,
682        last_lba: (offset + partition.aligned_size()) / SECTOR_SIZE - 1,
683        attributes: 0,
684        partition_name: partition_name.try_into().unwrap(),
685    }
686}
687
688/// Create one or more `ComponentDisk` proto messages for the given partition.
689fn create_component_disks(
690    partition: &PartitionInfo,
691    offset: u64,
692    zero_filler_path: &str,
693) -> Result<Vec<ComponentDisk>> {
694    let aligned_size = partition.aligned_size();
695
696    let mut component_disks = vec![ComponentDisk {
697        offset,
698        file_path: partition
699            .path
700            .to_str()
701            .ok_or_else(|| Error::InvalidPath(partition.path.to_owned()))?
702            .to_string(),
703        read_write_capability: if partition.writable {
704            ReadWriteCapability::READ_WRITE.into()
705        } else {
706            ReadWriteCapability::READ_ONLY.into()
707        },
708        ..ComponentDisk::new()
709    }];
710
711    if partition.size != aligned_size {
712        if partition.writable {
713            return Err(Error::UnalignedReadWrite(partition.to_owned()));
714        } else {
715            // Fill in the gap by reusing the zero filler file, because we know it is always bigger
716            // than the alignment size. Its size is 1 << PARTITION_SIZE_SHIFT (4k).
717            component_disks.push(ComponentDisk {
718                offset: offset + partition.size,
719                file_path: zero_filler_path.to_owned(),
720                read_write_capability: ReadWriteCapability::READ_ONLY.into(),
721                ..ComponentDisk::new()
722            });
723        }
724    }
725
726    Ok(component_disks)
727}
728
729/// Create a new composite disk image containing the given partitions, and write it out to the given
730/// files.
731pub fn create_composite_disk(
732    partitions: &[PartitionInfo],
733    zero_filler_path: &Path,
734    header_path: &Path,
735    header_file: &mut impl Write,
736    footer_path: &Path,
737    footer_file: &mut impl Write,
738    output_composite: &mut File,
739) -> Result<()> {
740    let zero_filler_path = zero_filler_path
741        .to_str()
742        .ok_or_else(|| Error::InvalidPath(zero_filler_path.to_owned()))?
743        .to_string();
744    let header_path = header_path
745        .to_str()
746        .ok_or_else(|| Error::InvalidPath(header_path.to_owned()))?
747        .to_string();
748    let footer_path = footer_path
749        .to_str()
750        .ok_or_else(|| Error::InvalidPath(footer_path.to_owned()))?
751        .to_string();
752
753    let mut composite_proto = CompositeDisk::new();
754    composite_proto.version = COMPOSITE_DISK_VERSION;
755    composite_proto.component_disks.push(ComponentDisk {
756        file_path: header_path,
757        offset: 0,
758        read_write_capability: ReadWriteCapability::READ_ONLY.into(),
759        ..ComponentDisk::new()
760    });
761
762    // Write partitions to a temporary buffer so that we can calculate the CRC, and construct the
763    // ComponentDisk proto messages at the same time.
764    let mut partitions_buffer =
765        [0u8; GPT_NUM_PARTITIONS as usize * GPT_PARTITION_ENTRY_SIZE as usize];
766    let mut writer: &mut [u8] = &mut partitions_buffer;
767    let mut next_disk_offset = GPT_BEGINNING_SIZE;
768    let mut labels = HashSet::with_capacity(partitions.len());
769    for partition in partitions {
770        let gpt_entry = create_gpt_entry(partition, next_disk_offset);
771        if !labels.insert(gpt_entry.partition_name) {
772            return Err(Error::DuplicatePartitionLabel(partition.label.clone()));
773        }
774        gpt_entry.write_bytes(&mut writer)?;
775
776        for component_disk in
777            create_component_disks(partition, next_disk_offset, &zero_filler_path)?
778        {
779            composite_proto.component_disks.push(component_disk);
780        }
781
782        next_disk_offset += partition.aligned_size();
783    }
784    // The secondary GPT needs to be at the very end of the file, but its size (0x4200) is not
785    // aligned to the chosen partition size (0x1000). We compensate for that by writing some
786    // padding to the start of the footer file.
787    const FOOTER_PADDING: u64 =
788        GPT_END_SIZE.next_multiple_of(1 << PARTITION_SIZE_SHIFT) - GPT_END_SIZE;
789    let footer_file_offset = next_disk_offset;
790    let secondary_table_offset = footer_file_offset + FOOTER_PADDING;
791    let disk_size = secondary_table_offset + GPT_END_SIZE;
792    composite_proto.component_disks.push(ComponentDisk {
793        file_path: footer_path,
794        offset: footer_file_offset,
795        read_write_capability: ReadWriteCapability::READ_ONLY.into(),
796        ..ComponentDisk::new()
797    });
798
799    // Calculate CRC32 of partition entries.
800    let mut hasher = Hasher::new();
801    hasher.update(&partitions_buffer);
802    let partition_entries_crc32 = hasher.finalize();
803
804    let disk_guid = Uuid::new_v4();
805    write_beginning(
806        header_file,
807        disk_guid,
808        &partitions_buffer,
809        partition_entries_crc32,
810        secondary_table_offset,
811        disk_size,
812    )?;
813
814    footer_file
815        .write_all(&[0; FOOTER_PADDING as usize])
816        .map_err(Error::WriteHeader)?;
817    write_end(
818        footer_file,
819        disk_guid,
820        &partitions_buffer,
821        partition_entries_crc32,
822        secondary_table_offset,
823    )?;
824
825    composite_proto.length = disk_size;
826    output_composite
827        .write_all(CDISK_MAGIC.as_bytes())
828        .map_err(Error::WriteHeader)?;
829    composite_proto
830        .write_to_writer(output_composite)
831        .map_err(Error::WriteProto)?;
832
833    Ok(())
834}
835
836/// Create a zero filler file which can be used to fill the gaps between partition files.
837/// The filler is sized to be big enough to fill the gaps. (1 << PARTITION_SIZE_SHIFT)
838pub fn create_zero_filler<P: AsRef<Path>>(zero_filler_path: P) -> Result<()> {
839    let f = OpenOptions::new()
840        .create(true)
841        .read(true)
842        .write(true)
843        .truncate(true)
844        .open(zero_filler_path.as_ref())
845        .map_err(Error::WriteZeroFiller)?;
846    f.set_len(1 << PARTITION_SIZE_SHIFT)
847        .map_err(Error::WriteZeroFiller)
848}
849
850#[cfg(test)]
851mod tests {
852    use std::fs::OpenOptions;
853    use std::io::Write;
854    use std::matches;
855
856    use base::AsRawDescriptor;
857    use tempfile::tempfile;
858
859    use super::*;
860
861    fn new_from_components(disks: Vec<ComponentDiskPart>) -> Result<CompositeDiskFile> {
862        CompositeDiskFile::new(disks, tempfile().unwrap())
863    }
864
865    #[test]
866    fn block_duplicate_offset_disks() {
867        let file1 = tempfile().unwrap();
868        let file2 = tempfile().unwrap();
869        let disk_part1 = ComponentDiskPart {
870            file: Box::new(file1),
871            offset: 0,
872            length: 100,
873            file_offset: 0,
874            needs_flush: AtomicBool::new(false),
875        };
876        let disk_part2 = ComponentDiskPart {
877            file: Box::new(file2),
878            offset: 0,
879            length: 100,
880            file_offset: 0,
881            needs_flush: AtomicBool::new(false),
882        };
883        assert!(new_from_components(vec![disk_part1, disk_part2]).is_err());
884    }
885
886    #[test]
887    fn get_len() {
888        let file1 = tempfile().unwrap();
889        let file2 = tempfile().unwrap();
890        let disk_part1 = ComponentDiskPart {
891            file: Box::new(file1),
892            offset: 0,
893            length: 100,
894            file_offset: 0,
895            needs_flush: AtomicBool::new(false),
896        };
897        let disk_part2 = ComponentDiskPart {
898            file: Box::new(file2),
899            offset: 100,
900            length: 100,
901            file_offset: 0,
902            needs_flush: AtomicBool::new(false),
903        };
904        let composite = new_from_components(vec![disk_part1, disk_part2]).unwrap();
905        let len = composite.get_len().unwrap();
906        assert_eq!(len, 200);
907    }
908
909    #[test]
910    fn async_get_len() {
911        let file1 = tempfile().unwrap();
912        let file2 = tempfile().unwrap();
913        let disk_part1 = ComponentDiskPart {
914            file: Box::new(file1),
915            offset: 0,
916            length: 100,
917            file_offset: 0,
918            needs_flush: AtomicBool::new(false),
919        };
920        let disk_part2 = ComponentDiskPart {
921            file: Box::new(file2),
922            offset: 100,
923            length: 100,
924            file_offset: 0,
925            needs_flush: AtomicBool::new(false),
926        };
927        let composite = new_from_components(vec![disk_part1, disk_part2]).unwrap();
928
929        let ex = Executor::new().unwrap();
930        let composite = Box::new(composite).to_async_disk(&ex).unwrap();
931        let len = composite.get_len().unwrap();
932        assert_eq!(len, 200);
933    }
934
935    #[test]
936    fn single_file_passthrough() {
937        let file = tempfile().unwrap();
938        let disk_part = ComponentDiskPart {
939            file: Box::new(file),
940            offset: 0,
941            length: 100,
942            file_offset: 0,
943            needs_flush: AtomicBool::new(false),
944        };
945        let composite = new_from_components(vec![disk_part]).unwrap();
946        let mut input_memory = [55u8; 5];
947        let input_volatile_memory = VolatileSlice::new(&mut input_memory[..]);
948        composite
949            .write_all_at_volatile(input_volatile_memory, 0)
950            .unwrap();
951        let mut output_memory = [0u8; 5];
952        let output_volatile_memory = VolatileSlice::new(&mut output_memory[..]);
953        composite
954            .read_exact_at_volatile(output_volatile_memory, 0)
955            .unwrap();
956        assert_eq!(input_memory, output_memory);
957    }
958
959    #[test]
960    fn single_file_passthrough_file_offset() {
961        let file = tempfile().unwrap();
962        let mut input_memory = [55u8, 56u8, 57u8, 58u8, 59u8];
963        let input_volatile_memory = VolatileSlice::new(&mut input_memory[..]);
964        file.write_all_at_volatile(input_volatile_memory, 0)
965            .unwrap();
966
967        let disk_part = ComponentDiskPart {
968            file: Box::new(file),
969            offset: 0,
970            length: 100,
971            file_offset: 2,
972            needs_flush: AtomicBool::new(false),
973        };
974        let composite = new_from_components(vec![disk_part]).unwrap();
975        let mut output_memory = [0u8; 3];
976        let output_volatile_memory = VolatileSlice::new(&mut output_memory[..]);
977        composite
978            .read_exact_at_volatile(output_volatile_memory, 0)
979            .unwrap();
980        assert_eq!(input_memory[2..], output_memory);
981    }
982
983    #[test]
984    fn async_single_file_passthrough() {
985        let file = tempfile().unwrap();
986        let disk_part = ComponentDiskPart {
987            file: Box::new(file),
988            offset: 0,
989            length: 100,
990            file_offset: 0,
991            needs_flush: AtomicBool::new(false),
992        };
993        let composite = new_from_components(vec![disk_part]).unwrap();
994        let ex = Executor::new().unwrap();
995        ex.run_until(async {
996            let composite = Box::new(composite).to_async_disk(&ex).unwrap();
997            let expected = [55u8; 5];
998            assert_eq!(
999                composite.write_double_buffered(0, &expected).await.unwrap(),
1000                5
1001            );
1002            let mut buf = [0u8; 5];
1003            assert_eq!(
1004                composite
1005                    .read_double_buffered(0, &mut buf[..])
1006                    .await
1007                    .unwrap(),
1008                5
1009            );
1010            assert_eq!(buf, expected);
1011        })
1012        .unwrap();
1013    }
1014
1015    #[test]
1016    fn async_single_file_passthrough_offset() {
1017        let file = tempfile().unwrap();
1018        let mut input_memory = [55u8, 56u8, 57u8, 58u8, 59u8];
1019        let input_volatile_memory = VolatileSlice::new(&mut input_memory[..]);
1020        file.write_all_at_volatile(input_volatile_memory, 0)
1021            .unwrap();
1022
1023        let disk_part = ComponentDiskPart {
1024            file: Box::new(file),
1025            offset: 0,
1026            length: 100,
1027            file_offset: 2,
1028            needs_flush: AtomicBool::new(false),
1029        };
1030        let composite = new_from_components(vec![disk_part]).unwrap();
1031        let ex = Executor::new().unwrap();
1032        ex.run_until(async {
1033            let composite = Box::new(composite).to_async_disk(&ex).unwrap();
1034            let mut buf = [0u8; 3];
1035            assert_eq!(
1036                composite
1037                    .read_double_buffered(0, &mut buf[..])
1038                    .await
1039                    .unwrap(),
1040                3
1041            );
1042            assert_eq!(input_memory[2..], buf);
1043        })
1044        .unwrap();
1045    }
1046
1047    #[test]
1048    fn triple_file_descriptors() {
1049        let file1 = tempfile().unwrap();
1050        let file2 = tempfile().unwrap();
1051        let file3 = tempfile().unwrap();
1052        let mut in_descriptors = vec![
1053            file1.as_raw_descriptor(),
1054            file2.as_raw_descriptor(),
1055            file3.as_raw_descriptor(),
1056        ];
1057        in_descriptors.sort_unstable();
1058        let disk_part1 = ComponentDiskPart {
1059            file: Box::new(file1),
1060            offset: 0,
1061            length: 100,
1062            file_offset: 0,
1063            needs_flush: AtomicBool::new(false),
1064        };
1065        let disk_part2 = ComponentDiskPart {
1066            file: Box::new(file2),
1067            offset: 100,
1068            length: 100,
1069            file_offset: 0,
1070            needs_flush: AtomicBool::new(false),
1071        };
1072        let disk_part3 = ComponentDiskPart {
1073            file: Box::new(file3),
1074            offset: 200,
1075            length: 100,
1076            file_offset: 0,
1077            needs_flush: AtomicBool::new(false),
1078        };
1079        let composite = new_from_components(vec![disk_part1, disk_part2, disk_part3]).unwrap();
1080        let mut out_descriptors = composite.as_raw_descriptors();
1081        out_descriptors.sort_unstable();
1082        assert_eq!(in_descriptors, out_descriptors);
1083    }
1084
1085    #[test]
1086    fn triple_file_passthrough() {
1087        let file1 = tempfile().unwrap();
1088        let file2 = tempfile().unwrap();
1089        let file3 = tempfile().unwrap();
1090        let disk_part1 = ComponentDiskPart {
1091            file: Box::new(file1),
1092            offset: 0,
1093            length: 100,
1094            file_offset: 0,
1095            needs_flush: AtomicBool::new(false),
1096        };
1097        let disk_part2 = ComponentDiskPart {
1098            file: Box::new(file2),
1099            offset: 100,
1100            length: 100,
1101            file_offset: 0,
1102            needs_flush: AtomicBool::new(false),
1103        };
1104        let disk_part3 = ComponentDiskPart {
1105            file: Box::new(file3),
1106            offset: 200,
1107            length: 100,
1108            file_offset: 0,
1109            needs_flush: AtomicBool::new(false),
1110        };
1111        let composite = new_from_components(vec![disk_part1, disk_part2, disk_part3]).unwrap();
1112        let mut input_memory = [55u8; 200];
1113        let input_volatile_memory = VolatileSlice::new(&mut input_memory[..]);
1114        composite
1115            .write_all_at_volatile(input_volatile_memory, 50)
1116            .unwrap();
1117        let mut output_memory = [0u8; 200];
1118        let output_volatile_memory = VolatileSlice::new(&mut output_memory[..]);
1119        composite
1120            .read_exact_at_volatile(output_volatile_memory, 50)
1121            .unwrap();
1122        assert!(input_memory.iter().eq(output_memory.iter()));
1123    }
1124
1125    #[test]
1126    fn async_triple_file_passthrough() {
1127        let file1 = tempfile().unwrap();
1128        let file2 = tempfile().unwrap();
1129        let file3 = tempfile().unwrap();
1130        let disk_part1 = ComponentDiskPart {
1131            file: Box::new(file1),
1132            offset: 0,
1133            length: 100,
1134            file_offset: 0,
1135            needs_flush: AtomicBool::new(false),
1136        };
1137        let disk_part2 = ComponentDiskPart {
1138            file: Box::new(file2),
1139            offset: 100,
1140            length: 100,
1141            file_offset: 0,
1142            needs_flush: AtomicBool::new(false),
1143        };
1144        let disk_part3 = ComponentDiskPart {
1145            file: Box::new(file3),
1146            offset: 200,
1147            length: 100,
1148            file_offset: 0,
1149            needs_flush: AtomicBool::new(false),
1150        };
1151        let composite = new_from_components(vec![disk_part1, disk_part2, disk_part3]).unwrap();
1152        let ex = Executor::new().unwrap();
1153        ex.run_until(async {
1154            let composite = Box::new(composite).to_async_disk(&ex).unwrap();
1155
1156            let expected = [55u8; 200];
1157            assert_eq!(
1158                composite.write_double_buffered(0, &expected).await.unwrap(),
1159                100
1160            );
1161            assert_eq!(
1162                composite
1163                    .write_double_buffered(100, &expected[100..])
1164                    .await
1165                    .unwrap(),
1166                100
1167            );
1168
1169            let mut buf = [0u8; 200];
1170            assert_eq!(
1171                composite
1172                    .read_double_buffered(0, &mut buf[..])
1173                    .await
1174                    .unwrap(),
1175                100
1176            );
1177            assert_eq!(
1178                composite
1179                    .read_double_buffered(100, &mut buf[100..])
1180                    .await
1181                    .unwrap(),
1182                100
1183            );
1184            assert_eq!(buf, expected);
1185        })
1186        .unwrap();
1187    }
1188
1189    #[test]
1190    fn async_triple_file_punch_hole() {
1191        let file1 = tempfile().unwrap();
1192        let file2 = tempfile().unwrap();
1193        let file3 = tempfile().unwrap();
1194        let disk_part1 = ComponentDiskPart {
1195            file: Box::new(file1),
1196            offset: 0,
1197            length: 100,
1198            file_offset: 0,
1199            needs_flush: AtomicBool::new(false),
1200        };
1201        let disk_part2 = ComponentDiskPart {
1202            file: Box::new(file2),
1203            offset: 100,
1204            length: 100,
1205            file_offset: 0,
1206            needs_flush: AtomicBool::new(false),
1207        };
1208        let disk_part3 = ComponentDiskPart {
1209            file: Box::new(file3),
1210            offset: 200,
1211            length: 100,
1212            file_offset: 0,
1213            needs_flush: AtomicBool::new(false),
1214        };
1215        let composite = new_from_components(vec![disk_part1, disk_part2, disk_part3]).unwrap();
1216        let ex = Executor::new().unwrap();
1217        ex.run_until(async {
1218            let composite = Box::new(composite).to_async_disk(&ex).unwrap();
1219
1220            let input = [55u8; 300];
1221            assert_eq!(
1222                composite.write_double_buffered(0, &input).await.unwrap(),
1223                100
1224            );
1225            assert_eq!(
1226                composite
1227                    .write_double_buffered(100, &input[100..])
1228                    .await
1229                    .unwrap(),
1230                100
1231            );
1232            assert_eq!(
1233                composite
1234                    .write_double_buffered(200, &input[200..])
1235                    .await
1236                    .unwrap(),
1237                100
1238            );
1239
1240            composite.punch_hole(50, 200).await.unwrap();
1241
1242            let mut buf = [0u8; 300];
1243            assert_eq!(
1244                composite
1245                    .read_double_buffered(0, &mut buf[..])
1246                    .await
1247                    .unwrap(),
1248                100
1249            );
1250            assert_eq!(
1251                composite
1252                    .read_double_buffered(100, &mut buf[100..])
1253                    .await
1254                    .unwrap(),
1255                100
1256            );
1257            assert_eq!(
1258                composite
1259                    .read_double_buffered(200, &mut buf[200..])
1260                    .await
1261                    .unwrap(),
1262                100
1263            );
1264
1265            let mut expected = input;
1266            expected[50..250].iter_mut().for_each(|x| *x = 0);
1267            assert_eq!(buf, expected);
1268        })
1269        .unwrap();
1270    }
1271
1272    #[test]
1273    fn async_triple_file_write_zeroes() {
1274        let file1 = tempfile().unwrap();
1275        let file2 = tempfile().unwrap();
1276        let file3 = tempfile().unwrap();
1277        let disk_part1 = ComponentDiskPart {
1278            file: Box::new(file1),
1279            offset: 0,
1280            length: 100,
1281            file_offset: 0,
1282            needs_flush: AtomicBool::new(false),
1283        };
1284        let disk_part2 = ComponentDiskPart {
1285            file: Box::new(file2),
1286            offset: 100,
1287            length: 100,
1288            file_offset: 0,
1289            needs_flush: AtomicBool::new(false),
1290        };
1291        let disk_part3 = ComponentDiskPart {
1292            file: Box::new(file3),
1293            offset: 200,
1294            length: 100,
1295            file_offset: 0,
1296            needs_flush: AtomicBool::new(false),
1297        };
1298        let composite = new_from_components(vec![disk_part1, disk_part2, disk_part3]).unwrap();
1299        let ex = Executor::new().unwrap();
1300        ex.run_until(async {
1301            let composite = Box::new(composite).to_async_disk(&ex).unwrap();
1302
1303            let input = [55u8; 300];
1304            assert_eq!(
1305                composite.write_double_buffered(0, &input).await.unwrap(),
1306                100
1307            );
1308            assert_eq!(
1309                composite
1310                    .write_double_buffered(100, &input[100..])
1311                    .await
1312                    .unwrap(),
1313                100
1314            );
1315            assert_eq!(
1316                composite
1317                    .write_double_buffered(200, &input[200..])
1318                    .await
1319                    .unwrap(),
1320                100
1321            );
1322
1323            composite.write_zeroes_at(50, 200).await.unwrap();
1324
1325            let mut buf = [0u8; 300];
1326            assert_eq!(
1327                composite
1328                    .read_double_buffered(0, &mut buf[..])
1329                    .await
1330                    .unwrap(),
1331                100
1332            );
1333            assert_eq!(
1334                composite
1335                    .read_double_buffered(100, &mut buf[100..])
1336                    .await
1337                    .unwrap(),
1338                100
1339            );
1340            assert_eq!(
1341                composite
1342                    .read_double_buffered(200, &mut buf[200..])
1343                    .await
1344                    .unwrap(),
1345                100
1346            );
1347
1348            let mut expected = input;
1349            expected[50..250].iter_mut().for_each(|x| *x = 0);
1350            assert_eq!(buf, expected);
1351        })
1352        .unwrap();
1353    }
1354
1355    // TODO: fsync on a RO file is legal, this test doesn't work as expected. Consider using a mock
1356    // DiskFile to detect the fsync calls.
1357    #[test]
1358    fn async_fsync_skips_unchanged_parts() {
1359        let mut rw_file = tempfile().unwrap();
1360        rw_file.write_all(&[0u8; 100]).unwrap();
1361        rw_file.seek(SeekFrom::Start(0)).unwrap();
1362        let mut ro_disk_image = tempfile::NamedTempFile::new().unwrap();
1363        ro_disk_image.write_all(&[0u8; 100]).unwrap();
1364        let ro_file = OpenOptions::new()
1365            .read(true)
1366            .open(ro_disk_image.path())
1367            .unwrap();
1368
1369        let rw_part = ComponentDiskPart {
1370            file: Box::new(rw_file),
1371            offset: 0,
1372            length: 100,
1373            file_offset: 0,
1374            needs_flush: AtomicBool::new(false),
1375        };
1376        let ro_part = ComponentDiskPart {
1377            file: Box::new(ro_file),
1378            offset: 100,
1379            length: 100,
1380            file_offset: 0,
1381            needs_flush: AtomicBool::new(false),
1382        };
1383        let composite = new_from_components(vec![rw_part, ro_part]).unwrap();
1384        let ex = Executor::new().unwrap();
1385        ex.run_until(async {
1386            let composite = Box::new(composite).to_async_disk(&ex).unwrap();
1387
1388            // Write to the RW part so that some fsync operation will occur.
1389            composite.write_zeroes_at(0, 20).await.unwrap();
1390
1391            // This is the test's assert. fsyncing should NOT touch a read-only disk part. On
1392            // Windows, this would be an error.
1393            composite.fsync().await.expect(
1394                "Failed to fsync composite disk. \
1395                     This can happen if the disk writable state is wrong.",
1396            );
1397        })
1398        .unwrap();
1399    }
1400
1401    #[test]
1402    fn beginning_size() {
1403        let mut buffer = vec![];
1404        let partitions = [0u8; GPT_NUM_PARTITIONS as usize * GPT_PARTITION_ENTRY_SIZE as usize];
1405        let disk_size = 1000 * SECTOR_SIZE;
1406        write_beginning(
1407            &mut buffer,
1408            Uuid::from_u128(0x12345678_1234_5678_abcd_12345678abcd),
1409            &partitions,
1410            42,
1411            disk_size - GPT_END_SIZE,
1412            disk_size,
1413        )
1414        .unwrap();
1415
1416        assert_eq!(buffer.len(), GPT_BEGINNING_SIZE as usize);
1417    }
1418
1419    #[test]
1420    fn end_size() {
1421        let mut buffer = vec![];
1422        let partitions = [0u8; GPT_NUM_PARTITIONS as usize * GPT_PARTITION_ENTRY_SIZE as usize];
1423        let disk_size = 1000 * SECTOR_SIZE;
1424        write_end(
1425            &mut buffer,
1426            Uuid::from_u128(0x12345678_1234_5678_abcd_12345678abcd),
1427            &partitions,
1428            42,
1429            disk_size - GPT_END_SIZE,
1430        )
1431        .unwrap();
1432
1433        assert_eq!(buffer.len(), GPT_END_SIZE as usize);
1434    }
1435
1436    /// Creates a composite disk image with no partitions.
1437    #[test]
1438    fn create_composite_disk_empty() {
1439        let mut header_image = tempfile().unwrap();
1440        let mut footer_image = tempfile().unwrap();
1441        let mut composite_image = tempfile().unwrap();
1442
1443        create_composite_disk(
1444            &[],
1445            Path::new("/zero_filler.img"),
1446            Path::new("/header_path.img"),
1447            &mut header_image,
1448            Path::new("/footer_path.img"),
1449            &mut footer_image,
1450            &mut composite_image,
1451        )
1452        .unwrap();
1453    }
1454
1455    /// Creates a composite disk image with two partitions.
1456    #[test]
1457    #[allow(clippy::unnecessary_to_owned)] // false positives
1458    fn create_composite_disk_success() {
1459        fn tmpfile(prefix: &str) -> tempfile::NamedTempFile {
1460            tempfile::Builder::new().prefix(prefix).tempfile().unwrap()
1461        }
1462
1463        let mut header_image = tmpfile("header");
1464        let mut footer_image = tmpfile("footer");
1465        let mut composite_image = tmpfile("composite");
1466
1467        // The test doesn't read these, just needs to be able to open them.
1468        let partition1 = tmpfile("partition1");
1469        let partition2 = tmpfile("partition1");
1470        let zero_filler = tmpfile("zero");
1471
1472        create_composite_disk(
1473            &[
1474                PartitionInfo {
1475                    label: "partition1".to_string(),
1476                    path: partition1.path().to_path_buf(),
1477                    partition_type: ImagePartitionType::LinuxFilesystem,
1478                    writable: false,
1479                    // Needs small amount of padding.
1480                    size: 4000,
1481                    part_guid: None,
1482                },
1483                PartitionInfo {
1484                    label: "partition2".to_string(),
1485                    path: partition2.path().to_path_buf(),
1486                    partition_type: ImagePartitionType::LinuxFilesystem,
1487                    writable: true,
1488                    // Needs no padding.
1489                    size: 4096,
1490                    part_guid: Some(Uuid::from_u128(0x4049C8DC_6C2B_C740_A95A_BDAA629D4378)),
1491                },
1492            ],
1493            zero_filler.path(),
1494            &header_image.path().to_path_buf(),
1495            header_image.as_file_mut(),
1496            &footer_image.path().to_path_buf(),
1497            footer_image.as_file_mut(),
1498            composite_image.as_file_mut(),
1499        )
1500        .unwrap();
1501
1502        // Check magic.
1503        composite_image.rewind().unwrap();
1504        let mut magic_space = [0u8; CDISK_MAGIC.len()];
1505        composite_image.read_exact(&mut magic_space[..]).unwrap();
1506        assert_eq!(magic_space, CDISK_MAGIC.as_bytes());
1507        // Check proto.
1508        let proto = CompositeDisk::parse_from_reader(&mut composite_image).unwrap();
1509        assert_eq!(
1510            proto,
1511            CompositeDisk {
1512                version: 2,
1513                component_disks: vec![
1514                    ComponentDisk {
1515                        file_path: header_image.path().to_str().unwrap().to_string(),
1516                        offset: 0,
1517                        read_write_capability: ReadWriteCapability::READ_ONLY.into(),
1518                        ..ComponentDisk::new()
1519                    },
1520                    ComponentDisk {
1521                        file_path: partition1.path().to_str().unwrap().to_string(),
1522                        offset: 0x5000, // GPT_BEGINNING_SIZE,
1523                        read_write_capability: ReadWriteCapability::READ_ONLY.into(),
1524                        ..ComponentDisk::new()
1525                    },
1526                    ComponentDisk {
1527                        file_path: zero_filler.path().to_str().unwrap().to_string(),
1528                        offset: 0x5fa0, // GPT_BEGINNING_SIZE + 4000,
1529                        read_write_capability: ReadWriteCapability::READ_ONLY.into(),
1530                        ..ComponentDisk::new()
1531                    },
1532                    ComponentDisk {
1533                        file_path: partition2.path().to_str().unwrap().to_string(),
1534                        offset: 0x6000, // GPT_BEGINNING_SIZE + 4096,
1535                        read_write_capability: ReadWriteCapability::READ_WRITE.into(),
1536                        ..ComponentDisk::new()
1537                    },
1538                    ComponentDisk {
1539                        file_path: footer_image.path().to_str().unwrap().to_string(),
1540                        offset: 0x7000, // GPT_BEGINNING_SIZE + 4096 + 4096,
1541                        read_write_capability: ReadWriteCapability::READ_ONLY.into(),
1542                        ..ComponentDisk::new()
1543                    },
1544                ],
1545                length: 0xc000,
1546                ..CompositeDisk::new()
1547            }
1548        );
1549
1550        // Open the file as a composite disk and do some basic GPT header/footer validation.
1551        let ex = Executor::new().unwrap();
1552        ex.run_until(async {
1553            let disk = Box::new(
1554                CompositeDiskFile::from_file(
1555                    composite_image.into_file(),
1556                    DiskFileParams {
1557                        path: "/foo".into(),
1558                        is_read_only: true,
1559                        lock: false,
1560                        ..Default::default()
1561                    },
1562                )
1563                .unwrap(),
1564            )
1565            .to_async_disk(&ex)
1566            .unwrap();
1567
1568            let header_offset = SECTOR_SIZE;
1569            let footer_offset = disk.get_len().unwrap() - SECTOR_SIZE;
1570
1571            let mut header_bytes = [0u8; SECTOR_SIZE as usize];
1572            assert_eq!(
1573                disk.read_double_buffered(header_offset, &mut header_bytes[..])
1574                    .await
1575                    .unwrap(),
1576                SECTOR_SIZE as usize
1577            );
1578
1579            let mut footer_bytes = [0u8; SECTOR_SIZE as usize];
1580            assert_eq!(
1581                disk.read_double_buffered(footer_offset, &mut footer_bytes[..])
1582                    .await
1583                    .unwrap(),
1584                SECTOR_SIZE as usize
1585            );
1586
1587            // Check the header and footer fields point to each other correctly.
1588            let header_current_lba = u64::from_le_bytes(header_bytes[24..32].try_into().unwrap());
1589            assert_eq!(header_current_lba * SECTOR_SIZE, header_offset);
1590            let header_backup_lba = u64::from_le_bytes(header_bytes[32..40].try_into().unwrap());
1591            assert_eq!(header_backup_lba * SECTOR_SIZE, footer_offset);
1592
1593            let footer_current_lba = u64::from_le_bytes(footer_bytes[24..32].try_into().unwrap());
1594            assert_eq!(footer_current_lba * SECTOR_SIZE, footer_offset);
1595            let footer_backup_lba = u64::from_le_bytes(footer_bytes[32..40].try_into().unwrap());
1596            assert_eq!(footer_backup_lba * SECTOR_SIZE, header_offset);
1597
1598            // Header and footer should be equal if we zero the pointers and CRCs.
1599            header_bytes[16..20].fill(0);
1600            header_bytes[24..40].fill(0);
1601            footer_bytes[16..20].fill(0);
1602            footer_bytes[24..40].fill(0);
1603            assert_eq!(header_bytes, footer_bytes);
1604        })
1605        .unwrap();
1606    }
1607
1608    /// Attempts to create a composite disk image with two partitions with the same label.
1609    #[test]
1610    fn create_composite_disk_duplicate_label() {
1611        let mut header_image = tempfile().unwrap();
1612        let mut footer_image = tempfile().unwrap();
1613        let mut composite_image = tempfile().unwrap();
1614
1615        let result = create_composite_disk(
1616            &[
1617                PartitionInfo {
1618                    label: "label".to_string(),
1619                    path: "/partition1.img".to_string().into(),
1620                    partition_type: ImagePartitionType::LinuxFilesystem,
1621                    writable: false,
1622                    size: 0,
1623                    part_guid: None,
1624                },
1625                PartitionInfo {
1626                    label: "label".to_string(),
1627                    path: "/partition2.img".to_string().into(),
1628                    partition_type: ImagePartitionType::LinuxFilesystem,
1629                    writable: true,
1630                    size: 0,
1631                    part_guid: None,
1632                },
1633            ],
1634            Path::new("/zero_filler.img"),
1635            Path::new("/header_path.img"),
1636            &mut header_image,
1637            Path::new("/footer_path.img"),
1638            &mut footer_image,
1639            &mut composite_image,
1640        );
1641        assert!(matches!(result, Err(Error::DuplicatePartitionLabel(label)) if label == "label"));
1642    }
1643}