314 lines
10 KiB
Rust
314 lines
10 KiB
Rust
use anyhow::{Context, anyhow};
|
||
use dns_parser::{Builder, Packet, QueryClass, QueryType, RData, ResponseCode};
|
||
use rand::seq::SliceRandom;
|
||
use std::io;
|
||
use std::net::{Ipv4Addr, Ipv6Addr, SocketAddr};
|
||
use std::str::FromStr;
|
||
use std::time::Duration;
|
||
use tokio::net::UdpSocket;
|
||
|
||
use rust_p2p_core::socket::LocalInterface;
|
||
|
||
pub async fn dns_query_txt(
|
||
domain: &str,
|
||
mut name_servers: Vec<String>,
|
||
default_interface: &Option<LocalInterface>,
|
||
) -> io::Result<Vec<String>> {
|
||
let mut err: Option<io::Error> = None;
|
||
if name_servers.is_empty() {
|
||
name_servers.push("223.5.5.5:53".into());
|
||
name_servers.push("114.114.114.114:53".into());
|
||
}
|
||
for name_server in name_servers {
|
||
match txt_dns(domain, name_server, default_interface).await {
|
||
Ok(addr) => {
|
||
if !addr.is_empty() {
|
||
return Ok(addr);
|
||
}
|
||
}
|
||
Err(e) => {
|
||
err.replace(e);
|
||
}
|
||
}
|
||
continue;
|
||
}
|
||
if let Some(e) = err {
|
||
Err(e)
|
||
} else {
|
||
Err(io::Error::other(format!("DNS query failed {domain:?}")))
|
||
}
|
||
}
|
||
pub async fn dns_query_one(
|
||
domain: &str,
|
||
name_servers: &Vec<String>,
|
||
default_interface: &Option<LocalInterface>,
|
||
) -> anyhow::Result<SocketAddr> {
|
||
let mut vec = dns_query_all(domain, name_servers, default_interface).await?;
|
||
if default_interface.is_some() {
|
||
// 出口网卡绑定目前应用于 IPv4 Socket；优先且只使用可绑定的 IPv4 地址。
|
||
vec.retain(SocketAddr::is_ipv4);
|
||
}
|
||
vec.shuffle(&mut rand::rng());
|
||
vec.pop().context("DNS query failed")
|
||
}
|
||
pub async fn dns_query_all(
|
||
domain: &str,
|
||
name_servers: &Vec<String>,
|
||
default_interface: &Option<LocalInterface>,
|
||
) -> anyhow::Result<Vec<SocketAddr>> {
|
||
match SocketAddr::from_str(domain) {
|
||
Ok(addr) => Ok(vec![addr]),
|
||
Err(_) => {
|
||
if name_servers.is_empty() {
|
||
let addrs: Vec<SocketAddr> = tokio::net::lookup_host(domain)
|
||
.await
|
||
.map_err(|e| io::Error::other(format!("DNS query failed: {domain:?},{e:?}")))?
|
||
.collect();
|
||
return Ok(addrs);
|
||
}
|
||
|
||
let mut err: Option<io::Error> = None;
|
||
for name_server in name_servers {
|
||
let end_index = domain
|
||
.rfind(':')
|
||
.ok_or_else(|| io::Error::other(format!("not port: {domain:?}")))?;
|
||
let host = &domain[..end_index];
|
||
let port = u16::from_str(&domain[end_index + 1..])
|
||
.map_err(|_| io::Error::other(format!("not port: {domain:?}")))?;
|
||
let th1 = {
|
||
let host = host.to_string();
|
||
let name_server = name_server.clone();
|
||
let default_interface = default_interface.clone();
|
||
tokio::spawn(a_dns(host, name_server, default_interface.clone()))
|
||
};
|
||
let th2 = {
|
||
let host = host.to_string();
|
||
let name_server = name_server.clone();
|
||
let default_interface = default_interface.clone();
|
||
tokio::spawn(aaaa_dns(host, name_server, default_interface.clone()))
|
||
};
|
||
let mut addr = Vec::new();
|
||
match th1.await? {
|
||
Ok(rs) => {
|
||
for ip in rs {
|
||
addr.push(SocketAddr::new(ip.into(), port));
|
||
}
|
||
}
|
||
Err(e) => {
|
||
err.replace(e);
|
||
}
|
||
}
|
||
match th2.await? {
|
||
Ok(rs) => {
|
||
for ip in rs {
|
||
addr.push(SocketAddr::new(ip.into(), port));
|
||
}
|
||
}
|
||
Err(e) => {
|
||
if addr.is_empty() {
|
||
err.replace(e);
|
||
continue;
|
||
}
|
||
}
|
||
}
|
||
if addr.is_empty() {
|
||
continue;
|
||
}
|
||
return Ok(addr);
|
||
}
|
||
if let Some(e) = err {
|
||
Err(e.into())
|
||
} else {
|
||
Err(anyhow!("DNS query failed {domain:?}"))
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 校验域名格式：label 非空且不超过 63 字节，全长不超过 253 字节。
|
||
/// dns-parser 的 add_question 对非法 label 直接 assert panic，
|
||
/// 必须在调用前拦截
|
||
fn is_valid_domain(domain: &str) -> bool {
|
||
let domain = domain.strip_suffix('.').unwrap_or(domain);
|
||
!domain.is_empty()
|
||
&& domain.len() <= 253
|
||
&& domain
|
||
.split('.')
|
||
.all(|label| !label.is_empty() && label.len() <= 63)
|
||
}
|
||
|
||
async fn query<'a>(
|
||
udp: &UdpSocket,
|
||
domain: &str,
|
||
name_server: SocketAddr,
|
||
record_type: QueryType,
|
||
buf: &'a mut [u8],
|
||
) -> io::Result<Packet<'a>> {
|
||
if !is_valid_domain(domain) {
|
||
return Err(io::Error::new(
|
||
io::ErrorKind::InvalidInput,
|
||
format!("invalid domain {domain:?}"),
|
||
));
|
||
}
|
||
let mut builder = Builder::new_query(1, true);
|
||
builder.add_question(domain, false, record_type, QueryClass::IN);
|
||
// 非法域名（如 label 超长）build 会失败，不能 unwrap panic
|
||
let packet = builder.build().map_err(|e| {
|
||
io::Error::new(
|
||
io::ErrorKind::InvalidInput,
|
||
format!("invalid domain {domain:?}: {e:?}"),
|
||
)
|
||
})?;
|
||
|
||
udp.connect(name_server).await?;
|
||
let mut count = 0;
|
||
let len = loop {
|
||
udp.send(&packet).await?;
|
||
|
||
match tokio::time::timeout(Duration::from_secs(3), udp.recv(buf)).await {
|
||
Ok(len) => {
|
||
break len?;
|
||
}
|
||
Err(_) => {
|
||
count += 1;
|
||
if count < 3 {
|
||
continue;
|
||
}
|
||
Err(io::Error::other(format!("DNS {name_server:?} recv error ")))?
|
||
}
|
||
};
|
||
};
|
||
|
||
let pkt = Packet::parse(&buf[..len]).map_err(|e| {
|
||
io::Error::other(format!(
|
||
"domain {domain:?} DNS {name_server:?} data error: {e}"
|
||
))
|
||
})?;
|
||
if pkt.header.response_code != ResponseCode::NoError {
|
||
return Err(io::Error::other(format!(
|
||
"response_code {} DNS {:?} domain {:?}",
|
||
pkt.header.response_code, name_server, domain
|
||
)));
|
||
}
|
||
if pkt.answers.is_empty() {
|
||
return Err(io::Error::other(format!(
|
||
"No records received DNS {name_server:?} domain {domain:?}"
|
||
)));
|
||
}
|
||
|
||
Ok(pkt)
|
||
}
|
||
|
||
pub async fn txt_dns(
|
||
domain: &str,
|
||
name_server: String,
|
||
default_interface: &Option<LocalInterface>,
|
||
) -> io::Result<Vec<String>> {
|
||
let name_server: SocketAddr = name_server
|
||
.parse()
|
||
.map_err(|e| io::Error::other(format!("dns {name_server} is error :{e:?}")))?;
|
||
let udp = bind_udp(name_server, default_interface)?;
|
||
let mut buf = vec![0u8; 65536];
|
||
let message = query(&udp, domain, name_server, QueryType::TXT, &mut buf).await?;
|
||
let mut rs = Vec::new();
|
||
for record in message.answers {
|
||
if let RData::TXT(txt) = record.data {
|
||
for x in txt.iter() {
|
||
let txt = std::str::from_utf8(x)
|
||
.map_err(|_| io::Error::other("record type txt is not string"))?;
|
||
rs.push(txt.to_string());
|
||
}
|
||
}
|
||
}
|
||
Ok(rs)
|
||
}
|
||
|
||
fn bind_udp(
|
||
name_server: SocketAddr,
|
||
default_interface: &Option<LocalInterface>,
|
||
) -> io::Result<UdpSocket> {
|
||
let addr: SocketAddr = if name_server.is_ipv4() {
|
||
SocketAddr::from(([0, 0, 0, 0], 0))
|
||
} else {
|
||
SocketAddr::from(([0; 8], 0))
|
||
};
|
||
let socket = rust_p2p_core::socket::bind_udp(addr, default_interface.as_ref())?;
|
||
UdpSocket::from_std(socket.into())
|
||
}
|
||
|
||
pub async fn a_dns(
|
||
domain: String,
|
||
name_server: String,
|
||
default_interface: Option<LocalInterface>,
|
||
) -> io::Result<Vec<Ipv4Addr>> {
|
||
let name_server: SocketAddr = name_server
|
||
.parse()
|
||
.map_err(|e| io::Error::other(format!("dns {name_server} is error :{e:?}")))?;
|
||
let udp = bind_udp(name_server, &default_interface)?;
|
||
let mut buf = vec![0u8; 65536];
|
||
let message = query(&udp, &domain, name_server, QueryType::A, &mut buf).await?;
|
||
let mut rs = Vec::new();
|
||
for record in message.answers {
|
||
if let RData::A(a) = record.data {
|
||
rs.push(a.0);
|
||
}
|
||
}
|
||
Ok(rs)
|
||
}
|
||
|
||
pub async fn aaaa_dns(
|
||
domain: String,
|
||
name_server: String,
|
||
default_interface: Option<LocalInterface>,
|
||
) -> io::Result<Vec<Ipv6Addr>> {
|
||
let name_server: SocketAddr = name_server
|
||
.parse()
|
||
.map_err(|e| io::Error::other(format!("dns {name_server} is error :{e:?}")))?;
|
||
let udp = bind_udp(name_server, &default_interface)?;
|
||
let mut buf = vec![0u8; 65536];
|
||
let message = query(&udp, &domain, name_server, QueryType::AAAA, &mut buf).await?;
|
||
let mut rs = Vec::new();
|
||
for record in message.answers {
|
||
if let RData::AAAA(a) = record.data {
|
||
rs.push(a.0);
|
||
}
|
||
}
|
||
Ok(rs)
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
/// 非法域名（label 超过 63 字节）必须返回错误而不是 panic
|
||
#[tokio::test]
|
||
async fn test_query_invalid_domain_no_panic() {
|
||
let udp = UdpSocket::bind("0.0.0.0:0").await.unwrap();
|
||
let mut buf = vec![0u8; 512];
|
||
let bad_domain = format!("{}.com", "a".repeat(64));
|
||
let rs = query(
|
||
&udp,
|
||
&bad_domain,
|
||
"127.0.0.1:53".parse().unwrap(),
|
||
QueryType::A,
|
||
&mut buf,
|
||
)
|
||
.await;
|
||
let err = rs.expect_err("invalid domain must be rejected");
|
||
assert_eq!(err.kind(), io::ErrorKind::InvalidInput);
|
||
}
|
||
|
||
#[test]
|
||
fn test_is_valid_domain() {
|
||
assert!(is_valid_domain("example.com"));
|
||
assert!(is_valid_domain("a-b_1.example.com"));
|
||
assert!(is_valid_domain("example.com.")); // FQDN 尾点合法
|
||
assert!(is_valid_domain(&format!("{}.com", "a".repeat(63))));
|
||
|
||
assert!(!is_valid_domain(""));
|
||
assert!(!is_valid_domain(&format!("{}.com", "a".repeat(64))));
|
||
assert!(!is_valid_domain("a..b"));
|
||
assert!(!is_valid_domain(&"a".repeat(254)));
|
||
}
|
||
}
|