rustybit_lib/storage/
piece_hash_verifier.rs

1use std::io::Write;
2
3use anyhow::Context;
4
5use super::util::{find_file_offsets_for_data, read_data_from_files};
6use super::{FileInfo, Storage};
7use crate::state::torrent::PieceState;
8use crate::util::piece_size_from_idx;
9
10pub struct PieceHashVerifier {
11    piece_length: usize,
12    buf: Vec<u8>,
13}
14
15impl PieceHashVerifier {
16    pub fn new(piece_length: usize) -> Self {
17        PieceHashVerifier {
18            piece_length,
19            buf: vec![0; piece_length],
20        }
21    }
22
23    pub fn check_all_pieces(
24        &mut self,
25        storage: &mut dyn Storage,
26        file_infos: &[FileInfo],
27        piece_hashes: &[[u8; 20]],
28        torrent_length: usize,
29    ) -> anyhow::Result<(usize, Vec<PieceState>)> {
30        let number_of_pieces = piece_hashes.len();
31        let mut pieces = (0..number_of_pieces)
32            .map(|_| PieceState::Queued)
33            .collect::<Vec<PieceState>>();
34        let mut verified_pieces = 0;
35        for piece_idx in 0..number_of_pieces {
36            let expected_piece_hash = piece_hashes.get(piece_idx).context("bug: piece with no hash")?;
37            match self
38                .verify_piece_hash(
39                    storage,
40                    file_infos,
41                    try_into!(piece_idx, u32)?,
42                    number_of_pieces,
43                    torrent_length,
44                    expected_piece_hash,
45                )
46                .context("piece hash verification failed")?
47            {
48                Some(has_matching_hash) => {
49                    if has_matching_hash {
50                        verified_pieces += 1;
51                        pieces[piece_idx] = PieceState::Verified;
52                    }
53                }
54                // Data is missing because the file is incomplete, skip further checks
55                None => break,
56            };
57        }
58
59        Ok((verified_pieces, pieces))
60    }
61
62    pub(super) fn verify_piece_hash(
63        &mut self,
64        storage: &mut dyn Storage,
65        file_infos: &[FileInfo],
66        piece_idx: u32,
67        number_of_pieces: usize,
68        torrent_length: usize,
69        expected_hash: &[u8; 20],
70    ) -> anyhow::Result<Option<bool>> {
71        let expected_piece_length = piece_size_from_idx(
72            number_of_pieces,
73            torrent_length,
74            self.piece_length,
75            try_into!(piece_idx, usize)?,
76        );
77        let file_offsets = find_file_offsets_for_data(file_infos, piece_idx, try_into!(self.piece_length, u64)?, None)
78            .context("error while finding offsets for a piece")?
79            .context("bug: failed to find a matching file for a piece?")?;
80
81        let had_enough_bytes =
82            read_data_from_files(storage, &mut self.buf, file_offsets, file_infos, expected_piece_length)
83                .context("error while reading piece from files")?;
84
85        if !had_enough_bytes {
86            // Skip hash verification, as it would fail inevitably
87            Ok(None)
88        } else {
89            let mut hasher = crypto_hash::Hasher::new(crypto_hash::Algorithm::SHA1);
90            hasher.write_all(&self.buf).context("error while updating hasher")?;
91            let mut calculated_hash = [0u8; 20];
92            calculated_hash.copy_from_slice(&hasher.finish());
93            Ok(Some(&calculated_hash == expected_hash))
94        }
95    }
96}