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/// SHA-256 digest of one deterministic artifact-input set.
15#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)]
16pub struct InputDigest([u8; 32]);
17
18impl InputDigest {
19    /// Borrow the raw SHA-256 bytes.
20    #[must_use]
21    pub const fn as_bytes(&self) -> &[u8; 32] {
22        &self.0
23    }
24
25    /// Render the digest as lowercase hexadecimal.
26    #[must_use]
27    pub fn to_hex(self) -> String {
28        let mut hex = String::with_capacity(64);
29        for byte in self.0 {
30            write!(hex, "{byte:02x}").expect("writing to a String cannot fail");
31        }
32        hex
33    }
34}
35
36impl std::fmt::Display for InputDigest {
37    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38        formatter.write_str(&self.to_hex())
39    }
40}
41
42pub(super) struct InputHasher(Sha256);
43
44impl InputHasher {
45    pub(super) fn new(domain: &str) -> Self {
46        let mut hasher = Self(Sha256::new());
47        hasher.field("domain", domain.as_bytes());
48        hasher
49    }
50
51    pub(super) fn field(&mut self, label: &str, value: &[u8]) {
52        self.field_header(
53            label,
54            u64::try_from(value.len()).expect("input value length must fit in u64"),
55        );
56        self.0.update(value);
57    }
58
59    fn field_header(&mut self, label: &str, value_len: u64) {
60        self.0.update(
61            u64::try_from(label.len())
62                .expect("input label length must fit in u64")
63                .to_le_bytes(),
64        );
65        self.0.update(label.as_bytes());
66        self.0.update(value_len.to_le_bytes());
67    }
68
69    fn file_field(&mut self, label: &str, path: &Path) -> io::Result<u64> {
70        let mut file = File::open(path)?;
71        let expected_len = file.metadata()?.len();
72        self.field_header(label, expected_len);
73
74        let mut actual_len = 0_u64;
75        let mut buffer = [0_u8; 16 * 1024];
76        loop {
77            let read = file.read(&mut buffer)?;
78            if read == 0 {
79                break;
80            }
81            actual_len = actual_len
82                .saturating_add(u64::try_from(read).expect("artifact read length must fit in u64"));
83            self.0.update(&buffer[..read]);
84        }
85        if actual_len != expected_len {
86            return Err(io::Error::new(
87                io::ErrorKind::InvalidData,
88                format!(
89                    "file changed size while hashing: expected {expected_len} bytes, read {actual_len}"
90                ),
91            ));
92        }
93        Ok(actual_len)
94    }
95
96    pub(super) fn finish(self) -> InputDigest {
97        InputDigest(self.0.finalize().into())
98    }
99}
100
101pub(super) fn digest_bytes(domain: &str, value: &[u8]) -> InputDigest {
102    let mut hasher = InputHasher::new(domain);
103    hasher.field("content", value);
104    hasher.finish()
105}
106
107pub(super) fn digest_file(domain: &str, path: &Path) -> io::Result<(u64, InputDigest)> {
108    let mut hasher = InputHasher::new(domain);
109    let bytes = hasher.file_field("content", path)?;
110    Ok((bytes, hasher.finish()))
111}
112
113pub(super) fn digest_labeled_paths(
114    domain: &str,
115    paths: &[(PathBuf, PathBuf)],
116    excluded_roots: &[PathBuf],
117) -> io::Result<InputDigest> {
118    let mut paths = paths.to_vec();
119    paths.sort_by(|(left, _), (right, _)| {
120        os_bytes(left.as_os_str()).cmp(&os_bytes(right.as_os_str()))
121    });
122
123    let excluded_roots = excluded_roots
124        .iter()
125        .filter_map(|path| path.canonicalize().ok())
126        .collect::<Vec<_>>();
127    let mut visited_directories = BTreeSet::new();
128    let mut hasher = InputHasher::new(domain);
129    for (label, path) in paths {
130        hash_path(
131            &mut hasher,
132            &label,
133            &path,
134            &excluded_roots,
135            &mut visited_directories,
136            true,
137        )?;
138    }
139    Ok(hasher.finish())
140}
141
142fn hash_path(
143    hasher: &mut InputHasher,
144    label: &Path,
145    path: &Path,
146    excluded_roots: &[PathBuf],
147    visited_directories: &mut BTreeSet<PathBuf>,
148    declared_root: bool,
149) -> io::Result<()> {
150    let canonical = path.canonicalize()?;
151    if excluded_roots
152        .iter()
153        .any(|excluded| canonical.starts_with(excluded))
154    {
155        if declared_root {
156            return Err(io::Error::new(
157                io::ErrorKind::InvalidInput,
158                format!(
159                    "declared input is located inside an excluded cache root: {}",
160                    path.display()
161                ),
162            ));
163        }
164        return Ok(());
165    }
166
167    let metadata = fs::metadata(path)?;
168    let label_bytes = os_bytes(label.as_os_str());
169    if metadata.is_file() {
170        hasher.field("file-path", &label_bytes);
171        hasher.file_field("file-content", path)?;
172        return Ok(());
173    }
174    if !metadata.is_dir() {
175        return Err(io::Error::new(
176            io::ErrorKind::InvalidInput,
177            format!(
178                "watched input is not a regular file or directory: {}",
179                path.display()
180            ),
181        ));
182    }
183
184    hasher.field("directory", &label_bytes);
185    if !visited_directories.insert(canonical) {
186        hasher.field("directory-already-visited", &label_bytes);
187        return Ok(());
188    }
189
190    let mut entries = fs::read_dir(path)?.collect::<Result<Vec<_>, _>>()?;
191    entries.sort_by_key(|entry| os_bytes(&entry.file_name()));
192    for entry in entries {
193        hash_path(
194            hasher,
195            &label.join(entry.file_name()),
196            &entry.path(),
197            excluded_roots,
198            visited_directories,
199            false,
200        )?;
201    }
202    Ok(())
203}
204
205pub(super) fn write_atomic(path: &Path, contents: &[u8]) -> io::Result<()> {
206    write_file_atomic(path, |file| file.write_all(contents))
207}
208
209pub(super) fn copy_file_atomic(source: &Path, destination: &Path) -> io::Result<u64> {
210    let mut source_file = File::open(source)?;
211    write_file_atomic(destination, |destination_file| {
212        io::copy(&mut source_file, destination_file)
213    })
214}
215
216fn write_file_atomic<T>(
217    path: &Path,
218    write: impl FnOnce(&mut File) -> io::Result<T>,
219) -> io::Result<T> {
220    let parent = path.parent().ok_or_else(|| {
221        io::Error::new(
222            io::ErrorKind::InvalidInput,
223            format!("atomic output path has no parent: {}", path.display()),
224        )
225    })?;
226    fs::create_dir_all(parent)?;
227
228    let file_name = path.file_name().ok_or_else(|| {
229        io::Error::new(
230            io::ErrorKind::InvalidInput,
231            format!("atomic output path has no file name: {}", path.display()),
232        )
233    })?;
234    let sequence = TEMP_FILE_SEQUENCE.fetch_add(1, Ordering::Relaxed);
235    let mut temp_name = file_name.to_os_string();
236    temp_name.push(format!(".tmp-{}-{sequence}", std::process::id()));
237    let temp_path = parent.join(temp_name);
238
239    let result = (|| {
240        let mut file = OpenOptions::new()
241            .create_new(true)
242            .write(true)
243            .open(&temp_path)?;
244        let value = write(&mut file)?;
245        file.sync_all()?;
246        fs::rename(&temp_path, path)?;
247        Ok(value)
248    })();
249    if result.is_err() {
250        let _ = fs::remove_file(&temp_path);
251    }
252    result
253}
254
255#[cfg(unix)]
256pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
257    use std::os::unix::ffi::OsStrExt as _;
258    value.as_bytes().to_vec()
259}
260
261#[cfg(windows)]
262pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
263    use std::os::windows::ffi::OsStrExt as _;
264    value
265        .encode_wide()
266        .flat_map(u16::to_le_bytes)
267        .collect::<Vec<_>>()
268}
269
270#[cfg(not(any(unix, windows)))]
271pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
272    value.to_string_lossy().as_bytes().to_vec()
273}
274
275#[cfg(test)]
276mod tests {
277    use super::{copy_file_atomic, digest_bytes, digest_file, write_atomic};
278    use crate::artifacts::test_support::unique_temp_directory;
279    use std::fs;
280
281    #[test]
282    fn streaming_digest_and_atomic_copy_preserve_exact_bytes() {
283        let root = unique_temp_directory("streaming-digest");
284        let source = root.join("source");
285        let destination = root.join("destination");
286        let mut contents = vec![0_u8; 192 * 1024];
287        for (index, byte) in contents.iter_mut().enumerate() {
288            *byte = u8::try_from(index % 251).expect("test byte must fit");
289        }
290        fs::write(&source, &contents).expect("write source");
291
292        let (bytes, streamed) = digest_file("streaming-test-v1", &source).expect("digest file");
293        assert_eq!(bytes, u64::try_from(contents.len()).unwrap());
294        assert_eq!(streamed, digest_bytes("streaming-test-v1", &contents));
295
296        write_atomic(&destination, b"old").expect("write original destination");
297        assert_eq!(
298            copy_file_atomic(&source, &destination).expect("copy source atomically"),
299            bytes
300        );
301        assert_eq!(fs::read(&destination).unwrap(), contents);
302        fs::remove_dir_all(root).expect("remove streaming-digest test directory");
303    }
304}
ic_testkit/artifacts/digest.rs

ic_testkit/artifacts/
digest.rs
