fix(peer): survive download token races on the listener path

A P-connection drop failed every download for its user, including ones
actively transferring on their own healthy F connection:
collect_failed_tokens now only fails queued downloads. And an F
connection could look its token up before the client-operations thread
had registered it: the listener now waits for the registration, and
claims it only for the matching peer and only when not already claimed.

Fixes #42
This commit is contained in:
Michel de Graaf
2026-08-17 10:23:21 +02:00
parent 1833717521
commit f9becf9680
3 changed files with 289 additions and 67 deletions
+58 -2
View File
@@ -142,8 +142,8 @@ impl DownloadStore {
}
}
/// Returns the tokens of downloads matching `username` (and optionally a
/// `filename`) after notifying their senders of `Failed`.
/// Returns the tokens of queued downloads matching `username` (and optionally
/// a `filename`) after notifying their senders of `Failed`.
///
/// Caller is responsible for then calling `update_status` and `remove` for
/// each token, typically under a write lock.
@@ -159,6 +159,7 @@ pub fn collect_failed_tokens(
.filter(|d| {
d.username == username && filename.is_none_or(|f| d.filename == *f)
})
.filter(|d| matches!(d.status, DownloadStatus::Queued))
.map(|d| {
let _ = d.sender.send(DownloadStatus::Failed(Some(
"The upload failed on the other side".to_string(),
@@ -305,6 +306,61 @@ mod tests {
assert!(!store.remove_by_file("peer", "song.mp3"), "idempotent");
}
#[test]
fn collect_failed_tokens_collects_only_queued_downloads() {
let mut store = DownloadStore::new();
let (tx_active, rx_active) = mpsc::channel();
let (tx_paused, rx_paused) = mpsc::channel();
let (tx_completed, rx_completed) = mpsc::channel();
let (tx_queued, _rx_queued) = mpsc::channel();
let mut active = make_download(
1,
DownloadStatus::InProgress {
bytes_downloaded: 25,
total_bytes: 100,
speed_bytes_per_sec: 10.0,
},
);
active.sender = tx_active;
let mut paused = make_download(
2,
DownloadStatus::Paused {
bytes_downloaded: 25,
total_bytes: 100,
},
);
paused.sender = tx_paused;
let mut completed = make_download(3, DownloadStatus::Completed);
completed.sender = tx_completed;
let mut queued = make_download(4, DownloadStatus::Queued);
queued.sender = tx_queued;
store.add(active);
store.add(paused);
store.add(completed);
store.add(queued);
let tokens = collect_failed_tokens(&store, "peer", None);
assert_eq!(tokens, vec![4], "only the queued download may fail");
assert!(
rx_active.try_recv().is_err(),
"an in-progress transfer must not be notified of failure"
);
assert!(
rx_paused.try_recv().is_err(),
"a paused transfer must not be notified of failure"
);
assert!(
rx_completed.try_recv().is_err(),
"a completed download must not be notified of failure"
);
}
#[test]
fn collect_failed_tokens_notifies_and_lists_matching() {
let mut store = DownloadStore::new();
+161 -38
View File
@@ -26,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 TOKEN_REGISTRATION_TIMEOUT: Duration = Duration::from_secs(10);
const MAX_HANDSHAKES: usize = 128;
const MAX_PEER_INIT_BYTES: usize = 4096;
@@ -116,41 +118,52 @@ fn wait_for_transfer_token(
Ok(())
}
fn extract_download_from_buffer(
reader: &mut MessageReader,
client_context: &Arc<RwLock<ClientContext>>,
username: &str,
peer_ip: &str,
peer_port: u32,
) -> Option<Download> {
if reader.buffer_len() == 0 {
return None;
}
fn take_transfer_token(reader: &mut MessageReader) -> Option<u32> {
let buffer = reader.get_buffer();
let token = u32::from_le_bytes(buffer.get(0..4)?.try_into().ok()?);
trace!(
"[listener:{}] got transfer_token: {} from data chunk",
username, token
);
Some(u32::from_le_bytes(buffer.get(0..4)?.try_into().ok()?))
}
let context = match client_context.read_safe() {
Ok(c) => c,
Err(e) => {
error!("[listener] client context lock: {}", e);
fn wait_for_download_registration(
client_context: &Arc<RwLock<ClientContext>>,
token: u32,
username: &str,
deadline: Instant,
) -> Option<Download> {
loop {
let download = match client_context.read_safe() {
Ok(context) => context
.get_download_by_token(token)
.filter(|d| {
d.username == username
&& !matches!(
d.status,
DownloadStatus::InProgress { .. }
| DownloadStatus::Paused { .. }
| DownloadStatus::Completed
)
})
.cloned(),
Err(e) => {
error!("[listener] client context lock: {}", e);
return None;
}
};
if download.is_some() {
return download;
}
if Instant::now() >= deadline {
if let Ok(context) = client_context.read_safe() {
trace!(
"[listener:{}] no claimable download for token {}, known tokens: {:?}",
username,
token,
context.get_download_tokens()
);
}
return None;
}
};
let download = context.get_download_by_token(token).cloned();
if download.is_none() {
let download_tokens = context.get_download_tokens();
trace!(
"[listener:{peer_ip}:{peer_port}] download token not found: {:?}, download tokens: {:?}",
token, download_tokens
);
thread::sleep(Duration::from_millis(50));
}
download
}
fn handle_peer_connection(
@@ -201,13 +214,15 @@ fn handle_file_connection(
return;
}
let download = extract_download_from_buffer(
&mut reader,
&context.client_context,
&peer.username,
&peer.host,
peer.port,
);
let download =
take_transfer_token(&mut reader).and_then(|transfer_token| {
wait_for_download_registration(
&context.client_context,
transfer_token,
&peer.username,
Instant::now() + TOKEN_REGISTRATION_TIMEOUT,
)
});
let failure_token = download.as_ref().map(|d| d.token);
let download_peer = DownloadPeer::new(
@@ -474,8 +489,116 @@ impl Listen {
#[cfg(test)]
mod tests {
use super::*;
use crate::types::DownloadMetadata;
use std::io::Write;
use std::net::TcpListener;
use std::sync::mpsc;
fn registered_download(token: u32) -> Download {
Download {
username: "peer".to_string(),
filename: "song.mp3".to_string(),
token,
size: 100,
download_directory: std::env::temp_dir().display().to_string(),
status: DownloadStatus::Queued,
sender: mpsc::channel().0,
queue_position: None,
metadata: DownloadMetadata::default(),
}
}
#[test]
fn a_download_registered_during_the_wait_is_found() {
let client_context = Arc::new(RwLock::new(ClientContext::new()));
let registering = {
let client_context = client_context.clone();
std::thread::spawn(move || {
std::thread::sleep(Duration::from_millis(100));
client_context
.write()
.unwrap()
.add_download(registered_download(7));
})
};
let download = wait_for_download_registration(
&client_context,
7,
"peer",
Instant::now() + Duration::from_secs(5),
);
assert!(
download.is_some(),
"a token registered during the wait must be found"
);
let _ = registering.join();
}
#[test]
fn an_unregistered_token_gives_up_at_the_deadline() {
let client_context = Arc::new(RwLock::new(ClientContext::new()));
let started = Instant::now();
let download = wait_for_download_registration(
&client_context,
7,
"peer",
started + Duration::from_millis(200),
);
assert!(download.is_none(), "an unknown token must not be invented");
assert!(
started.elapsed() < Duration::from_secs(2),
"the deadline must bound the wait"
);
}
#[test]
fn a_download_for_another_peer_is_not_claimed() {
let client_context = Arc::new(RwLock::new(ClientContext::new()));
client_context
.write()
.unwrap()
.add_download(registered_download(7));
let download = wait_for_download_registration(
&client_context,
7,
"impostor",
Instant::now() + Duration::from_millis(200),
);
assert!(
download.is_none(),
"a token must only match the peer it was registered for"
);
}
#[test]
fn an_already_claimed_download_is_not_reclaimed() {
let client_context = Arc::new(RwLock::new(ClientContext::new()));
let mut claimed = registered_download(7);
claimed.status = DownloadStatus::InProgress {
bytes_downloaded: 25,
total_bytes: 100,
speed_bytes_per_sec: 10.0,
};
client_context.write().unwrap().add_download(claimed);
let download = wait_for_download_registration(
&client_context,
7,
"peer",
Instant::now() + Duration::from_millis(200),
);
assert!(
download.is_none(),
"a transfer already in progress must not be claimed again"
);
}
#[test]
fn a_dribbling_handshake_hits_the_total_deadline() {
@@ -557,7 +680,7 @@ mod tests {
assert!(result.is_err(), "a silent peer must not park the thread");
assert!(
started.elapsed() < Duration::from_secs(5),
started.elapsed() < Duration::from_secs(2),
"the deadline must bound the wait"
);
let _ = writer.join();
+70 -27
View File
@@ -520,36 +520,43 @@ fn connect_retry(addr: &str, timeout: Duration) -> std::io::Result<TcpStream> {
}
}
fn run_mock_uploader(cfg: &MockUpload) -> std::io::Result<u64> {
// 1. P (control) connection: register ourselves with the downloader.
fn open_control_and_await_queue(
cfg: &MockUpload,
) -> std::io::Result<TcpStream> {
let mut p = connect_retry(&cfg.listen_addr, Duration::from_secs(5))?;
p.set_read_timeout(Some(Duration::from_secs(10)))?;
p.write_all(&peer_init_bytes(&cfg.peer_username, "P", 0))?;
p.flush()?;
let _ = cfg.ready.send(());
// 2. Wait for the downloader's QueueUpload (peer code 43).
loop {
let mut msg = read_framed(&mut p)?;
let msg = read_framed(&mut p)?;
if msg.get_message_code() == 43 {
msg.set_pointer(8);
let _requested = msg.read_string();
break;
}
}
Ok(p)
}
// 3. Offer the transfer with a TransferRequest (peer code 40). The size here
// becomes the download's expected size, so it must match the content.
/// Build a TransferRequest (peer code 40). The size becomes the download's
/// expected size, so it must match the content.
fn transfer_request(cfg: &MockUpload) -> Message {
let mut tr = Message::new();
tr.write_int32(40)
.write_int32(1) // direction: upload
.write_int32(cfg.token)
.write_string(&cfg.filename)
.write_int64(cfg.content.len() as u64);
p.write_all(&tr.get_buffer())?;
tr
}
fn run_mock_uploader(cfg: &MockUpload) -> std::io::Result<u64> {
let mut p = open_control_and_await_queue(cfg)?;
p.write_all(&transfer_request(cfg).get_buffer())?;
p.flush()?;
// 4. Wait for the downloader to allow it (TransferResponse, peer code 41).
// Wait for the downloader to allow it (TransferResponse, peer code 41).
loop {
let msg = read_framed(&mut p)?;
if msg.get_message_code() == 41 {
@@ -557,7 +564,6 @@ fn run_mock_uploader(cfg: &MockUpload) -> std::io::Result<u64> {
}
}
// 5. Open the F (file) connection to the downloader's listener and stream.
serve_file_over_f(
&cfg.listen_addr,
&cfg.peer_username,
@@ -567,22 +573,38 @@ fn run_mock_uploader(cfg: &MockUpload) -> std::io::Result<u64> {
)
}
fn run_mock_uploader_f_first(cfg: &MockUpload) -> std::io::Result<u64> {
let mut p = open_control_and_await_queue(cfg)?;
let listen_addr = cfg.listen_addr.clone();
let peer_username = cfg.peer_username.clone();
let token = cfg.token;
let content = cfg.content.clone();
let token_delay = cfg.token_delay;
let f = std::thread::spawn(move || {
serve_file_over_f(
&listen_addr,
&peer_username,
token,
&content,
token_delay,
)
});
std::thread::sleep(Duration::from_millis(300));
p.write_all(&transfer_request(cfg).get_buffer())?;
p.flush()?;
f.join()
.map_err(|_| std::io::Error::other("F connection thread panicked"))?
}
fn run_refusing_uploader(
cfg: &MockUpload,
refusal: &Message,
) -> std::io::Result<()> {
let mut p = connect_retry(&cfg.listen_addr, Duration::from_secs(5))?;
p.set_read_timeout(Some(Duration::from_secs(10)))?;
p.write_all(&peer_init_bytes(&cfg.peer_username, "P", 0))?;
p.flush()?;
let _ = cfg.ready.send(());
loop {
let msg = read_framed(&mut p)?;
if msg.get_message_code() == 43 {
break;
}
}
let mut p = open_control_and_await_queue(cfg)?;
p.write_all(&refusal.get_buffer())?;
p.flush()?;
@@ -724,7 +746,11 @@ fn unique_download_dir() -> PathBuf {
dir
}
fn file_downloads_over_p_and_f(peer_username: &str, token_delay: Duration) {
fn file_downloads_over_p_and_f(
peer_username: &str,
token_delay: Duration,
run_uploader: fn(&MockUpload) -> std::io::Result<u64>,
) {
let server = server_or_skip!();
// Downloader: connected to the server with its peer listener enabled.
@@ -755,7 +781,7 @@ fn file_downloads_over_p_and_f(peer_username: &str, token_delay: Duration) {
token_delay,
};
let uploader = std::thread::spawn(move || {
if let Err(e) = run_mock_uploader(&cfg) {
if let Err(e) = run_uploader(&cfg) {
eprintln!("[mock uploader] {e}");
}
});
@@ -792,12 +818,29 @@ fn file_downloads_over_p_and_f(peer_username: &str, token_delay: Duration) {
#[test]
fn a_file_downloads_from_a_peer_over_p_and_f_connections() {
file_downloads_over_p_and_f("e2e_mockpeer", Duration::ZERO);
file_downloads_over_p_and_f(
"e2e_mockpeer",
Duration::ZERO,
run_mock_uploader,
);
}
#[test]
fn a_file_downloads_when_the_transfer_token_arrives_late() {
file_downloads_over_p_and_f("e2e_latepeer", Duration::from_millis(300));
file_downloads_over_p_and_f(
"e2e_latepeer",
Duration::from_millis(300),
run_mock_uploader,
);
}
#[test]
fn a_download_completes_when_the_f_connection_beats_token_registration() {
file_downloads_over_p_and_f(
"e2e_earlypeer",
Duration::ZERO,
run_mock_uploader_f_first,
);
}
fn refusal_fails_the_download_quickly(