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 context =
292        |error: io::Error| io::Error::new(error.kind(), format!("{}: {error}", path.display()));
293    let canonical = path.canonicalize().map_err(context)?;
294    if let Some(trace) = &mut trace
295        && !canonical.starts_with(&trace.canonical_root)
296    {
297        trace.traversed_external_path = true;
298    }
299    if excluded_roots
300        .iter()
301        .any(|excluded| canonical.starts_with(excluded))
302    {
303        if declared_root {
304            return Err(io::Error::new(
305                io::ErrorKind::InvalidInput,
306                format!(
307                    "declared input is located inside an excluded cache root: {}",
308                    path.display()
309                ),
310            ));
311        }
312        return Ok(());
313    }
314
315    let metadata = fs::metadata(path).map_err(context)?;
316    let label_bytes = os_bytes(label.as_os_str());
317    if metadata.is_file() {
318        hasher.field("file-path", &label_bytes);
319        hasher.file_field("file-content", path).map_err(context)?;
320        return Ok(());
321    }
322    if !metadata.is_dir() {
323        return Err(io::Error::new(
324            io::ErrorKind::InvalidInput,
325            format!(
326                "watched input is not a regular file or directory: {}",
327                path.display()
328            ),
329        ));
330    }
331
332    hasher.field("directory", &label_bytes);
333    if !visited_directories.insert(canonical) {
334        hasher.field("directory-already-visited", &label_bytes);
335        return Ok(());
336    }
337
338    let mut entries = fs::read_dir(path)
339        .map_err(context)?
340        .collect::<Result<Vec<_>, _>>()
341        .map_err(context)?;
342    entries.sort_by_key(|entry| os_bytes(&entry.file_name()));
343    for entry in entries {
344        hash_path(
345            hasher,
346            &label.join(entry.file_name()),
347            &entry.path(),
348            excluded_roots,
349            visited_directories,
350            false,
351            trace.as_deref_mut(),
352        )?;
353    }
354    Ok(())
355}
356
357pub(super) fn write_atomic(path: &Path, contents: &[u8]) -> io::Result<()> {
358    write_file_atomic(path, |file| file.write_all(contents))
359}
360
361pub(super) fn copy_file_atomic(source: &Path, destination: &Path) -> io::Result<u64> {
362    let result = (|| {
363        let mut source_file = File::open(source)?;
364        write_file_atomic(destination, |destination_file| {
365            io::copy(&mut source_file, destination_file)
366        })
367    })();
368    result.map_err(|source_error| {
369        io::Error::new(
370            source_error.kind(),
371            AtomicCopyErrorContext {
372                source_path: source.to_owned(),
373                destination_path: destination.to_owned(),
374                source: source_error,
375            },
376        )
377    })
378}
379
380fn write_file_atomic<T>(
381    path: &Path,
382    write: impl FnOnce(&mut File) -> io::Result<T>,
383) -> io::Result<T> {
384    let parent = path.parent().ok_or_else(|| {
385        io::Error::new(
386            io::ErrorKind::InvalidInput,
387            format!("atomic output path has no parent: {}", path.display()),
388        )
389    })?;
390    fs::create_dir_all(parent)?;
391
392    let file_name = path.file_name().ok_or_else(|| {
393        io::Error::new(
394            io::ErrorKind::InvalidInput,
395            format!("atomic output path has no file name: {}", path.display()),
396        )
397    })?;
398    let sequence = TEMP_FILE_SEQUENCE.fetch_add(1, Ordering::Relaxed);
399    let mut temp_name = file_name.to_os_string();
400    temp_name.push(format!(".tmp-{}-{sequence}", std::process::id()));
401    let temp_path = parent.join(temp_name);
402
403    let result = (|| {
404        let mut file = OpenOptions::new()
405            .create_new(true)
406            .write(true)
407            .open(&temp_path)?;
408        let value = write(&mut file)?;
409        file.sync_all()?;
410        fs::rename(&temp_path, path)?;
411        Ok(value)
412    })();
413    if result.is_err() {
414        let _ = fs::remove_file(&temp_path);
415    }
416    result
417}
418
419#[cfg(unix)]
420pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
421    use std::os::unix::ffi::OsStrExt as _;
422    value.as_bytes().to_vec()
423}
424
425#[cfg(windows)]
426pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
427    use std::os::windows::ffi::OsStrExt as _;
428    value
429        .encode_wide()
430        .flat_map(u16::to_le_bytes)
431        .collect::<Vec<_>>()
432}
433
434#[cfg(not(any(unix, windows)))]
435pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
436    value.to_string_lossy().as_bytes().to_vec()
437}
438
439#[cfg(test)]
440mod tests {
441    use super::{
442        LabeledPathDigestCache, copy_file_atomic, digest_bytes, digest_file,
443        digest_labeled_paths_composable, write_atomic,
444    };
445    use crate::artifacts::test_support::unique_temp_directory;
446    use std::{fs, path::PathBuf};
447
448    #[test]
449    fn streaming_digest_and_atomic_copy_preserve_exact_bytes() {
450        let root = unique_temp_directory("streaming-digest");
451        let source = root.join("source");
452        let destination = root.join("destination");
453        let mut contents = vec![0_u8; 192 * 1024];
454        for (index, byte) in contents.iter_mut().enumerate() {
455            *byte = u8::try_from(index % 251).expect("test byte must fit");
456        }
457        fs::write(&source, &contents).expect("write source");
458
459        let (bytes, streamed) = digest_file("streaming-test-v1", &source).expect("digest file");
460        assert_eq!(
461            bytes,
462            u64::try_from(contents.len()).expect("fixture length must fit in u64")
463        );
464        assert_eq!(streamed, digest_bytes("streaming-test-v1", &contents));
465
466        write_atomic(&destination, b"old").expect("write original destination");
467        assert_eq!(
468            copy_file_atomic(&source, &destination).expect("copy source atomically"),
469            bytes
470        );
471        assert_eq!(
472            fs::read(&destination).expect("read copied destination"),
473            contents
474        );
475
476        let missing = root.join("missing");
477        let error = copy_file_atomic(&missing, &destination).expect_err("missing source must fail");
478        let message = error.to_string();
479        assert!(message.contains(&missing.display().to_string()));
480        assert!(message.contains(&destination.display().to_string()));
481        fs::remove_dir_all(root).expect("remove streaming-digest test directory");
482    }
483
484    #[test]
485    fn composable_digest_reuses_roots_across_irrelevant_exclusion_changes() {
486        let root = unique_temp_directory("composable-digest-cache");
487        let input = root.join("input");
488        fs::create_dir_all(&input).expect("create composable input");
489        fs::create_dir_all(root.join("generated-a")).expect("create first generated root");
490        fs::create_dir_all(root.join("generated-b")).expect("create second generated root");
491        fs::write(input.join("source"), b"source").expect("write composable input");
492        let paths = [(PathBuf::from("shared"), input)];
493        let mut cache = LabeledPathDigestCache::default();
494
495        let first = digest_labeled_paths_composable(
496            "composable-test-v1",
497            &paths,
498            &[root.join("generated-a")],
499            &mut cache,
500        )
501        .expect("hash first composable input");
502        let second = digest_labeled_paths_composable(
503            "composable-test-v1",
504            &paths,
505            &[root.join("generated-b")],
506            &mut cache,
507        )
508        .expect("reuse composable input root");
509
510        assert_eq!(first, second);
511        assert_eq!(cache.entries.len(), 1);
512        fs::remove_dir_all(root).expect("remove composable digest fixture");
513    }
514}
ic_testkit/artifacts/digest.rs

ic_testkit/artifacts/
digest.rs
