diff --git a/src/net/_session/disk.rs b/src/net/_session/disk.rs index f1af97b..609170c 100644 --- a/src/net/_session/disk.rs +++ b/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 where F: Read + Seek + Write { } impl DiskManager where F: Read + Seek + Write { - pub fn write_piece(&mut self, info_hash: &Hash, index: u32, piece: Vec) -> io::Result<()> { + + pub fn write_piece(&mut self, info_hash: &Hash, index: u32, mut piece: Vec) -> 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> { + 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 { + pub fn add_torrent(&mut self, metainfo: Arc) -> 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 where F: Read + Seek + Write { + metainfo: Arc, + files: Vec, +} + +impl Torrent where F: Read + Seek + Write { + // This function iterates over the files that are spanned by a piece. + fn iter_files_for_piece(&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 DiskManager 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 DiskManager where F: Read + Seek + Write { } } -pub struct Torrent where F: Read + Seek + Write { - metainfo: Arc, - files: Vec, -} - #[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 { - Ok(0) + self.read_len = buf.len(); + Ok(buf.len()) } } impl Seek for MockFile { fn seek(&mut self, pos: SeekFrom) -> io::Result { - 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 { 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, 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();