ic_testkit/artifacts/
digest.rs

1use sha2::{Digest, Sha256};
2use std::{
3    collections::BTreeSet,
4    ffi::OsStr,
5    fmt::Write as _,
6    fs::{self, File, OpenOptions},
7    io::{self, Read as _, Write as _},
8    path::{Path, PathBuf},
9    sync::atomic::{AtomicU64, Ordering},
10};
11
12static TEMP_FILE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
13
14#[derive(Debug)]
15struct AtomicCopyErrorContext {
16    source_path: PathBuf,
17    destination_path: PathBuf,
18    source: io::Error,
19}
20
21impl std::fmt::Display for AtomicCopyErrorContext {
22    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
23        write!(
24            formatter,
25            "failed to atomically copy {} to {}: {}",
26            self.source_path.display(),
27            self.destination_path.display(),
28            self.source
29        )
30    }
31}
32
33impl std::error::Error for AtomicCopyErrorContext {
34    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
35        Some(&self.source)
36    }
37}
38
39/// SHA-256 digest of one deterministic artifact-input set.
40#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
41pub struct InputDigest([u8; 32]);
42
43impl InputDigest {
44    /// Borrow the raw SHA-256 bytes.
45    #[must_use]
46    pub const fn as_bytes(&self) -> &[u8; 32] {
47        &self.0
48    }
49
50    /// Render the digest as lowercase hexadecimal.
51    #[must_use]
52    pub fn to_hex(self) -> String {
53        let mut hex = String::with_capacity(64);
54        for byte in self.0 {
55            write!(hex, "{byte:02x}").expect("writing to a String cannot fail");
56        }
57        hex
58    }
59}
60
61impl std::fmt::Display for InputDigest {
62    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        formatter.write_str(&self.to_hex())
64    }
65}
66
67pub(super) struct InputHasher(Sha256);
68
69impl InputHasher {
70    pub(super) fn new(domain: &str) -> Self {
71        let mut hasher = Self(Sha256::new());
72        hasher.field("domain", domain.as_bytes());
73        hasher
74    }
75
76    pub(super) fn field(&mut self, label: &str, value: &[u8]) {
77        self.field_header(
78            label,
79            u64::try_from(value.len()).expect("input value length must fit in u64"),
80        );
81        self.0.update(value);
82    }
83
84    fn field_header(&mut self, label: &str, value_len: u64) {
85        self.0.update(
86            u64::try_from(label.len())
87                .expect("input label length must fit in u64")
88                .to_le_bytes(),
89        );
90        self.0.update(label.as_bytes());
91        self.0.update(value_len.to_le_bytes());
92    }
93
94    fn file_field(&mut self, label: &str, path: &Path) -> io::Result<u64> {
95        let mut file = File::open(path)?;
96        let expected_len = file.metadata()?.len();
97        self.field_header(label, expected_len);
98
99        let mut actual_len = 0_u64;
100        let mut buffer = [0_u8; 16 * 1024];
101        loop {
102            let read = file.read(&mut buffer)?;
103            if read == 0 {
104                break;
105            }
106            actual_len = actual_len
107                .saturating_add(u64::try_from(read).expect("artifact read length must fit in u64"));
108            self.0.update(&buffer[..read]);
109        }
110        if actual_len != expected_len {
111            return Err(io::Error::new(
112                io::ErrorKind::InvalidData,
113                format!(
114                    "file changed size while hashing: expected {expected_len} bytes, read {actual_len}"
115                ),
116            ));
117        }
118        Ok(actual_len)
119    }
120
121    pub(super) fn finish(self) -> InputDigest {
122        InputDigest(self.0.finalize().into())
123    }
124}
125
126pub(super) fn digest_bytes(domain: &str, value: &[u8]) -> InputDigest {
127    let mut hasher = InputHasher::new(domain);
128    hasher.field("content", value);
129    hasher.finish()
130}
131
132pub(super) fn digest_file(domain: &str, path: &Path) -> io::Result<(u64, InputDigest)> {
133    let mut hasher = InputHasher::new(domain);
134    let bytes = hasher.file_field("content", path)?;
135    Ok((bytes, hasher.finish()))
136}
137
138pub(super) fn digest_labeled_paths(
139    domain: &str,
140    paths: &[(PathBuf, PathBuf)],
141    excluded_roots: &[PathBuf],
142) -> io::Result<InputDigest> {
143    let mut paths = paths.to_vec();
144    paths.sort_by(|(left, _), (right, _)| {
145        os_bytes(left.as_os_str()).cmp(&os_bytes(right.as_os_str()))
146    });
147
148    let excluded_roots = excluded_roots
149        .iter()
150        .filter_map(|path| path.canonicalize().ok())
151        .collect::<Vec<_>>();
152    let mut visited_directories = BTreeSet::new();
153    let mut hasher = InputHasher::new(domain);
154    for (label, path) in paths {
155        hash_path(
156            &mut hasher,
157            &label,
158            &path,
159            &excluded_roots,
160            &mut visited_directories,
161            true,
162            None,
163        )?;
164    }
165    Ok(hasher.finish())
166}
167
168#[derive(Default)]
169pub(super) struct LabeledPathDigestCache {
170    entries: Vec<LabeledPathDigestCacheEntry>,
171}
172
173struct LabeledPathDigestCacheEntry {
174    domain: String,
175    label: PathBuf,
176    path: PathBuf,
177    canonical_root: PathBuf,
178    excluded_roots: Vec<PathBuf>,
179    traversed_external_path: bool,
180    digest: InputDigest,
181}
182
183struct HashPathTrace {
184    canonical_root: PathBuf,
185    traversed_external_path: bool,
186}
187
188pub(super) fn digest_labeled_paths_composable(
189    domain: &str,
190    paths: &[(PathBuf, PathBuf)],
191    excluded_roots: &[PathBuf],
192    cache: &mut LabeledPathDigestCache,
193) -> io::Result<InputDigest> {
194    let mut paths = paths.to_vec();
195    paths.sort_by(|(left, _), (right, _)| {
196        os_bytes(left.as_os_str()).cmp(&os_bytes(right.as_os_str()))
197    });
198    let excluded_roots = excluded_roots
199        .iter()
200        .filter_map(|path| path.canonicalize().ok())
201        .collect::<Vec<_>>();
202    let mut hasher = InputHasher::new(&format!("{domain}/composable-v1"));
203    for (label, path) in paths {
204        let digest = cache.digest_root(domain, &label, &path, &excluded_roots)?;
205        hasher.field("input-label", &os_bytes(label.as_os_str()));
206        hasher.field("input-digest", digest.as_bytes());
207    }
208    Ok(hasher.finish())
209}
210
211impl LabeledPathDigestCache {
212    fn digest_root(
213        &mut self,
214        domain: &str,
215        label: &Path,
216        path: &Path,
217        excluded_roots: &[PathBuf],
218    ) -> io::Result<InputDigest> {
219        let canonical_root = path.canonicalize()?;
220        if let Some(entry) = self.entries.iter().find(|entry| {
221            entry.domain == domain
222                && entry.label == label
223                && entry.path == path
224                && entry.excluded_roots
225                    == effective_root_exclusions(
226                        &entry.canonical_root,
227                        excluded_roots,
228                        entry.traversed_external_path,
229                    )
230        }) {
231            return Ok(entry.digest);
232        }
233        let mut hasher = InputHasher::new(&format!("{domain}/root-v1"));
234        let mut trace = HashPathTrace {
235            canonical_root: canonical_root.clone(),
236            traversed_external_path: false,
237        };
238        hash_path(
239            &mut hasher,
240            label,
241            path,
242            excluded_roots,
243            &mut BTreeSet::new(),
244            true,
245            Some(&mut trace),
246        )?;
247        let digest = hasher.finish();
248        self.entries.push(LabeledPathDigestCacheEntry {
249            domain: domain.to_owned(),
250            label: label.to_owned(),
251            path: path.to_owned(),
252            canonical_root,
253            excluded_roots: effective_root_exclusions(
254                &trace.canonical_root,
255                excluded_roots,
256                trace.traversed_external_path,
257            ),
258            traversed_external_path: trace.traversed_external_path,
259            digest,
260        });
261        Ok(digest)
262    }
263}
264
265fn effective_root_exclusions(
266    canonical_root: &Path,
267    excluded_roots: &[PathBuf],
268    traversed_external_path: bool,
269) -> Vec<PathBuf> {
270    if traversed_external_path {
271        return excluded_roots.to_vec();
272    }
273    excluded_roots
274        .iter()
275        .filter(|excluded| {
276            excluded.starts_with(canonical_root) || canonical_root.starts_with(excluded)
277        })
278        .cloned()
279        .collect()
280}
281
282fn hash_path(
283    hasher: &mut InputHasher,
284    label: &Path,
285    path: &Path,
286    excluded_roots: &[PathBuf],
287    visited_directories: &mut BTreeSet<PathBuf>,
288    declared_root: bool,
289    mut trace: Option<&mut HashPathTrace>,
290) -> io::Result<()> {
291    let canonical = path.canonicalize()?;
292    if let Some(trace) = &mut trace
293        && !canonical.starts_with(&trace.canonical_root)
294    {
295        trace.traversed_external_path = true;
296    }
297    if excluded_roots
298        .iter()
299        .any(|excluded| canonical.starts_with(excluded))
300    {
301        if declared_root {
302            return Err(io::Error::new(
303                io::ErrorKind::InvalidInput,
304                format!(
305                    "declared input is located inside an excluded cache root: {}",
306                    path.display()
307                ),
308            ));
309        }
310        return Ok(());
311    }
312
313    let metadata = fs::metadata(path)?;
314    let label_bytes = os_bytes(label.as_os_str());
315    if metadata.is_file() {
316        hasher.field("file-path", &label_bytes);
317        hasher.file_field("file-content", path)?;
318        return Ok(());
319    }
320    if !metadata.is_dir() {
321        return Err(io::Error::new(
322            io::ErrorKind::InvalidInput,
323            format!(
324                "watched input is not a regular file or directory: {}",
325                path.display()
326            ),
327        ));
328    }
329
330    hasher.field("directory", &label_bytes);
331    if !visited_directories.insert(canonical) {
332        hasher.field("directory-already-visited", &label_bytes);
333        return Ok(());
334    }
335
336    let mut entries = fs::read_dir(path)?.collect::<Result<Vec<_>, _>>()?;
337    entries.sort_by_key(|entry| os_bytes(&entry.file_name()));
338    for entry in entries {
339        hash_path(
340            hasher,
341            &label.join(entry.file_name()),
342            &entry.path(),
343            excluded_roots,
344            visited_directories,
345            false,
346            trace.as_deref_mut(),
347        )?;
348    }
349    Ok(())
350}
351
352pub(super) fn write_atomic(path: &Path, contents: &[u8]) -> io::Result<()> {
353    write_file_atomic(path, |file| file.write_all(contents))
354}
355
356pub(super) fn copy_file_atomic(source: &Path, destination: &Path) -> io::Result<u64> {
357    let result = (|| {
358        let mut source_file = File::open(source)?;
359        write_file_atomic(destination, |destination_file| {
360            io::copy(&mut source_file, destination_file)
361        })
362    })();
363    result.map_err(|source_error| {
364        io::Error::new(
365            source_error.kind(),
366            AtomicCopyErrorContext {
367                source_path: source.to_owned(),
368                destination_path: destination.to_owned(),
369                source: source_error,
370            },
371        )
372    })
373}
374
375fn write_file_atomic<T>(
376    path: &Path,
377    write: impl FnOnce(&mut File) -> io::Result<T>,
378) -> io::Result<T> {
379    let parent = path.parent().ok_or_else(|| {
380        io::Error::new(
381            io::ErrorKind::InvalidInput,
382            format!("atomic output path has no parent: {}", path.display()),
383        )
384    })?;
385    fs::create_dir_all(parent)?;
386
387    let file_name = path.file_name().ok_or_else(|| {
388        io::Error::new(
389            io::ErrorKind::InvalidInput,
390            format!("atomic output path has no file name: {}", path.display()),
391        )
392    })?;
393    let sequence = TEMP_FILE_SEQUENCE.fetch_add(1, Ordering::Relaxed);
394    let mut temp_name = file_name.to_os_string();
395    temp_name.push(format!(".tmp-{}-{sequence}", std::process::id()));
396    let temp_path = parent.join(temp_name);
397
398    let result = (|| {
399        let mut file = OpenOptions::new()
400            .create_new(true)
401            .write(true)
402            .open(&temp_path)?;
403        let value = write(&mut file)?;
404        file.sync_all()?;
405        fs::rename(&temp_path, path)?;
406        Ok(value)
407    })();
408    if result.is_err() {
409        let _ = fs::remove_file(&temp_path);
410    }
411    result
412}
413
414#[cfg(unix)]
415pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
416    use std::os::unix::ffi::OsStrExt as _;
417    value.as_bytes().to_vec()
418}
419
420#[cfg(windows)]
421pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
422    use std::os::windows::ffi::OsStrExt as _;
423    value
424        .encode_wide()
425        .flat_map(u16::to_le_bytes)
426        .collect::<Vec<_>>()
427}
428
429#[cfg(not(any(unix, windows)))]
430pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
431    value.to_string_lossy().as_bytes().to_vec()
432}
433
434#[cfg(test)]
435mod tests {
436    use super::{
437        LabeledPathDigestCache, copy_file_atomic, digest_bytes, digest_file,
438        digest_labeled_paths_composable, write_atomic,
439    };
440    use crate::artifacts::test_support::unique_temp_directory;
441    use std::{fs, path::PathBuf};
442
443    #[test]
444    fn streaming_digest_and_atomic_copy_preserve_exact_bytes() {
445        let root = unique_temp_directory("streaming-digest");
446        let source = root.join("source");
447        let destination = root.join("destination");
448        let mut contents = vec![0_u8; 192 * 1024];
449        for (index, byte) in contents.iter_mut().enumerate() {
450            *byte = u8::try_from(index % 251).expect("test byte must fit");
451        }
452        fs::write(&source, &contents).expect("write source");
453
454        let (bytes, streamed) = digest_file("streaming-test-v1", &source).expect("digest file");
455        assert_eq!(
456            bytes,
457            u64::try_from(contents.len()).expect("fixture length must fit in u64")
458        );
459        assert_eq!(streamed, digest_bytes("streaming-test-v1", &contents));
460
461        write_atomic(&destination, b"old").expect("write original destination");
462        assert_eq!(
463            copy_file_atomic(&source, &destination).expect("copy source atomically"),
464            bytes
465        );
466        assert_eq!(
467            fs::read(&destination).expect("read copied destination"),
468            contents
469        );
470
471        let missing = root.join("missing");
472        let error = copy_file_atomic(&missing, &destination).expect_err("missing source must fail");
473        let message = error.to_string();
474        assert!(message.contains(&missing.display().to_string()));
475        assert!(message.contains(&destination.display().to_string()));
476        fs::remove_dir_all(root).expect("remove streaming-digest test directory");
477    }
478
479    #[test]
480    fn composable_digest_reuses_roots_across_irrelevant_exclusion_changes() {
481        let root = unique_temp_directory("composable-digest-cache");
482        let input = root.join("input");
483        fs::create_dir_all(&input).expect("create composable input");
484        fs::create_dir_all(root.join("generated-a")).expect("create first generated root");
485        fs::create_dir_all(root.join("generated-b")).expect("create second generated root");
486        fs::write(input.join("source"), b"source").expect("write composable input");
487        let paths = [(PathBuf::from("shared"), input)];
488        let mut cache = LabeledPathDigestCache::default();
489
490        let first = digest_labeled_paths_composable(
491            "composable-test-v1",
492            &paths,
493            &[root.join("generated-a")],
494            &mut cache,
495        )
496        .expect("hash first composable input");
497        let second = digest_labeled_paths_composable(
498            "composable-test-v1",
499            &paths,
500            &[root.join("generated-b")],
501            &mut cache,
502        )
503        .expect("reuse composable input root");
504
505        assert_eq!(first, second);
506        assert_eq!(cache.entries.len(), 1);
507        fs::remove_dir_all(root).expect("remove composable digest fixture");
508    }
509}
ic_testkit/artifacts/digest.rs

ic_testkit/artifacts/
digest.rs
