1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
// Copyright 2021-2022 System76 <info@system76.com>
// SPDX-License-Identifier: MPL-2.0

use crate::checksum::{Checksum, ChecksumError};
use futures::prelude::*;
use remem::Pool;
use std::{path::Path, sync::Arc};

/// Generates a stream of futures that validate checksums.
///
/// The caller can choose to distribute these futures across a thread pool.
///
/// ```
/// let mut stream = checksum_stream(checksums).map(tokio::spawn).buffered(8);
/// while let Some((path, result)) = stream.next().await {
///     eprintln!("{:?} checksum result: {:?}", path, result);
/// }
/// ```
pub fn checksum_stream<I: Stream<Item = (Arc<Path>, Checksum)> + Send + Unpin + 'static>(
    inputs: I,
) -> impl Stream<Item = impl Future<Output = (Arc<Path>, Result<(), ChecksumError>)>> {
    let buffer_pool = Pool::new(|| Box::new([0u8; 8 * 1024]));

    inputs.map(move |(dest, checksum)| {
        let pool = buffer_pool.clone();

        async {
            tokio::task::spawn_blocking(move || {
                let buf = &mut **pool.get();
                let result = validate_checksum(buf, &dest, &checksum);
                (dest, result)
            })
            .await
            .unwrap()
        }
    })
}

/// Validates the checksum of a single file
pub fn validate_checksum(
    buf: &mut [u8],
    dest: &Path,
    checksum: &Checksum,
) -> Result<(), ChecksumError> {
    let error = match std::fs::File::open(&*dest) {
        Ok(file) => match checksum.validate(file, buf) {
            Ok(()) => return Ok(()),
            Err(why) => why,
        },
        Err(why) => ChecksumError::from(why),
    };

    let _ = std::fs::remove_file(&*dest);
    Err(error)
}