feat: Exportable engine state (#166)

Author: artrixdotdev

Cargo.lock

@@ -27,17 +27,18 @@ num_enum = "^0.7"
 thiserror = "^2.0"
 async-trait = "^0.1"
 librqbit-utp = "^0.6"
-bitvec = "^1"
+bitvec = { version = "^1.0", features = ["serde", "alloc", "atomic"] }
 futures = "^0.3"
 bencode = "^0.1"
 serde_with = "^3.14"
 tokio-retry2 = { version = "^0.6", features = ["tracing"] }
 bytes = "^1"
 kameo = "^0.18"
-dashmap = "^6.1"
+dashmap = { version = "^6", features = ["serde"] }
 bon = "^3.7"
 tokio-util = "^0.7"
 
 [dev-dependencies]
 tracing-test = "0.2.5"
 tracing-subscriber = { workspace = true }
+serde_json = "^1"

crates/libtortillas/Cargo.toml

@@ -1,10 +1,11 @@
+use futures::future::join_all;
 use kameo::{
    Actor, Reply, mailbox,
    prelude::{ActorRef, Context, Message},
 };
 use tracing::{error, warn};
 
-use super::EngineActor;
+use super::{EngineActor, EngineExport};
 use crate::{
    actor_request_response,
    errors::EngineError,
@@ -14,7 +15,9 @@ use crate::{
       messages::PeerMessages,
       stream::{PeerRecv, PeerStream},
    },
-   torrent::{TorrentActor, TorrentActorArgs, TorrentMessage, TorrentState},
+   torrent::{
+      TorrentActor, TorrentActorArgs, TorrentMessage, TorrentRequest, TorrentResponse, TorrentState,
+   },
 };
 
 pub(crate) enum EngineMessage {
@@ -32,6 +35,9 @@ actor_request_response!(
    /// Creates a new [Torrent] actor.
    Torrent(Box<MetaInfo>) Torrent(ActorRef<TorrentActor>),
 
+   /// Exports the current state of the engine.
+   Export Export(EngineExport),
+
 );
 
 impl Message<EngineMessage> for EngineActor {
@@ -72,7 +78,6 @@ impl Message<EngineMessage> for EngineActor {
       };
    }
 }
-
 impl Message<EngineRequest> for EngineActor {
    type Reply = Result<EngineResponse, EngineError>;
 
@@ -108,7 +113,6 @@ impl Message<EngineRequest> for EngineActor {
                   sufficient_peers: self.sufficient_peers,
                   base_path: self.default_base_path.clone(),
                },
-               // if the size is 0, we use an unbounded mailbox
                match self.mailbox_size {
                   0 => {
                      warn!(
@@ -126,6 +130,30 @@ impl Message<EngineRequest> for EngineActor {
             self.torrents.insert(info_hash, torrent_ref.clone());
             Ok(EngineResponse::Torrent(torrent_ref))
          }
+         EngineRequest::Export => {
+            // Concurrently collect all torrent exports
+            let futures = self
+               .torrents
+               .iter()
+               .map(|torrent| {
+                  let torrent = torrent.clone();
+                  async move {
+                     match torrent
+                        .ask(TorrentRequest::Export)
+                        .await
+                        .expect("Failed to get torrent export")
+                     {
+                        TorrentResponse::Export(export) => *export,
+                        _ => unreachable!(),
+                     }
+                  }
+               })
+               .collect::<Vec<_>>();
+
+            let torrents = join_all(futures).await;
+
+            Ok(EngineResponse::Export(EngineExport { torrents }))
+         }
       }
    }
 }

crates/libtortillas/src/engine/messages.rs

@@ -43,13 +43,14 @@ pub(crate) use actor::*;
 use bon;
 use kameo::{Actor, actor::ActorRef};
 pub(crate) use messages::*;
+use serde::{Deserialize, Serialize};
 use tracing::error;
 
 use crate::{
    errors::EngineError,
    metainfo::{MetaInfo, TorrentFile},
    peer::PeerId,
-   torrent::{PieceStorageStrategy, Torrent},
+   torrent::{PieceStorageStrategy, Torrent, TorrentExport},
 };
 
 /// The main entry point for managing torrents.
@@ -264,10 +265,29 @@ impl Engine {
          .await
          .expect("Failed to start all torrents");
    }
+
+   /// Exports the current state of the engine.
+   /// See [`Torrent::export`] for more information.
+   pub async fn export(&self) -> EngineExport {
+      match self
+         .actor()
+         .ask(EngineRequest::Export)
+         .await
+         .expect("Failed to get torrents")
+      {
+         EngineResponse::Export(export) => export,
+         _ => unreachable!(),
+      }
+   }
 }
 
 impl Default for Engine {
    fn default() -> Self {
       Self::builder().build()
    }
 }
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct EngineExport {
+   pub torrents: Vec<TorrentExport>,
+}

crates/libtortillas/src/engine/mod.rs

@@ -201,13 +201,13 @@ impl<const N: usize> Serialize for Hash<N> {
    where
       S: Serializer,
    {
-      serializer.serialize_bytes(&self.0)
+      serializer.serialize_bytes(self.as_bytes())
    }
 }
 
 struct HashVisitor<const N: usize>;
 
-impl<const N: usize> Visitor<'_> for HashVisitor<N> {
+impl<'de, const N: usize> Visitor<'de> for HashVisitor<N> {
    type Value = Hash<N>;
 
    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
@@ -228,6 +228,36 @@ impl<const N: usize> Visitor<'_> for HashVisitor<N> {
          )))
       }
    }
+   fn visit_byte_buf<E>(self, v: Vec<u8>) -> Result<Self::Value, E>
+   where
+      E: de::Error,
+   {
+      if v.len() == N {
+         Ok(Hash(v.try_into().unwrap()))
+      } else {
+         Err(E::custom(format!(
+            "Expected length is {}, found {}",
+            N,
+            v.len(),
+         )))
+      }
+   }
+
+   fn visit_seq<A>(self, seq: A) -> Result<Self::Value, A::Error>
+   where
+      A: de::SeqAccess<'de>,
+   {
+      let mut seq = seq;
+      let mut bytes = Vec::with_capacity(N);
+      while let Some(byte) = seq.next_element()? {
+         bytes.push(byte);
+      }
+      if bytes.len() == N {
+         Ok(Hash(bytes.try_into().unwrap()))
+      } else {
+         Err(de::Error::invalid_length(bytes.len(), &self))
+      }
+   }
 }
 
 /// A collection of `Hash<N>` values that provides efficient storage and
@@ -420,7 +450,7 @@ impl<const N: usize> HashVecVisitor<N> {
    }
 }
 
-impl<const N: usize> Visitor<'_> for HashVecVisitor<N> {
+impl<'de, const N: usize> Visitor<'de> for HashVecVisitor<N> {
    type Value = HashVec<N>;
 
    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
@@ -449,6 +479,30 @@ impl<const N: usize> Visitor<'_> for HashVecVisitor<N> {
       }
       Ok(HashVec(vec))
    }
+   fn visit_seq<A>(self, seq: A) -> Result<Self::Value, A::Error>
+   where
+      A: de::SeqAccess<'de>,
+   {
+      let mut seq = seq;
+      let mut hashes = HashVec::<N>::new();
+      let mut bytes = Vec::new();
+      while let Some(byte) = seq.next_element()? {
+         bytes.push(byte);
+      }
+
+      if bytes.len().is_multiple_of(N) {
+         for chunk in bytes.chunks(N) {
+            hashes.push(Hash(
+               chunk
+                  .try_into()
+                  .unwrap_or_else(|_| panic!("guaranteed to be length {N}")),
+            ));
+         }
+         Ok(hashes)
+      } else {
+         Err(de::Error::invalid_length(bytes.len(), &self))
+      }
+   }
 }
 
 impl<const N: usize> From<Vec<Hash<N>>> for HashVec<N> {

crates/libtortillas/src/hashes.rs

@@ -13,7 +13,7 @@ use crate::{
    tracker::Tracker,
 };
 
-#[derive(Debug, Clone, Deserialize)]
+#[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct TorrentFile {
    /// The primary announce URI for the torrent.
    pub announce: Tracker,

crates/libtortillas/src/metainfo/file.rs

@@ -1,5 +1,5 @@
 use anyhow::Result;
-use serde::Deserialize;
+use serde::{Deserialize, Deserializer, Serialize, de};
 use serde_querystring;
 
 use crate::{
@@ -12,48 +12,71 @@ use crate::{
 #[derive(Debug, Clone, Deserialize)]
 pub struct MagnetUri {
    /// use `Self::info_hash` to get the info hash as a `Hash` struct.
-   #[serde(rename(deserialize = "xt"))]
+   #[serde(rename = "xt")]
    info_hash: String,
 
-   #[serde(rename(deserialize = "dn"))]
+   #[serde(rename = "dn")]
    pub name: String,
 
-   #[serde(rename(deserialize = "xl"))]
+   #[serde(rename = "xl")]
    pub length: Option<u32>,
 
-   #[serde(rename(deserialize = "tr"))]
+   #[serde(rename = "tr")]
    pub announce_list: Option<Vec<Tracker>>,
 
-   #[serde(rename(deserialize = "ws"), default)]
+   #[serde(rename = "ws", default)]
    pub web_seed: Vec<String>,
 
-   #[serde(rename(deserialize = "as"))]
+   #[serde(rename = "as")]
    pub source: Option<String>,
 
-   #[serde(rename(deserialize = "xs"))]
+   #[serde(rename = "xs")]
    pub exact_source: Option<String>,
 
-   #[serde(rename(deserialize = "kt"))]
+   #[serde(rename = "kt")]
    pub keywords: Option<Vec<String>>,
 
-   #[serde(rename(deserialize = "mt"))]
+   #[serde(rename = "mt")]
    pub manifest_topic: Option<String>,
 
-   #[serde(rename(deserialize = "so"))]
+   #[serde(rename = "so")]
    pub select_only: Option<Vec<String>>,
 
-   #[serde(rename(deserialize = "x.pe"))]
+   #[serde(rename = "x.pe")]
    pub peer: Option<String>,
+
+   #[serde(skip)]
+   uri: String,
+}
+
+impl Serialize for MagnetUri {
+   fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+   where
+      S: serde::Serializer,
+   {
+      serializer.serialize_str(&self.uri)
+   }
+}
+
+impl TryFrom<String> for MagnetUri {
+   type Error = anyhow::Error;
+   fn try_from(uri: String) -> Result<Self, Self::Error> {
+      let qs = uri
+         .split('?')
+         .nth(1)
+         .ok_or_else(|| anyhow::anyhow!("Invalid magnet URI"))?;
+
+      let mut magnet: Self =
+         serde_querystring::from_str(qs, serde_querystring::ParseMode::Duplicate)?;
+
+      magnet.uri = uri;
+      Ok(magnet)
+   }
 }
 
 impl MagnetUri {
    pub fn parse(uri: String) -> Result<MetaInfo> {
-      let qs = uri.split('?').next_back().unwrap(); // Turns magnet:?xt=... into xt=...
-      // Parse the modified query string
-      Ok(MetaInfo::MagnetUri(serde_querystring::from_str(
-         qs,
-         serde_querystring::ParseMode::Duplicate,
-      )?))
+      Ok(MetaInfo::MagnetUri(MagnetUri::try_from(uri)?))
    }
 
    pub fn announce_list(&self) -> Vec<Tracker> {
@@ -70,6 +93,14 @@ impl MagnetUri {
       Hash::from_hex(hex_part)
          .map_err(|e| anyhow::anyhow!("Failed to parse info_hash from hex: {}", e))
    }
+
+   pub fn deserialize<'de, D>(deserializer: D) -> Result<Self, D::Error>
+   where
+      D: Deserializer<'de>,
+   {
+      let uri = String::deserialize(deserializer)?;
+      Self::try_from(uri).map_err(de::Error::custom)
+   }
 }
 
 #[cfg(test)]