rework peer connection

src/net/session.rs

@@ -14,7 +14,7 @@ use sha1::Sha1;
 
 use metainfo::{Hash, Metainfo};
 use net::bitfield::BitField;
-use net::peer::{self, Packet};
+use net::peer::{Packet, PeerConnection};
 use tracker::http;
 
 const FRAGMENT_SIZE: u32 = 16 * 1024;
@@ -30,7 +30,7 @@ enum Signal {
     ConnectionOpened {
         info_hash: Hash,
         peer_id: Hash,
-        sock: TcpStream,
+        sock: PeerConnection,
     },
     ConnectionClosed {
         info_hash: Hash,
@@ -96,8 +96,8 @@ impl Session {
     }
 }
 
-struct PeerConnection {
-    sock: TcpStream,
+struct SessionPeer {
+    sock: PeerConnection,
     bitfield: BitField,
     am_choking: bool,
     am_interested: bool,
@@ -105,9 +105,9 @@ struct PeerConnection {
     peer_interested: bool,
 }
 
-impl PeerConnection {
-    pub fn new(sock: TcpStream, bitfield_len: u32) -> Self {
-        PeerConnection {
+impl SessionPeer {
+    pub fn new(sock: PeerConnection, bitfield_len: u32) -> Self {
+        SessionPeer {
             sock: sock,
             bitfield: BitField::with_capacity(bitfield_len),
             am_choking: true,
@@ -141,7 +141,7 @@ struct SessionPiece {
 struct SessionTorrent {
     metainfo: Metainfo,
     own_bitfield: BitField,
-    peers: HashMap<Hash, PeerConnection>,
+    peers: HashMap<Hash, SessionPeer>,
     pieces: BTreeMap<u32, SessionPiece>,
     files: Vec<File>,
 }
@@ -157,7 +157,7 @@ impl SessionTorrent {
         }
     }
 
-    pub fn get_peer(&mut self, peer_id: &Hash) -> &mut PeerConnection {
+    pub fn get_peer(&mut self, peer_id: &Hash) -> &mut SessionPeer {
         self.peers.get_mut(peer_id).unwrap()
     }
 
@@ -240,9 +240,9 @@ impl SessionTorrent {
     fn requeue(&mut self, peer_id: &Hash) {
         if let Some((index, begin, length)) = self.get_fragment(peer_id) {
             // println!("onto piece {} {}..{}", index, begin, begin + length);
-            peer::send_request(&mut self.get_peer(&peer_id).sock, index, begin,length);
+            self.get_peer(&peer_id).sock.send_request(index, begin,length);
         } else {
-            peer::send_not_interested(&mut self.get_peer(&peer_id).sock);
+            self.get_peer(&peer_id).sock.send_not_interested();
             println!("no fragment");
         }
     }
@@ -332,15 +332,15 @@ impl SessionNetworkThread {
                 Signal::AddTorrent(torrent) => self.signal_add_torrent(torrent),
                 Signal::RemoveTorrent(id) => {
                     if let Some(mut sess_torrent) = self.torrents.remove(&id) {
-                        for conn in sess_torrent.peers.values_mut() {
-                            let _ = conn.sock.shutdown(Shutdown::Both);
+                        for peer in sess_torrent.peers.values_mut() {
+                            peer.sock.shutdown();
                         }
                     }
                 }
                 Signal::ConnectionOpened { info_hash, peer_id, sock } => {
                     if let Some(sess_torrent) = self.torrents.get_mut(&info_hash) {
-                        let mut peer = PeerConnection::new(sock, sess_torrent.metainfo.pieces.len() as u32);
-                        peer::send_interested(&mut peer.sock);
+                        let mut peer = SessionPeer::new(sock, sess_torrent.metainfo.pieces.len() as u32);
+                        peer.sock.send_interested();
                         peer.am_interested = true;
                         sess_torrent.peers.insert(peer_id, peer);
                     }
@@ -376,20 +376,18 @@ impl SessionNetworkThread {
         if let Ok(resp) = http::get_peers(own_peer_id, self.session.port, &metainfo, 0, 0, 0) {
             self.torrents.insert(info_hash, SessionTorrent::new(metainfo, files));
 
-            println!("trying with {} peers", resp.peers.len());
-
             for &peer in resp.peers.iter() {
                 let sender = self.sender.clone();
 
                 thread::spawn(move || {
-                    if let Ok((mut sock, sock_clone, peer_id)) = peer::open_connection(peer, info_hash, own_peer_id) {
+                    if let Ok((mut sock, sock_clone, peer_id)) = PeerConnection::open_connection(peer, info_hash, own_peer_id) {
                         let _ = sender.send(Signal::ConnectionOpened {
                             info_hash: info_hash,
                             peer_id: peer_id,
                             sock: sock_clone,
                         });
 
-                        while let Ok(packet) = peer::read_packet(&mut sock) {
+                        while let Ok(packet) = sock.read_packet() {
                             let _ = sender.send(Signal::Packet {
                                 info_hash: info_hash,
                                 peer_id: peer_id,