From 70fcbfcb206b99e32cd9ef9ea6b2b8eb03773102 Mon Sep 17 00:00:00 2001 From: artrixdotdev Date: Mon, 27 Jul 2026 11:38:36 -0700 Subject: [PATCH 01/10] perf: eliminate peer message payload copies --- crates/libtortillas/examples/peer_transfer.rs | 63 +++++++++ crates/libtortillas/src/protocol/messages.rs | 118 ++++++++++++++--- crates/libtortillas/src/protocol/stream.rs | 124 ++++++++++++++++-- 3 files changed, 273 insertions(+), 32 deletions(-) create mode 100644 crates/libtortillas/examples/peer_transfer.rs diff --git a/crates/libtortillas/examples/peer_transfer.rs b/crates/libtortillas/examples/peer_transfer.rs new file mode 100644 index 00000000..01e9f225 --- /dev/null +++ b/crates/libtortillas/examples/peer_transfer.rs @@ -0,0 +1,63 @@ +use std::{error::Error, time::Instant}; + +use bytes::Bytes; +use libtortillas::protocol::{ + messages::PeerMessages, + stream::{PeerRecv, PeerSend, PeerStream}, +}; +use tokio::net::{TcpListener, TcpStream}; +use tracing::info; + +const BLOCK_LENGTH: usize = 16 * 1024; +const BLOCK_COUNT: usize = 8 * 1024; + +#[tokio::main(flavor = "current_thread")] +async fn main() -> Result<(), Box> { + tracing_subscriber::fmt().init(); + + let listener = TcpListener::bind("127.0.0.1:0").await?; + let address = listener.local_addr()?; + let receiver = tokio::spawn(async move { + let (stream, _) = listener.accept().await?; + let mut stream = PeerStream::tcp(stream); + let mut bytes_received = 0usize; + + for expected_index in 0..BLOCK_COUNT { + match stream.recv().await? { + PeerMessages::Piece(index, 0, block) + if index == expected_index as u32 && block.len() == BLOCK_LENGTH => + { + bytes_received += block.len(); + } + message => { + return Err(format!("received unexpected message: {message}").into()); + } + } + } + + Ok::<_, Box>(bytes_received) + }); + + let block = Bytes::from(vec![0xa5; BLOCK_LENGTH]); + let mut sender = PeerStream::tcp(TcpStream::connect(address).await?); + let started_at = Instant::now(); + + for index in 0..BLOCK_COUNT { + sender + .send(PeerMessages::Piece(index as u32, 0, block.clone())) + .await?; + } + + let bytes_transferred = receiver.await??; + let elapsed = started_at.elapsed(); + let mebibytes = bytes_transferred as f64 / (1024.0 * 1024.0); + let throughput = mebibytes / elapsed.as_secs_f64(); + info!( + bytes_transferred, + ?elapsed, + throughput_mib_per_second = throughput, + "completed local peer transfer" + ); + + Ok(()) +} diff --git a/crates/libtortillas/src/protocol/messages.rs b/crates/libtortillas/src/protocol/messages.rs index 6d165d95..303c7993 100644 --- a/crates/libtortillas/src/protocol/messages.rs +++ b/crates/libtortillas/src/protocol/messages.rs @@ -2,6 +2,7 @@ use core::hash; use std::{ collections::HashMap, fmt::Display, + mem::size_of, net::{IpAddr, Ipv4Addr, Ipv6Addr}, sync::{Arc, atomic::AtomicU8}, }; @@ -159,34 +160,70 @@ impl Display for PeerMessages { } } +const PIECE_HEADER_LENGTH: usize = 13; +const LENGTH_PREFIX_LENGTH: usize = size_of::(); + +pub(super) enum EncodedPeerMessage<'a> { + Contiguous(Bytes), + Piece { + header: [u8; PIECE_HEADER_LENGTH], + block: &'a Bytes, + }, +} + +impl EncodedPeerMessage<'_> { + fn into_bytes(self) -> Bytes { + match self { + Self::Contiguous(bytes) => bytes, + Self::Piece { header, block } => { + let mut message = BytesMut::with_capacity(header.len() + block.len()); + message.put_slice(&header); + message.put_slice(block); + message.freeze() + } + } + } +} + impl PeerMessages { pub fn to_bytes(&self) -> Result { + Ok(self.encode()?.into_bytes()) + } + + pub(super) fn encode(&self) -> Result, PeerActorError> { Ok(match self { - PeerMessages::Handshake(handshake) => handshake.to_bytes(), - PeerMessages::Choke => create_message_with_id(0, &[]), - PeerMessages::Unchoke => create_message_with_id(1, &[]), - PeerMessages::Interested => create_message_with_id(2, &[]), - PeerMessages::NotInterested => create_message_with_id(3, &[]), - PeerMessages::Have(index) => create_message_with_id(4, &index.to_be_bytes()), - PeerMessages::Bitfield(bits) => create_message_with_id(5, &encode_bitfield(bits)), + PeerMessages::Handshake(handshake) => { + EncodedPeerMessage::Contiguous(handshake.to_bytes()) + } + PeerMessages::Choke => EncodedPeerMessage::Contiguous(create_message_with_id(0, &[])), + PeerMessages::Unchoke => EncodedPeerMessage::Contiguous(create_message_with_id(1, &[])), + PeerMessages::Interested => { + EncodedPeerMessage::Contiguous(create_message_with_id(2, &[])) + } + PeerMessages::NotInterested => { + EncodedPeerMessage::Contiguous(create_message_with_id(3, &[])) + } + PeerMessages::Have(index) => { + EncodedPeerMessage::Contiguous(create_message_with_id(4, &index.to_be_bytes())) + } + PeerMessages::Bitfield(bits) => { + EncodedPeerMessage::Contiguous(create_message_with_id(5, &encode_bitfield(bits))) + } PeerMessages::Request(index, begin, length) // Code is identical | PeerMessages::Cancel(index, begin, length) => { let id = match self { PeerMessages::Request(..) => 6, _ => 8, }; - let mut payload = BytesMut::with_capacity(12); - payload.put_slice(&index.to_be_bytes()); - payload.put_slice(&begin.to_be_bytes()); - payload.put_slice(&length.to_be_bytes()); - create_message_with_id(id, &payload) + EncodedPeerMessage::Contiguous(create_message_with_triplet( + id, *index, *begin, *length, + )) } PeerMessages::Piece(index, begin, data) => { - let mut payload = BytesMut::with_capacity(8 + data.len()); - payload.put_slice(&index.to_be_bytes()); - payload.put_slice(&begin.to_be_bytes()); - payload.put_slice(data); - create_message_with_id(7, &payload) + EncodedPeerMessage::Piece { + header: create_piece_header(*index, *begin, data.len())?, + block: data, + } } PeerMessages::Extended(extended_id, handshake_message, metadata) => { let mut payload = BytesMut::new(); @@ -197,9 +234,11 @@ impl PeerMessages { if let Some(metadata) = metadata { payload.extend_from_slice(metadata); } - create_message_with_id(20, &payload) + EncodedPeerMessage::Contiguous(create_message_with_id(20, &payload)) + } + PeerMessages::KeepAlive => { + EncodedPeerMessage::Contiguous(Bytes::from_static(&[0u8; 4])) } - PeerMessages::KeepAlive => Bytes::from_static(&[0u8; 4]), }) } @@ -613,6 +652,33 @@ fn create_message_with_id(id: u8, payload: &[u8]) -> Bytes { message.freeze() } +fn create_message_with_triplet(id: u8, first: u32, second: u32, third: u32) -> Bytes { + let mut message = BytesMut::with_capacity(17); + message.put_u32(13); + message.put_u8(id); + message.put_u32(first); + message.put_u32(second); + message.put_u32(third); + message.freeze() +} + +fn create_piece_header( + index: u32, begin: u32, block_length: usize, +) -> Result<[u8; PIECE_HEADER_LENGTH], PeerActorError> { + let message_length = block_length + .checked_add(PIECE_HEADER_LENGTH - LENGTH_PREFIX_LENGTH) + .and_then(|length| u32::try_from(length).ok()) + .ok_or_else(|| PeerActorError::InvalidMessagePayload { + message_type: "Piece length exceeds the protocol limit".to_string(), + })?; + let mut header = [0; PIECE_HEADER_LENGTH]; + header[..4].copy_from_slice(&message_length.to_be_bytes()); + header[4] = 7; + header[5..9].copy_from_slice(&index.to_be_bytes()); + header[9..].copy_from_slice(&begin.to_be_bytes()); + Ok(header) +} + /// Encodes the canonical piece-index order using BEP 3's most-significant-bit /// first wire representation. Internal bit vectors intentionally keep their /// default ordering; unit conversion belongs at the protocol boundary. @@ -786,6 +852,20 @@ mod ipaddr_serde { mod tests { use super::*; + #[test] + fn piece_message_when_encoded_then_contains_header_and_original_block() { + let block = Bytes::from_static(b"payload"); + let message = PeerMessages::Piece(2, 4, block.clone()); + + let encoded = message.to_bytes().unwrap(); + + assert_eq!(&encoded[..4], &(16u32.to_be_bytes())); + assert_eq!(encoded[4], 7); + assert_eq!(&encoded[5..9], &(2u32.to_be_bytes())); + assert_eq!(&encoded[9..13], &(4u32.to_be_bytes())); + assert_eq!(&encoded[13..], block); + } + #[test] fn peer_bitfield_uses_most_significant_bit_first_piece_order() { let mut wire = vec![u8::MAX; 132]; diff --git a/crates/libtortillas/src/protocol/stream.rs b/crates/libtortillas/src/protocol/stream.rs index c4010371..f79599eb 100644 --- a/crates/libtortillas/src/protocol/stream.rs +++ b/crates/libtortillas/src/protocol/stream.rs @@ -1,6 +1,7 @@ use std::{ fmt, fmt::Display, + io::IoSlice, net::SocketAddr, pin::Pin, sync::Arc, @@ -18,7 +19,7 @@ use tokio::{ use tokio_util::sync::CancellationToken; use tracing::{error, instrument, trace}; -use super::messages::{Handshake, PeerMessages}; +use super::messages::{EncodedPeerMessage, Handshake, PeerMessages}; use crate::{ errors::PeerActorError, hashes::InfoHash, @@ -42,8 +43,15 @@ pub struct PeerStream { pub trait PeerSend: AsyncWrite + Unpin { /// Sends a PeerMessage to a peer. async fn send(&mut self, data: PeerMessages) -> Result<(), PeerActorError> { - let bytes = data.to_bytes()?; - self.write_all(&bytes).await.map_err(|e| { + let mut writer = self; + let result = match data.encode()? { + EncodedPeerMessage::Contiguous(bytes) => writer.write_all(&bytes).await, + EncodedPeerMessage::Piece { header, block } => { + let mut frame = Buf::chain(header.as_slice(), block.as_ref()); + AsyncWriteExt::write_all_buf(&mut writer, &mut frame).await + } + }; + result.map_err(|e| { error!(error = %e, "Failed to send message to peer"); PeerActorError::SendFailed(e.to_string()) }) @@ -89,7 +97,8 @@ pub trait PeerRecv: AsyncRead + Unpin { return Ok(PeerMessages::KeepAlive); } - let mut message_buf = BytesMut::with_capacity(4 + length as usize); + let frame_length = 4 + length as usize; + let mut message_buf = BytesMut::with_capacity(frame_length); message_buf.extend_from_slice(&length_buf); let mut message_type = [0u8; 1]; @@ -101,14 +110,12 @@ pub trait PeerRecv: AsyncRead + Unpin { message_buf.extend_from_slice(&message_type); // Read the rest of the message payload - let mut rest = vec![0u8; (length - 1) as usize]; + message_buf.resize(frame_length, 0); self - .read_exact(&mut rest) + .read_exact(&mut message_buf[5..]) .await .map_err(PeerActorError::ReceiveFailed)?; - message_buf.extend_from_slice(&rest); - PeerMessages::from_bytes(message_buf.freeze()) } /// Receives a message from a peer with cancellation support, returning @@ -194,12 +201,11 @@ impl PeerStream { /// Receives an incoming handshake from a peer. pub async fn recv_handshake_message(&mut self) -> Result { let protocol_len = self.read_u8().await?; - let mut buf = Vec::with_capacity(1 + protocol_len as usize + 8 + 40); + let handshake_length = 1 + protocol_len as usize + 8 + 40; + let mut buf = Vec::with_capacity(handshake_length); buf.push(protocol_len); - - let mut rest = vec![0u8; protocol_len as usize + 8 + 40]; - self.read_exact(&mut rest).await?; - buf.extend_from_slice(&rest); + buf.resize(handshake_length, 0); + self.read_exact(&mut buf[1..]).await?; Handshake::from_bytes(&buf).map_err(|e| PeerActorError::HandshakeFailed { reason: e.to_string(), @@ -352,6 +358,26 @@ impl AsyncWrite for PeerStream { result } + fn poll_write_vectored( + mut self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>], + ) -> Poll> { + let result = match &mut self.transport { + PeerTransport::Tcp(stream) => Pin::new(stream).poll_write_vectored(cx, bufs), + PeerTransport::Utp(stream) => Pin::new(stream).poll_write_vectored(cx, bufs), + }; + if let Poll::Ready(Ok(bytes_written)) = &result { + self.peer_state.increment_bytes_uploaded(*bytes_written); + } + result + } + + fn is_write_vectored(&self) -> bool { + match &self.transport { + PeerTransport::Tcp(stream) => stream.is_write_vectored(), + PeerTransport::Utp(stream) => stream.is_write_vectored(), + } + } + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { match &mut self.transport { PeerTransport::Tcp(stream) => Pin::new(stream).poll_flush(cx), @@ -425,6 +451,26 @@ impl AsyncWrite for PeerWriter { result } + fn poll_write_vectored( + mut self: Pin<&mut Self>, cx: &mut Context<'_>, bufs: &[IoSlice<'_>], + ) -> Poll> { + let result = match &mut self.writer { + PeerWriteHalf::Tcp(stream) => Pin::new(stream).poll_write_vectored(cx, bufs), + PeerWriteHalf::Utp(stream) => Pin::new(stream).poll_write_vectored(cx, bufs), + }; + if let Poll::Ready(Ok(bytes_written)) = &result { + self.peer_state.increment_bytes_uploaded(*bytes_written); + } + result + } + + fn is_write_vectored(&self) -> bool { + match &self.writer { + PeerWriteHalf::Tcp(stream) => stream.is_write_vectored(), + PeerWriteHalf::Utp(stream) => stream.is_write_vectored(), + } + } + fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { match &mut self.writer { PeerWriteHalf::Tcp(stream) => Pin::new(stream).poll_flush(cx), @@ -497,6 +543,7 @@ pub fn validate_handshake_protocol( mod tests { use std::time::Duration; + use bytes::Bytes; use tokio::{io::AsyncWriteExt, net::TcpListener, time::timeout}; use tracing_test::traced_test; @@ -508,6 +555,57 @@ mod tests { testing::{self, LocalPeer}, }; + #[derive(Default)] + struct VectoredWriter { + bytes: Vec, + vectored_writes: usize, + } + + impl AsyncWrite for VectoredWriter { + fn poll_write( + self: Pin<&mut Self>, _: &mut Context<'_>, _: &[u8], + ) -> Poll> { + panic!("piece messages should use vectored writes") + } + + fn poll_write_vectored( + mut self: Pin<&mut Self>, _: &mut Context<'_>, bufs: &[IoSlice<'_>], + ) -> Poll> { + self.vectored_writes += 1; + let written = bufs.iter().map(|buffer| buffer.len()).sum(); + for buffer in bufs { + self.bytes.extend_from_slice(buffer); + } + Poll::Ready(Ok(written)) + } + + fn is_write_vectored(&self) -> bool { + true + } + + fn poll_flush(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + + fn poll_shutdown(self: Pin<&mut Self>, _: &mut Context<'_>) -> Poll> { + Poll::Ready(Ok(())) + } + } + + impl PeerSend for VectoredWriter {} + + #[tokio::test] + async fn peer_send_when_message_is_piece_then_uses_vectored_write() { + let message = PeerMessages::Piece(2, 4, Bytes::from_static(b"payload")); + let expected = message.to_bytes().unwrap(); + let mut writer = VectoredWriter::default(); + + writer.send(message).await.unwrap(); + + assert_eq!(writer.vectored_writes, 1); + assert_eq!(writer.bytes, expected); + } + #[tokio::test] #[traced_test] async fn peer_stream_when_handshake_is_valid_then_returns_peer_id() { From 0b921d557abc4ee4fe671b1a45ab74b5b67f12e3 Mon Sep 17 00:00:00 2001 From: artrixdotdev Date: Mon, 27 Jul 2026 11:38:42 -0700 Subject: [PATCH 02/10] perf: avoid repeated tracker and piece work --- crates/libtortillas/src/pieces/piece_store.rs | 10 +++++++--- crates/libtortillas/src/tracker/http.rs | 10 +++++++++- 2 files changed, 16 insertions(+), 4 deletions(-) diff --git a/crates/libtortillas/src/pieces/piece_store.rs b/crates/libtortillas/src/pieces/piece_store.rs index 8714ae04..1c5554e9 100644 --- a/crates/libtortillas/src/pieces/piece_store.rs +++ b/crates/libtortillas/src/pieces/piece_store.rs @@ -2,7 +2,7 @@ use std::{io::Result as IoResult, path::PathBuf}; use bytes::Bytes; use kameo::{Actor, actor::ActorRef, messages}; -use tokio::fs::read; +use tokio::{fs::read, task::spawn_blocking}; use crate::{errors::TorrentError, hashes::Hash, torrent::util}; @@ -30,7 +30,11 @@ impl PieceStoreActor { pub(crate) async fn validate_and_read( &mut self, path: PathBuf, hash: Hash<20>, ) -> anyhow::Result { - util::validate_piece_file(path.clone(), hash).await?; - Ok(read(&path).await?.into()) + let data: Bytes = read(path).await?.into(); + spawn_blocking(move || -> anyhow::Result { + util::validate_piece_bytes(&data, hash)?; + Ok(data) + }) + .await? } } diff --git a/crates/libtortillas/src/tracker/http.rs b/crates/libtortillas/src/tracker/http.rs index b7ae854a..d073705e 100644 --- a/crates/libtortillas/src/tracker/http.rs +++ b/crates/libtortillas/src/tracker/http.rs @@ -9,6 +9,7 @@ use std::{ use anyhow::Result; use async_trait::async_trait; +use reqwest::Client; use serde::{ Deserialize, de::{self, Visitor}, @@ -130,6 +131,7 @@ impl TrackerRequest { #[derive(Clone, Debug)] pub struct HttpTracker { uri: String, + client: Client, pub peer_id: PeerId, info_hash: InfoHash, params: Arc>, @@ -186,6 +188,7 @@ impl HttpTracker { HttpTracker { interval: Arc::new(usize::MAX.into()), uri, + client: Client::new(), peer_id, params, info_hash, @@ -236,7 +239,12 @@ impl TrackerBase for HttpTracker { // HTTP request phase let request_start = Instant::now(); - let response = reqwest::get(&uri).await.map_err(TrackerActorError::Http)?; + let response = self + .client + .get(&uri) + .send() + .await + .map_err(TrackerActorError::Http)?; self.stats.increment_bytes_sent(uri.len()); let response_bytes = response.bytes().await.map_err(TrackerActorError::Http)?; From 8d9392291c1e4f76ca1363c81386cb35bc0b1d23 Mon Sep 17 00:00:00 2001 From: artrixdotdev Date: Mon, 27 Jul 2026 11:38:48 -0700 Subject: [PATCH 03/10] docs: simplify torrent actor link --- crates/libtortillas/src/engine/messages.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/libtortillas/src/engine/messages.rs b/crates/libtortillas/src/engine/messages.rs index 214aa976..fe57c280 100644 --- a/crates/libtortillas/src/engine/messages.rs +++ b/crates/libtortillas/src/engine/messages.rs @@ -151,7 +151,7 @@ pub(crate) mod commands { Ok(torrent) } - /// Creates a new [`Torrent`](crate::torrent::Torrent) actor. + /// Creates a new [`Torrent`] actor. #[message] pub(crate) async fn create_torrent( &mut self, request: CreateTorrentRequest, From b02b9467178da7a83f52f109e7c483cd6cab6b2d Mon Sep 17 00:00:00 2001 From: artrixdotdev Date: Mon, 27 Jul 2026 12:12:50 -0700 Subject: [PATCH 04/10] perf: eliminate protocol allocation churn --- crates/libtortillas/src/engine/messages.rs | 2 +- crates/libtortillas/src/lib.rs | 11 +- crates/libtortillas/src/protocol/messages.rs | 13 +- crates/libtortillas/src/protocol/stream.rs | 190 ++++++++++--------- crates/libtortillas/src/torrent/swarm.rs | 46 ++--- 5 files changed, 135 insertions(+), 127 deletions(-) diff --git a/crates/libtortillas/src/engine/messages.rs b/crates/libtortillas/src/engine/messages.rs index fe57c280..628a402f 100644 --- a/crates/libtortillas/src/engine/messages.rs +++ b/crates/libtortillas/src/engine/messages.rs @@ -85,7 +85,7 @@ pub(crate) mod commands { return; } - let info_hash = *handshake.info_hash; + let info_hash = handshake.info_hash; let mut peer = Peer::from_socket_addr(peer_addr); // Populate peer fields from parsed handshake. diff --git a/crates/libtortillas/src/lib.rs b/crates/libtortillas/src/lib.rs index 869027c3..863aa121 100644 --- a/crates/libtortillas/src/lib.rs +++ b/crates/libtortillas/src/lib.rs @@ -562,7 +562,7 @@ pub(crate) mod testing { let handshake = stream.recv_handshake_message().await?; handshakes.lock().await.push(handshake.clone()); - let response = Handshake::new(handshake.info_hash.clone(), peer_id); + let response = Handshake::new(handshake.info_hash, peer_id); stream.write_all(&response.to_bytes()).await?; for message in messages.iter() { @@ -590,7 +590,7 @@ pub(crate) mod testing { #[cfg(test)] mod tests { - use std::{net::Ipv4Addr, sync::Arc}; + use std::net::Ipv4Addr; use tokio::time::{Duration, timeout}; @@ -623,11 +623,8 @@ pub(crate) mod testing { let mut stream = PeerStream::connect(local_peer.peer().socket_addr(), None) .await .unwrap(); - let info_hash = Arc::new(test_info_hash()); - stream - .send_handshake(peer_id(), info_hash.clone()) - .await - .unwrap(); + let info_hash = test_info_hash(); + stream.send_handshake(peer_id(), info_hash).await.unwrap(); let (received_peer_id, _) = stream.recv_handshake().await.unwrap(); let message = timeout(Duration::from_secs(1), stream.recv()) diff --git a/crates/libtortillas/src/protocol/messages.rs b/crates/libtortillas/src/protocol/messages.rs index 303c7993..6e95f4a6 100644 --- a/crates/libtortillas/src/protocol/messages.rs +++ b/crates/libtortillas/src/protocol/messages.rs @@ -565,14 +565,14 @@ pub struct Handshake { /// Reserved bytes for protocol extensions pub reserved: [u8; 8], /// 20-byte SHA1 hash of the info dictionary - pub info_hash: Arc>, + pub info_hash: Hash<20>, /// 20-byte peer identifier pub peer_id: PeerId, } impl Handshake { /// Create a new handshake with the given info hash and peer ID - pub fn new(info_hash: Arc>, peer_id: PeerId) -> Self { + pub fn new(info_hash: Hash<20>, peer_id: PeerId) -> Self { let mut reserved = [0u8; 8]; // We support BEP 0010 @@ -613,14 +613,19 @@ impl Handshake { ensure!(bytes.len() >= total_expected_len, "handshake too short"); // Extract protocol string - let protocol = Bytes::copy_from_slice(&bytes[1..1 + protocol_len]); + let protocol_bytes = &bytes[1..1 + protocol_len]; + let protocol = if protocol_bytes == MAGIC_STRING { + Bytes::from_static(MAGIC_STRING) + } else { + Bytes::copy_from_slice(protocol_bytes) + }; // Extract reserved bytes let reserved = bytes[1 + protocol_len..1 + protocol_len + 8].try_into()?; // Extract info hash let info_hash_bytes = bytes[1 + protocol_len + 8..1 + protocol_len + 8 + 20].try_into()?; - let info_hash = Arc::new(Hash::new(info_hash_bytes)); + let info_hash = Hash::new(info_hash_bytes); // Extract peer ID let peer_id_bytes: [u8; 20] = diff --git a/crates/libtortillas/src/protocol/stream.rs b/crates/libtortillas/src/protocol/stream.rs index f79599eb..068606fc 100644 --- a/crates/libtortillas/src/protocol/stream.rs +++ b/crates/libtortillas/src/protocol/stream.rs @@ -1,6 +1,7 @@ use std::{ fmt, fmt::Display, + future::Future, io::IoSlice, net::SocketAddr, pin::Pin, @@ -9,7 +10,6 @@ use std::{ }; use anyhow::Result; -use async_trait::async_trait; use bytes::{Buf, BytesMut}; use librqbit_utp::{UtpSocketUdp, UtpStream, UtpStreamReadHalf, UtpStreamWriteHalf}; use tokio::{ @@ -39,97 +39,115 @@ pub struct PeerStream { peer_state: PeerState, } -#[async_trait] pub trait PeerSend: AsyncWrite + Unpin { /// Sends a PeerMessage to a peer. - async fn send(&mut self, data: PeerMessages) -> Result<(), PeerActorError> { - let mut writer = self; - let result = match data.encode()? { - EncodedPeerMessage::Contiguous(bytes) => writer.write_all(&bytes).await, - EncodedPeerMessage::Piece { header, block } => { - let mut frame = Buf::chain(header.as_slice(), block.as_ref()); - AsyncWriteExt::write_all_buf(&mut writer, &mut frame).await - } - }; - result.map_err(|e| { - error!(error = %e, "Failed to send message to peer"); - PeerActorError::SendFailed(e.to_string()) - }) + fn send(&mut self, data: PeerMessages) -> impl Future> + Send + where + Self: Send, + { + async move { + let mut writer = self; + let result = match data.encode()? { + EncodedPeerMessage::Contiguous(bytes) => writer.write_all(&bytes).await, + EncodedPeerMessage::Piece { header, block } => { + let mut frame = Buf::chain(header.as_slice(), block.as_ref()); + AsyncWriteExt::write_all_buf(&mut writer, &mut frame).await + } + }; + result.map_err(|e| { + error!(error = %e, "Failed to send message to peer"); + PeerActorError::SendFailed(e.to_string()) + }) + } } /// Sends a message to a peer with a cancellation support, returning an /// error if the operation is cancel - async fn send_with_cancel( + fn send_with_cancel( &mut self, data: PeerMessages, token: CancellationToken, - ) -> Result<(), PeerActorError> { - tokio::select! { - _ = token.cancelled() => { - trace!("Sending message to peer was cancelled"); - return Err(PeerActorError::MessageCancelled); - - }, - result = self.send(data) => { - result + ) -> impl Future> + Send + where + Self: Send, + { + async move { + tokio::select! { + _ = token.cancelled() => { + trace!("Sending message to peer was cancelled"); + Err(PeerActorError::MessageCancelled) + + }, + result = self.send(data) => { + result + } } } } } -#[async_trait] pub trait PeerRecv: AsyncRead + Unpin { /// Receives data from a peers stream. In other words, if you wish to /// directly contact a peer, use this function. - async fn recv(&mut self) -> Result { - // First 4 bytes is the big endian encoded length field and the 5th byte is a - // PeerMessage tag - let mut length_buf = [0u8; 4]; + fn recv(&mut self) -> impl Future> + Send + where + Self: Send, + { + async move { + // First 4 bytes is the big endian encoded length field and the 5th byte is a + // PeerMessage tag + let mut length_buf = [0u8; 4]; - self - .read_exact(&mut length_buf) - .await - .map_err(PeerActorError::ReceiveFailed)?; + self + .read_exact(&mut length_buf) + .await + .map_err(PeerActorError::ReceiveFailed)?; - let length = u32::from_be_bytes(length_buf); + let length = u32::from_be_bytes(length_buf); - // Safety check -- BitTorrent docs do not specify if KeepAlive messages have an - // ID (and I'm pretty sure they don't) - if length == 0 { - return Ok(PeerMessages::KeepAlive); - } + // Safety check -- BitTorrent docs do not specify if KeepAlive messages have an + // ID (and I'm pretty sure they don't) + if length == 0 { + return Ok(PeerMessages::KeepAlive); + } - let frame_length = 4 + length as usize; - let mut message_buf = BytesMut::with_capacity(frame_length); - message_buf.extend_from_slice(&length_buf); + let frame_length = 4 + length as usize; + let mut message_buf = BytesMut::with_capacity(frame_length); + message_buf.extend_from_slice(&length_buf); - let mut message_type = [0u8; 1]; - self.read_exact(&mut message_type).await.map_err(|e| { - error!(error = %e, "Failed to read message type from peer"); - PeerActorError::ReceiveFailed(e) - })?; + let mut message_type = [0u8; 1]; + self.read_exact(&mut message_type).await.map_err(|e| { + error!(error = %e, "Failed to read message type from peer"); + PeerActorError::ReceiveFailed(e) + })?; - message_buf.extend_from_slice(&message_type); + message_buf.extend_from_slice(&message_type); - // Read the rest of the message payload - message_buf.resize(frame_length, 0); - self - .read_exact(&mut message_buf[5..]) - .await - .map_err(PeerActorError::ReceiveFailed)?; + // Read the rest of the message payload + message_buf.resize(frame_length, 0); + self + .read_exact(&mut message_buf[5..]) + .await + .map_err(PeerActorError::ReceiveFailed)?; - PeerMessages::from_bytes(message_buf.freeze()) + PeerMessages::from_bytes(message_buf.freeze()) + } } /// Receives a message from a peer with cancellation support, returning /// an error if the operation is cancelled - async fn recv_with_cancel( + fn recv_with_cancel( &mut self, token: CancellationToken, - ) -> Result { - tokio::select! { - _ = token.cancelled() => { - trace!("Receiving message from peer was cancelled"); - return Err(PeerActorError::MessageCancelled); - }, - result = self.recv() => { - result + ) -> impl Future> + Send + where + Self: Send, + { + async move { + tokio::select! { + _ = token.cancelled() => { + trace!("Receiving message from peer was cancelled"); + Err(PeerActorError::MessageCancelled) + }, + result = self.recv() => { + result + } } } } @@ -190,9 +208,9 @@ impl PeerStream { /// Sends a handshake to a peer. Returns nothing if the handshake is sent /// without error. pub async fn send_handshake( - &mut self, our_id: PeerId, info_hash: Arc, + &mut self, our_id: PeerId, info_hash: InfoHash, ) -> Result<(), PeerActorError> { - let handshake = Handshake::new(info_hash.clone(), our_id); + let handshake = Handshake::new(info_hash, our_id); self.write_all(&handshake.to_bytes()).await?; Ok(()) @@ -279,7 +297,6 @@ impl Display for PeerStream { } impl PeerSend for PeerStream {} -#[async_trait] impl PeerRecv for PeerStream { async fn recv(&mut self) -> Result { loop { @@ -493,12 +510,12 @@ impl PeerSend for PeerWriter {} /// Takes in a received handshake and returns the handshake we should respond /// with as well as the new peer. It preassigns the our_id to the peer. pub fn validate_handshake( - received_handshake: &Handshake, peer_addr: SocketAddr, info_hash: Arc, + received_handshake: &Handshake, peer_addr: SocketAddr, info_hash: InfoHash, ) -> Result<(), PeerActorError> { validate_handshake_protocol(received_handshake, peer_addr)?; // Validate info hash - if info_hash.clone() != received_handshake.info_hash { + if info_hash != received_handshake.info_hash { error!( peer_addr = %peer_addr, received_info_hash = %received_handshake.info_hash.to_hex(), @@ -507,7 +524,7 @@ pub fn validate_handshake( ); return Err(PeerActorError::HandshakeInfoHashMismatch { received: received_handshake.info_hash.to_hex(), - expected: info_hash.clone().to_hex(), + expected: info_hash.to_hex(), }); } @@ -612,11 +629,11 @@ mod tests { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); - let info_hash = Arc::new(Hash::new([1u8; 20])); + let info_hash = Hash::new([1u8; 20]); let client_id = PeerId::new(); // Spawn client that sends handshake - let client_info_hash = info_hash.clone(); + let client_info_hash = info_hash; let client = tokio::spawn(async move { let mut stream = PeerStream::tcp(TcpStream::connect(addr).await.unwrap()); @@ -647,17 +664,15 @@ mod tests { async fn peer_stream_when_frames_are_exchanged_then_counts_every_wire_byte() { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); - let info_hash = Arc::new(Hash::new([1u8; 20])); + let info_hash = Hash::new([1u8; 20]); let client_id = PeerId::new(); let server_id = PeerId::new(); - let handshake_len = Handshake::new(info_hash.clone(), client_id) - .to_bytes() - .len(); + let handshake_len = Handshake::new(info_hash, client_id).to_bytes().len(); let interested_len = PeerMessages::Interested.to_bytes().unwrap().len(); let piece = PeerMessages::Piece(2, 4, b"payload".as_slice().into()); let piece_len = piece.to_bytes().unwrap().len(); - let server_info_hash = info_hash.clone(); + let server_info_hash = info_hash; let server = tokio::spawn(async move { let (stream, _) = listener.accept().await.unwrap(); let mut stream = PeerStream::tcp(stream); @@ -720,13 +735,10 @@ mod tests { let mut stream = PeerStream::connect(local_peer.peer().socket_addr(), None) .await .unwrap(); - let info_hash = Arc::new(testing::test_info_hash()); + let info_hash = testing::test_info_hash(); let client_id = PeerId::new(); - stream - .send_handshake(client_id, info_hash.clone()) - .await - .unwrap(); + stream.send_handshake(client_id, info_hash).await.unwrap(); let (received_peer_id, _) = timeout(Duration::from_secs(1), stream.recv_handshake()) .await .expect("handshake should arrive before timeout") @@ -748,7 +760,7 @@ mod tests { let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); let addr = listener.local_addr().unwrap(); - let info_hash = Arc::new(Hash::new([1u8; 20])); + let info_hash = Hash::new([1u8; 20]); let mut handshake = Handshake::new(info_hash, PeerId::new()); handshake.protocol = "not bittorrent".into(); let handshake_bytes = handshake.to_bytes(); @@ -776,8 +788,8 @@ mod tests { #[test] fn validate_handshake_when_protocol_is_invalid_then_returns_magic_mismatch() { - let info_hash = Arc::new(Hash::new([1u8; 20])); - let mut handshake = Handshake::new(info_hash.clone(), PeerId::new()); + let info_hash = Hash::new([1u8; 20]); + let mut handshake = Handshake::new(info_hash, PeerId::new()); handshake.protocol = "not bittorrent".into(); let error = @@ -791,7 +803,7 @@ mod tests { #[test] fn validate_handshake_protocol_when_protocol_is_invalid_then_returns_magic_mismatch() { - let info_hash = Arc::new(Hash::new([1u8; 20])); + let info_hash = Hash::new([1u8; 20]); let mut handshake = Handshake::new(info_hash, PeerId::new()); handshake.protocol = "not bittorrent".into(); @@ -806,8 +818,8 @@ mod tests { #[test] fn validate_handshake_when_info_hash_differs_then_returns_info_hash_mismatch() { - let expected_info_hash = Arc::new(Hash::new([1u8; 20])); - let received_info_hash = Arc::new(Hash::new([2u8; 20])); + let expected_info_hash = Hash::new([1u8; 20]); + let received_info_hash = Hash::new([2u8; 20]); let handshake = Handshake::new(received_info_hash, PeerId::new()); let error = validate_handshake( diff --git a/crates/libtortillas/src/torrent/swarm.rs b/crates/libtortillas/src/torrent/swarm.rs index ea65daf5..9401ef0f 100644 --- a/crates/libtortillas/src/torrent/swarm.rs +++ b/crates/libtortillas/src/torrent/swarm.rs @@ -1,5 +1,3 @@ -use std::sync::Arc; - use futures::{StreamExt, stream}; use kameo::{ actor::{ActorRef, Spawn}, @@ -24,7 +22,7 @@ use crate::{ impl TorrentActor { #[instrument(skip(self, peer, stream), fields(%self, peer_addr = ?peer.socket_addr(), torrent_id = %self.info_hash()))] pub(super) fn append_peer(&self, mut peer: Peer, stream: Option) { - let info_hash = Arc::new(self.info_hash()); + let info_hash = self.info_hash(); let actor_ref = self.actor_ref.clone(); let our_id = self.id; let utp_server = self.utp_server.clone(); @@ -33,7 +31,7 @@ impl TorrentActor { let mut id = peer.id; let stream = match stream { Some(mut stream) => { - let handshake = Handshake::new(info_hash.clone(), our_id); + let handshake = Handshake::new(info_hash, our_id); if let Err(err) = stream.send(PeerMessages::Handshake(handshake)).await { debug!(error = %err, peer_addr = %peer.socket_addr(), "Failed to send handshake to peer"); return; @@ -43,34 +41,30 @@ impl TorrentActor { None => { let stream = PeerStream::connect(peer.socket_addr(), Some(utp_server)).await; match stream { - Ok(mut stream) => { - match stream.send_handshake(our_id, Arc::clone(&info_hash)).await { - Ok(_) => match stream.recv_handshake_message().await { - Ok(handshake) => { - if let Err(err) = validate_handshake( - &handshake, - peer.socket_addr(), - Arc::clone(&info_hash), - ) { - trace!(error = %err, peer_addr = %peer.socket_addr(), "Failed to validate peer handshake; exiting"); - return; - } - id = Some(handshake.peer_id); - peer.reserved = handshake.reserved; - peer.determine_supported().await; - stream - } - Err(err) => { - trace!(error = %err, peer_addr = %peer.socket_addr(), "Failed to receive handshake from peer; exiting"); + Ok(mut stream) => match stream.send_handshake(our_id, info_hash).await { + Ok(_) => match stream.recv_handshake_message().await { + Ok(handshake) => { + if let Err(err) = + validate_handshake(&handshake, peer.socket_addr(), info_hash) + { + trace!(error = %err, peer_addr = %peer.socket_addr(), "Failed to validate peer handshake; exiting"); return; } - }, + id = Some(handshake.peer_id); + peer.reserved = handshake.reserved; + peer.determine_supported().await; + stream + } Err(err) => { - trace!(error = %err, peer_addr = %peer.socket_addr(), "Failed to send handshake to peer; exiting"); + trace!(error = %err, peer_addr = %peer.socket_addr(), "Failed to receive handshake from peer; exiting"); return; } + }, + Err(err) => { + trace!(error = %err, peer_addr = %peer.socket_addr(), "Failed to send handshake to peer; exiting"); + return; } - } + }, Err(err) => { trace!(error = %err, peer_addr = %peer.socket_addr(), "Failed to connect to peer; exiting"); return; From 4a46d59ba5a42bfed3fd3aa4683d9e61b9d8e67a Mon Sep 17 00:00:00 2001 From: artrixdotdev Date: Mon, 27 Jul 2026 12:13:12 -0700 Subject: [PATCH 05/10] perf: cache torrent info hashes --- crates/libtortillas/src/torrent/actor.rs | 17 +++++++---------- crates/libtortillas/src/torrent/piece_flow.rs | 2 ++ 2 files changed, 9 insertions(+), 10 deletions(-) diff --git a/crates/libtortillas/src/torrent/actor.rs b/crates/libtortillas/src/torrent/actor.rs index 1fab81cb..2be17fb1 100644 --- a/crates/libtortillas/src/torrent/actor.rs +++ b/crates/libtortillas/src/torrent/actor.rs @@ -110,6 +110,7 @@ pub(crate) struct TorrentActor { pub(crate) bitfield: BitVec, pub(super) id: PeerId, + pub(super) info_hash: InfoHash, /// Metadata resolved from a magnet source. `.torrent` metadata remains /// canonical inside `metainfo`. pub(super) resolved_magnet_info: Option, @@ -166,16 +167,7 @@ impl TorrentActor { } pub fn info_hash(&self) -> InfoHash { - if let Some(info) = &self.info_dict() { - info.hash().expect("Failed to compute info hash") - } else { - match &self.metainfo { - MetaInfo::Torrent(t) => t.info.hash().expect("Failed to compute info hash"), - MetaInfo::MagnetUri(m) => m - .info_hash() - .expect("Magnet URIs should always have info hashes"), - } - } + self.info_hash } /// Checks if the torrent is empty (we haven't downloaded any pieces yet) by /// checking if our bitfield is filled with zeros. @@ -833,6 +825,7 @@ impl Actor for TorrentActor { utp_server, trackers, id: peer_id, + info_hash: torrent_id, metainfo, resolved_magnet_info: None, actor_ref: us, @@ -1697,6 +1690,7 @@ mod tests { trackers: HashMap::new(), bitfield, id: peer_id, + info_hash, resolved_magnet_info: None, metainfo: metainfo.clone(), tracker_server: udp_server.clone(), @@ -1856,6 +1850,7 @@ mod tests { trackers: HashMap::new(), bitfield, id: peer_id, + info_hash, resolved_magnet_info: None, metainfo: metainfo.clone(), tracker_server: udp_server, @@ -2047,6 +2042,7 @@ mod tests { _ => unreachable!(), }; let piece_count = info_dict.piece_count(); + let info_hash = metainfo.info_hash().unwrap(); let peer_id = testing::peer_id(); let udp_server = testing::udp_server().await; @@ -2076,6 +2072,7 @@ mod tests { trackers: HashMap::new(), bitfield: BitVec::repeat(false, piece_count), id: peer_id, + info_hash, resolved_magnet_info: None, metainfo, tracker_server: udp_server, diff --git a/crates/libtortillas/src/torrent/piece_flow.rs b/crates/libtortillas/src/torrent/piece_flow.rs index 72ebcf64..0c8e6f5f 100644 --- a/crates/libtortillas/src/torrent/piece_flow.rs +++ b/crates/libtortillas/src/torrent/piece_flow.rs @@ -477,6 +477,7 @@ mod tests { piece_storage: PieceStorageStrategy, base_path: std::path::PathBuf, ) -> TorrentActor { let info = test_info(); + let info_hash = info.hash().unwrap(); let metainfo = test_metainfo(info.clone()); let peer_id = testing::peer_id(); let tracker_server = UdpServer::new(None).await.unwrap(); @@ -506,6 +507,7 @@ mod tests { trackers: HashMap::new(), bitfield: BitVec::::repeat(false, info.piece_count()), id: peer_id, + info_hash, resolved_magnet_info: None, metainfo, tracker_server, From 8ee5a79dc9b50665802e26b0206c9b5aed6d11ce Mon Sep 17 00:00:00 2001 From: artrixdotdev Date: Mon, 27 Jul 2026 12:13:19 -0700 Subject: [PATCH 06/10] perf: avoid redundant tracker work --- crates/libtortillas/src/lib.rs | 2 ++ crates/libtortillas/src/metainfo/file.rs | 27 ++++++++++++++++++++---- crates/libtortillas/src/torrent/actor.rs | 7 ++++++ 3 files changed, 32 insertions(+), 4 deletions(-) diff --git a/crates/libtortillas/src/lib.rs b/crates/libtortillas/src/lib.rs index 863aa121..ad7f3d21 100644 --- a/crates/libtortillas/src/lib.rs +++ b/crates/libtortillas/src/lib.rs @@ -194,6 +194,7 @@ //! transport-agnostic live views and event streams. Durable state is //! represented by [`EngineSnapshot`](engine::EngineSnapshot) and //! [`TorrentSnapshot`](torrent::TorrentSnapshot), never by live views. + // `cfg!` type-checks both branches; this drops disabled live code before name // resolution. macro_rules! live_only { @@ -263,6 +264,7 @@ pub(crate) mod testing { pub(crate) const BIG_BUCK_BUNNY_NAME: &str = "Big Buck Bunny"; pub(crate) const BIG_BUCK_BUNNY_INFO_HASH: &str = "dd8255ecdc7ca55fb0bbf81323d87062db1f6d1c"; pub(crate) const BIG_BUCK_BUNNY_TORRENT_FILE: &str = "big-buck-bunny.torrent"; + pub(crate) const WIRED_CD_TORRENT_FILE: &str = "wired-cd.torrent"; pub(crate) const KNOPPIX_TORRENT_FILE: &str = "KNOPPIX_V9.1DVD-2021-01-25-EN.torrent"; pub(crate) fn fixture_path(relative_path: &str) -> PathBuf { diff --git a/crates/libtortillas/src/metainfo/file.rs b/crates/libtortillas/src/metainfo/file.rs index d4fb3f90..7ebc8b2b 100644 --- a/crates/libtortillas/src/metainfo/file.rs +++ b/crates/libtortillas/src/metainfo/file.rs @@ -40,10 +40,12 @@ impl TorrentFile { } pub fn announce_list(&self) -> Vec { - let mut announce_list: Vec = self.announce.clone().into_iter().collect(); - if let Some(list) = self.announce_list.clone() { - for tracker in list.into_iter().flatten() { - announce_list.push(tracker); + let mut announce_list: Vec = self.announce.iter().cloned().collect(); + if let Some(list) = &self.announce_list { + for tracker in list.iter().flatten() { + if !announce_list.contains(tracker) { + announce_list.push(tracker.clone()); + } } } announce_list @@ -185,4 +187,21 @@ mod tests { assert!(torrent.announce.is_none()); assert!(torrent.announce_list().is_empty()); } + + #[tokio::test] + async fn torrent_file_when_announce_is_repeated_then_returns_it_once() { + let metainfo = testing::read_torrent_fixture(testing::WIRED_CD_TORRENT_FILE).await; + let MetaInfo::Torrent(torrent) = metainfo else { + panic!("Expected Torrent"); + }; + let trackers = torrent.announce_list(); + + assert_eq!( + trackers + .iter() + .filter(|tracker| Some(*tracker) == torrent.announce.as_ref()) + .count(), + 1 + ); + } } diff --git a/crates/libtortillas/src/torrent/actor.rs b/crates/libtortillas/src/torrent/actor.rs index 2be17fb1..8c16dff4 100644 --- a/crates/libtortillas/src/torrent/actor.rs +++ b/crates/libtortillas/src/torrent/actor.rs @@ -763,6 +763,13 @@ impl Actor for TorrentActor { let tracker_list = metainfo.announce_list(); let mut trackers = HashMap::new(); for tracker in tracker_list { + if matches!(tracker, Tracker::Websocket(_)) { + warn!( + tracker_uri = %tracker.uri(), + "Skipping unsupported websocket tracker" + ); + continue; + } #[cfg(feature = "live")] let endpoint = tracker.redacted_endpoint(); #[cfg(feature = "live")] From 43744976e43b50b41a3bd43921d63e6085549b6c Mon Sep 17 00:00:00 2001 From: artrixdotdev Date: Mon, 27 Jul 2026 12:23:16 -0700 Subject: [PATCH 07/10] fix: validate peer frame inputs --- crates/libtortillas/examples/peer_transfer.rs | 4 +- crates/libtortillas/src/protocol/stream.rs | 51 ++++++++++++++++++- crates/libtortillas/src/torrent/actor.rs | 2 +- crates/libtortillas/src/tracker/model.rs | 1 - 4 files changed, 54 insertions(+), 4 deletions(-) diff --git a/crates/libtortillas/examples/peer_transfer.rs b/crates/libtortillas/examples/peer_transfer.rs index 01e9f225..56930094 100644 --- a/crates/libtortillas/examples/peer_transfer.rs +++ b/crates/libtortillas/examples/peer_transfer.rs @@ -25,7 +25,9 @@ async fn main() -> Result<(), Box> { for expected_index in 0..BLOCK_COUNT { match stream.recv().await? { PeerMessages::Piece(index, 0, block) - if index == expected_index as u32 && block.len() == BLOCK_LENGTH => + if index == expected_index as u32 + && block.len() == BLOCK_LENGTH + && block.iter().all(|byte| *byte == 0xa5) => { bytes_received += block.len(); } diff --git a/crates/libtortillas/src/protocol/stream.rs b/crates/libtortillas/src/protocol/stream.rs index 068606fc..ba84496d 100644 --- a/crates/libtortillas/src/protocol/stream.rs +++ b/crates/libtortillas/src/protocol/stream.rs @@ -26,6 +26,10 @@ use crate::{ peer::{MAGIC_STRING, PeerId, PeerState}, }; +/// BEP 3 piece blocks and BEP 9 metadata pieces are 16 KiB. This leaves room +/// for extension dictionaries and large bitfields while bounding peer input. +const MAX_FRAME_PAYLOAD_LENGTH: usize = 1024 * 1024; + enum PeerTransport { Tcp(TcpStream), Utp(UtpStream), @@ -108,6 +112,7 @@ pub trait PeerRecv: AsyncRead + Unpin { if length == 0 { return Ok(PeerMessages::KeepAlive); } + validate_frame_payload_length(length as usize)?; let frame_length = 4 + length as usize; let mut message_buf = BytesMut::with_capacity(frame_length); @@ -328,12 +333,15 @@ fn buffered_message(read_buffer: &mut BytesMut) -> Option Option Result<(), PeerActorError> { + if length > MAX_FRAME_PAYLOAD_LENGTH { + return Err(PeerActorError::ProtocolViolation(format!( + "peer frame payload length {length} exceeds maximum {MAX_FRAME_PAYLOAD_LENGTH}" + ))); + } + Ok(()) +} + impl AsyncRead for PeerStream { fn poll_read( mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>, @@ -611,6 +628,18 @@ mod tests { impl PeerSend for VectoredWriter {} + async fn stream_with_frame_length(length: u32) -> PeerStream { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let client = tokio::spawn(async move { + let mut stream = TcpStream::connect(address).await.unwrap(); + stream.write_all(&length.to_be_bytes()).await.unwrap(); + }); + let (stream, _) = listener.accept().await.unwrap(); + client.await.unwrap(); + PeerStream::tcp(stream) + } + #[tokio::test] async fn peer_send_when_message_is_piece_then_uses_vectored_write() { let message = PeerMessages::Piece(2, 4, Bytes::from_static(b"payload")); @@ -623,6 +652,26 @@ mod tests { assert_eq!(writer.bytes, expected); } + #[tokio::test] + async fn peer_reader_when_frame_payload_exceeds_limit_then_rejects_header() { + let stream = stream_with_frame_length((MAX_FRAME_PAYLOAD_LENGTH + 1) as u32).await; + let (mut reader, _) = stream.split(); + + let error = reader.recv().await.unwrap_err(); + + assert!(matches!(error, PeerActorError::ProtocolViolation(_))); + } + + #[test] + fn buffered_message_when_frame_payload_exceeds_limit_then_rejects_header() { + let length = (MAX_FRAME_PAYLOAD_LENGTH + 1) as u32; + let mut buffer = BytesMut::from(length.to_be_bytes().as_slice()); + + let message = buffered_message(&mut buffer).unwrap(); + + assert!(matches!(message, Err(PeerActorError::ProtocolViolation(_)))); + } + #[tokio::test] #[traced_test] async fn peer_stream_when_handshake_is_valid_then_returns_peer_id() { diff --git a/crates/libtortillas/src/torrent/actor.rs b/crates/libtortillas/src/torrent/actor.rs index 8c16dff4..fa02a3be 100644 --- a/crates/libtortillas/src/torrent/actor.rs +++ b/crates/libtortillas/src/torrent/actor.rs @@ -765,7 +765,7 @@ impl Actor for TorrentActor { for tracker in tracker_list { if matches!(tracker, Tracker::Websocket(_)) { warn!( - tracker_uri = %tracker.uri(), + tracker_uri = %tracker.redacted_endpoint(), "Skipping unsupported websocket tracker" ); continue; diff --git a/crates/libtortillas/src/tracker/model.rs b/crates/libtortillas/src/tracker/model.rs index 50fbaa33..2b3f30d8 100644 --- a/crates/libtortillas/src/tracker/model.rs +++ b/crates/libtortillas/src/tracker/model.rs @@ -119,7 +119,6 @@ impl Tracker { } } -#[cfg(feature = "live")] impl Tracker { /// Returns a credential-free endpoint label for public views. pub(crate) fn redacted_endpoint(&self) -> String { From 5a15b644c0a30bf9e40b88603fa615e16f5b507a Mon Sep 17 00:00:00 2001 From: artrixdotdev Date: Mon, 27 Jul 2026 12:54:25 -0700 Subject: [PATCH 08/10] perf: reduce live transfer overhead --- crates/libtortillas/examples/live.rs | 46 +++- crates/libtortillas/src/peer/actor.rs | 13 +- .../src/pieces/piece_scheduler.rs | 202 ++++++++++++++---- 3 files changed, 207 insertions(+), 54 deletions(-) diff --git a/crates/libtortillas/examples/live.rs b/crates/libtortillas/examples/live.rs index 675009cb..07039d8a 100644 --- a/crates/libtortillas/examples/live.rs +++ b/crates/libtortillas/examples/live.rs @@ -1,10 +1,15 @@ -use std::path::PathBuf; +use std::{ + path::PathBuf, + time::{Duration, Instant}, +}; use libtortillas::prelude::{ Engine, EngineEventKind, EventStreamError, TorrentEventKind, TorrentSource, TorrentState, }; use tracing::{error, info, warn}; +const METRICS_LOG_INTERVAL: Duration = Duration::from_secs(1); + #[tokio::main] async fn main() -> Result<(), Box> { tracing_subscriber::fmt() @@ -22,16 +27,43 @@ async fn main() -> Result<(), Box> { let engine = Engine::default(); let mut listener = engine.listener(); let event_task = tokio::spawn(async move { + let mut last_metrics_log = None; loop { match listener.recv().await { Ok(event) => { let view = listener.view(); - info!( - sequence = event.sequence, - torrent_count = view.torrent_count(), - ?event.kind, - "received an engine event" - ); + match &event.kind { + EngineEventKind::Torrent { + torrent, + event: TorrentEventKind::MetricsChanged(metrics), + } => { + let complete = metrics + .progress + .remaining_bytes + .is_some_and(|bytes| bytes.0 == 0); + let now = Instant::now(); + if complete + || last_metrics_log + .is_none_or(|last| now.duration_since(last) >= METRICS_LOG_INTERVAL) + { + last_metrics_log = Some(now); + info!( + sequence = event.sequence, + torrent_id = %torrent.info_hash(), + downloaded_bytes = metrics.traffic.totals.downloaded.0, + verified_bytes = metrics.progress.verified_bytes.0, + total_bytes = ?metrics.progress.total_bytes.map(|bytes| bytes.0), + "received torrent metrics" + ); + } + } + _ => info!( + sequence = event.sequence, + torrent_count = view.torrent_count(), + ?event.kind, + "received an engine event" + ), + } if matches!(event.kind, EngineEventKind::Shutdown(_)) { break; } diff --git a/crates/libtortillas/src/peer/actor.rs b/crates/libtortillas/src/peer/actor.rs index 4e6a1381..319a71af 100644 --- a/crates/libtortillas/src/peer/actor.rs +++ b/crates/libtortillas/src/peer/actor.rs @@ -584,6 +584,16 @@ impl Message for PeerActor { &mut self, msg: PeerMessages, _: &mut KameoContext, ) -> Self::Reply { self.peer.update_last_message_received(); + #[cfg(feature = "live")] + let publish_live_state = matches!( + &msg, + PeerMessages::Choke + | PeerMessages::Unchoke + | PeerMessages::Interested + | PeerMessages::NotInterested + | PeerMessages::Have(_) + | PeerMessages::Bitfield(_) + ); match msg { PeerMessages::Piece(index, offset, data) => { trace!( @@ -742,7 +752,8 @@ impl Message for PeerActor { warn!("Received unexpected handshake from peer"); } } - crate::live_only! { + #[cfg(feature = "live")] + if publish_live_state { let samples = self.live_handle.view().metrics.transfer.samples; self .live_handle diff --git a/crates/libtortillas/src/pieces/piece_scheduler.rs b/crates/libtortillas/src/pieces/piece_scheduler.rs index 187b1303..916d7dd8 100644 --- a/crates/libtortillas/src/pieces/piece_scheduler.rs +++ b/crates/libtortillas/src/pieces/piece_scheduler.rs @@ -25,8 +25,10 @@ impl BlockRequest { pub(crate) struct PieceScheduler { completed_pieces: BitVec, completed_blocks: HashMap, - in_flight: HashMap<(usize, usize), InFlightBlock>, + in_flight: HashMap>>, + in_flight_per_peer: HashMap, peer_availability: HashMap>>, + request_cursors: HashMap, next_piece: usize, } @@ -42,7 +44,9 @@ impl PieceScheduler { completed_pieces: BitVec::repeat(false, piece_count), completed_blocks: HashMap::new(), in_flight: HashMap::new(), + in_flight_per_peer: HashMap::new(), peer_availability: HashMap::new(), + request_cursors: HashMap::new(), next_piece: 0, } } @@ -82,10 +86,9 @@ impl PieceScheduler { if block_index < blocks.len() { blocks.set(block_index, true); } - self - .in_flight - .remove(&(piece_index, block_index)) - .map(|request| request.peer_id) + let request = self.remove_in_flight(piece_index, block_index)?; + self.decrement_in_flight(request.peer_id); + Some(request.peer_id) } pub(crate) fn remove_piece_blocks(&mut self, piece_index: usize) -> Option { @@ -118,18 +121,40 @@ impl PieceScheduler { return requests; } + let piece_count = self.completed_pieces.len(); + if piece_count == 0 { + return requests; + } + let Some(available_pieces) = self.peer_availability.get(&peer_id) else { return requests; }; - let last_piece_index = self.completed_pieces.len().saturating_sub(1); + let blocks_per_piece = piece_length.div_ceil(BLOCK_SIZE); + if blocks_per_piece == 0 { + return requests; + } + + let last_piece_index = piece_count - 1; let last_piece_len = if total_length.is_multiple_of(piece_length) { piece_length } else { total_length % piece_length }; - - for piece_index in self.next_piece..self.completed_pieces.len() { + let first_slot = self.next_piece.saturating_mul(blocks_per_piece); + let total_slots = last_piece_index + .saturating_mul(blocks_per_piece) + .saturating_add(last_piece_len.div_ceil(BLOCK_SIZE)); + let cursor = self + .request_cursors + .get(&peer_id) + .copied() + .unwrap_or(first_slot) + .clamp(first_slot, total_slots); + + for slot in (cursor..total_slots).chain(first_slot..cursor) { + let piece_index = slot / blocks_per_piece; + let block_index = slot % blocks_per_piece; if self.completed_pieces[piece_index] || !available_pieces .get(piece_index) @@ -147,30 +172,41 @@ impl PieceScheduler { piece_length }; let total_blocks = piece_len.div_ceil(BLOCK_SIZE); + if block_index >= total_blocks { + continue; + } - for block_index in 0..total_blocks { - let key = (piece_index, block_index); - if self.in_flight.contains_key(&key) - || self - .completed_blocks - .get(&piece_index) - .and_then(|blocks| blocks.get(block_index).as_deref().copied()) - .unwrap_or(false) - { - continue; - } + if self.is_in_flight(piece_index, block_index) + || self + .completed_blocks + .get(&piece_index) + .and_then(|blocks| blocks.get(block_index).as_deref().copied()) + .unwrap_or(false) + { + continue; + } - self.in_flight.insert( - key, - InFlightBlock { - peer_id, - requested_at: Instant::now(), - }, - ); - requests.push(self.block_request(piece_index, block_index, piece_len)); - if requests.len() >= limit { - return requests; - } + let blocks = self.in_flight.entry(piece_index).or_default(); + if blocks.len() <= block_index { + blocks.resize_with(block_index + 1, || None); + } + blocks[block_index] = Some(InFlightBlock { + peer_id, + requested_at: Instant::now(), + }); + *self.in_flight_per_peer.entry(peer_id).or_default() += 1; + let next_cursor = slot.saturating_add(1); + self.request_cursors.insert( + peer_id, + if next_cursor < total_slots { + next_cursor + } else { + first_slot + }, + ); + requests.push(self.block_request(piece_index, block_index, piece_len)); + if requests.len() >= limit { + return requests; } } @@ -178,11 +214,7 @@ impl PieceScheduler { } pub(crate) fn in_flight_for_peer(&self, peer_id: PeerId) -> usize { - self - .in_flight - .values() - .filter(|request| request.peer_id == peer_id) - .count() + self.in_flight_per_peer.get(&peer_id).copied().unwrap_or(0) } pub(crate) fn update_peer_availability( @@ -192,31 +224,56 @@ impl PieceScheduler { } pub(crate) fn peer_disconnected(&mut self, peer_id: PeerId) { - self - .in_flight - .retain(|_, request| request.peer_id != peer_id); + self.in_flight.retain(|_, blocks| { + for block in blocks.iter_mut() { + if block + .as_ref() + .is_some_and(|request| request.peer_id == peer_id) + { + *block = None; + } + } + blocks.iter().any(Option::is_some) + }); + self.in_flight_per_peer.remove(&peer_id); self.peer_availability.remove(&peer_id); + self.request_cursors.remove(&peer_id); } pub(crate) fn release_stale_requests(&mut self, timeout: Duration) -> usize { - let before = self.in_flight.len(); let now = Instant::now(); - self - .in_flight - .retain(|_, request| now.saturating_duration_since(request.requested_at) < timeout); - before.saturating_sub(self.in_flight.len()) + let mut released = 0; + let mut in_flight_per_peer = HashMap::new(); + self.in_flight.retain(|_, blocks| { + for block in blocks.iter_mut() { + let Some(request) = block else { + continue; + }; + if now.saturating_duration_since(request.requested_at) >= timeout { + *block = None; + released += 1; + } else { + *in_flight_per_peer.entry(request.peer_id).or_default() += 1; + } + } + blocks.iter().any(Option::is_some) + }); + self.in_flight_per_peer = in_flight_per_peer; + released } pub(crate) fn release_peer_request( &mut self, peer_id: PeerId, piece_index: usize, offset: usize, ) { - let key = (piece_index, offset / BLOCK_SIZE); if self .in_flight - .get(&key) + .get(&piece_index) + .and_then(|blocks| blocks.get(offset / BLOCK_SIZE)) + .and_then(Option::as_ref) .is_some_and(|request| request.peer_id == peer_id) { - self.in_flight.remove(&key); + self.remove_in_flight(piece_index, offset / BLOCK_SIZE); + self.decrement_in_flight(peer_id); } } @@ -240,6 +297,39 @@ impl PieceScheduler { length, } } + + fn decrement_in_flight(&mut self, peer_id: PeerId) { + let remove = self + .in_flight_per_peer + .get_mut(&peer_id) + .is_some_and(|count| { + *count = count.saturating_sub(1); + *count == 0 + }); + if remove { + self.in_flight_per_peer.remove(&peer_id); + } + } + + fn is_in_flight(&self, piece_index: usize, block_index: usize) -> bool { + self + .in_flight + .get(&piece_index) + .and_then(|blocks| blocks.get(block_index)) + .is_some_and(Option::is_some) + } + + fn remove_in_flight(&mut self, piece_index: usize, block_index: usize) -> Option { + let (request, piece_has_requests) = { + let blocks = self.in_flight.get_mut(&piece_index)?; + let request = blocks.get_mut(block_index)?.take()?; + (request, blocks.iter().any(Option::is_some)) + }; + if !piece_has_requests { + self.in_flight.remove(&piece_index); + } + Some(request) + } } #[cfg(test)] @@ -277,6 +367,26 @@ mod tests { assert_eq!(scheduler.in_flight_for_peer(peer_id), 0); } + #[test] + fn scheduler_when_request_is_released_then_wraps_cursor() { + let peer_id = PeerId::Unknown([5; 20]); + let mut scheduler = PieceScheduler::new(2); + scheduler.update_peer_availability(peer_id, Arc::new([true, false].into_iter().collect())); + assert_eq!( + scheduler + .requests_for_peer(peer_id, 1, BLOCK_SIZE, BLOCK_SIZE * 2) + .len(), + 1 + ); + + scheduler.release_peer_request(peer_id, 0, 0); + let requests = scheduler.requests_for_peer(peer_id, 1, BLOCK_SIZE, BLOCK_SIZE * 2); + + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].piece_index, 0); + assert_eq!(requests[0].block_index, 0); + } + #[test] fn late_rejection_does_not_release_reassigned_request() { let original_peer = PeerId::Unknown([3; 20]); From 16ef5a56df103c9978d71b3f08a8ba7a83319f3a Mon Sep 17 00:00:00 2001 From: artrixdotdev Date: Mon, 27 Jul 2026 13:17:48 -0700 Subject: [PATCH 09/10] test: avoid duplicate tracker assumption --- crates/libtortillas/src/tracker/http.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/crates/libtortillas/src/tracker/http.rs b/crates/libtortillas/src/tracker/http.rs index d073705e..f8921889 100644 --- a/crates/libtortillas/src/tracker/http.rs +++ b/crates/libtortillas/src/tracker/http.rs @@ -605,8 +605,10 @@ mod tests { let announce_list = file.announce_list(); println!("announce_list: {:?}", announce_list); - // An HTTP tracker - let announce_uri = announce_list[1].uri(); + let announce_uri = announce_list + .first() + .expect("fixture should contain an HTTP tracker") + .uri(); let http_tracker = HttpTracker::new(announce_uri, info_hash.unwrap(), None, None); http_tracker .update(TrackerUpdate::Left(file.info.total_length())) From 52ef8d29f32efb7c28c0961c4978701be492fbbb Mon Sep 17 00:00:00 2001 From: artrixdotdev Date: Mon, 27 Jul 2026 14:05:17 -0700 Subject: [PATCH 10/10] fix: handle completed scheduler cursor --- .../src/pieces/piece_scheduler.rs | 48 ++++++++++++++++--- 1 file changed, 42 insertions(+), 6 deletions(-) diff --git a/crates/libtortillas/src/pieces/piece_scheduler.rs b/crates/libtortillas/src/pieces/piece_scheduler.rs index 916d7dd8..bbcd90ba 100644 --- a/crates/libtortillas/src/pieces/piece_scheduler.rs +++ b/crates/libtortillas/src/pieces/piece_scheduler.rs @@ -145,12 +145,16 @@ impl PieceScheduler { let total_slots = last_piece_index .saturating_mul(blocks_per_piece) .saturating_add(last_piece_len.div_ceil(BLOCK_SIZE)); - let cursor = self - .request_cursors - .get(&peer_id) - .copied() - .unwrap_or(first_slot) - .clamp(first_slot, total_slots); + let cursor = if self.next_piece == piece_count { + total_slots + } else { + self + .request_cursors + .get(&peer_id) + .copied() + .unwrap_or(first_slot) + .clamp(first_slot, total_slots) + }; for slot in (cursor..total_slots).chain(first_slot..cursor) { let piece_index = slot / blocks_per_piece; @@ -387,6 +391,38 @@ mod tests { assert_eq!(requests[0].block_index, 0); } + #[test] + fn scheduler_when_all_pieces_are_complete_then_returns_no_requests() { + let peer_id = PeerId::Unknown([6; 20]); + let piece_length = BLOCK_SIZE * 2; + let mut scheduler = PieceScheduler::new(2); + scheduler.update_peer_availability(peer_id, Arc::new([true, true].into_iter().collect())); + scheduler.mark_piece_complete(0); + scheduler.mark_piece_complete(1); + + let requests = + scheduler.requests_for_peer(peer_id, 1, piece_length, piece_length + BLOCK_SIZE / 2); + + assert!(requests.is_empty()); + } + + #[test] + fn scheduler_when_final_piece_is_partial_then_requests_available_block() { + let peer_id = PeerId::Unknown([7; 20]); + let piece_length = BLOCK_SIZE * 2; + let mut scheduler = PieceScheduler::new(2); + scheduler.update_peer_availability(peer_id, Arc::new([true, true].into_iter().collect())); + scheduler.mark_piece_complete(0); + + let requests = + scheduler.requests_for_peer(peer_id, 2, piece_length, piece_length + BLOCK_SIZE / 2); + + assert_eq!(requests.len(), 1); + assert_eq!(requests[0].piece_index, 1); + assert_eq!(requests[0].block_index, 0); + assert_eq!(requests[0].length, BLOCK_SIZE / 2); + } + #[test] fn late_rejection_does_not_release_reassigned_request() { let original_peer = PeerId::Unknown([3; 20]);