diff --git a/soulseek-rs-lib/examples/stress.rs b/soulseek-rs-lib/examples/stress.rs index 81c6a8f..31d086e 100644 --- a/soulseek-rs-lib/examples/stress.rs +++ b/soulseek-rs-lib/examples/stress.rs @@ -895,6 +895,16 @@ fn write_baseline(path: &Path, score: &Score, cfg: &Config, m: &Metrics) { /// Live thread count for this process, so the harness can show whether load is /// being absorbed or simply turned into threads. fn thread_count() -> usize { + // Linux keeps the count in /proc; `ps -M` is macOS's spelling and returns + // a single unrelated line there, which read as one thread all along and + // made this guard against thread explosions useless on CI. + if let Ok(status) = std::fs::read_to_string("/proc/self/status") { + return status + .lines() + .find_map(|line| line.strip_prefix("Threads:")) + .and_then(|count| count.trim().parse().ok()) + .unwrap_or(0); + } Command::new("ps") .args(["-M", &std::process::id().to_string()]) .output() diff --git a/soulseek-rs-lib/src/message/peer/folder_contents.rs b/soulseek-rs-lib/src/message/peer/folder_contents.rs index 893554c..48a4bf9 100644 --- a/soulseek-rs-lib/src/message/peer/folder_contents.rs +++ b/soulseek-rs-lib/src/message/peer/folder_contents.rs @@ -79,7 +79,7 @@ pub fn build_folder_contents( pub fn parse_folder_contents( message: &mut Message, ) -> Option<(u32, String, Vec)> { - let mut body = decompress_body(message)?; + let mut body = decompress_body(message, crate::utils::zlib::MAX_INFLATED)?; let token = body.read_int32(); let folder = body.read_string(); Some((token, folder, read_directories(&mut body))) diff --git a/soulseek-rs-lib/src/message/peer/shared_file_list.rs b/soulseek-rs-lib/src/message/peer/shared_file_list.rs index cc60817..60084e5 100644 --- a/soulseek-rs-lib/src/message/peer/shared_file_list.rs +++ b/soulseek-rs-lib/src/message/peer/shared_file_list.rs @@ -4,7 +4,7 @@ use crate::message::{Message, MessageHandler}; use crate::peer::PeerMessage; -use crate::utils::zlib::{deflate, inflate}; +use crate::utils::zlib::{MAX_INFLATED_LISTING, deflate, inflate_limited}; use std::sync::mpsc::Sender; /// One file in a shared directory: its basename, size and the `(code, value)` @@ -68,7 +68,8 @@ pub fn build_shared_file_list(dirs: &[SharedDirectory]) -> Message { /// Returns an empty listing if the payload is malformed. #[must_use] pub fn parse_shared_file_list(message: &mut Message) -> Vec { - decompress_body(message) + // A whole share listing is the one payload that legitimately runs large. + decompress_body(message, MAX_INFLATED_LISTING) .map_or_else(Vec::new, |mut body| read_directories(&mut body)) } @@ -135,11 +136,15 @@ pub fn read_directories(body: &mut Message) -> Vec { dirs } -/// Inflate a compressed peer payload positioned at the blob. -pub fn decompress_body(message: &mut Message) -> Option { +/// Inflate a compressed peer payload positioned at the blob, refusing one +/// that expands past `max_out` — a peer we have never met chose these bytes. +pub fn decompress_body( + message: &mut Message, + max_out: usize, +) -> Option { let pointer = message.get_pointer(); let size = message.get_size(); - inflate(&message.get_slice(pointer, size)) + inflate_limited(&message.get_slice(pointer, size), max_out) .ok() .map(Message::new_with_data) } diff --git a/soulseek-rs-lib/src/peer/upload_peer.rs b/soulseek-rs-lib/src/peer/upload_peer.rs index 36e8978..b999833 100644 --- a/soulseek-rs-lib/src/peer/upload_peer.rs +++ b/soulseek-rs-lib/src/peer/upload_peer.rs @@ -5,7 +5,7 @@ use std::io::{self, Read, Seek, SeekFrom, Write}; use std::net::{TcpStream, ToSocketAddrs}; use std::path::Path; use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; -use std::time::{Duration, Instant}; +use std::time::Duration; use crate::message::server::MessageFactory; use crate::peer::ConnectionType; @@ -99,41 +99,11 @@ pub fn serve_file( } stream.flush()?; - // Soulseek assigns closing the file-transfer connection to the downloader - // after it has read exactly the advertised number of bytes. Poll while - // waiting for that EOF so cancelling a finished-but-unacknowledged upload - // remains prompt rather than waiting for the full socket timeout. - stream.set_read_timeout(Some(Duration::from_millis(200)))?; - let close_deadline = Instant::now() + Duration::from_secs(30); - let mut unexpected = [0u8; 1]; - loop { - if cancel.load(Ordering::Relaxed) { - return Err(io::ErrorKind::Interrupted.into()); - } - match stream.read(&mut unexpected) { - Ok(0) => break, - Ok(_) => { - return Err(io::Error::new( - io::ErrorKind::InvalidData, - "unexpected data after download offset", - )); - } - Err(error) - if matches!( - error.kind(), - io::ErrorKind::WouldBlock | io::ErrorKind::TimedOut - ) => - { - if Instant::now() >= close_deadline { - return Err(io::Error::new( - io::ErrorKind::TimedOut, - "downloader did not close the file connection", - )); - } - } - Err(error) => return Err(error), - } - } + // The file is delivered. Soulseek leaves closing to the downloader, but + // waiting for that costs the upload slot the next peer is queued for — + // and a peer that leaves the socket open, or whose FIN a NAT swallows, + // would hold it for the whole wait. Nicotine+ finishes the transfer at + // delivery for the same reason; the close below is ours to make. trace!("[upload] served {} to {}:{}", path.display(), host, port); Ok(streamed) } @@ -285,14 +255,16 @@ mod tests { } #[test] - fn serve_file_cancels_while_waiting_for_the_downloader_to_close() { - let content: Vec = (0..4096u32).map(|i| (i % 251) as u8).collect(); - let path = scratch_file("cancel-close-wait", &content); + fn a_downloader_that_never_closes_still_leaves_the_upload_delivered() { + // The file is on the wire; whether the peer closes, or its FIN is + // eaten by a NAT, it has had every byte. Waiting for a close that may + // never come used to fail the transfer and hold its upload slot, + // which is what starved a queue of everyone behind it. + let content: Vec = (0..2048u32).map(|i| (i % 251) as u8).collect(); + let path = scratch_file("no-close", &content); let dir = path.parent().unwrap().to_path_buf(); let listener = TcpListener::bind("127.0.0.1:0").unwrap(); let port = u32::from(listener.local_addr().unwrap().port()); - let cancel = Arc::new(AtomicBool::new(false)); - let cancel_flag = cancel.clone(); let (result_tx, result_rx) = mpsc::channel(); let uploader = std::thread::spawn(move || { @@ -300,10 +272,10 @@ mod tests { "127.0.0.1", port, "me", - 782, + 783, &path, &AtomicU64::new(0), - &cancel_flag, + &AtomicBool::new(false), ); let _ = result_tx.send(result); }); @@ -313,12 +285,12 @@ mod tests { stream.read_exact(&mut received).unwrap(); assert_eq!(received, content); - cancel.store(true, Ordering::Relaxed); - let error = result_rx - .recv_timeout(Duration::from_secs(2)) - .expect("cancellation must wake the close wait") - .unwrap_err(); - assert_eq!(error.kind(), std::io::ErrorKind::Interrupted); + // Hold the socket open past the grace period and say nothing. + let served = result_rx + .recv_timeout(Duration::from_secs(5)) + .expect("delivery must finish the upload, not a wait for a close") + .expect("a delivered file is a finished upload"); + assert_eq!(served, content.len() as u64); drop(stream); let _ = uploader.join(); diff --git a/soulseek-rs-lib/src/utils/zlib.rs b/soulseek-rs-lib/src/utils/zlib.rs index 46a6bc8..bace9ae 100644 --- a/soulseek-rs-lib/src/utils/zlib.rs +++ b/soulseek-rs-lib/src/utils/zlib.rs @@ -59,8 +59,33 @@ impl BitReader { use crate::error::{Result, SoulseekRs}; +/// The most a peer payload may inflate to unless the caller says otherwise. +/// +/// Nicotine+ passes this ceiling to its decompressor for a search or folder +/// reply, and a larger one only for a browse listing. +pub const MAX_INFLATED: usize = 128 * 1024 * 1024; + +/// The ceiling for a whole browse listing, which a large sharer legitimately +/// fills. The same number Nicotine+ allows. +pub const MAX_INFLATED_LISTING: usize = 2 * 1024 * 1024 * 1024; + +/// What a payload that outgrows its ceiling is refused with. +const TOO_BIG: &str = "compressed payload expands past its limit"; + /// Inflate a zlib stream, checking its adler32 trailer. +/// +/// Bounded by [`MAX_INFLATED`]; use [`inflate_limited`] for a payload with a +/// different ceiling. pub fn inflate(input: &[u8]) -> Result> { + inflate_limited(input, MAX_INFLATED) +} + +/// Inflate a zlib stream, refusing to grow the output past `max_out`. +/// +/// The ceiling is the point: deflate reaches roughly 1000:1, so a few +/// kilobytes from a peer we have never met can otherwise expand into +/// gigabytes and take the process with it. +pub fn inflate_limited(input: &[u8], max_out: usize) -> Result> { let mut r = BitReader::new(input.to_vec()); let cmf = r.read_byte()?; let cm = cmf & 15; // Compression method @@ -84,7 +109,8 @@ pub fn inflate(input: &[u8]) -> Result> { "preset dictionary not supported".to_string(), )); } - let out = inflate_blocks(&mut r).map_err(SoulseekRs::CompressionError)?; + let out = inflate_blocks(&mut r, max_out) + .map_err(SoulseekRs::CompressionError)?; // `read_bytes` assembles little-endian; the zlib trailer is big-endian. if r.read_bytes(4)?.swap_bytes() != adler32(&out) { return Err(SoulseekRs::CompressionError( @@ -105,18 +131,24 @@ fn adler32(data: &[u8]) -> u32 { (b << 16) | a } -fn inflate_blocks(r: &mut BitReader) -> std::result::Result, String> { +fn inflate_blocks( + r: &mut BitReader, + max_out: usize, +) -> std::result::Result, String> { let mut bfinal = 0; let mut out = Vec::new(); while bfinal == 0 { bfinal = r.read_bit()?; let btype = r.read_bits(2)?; match btype { - 0 => inflate_block_no_compression(r, &mut out)?, - 1 => inflate_block_fixed(r, &mut out)?, - 2 => inflate_block_dynamic(r, &mut out)?, + 0 => inflate_block_no_compression(r, &mut out, max_out)?, + 1 => inflate_block_fixed(r, &mut out, max_out)?, + 2 => inflate_block_dynamic(r, &mut out, max_out)?, _ => return Err("invalid BTYPE".to_string()), } + if out.len() > max_out { + return Err(TOO_BIG.to_string()); + } } Ok(out) } @@ -124,9 +156,13 @@ fn inflate_blocks(r: &mut BitReader) -> std::result::Result, String> { fn inflate_block_no_compression( r: &mut BitReader, o: &mut Vec, + max_out: usize, ) -> std::result::Result<(), String> { let len = r.read_bytes(2)?; let _nlen = r.read_bytes(2)?; + if o.len().saturating_add(len as usize) > max_out { + return Err(TOO_BIG.to_string()); + } for _ in 0..len { o.push(r.read_byte()?); } @@ -214,8 +250,14 @@ fn inflate_block_data( literal_length_tree: &HuffmanTree, distance_tree: &HuffmanTree, out: &mut Vec, + max_out: usize, ) -> std::result::Result<(), String> { loop { + // Checked per symbol, not per block: a single block can emit output + // without bound, which is exactly how a bomb is built. + if out.len() > max_out { + return Err(TOO_BIG.to_string()); + } let sym = decode_symbol(r, literal_length_tree)?; if sym <= 255 { // Literal byte @@ -347,14 +389,16 @@ fn decode_trees( fn inflate_block_dynamic( r: &mut BitReader, o: &mut Vec, + max_out: usize, ) -> std::result::Result<(), String> { let (literal_length_tree, distance_tree) = decode_trees(r)?; - inflate_block_data(r, &literal_length_tree, &distance_tree, o) + inflate_block_data(r, &literal_length_tree, &distance_tree, o, max_out) } fn inflate_block_fixed( r: &mut BitReader, o: &mut Vec, + max_out: usize, ) -> std::result::Result<(), String> { let mut bl = Vec::new(); bl.extend(vec![8; 144]); // 0-143: 8 bits @@ -369,7 +413,7 @@ fn inflate_block_fixed( let distance_alphabet: Vec = (0..30).collect(); let distance_tree = bl_list_to_tree(&bl_dist, &distance_alphabet); - inflate_block_data(r, &literal_length_tree, &distance_tree, o) + inflate_block_data(r, &literal_length_tree, &distance_tree, o, max_out) } struct BitWriter { @@ -533,6 +577,31 @@ mod tests { use super::*; + #[test] + fn a_payload_that_expands_past_its_ceiling_is_refused() { + // Deflate reaches ~1000:1, so a few kilobytes from an unknown peer can + // become gigabytes. The ceiling is what stops that being our problem. + let bomb = deflate(&vec![0u8; 1024 * 1024]); + assert!( + bomb.len() < 8 * 1024, + "the test needs a payload that expands hugely, got {} bytes", + bomb.len() + ); + + let error = inflate_limited(&bomb, 64 * 1024) + .expect_err("a payload past the ceiling must be refused"); + assert!( + format!("{error}").contains("expands past its limit"), + "unexpected error: {error}" + ); + + // The same bytes are fine when they fit. + assert_eq!( + inflate_limited(&bomb, 2 * 1024 * 1024).unwrap().len(), + 1024 * 1024 + ); + } + #[test] fn a_corrupt_adler32_trailer_is_rejected() { let mut stream = deflate(b"the bytes a peer sent");