mdns-sd 0.15.1

//! A cache for DNS records.
//!
//! This is an internal implementation, not visible to the public API.

#[cfg(feature = "logging")]
use crate::log::{debug, trace};
use crate::{
    dns_parser::{DnsAddress, DnsPointer, DnsRecordBox, DnsSrv, InterfaceId, RRType},
    service_info::{split_sub_domain, MyIntf},
    ScopedIp,
};
use std::{
    collections::{HashMap, HashSet},
    time::SystemTime,
};

/// Associate a DnsRecord with the interface it was received on.
pub(crate) struct DnsRecordIntf {
    pub(crate) record: DnsRecordBox,
    pub(crate) src_intf: InterfaceId,
}

/// A cache for all types of DNS records.
pub(crate) struct DnsCache {
    /// DnsPointer records indexed by ty_domain
    ptr: HashMap<String, Vec<DnsRecordIntf>>,

    /// DnsSrv records indexed by the fullname of an instance
    srv: HashMap<String, Vec<DnsRecordIntf>>,

    /// DnsTxt records indexed by the fullname of an instance
    txt: HashMap<String, Vec<DnsRecordIntf>>,

    /// DnsAddr records indexed by the hostname in lowercase.
    addr: HashMap<String, Vec<DnsRecordIntf>>,

    /// A reverse lookup table from "instance fullname" to "subtype PTR name"
    subtype: HashMap<String, String>,

    /// Negative responses:
    /// A map from "instance fullname" to DnsNSec.
    nsec: HashMap<String, Vec<DnsRecordIntf>>,
}

impl DnsCache {
    pub(crate) fn new() -> Self {
        Self {
            ptr: HashMap::new(),
            srv: HashMap::new(),
            txt: HashMap::new(),
            addr: HashMap::new(),
            subtype: HashMap::new(),
            nsec: HashMap::new(),
        }
    }

    pub(crate) fn all_ptr(&self) -> &HashMap<String, Vec<DnsRecordIntf>> {
        &self.ptr
    }

    /// Count all PTR records in the cache.
    pub(crate) fn ptr_count(&self) -> usize {
        self.ptr.values().map(|v| v.len()).sum()
    }

    pub(crate) fn srv_count(&self) -> usize {
        self.srv.values().map(|v| v.len()).sum()
    }

    pub(crate) fn txt_count(&self) -> usize {
        self.txt.values().map(|v| v.len()).sum()
    }

    pub(crate) fn addr_count(&self) -> usize {
        self.addr.values().map(|v| v.len()).sum()
    }

    pub(crate) fn nsec_count(&self) -> usize {
        self.nsec.values().map(|v| v.len()).sum()
    }

    pub(crate) fn subtype_count(&self) -> usize {
        self.subtype.len()
    }

    pub(crate) fn get_ptr(&self, ty_domain: &str) -> Option<&Vec<DnsRecordIntf>> {
        self.ptr.get(ty_domain)
    }

    pub(crate) fn get_srv(&self, fullname: &str) -> Option<&Vec<DnsRecordIntf>> {
        self.srv.get(fullname)
    }

    pub(crate) fn get_txt(&self, fullname: &str) -> Option<&Vec<DnsRecordIntf>> {
        self.txt.get(fullname)
    }

    pub(crate) fn get_addr(&self, hostname: &str) -> Option<&Vec<DnsRecordIntf>> {
        self.addr.get(&hostname.to_lowercase())
    }

    /// A reverse lookup table from "instance fullname" to "subtype PTR name"
    pub(crate) fn get_subtype(&self, fullname: &str) -> Option<&String> {
        self.subtype.get(fullname)
    }

    /// Returns the list of instances that has `host` as its hostname.
    pub(crate) fn get_instances_on_host(&self, host: &str) -> Vec<String> {
        self.srv
            .iter()
            .filter_map(|(instance, srv_list)| {
                if let Some(item) = srv_list.first() {
                    if let Some(dns_srv) = item.record.any().downcast_ref::<DnsSrv>() {
                        if dns_srv.host() == host {
                            return Some(instance.clone());
                        }
                    }
                }
                None
            })
            .collect()
    }

    /// Returns a hashmap of hostnames and their addresses for a given `host`.
    ///
    /// Note that the keys in the returned HashMap are the same hostname, with different cases
    /// of letters (e.g. "example.local.", "Example.local.", "EXAMPLE.local.").
    pub(crate) fn get_addresses_for_host(&self, host: &str) -> HashMap<String, HashSet<ScopedIp>> {
        let hostname_lower = host.to_lowercase();
        let mut result = HashMap::new();

        if let Some(records) = self.addr.get(&hostname_lower) {
            for record in records {
                if let Some(dns_addr) = record.record.any().downcast_ref::<DnsAddress>() {
                    let record_name = record.record.get_name().to_string();
                    let address = dns_addr.address();

                    // Use the entry API to insert or update the HashSet for the record_name
                    result
                        .entry(record_name)
                        .or_insert_with(HashSet::new)
                        .insert(address);
                }
            }
        }

        result
    }

    /// Returns a list of resource records (name, rr_type) that need to be queried in order to
    /// verify the `instance`.
    ///
    /// If `expire_at` is not None, the resource records' expire time will be updated.
    pub(crate) fn service_verify_queries(
        &mut self,
        instance: &str,
        expire_at: Option<u64>,
    ) -> Vec<(String, RRType)> {
        let Some(srv_vec) = self.srv.get_mut(instance) else {
            return Vec::new();
        };

        let mut query_vec = vec![(instance.to_string(), RRType::SRV)];

        for srv in srv_vec {
            if let Some(new_expire) = expire_at {
                srv.record.set_expire_sooner(new_expire);
            }

            let Some(srv_record) = srv.record.any().downcast_ref::<DnsSrv>() else {
                continue;
            };

            // Will verify addresses for the hostname.
            query_vec.push((srv_record.host().to_string(), RRType::A));
            query_vec.push((srv_record.host().to_string(), RRType::AAAA));

            if let Some(new_expire) = expire_at {
                if let Some(addrs) = self.addr.get_mut(srv_record.host()) {
                    for addr in addrs {
                        addr.record.set_expire_sooner(new_expire);
                    }
                }
            }
        }

        query_vec
    }

    /// Update a DNSRecord TTL if already exists, otherwise insert a new record.
    ///
    /// Returns `None` if `incoming` is invalid / unrecognized, otherwise returns
    /// (a new record, true) or (existing record with TTL updated, false).
    ///
    /// If you need to add new timers for related records, push into `timers`.
    pub(crate) fn add_or_update(
        &mut self,
        intf: &MyIntf,
        incoming: DnsRecordBox,
        timers: &mut Vec<u64>,
        is_for_us: bool,
    ) -> Option<(&DnsRecordIntf, bool)> {
        let entry_name = incoming.get_name().to_string();

        // If it is PTR with subtype, store a mapping from the instance fullname
        // to the subtype in this cache.
        if incoming.get_type() == RRType::PTR && is_for_us {
            let (_, subtype_opt) = split_sub_domain(&entry_name);
            if let Some(subtype) = subtype_opt {
                if let Some(ptr) = incoming.any().downcast_ref::<DnsPointer>() {
                    if !self.subtype.contains_key(ptr.alias()) {
                        self.subtype
                            .insert(ptr.alias().to_string(), subtype.to_string());
                    }
                }
            }
        }

        // get the existing records for the type.
        let entry_name_lower = entry_name.to_lowercase();
        let record_vec = match incoming.get_type() {
            RRType::PTR => self.ptr.entry(entry_name).or_default(),
            RRType::SRV => self.srv.entry(entry_name).or_default(),
            RRType::TXT => self.txt.entry(entry_name).or_default(),
            RRType::A | RRType::AAAA => self.addr.entry(entry_name_lower).or_default(),
            RRType::NSEC => self.nsec.entry(entry_name).or_default(),
            _ => return None,
        };

        // No existing records for this name and type, and not for us.
        if record_vec.is_empty() && !is_for_us {
            trace!("add_or_update: not for us: {}", incoming.get_name());
            return None;
        }

        if incoming.get_cache_flush() {
            let now = current_time_millis();
            let class = incoming.get_class();
            let rtype = incoming.get_type();

            record_vec.iter_mut().for_each(|r| {
                // When cache flush is asked, we set expire date to 1 second in the future if:
                // - The record has the same rclass
                // - The record was created more than 1 second ago.
                // - The record expire is more than 1 second away.
                // Ref: RFC 6762 Section 10.2
                //
                // Note: when the updated record actually expires, it will trigger events properly.
                let mut should_flush = false;

                if class == r.record.get_class()
                    && rtype == r.record.get_type()
                    && now > r.record.get_created() + 1000
                    && r.record.get_expire() > now + 1000
                {
                    should_flush = true;

                    // additional checks for address records.
                    if rtype == RRType::A || rtype == RRType::AAAA {
                        if let Some(addr) = r.record.any().downcast_ref::<DnsAddress>() {
                            if let Some(addr_b) = incoming.any().downcast_ref::<DnsAddress>() {
                                should_flush = addr.interface_id.index == addr_b.interface_id.index;
                            }
                        }
                    }
                }

                if should_flush {
                    trace!("FLUSH one record: {:?}", &r.record);
                    let new_expire = now + 1000;
                    r.record.set_expire(new_expire);

                    // Add a timer so the run loop will handle this expire.
                    timers.push(new_expire);
                }
            });
        }

        // update TTL for existing record or create a new record.
        let (idx, updated) = match record_vec
            .iter_mut()
            .enumerate()
            .find(|(_idx, r)| r.record.matches(incoming.as_ref()))
        {
            Some((i, r)) => {
                // It is possible that this record was just updated in cache_flush
                // processing. That's okay. We can still reset here.
                r.record.reset_ttl(incoming.as_ref());
                (i, false)
            }
            None => {
                let new_record = DnsRecordIntf {
                    record: incoming,
                    src_intf: intf.into(),
                };
                record_vec.insert(0, new_record); // A new record.
                (0, true)
            }
        };

        Some((record_vec.get(idx).unwrap(), updated))
    }

    /// Remove a record from the cache if exists, otherwise no-op
    pub(crate) fn remove(&mut self, record: &DnsRecordBox) -> bool {
        let mut found = false;
        let record_name = record.get_name();
        let record_vec = match record.get_type() {
            RRType::PTR => self.ptr.get_mut(record_name),
            RRType::SRV => self.srv.get_mut(record_name),
            RRType::TXT => self.txt.get_mut(record_name),
            RRType::A | RRType::AAAA => self.addr.get_mut(record_name),
            _ => return found,
        };
        if let Some(record_vec) = record_vec {
            record_vec.retain(|x| match x.record.matches(record.as_ref()) {
                true => {
                    found = true;
                    false
                }
                false => true,
            });
        }
        found
    }

    /// Iterates all ADDR records and remove ones that expired.
    /// Returns the expired ones in a map of names and addresses.
    pub(crate) fn evict_expired_addr(&mut self, now: u64) -> HashMap<String, HashSet<ScopedIp>> {
        let mut removed = HashMap::new();

        self.addr.retain(|_, records| {
            records.retain(|addr| {
                let expired = addr.record.get_record().is_expired(now);
                if expired {
                    if let Some(addr_record) = addr.record.any().downcast_ref::<DnsAddress>() {
                        trace!("evict expired ADDR: {:?}", addr_record);
                        removed
                            .entry(addr.record.get_name().to_string())
                            .or_insert_with(HashSet::new)
                            .insert(addr_record.address());
                    }
                }
                !expired
            });

            !records.is_empty()
        });

        removed
    }

    /// Evicts expired PTR and SRV, TXT records for each ty_domain in the cache, and
    /// returns the set of expired instance names for each ty_domain.
    ///
    /// An instance in the returned set indicates its PTR and/or SRV record has expired.
    pub(crate) fn evict_expired_services(&mut self, now: u64) -> HashMap<String, HashSet<String>> {
        let mut expired_instances = HashMap::new();

        // Check all ty_domain in the cache by following all PTR records, regardless
        // if the ty_domain is actively queried or not.
        for (ty_domain, ptr_records) in self.ptr.iter_mut() {
            for ptr in ptr_records.iter() {
                if let Some(dns_ptr) = ptr.record.any().downcast_ref::<DnsPointer>() {
                    let instance_name = dns_ptr.alias();

                    // evict expired SRV records of this instance
                    if let Some(srv_records) = self.srv.get_mut(instance_name) {
                        srv_records.retain(|srv| {
                            let expired = srv.record.get_record().is_expired(now);
                            !expired
                        });

                        if srv_records.is_empty() {
                            debug!("expired SRV for {}: {:?}", ty_domain, instance_name);
                            expired_instances
                                .entry(ty_domain.to_string())
                                .or_insert_with(HashSet::new)
                                .insert(instance_name.to_string());

                            // don't keep empty value for this key.
                            self.srv.remove(instance_name);
                        }
                    }

                    // evict expired TXT records of this instance
                    if let Some(txt_records) = self.txt.get_mut(instance_name) {
                        txt_records.retain(|txt| !txt.record.get_record().is_expired(now))
                    }
                }
            }

            // evict expired PTR records
            ptr_records.retain(|x| {
                let expired = x.record.get_record().is_expired(now);
                if expired {
                    if let Some(dns_ptr) = x.record.any().downcast_ref::<DnsPointer>() {
                        trace!("expired PTR: domain:{ty_domain} record: {:?}", dns_ptr);
                        expired_instances
                            .entry(ty_domain.to_string())
                            .or_insert_with(HashSet::new)
                            .insert(dns_ptr.alias().to_string());
                    }
                }
                !expired
            });
        }

        expired_instances
    }

    /// Removes all records of a service type: PTR, SRV, TXT records and any ADDR records
    /// that are not referenced by any SRV record.
    pub(crate) fn remove_service_type(&mut self, ty_domain: &str) {
        let Some(ptr_records) = self.ptr.get_mut(ty_domain) else {
            return;
        };

        let mut hosts = HashSet::new();

        for ptr in ptr_records.iter() {
            if let Some(dns_ptr) = ptr.record.any().downcast_ref::<DnsPointer>() {
                let instance_name = dns_ptr.alias();

                // collect all hostnames from SRV records of this instance
                if let Some(srv_records) = self.srv.get_mut(instance_name) {
                    for srv in srv_records.iter() {
                        if let Some(dns_srv) = srv.record.any().downcast_ref::<DnsSrv>() {
                            hosts.insert(dns_srv.host().to_lowercase());
                        }
                    }
                }

                // remove all SRV records of this instance
                self.srv.remove(instance_name);

                // remove all TXT records of this instance
                self.txt.remove(instance_name);
            }
        }

        self.ptr.remove(ty_domain);

        // Check all hostnames in `hosts`: for each hostname, check if any SRV record
        // has `hostname` as its host. If no such SRV, remove the ADDR records of this hostname.
        for host in hosts {
            let mut has_srv = false;
            for srv_records in self.srv.values() {
                for srv in srv_records.iter() {
                    if let Some(dns_srv) = srv.record.any().downcast_ref::<DnsSrv>() {
                        if dns_srv.host().to_lowercase() == host {
                            has_srv = true;
                            break;
                        }
                    }
                }
                if has_srv {
                    break;
                }
            }

            if !has_srv {
                self.addr.remove(&host);
            }
        }
    }

    /// Checks refresh due for PTR records of `ty_domain`.
    /// Returns all updated refresh time.
    pub(crate) fn refresh_due_ptr(&mut self, ty_domain: &str) -> HashSet<u64> {
        let now = current_time_millis();

        // Check all PTR records for this ty_domain.
        self.ptr
            .get_mut(ty_domain)
            .into_iter()
            .flatten()
            .filter_map(|record| record.record.updated_refresh_time(now))
            .collect()
    }

    /// Returns a tuple of:
    /// 1. the map of instance names together with RRType(s) that are due for refresh
    ///    its SRV or TXT records.
    /// 2. the set of new timers that are due for refresh.
    pub(crate) fn refresh_due_srv_txt(
        &mut self,
        ty_domain: &str,
    ) -> (HashMap<String, Vec<RRType>>, HashSet<u64>) {
        let now = current_time_millis();

        let instances: Vec<_> = self
            .ptr
            .get(ty_domain)
            .into_iter()
            .flatten()
            .filter(|record| !record.record.get_record().is_expired(now))
            .filter_map(|record| {
                record
                    .record
                    .any()
                    .downcast_ref::<DnsPointer>()
                    .map(|ptr| ptr.alias())
            })
            .collect();

        let mut refresh_due: HashMap<String, Vec<RRType>> = HashMap::new();
        let mut new_timers = HashSet::new();
        for instance in instances {
            // Check SRV records.
            let refresh_timers: HashSet<u64> = self
                .srv
                .get_mut(instance)
                .into_iter()
                .flatten()
                .filter_map(|record| record.record.updated_refresh_time(now))
                .collect();

            if !refresh_timers.is_empty() {
                refresh_due
                    .entry(instance.to_string())
                    .and_modify(|v| v.push(RRType::SRV))
                    .or_insert(vec![RRType::SRV]);
                new_timers.extend(refresh_timers);
            }

            // Check TXT records.
            let refresh_timers: HashSet<u64> = self
                .txt
                .get_mut(instance)
                .into_iter()
                .flatten()
                .filter_map(|record| record.record.updated_refresh_time(now))
                .collect();

            if !refresh_timers.is_empty() {
                refresh_due
                    .entry(instance.to_string())
                    .and_modify(|v| v.push(RRType::TXT))
                    .or_insert(vec![RRType::TXT]);
                new_timers.extend(refresh_timers);
            }
        }

        (refresh_due, new_timers)
    }

    /// Returns the set of `host`, where refreshing the A / AAAA records is due
    /// for a `ty_domain`.
    pub(crate) fn refresh_due_hosts(&mut self, ty_domain: &str) -> (HashSet<String>, HashSet<u64>) {
        let now = current_time_millis();

        let instances: Vec<_> = self
            .ptr
            .get(ty_domain)
            .into_iter()
            .flatten()
            .filter(|record| !record.record.get_record().is_expired(now))
            .filter_map(|record| {
                record
                    .record
                    .any()
                    .downcast_ref::<DnsPointer>()
                    .map(|ptr| ptr.alias())
            })
            .collect();

        // Collect hostnames we have browsers for by SRV records.
        let mut hostnames_browsed = HashSet::new();
        for instance in instances {
            let hosts: HashSet<String> = self
                .srv
                .get(instance)
                .into_iter()
                .flatten()
                .filter_map(|record| {
                    record
                        .record
                        .any()
                        .downcast_ref::<DnsSrv>()
                        .map(|srv| srv.host().to_string())
                })
                .collect();

            hostnames_browsed.extend(hosts);
        }
        let mut refresh_due = HashSet::new();
        let mut new_timers = HashSet::new();
        for hostname in hostnames_browsed {
            let refresh_timers: HashSet<u64> = self
                .addr
                .get_mut(&hostname)
                .into_iter()
                .flatten()
                .filter_map(|record| record.record.updated_refresh_time(now))
                .collect();

            if !refresh_timers.is_empty() {
                refresh_due.insert(hostname);
                new_timers.extend(refresh_timers);
            }
        }
        (refresh_due, new_timers)
    }

    /// Returns the set of A/AAAA records that are due for refresh for a `hostname`.
    ///
    /// For these records, their refresh time will be updated so that they will not refresh again.
    pub(crate) fn refresh_due_hostname_resolutions(
        &mut self,
        hostname: &str,
    ) -> HashSet<(String, ScopedIp)> {
        let now = current_time_millis();

        self.addr
            .get_mut(hostname)
            .into_iter()
            .flatten()
            .filter_map(|record| {
                let rec = record.record.get_record_mut();
                if rec.is_expired(now) || !rec.refresh_due(now) {
                    return None;
                }
                rec.refresh_no_more();

                Some((
                    hostname.to_owned(),
                    record
                        .record
                        .any()
                        .downcast_ref::<DnsAddress>()
                        .unwrap()
                        .address(),
                ))
            })
            .collect()
    }

    /// Returns a list of Known Answer for a given question of `name` with `qtype`.
    /// The timestamp `now` is passed in to check TTL.
    ///
    /// Reference:  RFC 6762 section 7.1
    pub(crate) fn get_known_answers<'a>(
        &'a self,
        name: &str,
        qtype: RRType,
        now: u64,
    ) -> Vec<&'a DnsRecordIntf> {
        let records_opt = match qtype {
            RRType::PTR => self.get_ptr(name),
            RRType::SRV => self.get_srv(name),
            RRType::A | RRType::AAAA => self.get_addr(name),
            RRType::TXT => self.get_txt(name),
            _ => None,
        };

        let records = match records_opt {
            Some(items) => items,
            None => return Vec::new(),
        };

        // From RFC 6762 section 7.1:
        // ..Generally, this applies only to Shared records, not Unique records,..
        //
        // ..a Multicast DNS querier SHOULD NOT include
        // records in the Known-Answer list whose remaining TTL is less than
        // half of their original TTL.
        records
            .iter()
            .filter(move |r| {
                !r.record.get_record().is_unique() && !r.record.get_record().halflife_passed(now)
            })
            .collect()
    }

    pub(crate) fn remove_addrs_on_disabled_intf(&mut self, disabled_if_index: u32) {
        for (host, records) in self.addr.iter_mut() {
            records.retain(|record| {
                let Some(dns_addr) = record.record.any().downcast_ref::<DnsAddress>() else {
                    return false; // invalid address record.
                };

                // Remove the record if it is on this interface.
                if dns_addr.interface_id.index == disabled_if_index {
                    debug!(
                        "removing ADDR on disabled intf: {:?} host {host}",
                        dns_addr.interface_id.name
                    );
                    false
                } else {
                    true
                }
            });
        }
    }

    /// Removes all records that were received on `intf_id`.
    pub(crate) fn remove_records_on_intf(
        &mut self,
        intf_id: InterfaceId,
    ) -> HashMap<String, HashSet<String>> {
        let mut removed_instances = HashMap::new();

        self.ptr.iter_mut().for_each(|(name, records)| {
            let mut instances: HashSet<String> = HashSet::new();

            records.retain(|r| {
                if r.src_intf == intf_id {
                    if let Some(dns_ptr) = r.record.any().downcast_ref::<DnsPointer>() {
                        trace!("removing PTR on intf {:?}: {:?}", intf_id, dns_ptr);
                        instances.insert(dns_ptr.alias().to_string());
                    }
                    false
                } else {
                    true
                }
            });

            for instance in instances.drain() {
                // if no more record for this instance in `records`, it means its PTR record is removed.
                if !records.iter().any(|r| {
                    if let Some(dns_ptr) = r.record.any().downcast_ref::<DnsPointer>() {
                        dns_ptr.alias() == instance
                    } else {
                        false
                    }
                }) {
                    removed_instances
                        .entry(name.to_string())
                        .or_insert_with(HashSet::new)
                        .insert(instance);
                }
            }
        });

        self.ptr.retain(|_, records| !records.is_empty());

        self.srv.values_mut().for_each(|records| {
            records.retain(|r| r.src_intf != intf_id);
        });
        self.srv.retain(|_, records| !records.is_empty());

        self.txt.values_mut().for_each(|records| {
            records.retain(|r| r.src_intf != intf_id);
        });
        self.txt.retain(|_, records| !records.is_empty());

        self.addr.values_mut().for_each(|records| {
            records.retain(|r| r.src_intf != intf_id);
        });
        self.addr.retain(|_, records| !records.is_empty());

        self.nsec.values_mut().for_each(|records| {
            records.retain(|r| r.src_intf != intf_id);
        });
        self.nsec.retain(|_, records| !records.is_empty());

        removed_instances
    }
}

/// Returns UNIX time in millis
pub(crate) fn current_time_millis() -> u64 {
    SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .expect("failed to get current UNIX time")
        .as_millis() as u64
}