Add support for PTR query

mokeyish · Jul 12, 2024 · 51b1213 · 51b1213
1 parent a98144a
commit 51b1213
Showing 1 changed file with 85 additions and 8 deletions.
diff --git a/crates/resolver/src/hosts.rs b/crates/resolver/src/hosts.rs
@@ -2,11 +2,13 @@
 
 use std::collections::HashMap;
 use std::io;
+use std::net::IpAddr;
 use std::path::Path;
 use std::str::FromStr;
 use std::sync::Arc;
 
 use proto::op::Query;
+use proto::rr::rdata::PTR;
 use proto::rr::{Name, RecordType};
 use proto::rr::{RData, Record};
 use tracing::warn;
@@ -47,14 +49,62 @@ impl Hosts {
     /// Look up the addresses for the given host from the system hosts file.
     pub fn lookup_static_host(&self, query: &Query) -> Option<Lookup> {
         if !self.by_name.is_empty() {
-            if let Some(val) = self.by_name.get(query.name()) {
-                let result = match query.query_type() {
-                    RecordType::A => val.a.clone(),
-                    RecordType::AAAA => val.aaaa.clone(),
-                    _ => None,
-                };
-
-                return result;
+            match query.query_type() {
+                RecordType::A | RecordType::AAAA => {
+                    if let Some(val) = self.by_name.get(query.name()) {
+                        let result = match query.query_type() {
+                            RecordType::A => val.a.clone(),
+                            RecordType::AAAA => val.aaaa.clone(),
+                            _ => None,
+                        };
+
+                        return result;
+                    }
+                }
+                RecordType::PTR => {
+                    if let Ok(ip) = query.name().parse_arpa_name() {
+                        let ip = ip.addr();
+                        let mut hostnames = self
+                            .by_name
+                            .iter()
+                            .filter(|(_, v)| match ip {
+                                IpAddr::V4(ip) => match v.a.as_ref() {
+                                    Some(lookup) => lookup.iter().any(|r| {
+                                        r.ip_addr().map(|it| it == ip).unwrap_or_default()
+                                    }),
+                                    None => false,
+                                },
+                                IpAddr::V6(ip) => match v.aaaa.as_ref() {
+                                    Some(lookup) => lookup.iter().any(|r| {
+                                        r.ip_addr().map(|it| it == ip).unwrap_or_default()
+                                    }),
+                                    None => false,
+                                },
+                            })
+                            .map(|(n, _)| n)
+                            .collect::<Vec<_>>();
+
+                        if !hostnames.is_empty() {
+                            hostnames.sort();
+                            let records = hostnames
+                                .into_iter()
+                                .map(|hostname| {
+                                    Record::from_rdata(
+                                        query.name().clone(),
+                                        dns_lru::MAX_TTL,
+                                        RData::PTR(PTR(hostname.clone())),
+                                    )
+                                })
+                                .collect::<Vec<_>>();
+
+                            return Some(Lookup::new_with_max_ttl(
+                                query.clone(),
+                                Arc::from(records),
+                            ));
+                        }
+                    }
+                }
+                _ => (),
             }
         }
         None
@@ -254,5 +304,32 @@ mod tests {
             .map(ToOwned::to_owned)
             .collect::<Vec<RData>>();
         assert_eq!(rdatas, vec![RData::A(Ipv4Addr::new(10, 0, 1, 111).into())]);
+
+        let name = Name::from_str("111.1.0.10.in-addr.arpa.").unwrap();
+        let rdatas = hosts
+            .lookup_static_host(&Query::query(name, RecordType::PTR))
+            .unwrap()
+            .iter()
+            .map(ToOwned::to_owned)
+            .collect::<Vec<RData>>();
+        assert_eq!(
+            rdatas,
+            vec![
+                RData::PTR(PTR("a.example.com".parse().unwrap())),
+                RData::PTR(PTR("b.example.com".parse().unwrap()))
+            ]
+        );
+
+        let name = Name::from_str(
+            "1.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.0.ip6.arpa.",
+        )
+        .unwrap();
+        let rdatas = hosts
+            .lookup_static_host(&Query::query(name, RecordType::PTR))
+            .unwrap()
+            .iter()
+            .map(ToOwned::to_owned)
+            .collect::<Vec<RData>>();
+        assert_eq!(rdatas, vec![RData::PTR(PTR("localhost".parse().unwrap())),]);
     }
 }