v2

2026-02-10 18:20:39 +08:00
parent 71da72da34
commit b4301a8106
261 changed files with 20179 additions and 21547 deletions
@@ -0,0 +1,98 @@
+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()
+    }
+}