use std::io::{self, Read, Write}; use std::net::TcpStream; use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt}; use metainfo::Hash; use net::bitfield::BitField; use tracker::Peer; #[derive(Debug)] pub enum Packet { Choke, Unchoke, Interested, NotInterested, Have { index: u32, }, Bitfield { bitfield: Vec, }, Request { index: u32, begin: u32, length: u32, }, Piece { index: u32, begin: u32, block: Vec, }, Cancel { index: u32, begin: u32, length: u32, } } pub fn open_connection(peer: Peer, own_info_hash: Hash, own_peer_id: Hash) -> io::Result<(TcpStream, TcpStream, Hash)> { let mut sock = TcpStream::connect((peer.addr, peer.port))?; // send handshake sock.write_u8(19)?; sock.write(b"BitTorrent protocol")?; sock.write_u64::(0)?; sock.write(&own_info_hash)?; sock.write(&own_peer_id)?; sock.flush()?; // receive handshake let mut buf = [0u8; 68]; sock.read_exact(&mut buf)?; let peer_info_hash = Hash::from_slice(&buf[28..48]); let peer_id = Hash::from_slice(&buf[48..68]); if buf[0] != 19 || &buf[1..20] != b"BitTorrent protocol" || own_info_hash != peer_info_hash { return Err(io::Error::new(io::ErrorKind::Other, "invalid protocol")) } Ok((sock.try_clone()?, sock, peer_id)) } pub fn read_packet(sock: &mut TcpStream) -> io::Result { let len = sock.read_u32::()?; let id = sock.read_u8()?; Ok(match id { 0 => Packet::Choke, 1 => Packet::Unchoke, 2 => Packet::Interested, 3 => Packet::NotInterested, 4 => Packet::Have { index: sock.read_u32::()?, }, 5 => { let size = len as usize - 1; let mut bitfield = Vec::with_capacity(size); unsafe { bitfield.set_len(size); } sock.read_exact(&mut bitfield)?; Packet::Bitfield { bitfield: bitfield, } } 6 => Packet::Request { index: sock.read_u32::()?, begin: sock.read_u32::()?, length: sock.read_u32::()?, }, 7 => { let size = len as usize - 9; let index = sock.read_u32::()?; let begin = sock.read_u32::()?; let mut block = Vec::with_capacity(size); unsafe { block.set_len(size); } sock.read_exact(&mut block)?; Packet::Piece { index: index, begin: begin, block: block, } } 8 => Packet::Cancel { index: sock.read_u32::()?, begin: sock.read_u32::()?, length: sock.read_u32::()?, }, _ => return Err(io::Error::new(io::ErrorKind::Other, "invalid packet")), }) } pub fn send_keepalive(sock: &mut TcpStream) -> io::Result<()> { sock.write(b"")?; Ok(()) } pub fn send_choke(sock: &mut TcpStream) -> io::Result<()> { sock.write_u32::(1)?; sock.write_u8(0)?; Ok(()) } pub fn send_unchoke(sock: &mut TcpStream) -> io::Result<()> { sock.write_u32::(1)?; sock.write_u8(1)?; Ok(()) } pub fn send_interested(sock: &mut TcpStream) -> io::Result<()> { sock.write_u32::(1)?; sock.write_u8(2)?; Ok(()) } pub fn send_not_interested(sock: &mut TcpStream) -> io::Result<()> { sock.write_u32::(1)?; sock.write_u8(3)?; Ok(()) } pub fn send_have(sock: &mut TcpStream, piece: u32) -> io::Result<()> { sock.write_u32::(5)?; sock.write_u8(4)?; sock.write_u32::(piece)?; Ok(()) } pub fn send_request(sock: &mut TcpStream, piece: u32, begin: u32, length: u32) -> io::Result<()> { sock.write_u32::(13)?; sock.write_u8(6)?; sock.write_u32::(piece)?; sock.write_u32::(begin)?; sock.write_u32::(length)?; Ok(()) } pub fn send_piece(sock: &mut TcpStream, piece: u32, begin: u32, data: &[u8]) -> io::Result<()> { sock.write_u32::(9 + data.len() as u32)?; sock.write_u8(7)?; sock.write_u32::(piece)?; sock.write_u32::(begin)?; sock.write(data)?; Ok(()) } pub fn send_cancel(sock: &mut TcpStream, piece: u32, begin: u32, length: u32) -> io::Result<()> { sock.write_u32::(13)?; sock.write_u8(8)?; sock.write_u32::(piece)?; sock.write_u32::(begin)?; sock.write_u32::(length)?; Ok(()) }