session: add buffer utilities

src/net/buffers.rs

@@ -0,0 +1,68 @@
+use std::io::{self, Read};
+use std::ptr;
+
+use sha1::Sha1;
+
+use metainfo::Hash;
+
+/// Allocate a new buffer of the given size and read_exact into it.
+pub fn read_exact<R>(mut reader: R, size: usize) -> io::Result<Vec<u8>> where R: Read {
+    let mut buf = Vec::with_capacity(size);
+    unsafe {
+        buf.set_len(size);
+    }
+    reader.read_exact(&mut buf)?;
+    Ok(buf)
+}
+
+pub struct PieceBuffer {
+    fragment_count: u32,
+    num_fragments: u32,
+    buffer: Vec<u8>,
+}
+
+impl PieceBuffer {
+    pub fn new(num_fragments: u32, piece_length: u32) -> Self {
+        let mut buffer = Vec::with_capacity(piece_length as usize);
+        unsafe {
+            buffer.set_len(piece_length as usize);
+        }
+        PieceBuffer {
+            fragment_count: 0,
+            num_fragments: num_fragments,
+            buffer: buffer,
+        }
+    }
+
+    pub fn add_fragment(&mut self, begin: u32, fragment: &[u8]) {
+        let begin = begin as usize;
+        assert!(begin + fragment.len() <= self.buffer.len());
+        unsafe {
+            ptr::copy_nonoverlapping(fragment.as_ptr(), self.buffer[begin..].as_mut_ptr(), fragment.len());
+        }
+        self.fragment_count += 1;
+    }
+
+    #[inline]
+    pub fn complete(&self) -> bool {
+        self.fragment_count == self.num_fragments
+    }
+
+    pub fn matches_hash(&self, info_hash: &Hash) -> bool {
+        if !self.complete() {
+            panic!("trying to hash piece buffer before it's complete");
+        } else {
+            let mut m = Sha1::new();
+            m.update(&self.buffer);
+            m.digest().bytes() == &info_hash[..]
+        }
+    }
+
+    pub fn get(self) -> Vec<u8> {
+        if !self.complete() {
+            panic!("trying to get piece buffer before it's complete");
+        } else {
+            self.buffer
+        }
+    }
+}

src/net/mod.rs

@@ -1,3 +1,4 @@
 pub mod bitfield;
+mod buffers;
 pub mod peer;
 pub mod session;

src/net/peer.rs

@@ -4,6 +4,7 @@ use std::net::{Shutdown, TcpStream};
 use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
 
 use metainfo::Hash;
+use net::buffers;
 use tracker::Peer;
 
 #[derive(Debug)]
@@ -68,11 +69,8 @@ impl PeerConnection {
             },
             5 => {
                 let size = len as usize - 1;
-                let mut bitfield = Vec::with_capacity(size);
+                let bitfield = buffers::read_exact(&mut self.sock, size)?;
-                unsafe {
+
-                    bitfield.set_len(size);
-                }
-                self.sock.read_exact(&mut bitfield)?;
                 Packet::Bitfield {
                     bitfield: bitfield,
                 }
@@ -86,11 +84,7 @@ impl PeerConnection {
                 let size = len as usize - 9;
                 let index = self.sock.read_u32::<BigEndian>()?;
                 let begin = self.sock.read_u32::<BigEndian>()?;
-                let mut block = Vec::with_capacity(size);
+                let block = buffers::read_exact(&mut self.sock, size)?;
-                unsafe {
-                    block.set_len(size);
-                }
-                self.sock.read_exact(&mut block)?;
 
                 Packet::Piece {
                     index: index,

src/net/session.rs

@@ -9,10 +9,10 @@ use std::thread;
 use std::time::{Duration, Instant};
 
 use libc;
-use sha1::Sha1;
 
 use metainfo::{Hash, Metainfo};
 use net::bitfield::BitField;
+use net::buffers::PieceBuffer;
 use net::peer::{Packet, PeerConnection};
 use tracker::http;
 
@@ -132,9 +132,7 @@ pub struct SessionFragment {
 
 struct SessionPiece {
     fragments: BTreeMap<u32, SessionFragment>,
-    buffer: Vec<u8>,
+    buffer: PieceBuffer,
-    num_fragments: u32,
-    total_fragments: u32,
 }
 
 struct SessionTorrent {
@@ -197,9 +195,9 @@ impl SessionTorrent {
         for index in (0..self.own_bitfield.len()).filter(|index| !keys.contains(index)) {
             if !self.own_bitfield.is_set(index) && peer.bitfield.is_set(index) {
 
-                let total_fragments = f64::ceil(self.metainfo.piece_length as f64 / FRAGMENT_SIZE as f64) as u32;
+                let num_fragments = f64::ceil(self.metainfo.piece_length as f64 / FRAGMENT_SIZE as f64) as u32;
                 let mut fragments = BTreeMap::new();
-                for idx in 0..total_fragments {
+                for idx in 0..num_fragments {
                     let begin = idx * FRAGMENT_SIZE;
 
                     fragments.insert(begin ,SessionFragment {
@@ -218,15 +216,9 @@ impl SessionTorrent {
                     length = taken.length;
                 }
 
-                let mut buffer = Vec::with_capacity(self.metainfo.piece_length as usize);
-                unsafe {
-                    buffer.set_len(self.metainfo.piece_length as usize);
-                }
                 self.pieces.insert(index, SessionPiece {
                     fragments: fragments,
-                    buffer: buffer,
+                    buffer: PieceBuffer::new(num_fragments, self.metainfo.piece_length),
-                    num_fragments: 0,
-                    total_fragments: total_fragments,
                 });
 
                 return Some((index, begin, length));
@@ -246,7 +238,7 @@ impl SessionTorrent {
         }
     }
 
-    fn write_piece(&mut self, index: u32) {
+    fn write_piece(&mut self, index: u32, buffer: Vec<u8>) {
         let mut start = 0;
         let mut end = 0;
         let mut pos: u64 = 0;
@@ -259,7 +251,7 @@ impl SessionTorrent {
             if write < end {
                 let len = cmp::min(remaining, fileinfo.length as u64);
                 self.files[id].seek(SeekFrom::Start(write - start));
-                self.files[id].write_all(&self.pieces[&index].buffer[pos as usize..(pos+len) as usize]);
+                self.files[id].write_all(&buffer[pos as usize..(pos+len) as usize]);
                 write += len;
                 pos += len;
                 remaining -= len;
@@ -280,36 +272,39 @@ impl SessionTorrent {
 
     fn piece_reply(&mut self, peer_id: &Hash, index: u32, begin: u32, block: Vec<u8>) {
         let mut remove = false;
+        let mut reset = false;
         {
             if let Some(piece) = self.pieces.get_mut(&index) {
                 if let Some(fragment) = piece.fragments.get_mut(&begin) {
                     fragment.status = FragmentStatus::Complete;
 
-                    piece.buffer[begin as usize..(begin as usize)+block.len()].copy_from_slice(&block);
+                    piece.buffer.add_fragment(begin, &block);
-                    piece.num_fragments += 1;
+                    if piece.buffer.complete() {
-                    if piece.num_fragments == piece.total_fragments {
-                        // TODO check hash
                         println!("piece is done {}", index);
-                        let mut m = Sha1::new();
+                        if piece.buffer.matches_hash(&self.metainfo.pieces[index as usize]) {
-                        m.update(&piece.buffer);
-                        if m.digest().bytes() == &self.metainfo.pieces[index as usize][..] {
                             self.own_bitfield.set(index);
                             println!("it's a match!");
                             remove = true;
                         } else {
+                            reset = true;
                             println!("no match");
                         }
                     }
                 } else {
                     println!("could not find fragment {}", begin);
                 }
+                if reset {
+                    for fragment in piece.fragments.values_mut() {
+                        fragment.status = FragmentStatus::Available;
+                    }
+                }
             } else {
                 println!("could not find piece {}", index);
             }
         }
         if remove {
-            self.write_piece(index);
+            let piece = self.pieces.remove(&index).expect("told to remove piece that doesn't exist");
-            self.pieces.remove(&index);
+            self.write_piece(index, piece.buffer.get());
         }
         self.requeue(peer_id);
     }