read piece

src/net/_session/disk.rs

@@ -1,4 +1,5 @@
 use std::cmp;
+use std::fs::File;
 use std::io::{self, Read, Seek, SeekFrom, Write};
 use std::sync::Arc;
 
@@ -10,18 +11,73 @@ pub struct DiskManager<F> where F: Read + Seek + Write {
 }
 
 impl<F> DiskManager<F> where F: Read + Seek + Write {
-    pub fn write_piece(&mut self, info_hash: &Hash, index: u32, piece: Vec<u8>) -> io::Result<()> {
+
+    pub fn write_piece(&mut self, info_hash: &Hash, index: u32, mut piece: Vec<u8>) -> io::Result<()> {
         let torrent = self.torrents.get_mut(&info_hash).unwrap();
 
+        torrent.iter_files_for_piece(index, &mut piece, |file, seek, buf| {
+            file.seek(SeekFrom::Start(seek))?;
+            file.write_all(&buf)?;
+            Ok(())
+        })?;
+
+        Ok(())
+    }
+
+    pub fn read_piece(&mut self, info_hash: &Hash, index: u32) -> io::Result<Vec<u8>> {
+        let torrent = self.torrents.get_mut(&info_hash).unwrap();
+
+        let len = torrent.metainfo.piece_length as usize;
+        let mut piece = Vec::with_capacity(len);
+        unsafe {
+            piece.set_len(len);
+        }
+
+        torrent.iter_files_for_piece(index, &mut piece, |file, seek, buf| {
+            file.seek(SeekFrom::Start(seek))?;
+            file.read_exact(buf)?;
+            Ok(())
+        })?;
+
+        Ok(piece)
+    }
+}
+
+impl DiskManager<File> {
+    pub fn add_torrent(&mut self, metainfo: Arc<Metainfo>) -> io::Result<()> {
+        let mut files = Vec::new();
+        for fileinfo in metainfo.files.iter() {
+            let file = File::create(&fileinfo.path)?;
+            file.set_len(fileinfo.length)?;
+            files.push(file);
+        }
+        self.torrents.insert(metainfo.info_hash, Torrent {
+            files: files,
+            metainfo: metainfo,
+        });
+        Ok(())
+    }
+}
+
+pub struct Torrent<F> where F: Read + Seek + Write {
+    metainfo: Arc<Metainfo>,
+    files: Vec<F>,
+}
+
+impl<F> Torrent<F> where F: Read + Seek + Write {
+    // This function iterates over the files that are spanned by a piece.
+    fn iter_files_for_piece<C>(&mut self, index: u32, piece: &mut [u8], mut callback: C) -> io::Result<()>
+        where C: FnMut(&mut F, u64, &mut [u8]) -> io::Result<()>
+    {
         let buffer_len = piece.len();
         let mut buffer_pos: usize = 0;
         // All the files in a torrent can be seen as a single contiguous file. These
         // values are relative to this contiguous file.
-        let mut abs_pos = index as u64 * torrent.metainfo.piece_length as u64;
+        let mut abs_pos = index as u64 * self.metainfo.piece_length as u64;
         let mut abs_start = 0;
         let mut abs_end = 0;
 
-        for (idx, file) in torrent.metainfo.files.iter().enumerate() {
+        for (idx, file) in self.metainfo.files.iter().enumerate() {
             abs_end += file.length;
 
             if abs_start <= abs_pos && abs_pos < abs_end {
@@ -31,8 +87,7 @@ impl<F> DiskManager<F> where F: Read + Seek + Write {
                 let write_len = cmp::min(remaining, abs_end - abs_pos) as usize;
 
                 // Seek in the file. The position given is relative to the start of the file.
-                torrent.files[idx].seek(SeekFrom::Start(abs_pos - abs_start))?;
-                torrent.files[idx].write_all(&piece[buffer_pos..buffer_pos + write_len])?;
+                callback(&mut self.files[idx], abs_pos - abs_start, &mut piece[buffer_pos..buffer_pos + write_len])?;
 
                 abs_pos += write_len as u64;
                 buffer_pos += write_len;
@@ -51,11 +106,6 @@ impl<F> DiskManager<F> where F: Read + Seek + Write {
     }
 }
 
-pub struct Torrent<F> where F: Read + Seek + Write {
-    metainfo: Arc<Metainfo>,
-    files: Vec<F>,
-}
-
 #[cfg(test)]
 mod tests {
     use super::*;
@@ -69,6 +119,7 @@ mod tests {
 
     struct MockFile {
         seek: SeekFrom,
+        read_len: usize,
         write_len: usize,
     }
 
@@ -76,6 +127,7 @@ mod tests {
         pub fn new() -> Self {
             MockFile {
                 seek: SeekFrom::Start(0),
+                read_len: 0,
                 write_len: 0,
             }
         }
@@ -83,13 +135,13 @@ mod tests {
 
     impl Read for MockFile {
         fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
-            Ok(0)
+            self.read_len = buf.len();
+            Ok(buf.len())
         }
     }
 
     impl Seek for MockFile {
         fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
-            println!("{:?}", pos);
             self.seek = pos;
             Ok(0)
         }
@@ -98,12 +150,7 @@ mod tests {
     impl Write for MockFile {
         fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
             self.write_len = buf.len();
-            Ok(0)
-        }
-
-        fn write_all(&mut self, buf: &[u8]) -> io::Result<()> {
-            self.write_len = buf.len();
-            Ok(())
+            Ok(buf.len())
         }
 
         fn flush(&mut self) -> io::Result<()> {
@@ -111,8 +158,7 @@ mod tests {
         }
     }
 
-    #[test]
-    fn test_write_piece() {
+    fn build_dummy() -> (DiskManager<MockFile>, Hash) {
         let info_hash = Hash::new([1u8; 20]);
 
         let metainfo = Metainfo {
@@ -134,10 +180,46 @@ mod tests {
             files: vec![MockFile::new(), MockFile::new(), MockFile::new()],
         });
 
-        let mut dm = DiskManager {
+        let dm = DiskManager {
             torrents: torrents,
         };
 
+        (dm, info_hash)
+    }
+
+    #[test]
+    fn test_read_piece() {
+        let (mut dm, info_hash) = build_dummy();
+
+        {
+            // Read a piece in a single file that is larger than the piece.
+            dm.read_piece(&info_hash, 0).unwrap();
+            let torrent = &dm.torrents[&info_hash];
+
+            assert_eq!(torrent.files[0].seek, SeekFrom::Start(0));
+            assert_eq!(torrent.files[0].read_len, 10);
+        }
+
+        {
+            // Read a piece over multiple files that are smaller that a piece.
+            dm.read_piece(&info_hash, 1).unwrap();
+            let torrent = &dm.torrents[&info_hash];
+
+            assert_eq!(torrent.files[0].seek, SeekFrom::Start(10));
+            assert_eq!(torrent.files[0].read_len, 2);
+
+            assert_eq!(torrent.files[1].seek, SeekFrom::Start(0));
+            assert_eq!(torrent.files[1].read_len, 6);
+
+            assert_eq!(torrent.files[2].seek, SeekFrom::Start(0));
+            assert_eq!(torrent.files[2].read_len, 2);
+        }
+    }
+
+    #[test]
+    fn test_write_piece() {
+        let (mut dm, info_hash) = build_dummy();
+
         {
             // Write a piece in a single file that is larger than the piece.
             dm.write_piece(&info_hash, 0, vec![2u8; 10]).unwrap();