319 lines
9.4 KiB
Rust
319 lines
9.4 KiB
Rust
use crate::constants::CHUNK_SIZE;
|
|
|
|
const fn is_power_of_two(x: u64) -> bool {
|
|
(x != 0) && ((x & (x - 1)) == 0)
|
|
}
|
|
|
|
pub const fn ceil_div_u64(a: u64, b: u64) -> u64 {
|
|
(a + b - 1) / b
|
|
}
|
|
|
|
pub const fn last_element_size_u64(total: u64, chunk_size: u64) -> u64 {
|
|
let rem = total % chunk_size;
|
|
if rem == 0 {
|
|
return chunk_size;
|
|
}
|
|
rem
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub struct PieceInfo {
|
|
pub piece_index: ValidPieceIndex,
|
|
pub len: u32,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
|
pub struct ChunkInfo {
|
|
pub piece_index: ValidPieceIndex,
|
|
pub chunk_index: u32,
|
|
pub absolute_index: u32,
|
|
pub size: u32,
|
|
pub offset: u32,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Copy)]
|
|
pub struct Lengths {
|
|
chunk_length: u32,
|
|
total_length: u64,
|
|
piece_length: u32,
|
|
last_piece_id: u32,
|
|
last_piece_length: u32,
|
|
chunks_per_piece: u32,
|
|
}
|
|
|
|
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
|
|
pub struct ValidPieceIndex(u32);
|
|
impl std::fmt::Display for ValidPieceIndex {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
write!(f, "{}", self.0)
|
|
}
|
|
}
|
|
impl std::fmt::Debug for ValidPieceIndex {
|
|
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
|
write!(f, "{:?}", self.0)
|
|
}
|
|
}
|
|
|
|
impl ValidPieceIndex {
|
|
pub fn get(&self) -> u32 {
|
|
self.0
|
|
}
|
|
}
|
|
|
|
impl Lengths {
|
|
pub fn new(
|
|
total_length: u64,
|
|
piece_length: u32,
|
|
chunk_length: Option<u32>,
|
|
) -> anyhow::Result<Self> {
|
|
let chunk_length = chunk_length.unwrap_or(CHUNK_SIZE);
|
|
// I guess this is not needed? Don't recall why I put this check here.
|
|
//
|
|
// if !(is_power_of_two(piece_length as u64)) {
|
|
// anyhow::bail!("piece length {} is not a power of 2", piece_length);
|
|
// }
|
|
if !(is_power_of_two(chunk_length as u64)) {
|
|
anyhow::bail!("chunk length {} is not a power of 2", chunk_length);
|
|
}
|
|
if chunk_length >= piece_length {
|
|
anyhow::bail!(
|
|
"chunk length {} should be smaller than pice length {}",
|
|
chunk_length,
|
|
piece_length
|
|
);
|
|
}
|
|
if total_length == 0 {
|
|
anyhow::bail!("torrent with 0 length")
|
|
}
|
|
let total_pieces = ceil_div_u64(total_length, piece_length as u64) as u32;
|
|
Ok(Self {
|
|
chunk_length,
|
|
piece_length,
|
|
total_length,
|
|
chunks_per_piece: ceil_div_u64(piece_length as u64, chunk_length as u64) as u32,
|
|
last_piece_id: total_pieces - 1,
|
|
last_piece_length: last_element_size_u64(total_length, piece_length as u64) as u32,
|
|
})
|
|
}
|
|
pub const fn piece_bitfield_bytes(&self) -> usize {
|
|
ceil_div_u64(self.total_pieces() as u64, 8) as usize
|
|
}
|
|
pub const fn chunk_bitfield_bytes(&self) -> usize {
|
|
ceil_div_u64(self.total_chunks() as u64, 8) as usize
|
|
}
|
|
pub const fn total_length(&self) -> u64 {
|
|
self.total_length
|
|
}
|
|
pub const fn validate_piece_index(&self, index: u32) -> Option<ValidPieceIndex> {
|
|
if index > self.last_piece_id {
|
|
return None;
|
|
}
|
|
Some(ValidPieceIndex(index))
|
|
}
|
|
pub const fn default_piece_length(&self) -> u32 {
|
|
self.piece_length
|
|
}
|
|
pub const fn default_chunk_length(&self) -> u32 {
|
|
self.chunk_length
|
|
}
|
|
pub const fn default_chunks_per_piece(&self) -> u32 {
|
|
self.chunks_per_piece
|
|
}
|
|
pub const fn total_chunks(&self) -> u32 {
|
|
ceil_div_u64(self.total_length, self.chunk_length as u64) as u32
|
|
}
|
|
pub const fn last_piece_id(&self) -> ValidPieceIndex {
|
|
ValidPieceIndex(self.last_piece_id)
|
|
}
|
|
pub const fn total_pieces(&self) -> u32 {
|
|
self.last_piece_id + 1
|
|
}
|
|
pub const fn piece_length(&self, index: ValidPieceIndex) -> u32 {
|
|
if index.0 == self.last_piece_id {
|
|
return self.last_piece_length;
|
|
}
|
|
self.piece_length
|
|
}
|
|
pub const fn chunk_absolute_offset(&self, chunk_info: &ChunkInfo) -> u64 {
|
|
self.piece_offset(chunk_info.piece_index) + chunk_info.offset as u64
|
|
}
|
|
pub const fn piece_offset(&self, index: ValidPieceIndex) -> u64 {
|
|
index.0 as u64 * self.piece_length as u64
|
|
}
|
|
|
|
pub fn iter_piece_infos(&self) -> impl Iterator<Item = PieceInfo> {
|
|
let last_id = self.last_piece_id;
|
|
let last_len = self.last_piece_length;
|
|
let pl = self.piece_length;
|
|
(0..self.total_pieces()).map(move |idx| PieceInfo {
|
|
piece_index: ValidPieceIndex(idx),
|
|
len: if idx == last_id { last_len } else { pl },
|
|
})
|
|
}
|
|
|
|
pub fn iter_chunk_infos(&self, index: ValidPieceIndex) -> impl Iterator<Item = ChunkInfo> {
|
|
let mut remaining = self.piece_length(index);
|
|
let chunk_size = self.chunk_length;
|
|
let absolute_offset = index.0 * self.chunks_per_piece;
|
|
(0u32..).scan(0, move |offset, idx| {
|
|
if remaining == 0 {
|
|
return None;
|
|
}
|
|
let s = std::cmp::min(remaining, chunk_size);
|
|
let result = ChunkInfo {
|
|
piece_index: index,
|
|
chunk_index: idx,
|
|
absolute_index: absolute_offset + idx,
|
|
size: s,
|
|
offset: *offset,
|
|
};
|
|
*offset += s;
|
|
remaining -= s;
|
|
Some(result)
|
|
})
|
|
}
|
|
|
|
pub fn chunk_info_from_received_data(
|
|
&self,
|
|
piece_index: ValidPieceIndex,
|
|
begin: u32,
|
|
chunk_size: u32,
|
|
) -> Option<ChunkInfo> {
|
|
let index = begin / self.chunk_length;
|
|
let expected_chunk_size = self.chunk_size(piece_index, index)?;
|
|
let offset = self.chunk_offset_in_piece(piece_index, index)?;
|
|
if offset != begin {
|
|
return None;
|
|
}
|
|
if expected_chunk_size != chunk_size {
|
|
return None;
|
|
}
|
|
let absolute_index = self.chunks_per_piece * piece_index.get() + index;
|
|
Some(ChunkInfo {
|
|
piece_index,
|
|
chunk_index: index,
|
|
size: chunk_size,
|
|
offset,
|
|
absolute_index,
|
|
})
|
|
}
|
|
|
|
pub fn chunk_info_from_received_piece(
|
|
&self,
|
|
index: u32,
|
|
begin: u32,
|
|
block_len: u32,
|
|
) -> Option<ChunkInfo> {
|
|
self.chunk_info_from_received_data(self.validate_piece_index(index)?, begin, block_len)
|
|
}
|
|
pub const fn chunk_range(&self, index: ValidPieceIndex) -> std::ops::Range<usize> {
|
|
let start = index.0 * self.chunks_per_piece;
|
|
let end = start + self.chunks_per_piece(index);
|
|
start as usize..end as usize
|
|
}
|
|
pub const fn chunks_per_piece(&self, index: ValidPieceIndex) -> u32 {
|
|
if index.0 == self.last_piece_id {
|
|
return (self.last_piece_length + self.chunk_length - 1) / self.chunk_length;
|
|
}
|
|
self.chunks_per_piece
|
|
}
|
|
pub const fn chunk_offset_in_piece(
|
|
&self,
|
|
piece_index: ValidPieceIndex,
|
|
chunk_index: u32,
|
|
) -> Option<u32> {
|
|
if chunk_index >= self.chunks_per_piece(piece_index) {
|
|
return None;
|
|
}
|
|
Some(chunk_index * self.chunk_length)
|
|
}
|
|
pub fn chunk_size(&self, piece_index: ValidPieceIndex, chunk_index: u32) -> Option<u32> {
|
|
let chunks_per_piece = self.chunks_per_piece(piece_index);
|
|
let pl = self.piece_length(piece_index);
|
|
if chunk_index >= chunks_per_piece {
|
|
return None;
|
|
}
|
|
let offset = chunk_index * self.chunk_length;
|
|
Some(std::cmp::min(self.chunk_length, pl - offset))
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
fn make_lengths() -> Lengths {
|
|
Lengths::new(1174243328, 262144, None).unwrap()
|
|
}
|
|
|
|
#[test]
|
|
fn test_total_pieces() {
|
|
let l = make_lengths();
|
|
assert_eq!(l.total_pieces(), 4480);
|
|
}
|
|
|
|
#[test]
|
|
fn test_total_pieces_2() {
|
|
let l = Lengths::new(4148166656, 2097152, None).unwrap();
|
|
assert_eq!(l.total_pieces(), 1978);
|
|
}
|
|
|
|
#[test]
|
|
fn test_piece_length() {
|
|
let l = make_lengths();
|
|
let p = l.validate_piece_index(4479).unwrap();
|
|
|
|
assert_eq!(l.piece_length(l.validate_piece_index(0).unwrap()), 262144);
|
|
assert_eq!(l.piece_length(p), 100352);
|
|
}
|
|
|
|
#[test]
|
|
fn test_chunks_in_piece() {
|
|
let l = make_lengths();
|
|
let p = l.validate_piece_index(4479).unwrap();
|
|
|
|
assert_eq!(l.chunks_per_piece(l.validate_piece_index(0).unwrap()), 16);
|
|
assert_eq!(l.chunks_per_piece(p), 7);
|
|
}
|
|
|
|
#[test]
|
|
fn test_chunk_size() {
|
|
let l = make_lengths();
|
|
let p = l.validate_piece_index(4479).unwrap();
|
|
|
|
assert_eq!(l.chunk_size(p, 0), Some(16384));
|
|
assert_eq!(l.chunk_size(p, 6), Some(2048));
|
|
}
|
|
|
|
#[test]
|
|
fn test_chunk_infos() {
|
|
let l = make_lengths();
|
|
let p = l.validate_piece_index(4479).unwrap();
|
|
|
|
let mut it = l.iter_chunk_infos(p);
|
|
let first = it.next().unwrap();
|
|
let last = it.last().unwrap();
|
|
|
|
assert_eq!(
|
|
first,
|
|
ChunkInfo {
|
|
piece_index: p,
|
|
chunk_index: 0,
|
|
absolute_index: 71664,
|
|
size: 16384,
|
|
offset: 0,
|
|
}
|
|
);
|
|
|
|
assert_eq!(
|
|
last,
|
|
ChunkInfo {
|
|
piece_index: p,
|
|
chunk_index: 6,
|
|
absolute_index: 71670,
|
|
size: 2048,
|
|
offset: 98304,
|
|
}
|
|
);
|
|
}
|
|
}
|