bt/src/net/peer.rs

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(())
}