From fe7a1e09ba867f5eacf44caeb62532ae87684cd7 Mon Sep 17 00:00:00 2001
From: Igor Katson <igor.katson@gmail.com>
Date: Mon, 12 Aug 2024 23:43:23 +0100
Subject: [PATCH] Store torrent bytes in memory

---
 crates/librqbit/src/session.rs           | 42 ++++++++++++++++++++++--
 crates/librqbit/src/torrent_state/mod.rs |  5 +++
 2 files changed, 45 insertions(+), 2 deletions(-)

diff --git a/crates/librqbit/src/session.rs b/crates/librqbit/src/session.rs
index e201682..2ae66f5 100644
--- a/crates/librqbit/src/session.rs
+++ b/crates/librqbit/src/session.rs
@@ -119,6 +119,7 @@ impl SessionDatabase {
                                 .map(|u| u.to_string())
                                 .collect(),
                             info_hash: torrent.info_hash().as_string(),
+                            torrent_bytes: torrent.info.torrent_bytes.clone(),
                             info: torrent.info().info.clone(),
                             only_files: torrent.only_files().clone(),
                             is_paused: torrent
@@ -140,6 +141,12 @@ struct SerializedTorrent {
         deserialize_with = "deserialize_torrent"
     )]
     info: TorrentMetaV1Info<ByteBufOwned>,
+    #[serde(
+        serialize_with = "serialize_torrent_bytes",
+        deserialize_with = "deserialize_torrent_bytes",
+        default = "empty_bytes"
+    )]
+    torrent_bytes: ByteBufOwned,
     trackers: HashSet<String>,
     output_folder: PathBuf,
     only_files: Option<Vec<usize>>,
@@ -175,6 +182,32 @@ where
         .map_err(D::Error::custom)
 }
 
+fn serialize_torrent_bytes<S>(t: &ByteBufOwned, serializer: S) -> Result<S::Ok, S::Error>
+where
+    S: Serializer,
+{
+    use base64::{engine::general_purpose, Engine as _};
+    let s = general_purpose::STANDARD_NO_PAD.encode(&t.0);
+    s.serialize(serializer)
+}
+
+fn deserialize_torrent_bytes<'de, D>(deserializer: D) -> Result<ByteBufOwned, 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)?;
+    Ok(b.into())
+}
+
+fn empty_bytes() -> ByteBufOwned {
+    ByteBufOwned(Vec::new().into_boxed_slice())
+}
+
 #[derive(Serialize, Deserialize)]
 struct SerializedSessionDatabase {
     torrents: HashMap<usize, SerializedTorrent>,
@@ -1128,8 +1161,13 @@ impl Session {
             }));
         }
 
-        let mut builder =
-            ManagedTorrentBuilder::new(info, info_hash, output_folder, storage_factory);
+        let mut builder = ManagedTorrentBuilder::new(
+            info,
+            info_hash,
+            torrent_bytes,
+            output_folder,
+            storage_factory,
+        );
         builder
             .allow_overwrite(opts.overwrite)
             .spawner(self.spawner)
diff --git a/crates/librqbit/src/torrent_state/mod.rs b/crates/librqbit/src/torrent_state/mod.rs
index 6e544da..f653316 100644
--- a/crates/librqbit/src/torrent_state/mod.rs
+++ b/crates/librqbit/src/torrent_state/mod.rs
@@ -99,6 +99,7 @@ pub(crate) struct ManagedTorrentOptions {
 
 pub struct ManagedTorrentInfo {
     pub info: TorrentMetaV1Info<ByteBufOwned>,
+    pub torrent_bytes: ByteBufOwned,
     pub info_hash: Id20,
     pub(crate) spawner: BlockingSpawner,
     pub trackers: HashSet<String>,
@@ -501,6 +502,7 @@ pub(crate) struct ManagedTorrentBuilder {
     info: TorrentMetaV1Info<ByteBufOwned>,
     output_folder: PathBuf,
     info_hash: Id20,
+    torrent_bytes: ByteBufOwned,
     force_tracker_interval: Option<Duration>,
     peer_connect_timeout: Option<Duration>,
     peer_read_write_timeout: Option<Duration>,
@@ -518,12 +520,14 @@ impl ManagedTorrentBuilder {
     pub fn new(
         info: TorrentMetaV1Info<ByteBufOwned>,
         info_hash: Id20,
+        torrent_bytes: ByteBufOwned,
         output_folder: PathBuf,
         storage_factory: BoxStorageFactory,
     ) -> Self {
         Self {
             info,
             info_hash,
+            torrent_bytes,
             spawner: None,
             force_tracker_interval: None,
             peer_connect_timeout: None,
@@ -608,6 +612,7 @@ impl ManagedTorrentBuilder {
             span,
             file_infos,
             info: self.info,
+            torrent_bytes: self.torrent_bytes,
             info_hash: self.info_hash,
             trackers: self.trackers.into_iter().collect(),
             spawner: self.spawner.unwrap_or_default(),