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        )?;
163    }
164    Ok(hasher.finish())
165}
166
167fn hash_path(
168    hasher: &mut InputHasher,
169    label: &Path,
170    path: &Path,
171    excluded_roots: &[PathBuf],
172    visited_directories: &mut BTreeSet<PathBuf>,
173    declared_root: bool,
174) -> io::Result<()> {
175    let canonical = path.canonicalize()?;
176    if excluded_roots
177        .iter()
178        .any(|excluded| canonical.starts_with(excluded))
179    {
180        if declared_root {
181            return Err(io::Error::new(
182                io::ErrorKind::InvalidInput,
183                format!(
184                    "declared input is located inside an excluded cache root: {}",
185                    path.display()
186                ),
187            ));
188        }
189        return Ok(());
190    }
191
192    let metadata = fs::metadata(path)?;
193    let label_bytes = os_bytes(label.as_os_str());
194    if metadata.is_file() {
195        hasher.field("file-path", &label_bytes);
196        hasher.file_field("file-content", path)?;
197        return Ok(());
198    }
199    if !metadata.is_dir() {
200        return Err(io::Error::new(
201            io::ErrorKind::InvalidInput,
202            format!(
203                "watched input is not a regular file or directory: {}",
204                path.display()
205            ),
206        ));
207    }
208
209    hasher.field("directory", &label_bytes);
210    if !visited_directories.insert(canonical) {
211        hasher.field("directory-already-visited", &label_bytes);
212        return Ok(());
213    }
214
215    let mut entries = fs::read_dir(path)?.collect::<Result<Vec<_>, _>>()?;
216    entries.sort_by_key(|entry| os_bytes(&entry.file_name()));
217    for entry in entries {
218        hash_path(
219            hasher,
220            &label.join(entry.file_name()),
221            &entry.path(),
222            excluded_roots,
223            visited_directories,
224            false,
225        )?;
226    }
227    Ok(())
228}
229
230pub(super) fn write_atomic(path: &Path, contents: &[u8]) -> io::Result<()> {
231    write_file_atomic(path, |file| file.write_all(contents))
232}
233
234pub(super) fn copy_file_atomic(source: &Path, destination: &Path) -> io::Result<u64> {
235    let result = (|| {
236        let mut source_file = File::open(source)?;
237        write_file_atomic(destination, |destination_file| {
238            io::copy(&mut source_file, destination_file)
239        })
240    })();
241    result.map_err(|source_error| {
242        io::Error::new(
243            source_error.kind(),
244            AtomicCopyErrorContext {
245                source_path: source.to_owned(),
246                destination_path: destination.to_owned(),
247                source: source_error,
248            },
249        )
250    })
251}
252
253fn write_file_atomic<T>(
254    path: &Path,
255    write: impl FnOnce(&mut File) -> io::Result<T>,
256) -> io::Result<T> {
257    let parent = path.parent().ok_or_else(|| {
258        io::Error::new(
259            io::ErrorKind::InvalidInput,
260            format!("atomic output path has no parent: {}", path.display()),
261        )
262    })?;
263    fs::create_dir_all(parent)?;
264
265    let file_name = path.file_name().ok_or_else(|| {
266        io::Error::new(
267            io::ErrorKind::InvalidInput,
268            format!("atomic output path has no file name: {}", path.display()),
269        )
270    })?;
271    let sequence = TEMP_FILE_SEQUENCE.fetch_add(1, Ordering::Relaxed);
272    let mut temp_name = file_name.to_os_string();
273    temp_name.push(format!(".tmp-{}-{sequence}", std::process::id()));
274    let temp_path = parent.join(temp_name);
275
276    let result = (|| {
277        let mut file = OpenOptions::new()
278            .create_new(true)
279            .write(true)
280            .open(&temp_path)?;
281        let value = write(&mut file)?;
282        file.sync_all()?;
283        fs::rename(&temp_path, path)?;
284        Ok(value)
285    })();
286    if result.is_err() {
287        let _ = fs::remove_file(&temp_path);
288    }
289    result
290}
291
292#[cfg(unix)]
293pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
294    use std::os::unix::ffi::OsStrExt as _;
295    value.as_bytes().to_vec()
296}
297
298#[cfg(windows)]
299pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
300    use std::os::windows::ffi::OsStrExt as _;
301    value
302        .encode_wide()
303        .flat_map(u16::to_le_bytes)
304        .collect::<Vec<_>>()
305}
306
307#[cfg(not(any(unix, windows)))]
308pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
309    value.to_string_lossy().as_bytes().to_vec()
310}
311
312#[cfg(test)]
313mod tests {
314    use super::{copy_file_atomic, digest_bytes, digest_file, write_atomic};
315    use crate::artifacts::test_support::unique_temp_directory;
316    use std::fs;
317
318    #[test]
319    fn streaming_digest_and_atomic_copy_preserve_exact_bytes() {
320        let root = unique_temp_directory("streaming-digest");
321        let source = root.join("source");
322        let destination = root.join("destination");
323        let mut contents = vec![0_u8; 192 * 1024];
324        for (index, byte) in contents.iter_mut().enumerate() {
325            *byte = u8::try_from(index % 251).expect("test byte must fit");
326        }
327        fs::write(&source, &contents).expect("write source");
328
329        let (bytes, streamed) = digest_file("streaming-test-v1", &source).expect("digest file");
330        assert_eq!(bytes, u64::try_from(contents.len()).unwrap());
331        assert_eq!(streamed, digest_bytes("streaming-test-v1", &contents));
332
333        write_atomic(&destination, b"old").expect("write original destination");
334        assert_eq!(
335            copy_file_atomic(&source, &destination).expect("copy source atomically"),
336            bytes
337        );
338        assert_eq!(fs::read(&destination).unwrap(), contents);
339
340        let missing = root.join("missing");
341        let error = copy_file_atomic(&missing, &destination).expect_err("missing source must fail");
342        let message = error.to_string();
343        assert!(message.contains(&missing.display().to_string()));
344        assert!(message.contains(&destination.display().to_string()));
345        fs::remove_dir_all(root).expect("remove streaming-digest test directory");
346    }
347}
ic_testkit/artifacts/digest.rs

ic_testkit/artifacts/
digest.rs
