Merge pull request #28 from artrixdotdev/refactor/better-hashes

refactor: Move all info hashes to singular Hash struct

Author: artrixdotdev

lib/src/hashes.rs

@@ -1,3 +1,4 @@
 pub mod errors;
+pub mod hashes;
 pub mod parser;
 pub mod tracker;

lib/src/lib.rs

@@ -1,13 +1,14 @@
-use crate::{parser::MetaInfo, tracker::Tracker};
+use crate::{
+   hashes::{Hash, HashVec, InfoHash},
+   parser::MetaInfo,
+   tracker::Tracker,
+};
 
 use anyhow::Result;
-use serde::{
-   Deserialize, Serialize, Serializer,
-   de::{self, Visitor},
-};
+use serde::{Deserialize, Serialize};
 use serde_bencode as bencode;
 use sha1::{Digest, Sha1};
-use std::{fmt, path::PathBuf};
+use std::path::PathBuf;
 
 use tokio::fs;
 
@@ -44,7 +45,7 @@ pub struct Info {
    #[serde(rename = "piece length")]
    piece_length: u64,
    /// Binary string of concatenated 20-byte SHA-1 hash values
-   pieces: Hashes,
+   pieces: HashVec<20>,
    #[serde(flatten)]
    file: InfoKeys,
 
@@ -75,62 +76,11 @@ pub struct InfoFile {
 
 impl Info {
    /// Gets the file hash (xt, or exact topic) for the given Info struct
-   pub fn hash(&self) -> Result<String> {
+   pub fn hash(&self) -> Result<InfoHash> {
       let mut hasher = Sha1::new();
       hasher.update(serde_bencode::to_bytes(&self)?);
       let result = hasher.finalize();
-      Ok(hex::encode(result))
-   }
-}
-
-/// A custom type for serializing and deserializing a vector of 20-byte SHA-1 hashes.
-/// Credit to [Jon Gjengset](https://github.com/jonhoo/codecrafters-bittorrent-rust/)
-#[derive(Debug)]
-pub struct Hashes(pub Vec<[u8; 20]>);
-
-// Add this implementation
-impl<'de> Deserialize<'de> for Hashes {
-   fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
-   where
-      D: de::Deserializer<'de>,
-   {
-      deserializer.deserialize_bytes(HashesVisitor)
-   }
-}
-
-impl Serialize for Hashes {
-   fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
-   where
-      S: Serializer,
-   {
-      let single_slice = self.0.concat();
-
-      serializer.serialize_bytes(&single_slice)
-   }
-}
-
-struct HashesVisitor;
-
-impl Visitor<'_> for HashesVisitor {
-   type Value = Hashes;
-
-   fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
-      formatter.write_str("a byte string whose length is a multiple of 20")
-   }
-
-   fn visit_bytes<E>(self, v: &[u8]) -> Result<Self::Value, E>
-   where
-      E: de::Error,
-   {
-      if v.len() % 20 != 0 {
-         return Err(E::custom(format!("length is {}", v.len())));
-      }
-
-      Ok(Hashes(
-         v.chunks(20)
-            .map(|slice_20| slice_20.try_into().expect("guaranteed to be length 20"))
-            .collect(),
-      ))
+      Ok(Hash::from_hex(hex::encode(result))?)
    }
 }
 

lib/src/parser/file.rs

@@ -1,6 +1,10 @@
 use std::collections::HashMap;
 
-use crate::{parser::MetaInfo, tracker::Tracker};
+use crate::{
+   hashes::{Hash, InfoHash},
+   parser::MetaInfo,
+   tracker::Tracker,
+};
 
 use anyhow::Result;
 use serde::Deserialize;
@@ -9,8 +13,9 @@ use serde_qs;
 /// Magnet URI Spec: https://en.wikipedia.org/wiki/Magnet_URI_scheme or https://www.bittorrent.org/beps/bep_0053.html
 #[derive(Debug, Deserialize)]
 pub struct MagnetUri {
+   /// use `Self::info_hash` to get the info hash as a `Hash` struct.
    #[serde(rename(deserialize = "xt"))]
-   pub info_hash: String,
+   info_hash: String,
 
    #[serde(rename(deserialize = "dn"))]
    pub name: String,
@@ -82,6 +87,16 @@ impl MagnetUri {
       // Parse the modified query string
       Ok(MetaInfo::MagnetUri(serde_qs::from_str(&final_qs)?))
    }
+   pub fn info_hash(&self) -> Result<InfoHash, anyhow::Error> {
+      let hex_part = self
+         .info_hash
+         .split(":")
+         .last()
+         .ok_or_else(|| anyhow::anyhow!("Invalid info_hash format: no colon found"))?;
+
+      Hash::from_hex(hex_part)
+         .map_err(|e| anyhow::anyhow!("Failed to parse info_hash from hex: {}", e))
+   }
 }
 
 #[cfg(test)]

lib/src/parser/magnet.rs

@@ -5,6 +5,8 @@ mod magnet;
 pub use file::*;
 pub use magnet::*;
 
+use crate::hashes::InfoHash;
+
 /// Always utilize MetaInfo instead of directly using TorrentFile or MagnetUri
 #[derive(Debug, Deserialize)]
 pub enum MetaInfo {
@@ -17,12 +19,15 @@ impl MetaInfo {
    /// function will calculate and return the hash. If the enum is a [MagnetUri](MagnetUri), then this
    /// function will grab the existing hash and return it, as the MagnetUri spec already contains
    /// the hash.
-   pub fn info_hash(&self) -> String {
+   pub fn info_hash(&self) -> Result<InfoHash, anyhow::Error> {
       match &self {
-         MetaInfo::Torrent(torrent) => torrent.info.hash().unwrap(),
-         MetaInfo::MagnetUri(magnet_uri) => {
-            String::from(magnet_uri.info_hash.split(":").last().unwrap())
-         }
+         MetaInfo::Torrent(torrent) => torrent
+            .info
+            .hash()
+            .map_err(|e| anyhow::anyhow!("Failed to compute torrent info hash: {}", e)),
+         MetaInfo::MagnetUri(magnet_uri) => magnet_uri
+            .info_hash()
+            .map_err(|e| anyhow::anyhow!("Failed to extract magnet URI info hash: {}", e)),
       }
    }
 }
@@ -43,8 +48,10 @@ mod tests {
       let contents = tokio::fs::read_to_string(path).await.unwrap();
 
       let metainfo = MagnetUri::parse(contents).await.unwrap();
+
+      let info_hash = metainfo.info_hash().unwrap();
       assert_eq!(
-         metainfo.info_hash(),
+         info_hash.to_hex(),
          "dd8255ecdc7ca55fb0bbf81323d87062db1f6d1c"
       );
    }
@@ -56,7 +63,12 @@ mod tests {
          .unwrap()
          .join("tests/torrents/big-buck-bunny.torrent");
       let file = TorrentFile::parse(path).await.unwrap();
-      assert_eq!(file.info_hash(), "dd8255ecdc7ca55fb0bbf81323d87062db1f6d1c");
+
+      let info_hash = file.info_hash().unwrap();
+      assert_eq!(
+         info_hash.to_hex(),
+         "dd8255ecdc7ca55fb0bbf81323d87062db1f6d1c"
+      );
    }
 
    #[tokio::test]

lib/src/parser/mod.rs

@@ -1,5 +1,8 @@
 use super::{PeerAddr, TrackerTrait};
-use crate::errors::{HttpTrackerError, TrackerError};
+use crate::{
+   errors::{HttpTrackerError, TrackerError},
+   hashes::InfoHash,
+};
 use rand::distr::{Alphanumeric, SampleString};
 use serde::{
    Deserialize, Serialize,
@@ -55,13 +58,13 @@ impl TrackerRequest {
 pub struct HttpTracker {
    uri: String,
    peer_id: String,
-   info_hash: String,
+   info_hash: InfoHash,
    params: TrackerRequest,
 }
 
 impl HttpTracker {
    #[instrument(skip(info_hash), fields(uri = %uri))]
-   pub fn new(uri: String, info_hash: String) -> HttpTracker {
+   pub fn new(uri: String, info_hash: InfoHash) -> HttpTracker {
       let peer_id = Alphanumeric.sample_string(&mut rand::rng(), 20);
       debug!(peer_id = %peer_id, "Generated peer ID");
 
@@ -93,27 +96,14 @@ impl TrackerTrait for HttpTracker {
    async fn stream_peers(&mut self) -> anyhow::Result<Vec<PeerAddr>> {
       // Decode info_hash
       debug!("Decoding info hash");
-      let info_hash_part = self.info_hash.split("urn:btih:").last().ok_or_else(|| {
-         HttpTrackerError::InvalidInfoHash(format!("Invalid info_hash format: {}", self.info_hash))
-      })?;
-
-      let decoded_hash = hex::decode(info_hash_part).map_err(|e| {
-         error!(error = %e, "Failed to decode info_hash");
-         HttpTrackerError::InvalidInfoHash(format!("Failed to decode info_hash: {}", e))
-      })?;
-
-      let info_hash: [u8; 20] = decoded_hash.try_into().map_err(|_| {
-         error!("Info hash has incorrect length");
-         HttpTrackerError::InvalidInfoHash("Info hash must be 20 bytes".to_string())
-      })?;
 
       // Generate params + URL. Specifically using the compact format by adding "compact=1" to
       // params.
       debug!("Generating request parameters");
       let params = serde_qs::to_string(&self.params)
          .map_err(|e| HttpTrackerError::ParameterEncoding(e.to_string()))?;
 
-      let info_hash_encoded = urlencode(&info_hash);
+      let info_hash_encoded = urlencode(self.info_hash.as_bytes());
       trace!(encoded_hash = %info_hash_encoded, "URL-encoded info hash");
 
       let uri_params = format!(
@@ -238,10 +228,10 @@ mod tests {
       let metainfo = MagnetUri::parse(contents).await.unwrap();
       match metainfo {
          MetaInfo::MagnetUri(magnet) => {
+            let info_hash = magnet.info_hash();
             let announce_list = magnet.announce_list.unwrap();
             let announce_uri = announce_list[0].uri();
-            let info_hash = magnet.info_hash;
-            let mut http_tracker = HttpTracker::new(announce_uri, info_hash);
+            let mut http_tracker = HttpTracker::new(announce_uri, info_hash.unwrap());
 
             // Make request
             let res = HttpTracker::stream_peers(&mut http_tracker)

lib/src/tracker/http.rs

@@ -6,6 +6,8 @@ use serde::{
 };
 use std::{fmt, net::Ipv4Addr};
 use udp::UdpTracker;
+
+use crate::hashes::InfoHash;
 pub mod http;
 pub mod udp;
 // mod websocket;
@@ -37,21 +39,15 @@ pub enum Tracker {
 }
 
 impl Tracker {
-   pub async fn get_peers(&self, info_hash: String) -> Result<Vec<PeerAddr>> {
+   pub async fn get_peers(&self, info_hash: InfoHash) -> Result<Vec<PeerAddr>> {
       match self {
          Tracker::Http(uri) => {
             let mut tracker = HttpTracker::new(uri.clone(), info_hash);
 
             Ok(tracker.stream_peers().await.unwrap())
          }
          Tracker::Udp(uri) => {
-            let mut tracker = UdpTracker::new(
-               uri.clone(),
-               None,
-               hex::decode(info_hash).unwrap().try_into().unwrap(),
-            )
-            .await
-            .unwrap();
+            let mut tracker = UdpTracker::new(uri.clone(), None, info_hash).await.unwrap();
 
             Ok(tracker.stream_peers().await.unwrap())
          }
@@ -68,18 +64,18 @@ impl Tracker {
    }
 }
 
-struct AnnounceUriVisitor;
+struct TrackerVisitor;
 
 impl<'de> Deserialize<'de> for Tracker {
    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
    where
       D: serde::Deserializer<'de>,
    {
-      deserializer.deserialize_string(AnnounceUriVisitor)
+      deserializer.deserialize_string(TrackerVisitor)
    }
 }
 
-impl Visitor<'_> for AnnounceUriVisitor {
+impl Visitor<'_> for TrackerVisitor {
    type Value = Tracker;
 
    fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {