From 52883769e1047ee0b037ad4cba25679e57c84778 Mon Sep 17 00:00:00 2001 From: Igor Katson Date: Thu, 30 Nov 2023 00:48:57 +0000 Subject: [PATCH] Session persistence now saving full torrent contents --- Cargo.lock | 2 + crates/librqbit/Cargo.toml | 1 + crates/librqbit/src/session.rs | 69 +++++++++++++++++--- crates/librqbit_core/Cargo.toml | 5 +- crates/librqbit_core/src/torrent_metainfo.rs | 30 ++++++++- 5 files changed, 94 insertions(+), 13 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8e9c891..3460a27 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1005,6 +1005,7 @@ dependencies = [ "anyhow", "axum", "backoff", + "base64", "bincode", "bitvec", "byteorder", @@ -1076,6 +1077,7 @@ dependencies = [ "librqbit-clone-to-owned", "parking_lot", "serde", + "serde_json", "tokio", "tracing", "url", diff --git a/crates/librqbit/Cargo.toml b/crates/librqbit/Cargo.toml index 9e18974..4093071 100644 --- a/crates/librqbit/Cargo.toml +++ b/crates/librqbit/Cargo.toml @@ -61,6 +61,7 @@ url = "2" hex = "0.4" backoff = "0.4.0" dashmap = "5.5.3" +base64 = "0.21.5" [dev-dependencies] futures = {version = "0.3"} diff --git a/crates/librqbit/src/session.rs b/crates/librqbit/src/session.rs index 47bc9cc..55c67ee 100644 --- a/crates/librqbit/src/session.rs +++ b/crates/librqbit/src/session.rs @@ -10,6 +10,7 @@ use std::{ }; use anyhow::{bail, Context}; +use bencode::{bencode_serialize_to_writer, BencodeDeserializer}; use buffers::ByteString; use dht::{Dht, DhtBuilder, Id20, PersistentDht, PersistentDhtConfig}; use librqbit_core::{ @@ -19,7 +20,7 @@ use librqbit_core::{ }; use parking_lot::RwLock; use reqwest::Url; -use serde::{Deserialize, Serialize}; +use serde::{Deserialize, Deserializer, Serialize, Serializer}; use tokio_stream::StreamExt; use tracing::{debug, error, error_span, info, warn}; @@ -50,7 +51,7 @@ impl SessionDatabase { fn serialize(&self) -> SerializedSessionDatabase { SerializedSessionDatabase { - torrents: self + torrents_v2: self .torrents .values() .map(|torrent| SerializedTorrent { @@ -61,6 +62,7 @@ impl SessionDatabase { .map(|u| u.to_string()) .collect(), info_hash: torrent.info_hash().as_string(), + info: torrent.info().info.clone(), only_files: torrent.only_files.clone(), is_paused: torrent.with_state(|s| matches!(s, ManagedTorrentState::Paused(_))), output_folder: torrent.info().out_dir.clone(), @@ -73,15 +75,46 @@ impl SessionDatabase { #[derive(Serialize, Deserialize)] struct SerializedTorrent { info_hash: String, + #[serde( + serialize_with = "serialize_torrent", + deserialize_with = "deserialize_torrent" + )] + info: TorrentMetaV1Info, trackers: HashSet, output_folder: PathBuf, only_files: Option>, is_paused: bool, } +fn serialize_torrent(t: &TorrentMetaV1Info, serializer: S) -> Result +where + S: Serializer, +{ + use base64::{engine::general_purpose, Engine as _}; + use serde::ser::Error; + let mut writer = Vec::new(); + bencode_serialize_to_writer(t, &mut writer).map_err(S::Error::custom)?; + let s = general_purpose::STANDARD_NO_PAD.encode(&writer); + s.serialize(serializer) +} + +fn deserialize_torrent<'de, D>(deserializer: D) -> Result, D::Error> +where + D: Deserializer<'de>, +{ + use base64::{engine::general_purpose, Engine as _}; + use serde::de::Error; + let s = String::deserialize(deserializer)?; + let b = general_purpose::STANDARD_NO_PAD + .decode(s) + .map_err(D::Error::custom)?; + TorrentMetaV1Info::::deserialize(&mut BencodeDeserializer::new_from_buf(&b)) + .map_err(D::Error::custom) +} + #[derive(Serialize, Deserialize)] struct SerializedSessionDatabase { - torrents: Vec, + torrents_v2: Vec, } pub struct Session { @@ -171,6 +204,7 @@ pub fn read_local_file_including_stdin(filename: &str) -> anyhow::Result pub enum AddTorrent<'a> { Url(Cow<'a, str>), TorrentFileBytes(Cow<'a, [u8]>), + TorrentInfo(Box), } impl<'a> AddTorrent<'a> { @@ -201,6 +235,7 @@ impl<'a> AddTorrent<'a> { match self { Self::Url(s) => s.into_owned().into_bytes(), Self::TorrentFileBytes(b) => b.into_owned(), + Self::TorrentInfo(_) => unimplemented!(), } } } @@ -309,18 +344,33 @@ impl Session { let db: SerializedSessionDatabase = serde_json::from_reader(&mut rdr).context("error deserializing session database")?; let mut futures = Vec::new(); - for storrent in db.torrents.into_iter() { - let magnet = Magnet { - info_hash: Id20::from_str(&storrent.info_hash) - .context("error deserializing info_hash")?, - trackers: storrent.trackers.into_iter().collect(), + for storrent in db.torrents_v2.into_iter() { + let trackers: Vec = storrent + .trackers + .into_iter() + .map(|t| ByteString(t.into_bytes())) + .collect(); + let info = TorrentMetaV1Owned { + announce: trackers + .get(0) + .cloned() + .unwrap_or_else(|| ByteString(b"http://retracker.local/announce".into())), + announce_list: vec![trackers], + info: storrent.info, + comment: None, + created_by: None, + encoding: None, + publisher: None, + publisher_url: None, + creation_date: None, + info_hash: Id20::from_str(&storrent.info_hash)?, }; futures.push({ let session = self.clone(); async move { session .add_torrent( - AddTorrent::Url(Cow::Owned(magnet.to_string())), + AddTorrent::TorrentInfo(Box::new(info)), Some(AddTorrentOptions { paused: storrent.is_paused, output_folder: Some( @@ -442,6 +492,7 @@ impl Session { AddTorrent::TorrentFileBytes(bytes) => { torrent_from_bytes(&bytes).context("error decoding torrent")? } + AddTorrent::TorrentInfo(t) => *t, }; let dht_rx = match self.dht.as_ref() { diff --git a/crates/librqbit_core/Cargo.toml b/crates/librqbit_core/Cargo.toml index 4d1df54..0ffc7ef 100644 --- a/crates/librqbit_core/Cargo.toml +++ b/crates/librqbit_core/Cargo.toml @@ -18,7 +18,7 @@ sha1-rust = ["bencode/sha1-rust"] [dependencies] tracing = "0.1.40" -tokio = "1" +tokio = {version = "1", features = ["rt-multi-thread"]} hex = "0.4" anyhow = "1" url = "2" @@ -29,3 +29,6 @@ buffers = {path="../buffers", package="librqbit-buffers", version = "2.2.1"} bencode = {path = "../bencode", default-features=false, package="librqbit-bencode", version="2.2.1"} clone_to_owned = {path="../clone_to_owned", package="librqbit-clone-to-owned", version = "2.2.1"} itertools = "0.12" + +[dev-dependencies] +serde_json = "1" \ No newline at end of file diff --git a/crates/librqbit_core/src/torrent_metainfo.rs b/crates/librqbit_core/src/torrent_metainfo.rs index 6d8d765..c51ff06 100644 --- a/crates/librqbit_core/src/torrent_metainfo.rs +++ b/crates/librqbit_core/src/torrent_metainfo.rs @@ -5,7 +5,7 @@ use bencode::BencodeDeserializer; use buffers::{ByteBuf, ByteString}; use clone_to_owned::CloneToOwned; use itertools::Either; -use serde::Deserialize; +use serde::{Deserialize, Serialize}; use crate::id20::Id20; @@ -51,18 +51,23 @@ impl TorrentMetaV1 { } } -#[derive(Deserialize, Debug, Clone)] +#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)] pub struct TorrentMetaV1Info { + #[serde(skip_serializing_if = "Option::is_none")] pub name: Option, pub pieces: BufType, #[serde(rename = "piece length")] pub piece_length: u32, // Single-file mode + #[serde(skip_serializing_if = "Option::is_none")] pub length: Option, + + #[serde(skip_serializing_if = "Option::is_none")] pub md5sum: Option, // Multi-file mode + #[serde(skip_serializing_if = "Option::is_none")] pub files: Option>>, } @@ -180,7 +185,7 @@ impl> TorrentMetaV1Info { } } -#[derive(Deserialize, Debug, Clone)] +#[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq)] pub struct TorrentMetaV1File { pub length: u64, pub path: Vec, @@ -299,4 +304,23 @@ mod tests { "64a980abe6e448226bb930ba061592e44c3781a1" ); } + + #[test] + fn test_serialize_then_deserialize_bencode() { + let mut buf = Vec::new(); + std::fs::File::open(TORRENT_FILENAME) + .unwrap() + .read_to_end(&mut buf) + .unwrap(); + + let torrent: TorrentMetaV1Info = torrent_from_bytes(&buf).unwrap().info; + let mut writer = Vec::new(); + bencode::bencode_serialize_to_writer(&torrent, &mut writer).unwrap(); + let deserialized = TorrentMetaV1Info::::deserialize( + &mut BencodeDeserializer::new_from_buf(&writer), + ) + .unwrap(); + + assert_eq!(torrent, deserialized); + } }