test_assets_ureq/
lib.rs

1/*!
2Download test assets, managing them outside of git
3
4This library downloads test assets using http(s),
5and ensures integrity by comparing those assets to a hash.
6By managing the download separately, you can keep them
7out of VCS and don't make them bloat your repository.
8
9Usage example:
10
11```rust, no_run
12#[test]
13fn some_awesome_test() {
14    let asset_defs = [
15        TestAssetDef {
16            filename : format!("file_a.png"),
17            hash : format!("<sha256 here>"),
18            url : format!("https://url/to/a.png"),
19        },
20        TestAssetDef {
21            filename : format!("file_b.png"),
22            hash : format!("<sha256 here>"),
23            url : format!("https://url/to/a.png"),
24        },
25    ];
26    test_assets::dl_test_files(&asset_defs,
27        "test-assets", true).unwrap();
28    // use your files here
29    // with path under test-assets/file_a.png and test-assets/file_b.png
30}
31```
32
33If you have run the test once, it will re-use the files
34instead of re-downloading them.
35*/
36
37mod hash_list;
38
39use backon::BlockingRetryable;
40use backon::ExponentialBuilder;
41use hash_list::HashList;
42use sha2::digest::Digest;
43use sha2::Sha256;
44use std::fs::{create_dir_all, File};
45use std::io::{self, Read, Write};
46use std::time::Duration;
47use ureq::Agent;
48
49/// Definition for a test file
50///
51///
52pub struct TestAssetDef {
53    /// Name of the file on disk. This should be unique for the file.
54    pub filename: String,
55    /// Sha256 hash of the file's data in hexadecimal lowercase representation
56    pub hash: String,
57    /// The url the test file can be obtained from
58    pub url: String,
59}
60
61/// A type for a Sha256 hash value
62///
63/// Provides conversion functionality to hex representation and back
64#[derive(PartialEq, Eq, Hash, Clone)]
65pub struct Sha256Hash([u8; 32]);
66
67impl Sha256Hash {
68    #[must_use]
69    pub fn from_digest(sha: Sha256) -> Self {
70        let sha = sha.finalize();
71        let bytes = sha[..].try_into().unwrap();
72        Self(bytes)
73    }
74
75    /// Converts the hexadecimal string to a hash value
76    fn from_hex(s: &str) -> Result<Self, ()> {
77        let mut res = Self([0; 32]);
78        let mut idx = 0;
79        let mut iter = s.chars();
80        loop {
81            let upper = match iter.next().and_then(|c| c.to_digit(16)) {
82                Some(v) => v as u8,
83                None => return Err(()),
84            };
85            let lower = match iter.next().and_then(|c| c.to_digit(16)) {
86                Some(v) => v as u8,
87                None => return Err(()),
88            };
89            res.0[idx] = (upper << 4) | lower;
90            idx += 1;
91            if idx == 32 {
92                break;
93            }
94        }
95        Ok(res)
96    }
97    /// Converts the hash value to hexadecimal
98    #[must_use]
99    pub fn to_hex(&self) -> String {
100        let mut res = String::with_capacity(64);
101        for v in &self.0 {
102            use std::char::from_digit;
103            res.push(from_digit(u32::from(*v) >> 4, 16).unwrap());
104            res.push(from_digit(u32::from(*v) & 15, 16).unwrap());
105        }
106        res
107    }
108}
109
110#[derive(Debug)]
111pub enum TaError {
112    Io(io::Error),
113    DownloadFailed,
114    HashMismatch(String, String),
115    BadHashFormat,
116}
117
118impl From<io::Error> for TaError {
119    fn from(err: io::Error) -> Self {
120        Self::Io(err)
121    }
122}
123
124enum DownloadOutcome {
125    WithHash(Sha256Hash),
126}
127
128fn download_test_file(
129    agent: &mut Agent,
130    tfile: &TestAssetDef,
131    dir: &str,
132) -> Result<DownloadOutcome, TaError> {
133    let resp = match agent.get(&tfile.url).call() {
134        Ok(resp) => resp,
135        Err(e) => {
136            println!("{e:?}");
137            return Err(TaError::DownloadFailed);
138        }
139    };
140
141    let len: usize = resp.header("Content-Length").unwrap().parse().unwrap();
142
143    let mut bytes: Vec<u8> = Vec::with_capacity(len);
144    let read_len = resp.into_reader().take(10_000_000_000).read_to_end(&mut bytes)?;
145
146    if (bytes.len() != read_len) && (bytes.len() != len) {
147        return Err(TaError::DownloadFailed);
148    }
149
150    let file = File::create(format!("{}/{}", dir, tfile.filename))?;
151    let mut writer = io::BufWriter::new(file);
152    writer.write_all(&bytes).unwrap();
153
154    let mut hasher = Sha256::new();
155    hasher.update(&bytes);
156
157    Ok(DownloadOutcome::WithHash(Sha256Hash::from_digest(hasher)))
158}
159
160/// Downloads the test files into the passed directory.
161pub fn dl_test_files(defs: &[TestAssetDef], dir: &str, verbose: bool) -> Result<(), TaError> {
162    let mut agent = ureq::agent();
163
164    use std::io::ErrorKind;
165
166    let hash_list_path = format!("{dir}/hash_list");
167    let mut hash_list = match HashList::from_file(&hash_list_path) {
168        Ok(l) => l,
169        Err(TaError::Io(ref e)) if e.kind() == ErrorKind::NotFound => HashList::new(),
170        e => {
171            e?;
172            unreachable!()
173        }
174    };
175    create_dir_all(dir)?;
176    for tfile in defs.iter() {
177        let tfile_hash = Sha256Hash::from_hex(&tfile.hash).map_err(|_| TaError::BadHashFormat)?;
178        if hash_list.get_hash(&tfile.filename).map_or(false, |h| h == &tfile_hash) {
179            // Hash match
180            if verbose {
181                println!(
182                    "File {} has matching hash inside hash list, skipping download",
183                    tfile.filename
184                );
185            }
186            continue;
187        }
188        if verbose {
189            print!("Fetching file {} ...", tfile.filename);
190        }
191        let outcome = download_test_file(&mut agent, tfile, dir)?;
192        match outcome {
193            DownloadOutcome::WithHash(ref hash) => hash_list.add_entry(&tfile.filename, hash),
194        }
195        if verbose {
196            print!("  => ");
197        }
198        match outcome {
199            DownloadOutcome::WithHash(ref found_hash) => {
200                if found_hash == &tfile_hash {
201                    if verbose {
202                        println!("Success")
203                    }
204                } else {
205                    // if the hash mismatches after download, return error
206                    return Err(TaError::HashMismatch(found_hash.to_hex(), tfile.hash.clone()));
207                }
208            }
209        }
210    }
211    hash_list.to_file(&hash_list_path)?;
212    Ok(())
213}
214
215/// Download test-assets with backoff retries
216pub fn dl_test_files_backoff(
217    assets_defs: &[TestAssetDef],
218    test_path: &str,
219    verbose: bool,
220    max_delay: Duration,
221) -> Result<(), TaError> {
222    let strategy = ExponentialBuilder::default().with_max_delay(max_delay);
223
224    (|| dl_test_files(assets_defs, test_path, verbose)).retry(strategy).call().unwrap();
225
226    Ok(())
227}
test_assets_ureq/lib.rs

test_assets_ureq/
lib.rs
