From 60b9683facd10a6f7309420252b23d7a600b8bd0 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Sun, 20 Aug 2023 12:13:28 +0800
Subject: [PATCH 1/2] dns query from remote server

---
 src/dns.rs       | 19 ++++++++++++++++++-
 src/tun2proxy.rs | 14 +++++++++++---
 2 files changed, 29 insertions(+), 4 deletions(-)

diff --git a/src/dns.rs b/src/dns.rs
index 2ae1528..8aefd03 100644
--- a/src/dns.rs
+++ b/src/dns.rs
@@ -1,6 +1,9 @@
 #![allow(dead_code)]
 
-use std::{net::IpAddr, str::FromStr};
+use std::{
+    net::{IpAddr, Ipv4Addr, SocketAddr},
+    str::FromStr,
+};
 use trust_dns_proto::{
     op::{Message, ResponseCode},
     rr::{record_type::RecordType, Name, RData, Record},
@@ -90,3 +93,17 @@ pub fn parse_data_to_dns_message(data: &[u8], used_by_tcp: bool) -> Result<Messa
     let message = Message::from_vec(data).map_err(|e| e.to_string())?;
     Ok(message)
 }
+
+// FIXME: use IpAddr::is_global() instead when it's stable
+pub fn addr_is_private(addr: &SocketAddr) -> bool {
+    fn is_benchmarking(addr: &Ipv4Addr) -> bool {
+        addr.octets()[0] == 198 && (addr.octets()[1] & 0xfe) == 18
+    }
+    fn addr_v4_is_private(addr: &Ipv4Addr) -> bool {
+        is_benchmarking(addr) || addr.is_private() || addr.is_loopback() || addr.is_link_local()
+    }
+    match addr {
+        SocketAddr::V4(addr) => addr_v4_is_private(addr.ip()),
+        SocketAddr::V6(_) => false,
+    }
+}
diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs
index a537a3b..f4c0854 100644
--- a/src/tun2proxy.rs
+++ b/src/tun2proxy.rs
@@ -1,4 +1,4 @@
-use crate::{error::Error, error::Result, virtdevice::VirtualTunDevice, NetworkInterface, Options};
+use crate::{dns, error::Error, error::Result, virtdevice::VirtualTunDevice, NetworkInterface, Options};
 use mio::{event::Event, net::TcpStream, net::UdpSocket, unix::SourceFd, Events, Interest, Poll, Token};
 use smoltcp::{
     iface::{Config, Interface, SocketHandle, SocketSet},
@@ -468,7 +468,15 @@ impl<'a> TunToProxy<'a> {
             let (info, _first_packet, payload_offset, payload_size) = result?;
             let origin_dst = SocketAddr::try_from(&info.dst)?;
             let connection_info = match &mut self.options.virtual_dns {
-                None => info,
+                None => {
+                    let mut info = info;
+                    let port = origin_dst.port();
+                    if port == 53 && info.protocol == IpProtocol::Udp && dns::addr_is_private(&origin_dst) {
+                        let dns_addr: SocketAddr = "8.8.8.8:53".parse()?; // TODO: Configurable
+                        info.dst = Address::from(dns_addr);
+                    }
+                    info
+                }
                 Some(virtual_dns) => {
                     let dst_ip = origin_dst.ip();
                     virtual_dns.touch_ip(&dst_ip);
@@ -798,7 +806,7 @@ impl<'a> TunToProxy<'a> {
                         Ok(read_result) => read_result,
                         Err(error) => {
                             if error.kind() != std::io::ErrorKind::WouldBlock {
-                                log::error!("Read from proxy: {}", error);
+                                log::error!("{} Read from proxy: {}", conn_info.dst, error);
                             }
                             vecbuf.len()
                         }

From 6439cc7b4314f119f6633fcc3eb02d44b2878440 Mon Sep 17 00:00:00 2001
From: ssrlive <30760636+ssrlive@users.noreply.github.com>
Date: Sun, 20 Aug 2023 13:29:36 +0800
Subject: [PATCH 2/2] dns::remove_ipv6_entries

---
 src/dns.rs       |  6 ++++++
 src/tun2proxy.rs | 10 +++++++++-
 2 files changed, 15 insertions(+), 1 deletion(-)

diff --git a/src/dns.rs b/src/dns.rs
index 8aefd03..33be70c 100644
--- a/src/dns.rs
+++ b/src/dns.rs
@@ -50,6 +50,12 @@ pub fn build_dns_response(mut request: Message, domain: &str, ip: IpAddr, ttl: u
     Ok(request)
 }
 
+pub fn remove_ipv6_entries(message: &mut Message) {
+    message
+        .answers_mut()
+        .retain(|answer| !matches!(answer.data(), Some(RData::AAAA(_))));
+}
+
 pub fn extract_ipaddr_from_dns_message(message: &Message) -> Result<IpAddr, String> {
     if message.response_code() != ResponseCode::NoError {
         return Err(format!("{:?}", message.response_code()));
diff --git a/src/tun2proxy.rs b/src/tun2proxy.rs
index f4c0854..d68abdb 100644
--- a/src/tun2proxy.rs
+++ b/src/tun2proxy.rs
@@ -771,9 +771,17 @@ impl<'a> TunToProxy<'a> {
                 let buf = buf[..packet_size].to_vec();
                 let header = UdpHeader::retrieve_from_stream(&mut &buf[..])?;
 
+                let buf = if info.dst.port() == 53 {
+                    let mut message = dns::parse_data_to_dns_message(&buf[header.len()..], false)?;
+                    dns::remove_ipv6_entries(&mut message); // TODO: Configurable
+                    message.to_vec()?
+                } else {
+                    buf[header.len()..].to_vec()
+                };
+
                 // Write to client
                 let src = state.udp_origin_dst.ok_or("udp address")?;
-                self.send_udp_packet_to_client(src, info.src, &buf[header.len()..])?;
+                self.send_udp_packet_to_client(src, info.src, &buf)?;
             }
 
             return Ok(());