148 lines
4.2 KiB
Rust
148 lines
4.2 KiB
Rust
use parking_lot::RwLock;
|
|
use rust_p2p_core::nat::NatInfo;
|
|
use rust_p2p_core::route::Index;
|
|
use rust_p2p_core::tunnel::udp::UDPIndex;
|
|
use std::collections::HashMap;
|
|
use std::net::{Ipv4Addr, SocketAddr};
|
|
use std::sync::Arc;
|
|
|
|
#[derive(Clone, Default)]
|
|
pub struct MyNatInfo {
|
|
nat_info: Arc<RwLock<Option<NatInfo>>>,
|
|
}
|
|
impl MyNatInfo {
|
|
pub fn get(&self) -> Option<NatInfo> {
|
|
self.nat_info.read().clone()
|
|
}
|
|
pub fn update_public_addr(&self, index: Index, addr: SocketAddr) {
|
|
let (ip, port) = if let Some(r) = mapping_addr(addr) {
|
|
r
|
|
} else {
|
|
return;
|
|
};
|
|
log::debug!("public_addr:{},{},index={index:?}", ip, port);
|
|
let mut nat_info = self.nat_info.write();
|
|
let Some(nat_info) = nat_info.as_mut() else {
|
|
return;
|
|
};
|
|
if rust_p2p_core::extend::addr::is_ipv4_global(&ip) {
|
|
if !nat_info.public_ips.contains(&ip) {
|
|
nat_info.public_ips.push(ip);
|
|
}
|
|
match index {
|
|
Index::Udp(index) => {
|
|
let index = match index {
|
|
UDPIndex::MainV4(index) => index,
|
|
UDPIndex::MainV6(index) => index,
|
|
UDPIndex::SubV4(_) => return,
|
|
};
|
|
if let Some(p) = nat_info.public_udp_ports.get_mut(index) {
|
|
*p = port;
|
|
}
|
|
}
|
|
Index::Tcp(_) => {
|
|
nat_info.public_tcp_port = port;
|
|
}
|
|
_ => {}
|
|
}
|
|
} else {
|
|
log::debug!("not public addr: {addr:?}")
|
|
}
|
|
}
|
|
pub fn update_tcp_public_addr(&self, addr: SocketAddr) {
|
|
let SocketAddr::V4(addr) = addr else {
|
|
return;
|
|
};
|
|
let ip = *addr.ip();
|
|
let port = addr.port();
|
|
log::info!("tcp_public_addr, {}:{}", ip, port);
|
|
let mut nat_info = self.nat_info.write();
|
|
let Some(nat_info) = nat_info.as_mut() else {
|
|
return;
|
|
};
|
|
if ip.is_unspecified() && port == 0 {
|
|
nat_info.public_tcp_port = 0;
|
|
return;
|
|
}
|
|
if rust_p2p_core::extend::addr::is_ipv4_global(&ip) {
|
|
if !nat_info.public_ips.contains(&ip) {
|
|
nat_info.public_ips.push(ip);
|
|
}
|
|
nat_info.public_tcp_port = port;
|
|
} else {
|
|
log::debug!("not public addr: {addr:?}")
|
|
}
|
|
}
|
|
pub fn replace_nat_info(&self, nat_info: NatInfo) {
|
|
self.nat_info.write().replace(nat_info);
|
|
}
|
|
pub fn clear(&self) {
|
|
*self.nat_info.write() = None;
|
|
}
|
|
}
|
|
fn mapping_addr(addr: SocketAddr) -> Option<(Ipv4Addr, u16)> {
|
|
match addr {
|
|
SocketAddr::V4(addr) => Some((*addr.ip(), addr.port())),
|
|
SocketAddr::V6(addr) => addr.ip().to_ipv4_mapped().map(|ip| (ip, addr.port())),
|
|
}
|
|
}
|
|
|
|
#[derive(Copy, Clone, Debug)]
|
|
pub struct PunchState {
|
|
pub count: i64,
|
|
pub last_ts: i64,
|
|
}
|
|
|
|
#[derive(Clone, Default)]
|
|
pub struct PunchBackoff {
|
|
inner: Arc<RwLock<HashMap<Ipv4Addr, PunchState>>>,
|
|
}
|
|
|
|
impl PunchBackoff {
|
|
const MAX_BACKOFF_MS: i64 = 3_600_000; // 1h
|
|
const BASE_MS: i64 = 3000;
|
|
|
|
fn now() -> i64 {
|
|
crate::utils::time::now_ts_ms()
|
|
}
|
|
|
|
pub fn record(&self, ip: Ipv4Addr) {
|
|
let mut map = self.inner.write();
|
|
let entry = map.entry(ip).or_insert(PunchState {
|
|
count: 0,
|
|
last_ts: 0,
|
|
});
|
|
entry.count += 1;
|
|
entry.last_ts = Self::now();
|
|
}
|
|
#[allow(dead_code)]
|
|
pub fn reset(&self, ip: Ipv4Addr) {
|
|
self.inner.write().remove(&ip);
|
|
}
|
|
#[allow(dead_code)]
|
|
pub fn get(&self, ip: &Ipv4Addr) -> Option<PunchState> {
|
|
self.inner.read().get(ip).copied()
|
|
}
|
|
|
|
pub fn should_punch(&self, ip: Ipv4Addr) -> bool {
|
|
let map = self.inner.read();
|
|
let Some(state) = map.get(&ip) else {
|
|
return true;
|
|
};
|
|
|
|
let now = Self::now();
|
|
let elapsed = now - state.last_ts;
|
|
|
|
let mut backoff = Self::BASE_MS * state.count;
|
|
if backoff > Self::MAX_BACKOFF_MS {
|
|
backoff = Self::MAX_BACKOFF_MS;
|
|
}
|
|
|
|
elapsed >= backoff
|
|
}
|
|
#[allow(dead_code)]
|
|
pub fn clear(&self) {
|
|
self.inner.write().clear();
|
|
}
|
|
}
|