diff --git a/src/metainfo.rs b/src/metainfo.rs index f48f9e9..ceb7b8d 100644 --- a/src/metainfo.rs +++ b/src/metainfo.rs @@ -107,6 +107,15 @@ pub struct MetainfoFile { pub path: PathBuf, } +impl MetainfoFile { + pub fn new(length: u64, path: PathBuf) -> Self { + MetainfoFile { + length: length, + path: path, + } + } +} + fn sha1(bytes: &[u8]) -> Hash { let mut hasher = Sha1::new(); hasher.update(bytes); diff --git a/src/net/_session/disk.rs b/src/net/_session/disk.rs new file mode 100644 index 0000000..f1af97b --- /dev/null +++ b/src/net/_session/disk.rs @@ -0,0 +1,166 @@ +use std::cmp; +use std::io::{self, Read, Seek, SeekFrom, Write}; +use std::sync::Arc; + +use net::_session::TorrentMap; +use metainfo::{Hash, Metainfo}; + +pub struct DiskManager where F: Read + Seek + Write { + torrents: TorrentMap>, +} + +impl DiskManager where F: Read + Seek + Write { + pub fn write_piece(&mut self, info_hash: &Hash, index: u32, piece: Vec) -> io::Result<()> { + let torrent = self.torrents.get_mut(&info_hash).unwrap(); + + 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_start = 0; + let mut abs_end = 0; + + for (idx, file) in torrent.metainfo.files.iter().enumerate() { + abs_end += file.length; + + if abs_start <= abs_pos && abs_pos < abs_end { + // The amount of bytes to write is the remaining number of bytes in the buffer, + // or the size of the file, whichever is smaller. + let remaining = (buffer_len - buffer_pos) as u64; + 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])?; + + abs_pos += write_len as u64; + buffer_pos += write_len; + + // If the buffer position if equal to the length, we're done writing. + if buffer_pos == buffer_len { + break + } + } + + // The start of the next file is the end of this one. + abs_start = abs_end; + } + + Ok(()) + } +} + +pub struct Torrent where F: Read + Seek + Write { + metainfo: Arc, + files: Vec, +} + +#[cfg(test)] +mod tests { + use super::*; + + use std::io::{self, Read, Seek, SeekFrom, Write}; + use std::path::PathBuf; + use std::sync::Arc; + + use net::_session::TorrentMap; + use metainfo::{Hash, Metainfo, MetainfoFile}; + + struct MockFile { + seek: SeekFrom, + write_len: usize, + } + + impl MockFile { + pub fn new() -> Self { + MockFile { + seek: SeekFrom::Start(0), + write_len: 0, + } + } + } + + impl Read for MockFile { + fn read(&mut self, buf: &mut [u8]) -> io::Result { + Ok(0) + } + } + + impl Seek for MockFile { + fn seek(&mut self, pos: SeekFrom) -> io::Result { + println!("{:?}", pos); + self.seek = pos; + Ok(0) + } + } + + 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(()) + } + + fn flush(&mut self) -> io::Result<()> { + Ok(()) + } + } + + #[test] + fn test_write_piece() { + let info_hash = Hash::new([1u8; 20]); + + let metainfo = Metainfo { + announce: "".into(), + announce_list: vec![], + files: vec![MetainfoFile::new(12, PathBuf::new()), + MetainfoFile::new(6, PathBuf::new()), + MetainfoFile::new(2, PathBuf::new())], + info_hash: info_hash, + piece_length: 10, + num_pieces: 2, + pieces: vec![], + length: 20, + }; + + let mut torrents = TorrentMap::new(); + torrents.insert(info_hash, Torrent { + metainfo: Arc::new(metainfo), + files: vec![MockFile::new(), MockFile::new(), MockFile::new()], + }); + + let mut dm = DiskManager { + torrents: torrents, + }; + + { + // Write a piece in a single file that is larger than the piece. + dm.write_piece(&info_hash, 0, vec![2u8; 10]).unwrap(); + let torrent = &dm.torrents[&info_hash]; + + assert_eq!(torrent.files[0].seek, SeekFrom::Start(0)); + assert_eq!(torrent.files[0].write_len, 10); + } + + + { + // Write a piece over multiple files that are smaller that a piece. + dm.write_piece(&info_hash, 1, vec![2u8; 10]).unwrap(); + let torrent = &dm.torrents[&info_hash]; + + assert_eq!(torrent.files[0].seek, SeekFrom::Start(10)); + assert_eq!(torrent.files[0].write_len, 2); + + assert_eq!(torrent.files[1].seek, SeekFrom::Start(0)); + assert_eq!(torrent.files[1].write_len, 6); + + assert_eq!(torrent.files[2].seek, SeekFrom::Start(0)); + assert_eq!(torrent.files[2].write_len, 2); + } + } +} diff --git a/src/net/_session/mod.rs b/src/net/_session/mod.rs new file mode 100644 index 0000000..14c7334 --- /dev/null +++ b/src/net/_session/mod.rs @@ -0,0 +1,7 @@ +pub mod disk; + +use std::collections::HashMap; + +use metainfo::Hash; + +type TorrentMap = HashMap; diff --git a/src/net/mod.rs b/src/net/mod.rs index 00df1d7..e5067f4 100644 --- a/src/net/mod.rs +++ b/src/net/mod.rs @@ -2,3 +2,4 @@ pub mod bitfield; mod buffers; pub mod peer; pub mod session; +pub mod _session;