diff --git a/crates/libtortillas/src/lib.rs b/crates/libtortillas/src/lib.rs index b869f661..66e9ebff 100644 --- a/crates/libtortillas/src/lib.rs +++ b/crates/libtortillas/src/lib.rs @@ -16,6 +16,7 @@ pub(crate) mod testing { use tokio::fs::read_to_string; use crate::{ + hashes::Hash, metainfo::{MagnetUri, MetaInfo, TorrentFile}, peer::PeerId, tracker::udp::UdpServer, @@ -82,6 +83,14 @@ pub(crate) mod testing { PeerId::default() } + pub(crate) fn piece_hash(data: &[u8]) -> Hash<20> { + use sha1::{Digest, Sha1}; + + let mut hasher = Sha1::new(); + hasher.update(data); + Hash::from_bytes(hasher.finalize().into()) + } + pub(crate) async fn udp_server() -> UdpServer { UdpServer::new(None).await } diff --git a/crates/libtortillas/src/pieces/piece_manager.rs b/crates/libtortillas/src/pieces/piece_manager.rs index be7a1012..4ff290a7 100644 --- a/crates/libtortillas/src/pieces/piece_manager.rs +++ b/crates/libtortillas/src/pieces/piece_manager.rs @@ -5,7 +5,7 @@ use async_trait::async_trait; use bytes::Bytes; use tokio::{ fs::{OpenOptions, create_dir_all}, - io::{AsyncSeekExt, AsyncWriteExt}, + io::{AsyncReadExt, AsyncSeekExt, AsyncWriteExt}, }; use tracing::trace; @@ -140,38 +140,72 @@ impl FilePieceManager { pub fn path(&self) -> Option<&PathBuf> { self.0.as_ref() } -} -#[async_trait] -impl PieceManager for FilePieceManager { - fn info(&self) -> Option<&Info> { - self.1.as_ref() + pub(crate) async fn write_block( + &self, index: usize, offset: usize, block: Bytes, + ) -> anyhow::Result<()> { + self.write_piece_range(index, offset, &block).await } - async fn pre_start(&mut self, info_dict: Info) -> anyhow::Result<()> { - let base_path = self - .0 - .as_ref() - .ok_or_else(|| anyhow::anyhow!("path must be set before pre_start"))?; + /// Reads a complete piece back from the final torrent files. + /// + /// This is used by `PieceStorageStrategy::InFile` after all blocks for a + /// piece have arrived, so the torrent actor can hash the exact bytes that + /// were written to disk. + pub(crate) async fn read_piece(&self, index: usize) -> anyhow::Result { + let piece_bounds = self.piece_to_paths(index)?; + let piece_len = piece_bounds.iter().map(|(_, _, len)| len).sum(); + self.read_piece_range(index, 0, piece_len).await + } - let info_hash = info_dict.hash()?; + pub(crate) async fn read_piece_block( + &self, index: usize, offset: usize, length: usize, + ) -> anyhow::Result { + self.read_piece_range(index, offset, length).await + } - trace!(torrent_id = %info_hash, path = %base_path.display(), "Pre-starting piece manager"); + /// Converts a piece-relative byte range into file-relative segments. + /// + /// A single BitTorrent piece can cross file boundaries in multi-file + /// torrents. This helper centralizes that mapping so in-file reads and + /// writes cannot drift apart. + fn piece_range_segments( + &self, index: usize, offset: usize, length: usize, + ) -> anyhow::Result> { + let mut remaining = length; + let mut skipped = offset; + let mut segments = Vec::new(); + + if remaining == 0 { + return Ok(segments); + } - self.1 = Some(info_dict); - Ok(()) + for (path, file_offset, len) in self.piece_to_paths(index)? { + if skipped >= len { + skipped -= len; + continue; + } + + let segment_len = remaining.min(len - skipped); + segments.push((path, file_offset + skipped, segment_len)); + remaining -= segment_len; + skipped = 0; + + if remaining == 0 { + return Ok(segments); + } + } + + anyhow::bail!("piece range exceeds file lengths") } - #[allow(unused)] - async fn recv(&self, index: usize, data: Bytes) -> anyhow::Result<()> { + async fn write_piece_range( + &self, index: usize, offset: usize, data: &[u8], + ) -> anyhow::Result<()> { let base_path = self.path().ok_or_else(|| anyhow::anyhow!("path not set"))?; - let _info = self.info().ok_or_else(|| anyhow::anyhow!("info not set"))?; + let mut data_offset = 0; - let piece_bounds = self.piece_to_paths(index)?; - let mut data_offset = 0; // Track how much of `data` we've consumed - - for (path, file_offset, len) in piece_bounds { - // Ensure parent directories exist + for (path, file_offset, len) in self.piece_range_segments(index, offset, data.len())? { let full_path = base_path.join(&path); if let Some(parent) = full_path.parent() { create_dir_all(parent).await?; @@ -183,28 +217,87 @@ impl PieceManager for FilePieceManager { .truncate(false) .open(&full_path) .await?; - - // Position file correctly file.seek(SeekFrom::Start(file_offset as u64)).await?; - - // Write correct slice from the piece buffer file .write_all(&data[data_offset..data_offset + len]) .await?; data_offset += len; + } - trace!(index, offset = file_offset, len, path = %path.display(), "Wrote piece block to file"); + Ok(()) + } + + async fn read_piece_range( + &self, index: usize, offset: usize, length: usize, + ) -> anyhow::Result { + let base_path = self.path().ok_or_else(|| anyhow::anyhow!("path not set"))?; + let mut data = Vec::with_capacity(length); + + for (path, file_offset, len) in self.piece_range_segments(index, offset, length)? { + let full_path = base_path.join(&path); + let mut file = OpenOptions::new().read(true).open(&full_path).await?; + file.seek(SeekFrom::Start(file_offset as u64)).await?; + + let start = data.len(); + data.resize(start + len, 0); + file.read_exact(&mut data[start..]).await?; } + Ok(data.into()) + } +} + +#[async_trait] +impl PieceManager for FilePieceManager { + fn info(&self) -> Option<&Info> { + self.1.as_ref() + } + + async fn pre_start(&mut self, info_dict: Info) -> anyhow::Result<()> { + let base_path = self + .0 + .as_ref() + .ok_or_else(|| anyhow::anyhow!("path must be set before pre_start"))?; + + let info_hash = info_dict.hash()?; + + trace!(torrent_id = %info_hash, path = %base_path.display(), "Pre-starting piece manager"); + + self.1 = Some(info_dict); Ok(()) } + + #[allow(unused)] + async fn recv(&self, index: usize, data: Bytes) -> anyhow::Result<()> { + let _info = self.info().ok_or_else(|| anyhow::anyhow!("info not set"))?; + self.write_piece_range(index, 0, &data).await + } } #[cfg(test)] mod tests { use super::*; - use crate::{prelude::MetaInfo, testing}; + use crate::{hashes::HashVec, prelude::MetaInfo, testing}; + + fn single_file_info() -> Info { + Info { + name: "data.bin".to_string(), + piece_length: 4, + pieces: HashVec::from(vec![ + testing::piece_hash(b"abcd"), + testing::piece_hash(b"ef"), + ]), + file: InfoKeys::Single { + length: 6, + md5sum: None, + }, + is_private: None, + publisher: None, + publisher_url: None, + source: None, + } + } #[tokio::test] async fn file_piece_manager_when_mapping_first_piece_then_returns_expected_file() { @@ -225,4 +318,30 @@ mod tests { // The first file should always be "Big Buck Bunny.en.srt" assert_eq!(file_path.extension().unwrap(), "srt"); } + + #[tokio::test] + async fn file_piece_manager_when_writing_blocks_then_reads_piece_back() { + let base_path = testing::torrent_temp_path(); + let manager = FilePieceManager(Some(base_path.clone()), Some(single_file_info())); + + manager + .write_block(0, 0, Bytes::from_static(b"ab")) + .await + .unwrap(); + manager + .write_block(0, 2, Bytes::from_static(b"cd")) + .await + .unwrap(); + + assert_eq!( + manager.read_piece(0).await.unwrap(), + Bytes::from_static(b"abcd") + ); + assert_eq!( + manager.read_piece_block(0, 1, 2).await.unwrap(), + Bytes::from_static(b"bc") + ); + + tokio::fs::remove_dir_all(base_path).await.unwrap(); + } } diff --git a/crates/libtortillas/src/torrent/piece_flow.rs b/crates/libtortillas/src/torrent/piece_flow.rs index cba8bd81..92ed1d5b 100644 --- a/crates/libtortillas/src/torrent/piece_flow.rs +++ b/crates/libtortillas/src/torrent/piece_flow.rs @@ -7,12 +7,12 @@ use tokio::{ }; use tracing::{debug, info, trace, warn}; -use super::TorrentActor; +use super::{TorrentActor, util}; use crate::{ errors::TorrentError, peer::commands::{CancelPiece, Have, NeedPiece}, pieces::{PieceManager, ValidateAndRead, WriteBlock}, - torrent::{BLOCK_SIZE, PieceStorageStrategy, TorrentState}, + torrent::{BLOCK_SIZE, PieceStorageStrategy, TorrentState, actor::PieceManagerProxy}, tracker::{Announce, Event, TrackerUpdate}, }; @@ -157,6 +157,8 @@ impl TorrentActor { async fn write_block_to_storage(&self, index: usize, offset: usize, block: Bytes) -> bool { match &self.piece_storage { PieceStorageStrategy::Disk(_) => { + // Disk storage is the piece-cache strategy: blocks are assembled in + // standalone `.piece` files before being flushed to output files. let Ok(path) = self.get_piece_path(index) else { warn!(piece_index = index, "Failed to get piece path"); return false; @@ -182,13 +184,31 @@ impl TorrentActor { } } } - PieceStorageStrategy::InFile => { - warn!( - piece_index = index, - offset, "In-file block writes are not implemented yet" - ); - false - } + PieceStorageStrategy::InFile => match &self.piece_manager { + // In-file storage skips the `.piece` cache and writes the incoming + // block straight to its final file offset. + PieceManagerProxy::Default(manager) => { + match manager.write_block(index, offset, block).await { + Ok(()) => true, + Err(err) => { + warn!( + ?err, + piece_index = index, + offset, + "Failed to write in-file block" + ); + false + } + } + } + PieceManagerProxy::Custom(_) => { + warn!( + piece_index = index, + offset, "In-file storage requires the default piece manager" + ); + false + } + }, } } @@ -245,6 +265,8 @@ impl TorrentActor { match &self.piece_storage { PieceStorageStrategy::Disk(_) => { + // Cached pieces are validated as standalone files, then handed to + // the piece manager to populate the final output files. let Ok(path) = self.get_piece_path(index) else { warn!(index, "Failed to get piece path; re-requesting"); if let Some(blocks) = previous_blocks.as_ref() { @@ -286,11 +308,41 @@ impl TorrentActor { } } PieceStorageStrategy::InFile => { - warn!( - index, - "In-file piece validation is not implemented yet; re-requesting" - ); - return false; + // Blocks were already written into output files, so validation must + // hash the piece bytes read back from those files. + let PieceManagerProxy::Default(manager) = &self.piece_manager else { + warn!( + index, + "In-file storage requires the default piece manager; re-requesting" + ); + return false; + }; + + let data = match manager.read_piece(index).await { + Ok(data) => data, + Err(err) => { + warn!(?err, index, "Failed to read in-file piece; re-requesting"); + if let Some(blocks) = previous_blocks.as_ref() { + self + .piece_scheduler + .restore_piece_blocks(index, blocks.clone()); + } + self.request_blocks_from_peer(peer_id, 1).await; + return false; + } + }; + + if let Err(err) = util::validate_piece_bytes(&data, info_dict.pieces[index]) { + warn!( + ?err, + index, "Failed to validate in-file piece; re-requesting" + ); + if let Some(blocks) = previous_blocks { + self.piece_scheduler.restore_piece_blocks(index, blocks); + } + self.request_blocks_from_peer(peer_id, 1).await; + return false; + } } } true @@ -322,15 +374,196 @@ impl TorrentActor { pub(super) async fn read_piece_block( &self, index: usize, offset: usize, length: usize, ) -> anyhow::Result { - anyhow::ensure!( - matches!(self.piece_storage, PieceStorageStrategy::Disk(_)), - "in-file piece reads are not implemented yet" + match &self.piece_storage { + PieceStorageStrategy::Disk(_) => { + let mut file = File::open(self.get_piece_path(index)?).await?; + let mut buffer = vec![0; length]; + file.seek(SeekFrom::Start(offset as u64)).await?; + file.read_exact(&mut buffer).await?; + Ok(buffer.into()) + } + PieceStorageStrategy::InFile => match &self.piece_manager { + PieceManagerProxy::Default(manager) => { + manager.read_piece_block(index, offset, length).await + } + PieceManagerProxy::Custom(_) => { + anyhow::bail!("in-file reads require the default piece manager") + } + }, + } + } +} + +#[cfg(test)] +mod tests { + use std::{collections::HashMap, sync::atomic::AtomicU8}; + + use bitvec::vec::BitVec; + use bytes::Bytes; + use kameo::actor::Spawn; + use kameo_actors::scheduler::Scheduler; + use librqbit_utp::UtpSocket; + + use super::*; + use crate::{ + hashes::HashVec, + metainfo::{Info, InfoKeys, MetaInfo, TorrentFile}, + peer::PeerId, + pieces::{FilePieceManager, PieceScheduler, PieceStoreActor}, + testing, + torrent::actor::{PieceManagerProxy, TorrentActorArgs}, + tracker::{Tracker, udp::UdpServer}, + }; + + fn test_info() -> Info { + Info { + name: "data.bin".to_string(), + piece_length: 4, + pieces: HashVec::from(vec![testing::piece_hash(b"abcd")]), + file: InfoKeys::Single { + length: 4, + md5sum: None, + }, + is_private: None, + publisher: None, + publisher_url: None, + source: None, + } + } + + fn test_metainfo(info: Info) -> MetaInfo { + MetaInfo::Torrent(TorrentFile { + announce: Tracker::Http("http://127.0.0.1/announce".to_string()), + announce_list: None, + comment: None, + created_by: None, + creation_date: None, + encoding: None, + info, + url_list: None, + }) + } + + async fn test_actor( + piece_storage: PieceStorageStrategy, base_path: std::path::PathBuf, + ) -> TorrentActor { + let info = test_info(); + let metainfo = test_metainfo(info.clone()); + let peer_id = testing::peer_id(); + let tracker_server = UdpServer::new(None).await; + let utp_server = UtpSocket::new_udp(testing::ephemeral_socket_addr()) + .await + .unwrap(); + + let actor_ref = TorrentActor::spawn(TorrentActorArgs { + peer_id, + metainfo: metainfo.clone(), + utp_server: utp_server.clone(), + tracker_server: tracker_server.clone(), + primary_addr: None, + piece_storage: piece_storage.clone(), + autostart: Some(false), + sufficient_peers: Some(usize::MAX), + base_path: Some(base_path.clone()), + }); + + TorrentActor { + peers: HashMap::new(), + trackers: HashMap::new(), + bitfield: BitVec::::repeat(false, info.piece_count()), + id: peer_id, + info: Some(info.clone()), + metainfo, + tracker_server, + scheduler: Scheduler::spawn(Scheduler::new()), + utp_server, + actor_ref, + piece_storage, + piece_store: PieceStoreActor::spawn(()), + piece_manager: PieceManagerProxy::Default(FilePieceManager(Some(base_path), Some(info))), + state: TorrentState::Downloading, + piece_scheduler: PieceScheduler::new(1), + start_time: None, + sufficient_peers: 6, + autostart: false, + pending_start: false, + ready_hook: Vec::new(), + } + } + + #[tokio::test] + async fn torrent_actor_when_using_in_file_storage_then_downloads_and_reads_back_piece() { + let base_path = testing::torrent_temp_path(); + let mut actor = test_actor(PieceStorageStrategy::InFile, base_path.clone()).await; + + assert!( + actor + .write_block_to_storage(0, 0, Bytes::from_static(b"ab")) + .await + ); + assert!( + actor + .write_block_to_storage(0, 2, Bytes::from_static(b"cd")) + .await + ); + assert!( + actor + .validate_and_send_piece(PeerId::default(), 0, None) + .await + ); + + actor.bitfield.set_aliased(0, true); + assert_eq!( + actor.read_piece_block(0, 1, 2).await.unwrap(), + Bytes::from_static(b"bc") + ); + + actor.actor_ref.kill(); + actor.piece_store.kill(); + actor.scheduler.kill(); + tokio::fs::remove_dir_all(base_path).await.unwrap(); + } + + #[tokio::test] + async fn torrent_actor_when_using_disk_storage_then_keeps_piece_cache_and_writes_output_file() { + let piece_path = testing::torrent_temp_path(); + let base_path = testing::torrent_temp_path(); + let mut actor = test_actor( + PieceStorageStrategy::Disk(piece_path.clone()), + base_path.clone(), + ) + .await; + + assert!( + actor + .write_block_to_storage(0, 0, Bytes::from_static(b"ab")) + .await + ); + assert!( + actor + .write_block_to_storage(0, 2, Bytes::from_static(b"cd")) + .await + ); + assert!( + actor + .validate_and_send_piece(PeerId::default(), 0, None) + .await + ); + + actor.bitfield.set_aliased(0, true); + assert_eq!( + actor.read_piece_block(0, 1, 2).await.unwrap(), + Bytes::from_static(b"bc") + ); + assert_eq!( + tokio::fs::read(base_path.join("data.bin")).await.unwrap(), + b"abcd" ); - let mut file = File::open(self.get_piece_path(index)?).await?; - let mut buffer = vec![0; length]; - file.seek(SeekFrom::Start(offset as u64)).await?; - file.read_exact(&mut buffer).await?; - Ok(buffer.into()) + actor.actor_ref.kill(); + actor.piece_store.kill(); + actor.scheduler.kill(); + tokio::fs::remove_dir_all(piece_path).await.unwrap(); + tokio::fs::remove_dir_all(base_path).await.unwrap(); } } diff --git a/crates/libtortillas/src/torrent/util.rs b/crates/libtortillas/src/torrent/util.rs index 7fb1675d..78be9d52 100644 --- a/crates/libtortillas/src/torrent/util.rs +++ b/crates/libtortillas/src/torrent/util.rs @@ -120,6 +120,19 @@ pub async fn validate_piece_file( Ok(()) } +pub fn validate_piece_bytes(data: &[u8], hash: Hash<20>) -> anyhow::Result<()> { + let mut hasher = Sha1::new(); + hasher.update(data); + let piece_hash = Hash::from_bytes(hasher.finalize().into()); + + ensure!( + piece_hash == hash, + "Hashed bytes were not equal to hash from info dict" + ); + + Ok(()) +} + /// Creates a dir if it doesn't exist. This function can be called even if the /// directory exists -- if it does, nothing will happen. pub async fn create_dir(path: &PathBuf) -> Result<(), Error> {