bittorrent_udp: response_to_bytes: improve code

This commit is contained in:
Joakim Frostegård 2020-04-09 21:30:58 +02:00
parent f203250155
commit acfd2a75d8

View file

@ -1,26 +1,25 @@
use byteorder::{ReadBytesExt, WriteBytesExt, NetworkEndian}; use std::io::{self, Cursor, Write};
use std::io;
use std::io::{Cursor, Write};
use std::net::{IpAddr, Ipv6Addr, Ipv4Addr}; use std::net::{IpAddr, Ipv6Addr, Ipv4Addr};
use crate::types::{self, *}; use byteorder::{ReadBytesExt, WriteBytesExt, NetworkEndian};
use crate::types::*;
#[inline] #[inline]
pub fn response_to_bytes( pub fn response_to_bytes(
bytes: &mut impl Write, bytes: &mut impl Write,
response: types::Response, response: Response,
ip_version: types::IpVersion ip_version: IpVersion
){ ){
match response { match response {
types::Response::Connect(r) => { Response::Connect(r) => {
bytes.write_i32::<NetworkEndian>(0).unwrap(); bytes.write_i32::<NetworkEndian>(0).unwrap();
bytes.write_i32::<NetworkEndian>(r.transaction_id.0).unwrap(); bytes.write_i32::<NetworkEndian>(r.transaction_id.0).unwrap();
bytes.write_i64::<NetworkEndian>(r.connection_id.0).unwrap(); bytes.write_i64::<NetworkEndian>(r.connection_id.0).unwrap();
}, },
types::Response::Announce(r) => { Response::Announce(r) => {
bytes.write_i32::<NetworkEndian>(1).unwrap(); bytes.write_i32::<NetworkEndian>(1).unwrap();
bytes.write_i32::<NetworkEndian>(r.transaction_id.0).unwrap(); bytes.write_i32::<NetworkEndian>(r.transaction_id.0).unwrap();
bytes.write_i32::<NetworkEndian>(r.announce_interval.0).unwrap(); bytes.write_i32::<NetworkEndian>(r.announce_interval.0).unwrap();
@ -29,33 +28,24 @@ pub fn response_to_bytes(
// Write peer IPs and ports. Silently ignore peers with wrong // Write peer IPs and ports. Silently ignore peers with wrong
// IP version // IP version
for peer in r.peers { if ip_version == IpVersion::IPv4 {
let mut correct = false; for peer in r.peers {
if let IpAddr::V4(ip) = peer.ip_address {
match peer.ip_address { bytes.write_all(&ip.octets()).unwrap();
IpAddr::V4(ip) => { bytes.write_u16::<NetworkEndian>(peer.port.0).unwrap();
if let types::IpVersion::IPv4 = ip_version {
bytes.write_all(&ip.octets()).unwrap();
correct = true;
}
},
IpAddr::V6(ip) => {
if let types::IpVersion::IPv6 = ip_version {
bytes.write_all(&ip.octets()).unwrap();
correct = true;
}
} }
} }
} else {
if correct { for peer in r.peers {
bytes.write_u16::<NetworkEndian>(peer.port.0).unwrap(); if let IpAddr::V6(ip) = peer.ip_address {
bytes.write_all(&ip.octets()).unwrap();
bytes.write_u16::<NetworkEndian>(peer.port.0).unwrap();
}
} }
} }
}, },
types::Response::Scrape(r) => { Response::Scrape(r) => {
bytes.write_i32::<NetworkEndian>(2).unwrap(); bytes.write_i32::<NetworkEndian>(2).unwrap();
bytes.write_i32::<NetworkEndian>(r.transaction_id.0).unwrap(); bytes.write_i32::<NetworkEndian>(r.transaction_id.0).unwrap();
@ -69,7 +59,7 @@ pub fn response_to_bytes(
} }
}, },
types::Response::Error(r) => { Response::Error(r) => {
bytes.write_i32::<NetworkEndian>(3).unwrap(); bytes.write_i32::<NetworkEndian>(3).unwrap();
bytes.write_i32::<NetworkEndian>(r.transaction_id.0).unwrap(); bytes.write_i32::<NetworkEndian>(r.transaction_id.0).unwrap();