From b6c8a7441feb2cd6550d2d97ab89ffd569724d2c Mon Sep 17 00:00:00 2001
From: Zhang Jingqiang <zhangjingqiang@bytedance.com>
Date: Mon, 11 Mar 2024 10:40:49 +0800
Subject: [PATCH] g3-resolver: use returned answers before try cname

---
 demo/test-resolver/src/main.rs                |  5 +--
 .../src/escape/direct_fixed/udp_relay/send.rs |  9 ++---
 lib/g3-resolver/src/driver/hickory/client.rs  | 34 ++++++++++++-------
 3 files changed, 30 insertions(+), 18 deletions(-)

diff --git a/demo/test-resolver/src/main.rs b/demo/test-resolver/src/main.rs
index 76098cda..f14d6646 100644
--- a/demo/test-resolver/src/main.rs
+++ b/demo/test-resolver/src/main.rs
@@ -16,6 +16,7 @@
 
 use std::net::IpAddr;
 use std::str::FromStr;
+use std::sync::Arc;
 
 use log::info;
 use slog::{slog_o, Drain};
@@ -54,10 +55,10 @@ fn main() {
         };
         let resolver = ResolverBuilder::new(config).build().unwrap();
         let handle = resolver.get_handle();
-        let mut job = handle.get_v4("www.xjtu.edu.cn".to_string()).unwrap();
+        let mut job = handle.get_v4(Arc::from("www.xjtu.edu.cn")).unwrap();
         let data = job.recv().await;
         info!("v4 data: {:?}", data);
-        let mut job = handle.get_v6("www.xjtu.edu.cn".to_string()).unwrap();
+        let mut job = handle.get_v6(Arc::from("www.xjtu.edu.cn")).unwrap();
         let data = job.recv().await;
         info!("v6 data: {:?}", data);
     });
diff --git a/g3proxy/src/escape/direct_fixed/udp_relay/send.rs b/g3proxy/src/escape/direct_fixed/udp_relay/send.rs
index 4b10e779..6744f169 100644
--- a/g3proxy/src/escape/direct_fixed/udp_relay/send.rs
+++ b/g3proxy/src/escape/direct_fixed/udp_relay/send.rs
@@ -49,7 +49,7 @@ pub(crate) struct DirectUdpRelayRemoteSend<T> {
     resolver_handle: ArcIntegratedResolverHandle,
     resolve_strategy: ResolveStrategy,
     resolver_job: Option<ArriveFirstResolveJob>,
-    resolve_retry_domain: Option<String>,
+    resolve_retry_domain: Option<Arc<str>>,
     resolved_port: u16,
     resolved_ip: Option<IpAddr>,
 }
@@ -136,7 +136,7 @@ where
                                     let resolver_job = ArriveFirstResolveJob::new(
                                         &self.resolver_handle,
                                         self.resolve_strategy,
-                                        Arc::from(domain),
+                                        domain,
                                     )?;
                                     self.resolver_job = Some(resolver_job);
                                     // no retry by leaving resolve_retry_domain to None
@@ -165,13 +165,14 @@ where
             }
             Host::Domain(domain) => {
                 self.resolved_port = to.port();
+                let domain: Arc<str> = Arc::from(domain.as_str());
                 let resolver_job = ArriveFirstResolveJob::new(
                     &self.resolver_handle,
                     self.resolve_strategy,
-                    Arc::from(domain.as_str()),
+                    domain.clone(),
                 )?;
                 self.resolver_job = Some(resolver_job);
-                self.resolve_retry_domain = Some(domain.to_string());
+                self.resolve_retry_domain = Some(domain);
                 self.poll_send_packet(cx, buf, to)
             }
         }
diff --git a/lib/g3-resolver/src/driver/hickory/client.rs b/lib/g3-resolver/src/driver/hickory/client.rs
index 7ecb1ab4..e11bad0f 100644
--- a/lib/g3-resolver/src/driver/hickory/client.rs
+++ b/lib/g3-resolver/src/driver/hickory/client.rs
@@ -177,30 +177,40 @@ impl HickoryClientJob {
                     let mut ttl = 0;
                     for r in msg.take_answers() {
                         ttl = r.ttl();
-                        match r.data() {
-                            Some(RData::A(v)) => {
+                        let Some(rdata) = r.data() else {
+                            continue;
+                        };
+                        match rdata {
+                            RData::A(v) => {
                                 if req.rtype == RecordType::A {
                                     ips.push(IpAddr::V4(v.0));
                                 }
                             }
-                            Some(RData::AAAA(v)) => {
+                            RData::AAAA(v) => {
                                 if req.rtype == RecordType::AAAA {
                                     ips.push(IpAddr::V6(v.0))
                                 }
                             }
-                            Some(RData::CNAME(v)) => {
-                                has_cname = true;
-                                name = v.0.clone();
+                            RData::CNAME(v) => {
+                                if name.eq(r.name()) {
+                                    has_cname = true;
+                                    name = v.0.clone();
+                                }
                             }
                             _ => {}
                         }
                     }
-                    if has_cname {
-                        self.try_truncated = true;
-                        continue;
-                    }
-                    let ttl = ttl.clamp(self.config.positive_min_ttl, self.config.positive_max_ttl);
-                    return ResolvedRecord::resolved(req.domain, ttl, ips);
+                    return if ips.is_empty() {
+                        if has_cname {
+                            self.try_truncated = true;
+                            continue;
+                        }
+                        ResolvedRecord::resolved(req.domain, self.config.negative_ttl, ips)
+                    } else {
+                        let ttl =
+                            ttl.clamp(self.config.positive_min_ttl, self.config.positive_max_ttl);
+                        ResolvedRecord::resolved(req.domain, ttl, ips)
+                    };
                 }
                 Err(e) => {
                     self.state.add_failed();