271 lines
8.1 KiB
Rust
271 lines
8.1 KiB
Rust
use std::{
|
|
io::{Cursor, Write},
|
|
net::{Ipv4Addr, SocketAddrV4, SocketAddrV6},
|
|
};
|
|
|
|
use crate::bittorrent::{
|
|
info_hash::InfoHash,
|
|
peer::PeerId,
|
|
protocol::{Action, Event},
|
|
};
|
|
|
|
pub const CONNECT_REQUEST_SIZE: usize = 16;
|
|
pub const MIN_ANNOUNCE_REQUEST_SIZE: usize = 98;
|
|
pub const MIN_SCRAPE_REQUEST_SIZE: usize = 36;
|
|
|
|
#[derive(Debug)]
|
|
pub struct ConnectRequest {
|
|
pub protocol_id: i64,
|
|
pub transaction_id: i32,
|
|
}
|
|
|
|
impl ConnectRequest {
|
|
pub fn from_bytes(buf: &[u8]) -> Option<ConnectRequest> {
|
|
if buf.len() != CONNECT_REQUEST_SIZE {
|
|
return None;
|
|
}
|
|
|
|
Some(ConnectRequest {
|
|
protocol_id: i64::from_be_bytes(buf[0..8].try_into().ok()?),
|
|
transaction_id: i32::from_be_bytes(buf[12..16].try_into().ok()?),
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct ConnectResponse {
|
|
pub transaction_id: i32,
|
|
pub connection_id: i64,
|
|
}
|
|
|
|
impl ConnectResponse {
|
|
pub fn write_to_buf(&self, buf: &mut [u8]) -> usize {
|
|
let mut c = Cursor::new(buf);
|
|
let mut written: usize = 0;
|
|
|
|
const ACTION: i32 = Action::Connect as i32;
|
|
|
|
written += c.write(&ACTION.to_be_bytes()).unwrap();
|
|
written += c.write(&self.transaction_id.to_be_bytes()).unwrap();
|
|
written += c.write(&self.connection_id.to_be_bytes()).unwrap();
|
|
|
|
return written;
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct AnnounceRequest {
|
|
pub connection_id: i64,
|
|
pub transaction_id: i32,
|
|
/// Info hash of the torrent
|
|
pub info_hash: InfoHash,
|
|
/// Peer's ID
|
|
pub peer_id: PeerId,
|
|
/// Number of bytes peer has downloaded this session
|
|
pub downloaded: i64,
|
|
/// Number of bytes peer has remaining to complete the download
|
|
pub left: i64,
|
|
/// Number of bytes peer has uploaded this session
|
|
pub uploaded: i64,
|
|
/// One of 4 announce events. `Event::None` is used to update presence but not change peer's state.
|
|
pub event: Event,
|
|
/// Peer's desired IPv4 address to be sent the announce response on. `None` implies that the packet's source address should be used.
|
|
pub ipv4_address: Option<Ipv4Addr>,
|
|
/// Client's unique key for tracking a single client over multiple connections (sockets)
|
|
pub key: u32,
|
|
/// Maximum number of info hashes the client is requesting. `None` implies a tracker-defined default value should be used.
|
|
pub num_want: Option<u32>,
|
|
/// Peer's desired port to be the announce response on. If `ipv4_address` is `Some` then this should be used.
|
|
pub port: u16,
|
|
}
|
|
|
|
impl AnnounceRequest {
|
|
pub fn from_bytes(buf: &[u8]) -> Option<AnnounceRequest> {
|
|
if buf.len() < MIN_ANNOUNCE_REQUEST_SIZE {
|
|
return None;
|
|
}
|
|
|
|
Some(AnnounceRequest {
|
|
connection_id: i64::from_be_bytes(buf[0..8].try_into().ok()?),
|
|
transaction_id: i32::from_be_bytes(buf[12..16].try_into().ok()?),
|
|
info_hash: InfoHash::from_bytes(buf[16..36].try_into().ok()?)?,
|
|
peer_id: PeerId::from_bytes(buf[36..56].try_into().ok()?)?,
|
|
downloaded: i64::from_be_bytes(buf[56..64].try_into().ok()?),
|
|
left: i64::from_be_bytes(buf[64..72].try_into().ok()?),
|
|
uploaded: i64::from_be_bytes(buf[72..80].try_into().ok()?),
|
|
event: Event::from_i32(i32::from_be_bytes(buf[80..84].try_into().ok()?))?,
|
|
ipv4_address: if buf[84..88].iter().all(|b| *b == 0x00) {
|
|
None
|
|
} else {
|
|
Some(Ipv4Addr::from_bits(u32::from_be_bytes(
|
|
buf[84..88].try_into().ok()?,
|
|
)))
|
|
},
|
|
key: u32::from_be_bytes(buf[88..92].try_into().ok()?),
|
|
num_want: {
|
|
let want = i32::from_be_bytes(buf[92..96].try_into().ok()?);
|
|
if want < 0 { None } else { Some(want as u32) }
|
|
},
|
|
port: u16::from_be_bytes(buf[96..98].try_into().ok()?),
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum AnnounceResponse {
|
|
V4(AnnounceResponseV4),
|
|
V6(AnnounceResponseV6),
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct AnnounceResponseCommon {
|
|
pub transaction_id: i32,
|
|
/// Number of seconds peer should wait before re-announcing
|
|
pub interval: i32,
|
|
/// Number of peers still downloading
|
|
pub leechers: i32,
|
|
/// Number of peers who have completed the download and are seeding
|
|
pub seeders: i32,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct AnnounceResponseV4 {
|
|
pub common: AnnounceResponseCommon,
|
|
pub peers: Vec<SocketAddrV4>,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct AnnounceResponseV6 {
|
|
pub common: AnnounceResponseCommon,
|
|
pub peers: Vec<SocketAddrV6>,
|
|
}
|
|
|
|
impl AnnounceResponseCommon {
|
|
pub fn write_to_buf(&self, buf: &mut [u8]) -> usize {
|
|
let mut c = Cursor::new(buf);
|
|
let mut written: usize = 0;
|
|
|
|
const ACTION: i32 = Action::Announce as i32;
|
|
|
|
written += c.write(&ACTION.to_be_bytes()).unwrap();
|
|
written += c.write(&self.transaction_id.to_be_bytes()).unwrap();
|
|
written += c.write(&self.interval.to_be_bytes()).unwrap();
|
|
written += c.write(&self.leechers.to_be_bytes()).unwrap();
|
|
written += c.write(&self.seeders.to_be_bytes()).unwrap();
|
|
|
|
return written;
|
|
}
|
|
}
|
|
|
|
impl AnnounceResponseV4 {
|
|
pub fn write_to_buf(&self, buf: &mut [u8]) -> usize {
|
|
let mut written = self.common.write_to_buf(buf);
|
|
let mut c = Cursor::new(buf);
|
|
c.set_position(written as u64);
|
|
|
|
for peer in &self.peers {
|
|
written += c.write(&peer.ip().octets()).unwrap();
|
|
written += c.write(&peer.port().to_be_bytes()).unwrap();
|
|
}
|
|
|
|
return written;
|
|
}
|
|
}
|
|
|
|
impl AnnounceResponseV6 {
|
|
pub fn write_to_buf(&self, buf: &mut [u8]) -> usize {
|
|
let mut written = self.common.write_to_buf(buf);
|
|
let mut c = Cursor::new(buf);
|
|
c.set_position(written as u64);
|
|
|
|
for peer in &self.peers {
|
|
written += c.write(&peer.ip().octets()).unwrap();
|
|
written += c.write(&peer.port().to_be_bytes()).unwrap();
|
|
}
|
|
|
|
return written;
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct ScrapeRequest {
|
|
pub connection_id: i64,
|
|
pub transaction_id: i32,
|
|
pub info_hashes: Vec<InfoHash>,
|
|
}
|
|
|
|
impl ScrapeRequest {
|
|
pub fn from_bytes(buf: &[u8]) -> Option<ScrapeRequest> {
|
|
if buf.len() < MIN_SCRAPE_REQUEST_SIZE {
|
|
return None;
|
|
}
|
|
|
|
Some(ScrapeRequest {
|
|
connection_id: i64::from_be_bytes(buf[0..8].try_into().ok()?),
|
|
transaction_id: i32::from_be_bytes(buf[12..16].try_into().ok()?),
|
|
info_hashes: buf[16..]
|
|
.chunks_exact(20)
|
|
.map(|b| InfoHash {
|
|
bytes: b.try_into().unwrap(), // unwrap: `chunks_exact` guarantees the size
|
|
})
|
|
.collect(),
|
|
})
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Default)]
|
|
pub struct ScrapeStats {
|
|
/// Number of connected peers who have completed the download and are seeding
|
|
pub seeders: i32,
|
|
/// Number of peers who have ever completed the download
|
|
pub completed: i32,
|
|
/// Number of connected peers who have not completed the download and are seeding
|
|
pub leechers: i32,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct ScrapeResponse {
|
|
pub transaction_id: i32,
|
|
pub stats: Vec<ScrapeStats>,
|
|
}
|
|
|
|
impl ScrapeResponse {
|
|
pub fn write_to_buf(&self, buf: &mut [u8]) -> usize {
|
|
let mut c = Cursor::new(buf);
|
|
let mut written: usize = 0;
|
|
|
|
const ACTION: i32 = Action::Scrape as i32;
|
|
written += c.write(&ACTION.to_be_bytes()).unwrap();
|
|
written += c.write(&self.transaction_id.to_be_bytes()).unwrap();
|
|
|
|
for s in &self.stats {
|
|
written += c.write(&s.seeders.to_be_bytes()).unwrap();
|
|
written += c.write(&s.completed.to_be_bytes()).unwrap();
|
|
written += c.write(&s.leechers.to_be_bytes()).unwrap();
|
|
}
|
|
|
|
return written;
|
|
}
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub struct ErrorResponse {
|
|
pub transaction_id: i32,
|
|
pub message: String,
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum UdpRequest {
|
|
Connect(ConnectRequest),
|
|
Announce(AnnounceRequest),
|
|
Scrape(ScrapeRequest),
|
|
}
|
|
|
|
#[derive(Debug)]
|
|
pub enum UdpResponse {
|
|
Connect(ConnectResponse),
|
|
Announce(AnnounceResponse),
|
|
Scrape(ScrapeResponse),
|
|
Error(ErrorResponse),
|
|
}
|