From eeaae29fa8d7985c06f99ed69acc2cc08a372de7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Joakim=20Frosteg=C3=A5rd?= Date: Sun, 3 Apr 2022 16:26:41 +0200 Subject: [PATCH] http_private: handle announce requests in request worker --- Cargo.lock | 1 + aquatic_http_private/Cargo.toml | 1 + .../workers/{request.rs => request/common.rs} | 78 +------ .../src/workers/request/mod.rs | 210 ++++++++++++++++++ 4 files changed, 215 insertions(+), 75 deletions(-) rename aquatic_http_private/src/workers/{request.rs => request/common.rs} (59%) create mode 100644 aquatic_http_private/src/workers/request/mod.rs diff --git a/Cargo.lock b/Cargo.lock index 5df294d..a1ba986 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -153,6 +153,7 @@ dependencies = [ "hex", "log", "mimalloc", + "rand", "serde", "socket2 0.4.4", "sqlx", diff --git a/aquatic_http_private/Cargo.toml b/aquatic_http_private/Cargo.toml index 272fed8..56189a9 100644 --- a/aquatic_http_private/Cargo.toml +++ b/aquatic_http_private/Cargo.toml @@ -21,6 +21,7 @@ dotenv = "0.15" hex = "0.4" log = "0.4" mimalloc = { version = "0.1", default-features = false } +rand = { version = "0.8", features = ["small_rng"] } serde = { version = "1", features = ["derive"] } socket2 = { version = "0.4", features = ["all"] } sqlx = { version = "0.5", features = [ "runtime-tokio-rustls" , "mysql" ] } diff --git a/aquatic_http_private/src/workers/request.rs b/aquatic_http_private/src/workers/request/common.rs similarity index 59% rename from aquatic_http_private/src/workers/request.rs rename to aquatic_http_private/src/workers/request/common.rs index be09c0d..abbcbc3 100644 --- a/aquatic_http_private/src/workers/request.rs +++ b/aquatic_http_private/src/workers/request/common.rs @@ -1,19 +1,9 @@ -use std::cell::RefCell; use std::net::{Ipv4Addr, Ipv6Addr}; -use std::rc::Rc; use std::time::Instant; -use tokio::sync::mpsc::Receiver; -use tokio::task::LocalSet; -use tokio::time; - -use aquatic_common::{AmortizedIndexMap, CanonicalSocketAddr, ValidUntil}; +use aquatic_common::{AmortizedIndexMap, ValidUntil}; use aquatic_http_protocol::common::{AnnounceEvent, InfoHash, PeerId}; -use aquatic_http_protocol::request::AnnounceRequest; -use aquatic_http_protocol::response::{FailureResponse, Response, ResponsePeer}; - -use crate::common::ChannelAnnounceRequest; -use crate::config::Config; +use aquatic_http_protocol::response::ResponsePeer; pub trait Ip: ::std::fmt::Debug + Copy + Eq + ::std::hash::Hash {} @@ -63,7 +53,7 @@ impl Peer { #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct PeerMapKey { pub peer_id: PeerId, - pub ip: I, + pub ip_address: I, } pub type PeerMap = AmortizedIndexMap, Peer>; @@ -130,65 +120,3 @@ impl TorrentMaps { torrent_map.shrink_to_fit(); } } - -pub fn run_request_worker( - config: Config, - request_receiver: Receiver, -) -> anyhow::Result<()> { - let runtime = tokio::runtime::Builder::new_current_thread() - .enable_all() - .build()?; - - runtime.block_on(run_inner(config, request_receiver))?; - - Ok(()) -} - -async fn run_inner( - config: Config, - mut request_receiver: Receiver, -) -> anyhow::Result<()> { - let torrents = Rc::new(RefCell::new(TorrentMaps::default())); - - LocalSet::new().spawn_local(periodically_clean_torrents( - config.clone(), - torrents.clone(), - )); - - loop { - let request = request_receiver - .recv() - .await - .ok_or_else(|| anyhow::anyhow!("request channel closed"))?; - - let response = handle_announce_request( - &config, - &torrents, - request.source_addr, - request.request.into(), - ); - - let _ = request.response_sender.send(response); - } -} - -fn handle_announce_request( - config: &Config, - torrents: &Rc>, - source_addr: CanonicalSocketAddr, - request: AnnounceRequest, -) -> Response { - Response::Failure(FailureResponse::new("actually successful")) -} - -async fn periodically_clean_torrents(config: Config, torrents: Rc>) { - let mut interval = time::interval(time::Duration::from_secs( - config.cleaning.torrent_cleaning_interval, - )); - - loop { - interval.tick().await; - - torrents.borrow_mut().clean(); - } -} diff --git a/aquatic_http_private/src/workers/request/mod.rs b/aquatic_http_private/src/workers/request/mod.rs new file mode 100644 index 0000000..358ead6 --- /dev/null +++ b/aquatic_http_private/src/workers/request/mod.rs @@ -0,0 +1,210 @@ +mod common; + +use std::cell::RefCell; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; +use std::rc::Rc; + +use aquatic_http_protocol::request::AnnounceRequest; +use rand::prelude::SmallRng; +use rand::SeedableRng; +use tokio::sync::mpsc::Receiver; +use tokio::task::LocalSet; +use tokio::time; + +use aquatic_common::{extract_response_peers, CanonicalSocketAddr, ValidUntil}; +use aquatic_http_protocol::response::{ + AnnounceResponse, Response, ResponsePeer, ResponsePeerListV4, ResponsePeerListV6, +}; + +use crate::common::ChannelAnnounceRequest; +use crate::config::Config; + +use common::*; + +pub fn run_request_worker( + config: Config, + request_receiver: Receiver, +) -> anyhow::Result<()> { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + + runtime.block_on(run_inner(config, request_receiver))?; + + Ok(()) +} + +async fn run_inner( + config: Config, + mut request_receiver: Receiver, +) -> anyhow::Result<()> { + let torrents = Rc::new(RefCell::new(TorrentMaps::default())); + let mut rng = SmallRng::from_entropy(); + + LocalSet::new().spawn_local(periodically_clean_torrents( + config.clone(), + torrents.clone(), + )); + + loop { + let request = request_receiver + .recv() + .await + .ok_or_else(|| anyhow::anyhow!("request channel closed"))?; + + let valid_until = ValidUntil::new(config.cleaning.max_peer_age); + + let response = handle_announce_request( + &config, + &mut rng, + &mut torrents.borrow_mut(), + valid_until, + request.source_addr, + request.request.into(), + ); + + let _ = request.response_sender.send(Response::Announce(response)); + } +} + +async fn periodically_clean_torrents(config: Config, torrents: Rc>) { + let mut interval = time::interval(time::Duration::from_secs( + config.cleaning.torrent_cleaning_interval, + )); + + loop { + interval.tick().await; + + torrents.borrow_mut().clean(); + } +} + +fn handle_announce_request( + config: &Config, + rng: &mut SmallRng, + torrent_maps: &mut TorrentMaps, + valid_until: ValidUntil, + source_addr: CanonicalSocketAddr, + request: AnnounceRequest, +) -> AnnounceResponse { + match source_addr.get().ip() { + IpAddr::V4(source_ip) => { + let torrent_data: &mut TorrentData = + torrent_maps.ipv4.entry(request.info_hash).or_default(); + + let (seeders, leechers, response_peers) = upsert_peer_and_get_response_peers( + config, + rng, + torrent_data, + source_ip, + request, + valid_until, + ); + + let response = AnnounceResponse { + complete: seeders, + incomplete: leechers, + announce_interval: config.protocol.peer_announce_interval, + peers: ResponsePeerListV4(response_peers), + peers6: ResponsePeerListV6(vec![]), + warning_message: None, + }; + + response + } + IpAddr::V6(source_ip) => { + let torrent_data: &mut TorrentData = + torrent_maps.ipv6.entry(request.info_hash).or_default(); + + let (seeders, leechers, response_peers) = upsert_peer_and_get_response_peers( + config, + rng, + torrent_data, + source_ip, + request, + valid_until, + ); + + let response = AnnounceResponse { + complete: seeders, + incomplete: leechers, + announce_interval: config.protocol.peer_announce_interval, + peers: ResponsePeerListV4(vec![]), + peers6: ResponsePeerListV6(response_peers), + warning_message: None, + }; + + response + } + } +} + +/// Insert/update peer. Return num_seeders, num_leechers and response peers +pub fn upsert_peer_and_get_response_peers( + config: &Config, + rng: &mut SmallRng, + torrent_data: &mut TorrentData, + source_ip: I, + request: AnnounceRequest, + valid_until: ValidUntil, +) -> (usize, usize, Vec>) { + // Insert/update/remove peer who sent this request + + let peer_status = + PeerStatus::from_event_and_bytes_left(request.event, Some(request.bytes_left)); + + let peer = Peer { + ip_address: source_ip, + port: request.port, + status: peer_status, + valid_until, + }; + + let peer_map_key = PeerMapKey { + peer_id: request.peer_id, + ip_address: source_ip, + }; + + let opt_removed_peer = match peer_status { + PeerStatus::Leeching => { + torrent_data.num_leechers += 1; + + torrent_data.peers.insert(peer_map_key.clone(), peer) + } + PeerStatus::Seeding => { + torrent_data.num_seeders += 1; + + torrent_data.peers.insert(peer_map_key.clone(), peer) + } + PeerStatus::Stopped => torrent_data.peers.remove(&peer_map_key), + }; + + match opt_removed_peer.map(|peer| peer.status) { + Some(PeerStatus::Leeching) => { + torrent_data.num_leechers -= 1; + } + Some(PeerStatus::Seeding) => { + torrent_data.num_seeders -= 1; + } + _ => {} + } + + let max_num_peers_to_take = match request.numwant { + Some(0) | None => config.protocol.max_peers, + Some(numwant) => numwant.min(config.protocol.max_peers), + }; + + let response_peers: Vec> = extract_response_peers( + rng, + &torrent_data.peers, + max_num_peers_to_take, + peer_map_key, + Peer::to_response_peer, + ); + + ( + torrent_data.num_seeders, + torrent_data.num_leechers, + response_peers, + ) +}