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

petri/
disk_image.rs
