ws: allow peers to use multiple PeerIds, but only one per torrent

This commit is contained in:
Joakim Frostegård 2022-10-09 22:21:06 +02:00
parent 92794444f9
commit e0c4a4eaf7

View file

@ -24,7 +24,8 @@ 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, HashSet}; use hashbrown::hash_map::Entry;
use hashbrown::HashMap;
use slab::Slab; use slab::Slab;
use crate::config::Config; use crate::config::Config;
@ -45,8 +46,7 @@ struct ConnectionReference {
out_message_sender: Rc<LocalSender<(OutMessageMeta, OutMessage)>>, out_message_sender: Rc<LocalSender<(OutMessageMeta, 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: HashMap<InfoHash, PeerId>,
announced_info_hashes: HashSet<InfoHash>,
ip_version: IpVersion, ip_version: IpVersion,
} }
@ -145,7 +145,6 @@ pub async fn run_socket_worker(
server_start_instant, server_start_instant,
config.cleaning.max_connection_idle, config.cleaning.max_connection_idle,
), ),
peer_id: None,
announced_info_hashes: Default::default(), announced_info_hashes: Default::default(),
ip_version, ip_version,
}); });
@ -180,8 +179,7 @@ pub async fn run_socket_worker(
// Tell swarm workers to remove peer // Tell swarm workers to remove peer
if let Some(reference) = opt_reference { if let Some(reference) = opt_reference {
if let Some(peer_id) = reference.peer_id { for (info_hash, peer_id) in reference.announced_info_hashes {
for info_hash in reference.announced_info_hashes {
let message = SwarmControlMessage::ConnectionClosed { let message = SwarmControlMessage::ConnectionClosed {
info_hash, info_hash,
peer_id, peer_id,
@ -201,7 +199,6 @@ pub async fn run_socket_worker(
.unwrap(); .unwrap();
} }
} }
}
}), tq_regular) }), tq_regular)
.unwrap() .unwrap()
.detach(); .detach();
@ -499,26 +496,31 @@ impl<S: futures::AsyncRead + futures::AsyncWrite + Unpin> ConnectionReader<S> {
})?; })?;
// Store peer id / check if stored peer id matches // Store peer id / check if stored peer id matches
match &mut connection_reference.peer_id { match connection_reference
Some(peer_id) if *peer_id != announce_request.peer_id => { .announced_info_hashes
.entry(announce_request.info_hash)
{
Entry::Occupied(entry) => {
if *entry.get() != announce_request.peer_id {
// Drop Rc borrow before awaiting
drop(connection_slab);
self.send_error_response( self.send_error_response(
"Only one peer id can be used per connection".into(), "Only one peer id can be used per torrent".into(),
Some(ErrorResponseAction::Announce), Some(ErrorResponseAction::Announce),
Some(info_hash), Some(info_hash),
) )
.await?; .await?;
return Err(anyhow::anyhow!("Peer used more than one PeerId"));
}
Some(_) => (),
opt_peer_id @ None => {
*opt_peer_id = Some(announce_request.peer_id);
}
}
// Remember info hash for later return Err(anyhow::anyhow!(
connection_reference "Peer used more than one PeerId for a single torrent"
.announced_info_hashes ));
.insert(announce_request.info_hash); }
}
Entry::Vacant(entry) => {
entry.insert(announce_request.peer_id);
}
}
} }
let in_message = InMessage::AnnounceRequest(announce_request); let in_message = InMessage::AnnounceRequest(announce_request);