use lookup_host instead of aresolve

lazytiger · Jan 3, 2024 · bd89e07 · bd89e07
1 parent ed55059
commit bd89e07
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 51 deletions.
diff --git a/src/aserver/mod.rs b/src/aserver/mod.rs
@@ -10,7 +10,7 @@ use std::{
 use bytes::{Buf, BytesMut};
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
-    net::{TcpListener, TcpStream},
+    net::{lookup_host, TcpListener, TcpStream},
     runtime::{Handle, Runtime},
     spawn,
     sync::mpsc::{unbounded_channel, UnboundedSender},
@@ -28,7 +28,6 @@ use crate::{
     proto::{RequestParseResult, Sock5Address, TrojanRequest, CONNECT, PING, UDP_ASSOCIATE},
     server::{init_config, ping_backend::PingResult},
     types::{Result, TrojanError},
-    utils::aresolve,
 };
 
 mod ping;
@@ -109,14 +108,10 @@ async fn start_proxy_internal(
                             cmd,
                             match address {
                                 Sock5Address::Socket(addr) => addr,
-                                Sock5Address::Domain(domain, port) => {
-                                    let ip =
-                                        *aresolve(domain.as_str(), OPTIONS.system_dns.as_str())
-                                            .await?
-                                            .first()
-                                            .ok_or(TrojanError::Resolve)?;
-                                    SocketAddr::new(ip, port)
-                                }
+                                Sock5Address::Domain(domain, port) => lookup_host((domain, port))
+                                    .await?
+                                    .next()
+                                    .ok_or(TrojanError::Resolve)?,
                                 Sock5Address::None => *OPTIONS.back_addr.as_ref().unwrap(),
                                 _ => unreachable!(),
                             },

diff --git a/src/aserver/udp.rs b/src/aserver/udp.rs
@@ -3,7 +3,7 @@ use std::{collections::HashMap, io, net::SocketAddr, sync::Arc, time::Duration};
 use bytes::{Buf, BytesMut};
 use tokio::{
     io::{split, AsyncReadExt, AsyncWriteExt, WriteHalf},
-    net::{TcpStream, UdpSocket},
+    net::{lookup_host, TcpStream, UdpSocket},
     spawn,
     sync::mpsc::{channel, Receiver},
     time::{timeout, Instant},
@@ -14,7 +14,7 @@ use crate::{
     config::OPTIONS,
     proto::{Sock5Address, UdpAssociate, UdpParseResult},
     types::Result,
-    utils::{aresolve, is_private},
+    utils::is_private,
 };
 
 pub async fn start_udp(source: TlsStream<TcpStream>, mut buffer: BytesMut) -> Result<()> {
@@ -34,13 +34,14 @@ pub async fn start_udp(source: TlsStream<TcpStream>, mut buffer: BytesMut) -> Re
                         Sock5Address::Domain(domain, port) => {
                             if let Some(ip) = dns_cache_store.get(&domain) {
                                 SocketAddr::new(*ip, port)
-                            } else if let Ok(Some(ip)) =
-                                aresolve(domain.as_str(), OPTIONS.system_dns.as_str())
-                                    .await
-                                    .map(|ips| ips.get(0).copied())
-                            {
-                                dns_cache_store.insert(domain, ip);
-                                SocketAddr::new(ip, port)
+                            } else if let Ok(mut ret) = lookup_host((domain.clone(), port)).await {
+                                if let Some(addr) = ret.next() {
+                                    dns_cache_store.insert(domain, addr.ip());
+                                    addr
+                                } else {
+                                    log::error!("query {} failed", domain);
+                                    continue;
+                                }
                             } else {
                                 log::error!("query {} failed", domain);
                                 continue;

diff --git a/src/utils.rs b/src/utils.rs
@@ -8,7 +8,6 @@ use std::{
 
 use bytes::{Buf, BytesMut};
 use socket2::{Domain, Protocol, SockAddr, Socket, Type};
-use tokio::net::UdpSocket;
 use trust_dns_proto::{
     op::{Message, Query},
     rr::{DNSClass, Name, RecordType},
@@ -135,37 +134,6 @@ pub fn resolve(name: &str, dns_server_addr: &str) -> Result<Vec<IpAddr>> {
     }
 }
 
-pub async fn aresolve(name: &str, dns_server_addr: &str) -> Result<Vec<IpAddr>> {
-    let dns_server_addr: SocketAddr = dns_server_addr.parse()?;
-    let socket = UdpSocket::bind("0.0.0.0:0").await?;
-    let mut message = Message::new();
-    message.set_recursion_desired(true);
-    message.set_id(1);
-    let mut query = Query::new();
-    let name = Name::from_str(name)?;
-    query.set_name(name);
-    query.set_query_type(RecordType::A);
-    query.set_query_class(DNSClass::IN);
-    message.add_query(query);
-    let request = message.to_vec()?;
-    if request.len() != socket.send_to(request.as_slice(), &dns_server_addr).await? {
-        return Err(TrojanError::Dummy(()));
-    }
-    let mut response = vec![0u8; 1024];
-    let length = tokio::time::timeout(Duration::from_secs(3), socket.recv(response.as_mut_slice()))
-        .await??;
-    let message = Message::from_bytes(&response.as_slice()[..length])?;
-    if message.id() != 1 {
-        Err(TrojanError::Dummy(()))
-    } else {
-        Ok(message
-            .answers()
-            .iter()
-            .filter_map(|record| record.data().and_then(|data| data.ip_addr()))
-            .collect())
-    }
-}
-
 pub fn get_system_dns() -> types::Result<String> {
     let file = File::open("/etc/resolv.conf")?;
     let mut data = String::new();