From 29758c6d8f1ec248fd6c090d7d893c20f0ee9fec Mon Sep 17 00:00:00 2001 From: Igor Katson Date: Sun, 29 Sep 2024 10:13:51 +0100 Subject: [PATCH] Generic peer IP --- crates/librqbit/src/peer_connection.rs | 7 +- .../src/extended/handshake.rs | 78 ++--------- .../peer_binary_protocol/src/extended/ip.rs | 131 ++++++++++++++++++ .../peer_binary_protocol/src/extended/mod.rs | 4 + 4 files changed, 146 insertions(+), 74 deletions(-) create mode 100644 crates/peer_binary_protocol/src/extended/ip.rs diff --git a/crates/librqbit/src/peer_connection.rs b/crates/librqbit/src/peer_connection.rs index cd571695..3d54cf2e 100644 --- a/crates/librqbit/src/peer_connection.rs +++ b/crates/librqbit/src/peer_connection.rs @@ -14,10 +14,7 @@ use librqbit_core::{ }; use parking_lot::RwLock; use peer_binary_protocol::{ - extended::{ - handshake::{ExtendedHandshake, YourIP}, - ExtendedMessage, - }, + extended::{handshake::ExtendedHandshake, ExtendedMessage, PeerIP}, serialize_piece_preamble, Handshake, Message, MessageOwned, PIECE_MESSAGE_DEFAULT_LEN, }; use serde::{Deserialize, Serialize}; @@ -251,7 +248,7 @@ impl PeerConnection { if supports_extended { let mut my_extended = ExtendedHandshake::new(); my_extended.v = Some(ByteBuf(crate::client_name_and_version().as_bytes())); - my_extended.yourip = Some(YourIP(self.addr.ip())); + my_extended.yourip = Some(PeerIP(self.addr.ip())); self.handler .update_my_extended_handshake(&mut my_extended)?; let my_extended = Message::Extended(ExtendedMessage::Handshake(my_extended)); diff --git a/crates/peer_binary_protocol/src/extended/handshake.rs b/crates/peer_binary_protocol/src/extended/handshake.rs index 66a004fa..2a6824df 100644 --- a/crates/peer_binary_protocol/src/extended/handshake.rs +++ b/crates/peer_binary_protocol/src/extended/handshake.rs @@ -3,13 +3,13 @@ use std::{collections::HashMap, net::IpAddr}; use buffers::{ByteBuf, ByteBufT}; use bytes::Bytes; use clone_to_owned::CloneToOwned; -use serde::{Deserialize, Deserializer, Serialize}; +use serde::{Deserialize, Serialize}; use crate::{ EXTENDED_UT_METADATA_KEY, EXTENDED_UT_PEX_KEY, MY_EXTENDED_UT_METADATA, MY_EXTENDED_UT_PEX, }; -use super::PeerExtendedMessageIds; +use super::{PeerExtendedMessageIds, PeerIP4, PeerIP6, PeerIPAny}; #[derive(Deserialize, Serialize, Debug, Default)] pub struct ExtendedHandshake { @@ -20,11 +20,11 @@ pub struct ExtendedHandshake { #[serde(skip_serializing_if = "Option::is_none")] pub v: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub yourip: Option, + pub yourip: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub ipv6: Option, + pub ipv6: Option, #[serde(skip_serializing_if = "Option::is_none")] - pub ipv4: Option, + pub ipv4: Option, #[serde(skip_serializing_if = "Option::is_none")] pub reqq: Option, #[serde(skip_serializing_if = "Option::is_none")] @@ -72,18 +72,10 @@ where pub fn ip_addr(&self) -> Option { if let Some(ref b) = self.ipv4 { - let b = b.as_slice(); - if b.len() == 4 { - let ip_bytes: &[u8; 4] = b[0..4].try_into().unwrap(); // Safe to unwrap as we check slice length - return Some(IpAddr::from(*ip_bytes)); - } + return Some(b.0.into()); } if let Some(ref b) = self.ipv6 { - let b = b.as_slice(); - if b.len() == 16 { - let ip_bytes: &[u8; 16] = b[0..16].try_into().unwrap(); // Safe to unwrap as we check slice length - return Some(IpAddr::from(*ip_bytes)); - } + return Some(b.0.into()); } None } @@ -106,8 +98,8 @@ where p: self.p, v: self.v.clone_to_owned(within_buffer), yourip: self.yourip, - ipv6: self.ipv6.clone_to_owned(within_buffer), - ipv4: self.ipv4.clone_to_owned(within_buffer), + ipv6: self.ipv6, + ipv4: self.ipv4, reqq: self.reqq, metadata_size: self.metadata_size, complete_ago: self.complete_ago, @@ -115,55 +107,3 @@ where } } } - -#[derive(Debug, Clone, Copy)] -pub struct YourIP(pub IpAddr); - -impl Serialize for YourIP { - fn serialize(&self, serializer: S) -> Result - where - S: serde::Serializer, - { - match self.0 { - IpAddr::V4(ipv4) => { - let buf = ipv4.octets(); - serializer.serialize_bytes(&buf) - } - IpAddr::V6(ipv6) => { - let buf = ipv6.octets(); - serializer.serialize_bytes(&buf) - } - } - } -} - -impl<'de> Deserialize<'de> for YourIP { - fn deserialize(de: D) -> Result - where - D: Deserializer<'de>, - { - struct Visitor {} - impl<'de> serde::de::Visitor<'de> for Visitor { - type Value = YourIP; - - fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!(f, "expecting 4 bytes of ipv4 or 16 bytes of ipv6") - } - - fn visit_bytes(self, v: &[u8]) -> Result - where - E: serde::de::Error, - { - if v.len() == 4 { - let ip_bytes: &[u8; 4] = v[0..4].try_into().unwrap(); // Safe to unwrap as we check slice length - return Ok(YourIP(IpAddr::from(*ip_bytes))); - } else if v.len() == 16 { - let ip_bytes: &[u8; 16] = v[0..16].try_into().unwrap(); // Safe to unwrap as we check slice length - return Ok(YourIP(IpAddr::from(*ip_bytes))); - } - Err(E::custom("expected 4 or 16 byte address")) - } - } - de.deserialize_bytes(Visitor {}) - } -} diff --git a/crates/peer_binary_protocol/src/extended/ip.rs b/crates/peer_binary_protocol/src/extended/ip.rs new file mode 100644 index 00000000..ad68abfd --- /dev/null +++ b/crates/peer_binary_protocol/src/extended/ip.rs @@ -0,0 +1,131 @@ +use std::{ + marker::PhantomData, + net::{IpAddr, Ipv4Addr, Ipv6Addr}, +}; + +use serde::{Deserialize, Deserializer, Serialize}; + +enum IpOctets { + V4([u8; 4]), + V6([u8; 16]), +} + +impl IpOctets { + fn as_slice(&self) -> &[u8] { + match &self { + IpOctets::V4(s) => s, + IpOctets::V6(s) => s, + } + } +} + +#[derive(Debug, Clone, Copy)] +pub struct PeerIP(pub T); +pub type PeerIP4 = PeerIP; +pub type PeerIP6 = PeerIP; +pub type PeerIPAny = PeerIP; + +trait IpLike: Sized { + fn octets(&self) -> IpOctets; + fn try_from_slice(b: &[u8]) -> Option; + fn expecting() -> &'static str; +} + +impl IpLike for Ipv4Addr { + fn octets(&self) -> IpOctets { + IpOctets::V4(self.octets()) + } + + fn try_from_slice(b: &[u8]) -> Option { + let arr: [u8; 4] = b.try_into().ok()?; + Some(arr.into()) + } + + fn expecting() -> &'static str { + "expecting 4 bytes of ipv4" + } +} + +impl IpLike for Ipv6Addr { + fn octets(&self) -> IpOctets { + IpOctets::V6(self.octets()) + } + + fn try_from_slice(b: &[u8]) -> Option { + let arr: [u8; 16] = b.try_into().ok()?; + Some(arr.into()) + } + + fn expecting() -> &'static str { + "expecting 16 bytes of ipv6" + } +} + +impl IpLike for IpAddr { + fn octets(&self) -> IpOctets { + match self { + IpAddr::V4(ipv4_addr) => IpOctets::V4(ipv4_addr.octets()), + IpAddr::V6(ipv6_addr) => IpOctets::V6(ipv6_addr.octets()), + } + } + + fn try_from_slice(b: &[u8]) -> Option { + match b.len() { + 4 => Ipv4Addr::try_from_slice(b).map(Into::into), + 16 => Ipv6Addr::try_from_slice(b).map(Into::into), + _ => None, + } + } + + fn expecting() -> &'static str { + "expecting 4 or 16 bytes of ipv4 or ipv6" + } +} + +impl Serialize for PeerIP +where + T: IpLike, +{ + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + serializer.serialize_bytes(self.0.octets().as_slice()) + } +} + +impl<'de, T> Deserialize<'de> for PeerIP +where + T: IpLike, +{ + fn deserialize(deserializer: D) -> Result + where + D: Deserializer<'de>, + { + struct Visitor { + p: PhantomData, + } + impl<'de, T> serde::de::Visitor<'de> for Visitor + where + T: IpLike, + { + type Value = PeerIP; + + fn expecting(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + f.write_str(T::expecting()) + } + + fn visit_bytes(self, v: &[u8]) -> Result + where + E: serde::de::Error, + { + T::try_from_slice(v) + .map(PeerIP) + .ok_or_else(|| E::custom(T::expecting())) + } + } + deserializer.deserialize_bytes(Visitor { + p: Default::default(), + }) + } +} diff --git a/crates/peer_binary_protocol/src/extended/mod.rs b/crates/peer_binary_protocol/src/extended/mod.rs index 85296ba2..51c9459c 100644 --- a/crates/peer_binary_protocol/src/extended/mod.rs +++ b/crates/peer_binary_protocol/src/extended/mod.rs @@ -14,6 +14,10 @@ use self::{handshake::ExtendedHandshake, ut_metadata::UtMetadata}; use super::MessageDeserializeError; pub mod handshake; +mod ip; + +pub use ip::{PeerIP, PeerIP4, PeerIP6, PeerIPAny}; + pub mod ut_metadata; pub mod ut_pex;