code restructure

zung_torrent/src/sources/trackers/mod.rs

@@ -0,0 +1,289 @@
+//! For handleing torrent tracker requests and responses.
+//!
+//! The tracker is an HTTP/HTTPS service which responds to HTTP GET requests. The requests include
+//! metrics from clients that help the tracker keep overall statistics about the torrent. The
+//! response includes a peer list that helps the client participate in the torrent. The base URL
+//! consists of the "announce URL" as defined in the metainfo (.torrent) file. The parameters are
+//! then added to this URL, using standard CGI methods (i.e. a '?' after the announce URL, followed
+//! by 'param=value' sequences separated by '&').
+
+mod request;
+pub use request::*;
+
+use std::ops::Deref;
+use std::sync::Arc;
+
+use crate::meta_info::InfoHashEncoded;
+use crate::PeerID;
+use anyhow::{bail, Result};
+use futures::stream::FuturesUnordered;
+use tokio::task::JoinHandle;
+
+// TODO: Look into SmallStr
+#[derive(Debug)]
+pub enum Tracker {
+    Http(Arc<str>),
+    Udp(Arc<str>),
+    Invalid(Arc<str>),
+}
+
+impl Clone for Tracker {
+    fn clone(&self) -> Self {
+        match self {
+            Self::Http(arg0) => Self::Http(Arc::clone(arg0)),
+            Self::Udp(arg0) => Self::Udp(Arc::clone(arg0)),
+            Self::Invalid(arg0) => Self::Invalid(Arc::clone(arg0)),
+        }
+    }
+}
+
+impl Tracker {
+    pub fn new(tracker_url: &str) -> Self {
+        if tracker_url.starts_with("http") {
+            Self::Http(Arc::from(tracker_url))
+        } else if tracker_url.starts_with("udp") {
+            Self::Udp(Arc::from(tracker_url))
+        } else {
+            Self::Invalid(Arc::from(tracker_url))
+        }
+    }
+
+    pub fn url(&self) -> &str {
+        match self {
+            Tracker::Http(s) => s,
+            Tracker::Udp(s) => s,
+            Tracker::Invalid(s) => s,
+        }
+    }
+
+    pub async fn generate_request(
+        &self,
+        info_hash: InfoHashEncoded,
+        peer_id: PeerID,
+    ) -> Result<TrackerRequest> {
+        match self {
+            Tracker::Http(url) => Ok(TrackerRequest::Http {
+                url: url.clone(),
+                params: HttpTrackerRequestParams::new(info_hash, peer_id),
+            }),
+            Tracker::Udp(url) => {
+                let udp_url = url.strip_prefix("udp://").unwrap();
+                let udp_url = match udp_url.split_once("/") {
+                    Some(s) => s.0,
+                    None => udp_url,
+                };
+                let connection = UdpConnectRequest::new()
+                    .await?
+                    .connect_with(udp_url)
+                    .await?;
+
+                let connection_id = connection.connection_id();
+                Ok(TrackerRequest::Udp {
+                    url: url.clone(),
+                    connection_id,
+                    params: UdpTrackerRequestParams::new(connection_id, info_hash, peer_id),
+                })
+            }
+            Tracker::Invalid(url) => bail!("Unsupproted : {url}"),
+        }
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct TrackerList {
+    tracker_list: Vec<Tracker>,
+}
+
+impl TrackerList {
+    pub(crate) fn new(tracker_list: Vec<Tracker>) -> Self {
+        Self { tracker_list }
+    }
+
+    fn as_array(&self) -> &[Tracker] {
+        &self.tracker_list
+    }
+
+    /// Consumes the tracker list and returns the internal Vec of [`Tracker`]s.
+    pub fn into_vec(self) -> Vec<Tracker> {
+        self.tracker_list
+    }
+
+    /// Asyncly generates the [`TrackerRequest`]
+    ///
+    // TODO: Revisit this if there is a faster more efficient way.
+    pub fn generate_requests(
+        &self,
+        info_hash: InfoHashEncoded,
+        peer_id: PeerID,
+    ) -> FuturesUnordered<JoinHandle<Result<TrackerRequest>>> {
+        self.as_array()
+            .iter()
+            .cloned() // The clone here is just Arc::clone
+            .map(|tracker| {
+                tokio::spawn(async move { tracker.generate_request(info_hash, peer_id).await })
+            })
+            .collect()
+    }
+}
+
+impl Deref for TrackerList {
+    type Target = [Tracker];
+
+    fn deref(&self) -> &Self::Target {
+        self.as_array()
+    }
+}
+
+// Iterator implementation
+impl<'a> IntoIterator for &'a TrackerList {
+    type Item = &'a Tracker;
+    type IntoIter = std::slice::Iter<'a, Tracker>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.tracker_list.iter()
+    }
+}
+
+#[cfg(test)]
+mod tracker_tests {
+    use super::*;
+    use crate::meta_info::InfoHash;
+
+    // Test creation of a new TrackerRequest with default parameters.
+    #[tokio::test]
+    async fn test_tracker_request_creation() {
+        let sample_url = "http://example.com/announce";
+        let info_hash = InfoHash::new(b"test info_hash").as_encoded();
+        let peer_id = PeerID::default();
+        let tracker_request = Tracker::new(sample_url);
+        let tracker_request = tracker_request
+            .generate_request(info_hash, peer_id)
+            .await
+            .unwrap();
+
+        match tracker_request {
+            TrackerRequest::Http { url, params } => {
+                assert_eq!(url.as_ref(), sample_url);
+                assert_eq!(params.port, 6881);
+                assert_eq!(params.uploaded, 0);
+                assert_eq!(params.downloaded, 0);
+                assert_eq!(params.left, 0);
+                assert!(params.compact);
+                assert!(!params.no_peer_id);
+                assert_eq!(params.event, Some(Event::Started));
+                assert_eq!(params.numwant, Some(0));
+            }
+            TrackerRequest::Udp { .. } => {
+                unreachable!("Why is http being read as upd?")
+            }
+        }
+    }
+
+    // Test to_url method to check if URL is correctly formatted with query parameters.
+    #[tokio::test]
+    async fn test_tracker_request_to_url() {
+        let url = "http://example.com/announce";
+        let info_hash = InfoHash::new(b"test info_hash").as_encoded();
+        let peer_id = PeerID::default();
+        let tracker_request = Tracker::new(url);
+        let tracker_request = tracker_request
+            .generate_request(info_hash, peer_id)
+            .await
+            .unwrap();
+
+        // Generate the URL with query parameters
+        let generated_url = tracker_request.to_url().unwrap();
+
+        match tracker_request {
+            TrackerRequest::Http { params, .. } => {
+                // Verify that essential parts of the URL exist
+                assert!(generated_url.contains("http://example.com/announce"));
+                assert!(
+                    generated_url.contains(&format!("info_hash={}", info_hash.to_url_encoded()))
+                );
+                assert!(
+                    generated_url.contains(&format!("peer_id={}", params.peer_id.to_url_encoded()))
+                );
+            }
+            _ => panic!(),
+        }
+    }
+
+    // Test serialization of booleans as integers.
+    #[tokio::test]
+    async fn test_bool_as_int_serialization() {
+        let url = "http://example.com/announce";
+        let info_hash = InfoHash::new(b"test info_hash").as_encoded();
+
+        let peer_id = PeerID::default();
+        let tracker_request = Tracker::new(url);
+        let mut tracker_request = tracker_request
+            .generate_request(info_hash, peer_id)
+            .await
+            .unwrap();
+
+        match &mut tracker_request {
+            TrackerRequest::Http { params, .. } => {
+                // Set compact and no_peer_id to true to check if they serialize to 1
+                params.compact = true;
+                params.no_peer_id = true;
+            }
+            _ => panic!(),
+        }
+
+        let generated_url = tracker_request.to_url().unwrap();
+
+        // Check that the values serialize as integers (1 for true)
+        assert!(generated_url.contains("compact=1"));
+        assert!(generated_url.contains("no_peer_id=1"));
+
+        match &mut tracker_request {
+            TrackerRequest::Http { params, .. } => {
+                // Set them to false to check if they serialize to 0
+                params.compact = false;
+                params.no_peer_id = false;
+            }
+            _ => panic!(),
+        }
+
+        let generated_url = tracker_request.to_url().unwrap();
+
+        // Check that the values serialize as integers (0 for false)
+        assert!(generated_url.contains("compact=0"));
+        assert!(generated_url.contains("no_peer_id=0"));
+    }
+
+    // Test optional parameters like IP, numwant, key, and trackerid.
+    #[tokio::test]
+    async fn test_optional_parameters() {
+        let url = "http://example.com/announce";
+        let info_hash = InfoHash::new(b"test info_hash").as_encoded();
+        let peer_id = PeerID::default();
+        let tracker_request = Tracker::new(url);
+        let mut tracker_request = tracker_request
+            .generate_request(info_hash, peer_id)
+            .await
+            .unwrap();
+
+        match &mut tracker_request {
+            TrackerRequest::Http { params, .. } => {
+                // Set optional parameters
+                params.ip = Some("2001db81".to_string());
+                params.numwant = Some(25);
+                params.key = Some("unique-key".to_string());
+                params.trackerid = Some(TrackerID {
+                    id: "tracker-id-123".to_string(),
+                });
+
+                let generated_url = tracker_request.to_url().unwrap();
+
+                // Check that the optional parameters are included in the URL if provided
+                assert!(generated_url.contains("ip=2001db81"));
+                assert!(generated_url.contains("numwant=25"));
+                assert!(generated_url.contains("key=unique-key"));
+                assert!(generated_url.contains("trackerid=tracker-id-123"));
+            }
+            _ => panic!(),
+        }
+    }
+}