Skip to main content

start_command/
docker_snapshot_safety.rs

1//! Serialize snapshots separately from the execution store and bound disk usage.
2use crate::docker_cleanup::docker_command;
3use crate::execution_control::CommandRunner;
4use crate::isolation::isolation_log::shell_quote;
5use crate::store_lock::LockManager;
6use serde_json::Value;
7use std::path::PathBuf;
8
9const GIB: u64 = 1024 * 1024 * 1024;
10pub const SNAPSHOT_IMAGE_LABEL: &str = "start-command.snapshot-image";
11
12pub fn acquire_snapshot_lock() -> Result<LockManager, String> {
13    let path = std::env::var_os("START_SNAPSHOT_LOCK")
14        .map(PathBuf::from)
15        .unwrap_or_else(|| {
16            if cfg!(unix) {
17                PathBuf::from("/tmp/start-command-docker-snapshot.lock")
18            } else {
19                std::env::temp_dir().join("start-command-docker-snapshot.lock")
20            }
21        });
22    let mut lock = LockManager::new(path);
23    // Same-host LockManager ownership uses PID liveness, not the 60s stale
24    // age used for other hosts. Long commits keep their serialization lock.
25    if !lock.acquire(1000) {
26        return Err(
27            "A Docker snapshot resume is in progress on this host. Retry after it finishes.".into(),
28        );
29    }
30    Ok(lock)
31}
32
33pub fn containerd_root(config: &str) -> String {
34    for line in config.lines().map(str::trim) {
35        if line.starts_with('[') {
36            break;
37        }
38        if let Some((key, value)) = line.split_once('=') {
39            let value = value.trim();
40            if key.trim() == "root" && (value.starts_with('"') || value.starts_with('\'')) {
41                let quote = value.chars().next().unwrap();
42                if let Some(end) = value[1..].find(quote) {
43                    return value[1..end + 1].into();
44                }
45            }
46        }
47    }
48    "/var/lib/containerd".into()
49}
50
51fn default_containerd_root() -> Result<String, String> {
52    if let Ok(root) = std::env::var("START_CONTAINERD_ROOT") {
53        return Ok(root);
54    }
55    match std::fs::read_to_string("/etc/containerd/config.toml") {
56        Ok(config) => Ok(containerd_root(&config)),
57        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(containerd_root("")),
58        Err(e) => Err(format!(
59            "Cannot read containerd storage configuration: {}",
60            e
61        )),
62    }
63}
64
65pub fn preflight_snapshot<R: CommandRunner + ?Sized>(
66    name: &str,
67    runner: &R,
68) -> Result<String, String> {
69    preflight_snapshot_with(name, runner, None, &|root| {
70        let result = runner.run("df", &["-Pk".into(), root.into()]);
71        if !result.success {
72            return Err(format!(
73                "Cannot check free disk at {}: {}",
74                root,
75                result.stderr.trim()
76            ));
77        }
78        result
79            .stdout
80            .lines()
81            .last()
82            .and_then(|line| line.split_whitespace().nth(3))
83            .and_then(|s| s.parse::<u64>().ok())
84            .and_then(|n| n.checked_mul(1024))
85            .ok_or_else(|| format!("Cannot determine free disk at {}", root))
86    })
87}
88
89pub fn preflight_snapshot_with<R: CommandRunner + ?Sized>(
90    name: &str,
91    runner: &R,
92    containerd_override: Option<&str>,
93    free: &dyn Fn(&str) -> Result<u64, String>,
94) -> Result<String, String> {
95    let docker = docker_command().to_string_lossy().to_string();
96    let result = runner.run(
97        &docker,
98        &[
99            "inspect".into(),
100            "--size".into(),
101            "--format".into(),
102            "{{.SizeRw}}".into(),
103            name.into(),
104        ],
105    );
106    let size = result
107        .stdout
108        .trim()
109        .parse::<u64>()
110        .ok()
111        .filter(|_| result.success)
112        .ok_or("Cannot determine the writable layer size; refusing Docker snapshot resume.")?;
113    let result = runner.run(
114        &docker,
115        &["info".into(), "--format".into(), "{{json .}}".into()],
116    );
117    let info: Value = serde_json::from_str(&result.stdout).unwrap_or(Value::Null);
118    let root = info
119        .get("DockerRootDir")
120        .and_then(Value::as_str)
121        .filter(|root| !root.is_empty() && result.success)
122        .ok_or("Cannot determine Docker data root; refusing Docker snapshot resume.")?;
123    let mut roots = vec![root.to_string()];
124    if info
125        .get("DriverStatus")
126        .is_some_and(|d| d.to_string().contains("io.containerd.snapshotter"))
127    {
128        let root = match containerd_override {
129            Some(root) => root.to_string(),
130            None => default_containerd_root()?,
131        };
132        if !roots.contains(&root) {
133            roots.push(root);
134        }
135    }
136    let required = size
137        .checked_mul(2)
138        .and_then(|n| n.checked_add(10 * GIB))
139        .ok_or("Snapshot size exceeds the supported range.")?;
140    for root in roots {
141        let available = free(&root)?;
142        if available < required {
143            return Err(format!("Insufficient disk at {}: need {:.2} GiB (2 × writable layer plus reserve), available {:.2} GiB. Docker snapshot resume refused.", root, required as f64 / GIB as f64, available as f64 / GIB as f64));
144        }
145    }
146    Ok(format!(
147        "[Resume] snapshotting {:.2} GiB; disk preflight passed ({:.2} GiB required).",
148        size as f64 / GIB as f64,
149        required as f64 / GIB as f64
150    ))
151}
152
153pub fn snapshot_image_cleanup_snippets(name: &str, redirection: &str) -> (String, String) {
154    let format = shell_quote(&format!(
155        "{{{{index .Config.Labels \"{}\"}}}}",
156        SNAPSHOT_IMAGE_LABEL
157    ));
158    (format!("__start_snapshot_image=$(docker inspect -f {} {} 2>/dev/null || true)", format, shell_quote(name)),
159     format!("case \"$__start_snapshot_image\" in start-command-resume/*) docker rmi \"$__start_snapshot_image\" {} || true;; esac", redirection))
160}
161
162#[cfg(test)]
163mod tests {
164    use super::*;
165    use crate::execution_control::CommandRunOutput;
166    use std::cell::RefCell;
167
168    struct Runner;
169    impl CommandRunner for Runner {
170        fn run(&self, _bin: &str, args: &[String]) -> CommandRunOutput {
171            CommandRunOutput {
172                success: true,
173                stdout: if args[0] == "info" {
174                    "{\"DockerRootDir\":\"/docker\",\"DriverStatus\":[[\"driver-type\",\"io.containerd.snapshotter.v1\"]]}".into()
175                } else {
176                    (8 * GIB).to_string()
177                },
178                ..Default::default()
179            }
180        }
181    }
182
183    #[test]
184    fn checks_both_roots_and_refuses_insufficient_containerd_disk() {
185        let checked = RefCell::new(vec![]);
186        let error = preflight_snapshot_with("box", &Runner, Some("/containerd"), &|root| {
187            checked.borrow_mut().push(root.to_string());
188            Ok(if root == "/docker" {
189                100 * GIB
190            } else {
191                25 * GIB
192            })
193        })
194        .unwrap_err();
195        assert!(error.contains("/containerd"));
196        assert_eq!(*checked.borrow(), vec!["/docker", "/containerd"]);
197    }
198
199    #[test]
200    fn reads_only_top_level_containerd_root() {
201        assert_eq!(
202            containerd_root("root = '/custom'\n[plugins]\nroot = '/wrong'"),
203            "/custom"
204        );
205    }
206}