feat: Support in-file piece storage

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
    }

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<Bytes> {
+      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<Bytes> {
+      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<Vec<(PathBuf, usize, usize)>> {
+      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<Bytes> {
+      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();
+   }
 }