rework peer connection
This commit is contained in:
262
src/net/peer.rs
262
src/net/peer.rs
@@ -1,10 +1,9 @@
|
||||
use std::io::{self, Read, Write};
|
||||
use std::net::TcpStream;
|
||||
use std::net::{Shutdown, TcpStream};
|
||||
|
||||
use byteorder::{BigEndian, ReadBytesExt, WriteBytesExt};
|
||||
|
||||
use metainfo::Hash;
|
||||
use net::bitfield::BitField;
|
||||
use tracker::Peer;
|
||||
|
||||
#[derive(Debug)]
|
||||
@@ -22,143 +21,162 @@ pub enum Packet {
|
||||
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))?;
|
||||
#[derive(Debug)]
|
||||
pub struct PeerConnection {
|
||||
sock: TcpStream,
|
||||
}
|
||||
|
||||
// send handshake
|
||||
sock.write_u8(19)?;
|
||||
sock.write(b"BitTorrent protocol")?;
|
||||
sock.write_u64::<BigEndian>(0)?;
|
||||
sock.write(&own_info_hash)?;
|
||||
sock.write(&own_peer_id)?;
|
||||
sock.flush()?;
|
||||
impl PeerConnection {
|
||||
pub fn open_connection(peer: Peer, own_info_hash: Hash, own_peer_id: Hash) -> io::Result<(PeerConnection, PeerConnection, Hash)> {
|
||||
let mut sock = TcpStream::connect((peer.addr, peer.port))?;
|
||||
|
||||
// receive handshake
|
||||
let mut buf = [0u8; 68];
|
||||
sock.read_exact(&mut buf)?;
|
||||
// send handshake
|
||||
sock.write_u8(19)?;
|
||||
sock.write(b"BitTorrent protocol")?;
|
||||
sock.write_u64::<BigEndian>(0)?;
|
||||
sock.write(&own_info_hash)?;
|
||||
sock.write(&own_peer_id)?;
|
||||
sock.flush()?;
|
||||
|
||||
let peer_info_hash = Hash::from_slice(&buf[28..48]);
|
||||
let peer_id = Hash::from_slice(&buf[48..68]);
|
||||
// receive handshake
|
||||
let mut buf = [0u8; 68];
|
||||
sock.read_exact(&mut buf)?;
|
||||
|
||||
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"))
|
||||
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"))
|
||||
}
|
||||
|
||||
let mut conn = PeerConnection { sock: sock };
|
||||
|
||||
Ok((conn.try_clone()?, conn, peer_id))
|
||||
}
|
||||
|
||||
Ok((sock.try_clone()?, sock, peer_id))
|
||||
}
|
||||
pub fn read_packet(&mut self) -> io::Result<Packet> {
|
||||
let len = self.sock.read_u32::<BigEndian>()?;
|
||||
let id = self.sock.read_u8()?;
|
||||
|
||||
pub fn read_packet(sock: &mut TcpStream) -> io::Result<Packet> {
|
||||
let len = sock.read_u32::<BigEndian>()?;
|
||||
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::<BigEndian>()?,
|
||||
},
|
||||
5 => {
|
||||
let size = len as usize - 1;
|
||||
let mut bitfield = Vec::with_capacity(size);
|
||||
unsafe {
|
||||
bitfield.set_len(size);
|
||||
Ok(match id {
|
||||
0 => Packet::Choke,
|
||||
1 => Packet::Unchoke,
|
||||
2 => Packet::Interested,
|
||||
3 => Packet::NotInterested,
|
||||
4 => Packet::Have {
|
||||
index: self.sock.read_u32::<BigEndian>()?,
|
||||
},
|
||||
5 => {
|
||||
let size = len as usize - 1;
|
||||
let mut bitfield = Vec::with_capacity(size);
|
||||
unsafe {
|
||||
bitfield.set_len(size);
|
||||
}
|
||||
self.sock.read_exact(&mut bitfield)?;
|
||||
Packet::Bitfield {
|
||||
bitfield: bitfield,
|
||||
}
|
||||
}
|
||||
sock.read_exact(&mut bitfield)?;
|
||||
Packet::Bitfield {
|
||||
bitfield: bitfield,
|
||||
6 => Packet::Request {
|
||||
index: self.sock.read_u32::<BigEndian>()?,
|
||||
begin: self.sock.read_u32::<BigEndian>()?,
|
||||
length: self.sock.read_u32::<BigEndian>()?,
|
||||
},
|
||||
7 => {
|
||||
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);
|
||||
unsafe {
|
||||
block.set_len(size);
|
||||
}
|
||||
self.sock.read_exact(&mut block)?;
|
||||
|
||||
Packet::Piece {
|
||||
index: index,
|
||||
begin: begin,
|
||||
block: block,
|
||||
}
|
||||
}
|
||||
}
|
||||
6 => Packet::Request {
|
||||
index: sock.read_u32::<BigEndian>()?,
|
||||
begin: sock.read_u32::<BigEndian>()?,
|
||||
length: sock.read_u32::<BigEndian>()?,
|
||||
},
|
||||
7 => {
|
||||
let size = len as usize - 9;
|
||||
let index = sock.read_u32::<BigEndian>()?;
|
||||
let begin = sock.read_u32::<BigEndian>()?;
|
||||
let mut block = Vec::with_capacity(size);
|
||||
unsafe {
|
||||
block.set_len(size);
|
||||
}
|
||||
sock.read_exact(&mut block)?;
|
||||
8 => Packet::Cancel {
|
||||
index: self.sock.read_u32::<BigEndian>()?,
|
||||
begin: self.sock.read_u32::<BigEndian>()?,
|
||||
length: self.sock.read_u32::<BigEndian>()?,
|
||||
},
|
||||
_ => return Err(io::Error::new(io::ErrorKind::Other, "invalid packet")),
|
||||
})
|
||||
}
|
||||
|
||||
Packet::Piece {
|
||||
index: index,
|
||||
begin: begin,
|
||||
block: block,
|
||||
}
|
||||
}
|
||||
8 => Packet::Cancel {
|
||||
index: sock.read_u32::<BigEndian>()?,
|
||||
begin: sock.read_u32::<BigEndian>()?,
|
||||
length: sock.read_u32::<BigEndian>()?,
|
||||
},
|
||||
_ => return Err(io::Error::new(io::ErrorKind::Other, "invalid packet")),
|
||||
})
|
||||
}
|
||||
pub fn send_keepalive(&mut self) -> io::Result<()> {
|
||||
self.sock.write(b"")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn send_keepalive(sock: &mut TcpStream) -> io::Result<()> {
|
||||
sock.write(b"")?;
|
||||
Ok(())
|
||||
}
|
||||
pub fn send_choke(&mut self) -> io::Result<()> {
|
||||
self.sock.write_u32::<BigEndian>(1)?;
|
||||
self.sock.write_u8(0)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn send_choke(sock: &mut TcpStream) -> io::Result<()> {
|
||||
sock.write_u32::<BigEndian>(1)?;
|
||||
sock.write_u8(0)?;
|
||||
Ok(())
|
||||
}
|
||||
pub fn send_unchoke(&mut self) -> io::Result<()> {
|
||||
self.sock.write_u32::<BigEndian>(1)?;
|
||||
self.sock.write_u8(1)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn send_unchoke(sock: &mut TcpStream) -> io::Result<()> {
|
||||
sock.write_u32::<BigEndian>(1)?;
|
||||
sock.write_u8(1)?;
|
||||
Ok(())
|
||||
}
|
||||
pub fn send_interested(&mut self) -> io::Result<()> {
|
||||
self.sock.write_u32::<BigEndian>(1)?;
|
||||
self.sock.write_u8(2)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn send_interested(sock: &mut TcpStream) -> io::Result<()> {
|
||||
sock.write_u32::<BigEndian>(1)?;
|
||||
sock.write_u8(2)?;
|
||||
Ok(())
|
||||
}
|
||||
pub fn send_not_interested(&mut self) -> io::Result<()> {
|
||||
self.sock.write_u32::<BigEndian>(1)?;
|
||||
self.sock.write_u8(3)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn send_not_interested(sock: &mut TcpStream) -> io::Result<()> {
|
||||
sock.write_u32::<BigEndian>(1)?;
|
||||
sock.write_u8(3)?;
|
||||
Ok(())
|
||||
}
|
||||
pub fn send_have(&mut self, piece: u32) -> io::Result<()> {
|
||||
self.sock.write_u32::<BigEndian>(5)?;
|
||||
self.sock.write_u8(4)?;
|
||||
self.sock.write_u32::<BigEndian>(piece)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn send_have(sock: &mut TcpStream, piece: u32) -> io::Result<()> {
|
||||
sock.write_u32::<BigEndian>(5)?;
|
||||
sock.write_u8(4)?;
|
||||
sock.write_u32::<BigEndian>(piece)?;
|
||||
Ok(())
|
||||
}
|
||||
pub fn send_request(&mut self, piece: u32, begin: u32, length: u32) -> io::Result<()> {
|
||||
self.sock.write_u32::<BigEndian>(13)?;
|
||||
self.sock.write_u8(6)?;
|
||||
self.sock.write_u32::<BigEndian>(piece)?;
|
||||
self.sock.write_u32::<BigEndian>(begin)?;
|
||||
self.sock.write_u32::<BigEndian>(length)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn send_request(sock: &mut TcpStream, piece: u32, begin: u32, length: u32) -> io::Result<()> {
|
||||
sock.write_u32::<BigEndian>(13)?;
|
||||
sock.write_u8(6)?;
|
||||
sock.write_u32::<BigEndian>(piece)?;
|
||||
sock.write_u32::<BigEndian>(begin)?;
|
||||
sock.write_u32::<BigEndian>(length)?;
|
||||
Ok(())
|
||||
}
|
||||
pub fn send_piece(&mut self, piece: u32, begin: u32, data: &[u8]) -> io::Result<()> {
|
||||
self.sock.write_u32::<BigEndian>(9 + data.len() as u32)?;
|
||||
self.sock.write_u8(7)?;
|
||||
self.sock.write_u32::<BigEndian>(piece)?;
|
||||
self.sock.write_u32::<BigEndian>(begin)?;
|
||||
self.sock.write(data)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn send_piece(sock: &mut TcpStream, piece: u32, begin: u32, data: &[u8]) -> io::Result<()> {
|
||||
sock.write_u32::<BigEndian>(9 + data.len() as u32)?;
|
||||
sock.write_u8(7)?;
|
||||
sock.write_u32::<BigEndian>(piece)?;
|
||||
sock.write_u32::<BigEndian>(begin)?;
|
||||
sock.write(data)?;
|
||||
Ok(())
|
||||
}
|
||||
pub fn send_cancel(&mut self, piece: u32, begin: u32, length: u32) -> io::Result<()> {
|
||||
self.sock.write_u32::<BigEndian>(13)?;
|
||||
self.sock.write_u8(8)?;
|
||||
self.sock.write_u32::<BigEndian>(piece)?;
|
||||
self.sock.write_u32::<BigEndian>(begin)?;
|
||||
self.sock.write_u32::<BigEndian>(length)?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn send_cancel(sock: &mut TcpStream, piece: u32, begin: u32, length: u32) -> io::Result<()> {
|
||||
sock.write_u32::<BigEndian>(13)?;
|
||||
sock.write_u8(8)?;
|
||||
sock.write_u32::<BigEndian>(piece)?;
|
||||
sock.write_u32::<BigEndian>(begin)?;
|
||||
sock.write_u32::<BigEndian>(length)?;
|
||||
Ok(())
|
||||
pub fn try_clone(&mut self) -> io::Result<PeerConnection> {
|
||||
Ok(PeerConnection {
|
||||
sock: self.sock.try_clone()?,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn shutdown(&mut self) {
|
||||
let _ = self.sock.shutdown(Shutdown::Both);
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user