diff --git a/src/net/session.rs b/src/net/session.rs index d3630e1..9f1160e 100644 --- a/src/net/session.rs +++ b/src/net/session.rs @@ -49,6 +49,8 @@ struct InnerSession { port: u16, } +// == Session == + pub struct Session { // inner: Arc, sender: Sender, @@ -95,6 +97,8 @@ impl Session { } } +// == SessionPeer == + struct SessionPeer { sock: PeerConnection, bitfield: BitField, @@ -117,6 +121,8 @@ impl SessionPeer { } } +// == FragmentStatus == + #[derive(Copy, Clone, Debug)] enum FragmentStatus { Available, @@ -124,17 +130,23 @@ enum FragmentStatus { Taken(Instant), } +// == SessionFragment == + pub struct SessionFragment { begin: u32, length: u32, status: FragmentStatus, } +// == SessionPiece == + struct SessionPiece { fragments: BTreeMap, buffer: PieceBuffer, } +// == SessionTorrent == + struct SessionTorrent { metainfo: Metainfo, own_bitfield: BitField, @@ -310,6 +322,8 @@ impl SessionTorrent { } } +// == SessionNetworkThread == + struct SessionNetworkThread { session: Arc, sender: Sender, @@ -320,34 +334,13 @@ struct SessionNetworkThread { } impl SessionNetworkThread { - pub fn run(mut self, input: Receiver) { - for signal in input.iter() { + pub fn run(mut self, signals: Receiver) { + for signal in signals.iter() { match signal { Signal::AddTorrent(torrent) => self.signal_add_torrent(torrent), - Signal::RemoveTorrent(id) => { - if let Some(mut sess_torrent) = self.torrents.remove(&id) { - 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 = 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); - } - } - Signal::ConnectionClosed { info_hash, peer_id } => { - if let Some(sess_torrent) = self.torrents.get_mut(&info_hash) { - if let Some(mut peer) = sess_torrent.peers.remove(&peer_id) { - // if !peer.peer_choking && self.seeds > 1 { - // self.seeds -= 1; - // } - } - } - } + Signal::RemoveTorrent(id) => self.signal_remove_torrent(&id), + Signal::ConnectionOpened { info_hash, peer_id, sock } => self.signal_connection_opened(&info_hash, &peer_id, sock), + Signal::ConnectionClosed { info_hash, peer_id } => self.signal_connection_closed(&info_hash, &peer_id), Signal::Packet { info_hash, peer_id, packet } => self.signal_packet(info_hash, peer_id, packet), Signal::Stop => break, } @@ -399,6 +392,33 @@ impl SessionNetworkThread { } } + fn signal_remove_torrent(&mut self, id: &Hash) { + if let Some(mut sess_torrent) = self.torrents.remove(id) { + for peer in sess_torrent.peers.values_mut() { + peer.sock.shutdown(); + } + } + } + + fn signal_connection_opened(&mut self, info_hash: &Hash, peer_id: &Hash, conn: PeerConnection) { + if let Some(sess_torrent) = self.torrents.get_mut(info_hash) { + let mut peer = SessionPeer::new(conn, sess_torrent.metainfo.pieces.len() as u32); + peer.sock.send_interested(); + peer.am_interested = true; + sess_torrent.peers.insert(*peer_id, peer); + } + } + + fn signal_connection_closed(&mut self, info_hash: &Hash, peer_id: &Hash) { + if let Some(sess_torrent) = self.torrents.get_mut(&info_hash) { + if let Some(mut peer) = sess_torrent.peers.remove(&peer_id) { + // if !peer.peer_choking && self.seeds > 1 { + // self.seeds -= 1; + // } + } + } + } + fn signal_packet(&mut self, info_hash: Hash, peer_id: Hash, packet: Packet) { if let Some(torrent) = self.torrents.get_mut(&info_hash) { match packet {