xref: /aosp_15_r20/external/crosvm/vm_control/src/snapshot_format.rs (revision bb4ee6a4ae7042d18b07a98463b9c8b875e44b39)
1 // Copyright 2024 The ChromiumOS Authors
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4 
5 use std::fmt::Debug;
6 use std::fmt::Formatter;
7 use std::fs::File;
8 use std::io::Read;
9 use std::io::Write;
10 use std::path::Path;
11 use std::path::PathBuf;
12 
13 use anyhow::Context;
14 use anyhow::Result;
15 use crypto::CryptKey;
16 
17 // Use 4kB encrypted chunks by default (if encryption is used).
18 const DEFAULT_ENCRYPTED_CHUNK_SIZE_BYTES: usize = 1024 * 4;
19 
20 /// Writer of serialized VM snapshots.
21 ///
22 /// Each fragment is an opaque byte blob. Namespaces can be used to avoid fragment naming
23 /// collisions between devices.
24 ///
25 /// In the current implementation, fragments are files and namespaces are directories, but the API
26 /// is kept abstract so that we can potentially support something like a single file archive
27 /// output.
28 #[derive(Clone, serde::Serialize, serde::Deserialize)]
29 pub struct SnapshotWriter {
30     dir: PathBuf,
31     /// If encryption is used, the plaintext key will be stored here.
32     key: Option<CryptKey>,
33 }
34 
35 impl Debug for SnapshotWriter {
fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result36     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
37         f.debug_struct("SnapshotWriter")
38             .field("dir", &format!("{:?}", self.dir))
39             .field("key", if self.key.is_some() { &"Some" } else { &"None" })
40             .finish()
41     }
42 }
43 
44 impl SnapshotWriter {
45     /// Creates a new `SnapshotWriter` that will writes its data to a dir at `root`. The path must
46     /// not exist yet. If encryption is desired, set encrypt (Note: only supported downstream on
47     /// Windows).
48     // TODO(b/268094487): If the snapshot fails, we leave incomplete snapshot files at the
49     // requested path. Consider building up the snapshot dir somewhere else and moving it into
50     // place at the end.
new(root: PathBuf, encrypt: bool) -> Result<Self>51     pub fn new(root: PathBuf, encrypt: bool) -> Result<Self> {
52         std::fs::create_dir(&root)
53             .with_context(|| format!("failed to create snapshot root dir: {}", root.display()))?;
54 
55         if encrypt {
56             let key = crypto::generate_random_key();
57             // Creating an empty CryptWriter will still write header information
58             // to the file, and that header information is what we need. This
59             // ensures we use a single key for *all* snapshot files.
60             let mut writer = crypto::CryptWriter::new_from_key(
61                 File::create(root.join("enc_metadata")).context("failed to create enc_metadata")?,
62                 1024,
63                 &key,
64             )
65             .context("failed to create enc_metadata writer")?;
66             writer.flush().context("flush of enc_metadata failed")?;
67             return Ok(Self {
68                 dir: root,
69                 key: Some(key),
70             });
71         }
72 
73         Ok(Self {
74             dir: root,
75             key: None,
76         })
77     }
78 
79     /// Creates a snapshot fragment and get access to the `Write` impl representing it.
raw_fragment(&self, name: &str) -> Result<Box<dyn Write>>80     pub fn raw_fragment(&self, name: &str) -> Result<Box<dyn Write>> {
81         self.raw_fragment_with_chunk_size(name, DEFAULT_ENCRYPTED_CHUNK_SIZE_BYTES)
82     }
83 
84     /// When encryption is used, allows direct control of the encrypted chunk size.
raw_fragment_with_chunk_size( &self, name: &str, chunk_size_bytes: usize, ) -> Result<Box<dyn Write>>85     pub fn raw_fragment_with_chunk_size(
86         &self,
87         name: &str,
88         chunk_size_bytes: usize,
89     ) -> Result<Box<dyn Write>> {
90         let path = self.dir.join(name);
91         let file = File::options()
92             .write(true)
93             .create_new(true)
94             .open(&path)
95             .with_context(|| {
96                 format!(
97                     "failed to create snapshot fragment {name:?} at {}",
98                     path.display()
99                 )
100             })?;
101 
102         if let Some(key) = self.key.as_ref() {
103             return Ok(Box::new(crypto::CryptWriter::new_from_key(
104                 file,
105                 chunk_size_bytes,
106                 key,
107             )?));
108         }
109 
110         Ok(Box::new(file))
111     }
112 
113     /// Creates a snapshot fragment from a serialized representation of `v`.
write_fragment<T: serde::Serialize>(&self, name: &str, v: &T) -> Result<()>114     pub fn write_fragment<T: serde::Serialize>(&self, name: &str, v: &T) -> Result<()> {
115         let mut w = std::io::BufWriter::new(self.raw_fragment(name)?);
116         serde_json::to_writer(&mut w, v)?;
117         w.flush()?;
118         Ok(())
119     }
120 
121     /// Creates new namespace and returns a `SnapshotWriter` that writes to it. Namespaces can be
122     /// nested.
add_namespace(&self, name: &str) -> Result<Self>123     pub fn add_namespace(&self, name: &str) -> Result<Self> {
124         let dir = self.dir.join(name);
125         std::fs::create_dir(&dir).with_context(|| {
126             format!(
127                 "failed to create nested snapshot writer {name:?} at {}",
128                 dir.display()
129             )
130         })?;
131         Ok(Self {
132             dir,
133             key: self.key.clone(),
134         })
135     }
136 }
137 
138 /// Reads snapshots created by `SnapshotWriter`.
139 #[derive(Clone, serde::Serialize, serde::Deserialize)]
140 pub struct SnapshotReader {
141     dir: PathBuf,
142     /// If encryption is used, the plaintext key will be stored here.
143     key: Option<CryptKey>,
144 }
145 
146 impl Debug for SnapshotReader {
fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result147     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
148         f.debug_struct("SnapshotReader")
149             .field("dir", &format!("{:?}", self.dir))
150             .field("key", if self.key.is_some() { &"Some" } else { &"None" })
151             .finish()
152     }
153 }
154 
155 impl SnapshotReader {
156     /// Reads a snapshot at `root`. Set require_encrypted to require an encrypted snapshot.
new(root: &Path, require_encrypted: bool) -> Result<Self>157     pub fn new(root: &Path, require_encrypted: bool) -> Result<Self> {
158         let enc_metadata_path = root.join("enc_metadata");
159         if Path::exists(&enc_metadata_path) {
160             let key = Some(
161                 crypto::CryptReader::extract_key(
162                     File::open(&enc_metadata_path).context("failed to open encryption metadata")?,
163                 )
164                 .context("failed to load snapshot key")?,
165             );
166             return Ok(Self {
167                 dir: root.to_path_buf(),
168                 key,
169             });
170         } else if require_encrypted {
171             return Err(anyhow::anyhow!("snapshot was not encrypted"));
172         }
173 
174         Ok(Self {
175             dir: root.to_path_buf(),
176             key: None,
177         })
178     }
179 
180     /// Gets access to a `Read` impl that represents a fragment.
raw_fragment(&self, name: &str) -> Result<Box<dyn Read>>181     pub fn raw_fragment(&self, name: &str) -> Result<Box<dyn Read>> {
182         let path = self.dir.join(name);
183         let file = File::open(&path).with_context(|| {
184             format!(
185                 "failed to open snapshot fragment {name:?} at {}",
186                 path.display()
187             )
188         })?;
189         if let Some(key) = self.key.as_ref() {
190             return Ok(Box::new(crypto::CryptReader::from_file_and_key(file, key)?));
191         }
192 
193         Ok(Box::new(file))
194     }
195 
196     /// Reads a fragment.
read_fragment<T: serde::de::DeserializeOwned>(&self, name: &str) -> Result<T>197     pub fn read_fragment<T: serde::de::DeserializeOwned>(&self, name: &str) -> Result<T> {
198         Ok(serde_json::from_reader(std::io::BufReader::new(
199             self.raw_fragment(name)?,
200         ))?)
201     }
202 
203     /// Reads the names of all fragments in this namespace.
list_fragments(&self) -> Result<Vec<String>>204     pub fn list_fragments(&self) -> Result<Vec<String>> {
205         let mut result = Vec::new();
206         for entry in std::fs::read_dir(&self.dir)? {
207             let entry = entry?;
208             if entry.path().is_file() {
209                 if let Some(file_name) = entry.path().file_name() {
210                     result.push(file_name.to_string_lossy().into_owned());
211                 }
212             }
213         }
214         Ok(result)
215     }
216 
217     /// Open a namespace.
namespace(&self, name: &str) -> Result<Self>218     pub fn namespace(&self, name: &str) -> Result<Self> {
219         let dir = self.dir.join(name);
220         Ok(Self {
221             dir,
222             key: self.key.clone(),
223         })
224     }
225 
226     /// Reads the names of all child namespaces
list_namespaces(&self) -> Result<Vec<String>>227     pub fn list_namespaces(&self) -> Result<Vec<String>> {
228         let mut result = Vec::new();
229         for entry in std::fs::read_dir(&self.dir)? {
230             let entry = entry?;
231             if entry.path().is_dir() {
232                 if let Some(file_name) = entry.path().file_name() {
233                     result.push(file_name.to_string_lossy().into_owned());
234                 }
235             }
236         }
237         Ok(result)
238     }
239 }
240