proto: basic peer wire protocol that can download

src/net/peer.rs

@@ -0,0 +1,164 @@
+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<u8>, },
+    Request { index: u32, begin: u32, length: u32, },
+    Piece { index: u32, begin: u32, block: Vec<u8>, },
+    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::<BigEndian>(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<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);
+            }
+            sock.read_exact(&mut bitfield)?;
+            Packet::Bitfield {
+                bitfield: bitfield,
+            }
+        }
+        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)?;
+
+            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(sock: &mut TcpStream) -> io::Result<()> {
+    sock.write(b"")?;
+    Ok(())
+}
+
+pub fn send_choke(sock: &mut TcpStream) -> io::Result<()> {
+    sock.write_u32::<BigEndian>(1)?;
+    sock.write_u8(0)?;
+    Ok(())
+}
+
+pub fn send_unchoke(sock: &mut TcpStream) -> io::Result<()> {
+    sock.write_u32::<BigEndian>(1)?;
+    sock.write_u8(1)?;
+    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(sock: &mut TcpStream) -> io::Result<()> {
+    sock.write_u32::<BigEndian>(1)?;
+    sock.write_u8(3)?;
+    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(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(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(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(())
+}