Skip to main content

ic_testkit/artifacts/
digest.rs

1use sha2::{Digest, Sha256};
2use std::{
3    collections::BTreeSet,
4    ffi::OsStr,
5    fmt::Write as _,
6    fs::{self, OpenOptions},
7    io::{self, 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.0.update(
53            u64::try_from(label.len())
54                .expect("input label length must fit in u64")
55                .to_le_bytes(),
56        );
57        self.0.update(label.as_bytes());
58        self.0.update(
59            u64::try_from(value.len())
60                .expect("input value length must fit in u64")
61                .to_le_bytes(),
62        );
63        self.0.update(value);
64    }
65
66    pub(super) fn finish(self) -> InputDigest {
67        InputDigest(self.0.finalize().into())
68    }
69}
70
71pub(super) fn digest_bytes(domain: &str, value: &[u8]) -> InputDigest {
72    let mut hasher = InputHasher::new(domain);
73    hasher.field("content", value);
74    hasher.finish()
75}
76
77pub(super) fn digest_labeled_paths(
78    domain: &str,
79    paths: &[(PathBuf, PathBuf)],
80    excluded_roots: &[PathBuf],
81) -> io::Result<InputDigest> {
82    let mut paths = paths.to_vec();
83    paths.sort_by(|(left, _), (right, _)| {
84        os_bytes(left.as_os_str()).cmp(&os_bytes(right.as_os_str()))
85    });
86
87    let excluded_roots = excluded_roots
88        .iter()
89        .filter_map(|path| path.canonicalize().ok())
90        .collect::<Vec<_>>();
91    let mut visited_directories = BTreeSet::new();
92    let mut hasher = InputHasher::new(domain);
93    for (label, path) in paths {
94        hash_path(
95            &mut hasher,
96            &label,
97            &path,
98            &excluded_roots,
99            &mut visited_directories,
100        )?;
101    }
102    Ok(hasher.finish())
103}
104
105fn hash_path(
106    hasher: &mut InputHasher,
107    label: &Path,
108    path: &Path,
109    excluded_roots: &[PathBuf],
110    visited_directories: &mut BTreeSet<PathBuf>,
111) -> io::Result<()> {
112    let canonical = path.canonicalize()?;
113    if excluded_roots
114        .iter()
115        .any(|excluded| canonical.starts_with(excluded))
116    {
117        return Ok(());
118    }
119
120    let metadata = fs::metadata(path)?;
121    let label_bytes = os_bytes(label.as_os_str());
122    if metadata.is_file() {
123        hasher.field("file-path", &label_bytes);
124        hasher.field("file-content", &fs::read(path)?);
125        return Ok(());
126    }
127    if !metadata.is_dir() {
128        return Err(io::Error::new(
129            io::ErrorKind::InvalidInput,
130            format!(
131                "watched input is not a regular file or directory: {}",
132                path.display()
133            ),
134        ));
135    }
136
137    hasher.field("directory", &label_bytes);
138    if !visited_directories.insert(canonical) {
139        hasher.field("directory-already-visited", &label_bytes);
140        return Ok(());
141    }
142
143    let mut entries = fs::read_dir(path)?.collect::<Result<Vec<_>, _>>()?;
144    entries.sort_by_key(|entry| os_bytes(&entry.file_name()));
145    for entry in entries {
146        hash_path(
147            hasher,
148            &label.join(entry.file_name()),
149            &entry.path(),
150            excluded_roots,
151            visited_directories,
152        )?;
153    }
154    Ok(())
155}
156
157pub(super) fn write_atomic(path: &Path, contents: &[u8]) -> io::Result<()> {
158    let parent = path.parent().ok_or_else(|| {
159        io::Error::new(
160            io::ErrorKind::InvalidInput,
161            format!("atomic output path has no parent: {}", path.display()),
162        )
163    })?;
164    fs::create_dir_all(parent)?;
165
166    let file_name = path.file_name().ok_or_else(|| {
167        io::Error::new(
168            io::ErrorKind::InvalidInput,
169            format!("atomic output path has no file name: {}", path.display()),
170        )
171    })?;
172    let sequence = TEMP_FILE_SEQUENCE.fetch_add(1, Ordering::Relaxed);
173    let mut temp_name = file_name.to_os_string();
174    temp_name.push(format!(".tmp-{}-{sequence}", std::process::id()));
175    let temp_path = parent.join(temp_name);
176
177    let result = (|| {
178        let mut file = OpenOptions::new()
179            .create_new(true)
180            .write(true)
181            .open(&temp_path)?;
182        file.write_all(contents)?;
183        file.sync_all()?;
184        fs::rename(&temp_path, path)
185    })();
186    if result.is_err() {
187        let _ = fs::remove_file(&temp_path);
188    }
189    result
190}
191
192#[cfg(unix)]
193pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
194    use std::os::unix::ffi::OsStrExt as _;
195    value.as_bytes().to_vec()
196}
197
198#[cfg(windows)]
199pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
200    use std::os::windows::ffi::OsStrExt as _;
201    value
202        .encode_wide()
203        .flat_map(u16::to_le_bytes)
204        .collect::<Vec<_>>()
205}
206
207#[cfg(not(any(unix, windows)))]
208pub(super) fn os_bytes(value: &OsStr) -> Vec<u8> {
209    value.to_string_lossy().as_bytes().to_vec()
210}