use std::borrow::Cow;
use std::ffi::OsStr;
use std::io::{BufWriter, Read};
use std::path::Path;
use anyhow::Context;
use bencode::bencode_serialize_to_writer;
use buffers::ByteBufOwned;
use bytes::Bytes;
use librqbit_core::torrent_metainfo::{TorrentMetaV1File, TorrentMetaV1Info, TorrentMetaV1Owned};
use librqbit_core::Id20;
use sha1w::{ISha1, Sha1};
use crate::spawn_utils::BlockingSpawner;
#[derive(Debug, Clone, Default)]
pub struct CreateTorrentOptions<'a> {
pub name: Option<&'a str>,
pub piece_length: Option<u32>,
}
fn walk_dir_find_paths(dir: &Path, out: &mut Vec<Cow<'_, Path>>) -> anyhow::Result<()> {
out.extend(
walkdir::WalkDir::new(dir)
.into_iter()
.filter_map(|e| e.ok())
.filter(|e| e.file_type().is_file())
.map(|e| e.path().to_owned().into()),
);
Ok(())
}
fn compute_info_hash(t: &TorrentMetaV1Info<ByteBufOwned>) -> anyhow::Result<Id20> {
struct W {
hash: sha1w::Sha1,
}
impl std::io::Write for W {
fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
self.hash.update(buf);
Ok(buf.len())
}
fn flush(&mut self) -> std::io::Result<()> {
Ok(())
}
}
let mut writer = BufWriter::new(W { hash: Sha1::new() });
bencode_serialize_to_writer(t, &mut writer)?;
let hash = writer
.into_inner()
.map_err(|_| anyhow::anyhow!("into_inner errored"))?
.hash;
Ok(Id20::new(hash.finish()))
}
fn choose_piece_length(_input_files: &[Cow<'_, Path>]) -> u32 {
2 * 1024 * 1024
}
fn osstr_to_bytes(o: &OsStr) -> Vec<u8> {
o.to_str().unwrap().to_owned().into_bytes()
}
async fn create_torrent_raw<'a>(
path: &'a Path,
options: CreateTorrentOptions<'a>,
) -> anyhow::Result<TorrentMetaV1Info<ByteBufOwned>> {
path.try_exists()
.with_context(|| format!("path {:?} doesn't exist", path))?;
let basename = path
.file_name()
.ok_or_else(|| anyhow::anyhow!("cannot determine basename of {:?}", path))?;
let is_dir = path.is_dir();
let single_file_mode = !is_dir;
let name: ByteBufOwned = match options.name {
Some(name) => name.as_bytes().into(),
None => osstr_to_bytes(basename).into(),
};
let mut input_files: Vec<Cow<'a, Path>> = Default::default();
if is_dir {
walk_dir_find_paths(path, &mut input_files)
.with_context(|| format!("error walking {:?}", path))?;
} else {
input_files.push(Cow::Borrowed(path));
}
let piece_length = options
.piece_length
.unwrap_or_else(|| choose_piece_length(&input_files));
const READ_SIZE: u32 = 8192; let mut read_buf = vec![0; READ_SIZE as usize];
let mut length = 0;
let mut remaining_piece_length = piece_length;
let mut piece_checksum = sha1w::Sha1::new();
let mut piece_hashes = Vec::<u8>::new();
let mut output_files: Vec<TorrentMetaV1File<ByteBufOwned>> = Vec::new();
let spawner = BlockingSpawner::default();
'outer: for file in input_files {
let filename = &*file;
length = 0;
let mut fd = std::io::BufReader::new(
std::fs::File::open(&file).with_context(|| format!("error opening {:?}", filename))?,
);
loop {
let max_bytes_to_read = remaining_piece_length.min(READ_SIZE) as usize;
let size = spawner
.spawn_block_in_place(|| fd.read(&mut read_buf[..max_bytes_to_read]))
.with_context(|| format!("error reading {:?}", filename))?;
if size == 0 {
let filename = filename
.strip_prefix(path)
.context("internal error, can't strip prefix")?;
let path = filename
.components()
.map(|c| osstr_to_bytes(c.as_os_str()).into())
.collect();
output_files.push(TorrentMetaV1File {
length,
path,
attr: None,
sha1: None,
symlink_path: None,
});
continue 'outer;
}
length += size as u64;
piece_checksum.update(&read_buf[..size]);
remaining_piece_length -= TryInto::<u32>::try_into(size)?;
if remaining_piece_length == 0 {
remaining_piece_length = piece_length;
piece_hashes.extend_from_slice(&piece_checksum.finish());
piece_checksum = sha1w::Sha1::new();
}
}
}
if remaining_piece_length > 0 && length > 0 {
piece_hashes.extend_from_slice(&piece_checksum.finish());
}
Ok(TorrentMetaV1Info {
name: Some(name),
pieces: piece_hashes.into(),
piece_length,
length: if single_file_mode { Some(length) } else { None },
md5sum: None,
files: if single_file_mode {
None
} else {
Some(output_files)
},
attr: None,
sha1: None,
symlink_path: None,
private: false,
})
}
#[derive(Debug)]
pub struct CreateTorrentResult {
meta: TorrentMetaV1Owned,
}
impl CreateTorrentResult {
pub fn as_info(&self) -> &TorrentMetaV1Owned {
&self.meta
}
pub fn info_hash(&self) -> Id20 {
self.meta.info_hash
}
pub fn as_bytes(&self) -> anyhow::Result<Bytes> {
let mut b = Vec::new();
bencode_serialize_to_writer(&self.meta, &mut b).context("error serializing torrent")?;
Ok(b.into())
}
}
pub async fn create_torrent<'a>(
path: &'a Path,
options: CreateTorrentOptions<'a>,
) -> anyhow::Result<CreateTorrentResult> {
let info = create_torrent_raw(path, options).await?;
let info_hash = compute_info_hash(&info).context("error computing info hash")?;
Ok(CreateTorrentResult {
meta: TorrentMetaV1Owned {
announce: None,
announce_list: Vec::new(),
info,
comment: None,
created_by: None,
encoding: Some(b"utf-8"[..].into()),
publisher: None,
publisher_url: None,
creation_date: None,
info_hash,
},
})
}
#[cfg(test)]
mod tests {
use buffers::{ByteBuf, ByteBufOwned};
use librqbit_core::torrent_metainfo::torrent_from_bytes;
use crate::create_torrent;
#[tokio::test]
async fn test_create_torrent() {
use crate::tests::test_util;
let dir = test_util::create_default_random_dir_with_torrents(
3,
1000 * 1000,
Some("rqbit_test_create_torrent"),
);
let torrent = create_torrent(dir.path(), Default::default())
.await
.unwrap();
let bytes = torrent.as_bytes().unwrap();
let deserialized = torrent_from_bytes::<ByteBuf>(&bytes).unwrap();
assert_eq!(torrent.info_hash(), deserialized.info_hash);
let deserialized = torrent_from_bytes::<ByteBufOwned>(&bytes).unwrap();
assert_eq!(torrent.info_hash(), deserialized.info_hash);
}
}