Some cleanups for utpex
This commit is contained in:
Igor Katson
parent
5e09525dd5
commit
90bfb85bcc
5 changed files with 56 additions and 132 deletions
|
|
@ -3,6 +3,7 @@ use std::net::{IpAddr, SocketAddr};
|
|||
use byteorder::{ByteOrder, BE};
|
||||
use bytes::Bytes;
|
||||
use clone_to_owned::CloneToOwned;
|
||||
use itertools::{EitherOrBoth, Itertools};
|
||||
use serde::{Deserialize, Serialize};
|
||||
|
||||
#[derive(Debug)]
|
||||
|
|
@ -11,46 +12,18 @@ pub struct PexPeerInfo {
|
|||
pub addr: SocketAddr,
|
||||
}
|
||||
|
||||
impl PexPeerInfo {
|
||||
pub fn from_bytes(buf: &[u8], flags: Option<u8>) -> anyhow::Result<Self> {
|
||||
let (ip, port) = match buf.len() {
|
||||
6 => {
|
||||
let ip_bytes: &[u8; 4] = (&buf[0..4]).try_into()?;
|
||||
let ip = IpAddr::from(*ip_bytes);
|
||||
let port = BE::read_u16(&buf[4..6]);
|
||||
(ip, port)
|
||||
}
|
||||
18 => {
|
||||
let ip_bytes: &[u8; 16] = (&buf[0..16]).try_into()?;
|
||||
let ip = IpAddr::from(*ip_bytes);
|
||||
let port = BE::read_u16(&buf[16..18]);
|
||||
(ip, port)
|
||||
}
|
||||
_ => anyhow::bail!("invalid pex peer info"),
|
||||
};
|
||||
Ok(Self {
|
||||
flags: flags.unwrap_or(0),
|
||||
addr: (ip, port).into(),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize, Default, Deserialize)]
|
||||
pub struct UtPex<B> {
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
added: Option<B>,
|
||||
added: B,
|
||||
#[serde(rename = "added.f")]
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
added_f: Option<B>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
added6: Option<B>,
|
||||
added6: B,
|
||||
#[serde(rename = "added6.f")]
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
added6_f: Option<B>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
dropped: Option<B>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
dropped6: Option<B>,
|
||||
dropped: B,
|
||||
dropped6: B,
|
||||
}
|
||||
|
||||
impl<B> CloneToOwned for UtPex<B>
|
||||
|
|
@ -70,97 +43,45 @@ where
|
|||
}
|
||||
}
|
||||
|
||||
pub struct PeerAddrIterator<'a> {
|
||||
addrs: &'a [u8],
|
||||
flags: &'a [u8],
|
||||
offset: usize,
|
||||
addr_size: usize,
|
||||
}
|
||||
|
||||
|
||||
|
||||
impl<'a> Iterator for PeerAddrIterator<'a> {
|
||||
type Item = PexPeerInfo;
|
||||
fn next(&mut self) -> Option<Self::Item> {
|
||||
if self.offset*self.addr_size >= self.addrs.len() {
|
||||
return None;
|
||||
}
|
||||
|
||||
let addr = &self.addrs[self.offset*self.addr_size..(self.offset+1)*self.addr_size];
|
||||
let flags = self.flags.get(self.offset);
|
||||
self.offset += 1;
|
||||
Some(PexPeerInfo::from_bytes(addr, flags.cloned()).unwrap()) // safe to unwrap as we assure slice length
|
||||
|
||||
|
||||
}
|
||||
}
|
||||
|
||||
impl<B> UtPex<B>
|
||||
where
|
||||
B: AsRef<[u8]>,
|
||||
{
|
||||
pub fn added_peers<'a>(&'a self) -> anyhow::Result<Box<dyn Iterator<Item = PexPeerInfo> + 'a>> {
|
||||
if let Some(added) = &self.added {
|
||||
if added.as_ref().len() % 6 != 0 {
|
||||
anyhow::bail!("invalid pex added peers");
|
||||
}
|
||||
return Ok(Box::new(PeerAddrIterator {
|
||||
addrs: added.as_ref(),
|
||||
flags: self.added_f.as_ref().map(|f| f.as_ref()).unwrap_or(&[]),
|
||||
offset: 0,
|
||||
addr_size: 6,
|
||||
}));
|
||||
} else {
|
||||
return Ok(Box::new(std::iter::empty()));
|
||||
};
|
||||
fn added_peers_inner<'a>(
|
||||
&'a self,
|
||||
buf: &'a B,
|
||||
flags: &'a Option<B>,
|
||||
ip_len: usize,
|
||||
) -> impl Iterator<Item = PexPeerInfo> + 'a {
|
||||
let addrs = buf.as_ref().chunks_exact(ip_len + 2).map(move |c| {
|
||||
let ip = match ip_len {
|
||||
4 => IpAddr::from(TryInto::<[u8; 4]>::try_into(&c[..4]).unwrap()),
|
||||
16 => IpAddr::from(TryInto::<[u8; 16]>::try_into(&c[..16]).unwrap()),
|
||||
_ => unreachable!(),
|
||||
};
|
||||
let port = BE::read_u16(&c[ip_len..]);
|
||||
SocketAddr::new(ip, port)
|
||||
});
|
||||
let flags = flags
|
||||
.as_ref()
|
||||
.map(|b| b.as_ref().iter().copied())
|
||||
.into_iter()
|
||||
.flatten();
|
||||
addrs.zip_longest(flags).filter_map(|eob| match eob {
|
||||
EitherOrBoth::Both(addr, flags) => Some(PexPeerInfo { flags, addr }),
|
||||
EitherOrBoth::Left(addr) => Some(PexPeerInfo { flags: 0, addr }),
|
||||
EitherOrBoth::Right(_) => None,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn added_peers_v6<'a>(&'a self) -> anyhow::Result<Box<dyn Iterator<Item = PexPeerInfo> + 'a>> {
|
||||
if let Some(added) = &self.added6 {
|
||||
if added.as_ref().len() % 18 != 0 {
|
||||
anyhow::bail!("invalid pex added6 peers");
|
||||
}
|
||||
return Ok(Box::new(PeerAddrIterator {
|
||||
addrs: added.as_ref(),
|
||||
flags: self.added6_f.as_ref().map(|f| f.as_ref()).unwrap_or(&[]),
|
||||
offset: 0,
|
||||
addr_size: 18,
|
||||
}));
|
||||
} else {
|
||||
return Ok(Box::new(std::iter::empty()));
|
||||
};
|
||||
pub fn added_peers(&self) -> impl Iterator<Item = PexPeerInfo> + '_ {
|
||||
self.added_peers_inner(&self.added, &self.added_f, 4)
|
||||
.chain(self.added_peers_inner(&self.added6, &self.added6_f, 16))
|
||||
}
|
||||
|
||||
pub fn dropped_peers<'a>(&'a self) -> anyhow::Result<Box<dyn Iterator<Item = PexPeerInfo> + 'a>> {
|
||||
if let Some(dropped) = &self.dropped {
|
||||
if dropped.as_ref().len() % 6 != 0 {
|
||||
anyhow::bail!("invalid pex dropped peers");
|
||||
}
|
||||
return Ok(Box::new(PeerAddrIterator {
|
||||
addrs: dropped.as_ref(),
|
||||
flags: &[],
|
||||
offset: 0,
|
||||
addr_size: 6,
|
||||
}));
|
||||
} else {
|
||||
return Ok(Box::new(std::iter::empty()));
|
||||
};
|
||||
}
|
||||
|
||||
pub fn dropped_peers_v6<'a>(&'a self) -> anyhow::Result<Box<dyn Iterator<Item = PexPeerInfo> + 'a>> {
|
||||
if let Some(dropped) = &self.dropped6 {
|
||||
if dropped.as_ref().len() % 18 != 0 {
|
||||
anyhow::bail!("invalid pex dropped6 peers");
|
||||
}
|
||||
return Ok(Box::new(PeerAddrIterator {
|
||||
addrs: dropped.as_ref(),
|
||||
flags: &[],
|
||||
offset: 0,
|
||||
addr_size: 18,
|
||||
}));
|
||||
} else {
|
||||
return Ok(Box::new(std::iter::empty()));
|
||||
};
|
||||
pub fn dropped_peers(&self) -> impl Iterator<Item = PexPeerInfo> + '_ {
|
||||
self.added_peers_inner(&self.dropped, &None, 4)
|
||||
.chain(self.added_peers_inner(&self.dropped6, &None, 16))
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -184,11 +105,17 @@ mod tests {
|
|||
let msg = "64353a616464656431323ab99f9d14b56797f969861090373a61646465642e66323a0c00363a616464656436303a383a6164646564362e66303a373a64726f70706564303a383a64726f7070656436303a65";
|
||||
let bytes = decode_hex(msg);
|
||||
let pex = from_bytes::<UtPex<ByteBuf>>(&bytes).unwrap();
|
||||
let addrs: Vec<_> = pex.added_peers().unwrap().collect();
|
||||
let addrs: Vec<_> = pex.added_peers().collect();
|
||||
assert_eq!(2, addrs.len());
|
||||
assert_eq!("185.159.157.20:46439".parse::<SocketAddr>().unwrap(), addrs[0].addr);
|
||||
assert_eq!(
|
||||
"185.159.157.20:46439".parse::<SocketAddr>().unwrap(),
|
||||
addrs[0].addr
|
||||
);
|
||||
assert_eq!(12, addrs[0].flags);
|
||||
assert_eq!("151.249.105.134:4240".parse::<SocketAddr>().unwrap(), addrs[1].addr);
|
||||
assert_eq!(
|
||||
"151.249.105.134:4240".parse::<SocketAddr>().unwrap(),
|
||||
addrs[1].addr
|
||||
);
|
||||
assert_eq!(0, addrs[1].flags);
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -640,7 +640,7 @@ mod tests {
|
|||
fn test_extended_serialize() {
|
||||
let msg = Message::Extended(ExtendedMessage::Handshake(ExtendedHandshake::new()));
|
||||
let mut out = Vec::new();
|
||||
msg.serialize(&mut out, &|| Default::default()).unwrap();
|
||||
msg.serialize(&mut out, &Default::default).unwrap();
|
||||
dbg!(out);
|
||||
}
|
||||
|
||||
|
|
@ -656,7 +656,7 @@ mod tests {
|
|||
let (msg, size) = MessageBorrowed::deserialize(&buf).unwrap();
|
||||
assert_eq!(size, buf.len());
|
||||
let mut write_buf = Vec::new();
|
||||
msg.serialize(&mut write_buf, &|| Default::default()).unwrap();
|
||||
msg.serialize(&mut write_buf, &Default::default).unwrap();
|
||||
if buf != write_buf {
|
||||
{
|
||||
use std::io::Write;
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue