petri/
disk_image.rs

1// Copyright (c) Microsoft Corporation.
2// Licensed under the MIT License.
3
4//! Tools for building a disk image for a VM.
5
6use anyhow::Context;
7use fatfs::FormatVolumeOptions;
8use fatfs::FsOptions;
9use guid::Guid;
10use petri_artifacts_common::artifacts as common_artifacts;
11use petri_artifacts_common::tags::MachineArch;
12use petri_artifacts_common::tags::OsFlavor;
13use petri_artifacts_core::ArtifactResolver;
14use petri_artifacts_core::ResolvedArtifact;
15use std::io::Read;
16use std::io::Seek;
17use std::io::Write;
18use std::ops::Range;
19use std::path::Path;
20use std::path::PathBuf;
21use tempfile::TempPath;
22
23/// The description and artifacts needed to build a pipette disk image for a VM.
24#[derive(Debug)]
25pub struct AgentImage {
26    os_flavor: OsFlavor,
27    pipette: Option<ResolvedArtifact>,
28    extras: Vec<(String, ImageEntry)>,
29    free_space: u64,
30}
31
32#[derive(Debug)]
33/// A file or directory in the agent image
34pub enum ImageEntry {
35    /// Directory
36    Dir,
37    /// Persistent file
38    File(PathBuf),
39    /// Owned temporary file path
40    TempFile(TempPath),
41}
42
43impl<'a> ImageEntry {
44    fn as_content(&'a self) -> ImageEntryContent<'a> {
45        match self {
46            ImageEntry::Dir => ImageEntryContent::Dir,
47            ImageEntry::File(path_buf) => ImageEntryContent::Path(path_buf.as_path()),
48            ImageEntry::TempFile(temp_path) => ImageEntryContent::Path(temp_path.as_ref()),
49        }
50    }
51}
52
53/// Disk image type
54pub enum ImageType {
55    /// Raw image
56    Raw,
57    /// Fixed VHD1
58    Vhd,
59}
60
61impl AgentImage {
62    /// Resolves the artifacts needed to build a disk image for a VM.
63    pub fn new(os_flavor: OsFlavor) -> Self {
64        Self {
65            os_flavor,
66            pipette: None,
67            extras: Vec::new(),
68            free_space: 0,
69        }
70    }
71
72    /// Adds the appropriate pipette binary to the image
73    pub fn with_pipette(mut self, resolver: &ArtifactResolver<'_>, arch: MachineArch) -> Self {
74        self.pipette = match (self.os_flavor, arch) {
75            (OsFlavor::Windows, MachineArch::X86_64) => Some(
76                resolver
77                    .require(common_artifacts::PIPETTE_WINDOWS_X64)
78                    .erase(),
79            ),
80            (OsFlavor::Linux, MachineArch::X86_64) => Some(
81                resolver
82                    .require(common_artifacts::PIPETTE_LINUX_X64_MUSL)
83                    .erase(),
84            ),
85            (OsFlavor::Windows, MachineArch::Aarch64) => Some(
86                resolver
87                    .require(common_artifacts::PIPETTE_WINDOWS_AARCH64)
88                    .erase(),
89            ),
90            (OsFlavor::Linux, MachineArch::Aarch64) => Some(
91                resolver
92                    .require(common_artifacts::PIPETTE_LINUX_AARCH64_MUSL)
93                    .erase(),
94            ),
95            (OsFlavor::FreeBsd | OsFlavor::Uefi, _) => {
96                todo!("No pipette binary yet for os");
97            }
98        };
99        self
100    }
101
102    /// Check if the image contains pipette
103    pub fn contains_pipette(&self) -> bool {
104        self.pipette.is_some()
105    }
106
107    /// Check if the image has extra files beyond pipette and cloud-init.
108    pub fn has_extras(&self) -> bool {
109        !self.extras.is_empty()
110    }
111
112    /// Adds an extra file to the disk image.
113    pub fn add_file(&mut self, name: &str, file_path: impl AsRef<Path>) {
114        self.extras.push((
115            name.into(),
116            ImageEntry::File(file_path.as_ref().to_path_buf()),
117        ));
118    }
119
120    /// Add collection of image entries
121    pub fn add_extras(&mut self, extras: Vec<(String, ImageEntry)>) {
122        self.extras.extend(extras);
123    }
124
125    /// Allocate some extra free space in the image
126    pub fn allocate_free_space(&mut self, space: u64) {
127        self.free_space = space;
128    }
129
130    /// Builds a disk image containing pipette and any files needed for the guest VM
131    /// to run pipette.
132    pub fn build(&self, image_type: ImageType) -> anyhow::Result<Option<tempfile::NamedTempFile>> {
133        let mut files = self
134            .extras
135            .iter()
136            .map(|(name, file_path)| (name.as_str(), file_path.as_content()))
137            .collect::<Vec<_>>();
138        let volume_label = match self.os_flavor {
139            OsFlavor::Windows => {
140                // Windows doesn't use cloud-init, so we only need pipette
141                // (which is configured via the IMC hive).
142                if let Some(pipette) = self.pipette.as_ref() {
143                    files.push(("pipette.exe", ImageEntryContent::Path(pipette.as_ref())));
144                }
145                b"pipette    "
146            }
147            OsFlavor::Linux => {
148                if let Some(pipette) = self.pipette.as_ref() {
149                    files.push(("pipette", ImageEntryContent::Path(pipette.as_ref())));
150                }
151                // Linux uses cloud-init, so we need to include the cloud-init
152                // configuration files as well.
153                files.extend([
154                    (
155                        "meta-data",
156                        ImageEntryContent::Binary(include_bytes!("../guest-bootstrap/meta-data")),
157                    ),
158                    (
159                        "user-data",
160                        if self.pipette.is_some() {
161                            ImageEntryContent::Binary(include_bytes!(
162                                "../guest-bootstrap/user-data"
163                            ))
164                        } else {
165                            ImageEntryContent::Binary(include_bytes!(
166                                "../guest-bootstrap/user-data-no-agent"
167                            ))
168                        },
169                    ),
170                    // Specify a non-present NIC to work around https://github.com/canonical/cloud-init/issues/5511
171                    // TODO: support dynamically configuring the network based on vm configuration
172                    (
173                        "network-config",
174                        ImageEntryContent::Binary(include_bytes!(
175                            "../guest-bootstrap/network-config"
176                        )),
177                    ),
178                ]);
179                b"cidata     " // cloud-init looks for a volume label of "cidata",
180            }
181            // Nothing OS-specific yet for other flavors
182            _ => b"cidata     ",
183        };
184
185        let total_data_size: u64 = files
186            .iter()
187            .map(|(_, f)| {
188                Ok(match f {
189                    ImageEntryContent::Dir => 0,
190                    ImageEntryContent::Path(path) => fs_err::metadata(path)?.len(),
191                    ImageEntryContent::Binary(data) => data.len() as u64,
192                })
193            })
194            .collect::<Result<Vec<_>, std::io::Error>>()?
195            .into_iter()
196            .sum();
197
198        let total_size = total_data_size + self.free_space;
199
200        if files.is_empty() {
201            Ok(None)
202        } else {
203            let mut image_file = match image_type {
204                ImageType::Raw => tempfile::NamedTempFile::new()?,
205                ImageType::Vhd => tempfile::Builder::new().suffix(".vhd").tempfile()?,
206            };
207
208            image_file
209                .as_file()
210                .set_len(total_size.next_multiple_of(64 * 1024 * 1024))
211                .context("failed to set file size")?;
212
213            build_fat32_disk_image(&mut image_file, "CIDATA", volume_label, &files)?;
214
215            if matches!(image_type, ImageType::Vhd) {
216                disk_vhd1::Vhd1Disk::make_fixed(image_file.as_file())
217                    .context("failed to make vhd for agent image")?;
218            }
219
220            Ok(Some(image_file))
221        }
222    }
223}
224
225pub(crate) const SECTOR_SIZE: u64 = 512;
226
227pub(crate) enum ImageEntryContent<'a> {
228    Dir,
229    Path(&'a Path),
230    Binary(&'a [u8]),
231}
232
233pub(crate) fn build_fat32_disk_image(
234    file: &mut (impl Read + Write + Seek),
235    gpt_name: &str,
236    volume_label: &[u8; 11],
237    files: &[(&str, ImageEntryContent<'_>)],
238) -> anyhow::Result<()> {
239    let partition_range =
240        build_gpt(file, gpt_name).context("failed to construct partition table")?;
241    build_fat32(
242        &mut fscommon::StreamSlice::new(file, partition_range.start, partition_range.end)?,
243        volume_label,
244        files,
245    )
246    .context("failed to format volume")?;
247    Ok(())
248}
249
250fn build_gpt(file: &mut (impl Read + Write + Seek), name: &str) -> anyhow::Result<Range<u64>> {
251    let mut gpt = gptman::GPT::new_from(file, SECTOR_SIZE, Guid::new_random().into())?;
252
253    // Set up the "Protective" Master Boot Record
254    gptman::GPT::write_protective_mbr_into(file, SECTOR_SIZE)?;
255
256    // Set up the GPT Partition Table Header
257    gpt[1] = gptman::GPTPartitionEntry {
258        // Basic data partition guid
259        partition_type_guid: guid::guid!("EBD0A0A2-B9E5-4433-87C0-68B6B72699C7").into(),
260        unique_partition_guid: Guid::new_random().into(),
261        starting_lba: gpt.header.first_usable_lba,
262        ending_lba: gpt.header.last_usable_lba,
263        attribute_bits: 0,
264        partition_name: name.into(),
265    };
266    gpt.write_into(file)?;
267
268    // calculate the EFI partition's usable range
269    let partition_start_byte = gpt[1].starting_lba * SECTOR_SIZE;
270    let partition_num_bytes = (gpt[1].ending_lba - gpt[1].starting_lba) * SECTOR_SIZE;
271    Ok(partition_start_byte..partition_start_byte + partition_num_bytes)
272}
273
274fn build_fat32(
275    file: &mut (impl Read + Write + Seek),
276    volume_label: &[u8; 11],
277    files: &[(&str, ImageEntryContent<'_>)],
278) -> anyhow::Result<()> {
279    fatfs::format_volume(
280        &mut *file,
281        FormatVolumeOptions::new()
282            .volume_label(*volume_label)
283            .fat_type(fatfs::FatType::Fat32),
284    )
285    .context("failed to format volume")?;
286    let fs = fatfs::FileSystem::new(file, FsOptions::new()).context("failed to open fs")?;
287    for (path, src) in files {
288        if matches!(src, ImageEntryContent::Dir) {
289            fs.root_dir()
290                .create_dir(path)
291                .context("failed to create dir")?;
292            continue;
293        }
294
295        let mut dest = fs
296            .root_dir()
297            .create_file(path)
298            .context("failed to create file")?;
299        match *src {
300            ImageEntryContent::Path(src_path) => {
301                let mut src = fs_err::File::open(src_path)?;
302                std::io::copy(&mut src, &mut dest).context("failed to copy file")?;
303            }
304            ImageEntryContent::Binary(src_data) => {
305                dest.write_all(src_data).context("failed to write file")?;
306            }
307            ImageEntryContent::Dir => unreachable!(),
308        }
309        dest.flush().context("failed to flush file")?;
310    }
311    fs.unmount().context("failed to unmount fs")?;
312    Ok(())
313}
petri/disk_image.rs

petri/
disk_image.rs
