diff --git a/src/bittorrent/info_hash.rs b/src/bittorrent/info_hash.rs index e0b9918..bb2d0c9 100644 --- a/src/bittorrent/info_hash.rs +++ b/src/bittorrent/info_hash.rs @@ -8,6 +8,14 @@ pub struct InfoHash { pub bytes: [u8; INFO_HASH_SIZE], } +impl InfoHash { + pub fn from_bytes(buf: &[u8]) -> Option { + Some(InfoHash { + bytes: buf.try_into().ok()?, + }) + } +} + impl fmt::Display for InfoHash { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { const INFO_HASH_STR_LEN: usize = INFO_HASH_SIZE * 2; diff --git a/src/bittorrent/peer.rs b/src/bittorrent/peer.rs index b1dcf3c..d1467c0 100644 --- a/src/bittorrent/peer.rs +++ b/src/bittorrent/peer.rs @@ -8,6 +8,12 @@ pub struct PeerId { } impl PeerId { + pub fn from_bytes(buf: &[u8]) -> Option { + Some(PeerId { + bytes: buf.try_into().ok()?, + }) + } + pub fn to_string_lossy(&self) -> String { String::from_utf8_lossy(&self.bytes).into() } diff --git a/src/bittorrent/udp.rs b/src/bittorrent/udp.rs index 232372b..674bc26 100644 --- a/src/bittorrent/udp.rs +++ b/src/bittorrent/udp.rs @@ -1,6 +1,6 @@ use std::{ io::{Cursor, Write}, - net::{Ipv4Addr, Ipv6Addr, SocketAddrV4, SocketAddrV6}, + net::{Ipv4Addr, SocketAddrV4, SocketAddrV6}, }; use crate::bittorrent::{ @@ -9,18 +9,9 @@ use crate::bittorrent::{ 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; -pub const MIN_CONNECTION_RESPONSE_SIZE: usize = 16; -pub const MIN_ANNOUNCE_RESPONSE_SIZE: usize = 20; -pub const MIN_SCRAPE_RESPONSE_SIZE: usize = 8; -pub const SCRAPE_RESPONSE_ENTRY_SIZE: usize = 12; - -pub const IPV4_SIZE: usize = Ipv4Addr::BITS as usize / 8; -pub const IPV6_SIZE: usize = Ipv6Addr::BITS as usize / 8; -pub const PORT_SIZE: usize = size_of::(); -pub const IPV4_ADDR_PAIR_SIZE: usize = IPV4_SIZE + PORT_SIZE; -pub const IPV6_ADDR_PAIR_SIZE: usize = IPV6_SIZE + PORT_SIZE; #[derive(Debug)] pub struct ConnectRequest { @@ -28,6 +19,19 @@ pub struct ConnectRequest { pub transaction_id: i32, } +impl ConnectRequest { + pub fn from_bytes(buf: &[u8]) -> Option { + 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, @@ -75,6 +79,38 @@ pub struct AnnounceRequest { pub port: u16, } +impl AnnounceRequest { + pub fn from_bytes(buf: &[u8]) -> Option { + 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), @@ -158,6 +194,25 @@ pub struct ScrapeRequest { pub info_hashes: Vec, } +impl ScrapeRequest { + pub fn from_bytes(buf: &[u8]) -> Option { + 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 diff --git a/src/udp_server.rs b/src/udp_server.rs index 48473e6..c97e79e 100644 --- a/src/udp_server.rs +++ b/src/udp_server.rs @@ -131,84 +131,20 @@ fn try_parse_packet(buf: &[u8]) -> Option { } fn try_parse_connect(buf: &[u8]) -> Option { - // Buffer length is checked to be at least 16 in `try_parse_packet` + let conn = ConnectRequest::from_bytes(buf)?; - let protocol_id: i64 = i64::from_be_bytes(buf[0..8].try_into().ok()?); - - if protocol_id != UDP_MAGIC { + // Disregard the request immediately if `protocol_id` is not the valid magic number + if conn.protocol_id != UDP_MAGIC { return None; } - Some(ConnectRequest { - protocol_id: protocol_id, - transaction_id: i32::from_be_bytes(buf[12..16].try_into().ok()?), - }) + Some(conn) } fn try_parse_announce(buf: &[u8]) -> Option { - if buf.len() < MIN_ANNOUNCE_REQUEST_SIZE { - return None; - } - - let connection_id: i64 = i64::from_be_bytes(buf[0..8].try_into().ok()?); - let transaction_id: i32 = i32::from_be_bytes(buf[12..16].try_into().ok()?); - let info_hash: InfoHash = InfoHash { - bytes: buf[16..36].try_into().ok()?, - }; - let peer_id: PeerId = PeerId { - bytes: buf[36..56].try_into().unwrap_or_default(), - }; - let downloaded: i64 = i64::from_be_bytes(buf[56..64].try_into().ok()?); - let left: i64 = i64::from_be_bytes(buf[64..72].try_into().ok()?); - let uploaded: i64 = i64::from_be_bytes(buf[72..80].try_into().ok()?); - let event: Event = Event::from_i32(i32::from_be_bytes(buf[80..84].try_into().ok()?))?; - let ip_addr: Option = 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()?, - ))) - }; - let key: u32 = u32::from_be_bytes(buf[88..92].try_into().ok()?); - let num_want: Option = { - let want = i32::from_be_bytes(buf[92..96].try_into().ok()?); - if want < 0 { None } else { Some(want as u32) } - }; - let port: u16 = u16::from_be_bytes(buf[96..98].try_into().ok()?); - - Some(AnnounceRequest { - connection_id: connection_id, - transaction_id: transaction_id, - info_hash: info_hash, - peer_id: peer_id, - downloaded: downloaded, - left: left, - uploaded: uploaded, - event: event, - ipv4_address: ip_addr, - key: key, - num_want: num_want, - port: port, - }) + AnnounceRequest::from_bytes(buf) } fn try_parse_scrape(buf: &[u8]) -> Option { - if buf.len() < MIN_SCRAPE_REQUEST_SIZE { - return None; - } - - let connection_id: i64 = i64::from_be_bytes(buf[0..8].try_into().ok()?); - let transaction_id: i32 = i32::from_be_bytes(buf[12..16].try_into().ok()?); - let info_hashes: Vec = buf[16..] - .chunks_exact(20) - .map(|b| InfoHash { - bytes: b.try_into().unwrap(), // unwrap: `chunks_exact` guarantees the size - }) - .collect(); - - Some(ScrapeRequest { - connection_id: connection_id, - transaction_id: transaction_id, - info_hashes: info_hashes, - }) + ScrapeRequest::from_bytes(buf) }