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;
20
21/// The description and artifacts needed to build a pipette disk image for a VM.
22#[derive(Debug)]
23pub struct AgentImage {
24    os_flavor: OsFlavor,
25    pipette: Option<ResolvedArtifact>,
26    extras: Vec<(String, ResolvedArtifact)>,
27}
28
29/// Disk image type
30pub enum ImageType {
31    /// Raw image
32    Raw,
33    /// Fixed VHD1
34    Vhd,
35}
36
37impl AgentImage {
38    /// Resolves the artifacts needed to build a disk image for a VM.
39    pub fn new(os_flavor: OsFlavor) -> Self {
40        Self {
41            os_flavor,
42            pipette: None,
43            extras: Vec::new(),
44        }
45    }
46
47    /// Adds the appropriate pipette binary to the image
48    pub fn with_pipette(mut self, resolver: &ArtifactResolver<'_>, arch: MachineArch) -> Self {
49        self.pipette = match (self.os_flavor, arch) {
50            (OsFlavor::Windows, MachineArch::X86_64) => Some(
51                resolver
52                    .require(common_artifacts::PIPETTE_WINDOWS_X64)
53                    .erase(),
54            ),
55            (OsFlavor::Linux, MachineArch::X86_64) => Some(
56                resolver
57                    .require(common_artifacts::PIPETTE_LINUX_X64)
58                    .erase(),
59            ),
60            (OsFlavor::Windows, MachineArch::Aarch64) => Some(
61                resolver
62                    .require(common_artifacts::PIPETTE_WINDOWS_AARCH64)
63                    .erase(),
64            ),
65            (OsFlavor::Linux, MachineArch::Aarch64) => Some(
66                resolver
67                    .require(common_artifacts::PIPETTE_LINUX_AARCH64)
68                    .erase(),
69            ),
70            (OsFlavor::FreeBsd | OsFlavor::Uefi, _) => {
71                todo!("No pipette binary yet for os");
72            }
73        };
74        self
75    }
76
77    /// Check if the image contains pipette
78    pub fn contains_pipette(&self) -> bool {
79        self.pipette.is_some()
80    }
81
82    /// Check if the image has extra files beyond pipette and cloud-init.
83    pub fn has_extras(&self) -> bool {
84        !self.extras.is_empty()
85    }
86
87    /// Adds an extra file to the disk image.
88    pub fn add_file(&mut self, name: &str, artifact: ResolvedArtifact) {
89        self.extras.push((name.to_string(), artifact));
90    }
91
92    /// Builds a disk image containing pipette and any files needed for the guest VM
93    /// to run pipette.
94    pub fn build(&self, image_type: ImageType) -> anyhow::Result<Option<tempfile::NamedTempFile>> {
95        let mut files = self
96            .extras
97            .iter()
98            .map(|(name, artifact)| (name.as_str(), PathOrBinary::Path(artifact.as_ref())))
99            .collect::<Vec<_>>();
100        let volume_label = match self.os_flavor {
101            OsFlavor::Windows => {
102                // Windows doesn't use cloud-init, so we only need pipette
103                // (which is configured via the IMC hive).
104                if let Some(pipette) = self.pipette.as_ref() {
105                    files.push(("pipette.exe", PathOrBinary::Path(pipette.as_ref())));
106                }
107                b"pipette    "
108            }
109            OsFlavor::Linux => {
110                if let Some(pipette) = self.pipette.as_ref() {
111                    files.push(("pipette", PathOrBinary::Path(pipette.as_ref())));
112                }
113                // Linux uses cloud-init, so we need to include the cloud-init
114                // configuration files as well.
115                files.extend([
116                    (
117                        "meta-data",
118                        PathOrBinary::Binary(include_bytes!("../guest-bootstrap/meta-data")),
119                    ),
120                    (
121                        "user-data",
122                        if self.pipette.is_some() {
123                            PathOrBinary::Binary(include_bytes!("../guest-bootstrap/user-data"))
124                        } else {
125                            PathOrBinary::Binary(include_bytes!(
126                                "../guest-bootstrap/user-data-no-agent"
127                            ))
128                        },
129                    ),
130                    // Specify a non-present NIC to work around https://github.com/canonical/cloud-init/issues/5511
131                    // TODO: support dynamically configuring the network based on vm configuration
132                    (
133                        "network-config",
134                        PathOrBinary::Binary(include_bytes!("../guest-bootstrap/network-config")),
135                    ),
136                ]);
137                b"cidata     " // cloud-init looks for a volume label of "cidata",
138            }
139            // Nothing OS-specific yet for other flavors
140            _ => b"cidata     ",
141        };
142
143        if files.is_empty() {
144            Ok(None)
145        } else {
146            let mut image_file = match image_type {
147                ImageType::Raw => tempfile::NamedTempFile::new()?,
148                ImageType::Vhd => tempfile::Builder::new().suffix(".vhd").tempfile()?,
149            };
150
151            image_file
152                .as_file()
153                .set_len(64 * 1024 * 1024)
154                .context("failed to set file size")?;
155
156            build_fat32_disk_image(&mut image_file, "CIDATA", volume_label, &files)?;
157
158            if matches!(image_type, ImageType::Vhd) {
159                disk_vhd1::Vhd1Disk::make_fixed(image_file.as_file())
160                    .context("failed to make vhd for agent image")?;
161            }
162
163            Ok(Some(image_file))
164        }
165    }
166}
167
168pub(crate) const SECTOR_SIZE: u64 = 512;
169
170pub(crate) enum PathOrBinary<'a> {
171    Path(&'a Path),
172    Binary(&'a [u8]),
173}
174
175pub(crate) fn build_fat32_disk_image(
176    file: &mut (impl Read + Write + Seek),
177    gpt_name: &str,
178    volume_label: &[u8; 11],
179    files: &[(&str, PathOrBinary<'_>)],
180) -> anyhow::Result<()> {
181    let partition_range =
182        build_gpt(file, gpt_name).context("failed to construct partition table")?;
183    build_fat32(
184        &mut fscommon::StreamSlice::new(file, partition_range.start, partition_range.end)?,
185        volume_label,
186        files,
187    )
188    .context("failed to format volume")?;
189    Ok(())
190}
191
192fn build_gpt(file: &mut (impl Read + Write + Seek), name: &str) -> anyhow::Result<Range<u64>> {
193    let mut gpt = gptman::GPT::new_from(file, SECTOR_SIZE, Guid::new_random().into())?;
194
195    // Set up the "Protective" Master Boot Record
196    gptman::GPT::write_protective_mbr_into(file, SECTOR_SIZE)?;
197
198    // Set up the GPT Partition Table Header
199    gpt[1] = gptman::GPTPartitionEntry {
200        // Basic data partition guid
201        partition_type_guid: guid::guid!("EBD0A0A2-B9E5-4433-87C0-68B6B72699C7").into(),
202        unique_partition_guid: Guid::new_random().into(),
203        starting_lba: gpt.header.first_usable_lba,
204        ending_lba: gpt.header.last_usable_lba,
205        attribute_bits: 0,
206        partition_name: name.into(),
207    };
208    gpt.write_into(file)?;
209
210    // calculate the EFI partition's usable range
211    let partition_start_byte = gpt[1].starting_lba * SECTOR_SIZE;
212    let partition_num_bytes = (gpt[1].ending_lba - gpt[1].starting_lba) * SECTOR_SIZE;
213    Ok(partition_start_byte..partition_start_byte + partition_num_bytes)
214}
215
216fn build_fat32(
217    file: &mut (impl Read + Write + Seek),
218    volume_label: &[u8; 11],
219    files: &[(&str, PathOrBinary<'_>)],
220) -> anyhow::Result<()> {
221    fatfs::format_volume(
222        &mut *file,
223        FormatVolumeOptions::new()
224            .volume_label(*volume_label)
225            .fat_type(fatfs::FatType::Fat32),
226    )
227    .context("failed to format volume")?;
228    let fs = fatfs::FileSystem::new(file, FsOptions::new()).context("failed to open fs")?;
229    for (path, src) in files {
230        let mut dest = fs
231            .root_dir()
232            .create_file(path)
233            .context("failed to create file")?;
234        match *src {
235            PathOrBinary::Path(src_path) => {
236                let mut src = fs_err::File::open(src_path)?;
237                std::io::copy(&mut src, &mut dest).context("failed to copy file")?;
238            }
239            PathOrBinary::Binary(src_data) => {
240                dest.write_all(src_data).context("failed to write file")?;
241            }
242        }
243        dest.flush().context("failed to flush file")?;
244    }
245    fs.unmount().context("failed to unmount fs")?;
246    Ok(())
247}
petri/disk_image.rs

petri/
disk_image.rs
