Move everything to workspaces

This commit is contained in:
Igor Katson 2021-07-03 19:10:59 +01:00
parent 75547d3000
commit ad867e8e3c
42 changed files with 338 additions and 168 deletions

View file

@ -0,0 +1 @@
pub const CHUNK_SIZE: u32 = 16384;

View file

@ -0,0 +1,7 @@
pub type InfoHash = [u8; 20];
pub fn decode_info_hash(hash_str: &str) -> anyhow::Result<InfoHash> {
let mut hash_arr = [0u8; 20];
hex::decode_to_slice(hash_str, &mut hash_arr)?;
Ok(hash_arr)
}

View file

@ -0,0 +1,304 @@
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);
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
);
}
Ok(Self {
chunk_length,
piece_length,
total_length,
chunks_per_piece: piece_length / chunk_length,
last_piece_id: ((total_length + 1) / piece_length as u64) as u32,
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 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_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,
}
);
}
}

View file

@ -0,0 +1,7 @@
pub mod constants;
pub mod info_hash;
pub mod lengths;
pub mod magnet;
pub mod peer_id;
pub mod speed_estimator;
pub mod torrent_metainfo;

View file

@ -0,0 +1,48 @@
use crate::info_hash::{decode_info_hash, InfoHash};
use anyhow::Context;
pub struct Magnet {
pub info_hash: InfoHash,
pub trackers: Vec<String>,
}
impl Magnet {
pub fn parse(url: &str) -> anyhow::Result<Magnet> {
let url = url::Url::parse(url).context("magnet link must be a valid URL")?;
if url.scheme() != "magnet" {
anyhow::bail!("expected scheme magnet");
}
let mut info_hash: Option<InfoHash> = None;
let mut trackers = Vec::<String>::new();
for (key, value) in url.query_pairs() {
match key.as_ref() {
"xt" => match value.as_ref().strip_prefix("urn:btih:") {
Some(infohash) => {
info_hash.replace(decode_info_hash(infohash)?);
}
None => anyhow::bail!("expected xt to start with urn:btih:"),
},
"tr" => trackers.push(value.into()),
_ => {}
}
}
match info_hash {
Some(info_hash) => Ok(Magnet {
info_hash,
trackers,
}),
None => {
anyhow::bail!("did not find infohash")
}
}
}
}
#[cfg(test)]
mod tests {
#[test]
fn test_parse_magnet_as_url() {
let magnet = "magnet:?xt=urn:btih:a621779b5e3d486e127c3efbca9b6f8d135f52e5&dn=rutor.info_%D0%92%D0%BE%D0%B9%D0%BD%D0%B0+%D0%B1%D1%83%D0%B4%D1%83%D1%89%D0%B5%D0%B3%D0%BE+%2F+The+Tomorrow+War+%282021%29+WEB-DLRip+%D0%BE%D1%82+MegaPeer+%7C+P+%7C+NewComers&tr=udp://opentor.org:2710&tr=udp://opentor.org:2710&tr=http://retracker.local/announce";
dbg!(url::Url::parse(magnet).unwrap());
}
}

View file

@ -0,0 +1,56 @@
#[derive(Debug)]
pub enum AzureusStyleKind {
Deluge,
LibTorrent,
Transmission,
Other([char; 2]),
}
#[derive(Debug)]
pub struct AzureusStyle {
pub kind: AzureusStyleKind,
pub version: [char; 4],
}
impl AzureusStyleKind {
pub const fn from_bytes(b1: u8, b2: u8) -> Self {
match &[b1, b2] {
b"DE" => AzureusStyleKind::Deluge,
b"lt" | b"LT" => AzureusStyleKind::LibTorrent,
b"TR" => AzureusStyleKind::Transmission,
_ => AzureusStyleKind::Other([b1 as char, b2 as char]),
}
}
}
fn try_decode_azureus_style(p: &[u8; 20]) -> Option<AzureusStyle> {
if !(p[0] == b'-' && p[7] == b'-') {
return None;
}
let mut version = ['0'; 4];
for (i, c) in (&p[3..7]).iter().copied().enumerate() {
version[i] = c as char;
}
let kind = AzureusStyleKind::from_bytes(p[1], p[2]);
Some(AzureusStyle { kind, version })
}
#[derive(Debug)]
pub enum PeerId {
AzureusStyle(AzureusStyle),
}
pub fn try_decode_peer_id(p: [u8; 20]) -> Option<PeerId> {
Some(PeerId::AzureusStyle(try_decode_azureus_style(&p)?))
}
pub fn generate_peer_id() -> [u8; 20] {
let mut peer_id = [0u8; 20];
let u = uuid::Uuid::new_v4();
(&mut peer_id[4..20]).copy_from_slice(&u.as_bytes()[..]);
(&mut peer_id[..8]).copy_from_slice(b"-rQ0001-");
peer_id
}

View file

@ -0,0 +1,77 @@
use std::{
collections::VecDeque,
sync::atomic::{AtomicU64, Ordering},
time::{Duration, Instant},
};
use parking_lot::Mutex;
struct ProgressSnapshot {
downloaded_bytes: u64,
instant: Instant,
}
pub struct SpeedEstimator {
latest_per_second_snapshots: Mutex<VecDeque<ProgressSnapshot>>,
download_bytes_per_second: AtomicU64,
time_remaining_millis: AtomicU64,
}
impl SpeedEstimator {
pub fn new(window_seconds: usize) -> Self {
assert!(window_seconds > 1);
Self {
latest_per_second_snapshots: Mutex::new(VecDeque::with_capacity(window_seconds)),
download_bytes_per_second: Default::default(),
time_remaining_millis: Default::default(),
}
}
pub fn time_remaining(&self) -> Option<Duration> {
let tr = self.time_remaining_millis.load(Ordering::Relaxed);
if tr == 0 {
return None;
}
Some(Duration::from_millis(tr))
}
pub fn download_bps(&self) -> u64 {
self.download_bytes_per_second.load(Ordering::Relaxed)
}
pub fn download_mbps(&self) -> f64 {
self.download_bps() as f64 / 1024f64 / 1024f64
}
pub fn add_snapshot(&self, downloaded_bytes: u64, remaining_bytes: u64, instant: Instant) {
let mut g = self.latest_per_second_snapshots.lock();
if g.len() < g.capacity() {
g.push_back(ProgressSnapshot {
downloaded_bytes,
instant,
});
return;
}
let first = g.pop_front().unwrap();
let downloaded_bytes_diff = downloaded_bytes - first.downloaded_bytes;
let elapsed = instant - first.instant;
let bps = downloaded_bytes_diff as f64 / elapsed.as_secs_f64();
let time_remaining_millis_rounded: u64 = if downloaded_bytes_diff > 0 {
let time_remaining_secs = remaining_bytes as f64 / bps;
(time_remaining_secs * 1000f64) as u64
} else {
0
};
self.time_remaining_millis
.store(time_remaining_millis_rounded, Ordering::Relaxed);
self.download_bytes_per_second
.store(bps as u64, Ordering::Relaxed);
g.push_back(ProgressSnapshot {
downloaded_bytes,
instant,
});
}
}

View file

@ -0,0 +1,283 @@
use std::{fmt::Write, ops::Deref, path::PathBuf};
use bencode::BencodeDeserializer;
use buffers::{ByteBuf, ByteString};
use clone_to_owned::CloneToOwned;
use serde::Deserialize;
pub type TorrentMetaV1Borrowed<'a> = TorrentMetaV1<ByteBuf<'a>>;
pub type TorrentMetaV1Owned = TorrentMetaV1<ByteString>;
pub fn torrent_from_bytes<'de, ByteBuf: Clone + Deserialize<'de>>(
buf: &'de [u8],
) -> anyhow::Result<TorrentMetaV1<ByteBuf>> {
let mut de = BencodeDeserializer::new_from_buf(buf);
de.is_torrent_info = true;
let mut t = TorrentMetaV1::deserialize(&mut de)?;
t.info_hash = de.torrent_info_digest.unwrap();
Ok(t)
}
#[derive(Deserialize, Debug, Clone)]
pub struct TorrentMetaV1<BufType: Clone> {
pub announce: BufType,
#[serde(rename = "announce-list")]
pub announce_list: Vec<Vec<BufType>>,
pub info: TorrentMetaV1Info<BufType>,
pub comment: Option<BufType>,
#[serde(rename = "created by")]
pub created_by: Option<BufType>,
pub encoding: Option<BufType>,
pub publisher: Option<BufType>,
#[serde(rename = "publisher-url")]
pub publisher_url: Option<BufType>,
#[serde(rename = "creation date")]
pub creation_date: Option<usize>,
#[serde(skip)]
pub info_hash: [u8; 20],
}
impl<BufType: Clone> TorrentMetaV1<BufType> {
pub fn iter_announce(&self) -> impl Iterator<Item = &BufType> {
std::iter::once(&self.announce).chain(self.announce_list.iter().flatten())
}
}
#[derive(Deserialize, Debug, Clone)]
pub struct TorrentMetaV1Info<BufType: Clone> {
pub name: Option<BufType>,
pub pieces: BufType,
#[serde(rename = "piece length")]
pub piece_length: u32,
// Single-file mode
pub length: Option<u64>,
pub md5sum: Option<BufType>,
// Multi-file mode
pub files: Option<Vec<TorrentMetaV1File<BufType>>>,
}
pub enum FileIteratorName<'a, ByteBuf> {
Single(Option<&'a ByteBuf>),
Tree(&'a [ByteBuf]),
}
impl<'a, ByteBuf> std::fmt::Debug for FileIteratorName<'a, ByteBuf>
where
ByteBuf: AsRef<[u8]>,
{
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
for (idx, item) in self.iter_components().enumerate() {
if idx > 0 {
f.write_char(std::path::MAIN_SEPARATOR)?;
}
match item {
Some(bit) => {
f.write_str(std::str::from_utf8(bit.as_ref()).unwrap_or("<INVALID UTF-8>"))?;
}
None => f.write_str("output")?,
}
}
Ok(())
}
}
impl<'a, ByteBuf> FileIteratorName<'a, ByteBuf> {
pub fn to_pathbuf(&self) -> anyhow::Result<PathBuf>
where
ByteBuf: AsRef<[u8]>,
{
let mut buf = PathBuf::new();
for part in self.iter_components() {
if let Some(part) = part {
buf.push(std::str::from_utf8(part.as_ref())?)
} else {
buf.push("output");
break;
}
}
Ok(buf)
}
pub fn iter_components(&self) -> impl Iterator<Item = Option<&'a ByteBuf>> {
let single_it = std::iter::once(match self {
FileIteratorName::Single(n) => Some(*n),
FileIteratorName::Tree(_) => None,
});
let multi_it = match self {
FileIteratorName::Single(_) => &[],
FileIteratorName::Tree(t) => *t,
}
.iter()
.map(|p| Some(Some(p)));
single_it.chain(multi_it).flatten()
}
}
impl<BufType: Clone + Deref<Target = [u8]>> TorrentMetaV1Info<BufType> {
pub fn get_hash(&self, piece: u32) -> Option<&[u8]> {
let start = piece as usize * 20;
let end = start + 20;
let expected_hash = self.pieces.deref().get(start..end)?;
Some(expected_hash)
}
pub fn compare_hash(&self, piece: u32, hash: [u8; 20]) -> Option<bool> {
let start = piece as usize * 20;
let end = start + 20;
let expected_hash = self.pieces.deref().get(start..end)?;
Some(expected_hash == hash)
}
pub fn iter_filenames_and_lengths(
&self,
) -> impl Iterator<Item = (FileIteratorName<'_, BufType>, u64)> {
let single_it = std::iter::once(match (self.name.as_ref(), self.length) {
(Some(n), Some(l)) => Some((FileIteratorName::Single(Some(n)), l)),
_ => None,
});
let multi_it = self
.files
.as_deref()
.unwrap_or_default()
.iter()
.map(|f| Some((FileIteratorName::Tree(&f.path), f.length)));
single_it.chain(multi_it).flatten()
}
pub fn iter_file_lengths(&self) -> impl Iterator<Item = u64> + '_ {
std::iter::once(self.length)
.chain(
self.files
.as_deref()
.unwrap_or_default()
.iter()
.map(|f| Some(f.length)),
)
.flatten()
}
}
#[derive(Deserialize, Debug, Clone)]
pub struct TorrentMetaV1File<BufType: Clone> {
pub length: u64,
pub path: Vec<BufType>,
}
impl<BufType> TorrentMetaV1File<BufType>
where
BufType: Clone + AsRef<[u8]>,
{
pub fn full_path(&self, parent: &mut PathBuf) -> anyhow::Result<()> {
for p in self.path.iter() {
let bit = std::str::from_utf8(p.as_ref())?;
parent.push(bit);
}
Ok(())
}
}
impl<ByteBuf> CloneToOwned for TorrentMetaV1File<ByteBuf>
where
ByteBuf: CloneToOwned + Clone,
<ByteBuf as CloneToOwned>::Target: Clone,
{
type Target = TorrentMetaV1File<<ByteBuf as CloneToOwned>::Target>;
fn clone_to_owned(&self) -> Self::Target {
TorrentMetaV1File {
length: self.length,
path: self.path.clone_to_owned(),
}
}
}
impl<ByteBuf> CloneToOwned for TorrentMetaV1Info<ByteBuf>
where
ByteBuf: CloneToOwned + Clone,
<ByteBuf as CloneToOwned>::Target: Clone,
{
type Target = TorrentMetaV1Info<<ByteBuf as CloneToOwned>::Target>;
fn clone_to_owned(&self) -> Self::Target {
TorrentMetaV1Info {
name: self.name.clone_to_owned(),
pieces: self.pieces.clone_to_owned(),
piece_length: self.piece_length,
length: self.length,
md5sum: self.md5sum.clone_to_owned(),
files: self.files.clone_to_owned(),
}
}
}
impl<ByteBuf> CloneToOwned for TorrentMetaV1<ByteBuf>
where
ByteBuf: CloneToOwned + Clone,
<ByteBuf as CloneToOwned>::Target: Clone,
{
type Target = TorrentMetaV1<<ByteBuf as CloneToOwned>::Target>;
fn clone_to_owned(&self) -> Self::Target {
TorrentMetaV1 {
announce: self.announce.clone_to_owned(),
announce_list: self.announce_list.clone_to_owned(),
info: self.info.clone_to_owned(),
comment: self.comment.clone_to_owned(),
created_by: self.created_by.clone_to_owned(),
encoding: self.encoding.clone_to_owned(),
publisher: self.publisher.clone_to_owned(),
publisher_url: self.publisher_url.clone_to_owned(),
creation_date: self.creation_date,
info_hash: self.info_hash,
}
}
}
#[cfg(test)]
mod tests {
use std::io::Read;
use super::*;
#[test]
fn test_deserialize_torrent_owned() {
let mut buf = Vec::new();
let filename = "resources/ubuntu-21.04-desktop-amd64.iso.torrent";
std::fs::File::open(filename)
.unwrap()
.read_to_end(&mut buf)
.unwrap();
let torrent: TorrentMetaV1Owned = torrent_from_bytes(&buf).unwrap();
dbg!(torrent);
}
#[test]
fn test_deserialize_torrent_borrowed() {
let mut buf = Vec::new();
let filename = "resources/ubuntu-21.04-desktop-amd64.iso.torrent";
std::fs::File::open(filename)
.unwrap()
.read_to_end(&mut buf)
.unwrap();
let torrent: TorrentMetaV1Borrowed = torrent_from_bytes(&buf).unwrap();
dbg!(torrent);
}
#[test]
fn test_deserialize_torrent_with_info_hash() {
let mut buf = Vec::new();
let filename = "resources/ubuntu-21.04-desktop-amd64.iso.torrent";
std::fs::File::open(filename)
.unwrap()
.read_to_end(&mut buf)
.unwrap();
let torrent: TorrentMetaV1Borrowed = torrent_from_bytes(&buf).unwrap();
assert_eq!(
torrent.info_hash,
*b"\x64\xa9\x80\xab\xe6\xe4\x48\x22\x6b\xb9\x30\xba\x06\x15\x92\xe4\x4c\x37\x81\xa1"
);
}
}