1mod rust_criu_protobuf;
2
3use anyhow::{Context, Result};
4use protobuf::Message;
5use rust_criu_protobuf::rpc;
6use std::error::Error;
7use std::fs::File;
8use std::io::{Read, Write};
9use std::os::unix::io::FromRawFd;
10use std::process::Command;
11
12#[derive(Clone)]
13pub enum CgMode {
14 IGNORE = 0,
15 NONE = 1,
16 PROPS = 2,
17 SOFT = 3,
18 FULL = 4,
19 STRICT = 5,
20 DEFAULT = 6,
21}
22
23impl CgMode {
24 pub fn from(value: i32) -> CgMode {
25 match value {
26 0 => Self::IGNORE,
27 1 => Self::NONE,
28 2 => Self::PROPS,
29 3 => Self::SOFT,
30 4 => Self::FULL,
31 5 => Self::STRICT,
32 6 => Self::DEFAULT,
33 _ => Self::DEFAULT,
34 }
35 }
36}
37
38#[derive(Clone)]
39pub struct Criu {
40 criu_path: String,
41 sv: [i32; 2],
42 pid: i32,
43 images_dir_fd: i32,
44 log_level: i32,
45 log_file: Option<String>,
46 external_mounts: Vec<(String, String)>,
47 orphan_pts_master: Option<bool>,
48 root: Option<String>,
49 leave_running: Option<bool>,
50 ext_unix_sk: Option<bool>,
51 shell_job: Option<bool>,
52 tcp_established: Option<bool>,
53 file_locks: Option<bool>,
54 manage_cgroups: Option<bool>,
55 work_dir_fd: i32,
56 freeze_cgroup: Option<String>,
57 cgroups_mode: Option<CgMode>,
58 cgroup_props: Option<String>,
59}
60
61impl Criu {
62 pub fn new() -> Result<Self, Box<dyn Error>> {
63 Criu::new_with_criu_path(String::from("criu"))
64 }
65
66 pub fn new_with_criu_path(path_to_criu: String) -> Result<Self, Box<dyn Error>> {
67 Ok(Self {
68 criu_path: path_to_criu,
69 sv: [-1, -1],
70 pid: -1,
71 images_dir_fd: -1,
72 log_level: -1,
73 log_file: None,
74 external_mounts: Vec::new(),
75 orphan_pts_master: None,
76 root: None,
77 leave_running: None,
78 ext_unix_sk: None,
79 shell_job: None,
80 tcp_established: None,
81 file_locks: None,
82 manage_cgroups: None,
83 work_dir_fd: -1,
84 freeze_cgroup: None,
85 cgroups_mode: None,
86 cgroup_props: None,
87 })
88 }
89
90 pub fn get_criu_version(&mut self) -> Result<u32, Box<dyn Error>> {
91 let response = self.do_swrk_with_response(rpc::Criu_req_type::VERSION, None)?;
92
93 let mut version: u32 = (response.version.major_number() * 10000)
94 .try_into()
95 .context("parsing criu version failed")?;
96 version += (response.version.minor_number() * 100) as u32;
97 version += response.version.sublevel() as u32;
98
99 if response.version.has_gitid() {
100 version -= version % 100;
102 version += 100;
103 }
104
105 Ok(version)
106 }
107
108 fn do_swrk_with_response(
109 &mut self,
110 request_type: rpc::Criu_req_type,
111 criu_opts: Option<rpc::Criu_opts>,
112 ) -> Result<rpc::Criu_resp, Box<dyn Error>> {
113 if unsafe {
114 libc::socketpair(
115 libc::AF_LOCAL,
116 libc::SOCK_SEQPACKET,
117 0,
118 self.sv.as_mut_ptr(),
119 ) != 0
120 } {
121 return Err("libc::socketpair failed".into());
122 }
123
124 let mut criu = Command::new(self.criu_path.clone())
125 .arg("swrk")
126 .arg(format!("{}", self.sv[1]))
127 .spawn()
128 .with_context(|| {
129 format!(
130 "executing criu binary for swrk using path {:?} failed",
131 self.criu_path
132 )
133 })?;
134
135 let mut req = rpc::Criu_req::new();
136 req.set_type(request_type);
137
138 if let Some(co) = criu_opts {
139 req.opts = protobuf::MessageField::some(co);
140 }
141
142 let mut f = unsafe { File::from_raw_fd(self.sv[0]) };
143
144 f.write_all(
145 &req.write_to_bytes()
146 .context("writing protobuf request to byte vec failed")?,
147 )
148 .with_context(|| {
149 format!(
150 "writing protobuf request to file (fd : {}) failed",
151 self.sv[0]
152 )
153 })?;
154
155 let mut buffer = [0; 2 * 4096];
157
158 let read = f.read(&mut buffer[..]).with_context(|| {
159 format!(
160 "reading criu response from file (fd :{}) failed",
161 self.sv[0]
162 )
163 })?;
164
165 let response: rpc::Criu_resp =
166 Message::parse_from_bytes(&buffer[..read]).context("parsing criu response failed")?;
167
168 if !response.success() {
169 criu.kill()
170 .context("killing criu process (due to failed request) failed")?;
171 return Result::Err(
172 format!(
173 "CRIU RPC request failed with message:{} error:{}",
174 response.cr_errmsg(),
175 response.cr_errno()
176 )
177 .into(),
178 );
179 }
180
181 if response.type_() != request_type {
182 criu.kill()
183 .context("killing criu process (due to incorrect response) failed")?;
184 return Result::Err(
185 format!("Unexpected CRIU RPC response ({:?})", response.type_()).into(),
186 );
187 }
188
189 criu.kill().context("killing criu process failed")?;
190 Result::Ok(response)
191 }
192
193 pub fn set_pid(&mut self, pid: i32) {
194 self.pid = pid;
195 }
196
197 pub fn set_images_dir_fd(&mut self, fd: i32) {
198 self.images_dir_fd = fd;
199 }
200
201 pub fn set_log_level(&mut self, log_level: i32) {
202 self.log_level = log_level;
203 }
204
205 pub fn set_log_file(&mut self, log_file: String) {
206 self.log_file = Some(log_file);
207 }
208
209 pub fn set_external_mount(&mut self, key: String, value: String) {
210 self.external_mounts.push((key, value));
211 }
212
213 pub fn set_orphan_pts_master(&mut self, orphan_pts_master: bool) {
214 self.orphan_pts_master = Some(orphan_pts_master);
215 }
216
217 pub fn set_root(&mut self, root: String) {
218 self.root = Some(root);
219 }
220
221 pub fn set_leave_running(&mut self, leave_running: bool) {
222 self.leave_running = Some(leave_running);
223 }
224
225 pub fn set_ext_unix_sk(&mut self, ext_unix_sk: bool) {
226 self.ext_unix_sk = Some(ext_unix_sk);
227 }
228
229 pub fn set_shell_job(&mut self, shell_job: bool) {
230 self.shell_job = Some(shell_job);
231 }
232
233 pub fn set_tcp_established(&mut self, tcp_established: bool) {
234 self.tcp_established = Some(tcp_established);
235 }
236
237 pub fn set_file_locks(&mut self, file_locks: bool) {
238 self.file_locks = Some(file_locks);
239 }
240
241 pub fn set_manage_cgroups(&mut self, manage_cgroups: bool) {
242 self.manage_cgroups = Some(manage_cgroups);
243 }
244
245 pub fn set_work_dir_fd(&mut self, fd: i32) {
246 self.work_dir_fd = fd;
247 }
248
249 pub fn set_freeze_cgroup(&mut self, freeze_cgroup: String) {
250 self.freeze_cgroup = Some(freeze_cgroup);
251 }
252
253 pub fn cgroups_mode(&mut self, mode: CgMode) {
254 self.cgroups_mode = Some(mode);
255 }
256
257 pub fn set_cgroup_props(&mut self, props: String) {
258 self.cgroup_props = Some(props);
259 }
260
261 fn fill_criu_opts(&mut self, criu_opts: &mut rpc::Criu_opts) {
262 if self.pid != -1 {
263 criu_opts.set_pid(self.pid);
264 }
265
266 if self.images_dir_fd != -1 {
267 criu_opts.set_images_dir_fd(self.images_dir_fd);
268 }
269
270 if self.log_level != -1 {
271 criu_opts.set_log_level(self.log_level);
272 }
273
274 if self.log_file.is_some() {
275 criu_opts.set_log_file(self.log_file.clone().unwrap());
276 }
277
278 if !self.external_mounts.is_empty() {
279 let mut external_mounts = Vec::new();
280 for e in &self.external_mounts {
281 let mut external_mount = rpc::Ext_mount_map::new();
282 external_mount.set_key(e.0.clone());
283 external_mount.set_val(e.1.clone());
284 external_mounts.push(external_mount);
285 }
286 self.external_mounts.clear();
287 criu_opts.ext_mnt = external_mounts;
288 }
289
290 if self.orphan_pts_master.is_some() {
291 criu_opts.set_orphan_pts_master(self.orphan_pts_master.unwrap());
292 }
293
294 if self.root.is_some() {
295 criu_opts.set_root(self.root.clone().unwrap());
296 }
297
298 if self.leave_running.is_some() {
299 criu_opts.set_leave_running(self.leave_running.unwrap());
300 }
301
302 if self.ext_unix_sk.is_some() {
303 criu_opts.set_ext_unix_sk(self.ext_unix_sk.unwrap());
304 }
305
306 if self.shell_job.is_some() {
307 criu_opts.set_shell_job(self.shell_job.unwrap());
308 }
309
310 if self.tcp_established.is_some() {
311 criu_opts.set_tcp_established(self.tcp_established.unwrap());
312 }
313
314 if self.file_locks.is_some() {
315 criu_opts.set_file_locks(self.file_locks.unwrap());
316 }
317
318 if self.manage_cgroups.is_some() {
319 criu_opts.set_manage_cgroups(self.manage_cgroups.unwrap());
320 }
321
322 if self.work_dir_fd != -1 {
323 criu_opts.set_work_dir_fd(self.work_dir_fd);
324 }
325
326 if self.freeze_cgroup.is_some() {
327 criu_opts.set_freeze_cgroup(self.freeze_cgroup.clone().unwrap());
328 }
329
330 if self.cgroups_mode.is_some() {
331 let mode = match self.cgroups_mode.as_ref().unwrap() {
332 CgMode::IGNORE => rpc::Criu_cg_mode::IGNORE,
333 CgMode::NONE => rpc::Criu_cg_mode::CG_NONE,
334 CgMode::PROPS => rpc::Criu_cg_mode::PROPS,
335 CgMode::SOFT => rpc::Criu_cg_mode::SOFT,
336 CgMode::FULL => rpc::Criu_cg_mode::FULL,
337 CgMode::STRICT => rpc::Criu_cg_mode::STRICT,
338 CgMode::DEFAULT => rpc::Criu_cg_mode::DEFAULT,
339 };
340 criu_opts.set_manage_cgroups_mode(mode);
341 }
342
343 if self.cgroup_props.is_some() {
344 criu_opts.set_cgroup_props(self.cgroup_props.clone().unwrap());
345 }
346 }
347
348 fn clear(&mut self) {
349 self.pid = -1;
350 self.images_dir_fd = -1;
351 self.log_level = -1;
352 self.log_file = None;
353 self.external_mounts = Vec::new();
354 self.orphan_pts_master = None;
355 self.root = None;
356 self.leave_running = None;
357 self.ext_unix_sk = None;
358 self.shell_job = None;
359 self.tcp_established = None;
360 self.file_locks = None;
361 self.manage_cgroups = None;
362 self.work_dir_fd = -1;
363 self.freeze_cgroup = None;
364 self.cgroups_mode = None;
365 self.cgroup_props = None;
366 }
367
368 pub fn dump(&mut self) -> Result<(), Box<dyn Error>> {
369 let mut criu_opts = rpc::Criu_opts::default();
370 self.fill_criu_opts(&mut criu_opts);
371 self.do_swrk_with_response(rpc::Criu_req_type::DUMP, Some(criu_opts))?;
372 self.clear();
373
374 Ok(())
375 }
376
377 pub fn restore(&mut self) -> Result<(), Box<dyn Error>> {
378 let mut criu_opts = rpc::Criu_opts::default();
379 self.fill_criu_opts(&mut criu_opts);
380 self.do_swrk_with_response(rpc::Criu_req_type::RESTORE, Some(criu_opts))?;
381 self.clear();
382
383 Ok(())
384 }
385}
386
387impl Drop for Criu {
388 fn drop(&mut self) {
389 unsafe { libc::close(self.sv[0]) };
390 unsafe { libc::close(self.sv[1]) };
391 }
392}