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

petri/
disk_image.rs
