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

petri/
disk_image.rs
