fix(client): bound the threads a peer flood can create
One popular search produced a ConnectToPeer for every responder, and each one cost a 2 MiB wrapper thread, an unbounded 256 KiB actor, a listener thread per inbound socket, and on failure a 2 MiB thread that slept for 20 seconds. A container with a pids limit died on one album search (issues #32, #33). Four bounds, no behavior change for a quiet session: - P-type registrations run inline on the ops loop; only F transfers keep a worker thread. The dial already ran on the actor thread. - The broker timeout is a deadline swept by the ops loop instead of a sleeping thread per failed dial. - The registry caps live actors (default 512, Client::set_max_peers). At capacity the oldest peer outside the caller's protected set (live downloads, uploads, pending serves, queued messages, browses in flight) is evicted; if none qualify the registration is refused before any thread is spawned. - Concurrent listener handshakes are bounded by a 128-permit semaphore released once the peer-init arrives, so slow handshakes cannot multiply threads while transfers stay unthrottled. A replayed TransferResponse for an InProgress or Completed token is dropped instead of starting a second transfer thread over the same .part file. Fixes #32 Fixes #33
This commit is contained in:
@@ -6,11 +6,12 @@ use crate::peer::Peer;
|
||||
use crate::utils::lock::MutexExt;
|
||||
use crate::{debug, error};
|
||||
|
||||
use std::collections::HashMap;
|
||||
use std::collections::{HashMap, HashSet};
|
||||
use std::net::TcpStream;
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering};
|
||||
use std::sync::mpsc::Sender;
|
||||
use std::sync::{Arc, Mutex};
|
||||
use std::time::Instant;
|
||||
|
||||
/// Source of unique per-actor ids so terminal-outcome eviction can be made
|
||||
/// identity-aware (a replaced actor must not evict its replacement).
|
||||
@@ -18,13 +19,14 @@ static NEXT_PEER_ID: AtomicU64 = AtomicU64::new(1);
|
||||
|
||||
/// Registered peers keyed by username, each stored with the unique id of the
|
||||
/// actor currently occupying the slot.
|
||||
type PeerMap = HashMap<String, (u64, ActorHandle<PeerMessage>)>;
|
||||
type PeerMap = HashMap<String, (u64, ActorHandle<PeerMessage>, Instant)>;
|
||||
|
||||
pub struct PeerRegistry {
|
||||
peers: Arc<Mutex<PeerMap>>,
|
||||
actor_system: Arc<ActorSystem>,
|
||||
client_channel: Sender<ClientOperation>,
|
||||
own_username: String,
|
||||
capacity: Arc<AtomicUsize>,
|
||||
}
|
||||
|
||||
impl PeerRegistry {
|
||||
@@ -33,12 +35,28 @@ impl PeerRegistry {
|
||||
actor_system: Arc<ActorSystem>,
|
||||
client_channel: Sender<ClientOperation>,
|
||||
own_username: String,
|
||||
) -> Self {
|
||||
Self::with_capacity(
|
||||
actor_system,
|
||||
client_channel,
|
||||
own_username,
|
||||
Arc::new(AtomicUsize::new(usize::MAX)),
|
||||
)
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub(crate) fn with_capacity(
|
||||
actor_system: Arc<ActorSystem>,
|
||||
client_channel: Sender<ClientOperation>,
|
||||
own_username: String,
|
||||
capacity: Arc<AtomicUsize>,
|
||||
) -> Self {
|
||||
Self {
|
||||
peers: Arc::new(Mutex::new(HashMap::new())),
|
||||
actor_system,
|
||||
client_channel,
|
||||
own_username,
|
||||
capacity,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -47,6 +65,16 @@ impl PeerRegistry {
|
||||
peer: Peer,
|
||||
stream: Option<TcpStream>,
|
||||
reader: Option<MessageReader>,
|
||||
) -> Result<ActorHandle<PeerMessage>, String> {
|
||||
self.register_peer_protected(peer, stream, reader, &HashSet::new())
|
||||
}
|
||||
|
||||
pub(crate) fn register_peer_protected(
|
||||
&self,
|
||||
peer: Peer,
|
||||
stream: Option<TcpStream>,
|
||||
reader: Option<MessageReader>,
|
||||
protected: &HashSet<String>,
|
||||
) -> Result<ActorHandle<PeerMessage>, String> {
|
||||
let username = peer.username.clone();
|
||||
let id = NEXT_PEER_ID.fetch_add(1, Ordering::Relaxed);
|
||||
@@ -71,6 +99,28 @@ impl PeerRegistry {
|
||||
.lock_safe()
|
||||
.map_err(|e| format!("peer registry lock poisoned: {e}"))?;
|
||||
|
||||
let capacity = self.capacity.load(Ordering::Relaxed).max(1);
|
||||
if peers.len() >= capacity && !peers.contains_key(&username) {
|
||||
let victim = peers
|
||||
.iter()
|
||||
.filter(|(name, _)| !protected.contains(name.as_str()))
|
||||
.min_by_key(|(_, (_, _, registered))| *registered)
|
||||
.map(|(name, _)| name.clone());
|
||||
let Some(victim) = victim else {
|
||||
return Err(format!(
|
||||
"peer registry at capacity ({capacity}) and every peer \
|
||||
is busy; refusing {username}"
|
||||
));
|
||||
};
|
||||
if let Some((_, handle, _)) = peers.remove(&victim) {
|
||||
let _ = handle.stop();
|
||||
debug!(
|
||||
"[peer_registry] evicted idle peer {} to admit {}",
|
||||
victim, username
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
let handle = self
|
||||
.actor_system
|
||||
.try_spawn_with_handle(actor, |actor, handle| {
|
||||
@@ -81,8 +131,8 @@ impl PeerRegistry {
|
||||
// become an orphan pinning a pool worker forever. Eviction on the
|
||||
// replaced actor's later shutdown is identity-aware (keyed on its id),
|
||||
// so stopping it here cannot evict this new connection.
|
||||
if let Some((_, old_handle)) =
|
||||
peers.insert(username.clone(), (id, handle.clone()))
|
||||
if let Some((_, old_handle, _)) =
|
||||
peers.insert(username.clone(), (id, handle.clone(), Instant::now()))
|
||||
{
|
||||
let _ = old_handle.stop();
|
||||
debug!(
|
||||
@@ -97,7 +147,9 @@ impl PeerRegistry {
|
||||
#[must_use]
|
||||
pub fn get_peer(&self, username: &str) -> Option<ActorHandle<PeerMessage>> {
|
||||
match self.peers.lock_safe() {
|
||||
Ok(peers) => peers.get(username).map(|(_, handle)| handle.clone()),
|
||||
Ok(peers) => {
|
||||
peers.get(username).map(|(_, handle, _)| handle.clone())
|
||||
}
|
||||
Err(e) => {
|
||||
error!("[peer_registry] get_peer: {}", e);
|
||||
None
|
||||
@@ -123,7 +175,7 @@ impl PeerRegistry {
|
||||
debug!("[peer_registry] Removed peer actor for {}", username);
|
||||
}
|
||||
|
||||
removed.map(|(_, handle)| handle)
|
||||
removed.map(|(_, handle, _)| handle)
|
||||
}
|
||||
|
||||
/// Remove and return the actor for `username` only if it is still the actor
|
||||
@@ -142,8 +194,11 @@ impl PeerRegistry {
|
||||
return None;
|
||||
}
|
||||
};
|
||||
if peers.get(username).is_some_and(|(stored, _)| *stored == id) {
|
||||
let removed = peers.remove(username).map(|(_, handle)| handle);
|
||||
if peers
|
||||
.get(username)
|
||||
.is_some_and(|(stored, _, _)| *stored == id)
|
||||
{
|
||||
let removed = peers.remove(username).map(|(_, handle, _)| handle);
|
||||
debug!(
|
||||
"[peer_registry] Removed peer actor {} for {}",
|
||||
id, username
|
||||
@@ -192,6 +247,7 @@ impl Clone for PeerRegistry {
|
||||
actor_system: self.actor_system.clone(),
|
||||
client_channel: self.client_channel.clone(),
|
||||
own_username: self.own_username.clone(),
|
||||
capacity: self.capacity.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -201,8 +257,112 @@ mod tests {
|
||||
use super::PeerRegistry;
|
||||
use crate::actor::ActorSystem;
|
||||
use crate::peer::{ConnectionType, Peer};
|
||||
use std::collections::HashSet;
|
||||
use std::net::{TcpListener, TcpStream};
|
||||
use std::sync::Arc;
|
||||
use std::sync::atomic::AtomicUsize;
|
||||
use std::time::Duration;
|
||||
|
||||
fn loopback_peer(name: &str) -> (Peer, TcpStream, TcpStream) {
|
||||
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let stream = TcpStream::connect(addr).unwrap();
|
||||
stream.set_nonblocking(true).unwrap();
|
||||
let far_end = listener.accept().unwrap().0;
|
||||
let peer = Peer::new(
|
||||
name.to_string(),
|
||||
ConnectionType::P,
|
||||
"127.0.0.1".to_string(),
|
||||
u32::from(addr.port()),
|
||||
None,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
);
|
||||
(peer, stream, far_end)
|
||||
}
|
||||
|
||||
fn capped_registry(
|
||||
capacity: usize,
|
||||
) -> (
|
||||
PeerRegistry,
|
||||
std::sync::mpsc::Receiver<crate::client::ClientOperation>,
|
||||
) {
|
||||
let system = Arc::new(ActorSystem::new());
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
let registry = PeerRegistry::with_capacity(
|
||||
system,
|
||||
tx,
|
||||
"me".to_string(),
|
||||
Arc::new(AtomicUsize::new(capacity)),
|
||||
);
|
||||
(registry, rx)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capacity_evicts_the_oldest_idle_peer() {
|
||||
let (registry, _rx) = capped_registry(2);
|
||||
|
||||
for name in ["a", "b", "c"] {
|
||||
let (peer, stream, _far_end) = loopback_peer(name);
|
||||
registry.register_peer(peer, Some(stream), None).unwrap();
|
||||
std::thread::sleep(Duration::from_millis(5));
|
||||
}
|
||||
|
||||
assert!(!registry.contains("a"), "the oldest peer must be evicted");
|
||||
assert!(registry.contains("b"));
|
||||
assert!(registry.contains("c"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn capacity_refuses_when_every_peer_is_protected() {
|
||||
let (registry, _rx) = capped_registry(1);
|
||||
|
||||
let (busy, busy_stream, _busy_far) = loopback_peer("busy");
|
||||
registry
|
||||
.register_peer(busy, Some(busy_stream), None)
|
||||
.unwrap();
|
||||
|
||||
let protected: HashSet<String> =
|
||||
std::iter::once("busy".to_string()).collect();
|
||||
let (newcomer, stream, _far) = loopback_peer("newcomer");
|
||||
let result = registry.register_peer_protected(
|
||||
newcomer,
|
||||
Some(stream),
|
||||
None,
|
||||
&protected,
|
||||
);
|
||||
|
||||
assert!(result.is_err(), "a full registry of busy peers must refuse");
|
||||
assert!(registry.contains("busy"));
|
||||
assert!(!registry.contains("newcomer"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_returning_username_is_replaced_not_refused_at_capacity() {
|
||||
let (registry, _rx) = capped_registry(1);
|
||||
|
||||
let (first, first_stream, _first_far) = loopback_peer("bob");
|
||||
registry
|
||||
.register_peer(first, Some(first_stream), None)
|
||||
.unwrap();
|
||||
|
||||
let protected: HashSet<String> =
|
||||
std::iter::once("bob".to_string()).collect();
|
||||
let (again, stream, _far) = loopback_peer("bob");
|
||||
let result = registry.register_peer_protected(
|
||||
again,
|
||||
Some(stream),
|
||||
None,
|
||||
&protected,
|
||||
);
|
||||
|
||||
assert!(
|
||||
result.is_ok(),
|
||||
"a reconnecting username must replace itself"
|
||||
);
|
||||
assert!(registry.contains("bob"));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn remove_peer_if_respects_actor_identity() {
|
||||
|
||||
@@ -50,10 +50,11 @@ impl Client {
|
||||
) = mpsc::channel();
|
||||
|
||||
let mut ctx = self.context.write_safe()?;
|
||||
let peer_registry = PeerRegistry::new(
|
||||
let peer_registry = PeerRegistry::with_capacity(
|
||||
ctx.actor_system.clone(),
|
||||
sender.clone(),
|
||||
self.username.clone(),
|
||||
ctx.max_peers.clone(),
|
||||
);
|
||||
ctx.peer_registry = Some(peer_registry);
|
||||
|
||||
@@ -175,6 +176,13 @@ impl Client {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn set_max_peers(&self, max_peers: usize) {
|
||||
if let Ok(ctx) = self.context.read_safe() {
|
||||
ctx.max_peers
|
||||
.store(max_peers.max(1), std::sync::atomic::Ordering::Relaxed);
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub fn remove_peer(&self, username: &str) {
|
||||
let context = match self.context.read_safe() {
|
||||
@@ -214,7 +222,10 @@ impl Client {
|
||||
}
|
||||
};
|
||||
if let Some(ref registry) = context.peer_registry {
|
||||
match registry.register_peer(peer_clone, stream, None) {
|
||||
let protected = context.protected_peers();
|
||||
match registry.register_peer_protected(
|
||||
peer_clone, stream, None, &protected,
|
||||
) {
|
||||
Ok(_) => (),
|
||||
Err(e) => {
|
||||
trace!(
|
||||
|
||||
@@ -23,7 +23,7 @@ use std::{
|
||||
net::TcpStream,
|
||||
sync::{
|
||||
Arc, RwLock,
|
||||
atomic::{AtomicBool, AtomicU32, Ordering},
|
||||
atomic::{AtomicBool, AtomicU32, AtomicUsize, Ordering},
|
||||
mpsc::{self, Receiver, Sender},
|
||||
},
|
||||
thread::{self, sleep},
|
||||
@@ -58,6 +58,10 @@ pub const DEFAULT_UPLOAD_SLOTS: usize = 10;
|
||||
/// before giving up and failing the download. Matches the direct-dial timeout.
|
||||
const BROKER_CONNECT_TIMEOUT: Duration = Duration::from_secs(20);
|
||||
|
||||
const DEFAULT_MAX_PEERS: usize = 512;
|
||||
|
||||
const BROWSE_PROTECT_WINDOW: Duration = Duration::from_mins(1);
|
||||
|
||||
/// Source of non-zero correlation tokens for server-brokered connections.
|
||||
static NEXT_CONNECT_TOKEN: AtomicU32 = AtomicU32::new(1);
|
||||
|
||||
@@ -295,7 +299,8 @@ pub struct ClientContext {
|
||||
private_messages: Vec<UserMessage>,
|
||||
/// Correlation tokens for server-brokered (firewalled) connections, mapping
|
||||
/// a token we sent in a ConnectToPeer to the peer we expect back.
|
||||
pending_connect_tokens: HashMap<u32, String>,
|
||||
pending_connect_tokens: HashMap<u32, (String, Instant)>,
|
||||
max_peers: Arc<AtomicUsize>,
|
||||
/// Files we share with peers (read-only after connect).
|
||||
pub shares: Arc<Shares>,
|
||||
/// The directories the current share index was built from.
|
||||
@@ -311,6 +316,7 @@ pub struct ClientContext {
|
||||
pending_serves: HashMap<String, Vec<u32>>,
|
||||
/// Shared-file listings received from peers we browsed.
|
||||
browse_results: HashMap<String, Vec<SharedDirectory>>,
|
||||
pending_browses: HashMap<String, Instant>,
|
||||
/// Latest snapshot of the public chat-room list (from `RoomList`, code 64).
|
||||
room_list: Vec<RoomInfo>,
|
||||
/// Chat-room events awaiting consumption by the client/UI.
|
||||
@@ -395,6 +401,7 @@ impl ClientContext {
|
||||
searches: HashMap::new(),
|
||||
private_messages: Vec::new(),
|
||||
pending_connect_tokens: HashMap::new(),
|
||||
max_peers: Arc::new(AtomicUsize::new(DEFAULT_MAX_PEERS)),
|
||||
shares: Arc::new(Shares::empty()),
|
||||
shared_directories: Vec::new(),
|
||||
peer_addresses: HashMap::new(),
|
||||
@@ -403,6 +410,7 @@ impl ClientContext {
|
||||
active_uploads: HashMap::new(),
|
||||
pending_serves: HashMap::new(),
|
||||
browse_results: HashMap::new(),
|
||||
pending_browses: HashMap::new(),
|
||||
room_list: Vec::new(),
|
||||
room_events: Vec::new(),
|
||||
room_members: HashMap::new(),
|
||||
@@ -563,9 +571,17 @@ impl ClientContext {
|
||||
username: String,
|
||||
directories: Vec<SharedDirectory>,
|
||||
) {
|
||||
self.pending_browses.remove(&username);
|
||||
self.browse_results.insert(username, directories);
|
||||
}
|
||||
|
||||
pub(crate) fn mark_browse_pending(&mut self, username: &str) {
|
||||
let now = Instant::now();
|
||||
self.pending_browses.retain(|_, deadline| *deadline > now);
|
||||
self.pending_browses
|
||||
.insert(username.to_string(), now + BROWSE_PROTECT_WINDOW);
|
||||
}
|
||||
|
||||
/// Remove and return the shared-file listing browsed from `username`.
|
||||
pub fn take_browse_result(
|
||||
&mut self,
|
||||
@@ -577,12 +593,57 @@ impl ClientContext {
|
||||
/// Remember that a server-brokered connection to `username` is pending under
|
||||
/// `token`; the peer will quote it back in a PierceFirewall.
|
||||
pub fn add_pending_connect(&mut self, token: u32, username: String) {
|
||||
self.pending_connect_tokens.insert(token, username);
|
||||
self.pending_connect_tokens
|
||||
.insert(token, (username, Instant::now() + BROKER_CONNECT_TIMEOUT));
|
||||
}
|
||||
|
||||
/// Resolve and consume the peer expected for a brokered connection `token`.
|
||||
pub fn take_pending_connect(&mut self, token: u32) -> Option<String> {
|
||||
self.pending_connect_tokens.remove(&token)
|
||||
self.pending_connect_tokens
|
||||
.remove(&token)
|
||||
.map(|(username, _)| username)
|
||||
}
|
||||
|
||||
pub fn take_expired_connects(&mut self, now: Instant) -> Vec<String> {
|
||||
let mut expired = Vec::new();
|
||||
self.pending_connect_tokens
|
||||
.retain(|_, (username, deadline)| {
|
||||
if now >= *deadline {
|
||||
expired.push(std::mem::take(username));
|
||||
return false;
|
||||
}
|
||||
true
|
||||
});
|
||||
expired
|
||||
}
|
||||
|
||||
pub(crate) fn protected_peers(&self) -> HashSet<String> {
|
||||
let mut protected: HashSet<String> = self
|
||||
.downloads
|
||||
.list()
|
||||
.iter()
|
||||
.filter(|d| {
|
||||
matches!(
|
||||
d.status,
|
||||
DownloadStatus::Queued
|
||||
| DownloadStatus::InProgress { .. }
|
||||
| DownloadStatus::Paused { .. }
|
||||
)
|
||||
})
|
||||
.map(|d| d.username.clone())
|
||||
.collect();
|
||||
protected
|
||||
.extend(self.active_uploads.values().map(|u| u.username.clone()));
|
||||
protected.extend(self.pending_serves.keys().cloned());
|
||||
protected.extend(self.pending_peer_messages.keys().cloned());
|
||||
let now = Instant::now();
|
||||
protected.extend(
|
||||
self.pending_browses
|
||||
.iter()
|
||||
.filter(|(_, deadline)| **deadline > now)
|
||||
.map(|(username, _)| username.clone()),
|
||||
);
|
||||
protected
|
||||
}
|
||||
|
||||
/// Record a private message received from another user.
|
||||
|
||||
@@ -1,12 +1,14 @@
|
||||
use super::{
|
||||
Arc, BROKER_CONNECT_TIMEOUT, Client, ClientContext, ClientOperation,
|
||||
ConnectionType, Download, DownloadPeer, DownloadStatus, Peer, PeerMessage,
|
||||
Arc, Client, ClientContext, ClientOperation, ConnectionType, Download,
|
||||
DownloadPeer, DownloadStatus, Duration, Instant, Peer, PeerMessage,
|
||||
PeerRegistry, Receiver, RwLock, RwLockExt, ServerMessage,
|
||||
build_search_response, debug, error, info, next_connect_token, sleep,
|
||||
build_search_response, debug, error, info, mpsc, next_connect_token,
|
||||
thread, trace, warn,
|
||||
};
|
||||
use crate::message::server::MessageFactory;
|
||||
|
||||
const CONNECT_SWEEP_INTERVAL: Duration = Duration::from_secs(1);
|
||||
|
||||
impl Client {
|
||||
pub(crate) fn listen_to_client_operations(
|
||||
reader: Receiver<ClientOperation>,
|
||||
@@ -14,9 +16,21 @@ impl Client {
|
||||
own_username: String,
|
||||
) {
|
||||
thread::spawn(move || {
|
||||
for operation in reader {
|
||||
let mut last_sweep = Instant::now();
|
||||
loop {
|
||||
let next = reader.recv_timeout(CONNECT_SWEEP_INTERVAL);
|
||||
if last_sweep.elapsed() >= CONNECT_SWEEP_INTERVAL {
|
||||
last_sweep = Instant::now();
|
||||
Self::sweep_expired_connects(&client_context);
|
||||
}
|
||||
let operation = match next {
|
||||
Ok(operation) => operation,
|
||||
Err(mpsc::RecvTimeoutError::Timeout) => continue,
|
||||
Err(mpsc::RecvTimeoutError::Disconnected) => break,
|
||||
};
|
||||
match operation {
|
||||
ClientOperation::ConnectToPeer(peer) => {
|
||||
if matches!(peer.connection_type, ConnectionType::F) {
|
||||
let client_context_clone = client_context.clone();
|
||||
let own_username_clone = own_username.clone();
|
||||
|
||||
@@ -28,6 +42,14 @@ impl Client {
|
||||
None,
|
||||
);
|
||||
});
|
||||
} else {
|
||||
Self::connect_to_peer(
|
||||
peer,
|
||||
client_context.clone(),
|
||||
own_username.clone(),
|
||||
None,
|
||||
);
|
||||
}
|
||||
}
|
||||
ClientOperation::SearchResult(search_result) => {
|
||||
trace!("[client] SearchResult {:?}", search_result);
|
||||
@@ -122,6 +144,28 @@ impl Client {
|
||||
);
|
||||
continue;
|
||||
};
|
||||
if matches!(
|
||||
download.status,
|
||||
DownloadStatus::InProgress { .. }
|
||||
| DownloadStatus::Completed
|
||||
) {
|
||||
debug!(
|
||||
"[client] transfer token {} already claimed; \
|
||||
ignoring replayed TransferResponse",
|
||||
token
|
||||
);
|
||||
continue;
|
||||
}
|
||||
if let Ok(mut ctx) = client_context.write_safe() {
|
||||
ctx.update_download_with_status(
|
||||
token,
|
||||
DownloadStatus::InProgress {
|
||||
bytes_downloaded: 0,
|
||||
total_bytes: download.size,
|
||||
speed_bytes_per_sec: 0.0,
|
||||
},
|
||||
);
|
||||
}
|
||||
|
||||
thread::spawn(move || {
|
||||
let download_peer = DownloadPeer::new(
|
||||
@@ -285,17 +329,12 @@ impl Client {
|
||||
u8::try_from(obfuscation_type).unwrap_or(0),
|
||||
obfuscated_port,
|
||||
);
|
||||
let client_context_clone = client_context.clone();
|
||||
let own_username_clone = own_username.clone();
|
||||
|
||||
thread::spawn(move || {
|
||||
Self::connect_to_peer(
|
||||
peer,
|
||||
client_context_clone,
|
||||
own_username_clone,
|
||||
client_context.clone(),
|
||||
own_username.clone(),
|
||||
None,
|
||||
);
|
||||
});
|
||||
}
|
||||
}
|
||||
ClientOperation::UpdateDownloadTokens(
|
||||
@@ -751,30 +790,19 @@ impl Client {
|
||||
ConnectionType::P,
|
||||
);
|
||||
let _ = sender.send(ServerMessage::SendMessage(msg));
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Bound the brokered attempt: if no PierceFirewall
|
||||
// consumes the token, fail the peer's queued
|
||||
// downloads (so the caller's Receiver unblocks)
|
||||
// and reclaim the token. A successful pierce
|
||||
// takes the token first, making this a no-op.
|
||||
let timeout_ctx = client_context.clone();
|
||||
let timeout_user = username.clone();
|
||||
thread::spawn(move || {
|
||||
sleep(BROKER_CONNECT_TIMEOUT);
|
||||
let still_pending =
|
||||
timeout_ctx.write_safe().is_ok_and(|mut c| {
|
||||
c.take_pending_connect(token).is_some()
|
||||
});
|
||||
if still_pending {
|
||||
Self::fail_queued_downloads(
|
||||
&timeout_ctx,
|
||||
&timeout_user,
|
||||
);
|
||||
fn sweep_expired_connects(client_context: &Arc<RwLock<ClientContext>>) {
|
||||
let expired = match client_context.write_safe() {
|
||||
Ok(mut ctx) => ctx.take_expired_connects(Instant::now()),
|
||||
Err(_) => return,
|
||||
};
|
||||
for username in expired {
|
||||
Self::fail_queued_downloads(client_context, &username);
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
@@ -177,7 +177,8 @@ impl Client {
|
||||
let request =
|
||||
crate::message::server::MessageFactory::build_get_share_file_list();
|
||||
let (connected, registry) = {
|
||||
let ctx = self.context.read_safe()?;
|
||||
let mut ctx = self.context.write_safe()?;
|
||||
ctx.mark_browse_pending(username);
|
||||
(
|
||||
ctx.peer_registry
|
||||
.as_ref()
|
||||
|
||||
@@ -395,3 +395,51 @@ fn a_clean_disconnect_keeps_queued_uploads_an_error_drops_them() {
|
||||
thread::sleep(Duration::from_millis(10));
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn a_pending_connect_expires_after_the_broker_timeout() {
|
||||
let mut context = ClientContext::new();
|
||||
context.add_pending_connect(7, "ghost".to_string());
|
||||
|
||||
assert!(context.take_expired_connects(Instant::now()).is_empty());
|
||||
|
||||
let after_deadline =
|
||||
Instant::now() + BROKER_CONNECT_TIMEOUT + Duration::from_secs(1);
|
||||
assert_eq!(
|
||||
context.take_expired_connects(after_deadline),
|
||||
vec!["ghost".to_string()]
|
||||
);
|
||||
assert!(context.take_pending_connect(7).is_none());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn protected_peers_covers_downloads_uploads_and_pending_serves() {
|
||||
let mut context = ClientContext::new();
|
||||
context.add_download(download(
|
||||
"downloader",
|
||||
"song.mp3",
|
||||
1,
|
||||
DownloadStatus::Queued,
|
||||
mpsc::channel().0,
|
||||
));
|
||||
context.add_download(download(
|
||||
"done",
|
||||
"old.mp3",
|
||||
2,
|
||||
DownloadStatus::Completed,
|
||||
mpsc::channel().0,
|
||||
));
|
||||
context
|
||||
.pending_serves
|
||||
.insert("waiting".to_string(), vec![9]);
|
||||
context.mark_browse_pending("browsed");
|
||||
|
||||
let protected = context.protected_peers();
|
||||
assert!(protected.contains("downloader"));
|
||||
assert!(protected.contains("waiting"));
|
||||
assert!(protected.contains("browsed"));
|
||||
assert!(!protected.contains("done"));
|
||||
|
||||
context.store_browse_result("browsed".to_string(), Vec::new());
|
||||
assert!(!context.protected_peers().contains("browsed"));
|
||||
}
|
||||
|
||||
@@ -11,6 +11,7 @@ use crate::message::{Message, MessageReader};
|
||||
use crate::peer::{ConnectionType, DownloadPeer, Peer};
|
||||
use crate::types::Download;
|
||||
use crate::utils::lock::RwLockExt;
|
||||
use crate::utils::semaphore::{Permit, Semaphore};
|
||||
use crate::{DownloadStatus, debug, error, info, trace};
|
||||
|
||||
/// How long to wait before accepting again after a failure.
|
||||
@@ -25,6 +26,8 @@ const PEER_INIT_MESSAGE_CODE: u8 = 1;
|
||||
/// How long an accepted peer gets to send its peer-init handshake.
|
||||
const PEER_INIT_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
|
||||
const MAX_HANDSHAKES: usize = 128;
|
||||
|
||||
#[derive(Clone)]
|
||||
struct ConnectionContext {
|
||||
client_sender: Sender<ClientOperation>,
|
||||
@@ -140,7 +143,13 @@ fn handle_peer_connection(
|
||||
}
|
||||
};
|
||||
if let Some(ref registry) = client_context.peer_registry {
|
||||
match registry.register_peer(peer.clone(), Some(stream), Some(reader)) {
|
||||
let protected = client_context.protected_peers();
|
||||
match registry.register_peer_protected(
|
||||
peer.clone(),
|
||||
Some(stream),
|
||||
Some(reader),
|
||||
&protected,
|
||||
) {
|
||||
Ok(_) => (),
|
||||
Err(e) => {
|
||||
error!(
|
||||
@@ -283,6 +292,7 @@ fn handle_pierce_firewall(
|
||||
fn handle_incoming_connection(
|
||||
mut stream: TcpStream,
|
||||
context: ConnectionContext,
|
||||
permit: Permit,
|
||||
) {
|
||||
let Ok(peer_addr) = stream.peer_addr() else {
|
||||
error!("[listener] failed to get peer address");
|
||||
@@ -303,6 +313,7 @@ fn handle_incoming_connection(
|
||||
return;
|
||||
}
|
||||
};
|
||||
drop(permit);
|
||||
|
||||
// A firewalled peer brokered through the server connects back with a
|
||||
// PierceFirewall (code 0) instead of a PeerInit (code 1).
|
||||
@@ -406,6 +417,7 @@ impl Listen {
|
||||
client_context,
|
||||
own_username,
|
||||
};
|
||||
let handshakes = Arc::new(Semaphore::new(MAX_HANDSHAKES));
|
||||
|
||||
for stream in listener.incoming() {
|
||||
let Ok(stream) = stream else {
|
||||
@@ -421,11 +433,14 @@ impl Listen {
|
||||
continue;
|
||||
};
|
||||
|
||||
let permit = handshakes.acquire();
|
||||
let context = context.clone();
|
||||
// One thread per connection: the peer-init handshake blocks, and a
|
||||
// peer that is slow to send one must not stop us accepting anybody
|
||||
// else — a wedged accept loop makes us unreachable to every peer.
|
||||
thread::spawn(move || handle_incoming_connection(stream, context));
|
||||
thread::spawn(move || {
|
||||
handle_incoming_connection(stream, context, permit);
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,6 +3,7 @@ pub mod logger;
|
||||
pub mod lock;
|
||||
pub mod md5;
|
||||
pub mod path;
|
||||
pub(crate) mod semaphore;
|
||||
pub mod zlib;
|
||||
|
||||
// Re-export commonly used items
|
||||
|
||||
@@ -0,0 +1,73 @@
|
||||
use std::sync::{Arc, Condvar, Mutex, PoisonError};
|
||||
|
||||
pub struct Semaphore {
|
||||
permits: Mutex<usize>,
|
||||
available: Condvar,
|
||||
}
|
||||
|
||||
pub struct Permit(Arc<Semaphore>);
|
||||
|
||||
impl Semaphore {
|
||||
#[must_use]
|
||||
pub const fn new(permits: usize) -> Self {
|
||||
Self {
|
||||
permits: Mutex::new(permits),
|
||||
available: Condvar::new(),
|
||||
}
|
||||
}
|
||||
|
||||
#[must_use]
|
||||
pub fn acquire(self: &Arc<Self>) -> Permit {
|
||||
let mut permits =
|
||||
self.permits.lock().unwrap_or_else(PoisonError::into_inner);
|
||||
while *permits == 0 {
|
||||
permits = self
|
||||
.available
|
||||
.wait(permits)
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
}
|
||||
*permits -= 1;
|
||||
Permit(Arc::clone(self))
|
||||
}
|
||||
}
|
||||
|
||||
impl Drop for Permit {
|
||||
fn drop(&mut self) {
|
||||
let mut permits = self
|
||||
.0
|
||||
.permits
|
||||
.lock()
|
||||
.unwrap_or_else(PoisonError::into_inner);
|
||||
*permits += 1;
|
||||
self.0.available.notify_one();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use std::time::Duration;
|
||||
|
||||
#[test]
|
||||
fn permits_bound_concurrent_holders() {
|
||||
let semaphore = Arc::new(Semaphore::new(2));
|
||||
let first = semaphore.acquire();
|
||||
let _second = semaphore.acquire();
|
||||
|
||||
let (tx, rx) = std::sync::mpsc::channel();
|
||||
std::thread::spawn(move || {
|
||||
let _third = semaphore.acquire();
|
||||
tx.send(()).unwrap();
|
||||
});
|
||||
|
||||
assert!(
|
||||
rx.recv_timeout(Duration::from_millis(100)).is_err(),
|
||||
"a third acquire must block while both permits are held"
|
||||
);
|
||||
drop(first);
|
||||
assert!(
|
||||
rx.recv_timeout(Duration::from_secs(2)).is_ok(),
|
||||
"releasing a permit must wake the waiter"
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user