starfield_datastore/
check.rs

1use sha2::{Digest, Sha256};
2use std::{
3    fs::File,
4    io::{Read, Seek, SeekFrom},
5    path::Path,
6    sync::Arc,
7};
8
9pub type Predicate = dyn Fn(&[u8]) -> std::result::Result<(), String> + Send + Sync;
10
11#[derive(Clone)]
12pub enum ContentCheck {
13    Sha256(String),
14    Magic {
15        prefixes: Vec<Vec<u8>>,
16        trim_leading_whitespace: bool,
17    },
18    MinBytes(u64),
19    NotHtml,
20    Custom(Arc<Predicate>),
21    All(Vec<ContentCheck>),
22    None,
23}
24
25impl std::fmt::Debug for ContentCheck {
26    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
27        match self {
28            Self::Sha256(s) => f.debug_tuple("Sha256").field(s).finish(),
29            Self::Magic {
30                prefixes,
31                trim_leading_whitespace,
32            } => f
33                .debug_struct("Magic")
34                .field("prefixes", prefixes)
35                .field("trim_leading_whitespace", trim_leading_whitespace)
36                .finish(),
37            Self::MinBytes(n) => f.debug_tuple("MinBytes").field(n).finish(),
38            Self::NotHtml => f.write_str("NotHtml"),
39            Self::Custom(_) => f.write_str("Custom(<predicate>)"),
40            Self::All(checks) => f.debug_tuple("All").field(checks).finish(),
41            Self::None => f.write_str("None"),
42        }
43    }
44}
45
46#[derive(Debug, Clone)]
47pub struct CheckFailure {
48    pub check: String,
49    pub got: String,
50}
51
52impl ContentCheck {
53    pub fn magic(prefixes: Vec<Vec<u8>>, trim_leading_whitespace: bool) -> Self {
54        Self::Magic {
55            prefixes,
56            trim_leading_whitespace,
57        }
58    }
59    pub fn custom(f: Arc<Predicate>) -> Self {
60        Self::Custom(f)
61    }
62    pub fn default_binary() -> Self {
63        Self::All(vec![Self::NotHtml, Self::MinBytes(1024)])
64    }
65
66    pub fn check(&self, bytes: &[u8]) -> std::result::Result<(), CheckFailure> {
67        self.check_reader(&mut std::io::Cursor::new(bytes), bytes.len() as u64)
68    }
69
70    /// Reject a wrong signature or HTML before downloading the remaining body.
71    /// Inconclusive short prefixes pass; `check_file` remains authoritative.
72    pub fn check_prefix(&self, head: &[u8]) -> std::result::Result<(), CheckFailure> {
73        match self {
74            Self::NotHtml => self.check(head),
75            Self::Magic {
76                prefixes,
77                trim_leading_whitespace,
78            } => {
79                let head = if *trim_leading_whitespace {
80                    head.trim_ascii_start()
81                } else {
82                    head
83                };
84                if prefixes
85                    .iter()
86                    .any(|prefix| head.starts_with(prefix) || prefix.starts_with(head))
87                {
88                    Ok(())
89                } else {
90                    Err(CheckFailure {
91                        check: "Magic".into(),
92                        got: "unexpected signature (body redacted)".into(),
93                    })
94                }
95            }
96            Self::All(checks) => {
97                for check in checks {
98                    check.check_prefix(head)?;
99                }
100                Ok(())
101            }
102            _ => Ok(()),
103        }
104    }
105
106    /// Built-in checks stream files; only a custom predicate loads the full body.
107    pub fn check_file(&self, path: &Path) -> std::result::Result<(), CheckFailure> {
108        if let Self::Sha256(expected) = self {
109            let actual = digest_file(path).map_err(io_failure)?;
110            return if &actual == expected {
111                Ok(())
112            } else {
113                Err(CheckFailure {
114                    check: "Sha256".into(),
115                    got: format!("digest {actual}"),
116                })
117            };
118        }
119        let mut file = File::open(path).map_err(io_failure)?;
120        let len = file.metadata().map_err(io_failure)?.len();
121        self.check_reader(&mut file, len)
122    }
123
124    fn check_reader(
125        &self,
126        reader: &mut (impl Read + Seek),
127        len: u64,
128    ) -> std::result::Result<(), CheckFailure> {
129        reader.seek(SeekFrom::Start(0)).map_err(io_failure)?;
130        let fail = |name: &str, got: String| {
131            Err(CheckFailure {
132                check: name.into(),
133                got,
134            })
135        };
136        match self {
137            Self::None => Ok(()),
138            Self::All(checks) => {
139                for check in checks {
140                    check.check_reader(reader, len)?;
141                }
142                Ok(())
143            }
144            Self::MinBytes(min) if len < *min => {
145                fail("MinBytes", format!("{len} bytes, expected at least {min}"))
146            }
147            Self::MinBytes(_) => Ok(()),
148            Self::Sha256(expected) => {
149                let mut hasher = Sha256::new();
150                std::io::copy(reader, &mut hasher).map_err(io_failure)?;
151                let actual = format!("{:x}", hasher.finalize());
152                if &actual == expected {
153                    Ok(())
154                } else {
155                    fail("Sha256", format!("digest {actual}"))
156                }
157            }
158            Self::Magic {
159                prefixes,
160                trim_leading_whitespace,
161            } => {
162                let max = prefixes.iter().map(Vec::len).max().unwrap_or(0);
163                if prefixes.is_empty() || prefixes.iter().any(Vec::is_empty) {
164                    return fail("Magic", "no nonempty signature configured".into());
165                }
166                let mut buffered = std::io::BufReader::new(reader);
167                let mut first = [0];
168                let mut head = Vec::with_capacity(max);
169                while buffered.read(&mut first).map_err(io_failure)? != 0 {
170                    if !*trim_leading_whitespace || !first[0].is_ascii_whitespace() {
171                        head.push(first[0]);
172                        break;
173                    }
174                }
175                buffered
176                    .take(max.saturating_sub(head.len()) as u64)
177                    .read_to_end(&mut head)
178                    .map_err(io_failure)?;
179                if prefixes.iter().any(|p| head.starts_with(p)) {
180                    Ok(())
181                } else {
182                    fail("Magic", format!("{len} bytes with an unexpected signature"))
183                }
184            }
185            Self::NotHtml => {
186                let mut head = Vec::new();
187                reader
188                    .take(8192)
189                    .read_to_end(&mut head)
190                    .map_err(io_failure)?;
191                let head = String::from_utf8_lossy(&head).to_ascii_lowercase();
192                let head = head.trim_start_matches('\u{feff}').trim_start();
193                if head.starts_with("<!doctype html")
194                    || head.starts_with("<html")
195                    || head.starts_with("<head")
196                    || head.starts_with("<body")
197                    || head.contains("<form")
198                    || head.contains("<html")
199                {
200                    fail(
201                        "NotHtml",
202                        "HTML document or login form (body redacted)".into(),
203                    )
204                } else {
205                    Ok(())
206                }
207            }
208            Self::Custom(f) => {
209                let mut body = Vec::new();
210                reader.read_to_end(&mut body).map_err(io_failure)?;
211                f(&body).map_err(|got| CheckFailure {
212                    check: "Custom".into(),
213                    got,
214                })
215            }
216        }
217    }
218}
219
220fn io_failure(error: std::io::Error) -> CheckFailure {
221    CheckFailure {
222        check: "Read".into(),
223        got: error.to_string(),
224    }
225}
226
227pub(crate) fn digest_file(path: &Path) -> std::io::Result<String> {
228    let mut file = File::open(path)?;
229    let mut hash = Sha256::new();
230    std::io::copy(&mut file, &mut hash)?;
231    Ok(format!("{:x}", hash.finalize()))
232}
233
234#[cfg(test)]
235mod tests {
236    use super::*;
237    #[test]
238    fn default_rejects_html_and_short_bodies() {
239        assert!(ContentCheck::default_binary().check(&vec![0; 1024]).is_ok());
240        assert!(ContentCheck::default_binary().check(b"small").is_err());
241        assert!(ContentCheck::default_binary()
242            .check(format!(" \n<html><form>{}", "x".repeat(2000)).as_bytes())
243            .is_err());
244    }
245    #[test]
246    fn magic_trim_is_explicit_and_custom_captures() {
247        assert!(ContentCheck::magic(vec![b"ABC".to_vec()], true)
248            .check(b"\n ABC")
249            .is_ok());
250        assert!(ContentCheck::magic(vec![b"ABC".to_vec()], false)
251            .check(b" ABC")
252            .is_err());
253        let prefix = vec![1, 2];
254        let check = ContentCheck::custom(Arc::new(move |bytes| {
255            if bytes.starts_with(&prefix) {
256                Ok(())
257            } else {
258                Err("wrong prefix".into())
259            }
260        }));
261        assert!(check.check(&[1, 2, 3]).is_ok());
262        assert_eq!(format!("{check:?}"), "Custom(<predicate>)");
263    }
264
265    #[test]
266    fn prefix_checks_reject_wrong_kind_and_defer_incomplete_validation() {
267        let check = ContentCheck::All(vec![
268            ContentCheck::NotHtml,
269            ContentCheck::MinBytes(10000),
270            ContentCheck::magic(vec![b"ABC".to_vec()], true),
271        ]);
272        assert!(check.check_prefix(b" \nA").is_ok());
273        assert!(check.check_prefix(b"XYZ").is_err());
274        assert!(check.check_prefix(b"<html>login</html>").is_err());
275        assert!(check.check_prefix(b"ABC").is_ok());
276        assert!(check.check(b"ABC").is_err());
277    }
278}
starfield_datastore/check.rs

starfield_datastore/
check.rs
