99 lines
3.4 KiB
Rust
99 lines
3.4 KiB
Rust
mod proto {
|
|
include!(concat!(env!("OUT_DIR"), "/protocol.client.rs"));
|
|
}
|
|
|
|
use anyhow::bail;
|
|
use bytes::BytesMut;
|
|
use prost::Message;
|
|
use std::net::{Ipv4Addr, Ipv6Addr};
|
|
|
|
use crate::protocol::ProtoToBytesMut;
|
|
pub use proto::*;
|
|
|
|
pub fn encode_nat_info(nat_info: &rust_p2p_core::nat::NatInfo) -> proto::NatInfo {
|
|
let nat_type = match nat_info.nat_type {
|
|
rust_p2p_core::nat::NatType::Cone => proto::NatType::Cone,
|
|
rust_p2p_core::nat::NatType::Symmetric => proto::NatType::Symmetric,
|
|
};
|
|
|
|
proto::NatInfo {
|
|
nat_type: nat_type.into(),
|
|
public_ips: nat_info.public_ips.iter().map(|v| (*v).into()).collect(),
|
|
public_udp_ports: nat_info
|
|
.public_udp_ports
|
|
.iter()
|
|
.map(|v| (*v).into())
|
|
.collect(),
|
|
public_port_range: nat_info.public_port_range.into(),
|
|
local_ipv4s: nat_info.local_ipv4s.iter().map(|v| (*v).into()).collect(),
|
|
ipv6: nat_info.ipv6.map(|v| v.octets().to_vec()),
|
|
local_udp_ports: nat_info
|
|
.local_udp_ports
|
|
.iter()
|
|
.map(|v| (*v).into())
|
|
.collect(),
|
|
local_tcp_port: nat_info.local_tcp_port.into(),
|
|
public_tcp_port: nat_info.public_tcp_port.into(),
|
|
}
|
|
}
|
|
pub fn decode_nat_info(msg: proto::NatInfo) -> anyhow::Result<rust_p2p_core::nat::NatInfo> {
|
|
let nat_type = match msg.nat_type() {
|
|
proto::NatType::Cone => rust_p2p_core::nat::NatType::Cone,
|
|
proto::NatType::Symmetric => rust_p2p_core::nat::NatType::Symmetric,
|
|
};
|
|
let ipv6: Option<[u8; 16]> = msg.ipv6.and_then(|v| v.as_slice().try_into().ok());
|
|
|
|
// Validate all ports fit in u16
|
|
let validate_port = |p: u32| -> anyhow::Result<u16> {
|
|
u16::try_from(p).map_err(|_| anyhow::anyhow!("invalid port number: {}", p))
|
|
};
|
|
|
|
let public_udp_ports: Result<Vec<_>, _> = msg
|
|
.public_udp_ports
|
|
.into_iter()
|
|
.map(validate_port)
|
|
.collect();
|
|
let local_udp_ports: Result<Vec<_>, _> =
|
|
msg.local_udp_ports.into_iter().map(validate_port).collect();
|
|
|
|
Ok(rust_p2p_core::nat::NatInfo {
|
|
nat_type,
|
|
public_ips: msg.public_ips.into_iter().map(|v| v.into()).collect(),
|
|
public_udp_ports: public_udp_ports?,
|
|
mapping_tcp_addr: vec![],
|
|
mapping_udp_addr: vec![],
|
|
public_port_range: validate_port(msg.public_port_range)?,
|
|
local_ipv4: msg
|
|
.local_ipv4s
|
|
.first()
|
|
.map(|v| (*v).into())
|
|
.unwrap_or(Ipv4Addr::UNSPECIFIED),
|
|
local_ipv4s: msg.local_ipv4s.into_iter().map(|v| v.into()).collect(),
|
|
ipv6: ipv6.map(Ipv6Addr::from),
|
|
local_udp_ports: local_udp_ports?,
|
|
local_tcp_port: validate_port(msg.local_tcp_port)?,
|
|
public_tcp_port: validate_port(msg.public_tcp_port)?,
|
|
})
|
|
}
|
|
#[derive(Clone, Debug)]
|
|
pub struct PunchInfo {
|
|
pub nat_info: rust_p2p_core::nat::NatInfo,
|
|
}
|
|
|
|
impl PunchInfo {
|
|
pub fn from_slice(buf: &[u8]) -> anyhow::Result<Self> {
|
|
let msg = proto::PunchInfo::decode(buf)?;
|
|
let Some(nat_info) = msg.nat_info else {
|
|
bail!("Punched info decode failed.");
|
|
};
|
|
let nat_info = decode_nat_info(nat_info)?;
|
|
Ok(Self { nat_info })
|
|
}
|
|
pub fn encode(&self) -> BytesMut {
|
|
let message = proto::PunchInfo {
|
|
nat_info: Some(encode_nat_info(&self.nat_info)),
|
|
};
|
|
message.encode_bytes_mut()
|
|
}
|
|
}
|