Session persistence now saving full torrent contents

This commit is contained in:
Igor Katson 2023-11-30 00:48:57 +00:00
parent 7d02d79ff5
commit 52883769e1
No known key found for this signature in database
GPG key ID: B4EC22B66D61A3F5
5 changed files with 94 additions and 13 deletions

2
Cargo.lock generated
View file

@ -1005,6 +1005,7 @@ dependencies = [
"anyhow", "anyhow",
"axum", "axum",
"backoff", "backoff",
"base64",
"bincode", "bincode",
"bitvec", "bitvec",
"byteorder", "byteorder",
@ -1076,6 +1077,7 @@ dependencies = [
"librqbit-clone-to-owned", "librqbit-clone-to-owned",
"parking_lot", "parking_lot",
"serde", "serde",
"serde_json",
"tokio", "tokio",
"tracing", "tracing",
"url", "url",

View file

@ -61,6 +61,7 @@ url = "2"
hex = "0.4" hex = "0.4"
backoff = "0.4.0" backoff = "0.4.0"
dashmap = "5.5.3" dashmap = "5.5.3"
base64 = "0.21.5"
[dev-dependencies] [dev-dependencies]
futures = {version = "0.3"} futures = {version = "0.3"}

View file

@ -10,6 +10,7 @@ use std::{
}; };
use anyhow::{bail, Context}; use anyhow::{bail, Context};
use bencode::{bencode_serialize_to_writer, BencodeDeserializer};
use buffers::ByteString; use buffers::ByteString;
use dht::{Dht, DhtBuilder, Id20, PersistentDht, PersistentDhtConfig}; use dht::{Dht, DhtBuilder, Id20, PersistentDht, PersistentDhtConfig};
use librqbit_core::{ use librqbit_core::{
@ -19,7 +20,7 @@ use librqbit_core::{
}; };
use parking_lot::RwLock; use parking_lot::RwLock;
use reqwest::Url; use reqwest::Url;
use serde::{Deserialize, Serialize}; use serde::{Deserialize, Deserializer, Serialize, Serializer};
use tokio_stream::StreamExt; use tokio_stream::StreamExt;
use tracing::{debug, error, error_span, info, warn}; use tracing::{debug, error, error_span, info, warn};
@ -50,7 +51,7 @@ impl SessionDatabase {
fn serialize(&self) -> SerializedSessionDatabase { fn serialize(&self) -> SerializedSessionDatabase {
SerializedSessionDatabase { SerializedSessionDatabase {
torrents: self torrents_v2: self
.torrents .torrents
.values() .values()
.map(|torrent| SerializedTorrent { .map(|torrent| SerializedTorrent {
@ -61,6 +62,7 @@ impl SessionDatabase {
.map(|u| u.to_string()) .map(|u| u.to_string())
.collect(), .collect(),
info_hash: torrent.info_hash().as_string(), info_hash: torrent.info_hash().as_string(),
info: torrent.info().info.clone(),
only_files: torrent.only_files.clone(), only_files: torrent.only_files.clone(),
is_paused: torrent.with_state(|s| matches!(s, ManagedTorrentState::Paused(_))), is_paused: torrent.with_state(|s| matches!(s, ManagedTorrentState::Paused(_))),
output_folder: torrent.info().out_dir.clone(), output_folder: torrent.info().out_dir.clone(),
@ -73,15 +75,46 @@ impl SessionDatabase {
#[derive(Serialize, Deserialize)] #[derive(Serialize, Deserialize)]
struct SerializedTorrent { struct SerializedTorrent {
info_hash: String, info_hash: String,
#[serde(
serialize_with = "serialize_torrent",
deserialize_with = "deserialize_torrent"
)]
info: TorrentMetaV1Info<ByteString>,
trackers: HashSet<String>, trackers: HashSet<String>,
output_folder: PathBuf, output_folder: PathBuf,
only_files: Option<Vec<usize>>, only_files: Option<Vec<usize>>,
is_paused: bool, is_paused: bool,
} }
fn serialize_torrent<S>(t: &TorrentMetaV1Info<ByteString>, serializer: S) -> Result<S::Ok, S::Error>
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<TorrentMetaV1Info<ByteString>, 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::<ByteString>::deserialize(&mut BencodeDeserializer::new_from_buf(&b))
.map_err(D::Error::custom)
}
#[derive(Serialize, Deserialize)] #[derive(Serialize, Deserialize)]
struct SerializedSessionDatabase { struct SerializedSessionDatabase {
torrents: Vec<SerializedTorrent>, torrents_v2: Vec<SerializedTorrent>,
} }
pub struct Session { pub struct Session {
@ -171,6 +204,7 @@ pub fn read_local_file_including_stdin(filename: &str) -> anyhow::Result<Vec<u8>
pub enum AddTorrent<'a> { pub enum AddTorrent<'a> {
Url(Cow<'a, str>), Url(Cow<'a, str>),
TorrentFileBytes(Cow<'a, [u8]>), TorrentFileBytes(Cow<'a, [u8]>),
TorrentInfo(Box<TorrentMetaV1Owned>),
} }
impl<'a> AddTorrent<'a> { impl<'a> AddTorrent<'a> {
@ -201,6 +235,7 @@ impl<'a> AddTorrent<'a> {
match self { match self {
Self::Url(s) => s.into_owned().into_bytes(), Self::Url(s) => s.into_owned().into_bytes(),
Self::TorrentFileBytes(b) => b.into_owned(), Self::TorrentFileBytes(b) => b.into_owned(),
Self::TorrentInfo(_) => unimplemented!(),
} }
} }
} }
@ -309,18 +344,33 @@ impl Session {
let db: SerializedSessionDatabase = let db: SerializedSessionDatabase =
serde_json::from_reader(&mut rdr).context("error deserializing session database")?; serde_json::from_reader(&mut rdr).context("error deserializing session database")?;
let mut futures = Vec::new(); let mut futures = Vec::new();
for storrent in db.torrents.into_iter() { for storrent in db.torrents_v2.into_iter() {
let magnet = Magnet { let trackers: Vec<ByteString> = storrent
info_hash: Id20::from_str(&storrent.info_hash) .trackers
.context("error deserializing info_hash")?, .into_iter()
trackers: storrent.trackers.into_iter().collect(), .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({ futures.push({
let session = self.clone(); let session = self.clone();
async move { async move {
session session
.add_torrent( .add_torrent(
AddTorrent::Url(Cow::Owned(magnet.to_string())), AddTorrent::TorrentInfo(Box::new(info)),
Some(AddTorrentOptions { Some(AddTorrentOptions {
paused: storrent.is_paused, paused: storrent.is_paused,
output_folder: Some( output_folder: Some(
@ -442,6 +492,7 @@ impl Session {
AddTorrent::TorrentFileBytes(bytes) => { AddTorrent::TorrentFileBytes(bytes) => {
torrent_from_bytes(&bytes).context("error decoding torrent")? torrent_from_bytes(&bytes).context("error decoding torrent")?
} }
AddTorrent::TorrentInfo(t) => *t,
}; };
let dht_rx = match self.dht.as_ref() { let dht_rx = match self.dht.as_ref() {

View file

@ -18,7 +18,7 @@ sha1-rust = ["bencode/sha1-rust"]
[dependencies] [dependencies]
tracing = "0.1.40" tracing = "0.1.40"
tokio = "1" tokio = {version = "1", features = ["rt-multi-thread"]}
hex = "0.4" hex = "0.4"
anyhow = "1" anyhow = "1"
url = "2" 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"} 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"} clone_to_owned = {path="../clone_to_owned", package="librqbit-clone-to-owned", version = "2.2.1"}
itertools = "0.12" itertools = "0.12"
[dev-dependencies]
serde_json = "1"

View file

@ -5,7 +5,7 @@ use bencode::BencodeDeserializer;
use buffers::{ByteBuf, ByteString}; use buffers::{ByteBuf, ByteString};
use clone_to_owned::CloneToOwned; use clone_to_owned::CloneToOwned;
use itertools::Either; use itertools::Either;
use serde::Deserialize; use serde::{Deserialize, Serialize};
use crate::id20::Id20; use crate::id20::Id20;
@ -51,18 +51,23 @@ impl<BufType> TorrentMetaV1<BufType> {
} }
} }
#[derive(Deserialize, Debug, Clone)] #[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq)]
pub struct TorrentMetaV1Info<BufType> { pub struct TorrentMetaV1Info<BufType> {
#[serde(skip_serializing_if = "Option::is_none")]
pub name: Option<BufType>, pub name: Option<BufType>,
pub pieces: BufType, pub pieces: BufType,
#[serde(rename = "piece length")] #[serde(rename = "piece length")]
pub piece_length: u32, pub piece_length: u32,
// Single-file mode // Single-file mode
#[serde(skip_serializing_if = "Option::is_none")]
pub length: Option<u64>, pub length: Option<u64>,
#[serde(skip_serializing_if = "Option::is_none")]
pub md5sum: Option<BufType>, pub md5sum: Option<BufType>,
// Multi-file mode // Multi-file mode
#[serde(skip_serializing_if = "Option::is_none")]
pub files: Option<Vec<TorrentMetaV1File<BufType>>>, pub files: Option<Vec<TorrentMetaV1File<BufType>>>,
} }
@ -180,7 +185,7 @@ impl<BufType: AsRef<[u8]>> TorrentMetaV1Info<BufType> {
} }
} }
#[derive(Deserialize, Debug, Clone)] #[derive(Deserialize, Serialize, Debug, Clone, PartialEq, Eq)]
pub struct TorrentMetaV1File<BufType> { pub struct TorrentMetaV1File<BufType> {
pub length: u64, pub length: u64,
pub path: Vec<BufType>, pub path: Vec<BufType>,
@ -299,4 +304,23 @@ mod tests {
"64a980abe6e448226bb930ba061592e44c3781a1" "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<ByteBuf> = torrent_from_bytes(&buf).unwrap().info;
let mut writer = Vec::new();
bencode::bencode_serialize_to_writer(&torrent, &mut writer).unwrap();
let deserialized = TorrentMetaV1Info::<ByteBuf>::deserialize(
&mut BencodeDeserializer::new_from_buf(&writer),
)
.unwrap();
assert_eq!(torrent, deserialized);
}
} }