leyoda/rqbit
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Nothing, small cleanups

Browse source
This commit is contained in:
Igor Katson 2021-07-16 13:58:10 +01:00
parent ca11caf902
commit bd0df49591
5 changed files with 87 additions and 53 deletions
Download patch file Download diff file Expand all files Collapse all files

8
crates/librqbit/src/file_ops.rs
View file

@ -98,7 +98,7 @@ impl<'a, Sha1Impl: ISha1> FileOps<'a, Sha1Impl> {
let mut file_iterator = self
.files
.iter()
.zip(self.torrent.iter_filenames_and_lengths())
.zip(self.torrent.iter_filenames_and_lengths()?)
.enumerate()
.map(|(idx, (fd, (name, len)))| {
let full_file_required = if let Some(only_files) = only_files {
@ -228,7 +228,7 @@ impl<'a, Sha1Impl: ISha1> FileOps<'a, Sha1Impl> {
let mut piece_remaining_bytes = piece_length as usize;
for (file_idx, (name, file_len)) in self.torrent.iter_filenames_and_lengths().enumerate() {
for (file_idx, (name, file_len)) in self.torrent.iter_filenames_and_lengths()?.enumerate() {
if absolute_offset > file_len {
absolute_offset -= file_len;
continue;
@ -297,7 +297,7 @@ impl<'a, Sha1Impl: ISha1> FileOps<'a, Sha1Impl> {
let mut absolute_offset = self.lengths.chunk_absolute_offset(&chunk_info);
let mut buf = result_buf;
for (file_idx, file_len) in self.torrent.iter_file_lengths().enumerate() {
for (file_idx, file_len) in self.torrent.iter_file_lengths()?.enumerate() {
if absolute_offset > file_len {
absolute_offset -= file_len;
continue;
@ -351,7 +351,7 @@ impl<'a, Sha1Impl: ISha1> FileOps<'a, Sha1Impl> {
let mut buf = data.block.as_ref();
let mut absolute_offset = self.lengths.chunk_absolute_offset(&chunk_info);
for (file_idx, (name, file_len)) in self.torrent.iter_filenames_and_lengths().enumerate() {
for (file_idx, (name, file_len)) in self.torrent.iter_filenames_and_lengths()?.enumerate() {
if absolute_offset > file_len {
absolute_offset -= file_len;
continue;

1
crates/librqbit/src/http_api.rs
View file

@ -106,6 +106,7 @@ impl ApiInternal {
.torrent_state()
.info()
.iter_filenames_and_lengths()
.unwrap()
.map(|(filename_it, length)| {
let name = filename_it.to_string().ok();
TorrentDetailsResponseFile { name, length }

6
crates/librqbit/src/torrent_manager.rs
View file

@ -148,7 +148,7 @@ struct TorrentManager {
fn make_lengths<ByteBuf: AsRef<[u8]>>(
torrent: &TorrentMetaV1Info<ByteBuf>,
) -> anyhow::Result<Lengths> {
let total_length = torrent.iter_file_lengths().sum();
let total_length = torrent.iter_file_lengths()?.sum();
Lengths::new(total_length, torrent.piece_length, None)
}
@ -163,9 +163,9 @@ impl TorrentManager {
let options = options.unwrap_or_default();
let files = {
let mut files =
Vec::<Arc<Mutex<File>>>::with_capacity(info.iter_file_lengths().count());
Vec::<Arc<Mutex<File>>>::with_capacity(info.iter_file_lengths()?.count());
for (path_bits, _) in info.iter_filenames_and_lengths() {
for (path_bits, _) in info.iter_filenames_and_lengths()? {
let mut full_path = out.as_ref().to_owned();
for bit in path_bits.iter_components() {
full_path.push(

104
crates/librqbit_core/src/torrent_metainfo.rs
View file

@ -1,5 +1,6 @@
use std::{fmt::Write, path::PathBuf};
use std::path::PathBuf;
use anyhow::Context;
use bencode::BencodeDeserializer;
use buffers::{ByteBuf, ByteString};
use clone_to_owned::CloneToOwned;
@ -74,18 +75,10 @@ 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")?,
}
match self.to_string() {
Ok(s) => write!(f, "{:?}", s),
Err(e) => write!(f, "<{:?}>", e),
}
Ok(())
}
}
@ -94,19 +87,13 @@ impl<'a, ByteBuf> FileIteratorName<'a, ByteBuf> {
where
ByteBuf: AsRef<[u8]>,
{
let mut it = self.iter_components();
let mut buf = it
.next()
.and_then(|v| v)
.map(|v| std::str::from_utf8(v.as_ref()))
.ok_or_else(|| anyhow::anyhow!("empty filename"))??
.to_string();
for bit in it {
buf.push('/');
let bit = bit
.map(|v| std::str::from_utf8(v.as_ref()))
.ok_or_else(|| anyhow::anyhow!("empty filename"))??;
buf.push_str(bit);
let mut buf = String::new();
for (idx, bit) in self.iter_components().enumerate() {
let bit = bit?;
if idx > 0 {
buf.push(std::path::MAIN_SEPARATOR);
}
buf.push_str(bit)
}
Ok(buf)
}
@ -115,17 +102,16 @@ impl<'a, ByteBuf> FileIteratorName<'a, ByteBuf> {
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;
}
for bit in self.iter_components() {
let bit = bit?;
buf.push(bit)
}
Ok(buf)
}
pub fn iter_components(&self) -> impl Iterator<Item = Option<&'a ByteBuf>> {
pub fn iter_components(&self) -> impl Iterator<Item = anyhow::Result<&'a str>>
where
ByteBuf: AsRef<[u8]>,
{
let single_it = std::iter::once(match self {
FileIteratorName::Single(n) => Some(*n),
FileIteratorName::Tree(_) => None,
@ -137,7 +123,27 @@ impl<'a, ByteBuf> FileIteratorName<'a, ByteBuf> {
.iter()
.map(|p| Some(Some(p)));
single_it.chain(multi_it).flatten()
let it = single_it.chain(multi_it).flatten();
it.map(|part| {
let part = match part {
Some(part) => part,
None => return Ok("torrent-content"),
};
let bit = std::str::from_utf8(part.as_ref())
.context("cannot decode filename bit as UTF-8")?;
if bit.contains("..") {
anyhow::bail!("path traversal detected, \"..\" in filename bit {:?}", bit);
}
if bit.contains(std::path::MAIN_SEPARATOR) {
anyhow::bail!(
"suspicios separator {:?} in filename bit {:?}",
std::path::MAIN_SEPARATOR,
bit
);
}
Ok(bit)
})
}
}
@ -154,11 +160,27 @@ impl<BufType: AsRef<[u8]>> TorrentMetaV1Info<BufType> {
let expected_hash = self.pieces.as_ref().get(start..end)?;
Some(expected_hash == hash)
}
fn is_single_file(&self) -> anyhow::Result<bool> {
match (self.length, self.files.as_ref()) {
(None, Some(files)) => {
if files.is_empty() {
anyhow::bail!("expected multi-file torrent to have at least one file")
}
Ok(false)
}
(Some(_), None) => Ok(true),
_ => anyhow::bail!("torrent can't be both in single and multi-file mode"),
}
}
pub fn iter_filenames_and_lengths(
&self,
) -> impl Iterator<Item = (FileIteratorName<'_, BufType>, u64)> {
) -> anyhow::Result<impl Iterator<Item = (FileIteratorName<'_, BufType>, u64)>> {
self.is_single_file()?;
let single_it = std::iter::once(match (self.name.as_ref(), self.length) {
(Some(n), Some(l)) => Some((FileIteratorName::Single(Some(n)), l)),
(n, Some(l)) => Some((FileIteratorName::Single(n), l)),
_ => None,
});
let multi_it = self
@ -167,10 +189,11 @@ impl<BufType: AsRef<[u8]>> TorrentMetaV1Info<BufType> {
.unwrap_or_default()
.iter()
.map(|f| Some((FileIteratorName::Tree(&f.path), f.length)));
single_it.chain(multi_it).flatten()
Ok(single_it.chain(multi_it).flatten())
}
pub fn iter_file_lengths(&self) -> impl Iterator<Item = u64> + '_ {
std::iter::once(self.length)
pub fn iter_file_lengths(&self) -> anyhow::Result<impl Iterator<Item = u64> + '_> {
self.is_single_file()?;
let it = std::iter::once(self.length)
.chain(
self.files
.as_deref()
@ -178,7 +201,8 @@ impl<BufType: AsRef<[u8]>> TorrentMetaV1Info<BufType> {
.iter()
.map(|f| Some(f.length)),
)
.flatten()
.flatten();
Ok(it)
}
}

21
crates/rqbit/src/main.rs
View file

@ -3,7 +3,7 @@ use std::{fs::File, io::Read, net::SocketAddr, str::FromStr, time::Duration};
use anyhow::Context;
use clap::Clap;
use dht::{Dht, Id20};
use futures::{Stream, StreamExt};
use futures::StreamExt;
use librqbit::{
dht_utils::{read_metainfo_from_peer_receiver, ReadMetainfoResult},
generate_peer_id,
@ -121,7 +121,7 @@ fn compute_only_files<ByteBuf: AsRef<[u8]>>(
) -> anyhow::Result<Vec<usize>> {
let filename_re = regex::Regex::new(&filename_re).context("filename regex is incorrect")?;
let mut only_files = Vec::new();
for (idx, (filename, _)) in torrent.iter_filenames_and_lengths().enumerate() {
for (idx, (filename, _)) in torrent.iter_filenames_and_lengths()?.enumerate() {
let full_path = filename
.to_pathbuf()
.with_context(|| format!("filename of file {} is not valid utf8", idx))?;
@ -303,15 +303,24 @@ async fn main_torrent_info(
spawner: BlockingSpawner,
) -> anyhow::Result<()> {
info!("Torrent info: {:#?}", &info);
if opts.list {
return Ok(());
}
let only_files = if let Some(filename_re) = opts.only_files_matching_regex {
Some(compute_only_files(&info, &filename_re)?)
let only_files = compute_only_files(&info, &filename_re)?;
for (idx, (filename, _)) in info.iter_filenames_and_lengths()?.enumerate() {
if !only_files.contains(&idx) {
continue;
}
info!("Will download {:?}", filename);
}
Some(only_files)
} else {
None
};
if opts.list {
info!("--list was passed, nothing to do, exiting.");
return Ok(());
}
let http_api_listen_addr = opts.http_api_listen_addr;
let mut builder = TorrentManagerBuilder::new(info, info_hash, opts.output_folder);