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

petri/
disk_image.rs
