Skip to main content

microsandbox_image/checkpoint/
compact.rs

1//! Materialize a pinned immutable disk prefix without following ambient backing paths.
2
3use std::io;
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6
7use imago::file::File as ImagoFile;
8use imago::qcow2::Qcow2;
9use imago::raw::Raw;
10use imago::{
11    DenyImplicitOpenGate, DynStorage, FormatAccess, FormatCreateBuilder, FormatDriverBuilder,
12    Mapping,
13};
14use tokio::io::{AsyncSeekExt, AsyncWriteExt};
15
16//--------------------------------------------------------------------------------------------------
17// Types
18//--------------------------------------------------------------------------------------------------
19
20type SharedImage = Arc<FormatAccess<Box<dyn DynStorage>>>;
21
22/// One caller-resolved immutable member of a compaction prefix.
23#[derive(Clone, Debug)]
24pub struct CompactLayer {
25    /// File pinned by the caller's disk mutation lease.
26    pub path: PathBuf,
27    /// Whether this file is qcow2; otherwise it is raw.
28    pub qcow2: bool,
29}
30
31/// Work performed while materializing a consolidated base.
32#[derive(Clone, Debug)]
33pub struct CompactMaterialization {
34    /// Guest-visible capacity, not the qcow2 container length.
35    pub virtual_size: u64,
36    /// Guest bytes written into materialized runs, including zeros within those runs.
37    pub materialized_bytes: u64,
38}
39
40//--------------------------------------------------------------------------------------------------
41// Functions
42//--------------------------------------------------------------------------------------------------
43
44/// Open a complete, explicitly supplied immutable chain read-only.
45async fn open_chain(layers: &[CompactLayer]) -> io::Result<SharedImage> {
46    open_chain_access(layers, false).await
47}
48
49/// Resolve an owned chain with writes confined to its caller-private staging head.
50pub(crate) async fn open_writable_chain(layers: &[CompactLayer]) -> io::Result<SharedImage> {
51    open_chain_access(layers, true).await
52}
53
54async fn open_chain_access(
55    layers: &[CompactLayer],
56    writable_head: bool,
57) -> io::Result<SharedImage> {
58    if layers.is_empty() || layers.iter().skip(1).any(|layer| !layer.qcow2) {
59        return Err(io::Error::new(
60            io::ErrorKind::InvalidInput,
61            "invalid raw/qcow2 compaction prefix",
62        ));
63    }
64    let mut backing: Option<SharedImage> = None;
65    for (index, layer) in layers.iter().enumerate() {
66        let writable = writable_head && index + 1 == layers.len();
67        let storage: Box<dyn DynStorage> = Box::new(ImagoFile::try_from(
68            std::fs::OpenOptions::new()
69                .read(true)
70                .write(writable)
71                .open(&layer.path)?,
72        )?);
73        backing = Some(if layer.qcow2 {
74            let image = Qcow2::<Box<dyn DynStorage>, SharedImage>::builder(storage)
75                .write(writable)
76                .backing(backing)
77                .data_file(None)
78                .open(DenyImplicitOpenGate::default())
79                .await?;
80            if image.requires_external_data_file() {
81                return Err(io::Error::new(
82                    io::ErrorKind::Unsupported,
83                    "external qcow2 data files are not supported",
84                ));
85            }
86            Arc::new(FormatAccess::new(image))
87        } else {
88            Arc::new(FormatAccess::new(
89                Raw::<Box<dyn DynStorage>>::builder(storage)
90                    .write(writable)
91                    .open(DenyImplicitOpenGate::default())
92                    .await?,
93            ))
94        });
95    }
96    Ok(backing.expect("nonempty chain checked above"))
97}
98
99/// Copy the resolved guest bytes into a new standalone sparse qcow2 image.
100///
101/// The caller owns staging and removes it on cancellation/error. The destination must not exist.
102/// All source layers must remain immutable for the duration. Neither header paths nor guest paths
103/// can open additional files; only the supplied chain is used. Run outside the VM pause window.
104pub async fn materialize_compact_prefix(
105    layers: &[CompactLayer],
106    destination: &Path,
107) -> io::Result<CompactMaterialization> {
108    let source = open_chain(layers).await?;
109    let virtual_size = source.size();
110    if virtual_size == 0 || !virtual_size.is_multiple_of(512) {
111        return Err(io::Error::new(
112            io::ErrorKind::InvalidData,
113            "invalid compaction disk capacity",
114        ));
115    }
116    let file = std::fs::OpenOptions::new()
117        .read(true)
118        .write(true)
119        .create_new(true)
120        .open(destination)?;
121    Qcow2::<ImagoFile>::create_builder(ImagoFile::try_from(file)?)
122        .size(virtual_size)
123        .cluster_size(65536)
124        .create()
125        .await?;
126    let file = std::fs::OpenOptions::new()
127        .read(true)
128        .write(true)
129        .open(destination)?;
130    let target = FormatAccess::new(
131        Qcow2::<ImagoFile>::builder(ImagoFile::try_from(file)?)
132            .write(true)
133            .backing(None)
134            .data_file(None)
135            .open(DenyImplicitOpenGate::default())
136            .await?,
137    );
138    let mut buffer = vec![0u8; 1024 * 1024];
139    let mut offset = 0;
140    let mut materialized_bytes = 0;
141    while offset < virtual_size {
142        let (mapping, length) = source.get_mapping(offset, virtual_size - offset).await?;
143        if length == 0 {
144            return Err(io::Error::new(
145                io::ErrorKind::InvalidData,
146                "compaction mapping made no progress",
147            ));
148        }
149        if matches!(mapping, Mapping::Zero { .. }) {
150            offset += length;
151            continue;
152        }
153        let count = length.min(buffer.len() as u64) as usize;
154        source.read(&mut buffer[..count], offset).await?;
155        // Raw sources may represent sparse holes as data mappings. Preserve sparseness even
156        // there, without depending on platform-specific host extent reporting.
157        if buffer[..count].iter().any(|byte| *byte != 0) {
158            target.write(&buffer[..count], offset).await?;
159            materialized_bytes += count as u64;
160        }
161        offset += count as u64;
162    }
163    target.flush().await?;
164    target.sync().await?;
165    Ok(CompactMaterialization {
166        virtual_size,
167        materialized_bytes,
168    })
169}
170
171/// Materialize a complete explicit chain as a standalone sparse raw disk.
172///
173/// Used when a consumer needs the resolved disk rather than a layered artifact.
174/// The destination must not exist; callers own staging and cleanup on failure.
175/// Source layers must remain immutable, and embedded backing paths are never opened.
176pub async fn materialize_raw_prefix(
177    layers: &[CompactLayer],
178    destination: &Path,
179) -> io::Result<CompactMaterialization> {
180    let source = open_chain(layers).await?;
181    let virtual_size = source.size();
182    if virtual_size == 0 || !virtual_size.is_multiple_of(512) {
183        return Err(io::Error::new(
184            io::ErrorKind::InvalidData,
185            "invalid raw disk capacity",
186        ));
187    }
188    let mut target = tokio::fs::OpenOptions::new()
189        .write(true)
190        .create_new(true)
191        .open(destination)
192        .await?;
193    target.set_len(virtual_size).await?;
194    let mut buffer = vec![0u8; 1024 * 1024];
195    let mut offset = 0;
196    let mut materialized_bytes = 0;
197    while offset < virtual_size {
198        let (mapping, length) = source.get_mapping(offset, virtual_size - offset).await?;
199        if length == 0 {
200            return Err(io::Error::new(
201                io::ErrorKind::InvalidData,
202                "raw disk mapping made no progress",
203            ));
204        }
205        if matches!(mapping, Mapping::Zero { .. }) {
206            offset += length;
207            continue;
208        }
209        let count = length.min(buffer.len() as u64) as usize;
210        source.read(&mut buffer[..count], offset).await?;
211        if buffer[..count].iter().any(|byte| *byte != 0) {
212            target.seek(io::SeekFrom::Start(offset)).await?;
213            target.write_all(&buffer[..count]).await?;
214            materialized_bytes += count as u64;
215        }
216        offset += count as u64;
217    }
218    target.sync_all().await?;
219    Ok(CompactMaterialization {
220        virtual_size,
221        materialized_bytes,
222    })
223}
224
225/// Read a raw or qcow2 file's declared capacity without opening its backing filename.
226pub async fn compact_layer_capacity(layer: CompactLayer) -> io::Result<u64> {
227    Ok(open_chain(&[layer]).await?.size())
228}
229
230/// Read the capacities of a pinned closure from synchronous descriptor-building code.
231/// A separate current-thread executor also permits use by callers already inside Tokio.
232pub fn layer_capacities(layers: Vec<CompactLayer>) -> io::Result<Vec<u64>> {
233    std::thread::spawn(move || {
234        tokio::runtime::Builder::new_current_thread()
235            .enable_all()
236            .build()?
237            .block_on(async {
238                let mut capacities = Vec::with_capacity(layers.len());
239                for layer in layers {
240                    capacities.push(compact_layer_capacity(layer).await?);
241                }
242                Ok(capacities)
243            })
244    })
245    .join()
246    .map_err(|_| io::Error::other("layer-capacity worker panicked"))?
247}
248
249/// Validate a complete explicit chain without linking a VM runner or following header paths.
250///
251/// The caller must prevent concurrent mutation until publication completes.
252pub async fn validate_compact_chain(layers: &[CompactLayer]) -> io::Result<()> {
253    open_chain(layers).await?;
254    Ok(())
255}
256
257//--------------------------------------------------------------------------------------------------
258// Tests
259//--------------------------------------------------------------------------------------------------
260
261#[cfg(test)]
262mod tests {
263    use super::*;
264    use crate::checkpoint::create_qcow2_overlay;
265
266    #[tokio::test]
267    async fn compacted_prefix_preserves_overwrites_zeroes_and_grown_tail() {
268        let dir = tempfile::tempdir().unwrap();
269        let raw = dir.path().join("base.raw");
270        std::fs::write(&raw, vec![17u8; 131072]).unwrap();
271        let overlay = dir.path().join("layer.qcow2");
272        create_qcow2_overlay(&overlay, 262144, &raw, "raw")
273            .await
274            .unwrap();
275        let image = FormatAccess::new(
276            Qcow2::<ImagoFile>::builder(
277                ImagoFile::try_from(
278                    std::fs::OpenOptions::new()
279                        .read(true)
280                        .write(true)
281                        .open(&overlay)
282                        .unwrap(),
283                )
284                .unwrap(),
285            )
286            .write(true)
287            .backing(None)
288            .data_file(None)
289            .open(DenyImplicitOpenGate::default())
290            .await
291            .unwrap(),
292        );
293        image.write(&vec![29u8; 65536][..], 0).await.unwrap();
294        image.write_zeroes(65536, 65536).await.unwrap();
295        image.flush().await.unwrap();
296        image.sync().await.unwrap();
297        drop(image);
298        let layers = vec![
299            CompactLayer {
300                path: raw.clone(),
301                qcow2: false,
302            },
303            CompactLayer {
304                path: overlay.clone(),
305                qcow2: true,
306            },
307        ];
308        let original = std::fs::read(&overlay).unwrap();
309        let destination = dir.path().join("compact.qcow2");
310        let result = materialize_compact_prefix(&layers, &destination)
311            .await
312            .unwrap();
313        assert_eq!(result.virtual_size, 262144);
314        let input = open_chain(&layers).await.unwrap();
315        let output = open_chain(&[CompactLayer {
316            path: destination.clone(),
317            qcow2: true,
318        }])
319        .await
320        .unwrap();
321        let mut before = vec![0; 262144];
322        let mut after = before.clone();
323        input.read(&mut before[..], 0).await.unwrap();
324        output.read(&mut after[..], 0).await.unwrap();
325        assert!(before == after, "compaction changed guest bytes");
326        assert!(after[..65536].iter().all(|byte| *byte == 29));
327        assert!(after[65536..].iter().all(|byte| *byte == 0));
328        let raw_destination = dir.path().join("flattened.raw");
329        let raw_result = materialize_raw_prefix(&layers, &raw_destination)
330            .await
331            .unwrap();
332        assert_eq!(raw_result.virtual_size, 262144);
333        assert_eq!(std::fs::read(&raw_destination).unwrap(), before);
334        assert!(
335            materialize_raw_prefix(&layers, &raw_destination)
336                .await
337                .is_err()
338        );
339        assert_eq!(std::fs::read(overlay).unwrap(), original);
340        assert!(
341            materialize_compact_prefix(&layers, &destination)
342                .await
343                .is_err()
344        );
345    }
346}