Initial Version
This commit is contained in:
sijie.sun
@@ -0,0 +1,433 @@
|
||||
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr, SocketAddrV4, SocketAddrV6};
|
||||
use std::sync::Arc;
|
||||
use std::time::Duration;
|
||||
|
||||
use crossbeam::atomic::AtomicCell;
|
||||
use easytier_rpc::{NatType, StunInfo};
|
||||
use stun_format::Attr;
|
||||
use tokio::net::{lookup_host, UdpSocket};
|
||||
use tokio::sync::RwLock;
|
||||
use tokio::task::JoinSet;
|
||||
|
||||
use crate::common::error::Error;
|
||||
|
||||
struct Stun {
|
||||
stun_server: String,
|
||||
req_repeat: u8,
|
||||
resp_timeout: Duration,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy)]
|
||||
struct BindRequestResponse {
|
||||
source_addr: SocketAddr,
|
||||
mapped_socket_addr: Option<SocketAddr>,
|
||||
changed_socket_addr: Option<SocketAddr>,
|
||||
|
||||
ip_changed: bool,
|
||||
port_changed: bool,
|
||||
}
|
||||
|
||||
impl BindRequestResponse {
|
||||
pub fn get_mapped_addr_no_check(&self) -> &SocketAddr {
|
||||
self.mapped_socket_addr.as_ref().unwrap()
|
||||
}
|
||||
}
|
||||
|
||||
impl Stun {
|
||||
pub fn new(stun_server: String) -> Self {
|
||||
Self {
|
||||
stun_server,
|
||||
req_repeat: 3,
|
||||
resp_timeout: Duration::from_millis(3000),
|
||||
}
|
||||
}
|
||||
|
||||
async fn wait_stun_response<'a, const N: usize>(
|
||||
&self,
|
||||
buf: &'a mut [u8; N],
|
||||
udp: &UdpSocket,
|
||||
tids: &Vec<u128>,
|
||||
) -> Result<(stun_format::Msg<'a>, SocketAddr), Error> {
|
||||
let mut now = tokio::time::Instant::now();
|
||||
let deadline = now + self.resp_timeout;
|
||||
|
||||
while now < deadline {
|
||||
let mut udp_buf = [0u8; 1500];
|
||||
let (len, remote_addr) =
|
||||
tokio::time::timeout(deadline - now, udp.recv_from(udp_buf.as_mut_slice()))
|
||||
.await??;
|
||||
now = tokio::time::Instant::now();
|
||||
|
||||
if len < 20 {
|
||||
continue;
|
||||
}
|
||||
|
||||
// TODO:: we cannot borrow `buf` directly in udp recv_from, so we copy it here
|
||||
unsafe { std::ptr::copy(udp_buf.as_ptr(), buf.as_ptr() as *mut u8, len) };
|
||||
|
||||
let msg = stun_format::Msg::<'a>::from(&buf[..]);
|
||||
tracing::trace!(b = ?&udp_buf[..len], ?msg, ?tids, "recv stun response");
|
||||
|
||||
if msg.typ().is_none() || msg.tid().is_none() {
|
||||
continue;
|
||||
}
|
||||
|
||||
if matches!(
|
||||
msg.typ().as_ref().unwrap(),
|
||||
stun_format::MsgType::BindingResponse
|
||||
) && tids.contains(msg.tid().as_ref().unwrap())
|
||||
{
|
||||
return Ok((msg, remote_addr));
|
||||
}
|
||||
}
|
||||
|
||||
Err(Error::Unknown)
|
||||
}
|
||||
|
||||
fn stun_addr(addr: stun_format::SocketAddr) -> SocketAddr {
|
||||
match addr {
|
||||
stun_format::SocketAddr::V4(ip, port) => {
|
||||
SocketAddr::V4(SocketAddrV4::new(Ipv4Addr::from(ip), port))
|
||||
}
|
||||
stun_format::SocketAddr::V6(ip, port) => {
|
||||
SocketAddr::V6(SocketAddrV6::new(Ipv6Addr::from(ip), port, 0, 0))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn extrace_mapped_addr(msg: &stun_format::Msg) -> Option<SocketAddr> {
|
||||
let mut mapped_addr = None;
|
||||
for x in msg.attrs_iter() {
|
||||
match x {
|
||||
Attr::MappedAddress(addr) => {
|
||||
if mapped_addr.is_none() {
|
||||
let _ = mapped_addr.insert(Self::stun_addr(addr));
|
||||
}
|
||||
}
|
||||
Attr::XorMappedAddress(addr) => {
|
||||
if mapped_addr.is_none() {
|
||||
let _ = mapped_addr.insert(Self::stun_addr(addr));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
mapped_addr
|
||||
}
|
||||
|
||||
fn extract_changed_addr(msg: &stun_format::Msg) -> Option<SocketAddr> {
|
||||
let mut changed_addr = None;
|
||||
for x in msg.attrs_iter() {
|
||||
match x {
|
||||
Attr::ChangedAddress(addr) => {
|
||||
if changed_addr.is_none() {
|
||||
let _ = changed_addr.insert(Self::stun_addr(addr));
|
||||
}
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
}
|
||||
changed_addr
|
||||
}
|
||||
|
||||
pub async fn bind_request(
|
||||
&self,
|
||||
source_port: u16,
|
||||
change_ip: bool,
|
||||
change_port: bool,
|
||||
) -> Result<BindRequestResponse, Error> {
|
||||
let stun_host = lookup_host(&self.stun_server)
|
||||
.await?
|
||||
.next()
|
||||
.ok_or(Error::NotFound)?;
|
||||
// let udp_socket = socket2::Socket::new(
|
||||
// match stun_host {
|
||||
// SocketAddr::V4(..) => socket2::Domain::IPV4,
|
||||
// SocketAddr::V6(..) => socket2::Domain::IPV6,
|
||||
// },
|
||||
// socket2::Type::DGRAM,
|
||||
// Some(socket2::Protocol::UDP),
|
||||
// )?;
|
||||
// udp_socket.set_reuse_port(true)?;
|
||||
// udp_socket.set_reuse_address(true)?;
|
||||
let udp = UdpSocket::bind(format!("0.0.0.0:{}", source_port)).await?;
|
||||
|
||||
// repeat req in case of packet loss
|
||||
let mut tids = vec![];
|
||||
for _ in 0..self.req_repeat {
|
||||
let mut buf = [0u8; 28];
|
||||
// memset buf
|
||||
unsafe { std::ptr::write_bytes(buf.as_mut_ptr(), 0, buf.len()) };
|
||||
let mut msg = stun_format::MsgBuilder::from(buf.as_mut_slice());
|
||||
msg.typ(stun_format::MsgType::BindingRequest).unwrap();
|
||||
let tid = rand::random::<u32>();
|
||||
msg.tid(tid as u128).unwrap();
|
||||
if change_ip || change_port {
|
||||
msg.add_attr(Attr::ChangeRequest {
|
||||
change_ip,
|
||||
change_port,
|
||||
})
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
tids.push(tid as u128);
|
||||
tracing::trace!(b = ?msg.as_bytes(), tid, "send stun request");
|
||||
udp.send_to(msg.as_bytes(), &stun_host).await?;
|
||||
}
|
||||
|
||||
tracing::trace!("waiting stun response");
|
||||
let mut buf = [0; 1620];
|
||||
let (msg, recv_addr) = self.wait_stun_response(&mut buf, &udp, &tids).await?;
|
||||
|
||||
let changed_socket_addr = Self::extract_changed_addr(&msg);
|
||||
let ip_changed = stun_host.ip() != recv_addr.ip();
|
||||
let port_changed = stun_host.port() != recv_addr.port();
|
||||
|
||||
let resp = BindRequestResponse {
|
||||
source_addr: udp.local_addr()?,
|
||||
mapped_socket_addr: Self::extrace_mapped_addr(&msg),
|
||||
changed_socket_addr,
|
||||
ip_changed,
|
||||
port_changed,
|
||||
};
|
||||
|
||||
tracing::info!(
|
||||
?stun_host,
|
||||
?recv_addr,
|
||||
?changed_socket_addr,
|
||||
"finish stun bind request"
|
||||
);
|
||||
|
||||
Ok(resp)
|
||||
}
|
||||
}
|
||||
|
||||
struct UdpNatTypeDetector {
|
||||
stun_servers: Vec<String>,
|
||||
}
|
||||
|
||||
impl UdpNatTypeDetector {
|
||||
pub fn new(stun_servers: Vec<String>) -> Self {
|
||||
Self { stun_servers }
|
||||
}
|
||||
|
||||
async fn get_udp_nat_type(&self, mut source_port: u16) -> NatType {
|
||||
// Like classic STUN (rfc3489). Detect NAT behavior for UDP.
|
||||
// Modified from rfc3489. Requires at least two STUN servers.
|
||||
let mut ret_test1_1 = None;
|
||||
let mut ret_test1_2 = None;
|
||||
let mut ret_test2 = None;
|
||||
let mut ret_test3 = None;
|
||||
|
||||
if source_port == 0 {
|
||||
let udp = UdpSocket::bind("0.0.0.0:0").await.unwrap();
|
||||
source_port = udp.local_addr().unwrap().port();
|
||||
}
|
||||
|
||||
let mut succ = false;
|
||||
for server_ip in &self.stun_servers {
|
||||
let stun = Stun::new(server_ip.clone());
|
||||
let ret = stun.bind_request(source_port, false, false).await;
|
||||
if ret.is_err() {
|
||||
// Try another STUN server
|
||||
continue;
|
||||
}
|
||||
if ret_test1_1.is_none() {
|
||||
ret_test1_1 = ret.ok();
|
||||
continue;
|
||||
}
|
||||
ret_test1_2 = ret.ok();
|
||||
let ret = stun.bind_request(source_port, true, true).await;
|
||||
if let Ok(resp) = ret {
|
||||
if !resp.ip_changed || !resp.port_changed {
|
||||
// Try another STUN server
|
||||
continue;
|
||||
}
|
||||
}
|
||||
ret_test2 = ret.ok();
|
||||
ret_test3 = stun.bind_request(source_port, false, true).await.ok();
|
||||
succ = true;
|
||||
break;
|
||||
}
|
||||
|
||||
if !succ {
|
||||
return NatType::Unknown;
|
||||
}
|
||||
|
||||
tracing::info!(
|
||||
?ret_test1_1,
|
||||
?ret_test1_2,
|
||||
?ret_test2,
|
||||
?ret_test3,
|
||||
"finish stun test, try to detect nat type"
|
||||
);
|
||||
|
||||
let ret_test1_1 = ret_test1_1.unwrap();
|
||||
let ret_test1_2 = ret_test1_2.unwrap();
|
||||
|
||||
if ret_test1_1.mapped_socket_addr != ret_test1_2.mapped_socket_addr {
|
||||
return NatType::Symmetric;
|
||||
}
|
||||
|
||||
if ret_test1_1.mapped_socket_addr.is_some()
|
||||
&& ret_test1_1.source_addr == ret_test1_1.mapped_socket_addr.unwrap()
|
||||
{
|
||||
if !ret_test2.is_none() {
|
||||
return NatType::OpenInternet;
|
||||
} else {
|
||||
return NatType::SymUdpFirewall;
|
||||
}
|
||||
} else {
|
||||
if let Some(ret_test2) = ret_test2 {
|
||||
if source_port == ret_test2.get_mapped_addr_no_check().port()
|
||||
&& source_port == ret_test1_1.get_mapped_addr_no_check().port()
|
||||
{
|
||||
return NatType::NoPat;
|
||||
} else {
|
||||
return NatType::FullCone;
|
||||
}
|
||||
} else {
|
||||
if !ret_test3.is_none() {
|
||||
return NatType::Restricted;
|
||||
} else {
|
||||
return NatType::PortRestricted;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
#[auto_impl::auto_impl(&, Arc, Box)]
|
||||
pub trait StunInfoCollectorTrait: Send + Sync {
|
||||
fn get_stun_info(&self) -> StunInfo;
|
||||
async fn get_udp_port_mapping(&self, local_port: u16) -> Result<SocketAddr, Error>;
|
||||
}
|
||||
|
||||
pub struct StunInfoCollector {
|
||||
stun_servers: Arc<RwLock<Vec<String>>>,
|
||||
udp_nat_type: Arc<AtomicCell<(NatType, std::time::Instant)>>,
|
||||
redetect_notify: Arc<tokio::sync::Notify>,
|
||||
tasks: JoinSet<()>,
|
||||
}
|
||||
|
||||
#[async_trait::async_trait]
|
||||
impl StunInfoCollectorTrait for StunInfoCollector {
|
||||
fn get_stun_info(&self) -> StunInfo {
|
||||
let (typ, time) = self.udp_nat_type.load();
|
||||
StunInfo {
|
||||
udp_nat_type: typ as i32,
|
||||
tcp_nat_type: 0,
|
||||
last_update_time: time.elapsed().as_secs() as i64,
|
||||
}
|
||||
}
|
||||
|
||||
async fn get_udp_port_mapping(&self, local_port: u16) -> Result<SocketAddr, Error> {
|
||||
let stun_servers = self.stun_servers.read().await.clone();
|
||||
for server in stun_servers.iter() {
|
||||
let stun = Stun::new(server.clone());
|
||||
let Ok(ret) = stun.bind_request(local_port, false, false).await else {
|
||||
tracing::warn!(?server, "stun bind request failed");
|
||||
continue;
|
||||
};
|
||||
if let Some(mapped_addr) = ret.mapped_socket_addr {
|
||||
return Ok(mapped_addr);
|
||||
}
|
||||
}
|
||||
Err(Error::NotFound)
|
||||
}
|
||||
}
|
||||
|
||||
impl StunInfoCollector {
|
||||
pub fn new(stun_servers: Vec<String>) -> Self {
|
||||
let mut ret = Self {
|
||||
stun_servers: Arc::new(RwLock::new(stun_servers)),
|
||||
udp_nat_type: Arc::new(AtomicCell::new((
|
||||
NatType::Unknown,
|
||||
std::time::Instant::now(),
|
||||
))),
|
||||
redetect_notify: Arc::new(tokio::sync::Notify::new()),
|
||||
tasks: JoinSet::new(),
|
||||
};
|
||||
|
||||
ret.start_stun_routine();
|
||||
|
||||
ret
|
||||
}
|
||||
|
||||
fn start_stun_routine(&mut self) {
|
||||
let stun_servers = self.stun_servers.clone();
|
||||
let udp_nat_type = self.udp_nat_type.clone();
|
||||
let redetect_notify = self.redetect_notify.clone();
|
||||
self.tasks.spawn(async move {
|
||||
loop {
|
||||
let detector = UdpNatTypeDetector::new(stun_servers.read().await.clone());
|
||||
let ret = detector.get_udp_nat_type(0).await;
|
||||
udp_nat_type.store((ret, std::time::Instant::now()));
|
||||
|
||||
let sleep_sec = match ret {
|
||||
NatType::Unknown => 15,
|
||||
_ => 60,
|
||||
};
|
||||
tracing::info!(?ret, ?sleep_sec, "finish udp nat type detect");
|
||||
|
||||
tokio::select! {
|
||||
_ = redetect_notify.notified() => {}
|
||||
_ = tokio::time::sleep(Duration::from_secs(sleep_sec)) => {}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
pub fn update_stun_info(&self) {
|
||||
self.redetect_notify.notify_one();
|
||||
}
|
||||
|
||||
pub async fn set_stun_servers(&self, stun_servers: Vec<String>) {
|
||||
*self.stun_servers.write().await = stun_servers;
|
||||
self.update_stun_info();
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_stun_bind_request() {
|
||||
// miwifi / qq seems not correctly responde to change_ip and change_port, they always try to change the src ip and port.
|
||||
// let stun = Stun::new("stun.counterpath.com:3478".to_string());
|
||||
let stun = Stun::new("180.235.108.91:3478".to_string());
|
||||
// let stun = Stun::new("193.22.2.248:3478".to_string());
|
||||
|
||||
// let stun = Stun::new("stun.chat.bilibili.com:3478".to_string());
|
||||
// let stun = Stun::new("stun.miwifi.com:3478".to_string());
|
||||
|
||||
let rs = stun.bind_request(12345, true, true).await.unwrap();
|
||||
assert!(rs.ip_changed);
|
||||
assert!(rs.port_changed);
|
||||
|
||||
let rs = stun.bind_request(12345, true, false).await.unwrap();
|
||||
assert!(rs.ip_changed);
|
||||
assert!(!rs.port_changed);
|
||||
|
||||
let rs = stun.bind_request(12345, false, true).await.unwrap();
|
||||
assert!(!rs.ip_changed);
|
||||
assert!(rs.port_changed);
|
||||
|
||||
let rs = stun.bind_request(12345, false, false).await.unwrap();
|
||||
assert!(!rs.ip_changed);
|
||||
assert!(!rs.port_changed);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_udp_nat_type_detect() {
|
||||
let detector = UdpNatTypeDetector::new(vec![
|
||||
"stun.counterpath.com:3478".to_string(),
|
||||
"180.235.108.91:3478".to_string(),
|
||||
]);
|
||||
let ret = detector.get_udp_nat_type(0).await;
|
||||
|
||||
assert_eq!(ret, NatType::FullCone);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user