cdns_rs/a_sync/

query.rs

/*-
 * cdns-rs - a simple sync/async DNS query library
 * 
 * Copyright (C) 2020  Aleksandr Morozov
 * 
 * Copyright 2025 Aleksandr Morozov
 * 
 * Licensed under the EUPL, Version 1.2 or - as soon they will be approved by
 * the European Commission - subsequent versions of the EUPL (the "Licence").
 * 
 * You may not use this work except in compliance with the Licence.
 * 
 * You may obtain a copy of the Licence at:
 * 
 *    https://joinup.ec.europa.eu/software/page/eupl
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the Licence is distributed on an "AS IS" basis, WITHOUT
 * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
 * Licence for the specific language governing permissions and limitations
 * under the Licence.
 */


use std::collections::{BTreeSet, HashMap};
use std::convert::TryFrom;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;

use async_recursion::async_recursion;
use tokio::net::{TcpStream, UdpSocket};

use crate::cfg_resolv_parser::{ConfigEntryTls, ResolveConfEntry, ResolveConfigFamily};
use crate::common::{CDdnsGlobals, DnsRequestAnswer, DnsRequestHeader};
use crate::{error::*, DnsResponsePayload, QDnsQueryResult};
use crate::query::QDnsQuery;
use crate::{write_error, internal_error};
use crate::query_private::QDnsReq;

use super::caches::CACHE;
use super::network::{NetworkTap, NetworkTapType, SocketTap, TcpTlsConnection};
use super::{QDnsName, QType, QuerySetup, ResolveConfig};

/// A main instance which contains all common logic.
#[derive(Clone, Debug)]
pub struct QDns
{
    /// An instance of the parser /etc/resolv.conf or custom config
    resolvers: Arc<ResolveConfig>,

    /// A pre-ordered list of the requests, if more than one
    ordered_req_list: Vec<QDnsReq>,

    /// Override options
    opts: QuerySetup,
}


impl QDns
{
    /// Initializes new empty storage for requests.
    /// 
    /// In some cases it is good idea to combine different requests, because by default
    /// all requests are performed in parallel. But in some cases it is bad idea.
    /// 
    /// # Arguments
    /// 
    /// * `resolvers` - an [Arc] [ResolveConfig] which contains configuration i.e nameservers
    /// 
    /// * `planned_reqs_len` - how many requests are planned
    /// 
    /// * `opts` - [QuerySetup] additional options or overrides. Use default() for default
    ///     values.
    /// 
    /// # Returns
    /// 
    /// Never panics. Returns Self.
    pub async 
    fn make_empty(resolvers: Option<Arc<ResolveConfig>>, opts: QuerySetup) -> CDnsResult<Self>
    {
        return Ok(
            Self
            {
                resolvers: resolvers.unwrap_or(CACHE.clone_resolve_list().await?),
                ordered_req_list: Vec::new(),
                opts: opts,
            }
        );
    }

    /// Adds new request to previously created empty storage for request with [QDns::make_empty].
    /// 
    /// # Arguemnts
    /// 
    /// * `qtype` - a [QType] type of the request
    /// 
    /// * `req_name` - a [Into] [QDnsName] which is target. i.e 'localhost' or 'domain.tld'
    /// 
    /// # Returns
    /// 
    /// * [CDnsResult] - Ok with nothing in inner type
    /// 
    /// * [CDnsResult] - Err with error description
    pub 
    fn add_request(&mut self, qtype: QType, req_name: impl Into<QDnsName>)
    {
        let qr = QDnsReq::new_into(req_name, qtype);

        self.ordered_req_list.push(qr);

        return;
    }

    /// This is helper which makes for you an A, AAAA query. The order of A and AAAA and
    /// which from both are allowed is defined in the [ResolveConfig].
    /// 
    /// Use this function directly. Do not use [QDns::make_empty]
    /// 
    /// # Arguments
    /// 
    /// * `resolvers` - an [Option] value [Arc] [ResolveConfig] which can be used to 
    ///     override the system's `resolv.conf`
    /// 
    /// * `req_name` - a [Into] [QDnsName] which is target i.e 'localhost' or 'domain.tld'
    /// 
    /// * `opts` - [QuerySetup] additional options or overrides. Use default() for default
    ///     values.
    /// 
    /// # Returns
    /// 
    /// A [CDnsResult] is returned;
    /// 
    /// * [Result::Ok] - with Self as inner type
    /// 
    /// * [Result::Err] - with error description
    pub async 
    fn make_a_aaaa_request(resolvers_opt: Option<Arc<ResolveConfig>>, req_name: impl Into<QDnsName>, 
        opts: QuerySetup) -> CDnsResult<Self>
    {
        let resolvers = resolvers_opt.unwrap_or(CACHE.clone_resolve_list().await?);

        // store the A and AAAA depending on order
        let reqs: Vec<QDnsReq> = 
            match resolvers.family
            {
                ResolveConfigFamily::INET4_INET6 => 
                {
                    let req_n: QDnsName = req_name.into();

                    vec![
                        QDnsReq::new(req_n.clone(), QType::A),
                        QDnsReq::new(req_n, QType::AAAA),
                    ]
                },
                ResolveConfigFamily::INET6_INET4 => 
                {
                    let req_n: QDnsName = req_name.into();

                    vec![
                        QDnsReq::new(req_n.clone(), QType::AAAA),
                        QDnsReq::new(req_n, QType::A),
                    ]
                },
                ResolveConfigFamily::INET6 => 
                {
                    vec![
                        QDnsReq::new(req_name.into(), QType::AAAA),
                    ]
                },
                ResolveConfigFamily::INET4 => 
                {
                    vec![
                        QDnsReq::new(req_name.into(), QType::A),
                    ]
                }
                _ =>
                {
                    // set default
                    let req_n: QDnsName = req_name.into();

                    vec![
                        QDnsReq::new(req_n.clone(), QType::A),
                        QDnsReq::new(req_n, QType::AAAA),
                    ]
                }
            };

        
        
        return Ok(
            Self
            {
                resolvers: resolvers,
                ordered_req_list: reqs,
                opts: opts,
            }
        );
    }

    /// Runs the created query/ies
    /// 
    /// # Returns
    /// 
    /// * [CDnsResult] with [QDnsQueryResult] which may contain results
    /// 
    /// * [CDnsResult] with error
    /// 
    /// Should not panic. MT-safe.
    pub async 
    fn query(mut self) -> QDnsQueryResult
    {
        // check if we need to measure time
        let now = 
            if self.opts.measure_time == true
            {
                Some(Instant::now())
            }
            else
            {
                None
            };

        // determine where to look firstly i.e file -> bind, bind -> file
        // or bind only, or file only
        if self.resolvers.lookup.is_file_first()
        {
            let mut query_res = 
                match self.lookup_file(now.as_ref()).await
                {
                    Ok(file) => 
                    {
                        if file.is_empty() == false
                        {
                            // remove records which was resolved
                            self.ordered_req_list.retain(|req| 
                                {
                                    return !file.contains_dnsreq(req);
                                }
                            );
                        }

                        file
                    },
                    Err(e) => 
                    {
                        write_error!(e);

                        QDnsQueryResult::default()
                    }
                };


            // if something left unresolved, try ask internet
            if self.ordered_req_list.is_empty() == false && self.resolvers.lookup.is_bind() == true
            {
                let res = self.process_request(now.as_ref()).await;
                
                query_res.extend(res);
            }

            return query_res;
        }
        else
        {
            let mut dns_res = self.process_request(now.as_ref()).await;
            if dns_res.is_empty() == false
            {
                // remove records which was resolved
                self.ordered_req_list.retain(|req| 
                    {
                        return !dns_res.contains_dnsreq(req);
                    }
                );
            }



            if self.ordered_req_list.is_empty() == false && self.resolvers.lookup.is_file() == true
            {
                match self.lookup_file(now.as_ref()).await
                {
                    Ok(res) => 
                    {
                        dns_res.extend(res);
                    },
                    Err(e) =>
                    {
                        write_error!(e);
                    }
                }
            }

            return dns_res;
        }
    }

    /// Returns timeout
    fn get_timeout(&self) -> Duration
    {
        if let Some(timeout) = self.opts.timeout
        {
            return Duration::from_secs(timeout as u64);
        }
        else
        {
            return Duration::from_secs(self.resolvers.timeout as u64);
        }
    }

    /// Searches in /etc/hosts
    async 
    fn lookup_file(&mut self, now: Option<&Instant>) -> CDnsResult<QDnsQueryResult>
    {
        let mut dnsquries: QDnsQueryResult = QDnsQueryResult::default();

        // check if the it is overriden
        if self.opts.ign_hosts == false
        {
            let hlist = CACHE.clone_host_list().await?;

            for req in self.ordered_req_list.iter()
            {
                match *req.get_type()
                {
                    QType::A | QType::AAAA => 
                    {
                        let req_name = String::from(req.get_req_name());

                        let Some(host_name_ent) = hlist.search_by_fqdn(req.get_type(), req_name.as_str())
                            else { continue };
                        
                        let Some(drp) = DnsResponsePayload::new_local(*req.get_type(), host_name_ent)
                            else { continue };
                        
                        // store to list
                        dnsquries.push(req.clone(), Ok(QDnsQuery::from_local(drp, now)));
                    },
                    QType::PTR => 
                    {
                        let Ok(ip) = IpAddr::try_from(req.get_req_name())
                            else { continue };

                        let Some(host_name_ent) = hlist.search_by_ip(&ip)
                            else { continue };

                        let Some(drp) = DnsResponsePayload::new_local(*req.get_type(), host_name_ent)
                            else { continue };    

                        dnsquries.push(req.clone(), Ok(QDnsQuery::from_local(drp, now))); 
                    },
                    _ => 
                        continue,
                }
            }

        }

        return Ok(dnsquries);
    }

/// Creates socket based on config and flag. If `force_tcp` is set then Tcp tap will
    /// be created.
    #[inline]
    fn create_socket(&self, force_tcp: bool, resolver: Arc<ResolveConfEntry>) -> CDnsResult<Box<NetworkTapType>>
    {
//        if self.resolvers.option_flags.contains(OptionFlags::OPT_SINGLE_REQUEST_REOPEN)
        if resolver.get_tls_type() != ConfigEntryTls::None
        {
            return NetworkTap::<TcpTlsConnection>::new(resolver, self.get_timeout());
        }
        else if self.resolvers.option_flags.is_force_tcp() == true || force_tcp == true
        {
            return NetworkTap::<TcpStream>::new(resolver, self.get_timeout());
        }
        else
        {
            return NetworkTap::<UdpSocket>::new(resolver, self.get_timeout());
        }
    }

    /// Quering all nameservers
    async 
    fn process_request(&mut self, now: Option<&Instant>) -> QDnsQueryResult
    {
        let mut responses: QDnsQueryResult = QDnsQueryResult::with_capacity(self.ordered_req_list.len());

        if self.resolvers.option_flags.is_no_parallel() == true
        {
            for req in self.ordered_req_list.iter()
            {
                let mut last_resp: Option<CDnsResult<QDnsQuery>> = None;

                for resolver in self.resolvers.get_resolvers_iter()
                {
                    match self.query_exec_seq(now, resolver.clone(), req, None).await
                    {
                        Ok(resp) =>
                        {
                            if resp.should_check_next_ns() == true
                            {
                                last_resp = Some(Ok(resp));

                                continue;
                            }
                            else
                            {
                                responses.push(req.clone(), Ok(resp));

                                let _ = last_resp.take();

                                break;
                            }
                        },
                        Err(e) =>
                        {
                            if last_resp.is_none() == true
                            {
                                last_resp = Some(Err(e));
                            }                      

                            continue;
                        }
                    }
                } // for

                responses.push(req.clone(), last_resp.take().unwrap());
            }// for
        }
        else
        {
           // let mut responses: QDnsQueryResult = QDnsQueryResult::with_capacity(self.ordered_req_list.len());

            for resolver in self.resolvers.get_resolvers_iter()
            {
                if self.ordered_req_list.is_empty() == true
                {
                    break;
                }

                match self.query_exec_pipelined(now, resolver.clone(), None).await
                {
                    Ok(resp) =>
                    {
                        for (qdns_res, qdns_que) in resp
                        {
                            if let Ok(ref resp) = qdns_que
                            {
                                if resp.should_check_next_ns() == false
                                    {
                                        self
                                            .ordered_req_list
                                            .retain(
                                                |req_item|
                                                req_item != &qdns_res
                                            );
                                    }
                            }
                            
                            responses.push(qdns_res, qdns_que);
                        }
                    },
                    Err(e) =>
                    {
                        write_error!(e);

                        continue;
                    }
                }
            }
        }

        return responses;
    }

    #[async_recursion]
    async 
    fn query_exec_pipelined(
        &self, 
        now: Option<&Instant>, 
        resolver: Arc<ResolveConfEntry>,
        requery: Option<HashMap<DnsRequestHeader, QDnsReq>>,
    ) -> CDnsResult<QDnsQueryResult>
    {
        let force_tcp = self.resolvers.option_flags.is_force_tcp() || requery.is_some();

        let mut query_headers: HashMap<DnsRequestHeader, QDnsReq> = 
            if let Some(requer) = requery
            {
                let pkts_ids = requer.iter().map(|q| q.0.get_id()).collect::<BTreeSet<u16>>();

                // regenerate que ID and collect
                requer
                    .into_iter()
                    .map(
                        |(mut qrr, qdr)| 
                        { 
                            loop
                            {
                                qrr.regenerate_id(); 

                                if pkts_ids.contains(&qrr.get_id()) == false
                                {
                                    break;
                                }
                            }
                            (qrr, qdr)
                        })
                    .collect::<HashMap<DnsRequestHeader, QDnsReq>>()
            }
            else    
            {
                let mut pkts_ids: BTreeSet<u16> = BTreeSet::new();

                // create que
                self
                    .ordered_req_list
                    .iter()
                    .map(
                        |query| 
                        {
                            
                            let mut drh_res = DnsRequestHeader::try_from(query);
                            
                            loop
                            {
                                if let Ok(ref mut drh) = drh_res
                                {
                                    if pkts_ids.contains(&drh.get_id()) == true
                                    {
                                        drh.regenerate_id();

                                        continue;
                                    }
                                    else
                                    {
                                        pkts_ids.insert(drh.get_id());
                                        break;
                                    }
                                }
                                else
                                {
                                    break;
                                }
                            }

                            drh_res.map(|dh| (dh, query.clone()))
                        }
                    )
                    .collect::<CDnsResult<HashMap<DnsRequestHeader, QDnsReq>>>()?
            };

        // create socket and store
        let mut tap = self.create_socket(force_tcp, resolver.clone())?;

        tap.connect(CDdnsGlobals::get_tcp_conn_timeout()).await?;

        // send everything
        for qh in query_headers.iter()
        {
            let pkt = qh.0.to_bytes(tap.should_append_len())?;

            tap.send(pkt.as_slice()).await?;
        }

        let mut resp: QDnsQueryResult = QDnsQueryResult::with_capacity(self.ordered_req_list.len());
        let mut requery: HashMap<DnsRequestHeader, QDnsReq> = HashMap::new();
        // poll socket
        
        loop
        {
            if query_headers.len() == 0
            {
                break;
            }
            
            tap.poll_read().await?;

            let ans = self.read_response(tap.as_mut()).await?;

            let Some((query_header, qdnsreq)) = query_headers.remove_entry(&ans.req_header)
                else
                {
                    internal_error!(CDnsErrorType::IoError, "can not find response with request: {}", ans.req_header);
                };

            ans.verify(&query_header)?;

            // verified
            let qdns_resp = QDnsQuery::from_response(tap.get_remote_addr(), ans, now);

            if let Ok(ref qdns) = qdns_resp
            {
                if qdns.get_status().should_try_tcp() == true && force_tcp == false
                {
                    requery.insert(query_header, qdnsreq);
                }
                else
                {
                    resp.push(qdnsreq, qdns_resp);
                }
            }
            else
            {
                resp.push(qdnsreq, qdns_resp);
            }
        }

        if requery.is_empty() == false
        {
            let res = self.query_exec_pipelined(now, resolver, Some(requery)).await?;

            resp.extend(res);
        }
        
        return Ok(resp);   
    }

    /// Processes request/s for nameserver
    #[async_recursion]
    async 
    fn query_exec_seq(
        &self, 
        now: Option<&Instant>, 
        resolver: Arc<ResolveConfEntry>,
        query: &QDnsReq,
        requery: Option<DnsRequestHeader>,
    ) -> CDnsResult<QDnsQuery>
    {
        let force_tcp = self.resolvers.option_flags.is_force_tcp() || requery.is_some();

        // create socket and store
        let mut tap = self.create_socket(force_tcp, resolver.clone())?;

        let query_header = 
            // form the list of reqests binded to taps
            if let Some(mut requery) = requery
            {
                requery.regenerate_id();

                requery
            }
            else
            {
                let drh_req = DnsRequestHeader::try_from(query)?;
                
                drh_req
            };

        let res = 
            {
                // open connection, with timeout because working in non parallel mode
                tap.connect(CDdnsGlobals::get_tcp_conn_timeout()).await?;

                // convert to byte packet (the error returned is global error i.e OOM)
                let pkt = query_header.to_bytes(tap.should_append_len())?;

                // send to DNS server, if error happens, then stop processing this server
                tap.send(pkt.as_slice()).await?;

                let ans = self.read_response(tap.as_mut()).await?;

                ans.verify(&query_header)?;

                // verified
                let resp = QDnsQuery::from_response(tap.get_remote_addr(), ans, now)?;

                Ok(resp)
            };

        if (res.is_ok() == true && res.as_ref().unwrap().status.should_try_tcp() == false) || 
            (res.is_err() == true && force_tcp == true)
        {
            return res;
        }
        

        return 
            self.query_exec_seq(now, resolver.clone(), query, Some(query_header)).await;
    }

    /// Reads response from tap
    async 
    fn read_response(&self, socktap: &mut (dyn SocketTap)) -> CDnsResult<DnsRequestAnswer>
    {
        if socktap.is_tcp() == false
        {
            let mut rcvbuf = vec![0_u8; 1457];

            // receive message
            let n = socktap.recv(rcvbuf.as_mut_slice()).await?;

            // parsing response to structure
            return DnsRequestAnswer::parse(&rcvbuf); //DnsRequestAnswer::try_from(rcvbuf.as_slice(), n, socktap.is_tcp())?;
        }
        else
        {
            let mut pkg_pen: [u8; 2] = [0, 0];
            let n = socktap.recv(&mut pkg_pen).await?;

            if n == 0
            {
                internal_error!(CDnsErrorType::IoError, "tcp received zero len message!");
            }
            else if n != 2
            {
                internal_error!(CDnsErrorType::IoError, "tcp expected 2 bytes to be read!");
            }

            let ln = u16::from_be_bytes(pkg_pen);

            let mut rcvbuf = vec![0_u8; ln as usize];

            // receive message
            let mut n = socktap.recv(rcvbuf.as_mut_slice()).await?;

            if n == 0
            {
                internal_error!(CDnsErrorType::IoError, "tcp received zero len message!");
            }
            else if n == 1
            {
                n = socktap.recv(&mut rcvbuf[1..]).await?;

                if n == 0
                {
                    internal_error!(CDnsErrorType::IoError, "tcp received zero len message again!");
                }

                n += 1;
            }

            return DnsRequestAnswer::parse(&rcvbuf); 
        }
    }
}

#[cfg(test)]
mod tests
{
    use std::net::IpAddr;

    use crate::{a_sync::query::QDns, common::{byte2hexchar, ip2pkt, RecordPTR, RecordReader}, QDnsQueryRec, QType, QuerySetup};

    #[tokio::test]
    async fn test_ip2pkt()
    {
        use tokio::time::Instant;
        use std::net::{IpAddr, Ipv4Addr};
        
        let test = IpAddr::V4(Ipv4Addr::new(8, 8, 8, 8));

        let now = Instant::now();

        let res = ip2pkt(&test);

        let elapsed = now.elapsed();
        println!("Elapsed: {:.2?}", elapsed);

        assert_eq!(res.is_ok(), true, "err: {}", res.err().unwrap());

        let res = res.unwrap();
        let ctrl = b"\x01\x38\x01\x38\x01\x38\x01\x38\x07\x69\x6e\x2d\x61\x64\x64\x72\x04\x61\x72\x70\x61\x00";

        assert_eq!(res.as_slice(), ctrl);
    }


    #[tokio::test]
    async fn test_byte2hexchar()
    {
        assert_eq!(byte2hexchar(1), 0x31);
        assert_eq!(byte2hexchar(9), 0x39);
        assert_eq!(byte2hexchar(10), 'a' as u8);
        assert_eq!(byte2hexchar(15), 'f' as u8);
    }


    #[tokio::test]
    async fn reverse_lookup_test()
    {
        use tokio::time::Instant;
        
        let ipp: IpAddr = "8.8.8.8".parse().unwrap();
        //let test = QDnsName::from(&ipp);

        let mut query_setup = QuerySetup::default();
        query_setup.set_measure_time(true);

        let now = Instant::now();

        let mut dns_req = 
            QDns::make_empty(None, query_setup).await.unwrap();

        dns_req.add_request(QType::PTR, ipp);

        let res = dns_req.query().await;

        let elapsed = now.elapsed();
        println!("Elapsed: {:.2?}", elapsed);

        println!("{}", res);

        assert_eq!(res.is_empty(), false);

        let recs = res.collect_ok();
        let rec = &recs[0];
        //assert_eq!(rec.server.as_str(), "/etc/hosts");
        assert_eq!(rec.status, QDnsQueryRec::Ok);

        assert_eq!(rec.resp.len(), 1);
        assert_eq!(rec.resp[0].rdata, RecordPTR::wrap(RecordPTR{ fqdn: "dns.google".to_string() }));
    }

    #[tokio::test]
    async fn reverse_lookup_hosts_test()
    {
        use tokio::time::Instant;

        let ipp: IpAddr = "127.0.0.1".parse().unwrap();
        //let test = QDnsName::from(&ipp);
        
        let now = Instant::now();

        let mut query_setup = QuerySetup::default();
        query_setup.set_measure_time(true);

        let mut dns_req = 
            QDns::make_empty(None, query_setup).await.unwrap();

        dns_req.add_request(QType::PTR, ipp);

        let res = dns_req.query().await;
        
        let elapsed = now.elapsed();
        println!("Elapsed: {:.2?}", elapsed);

        println!("{}", res);

        assert_eq!(res.is_empty(), false);

        let recs = res.collect_ok();
        let rec = &recs[0];

        assert_eq!(rec.server.as_str(), "/etc/hosts");
        assert_eq!(rec.status, QDnsQueryRec::Ok);

        assert_eq!(rec.resp.len(), 1);
        assert_eq!(rec.resp[0].rdata, RecordPTR::wrap(RecordPTR{ fqdn: "localhost".to_string() }));
    }


    #[tokio::test]
    async fn reverse_lookup_a()
    {
        use tokio::time::Instant;

       // let test = QDnsName::from("dns.google");

        let mut query_setup = QuerySetup::default();
        query_setup.set_measure_time(true);

        
        let res = QDns::make_a_aaaa_request(None, "dns.google", query_setup).await.unwrap();
        
        
        let now = Instant::now();
        let res = res.query().await;
        

        let elapsed = now.elapsed();
        println!("Elapsed: {:.2?}", elapsed);

        println!("{}", res);
    }
}