simplify and unify request parsing

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<InfoHash> {
+        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;

src/bittorrent/peer.rs

@@ -8,6 +8,12 @@ pub struct PeerId {
 }
 
 impl PeerId {
+    pub fn from_bytes(buf: &[u8]) -> Option<PeerId> {
+        Some(PeerId {
+            bytes: buf.try_into().ok()?,
+        })
+    }
+
     pub fn to_string_lossy(&self) -> String {
         String::from_utf8_lossy(&self.bytes).into()
     }

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::<u16>();
-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<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,
@@ -75,6 +79,38 @@ pub struct AnnounceRequest {
     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),
@@ -106,48 +142,48 @@ pub struct AnnounceResponseV6 {
 
 impl AnnounceResponseCommon {
     pub fn write_to_buf(&self, buf: &mut [u8]) -> usize {
-        const ACTION: i32 = Action::Announce as i32;
-        buf[0..4].copy_from_slice(&ACTION.to_be_bytes());
-        buf[4..8].copy_from_slice(&self.transaction_id.to_be_bytes());
-        buf[8..12].copy_from_slice(&self.interval.to_be_bytes());
-        buf[12..16].copy_from_slice(&self.leechers.to_be_bytes());
-        buf[16..20].copy_from_slice(&self.seeders.to_be_bytes());
+        let mut c = Cursor::new(buf);
+        let mut written: usize = 0;
 
-        return MIN_ANNOUNCE_RESPONSE_SIZE;
+        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 written = self.common.write_to_buf(buf);
+        let mut written = self.common.write_to_buf(buf);
+        let mut c = Cursor::new(buf);
+        c.set_position(written as u64);
 
-        for (offset, addr) in (written..buf.len())
-            .step_by(IPV4_ADDR_PAIR_SIZE)
-            .zip(&self.peers)
-        {
-            buf[offset..offset + IPV4_SIZE].copy_from_slice(&addr.ip().octets());
-            buf[offset + IPV4_SIZE..offset + IPV4_ADDR_PAIR_SIZE]
-                .copy_from_slice(&addr.port().to_be_bytes());
+        for peer in &self.peers {
+            written += c.write(&peer.ip().octets()).unwrap();
+            written += c.write(&peer.port().to_be_bytes()).unwrap();
         }
 
-        return written + self.peers.len() * IPV4_ADDR_PAIR_SIZE;
+        return written;
     }
 }
 
 impl AnnounceResponseV6 {
     pub fn write_to_buf(&self, buf: &mut [u8]) -> usize {
-        let written = self.common.write_to_buf(buf);
+        let mut written = self.common.write_to_buf(buf);
+        let mut c = Cursor::new(buf);
+        c.set_position(written as u64);
 
-        for (offset, addr) in (written..buf.len())
-            .step_by(IPV6_ADDR_PAIR_SIZE)
-            .zip(&self.peers)
-        {
-            buf[offset..offset + IPV6_SIZE].copy_from_slice(&addr.ip().octets());
-            buf[offset + IPV6_SIZE..offset + IPV6_ADDR_PAIR_SIZE]
-                .copy_from_slice(&addr.port().to_be_bytes());
+        for peer in &self.peers {
+            written += c.write(&peer.ip().octets()).unwrap();
+            written += c.write(&peer.port().to_be_bytes()).unwrap();
         }
 
-        return written + self.peers.len() * IPV6_ADDR_PAIR_SIZE;
+        return written;
     }
 }
 
@@ -158,6 +194,25 @@ pub struct ScrapeRequest {
     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

src/udp_server.rs

@@ -131,84 +131,20 @@ fn try_parse_packet(buf: &[u8]) -> Option<UdpRequest> {
 }
 
 fn try_parse_connect(buf: &[u8]) -> Option<ConnectRequest> {
-    // 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<AnnounceRequest> {
-    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<Ipv4Addr> = 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<u32> = {
-        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<ScrapeRequest> {
-    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<InfoHash> = 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)
 }