dumbgoy/soulseek-rs-modded
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Merge pull request #37 from michel/issue-32-33-thread-flood

Browse Source 
fix(client): bound the threads a peer flood can create
This commit is contained in:
Michel de Graaf
2026-08-13 14:30:43 +02:00
committed by GitHub
parent fd1d168ca6 492b961766
commit 0aec53b40a
10 changed files with 717 additions and 129 deletions
Download Patch File Download Diff File Expand all files Collapse all files
soulseek-rs-lib/src/actor/peer_registry.rs
+169 -9
View File
@@ -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() {
soulseek-rs-lib/src/client/connection.rs
+37 -18
View File
@@ -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() {
@@ -206,25 +214,36 @@ impl Client {
ConnectionType::P => {
let username = peer.username;
let context = match client_context.read_safe() {
Ok(c) => c,
Err(e) => {
error!("[client] connect_to_peer read: {}", e);
return;
let refused = {
let context = match client_context.read_safe() {
Ok(c) => c,
Err(e) => {
error!("[client] connect_to_peer read: {}", e);
return;
}
};
if let Some(ref registry) = context.peer_registry {
let protected = context.protected_peers();
match registry.register_peer_protected(
peer_clone, stream, None, &protected,
) {
Ok(_) => false,
Err(e) => {
warn!(
"[client] refusing peer connection for \
{:?}: {:?}",
username, e
);
true
}
}
} else {
trace!("PeerRegistry not initialized");
false
}
};
if let Some(ref registry) = context.peer_registry {
match registry.register_peer(peer_clone, stream, None) {
Ok(_) => (),
Err(e) => {
trace!(
"Failed to spawn peer actor for {:?}: {:?}",
username, e
);
}
}
} else {
trace!("PeerRegistry not initialized");
if refused {
Self::fail_queued_downloads(&client_context, &username);
}
}
soulseek-rs-lib/src/client/mod.rs
+68 -4
View File
@@ -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(5);
/// 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,60 @@ 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.uploads.values().map(|j| j.downloader.clone()));
protected
.extend(self.upload_queue.iter().map(|q| q.downloader.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.
soulseek-rs-lib/src/client/operations.rs
+70 -42
View File
@@ -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,20 +16,40 @@ 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) => {
let client_context_clone = client_context.clone();
let own_username_clone = own_username.clone();
if matches!(peer.connection_type, ConnectionType::F) {
let client_context_clone = client_context.clone();
let own_username_clone = own_username.clone();
thread::spawn(move || {
thread::spawn(move || {
Self::connect_to_peer(
peer,
client_context_clone,
own_username_clone,
None,
);
});
} else {
Self::connect_to_peer(
peer,
client_context_clone,
own_username_clone,
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,
None,
);
});
Self::connect_to_peer(
peer,
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);
}
}
}
soulseek-rs-lib/src/client/rooms.rs
+2 -1
View File
@@ -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()
soulseek-rs-lib/src/client/tests.rs
+143
View File
@@ -395,3 +395,146 @@ 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"));
}
#[test]
fn an_expired_broker_connect_fails_the_queued_downloads() {
let client = Client::new("u", "p");
let (sender, receiver) = mpsc::channel();
{
let mut ctx = client.context.write().unwrap();
ctx.add_download(download(
"ghost",
"f.mp3",
7,
DownloadStatus::Queued,
sender,
));
ctx.pending_connect_tokens
.insert(7, ("ghost".to_string(), Instant::now()));
}
let (_ops_tx, ops_rx) = mpsc::channel();
Client::listen_to_client_operations(
ops_rx,
client.context,
"u".to_string(),
);
let status = receiver.recv_timeout(Duration::from_secs(5));
assert!(
matches!(status, Ok(DownloadStatus::Failed(_))),
"the sweep must fail the queued download, got {status:?}"
);
}
#[test]
fn a_replayed_transfer_response_does_not_start_a_second_transfer() {
let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap();
listener.set_nonblocking(true).unwrap();
let port = u32::from(listener.local_addr().unwrap().port());
let client = Client::new("u", "p");
let (sender, _receiver) = mpsc::channel();
client.context.write().unwrap().add_download(download(
"peer",
"f.mp3",
9,
DownloadStatus::Queued,
sender,
));
let (ops_tx, ops_rx) = mpsc::channel();
Client::listen_to_client_operations(
ops_rx,
client.context,
"u".to_string(),
);
let peer = Peer::new(
"peer".to_string(),
ConnectionType::F,
"127.0.0.1".to_string(),
port,
None,
0,
0,
0,
);
ops_tx
.send(ClientOperation::DownloadFromPeer(9, peer.clone(), true))
.unwrap();
ops_tx
.send(ClientOperation::DownloadFromPeer(9, peer, true))
.unwrap();
let deadline = Instant::now() + Duration::from_secs(5);
let mut connections = 0;
while Instant::now() < deadline {
match listener.accept() {
Ok(_) => {
connections += 1;
if connections == 1 {
thread::sleep(Duration::from_millis(500));
}
}
Err(ref e) if e.kind() == std::io::ErrorKind::WouldBlock => {
if connections >= 1 {
break;
}
thread::sleep(Duration::from_millis(20));
}
Err(_) => break,
}
}
assert_eq!(
connections, 1,
"a replayed TransferResponse must not dial the peer again"
);
}
soulseek-rs-lib/src/peer/download_peer.rs
-8
View File
@@ -484,14 +484,6 @@ impl DownloadPeer {
self.no_pierce
);
if let Some(ref dl) = download {
let _ = dl.sender.send(DownloadStatus::Queued);
client_context
.write()
.map_err(|_| DownloadError::LockPoisoned)?
.update_download_with_status(dl.token, DownloadStatus::Queued);
}
let mut stream = match stream {
Some(s) => {
Self::apply_transfer_deadlines(&s)?;
soulseek-rs-lib/src/peer/listen.rs
+154 -47
View File
@@ -3,15 +3,16 @@ use std::net::{TcpListener, TcpStream};
use std::sync::mpsc::Sender;
use std::sync::{Arc, RwLock};
use std::thread;
use std::time::Duration;
use std::time::{Duration, Instant};
use crate::client::{ClientContext, ClientOperation};
use crate::client::{Client, ClientContext, ClientOperation};
use crate::message::{Message, MessageReader};
use crate::peer::{ConnectionType, DownloadPeer, Peer};
use crate::types::Download;
use crate::utils::lock::RwLockExt;
use crate::{DownloadStatus, debug, error, info, trace};
use crate::utils::semaphore::{Permit, Semaphore};
use crate::{DownloadStatus, debug, error, info, trace, warn};
/// How long to wait before accepting again after a failure.
///
@@ -25,6 +26,10 @@ 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;
const MAX_PEER_INIT_BYTES: usize = 4096;
#[derive(Clone)]
struct ConnectionContext {
client_sender: Sender<ClientOperation>,
@@ -41,16 +46,30 @@ struct PeerInitData {
fn read_peer_init_message(
stream: &mut TcpStream,
reader: &mut MessageReader,
deadline: Instant,
) -> io::Result<Message> {
// An untrusted peer gets a bounded handshake. Without this a peer that
// connects and stays silent parks this read forever, pinning a thread that
// owes us a peer init.
stream.set_read_timeout(Some(PEER_INIT_TIMEOUT))?;
let message = loop {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Err(io::Error::new(
io::ErrorKind::TimedOut,
"peer init deadline passed",
));
}
stream.set_read_timeout(Some(remaining))?;
reader.read_from_socket(stream)?;
if let Ok(Some(msg)) = reader.extract_message() {
break msg;
}
if reader.buffer_len() > MAX_PEER_INIT_BYTES {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
"peer init larger than any handshake",
));
}
};
// Whatever this connection turns into reads on its own terms from here.
@@ -119,7 +138,7 @@ fn handle_peer_connection(
stream: TcpStream,
reader: MessageReader,
context: &ConnectionContext,
) {
) -> bool {
// The peer actor multiplexes socket reads with its mailbox on a single
// thread: `tick()` reads the socket, but outgoing messages (e.g. a queued
// QueueUpload for a download) are delivered through the mailbox between
@@ -128,30 +147,44 @@ fn handle_peer_connection(
// outbound path and drive this connection non-blocking.
if let Err(e) = stream.set_nonblocking(true) {
error!("[listener] failed to set peer stream non-blocking: {}", e);
return;
return false;
}
stream.set_nodelay(true).ok();
let client_context = match context.client_context.read_safe() {
Ok(c) => c,
Err(e) => {
error!("[listener] handle_peer_connection lock: {}", e);
return;
let registered = {
let client_context = match context.client_context.read_safe() {
Ok(c) => c,
Err(e) => {
error!("[listener] handle_peer_connection lock: {}", e);
return false;
}
};
if let Some(ref registry) = client_context.peer_registry {
let protected = client_context.protected_peers();
match registry.register_peer_protected(
peer.clone(),
Some(stream),
Some(reader),
&protected,
) {
Ok(_) => true,
Err(e) => {
warn!(
"[listener] refusing peer connection for {:?}: {:?}",
peer.username, e
);
false
}
}
} else {
error!("PeerRegistry not initialized");
false
}
};
if let Some(ref registry) = client_context.peer_registry {
match registry.register_peer(peer.clone(), Some(stream), Some(reader)) {
Ok(_) => (),
Err(e) => {
error!(
"Failed to spawn peer actor for {:?}: {:?}",
peer.username, e
);
}
}
} else {
error!("PeerRegistry not initialized");
if !registered {
Client::fail_queued_downloads(&context.client_context, &peer.username);
}
registered
}
fn handle_file_connection(
@@ -271,18 +304,19 @@ fn handle_pierce_firewall(
0,
0,
);
handle_peer_connection(peer, stream, reader, context);
// Inbound peers don't self-announce, so nudge the client to flush any
// downloads queued for this now-connected peer.
let _ = context
.client_sender
.send(ClientOperation::PeerConnected(username));
if handle_peer_connection(peer, stream, reader, context) {
// Inbound peers don't self-announce, so nudge the client to flush any
// downloads queued for this now-connected peer.
let _ = context
.client_sender
.send(ClientOperation::PeerConnected(username));
}
}
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");
@@ -294,15 +328,18 @@ fn handle_incoming_connection(
let mut reader = MessageReader::new();
// A peer that dials and then goes away is routine, not an error.
let message = match read_peer_init_message(&mut stream, &mut reader) {
Ok(message) => message,
Err(e) => {
debug!(
"[listener:{peer_ip}:{peer_port}] no peer init message: {e}"
);
return;
}
};
let deadline = Instant::now() + PEER_INIT_TIMEOUT;
let message =
match read_peer_init_message(&mut stream, &mut reader, deadline) {
Ok(message) => message,
Err(e) => {
debug!(
"[listener:{peer_ip}:{peer_port}] no peer init message: {e}"
);
return;
}
};
drop(permit);
// A firewalled peer brokered through the server connects back with a
// PierceFirewall (code 0) instead of a PeerInit (code 1).
@@ -343,13 +380,14 @@ fn handle_incoming_connection(
match init_data.connection_type {
ConnectionType::P => {
handle_peer_connection(peer, stream, reader, &context);
// Inbound peers don't self-announce, so nudge the client to flush
// anything queued for this now-connected peer — same as the
// PierceFirewall path.
let _ = context
.client_sender
.send(ClientOperation::PeerConnected(init_data.username));
if handle_peer_connection(peer, stream, reader, &context) {
// Inbound peers don't self-announce, so nudge the client to
// flush anything queued for this now-connected peer — same as
// the PierceFirewall path.
let _ = context
.client_sender
.send(ClientOperation::PeerConnected(init_data.username));
}
}
ConnectionType::F => handle_file_connection(
@@ -406,6 +444,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 +460,79 @@ 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);
});
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use std::io::Write;
use std::net::TcpListener;
#[test]
fn a_dribbling_handshake_hits_the_total_deadline() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let writer = std::thread::spawn(move || {
let mut socket = TcpStream::connect(addr).unwrap();
let _ = socket.write_all(&1000u32.to_le_bytes());
for _ in 0..20 {
if socket.write_all(&[0u8]).is_err() {
return;
}
std::thread::sleep(Duration::from_millis(50));
}
});
let (mut stream, _) = listener.accept().unwrap();
let mut reader = MessageReader::new();
let started = Instant::now();
let result = read_peer_init_message(
&mut stream,
&mut reader,
started + Duration::from_millis(300),
);
assert!(result.is_err(), "a dribbled handshake must not succeed");
assert!(
started.elapsed() < Duration::from_secs(5),
"the deadline must bound the whole handshake, not each read"
);
let _ = writer.join();
}
#[test]
fn an_oversized_handshake_is_rejected() {
let listener = TcpListener::bind("127.0.0.1:0").unwrap();
let addr = listener.local_addr().unwrap();
let writer = std::thread::spawn(move || {
let mut socket = TcpStream::connect(addr).unwrap();
let junk = vec![0xFFu8; MAX_PEER_INIT_BYTES * 2];
let _ = socket.write_all(&junk);
std::thread::sleep(Duration::from_millis(500));
});
let (mut stream, _) = listener.accept().unwrap();
let mut reader = MessageReader::new();
let result = read_peer_init_message(
&mut stream,
&mut reader,
Instant::now() + Duration::from_secs(5),
);
assert!(
result.is_err(),
"an oversized handshake must be rejected, got {result:?}"
);
let _ = writer.join();
}
}
soulseek-rs-lib/src/utils/mod.rs
+1
View File
@@ -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
soulseek-rs-lib/src/utils/semaphore.rs
+73
View File
@@ -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"
);
}
}