YGGverse/aquatic
1
0
Fork 0
mirror of https://github.com/YGGverse/aquatic.git synced 2026-04-01 10:15:31 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

ws: store peer_id and announced info hashes in ConnectionReference

Browse source
This commit is contained in:
Joakim Frostegård 2022-07-05 12:03:51 +02:00
parent b06c12e9a5
commit b30da1a930
1 changed files with 56 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

64
aquatic_ws/src/workers/socket.rs
View file

@ -25,7 +25,7 @@ use glommio::net::{TcpListener, TcpStream};
use glommio::task::JoinHandle; use glommio::task::JoinHandle;
use glommio::timer::{sleep, timeout, TimerActionRepeat}; use glommio::timer::{sleep, timeout, TimerActionRepeat};
use glommio::{enclose, prelude::*}; use glommio::{enclose, prelude::*};
use hashbrown::HashMap; use hashbrown::{HashMap, HashSet};
use slab::Slab; use slab::Slab;
use crate::config::Config; use crate::config::Config;
@ -46,6 +46,8 @@ struct ConnectionReference {
out_message_sender: Rc<LocalSender<(ConnectionMeta, OutMessage)>>, out_message_sender: Rc<LocalSender<(ConnectionMeta, OutMessage)>>,
/// Updated after sending message to peer /// Updated after sending message to peer
valid_until: ValidUntil, valid_until: ValidUntil,
peer_id: Option<PeerId>,
announced_info_hashes: HashSet<InfoHash>,
} }
pub async fn run_socket_worker( pub async fn run_socket_worker(
@ -112,6 +114,8 @@ pub async fn run_socket_worker(
task_handle: None, task_handle: None,
out_message_sender: out_message_sender.clone(), out_message_sender: out_message_sender.clone(),
valid_until: ValidUntil::new(config.cleaning.max_connection_idle), valid_until: ValidUntil::new(config.cleaning.max_connection_idle),
peer_id: None,
announced_info_hashes: Default::default(),
}); });
::log::info!("accepting stream: {}", key); ::log::info!("accepting stream: {}", key);
@ -240,10 +244,11 @@ async fn run_connection(
let access_list_cache = create_access_list_cache(&access_list); let access_list_cache = create_access_list_cache(&access_list);
let reader_handle = spawn_local_into( let reader_handle = spawn_local_into(
enclose!((config, pending_scrape_slab) async move { enclose!((config, connection_slab, pending_scrape_slab) async move {
let mut reader = ConnectionReader { let mut reader = ConnectionReader {
config, config,
access_list_cache, access_list_cache,
connection_slab,
in_message_senders, in_message_senders,
out_message_sender, out_message_sender,
pending_scrape_slab, pending_scrape_slab,
@ -289,6 +294,7 @@ async fn run_connection(
struct ConnectionReader { struct ConnectionReader {
config: Rc<Config>, config: Rc<Config>,
access_list_cache: AccessListCache, access_list_cache: AccessListCache,
connection_slab: Rc<RefCell<Slab<ConnectionReference>>>,
in_message_senders: Rc<Senders<(ConnectionMeta, InMessage)>>, in_message_senders: Rc<Senders<(ConnectionMeta, InMessage)>>,
out_message_sender: Rc<LocalSender<(ConnectionMeta, OutMessage)>>, out_message_sender: Rc<LocalSender<(ConnectionMeta, OutMessage)>>,
pending_scrape_slab: Rc<RefCell<Slab<PendingScrapeResponse>>>, pending_scrape_slab: Rc<RefCell<Slab<PendingScrapeResponse>>>,
@ -320,7 +326,7 @@ impl ConnectionReader {
Err(err) => { Err(err) => {
::log::debug!("Couldn't parse in_message: {:?}", err); ::log::debug!("Couldn't parse in_message: {:?}", err);
self.send_error_response("Invalid request".into(), None) self.send_error_response("Invalid request".into(), None, None)
.await?; .await?;
} }
} }
@ -339,6 +345,39 @@ impl ConnectionReader {
.load() .load()
.allows(self.config.access_list.mode, &info_hash.0) .allows(self.config.access_list.mode, &info_hash.0)
{ {
{
let mut connection_slab = self.connection_slab.borrow_mut();
let connection_reference = connection_slab
.get_mut(self.connection_id.0)
.ok_or_else(|| {
anyhow::anyhow!(
"connection reference {} not found in slab",
self.connection_id.0
)
})?;
match &mut connection_reference.peer_id {
Some(peer_id) if *peer_id != announce_request.peer_id => {
self.send_error_response(
"Only one peer id can be used per connection".into(),
Some(ErrorResponseAction::Announce),
Some(info_hash),
)
.await?;
return Err(anyhow::anyhow!("Peer used more than one PeerId"));
}
Some(_) => (),
opt_peer_id @ None => {
*opt_peer_id = Some(announce_request.peer_id);
}
}
connection_reference
.announced_info_hashes
.insert(announce_request.info_hash);
}
let in_message = InMessage::AnnounceRequest(announce_request); let in_message = InMessage::AnnounceRequest(announce_request);
let consumer_index = let consumer_index =
@ -354,8 +393,12 @@ impl ConnectionReader {
.unwrap(); .unwrap();
::log::info!("sent message to swarm worker"); ::log::info!("sent message to swarm worker");
} else { } else {
self.send_error_response("Info hash not allowed".into(), Some(info_hash)) self.send_error_response(
.await?; "Info hash not allowed".into(),
Some(ErrorResponseAction::Announce),
Some(info_hash),
)
.await?;
} }
} }
InMessage::ScrapeRequest(ScrapeRequest { info_hashes, .. }) => { InMessage::ScrapeRequest(ScrapeRequest { info_hashes, .. }) => {
@ -364,8 +407,12 @@ impl ConnectionReader {
} else { } else {
// If request.info_hashes is empty, don't return scrape for all // If request.info_hashes is empty, don't return scrape for all
// torrents, even though reference server does it. It is too expensive. // torrents, even though reference server does it. It is too expensive.
self.send_error_response("Full scrapes are not allowed".into(), None) self.send_error_response(
.await?; "Full scrapes are not allowed".into(),
Some(ErrorResponseAction::Scrape),
None,
)
.await?;
return Ok(()); return Ok(());
}; };
@ -415,10 +462,11 @@ impl ConnectionReader {
async fn send_error_response( async fn send_error_response(
&self, &self,
failure_reason: Cow<'static, str>, failure_reason: Cow<'static, str>,
action: Option<ErrorResponseAction>,
info_hash: Option<InfoHash>, info_hash: Option<InfoHash>,
) -> anyhow::Result<()> { ) -> anyhow::Result<()> {
let out_message = OutMessage::ErrorResponse(ErrorResponse { let out_message = OutMessage::ErrorResponse(ErrorResponse {
action: Some(ErrorResponseAction::Scrape), action,
failure_reason, failure_reason,
info_hash, info_hash,
}); });