From c50dec349610f9d2a94cd97b485fca8b6b652a0a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= <joakim.frostegard@gmail.com>
Date: Sat, 4 Jul 2020 12:00:36 +0200
Subject: [PATCH] aquatic_http: add PeerMapKey with PeerID and Either<Ip,
 announce key>

---
 aquatic_http/src/lib/common.rs           | 21 ++++++++++-----------
 aquatic_http/src/lib/handler.rs          | 18 +++++++++++++++---
 aquatic_http/src/lib/protocol/request.rs | 11 +++++++++++
 3 files changed, 36 insertions(+), 14 deletions(-)

diff --git a/aquatic_http/src/lib/common.rs b/aquatic_http/src/lib/common.rs
index 21ccb85..980d1c3 100644
--- a/aquatic_http/src/lib/common.rs
+++ b/aquatic_http/src/lib/common.rs
@@ -1,6 +1,7 @@
-use std::net::SocketAddr;
+use std::net::{IpAddr, SocketAddr};
 use std::sync::Arc;
 
+use either::Either;
 use flume::{Sender, Receiver};
 use hashbrown::HashMap;
 use indexmap::IndexMap;
@@ -15,7 +16,6 @@ use crate::protocol::request::Request;
 use crate::protocol::response::Response;
 
 
-// identical to ws version
 #[derive(Clone, Copy, Debug)]
 pub struct ConnectionMeta {
     /// Index of socket worker responsible for this connection. Required for
@@ -26,7 +26,6 @@ pub struct ConnectionMeta {
 }
 
 
-// identical to ws version
 #[derive(PartialEq, Eq, Clone, Copy, Debug)]
 pub enum PeerStatus {
     Seeding,
@@ -35,7 +34,6 @@ pub enum PeerStatus {
 }
 
 
-// almost identical to ws version - FIXME only if bytes left is optional
 impl PeerStatus {
     /// Determine peer status from announce event and number of bytes left.
     /// 
@@ -65,11 +63,16 @@ pub struct Peer {
 }
 
 
-// identical to ws version
-pub type PeerMap = IndexMap<PeerId, Peer>;
+#[derive(Debug, Clone, PartialEq, Eq, Hash)]
+pub struct PeerMapKey {
+    pub peer_id: PeerId,
+    pub ip_or_key: Either<IpAddr, String>
+}
+
+
+pub type PeerMap = IndexMap<PeerMapKey, Peer>;
 
 
-// identical to ws version
 pub struct TorrentData {
     pub peers: PeerMap,
     pub num_seeders: usize,
@@ -77,7 +80,6 @@ pub struct TorrentData {
 }
 
 
-// identical to ws version
 impl Default for TorrentData {
     #[inline]
     fn default() -> Self {
@@ -90,11 +92,9 @@ impl Default for TorrentData {
 }
 
 
-// identical to ws version
 pub type TorrentMap = HashMap<InfoHash, TorrentData>;
 
 
-// identical to ws version
 #[derive(Default)]
 pub struct TorrentMaps {
     pub ipv4: TorrentMap,
@@ -102,7 +102,6 @@ pub struct TorrentMaps {
 }
 
 
-// identical to ws version
 #[derive(Clone)]
 pub struct State {
     pub torrent_maps: Arc<Mutex<TorrentMaps>>,
diff --git a/aquatic_http/src/lib/handler.rs b/aquatic_http/src/lib/handler.rs
index d18761b..f35d125 100644
--- a/aquatic_http/src/lib/handler.rs
+++ b/aquatic_http/src/lib/handler.rs
@@ -1,6 +1,7 @@
 use std::time::Duration;
 use std::vec::Drain;
 
+use either::Either;
 use hashbrown::HashMap;
 use parking_lot::MutexGuard;
 use rand::{Rng, SeedableRng, rngs::SmallRng};
@@ -116,19 +117,30 @@ pub fn handle_announce_requests(
                 valid_until,
             };
 
+            let ip_or_key = request.key
+                .map(|k| Either::Right(k))
+                .unwrap_or_else(||
+                    Either::Left(request_sender_meta.peer_addr.ip())
+                );
+
+            let peer_map_key = PeerMapKey {
+                peer_id: request.peer_id,
+                ip_or_key,
+            };
+
             let opt_removed_peer = match peer_status {
                 PeerStatus::Leeching => {
                     torrent_data.num_leechers += 1;
 
-                    torrent_data.peers.insert(request.peer_id, peer)
+                    torrent_data.peers.insert(peer_map_key, peer)
                 },
                 PeerStatus::Seeding => {
                     torrent_data.num_seeders += 1;
 
-                    torrent_data.peers.insert(request.peer_id, peer)
+                    torrent_data.peers.insert(peer_map_key, peer)
                 },
                 PeerStatus::Stopped => {
-                    torrent_data.peers.remove(&request.peer_id)
+                    torrent_data.peers.remove(&peer_map_key)
                 }
             };
 
diff --git a/aquatic_http/src/lib/protocol/request.rs b/aquatic_http/src/lib/protocol/request.rs
index 96eaa18..98db6e3 100644
--- a/aquatic_http/src/lib/protocol/request.rs
+++ b/aquatic_http/src/lib/protocol/request.rs
@@ -15,6 +15,7 @@ pub struct AnnounceRequest {
     pub compact: bool,
     /// Number of response peers wanted
     pub numwant: Option<usize>,
+    pub key: Option<String>,
 }
 
 
@@ -78,6 +79,15 @@ impl Request {
             } else {
                 None
             };
+            let key = if let Some(s) = data.get("key"){
+                if s.len() > 100 {
+                    return Err(anyhow::anyhow!("'key' is too long"))
+                }
+
+                Some(s.clone())
+            } else {
+                None
+            };
 
             let request = AnnounceRequest {
                 info_hash: info_hashes.get(0)
@@ -103,6 +113,7 @@ impl Request {
                     .map(|s| s == "1")
                     .unwrap_or(true),
                 numwant,
+                key,
             };
 
             Ok(Request::Announce(request))