use crate::metainfo::{InfoHash, TorrentKind};
use alloc::vec::Vec;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum HashKind {
Sha1,
Sha256Merkle,
}
impl From<TorrentKind> for HashKind {
fn from(k: TorrentKind) -> Self {
match k {
TorrentKind::V1 => HashKind::Sha1,
TorrentKind::V2 | TorrentKind::Hybrid => HashKind::Sha256Merkle,
}
}
}
pub struct VerifyJob {
pub torrent: InfoHash,
pub piece: u32,
pub len: u32,
pub kind: HashKind,
pub expect: Vec<u8>,
pub data: Vec<u8>,
}
pub struct VerifyResult {
pub torrent: InfoHash,
pub piece: u32,
pub ok: bool,
pub data: Vec<u8>,
}
pub fn verify_piece(kind: HashKind, len: u32, data: &[u8], expect: &[u8]) -> bool {
if data.len() != len as usize {
return false;
}
match kind {
HashKind::Sha1 => expect.len() == 20 && crate::crypto::Sha1::digest(data) == expect[..20],
HashKind::Sha256Merkle => {
if expect.len() != 32 {
return false;
}
let blocks: Vec<[u8; 32]> = data
.chunks(crate::consts::BLOCK_LEN as usize)
.map(crate::crypto::Sha256::digest)
.collect();
crate::metainfo::merkle_root(&blocks) == expect[..32]
}
}
}
#[cfg(feature = "std")]
mod pool {
use super::{VerifyJob, VerifyResult};
use alloc::vec::Vec;
use std::sync::mpsc;
pub(crate) struct VerifyPoolInner {
txs: Vec<mpsc::Sender<VerifyJob>>,
rxs: Vec<mpsc::Receiver<VerifyResult>>,
workers: Vec<std::thread::JoinHandle<()>>,
}
impl VerifyPoolInner {
pub(crate) fn spawn(workers: usize) -> Self {
let n = workers.max(1);
let mut txs = Vec::with_capacity(n);
let mut rxs = Vec::with_capacity(n);
let mut handles = Vec::with_capacity(n);
for _ in 0..n {
let (job_tx, job_rx) = mpsc::channel::<VerifyJob>();
let (res_tx, res_rx) = mpsc::channel::<VerifyResult>();
txs.push(job_tx);
rxs.push(res_rx);
handles.push(std::thread::spawn(move || {
while let Ok(job) = job_rx.recv() {
let ok = super::verify_piece(job.kind, job.len, &job.data, &job.expect);
if res_tx
.send(VerifyResult {
torrent: job.torrent,
piece: job.piece,
ok,
data: job.data,
})
.is_err()
{
break;
}
}
}));
}
VerifyPoolInner {
txs,
rxs,
workers: handles,
}
}
pub(crate) fn submit(&self, job: VerifyJob) {
let idx = (job.piece as usize) % self.txs.len();
let _ = self.txs[idx].send(job);
}
pub(crate) fn poll(&self) -> Option<VerifyResult> {
for rx in &self.rxs {
if let Ok(res) = rx.try_recv() {
return Some(res);
}
}
None
}
}
impl Drop for VerifyPoolInner {
fn drop(&mut self) {
self.txs.clear();
for w in self.workers.drain(..) {
let _ = w.join();
}
}
}
}
pub struct VerifyPool {
#[cfg(feature = "std")]
inner: pool::VerifyPoolInner,
}
impl VerifyPool {
pub fn spawn(workers: usize) -> Self {
#[cfg(feature = "std")]
{
VerifyPool {
inner: pool::VerifyPoolInner::spawn(workers),
}
}
#[cfg(not(feature = "std"))]
{
let _ = workers;
VerifyPool {}
}
}
pub fn submit(&self, job: VerifyJob) {
#[cfg(feature = "std")]
self.inner.submit(job);
#[cfg(not(feature = "std"))]
{
let _ = job;
}
}
pub fn poll(&self) -> Option<VerifyResult> {
#[cfg(feature = "std")]
{
self.inner.poll()
}
#[cfg(not(feature = "std"))]
{
None
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn sha1_verify_matches_and_rejects() {
let data = b"the quick brown fox".to_vec();
let mut expect = crate::crypto::Sha1::digest(&data).to_vec();
assert!(verify_piece(
HashKind::Sha1,
data.len() as u32,
&data,
&expect
));
expect[0] ^= 0xff;
assert!(!verify_piece(
HashKind::Sha1,
data.len() as u32,
&data,
&expect
));
}
#[test]
fn v2_merkle_verify_matches_reference() {
let data: Vec<u8> = (0..(48 * 1024) as u32).map(|i| (i % 251) as u8).collect();
let blocks: Vec<[u8; 32]> = data
.chunks(crate::consts::BLOCK_LEN as usize)
.map(crate::crypto::Sha256::digest)
.collect();
let root = crate::metainfo::merkle_root(&blocks);
assert!(verify_piece(
HashKind::Sha256Merkle,
data.len() as u32,
&data,
&root
));
assert!(!verify_piece(
HashKind::Sha256Merkle,
(data.len() - 1) as u32,
&data,
&root
));
let mut bad = root;
bad[0] ^= 1;
assert!(!verify_piece(
HashKind::Sha256Merkle,
data.len() as u32,
&data,
&bad
));
}
#[test]
fn hash_kind_maps_torrent_kind() {
assert_eq!(HashKind::from(TorrentKind::V1), HashKind::Sha1);
assert_eq!(HashKind::from(TorrentKind::V2), HashKind::Sha256Merkle);
assert_eq!(HashKind::from(TorrentKind::Hybrid), HashKind::Sha256Merkle);
}
#[cfg(feature = "std")]
#[test]
fn pool_verifies_across_workers() {
let pool = VerifyPool::spawn(2);
let torrent = InfoHash::v1([1u8; 20]);
let data: Vec<u8> = (0..(16 * 1024) as u32).map(|i| (i % 251) as u8).collect();
let expect = crate::crypto::Sha1::digest(&data).to_vec();
for p in 0..4u32 {
pool.submit(VerifyJob {
torrent,
piece: p,
len: data.len() as u32,
kind: HashKind::Sha1,
expect: expect.clone(),
data: data.clone(),
});
}
let mut got = 0;
let mut attempts = 0u32;
while got < 4 && attempts < 100_000 {
if let Some(res) = pool.poll() {
assert!(res.ok, "piece {} should verify", res.piece);
assert_eq!(res.torrent, torrent);
assert_eq!(res.data, data);
got += 1;
} else {
std::thread::yield_now();
attempts += 1;
}
}
assert_eq!(got, 4, "all four jobs must complete");
}
}