lazydns 0.2.63

use crate::config::PluginConfig;
use crate::dns::{Message, RData, ResourceRecord};
use crate::plugin::{Context, Plugin};
use crate::{RegisterPlugin, Result};
use async_trait::async_trait;
use dashmap::DashMap;
use serde_yaml::Value;
use std::fmt;
use std::net::IpAddr;
use std::sync::Arc;
use std::time::{Duration, Instant};

/// Configuration arguments for `ReverseLookupPlugin`.
///
/// - `size`: approximate capacity for the in-memory reverse cache (unused
///   beyond documentation parity with upstream implementations).
/// - `handle_ptr`: whether the plugin should attempt to answer PTR queries
///   directly from the cache.
/// - `ttl`: maximum TTL (in seconds) to honor when storing answers from
///   responses; saved entries will use the minimum of record TTL and this
///   configured value.
#[derive(Debug, Clone)]
pub struct ReverseLookupArgs {
    pub size: usize,
    pub handle_ptr: bool,
    pub ttl: u32,
}

impl Default for ReverseLookupArgs {
    fn default() -> Self {
        Self {
            size: 64 * 1024,
            handle_ptr: true,
            ttl: 7200,
        }
    }
}

/// Cache entry stored in the reverse lookup table: `(fqdn, expiration)`.
type Entry = (String, Instant);

/// Reverse lookup plugin.
///
/// The plugin collects A/AAAA answers from responses (via `save_ips_after`)
/// and keeps a small in-memory mapping from IP -> owner name. When a PTR
/// query is received and `handle_ptr` is enabled, the plugin will attempt to
/// translate the reverse name into an IP and reply from the cache if a valid
/// (non-expired) entry exists.
#[derive(RegisterPlugin)]
pub struct ReverseLookupPlugin {
    cache: Arc<DashMap<String, Entry>>,
    args: ReverseLookupArgs,
}

// Auto-register plugin builder for config-based construction

impl ReverseLookupPlugin {
    pub fn new(args: ReverseLookupArgs) -> Self {
        Self {
            cache: Arc::new(DashMap::new()),
            args,
        }
    }

    /// Parse a PTR-style qname into an `IpAddr`.
    ///
    /// Supports IPv4 reverse names under `in-addr.arpa` and IPv6 nibble
    /// reverse names under `ip6.arpa`. Returns `None` if parsing fails.
    ///
    /// This is a helper used by `execute` to detect the query target.
    fn parse_ptr(qname: &str) -> Option<IpAddr> {
        let s = qname.trim_end_matches('.').to_ascii_lowercase();
        if s.ends_with(".in-addr.arpa") {
            let without = s.trim_end_matches(".in-addr.arpa").trim_end_matches('.');
            let parts: Vec<&str> = without.split('.').collect();
            if parts.len() == 4 {
                let rev: Vec<&str> = parts.into_iter().rev().collect();
                let ip = rev.join(".");
                return ip.parse::<IpAddr>().ok();
            }
        } else if s.ends_with(".ip6.arpa") {
            // attempt to parse nibble format: labels are nibbles in reverse order
            let without = s.trim_end_matches(".ip6.arpa").trim_end_matches('.');
            let labels: Vec<&str> = without.split('.').collect();
            let mut hex = String::new();
            for label in labels.into_iter().rev() {
                if label.len() != 1 {
                    return None;
                }
                hex.push_str(label);
            }
            // group into bytes
            if !hex.len().is_multiple_of(2) {
                hex.push('0');
            }
            let bytes_res: std::result::Result<Vec<u8>, std::num::ParseIntError> = (0..hex.len())
                .step_by(2)
                .map(|i| u8::from_str_radix(&hex[i..i + 2], 16))
                .collect();
            if let Ok(bytes) = bytes_res
                && bytes.len() == 16
            {
                use std::net::Ipv6Addr;
                let mut arr = [0u8; 16];
                arr.copy_from_slice(&bytes);
                return Some(IpAddr::V6(Ipv6Addr::from(arr)));
            }
        }
        None
    }

    fn lookup(&self, ip: &IpAddr) -> Option<String> {
        let key = ip.to_string();
        if let Some(v) = self.cache.get(&key) {
            if v.value().1 > Instant::now() {
                return Some(v.value().0.clone());
            } else {
                // expired
                self.cache.remove(&key);
            }
        }
        None
    }

    fn save_from_response(&self, req: &Message, resp: &Message) {
        if resp.answer_count() == 0 {
            return;
        }
        let now = Instant::now();
        for rr in resp.answers() {
            match rr.rdata() {
                RData::A(ipv4) => {
                    let ip = IpAddr::V4(*ipv4);
                    let ttl = rr.ttl().min(self.args.ttl);
                    let exp = now + Duration::from_secs(ttl as u64);
                    let name = if req.question_count() == 1 {
                        req.questions()[0].qname().to_string()
                    } else {
                        rr.name().to_string()
                    };
                    self.cache.insert(ip.to_string(), (name, exp));
                }
                RData::AAAA(ipv6) => {
                    let ip = IpAddr::V6(*ipv6);
                    let ttl = rr.ttl().min(self.args.ttl);
                    let exp = now + Duration::from_secs(ttl as u64);
                    let name = if req.question_count() == 1 {
                        req.questions()[0].qname().to_string()
                    } else {
                        rr.name().to_string()
                    };
                    self.cache.insert(ip.to_string(), (name, exp));
                }
                _ => {}
            }
        }
    }
}

/// Public API and helpers for `ReverseLookupPlugin`.
impl ReverseLookupPlugin {
    /// Save any A/AAAA answers from a response into the internal cache.
    ///
    /// This is a convenience wrapper around `save_from_response` intended for
    /// callers that manage the sequence execution and want to record names for
    /// later PTR responses.
    #[allow(dead_code)]
    pub fn save_response(&self, req: &Message, resp: &Message) {
        self.save_from_response(req, resp);
    }
}

impl fmt::Debug for ReverseLookupPlugin {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("ReverseLookup").finish()
    }
}

#[async_trait]
impl Plugin for ReverseLookupPlugin {
    fn name(&self) -> &str {
        "reverse_lookup"
    }

    fn as_any(&self) -> &dyn std::any::Any {
        self
    }

    fn init(config: &PluginConfig) -> crate::Result<std::sync::Arc<dyn crate::plugin::Plugin>> {
        // Parse args with sensible defaults
        let args_map = config.effective_args();

        let mut args = ReverseLookupArgs::default();

        if let Some(v) = args_map.get("size") {
            match v {
                Value::Number(n) => {
                    if let Some(u) = n.as_u64() {
                        args.size = u as usize;
                    }
                }
                Value::String(s) => {
                    if let Ok(u) = s.parse::<usize>() {
                        args.size = u;
                    }
                }
                _ => {}
            }
        }

        if let Some(Value::Bool(b)) = args_map.get("handle_ptr") {
            args.handle_ptr = *b;
        }

        if let Some(v) = args_map.get("ttl") {
            match v {
                Value::Number(n) => {
                    if let Some(u) = n.as_u64() {
                        args.ttl = u as u32;
                    }
                }
                Value::String(s) => {
                    if let Ok(u) = s.parse::<u32>() {
                        args.ttl = u;
                    }
                }
                _ => {}
            }
        }

        Ok(std::sync::Arc::new(ReverseLookupPlugin::new(args)))
    }

    async fn execute(&self, ctx: &mut Context) -> Result<()> {
        // If PTR and configured, attempt to reply from cache
        let req = ctx.request();
        if self.args.handle_ptr
            && let Some(q) = req.questions().first()
            && q.qtype() == crate::dns::types::RecordType::PTR
            && let Some(ip) = Self::parse_ptr(q.qname())
            && let Some(fqdn) = self.lookup(&ip)
        {
            let mut r = Message::new();
            r.set_id(req.id());
            r.set_response(true);
            r.add_question(q.clone());
            r.add_answer(ResourceRecord::new(
                q.qname().to_string(),
                crate::dns::types::RecordType::PTR,
                crate::dns::types::RecordClass::IN,
                5,
                RData::PTR(fqdn),
            ));
            ctx.set_response(Some(r));
            return Ok(());
        }

        // Otherwise do nothing here; saving of IPs is expected to be triggered
        // by sequence runner after response is available via `save_ips_after`.
        Ok(())
    }
}

impl ReverseLookupPlugin {
    /// Helper to be called by sequence runner after response is populated.
    ///
    /// This method extracts A/AAAA answers from `resp` and records mappings
    /// from IP -> owner name in the internal cache. It is intended to be
    /// invoked by higher-level executors after the response is available.
    pub fn save_ips_after(&self, req: &Message, resp: &Message) {
        self.save_from_response(req, resp);
    }

    /// Expose lookup for tests: return the cached fqdn for an IP if present.
    ///
    /// Returns `Some(fqdn)` when the entry exists and has not yet expired.
    pub fn lookup_cached(&self, ip: &IpAddr) -> Option<String> {
        self.lookup(ip)
    }

    /// Create a quick-setup instance from a size string (upstream-compatible
    /// helper). If `s` parses as a positive integer it will be used as the
    /// cache size; otherwise defaults are applied.
    pub fn quick_setup(s: &str) -> Self {
        let mut args = ReverseLookupArgs::default();
        if !s.is_empty()
            && let Ok(n) = s.parse::<usize>()
        {
            args.size = n;
        }
        Self::new(args)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    // test only uses parse_ptr

    #[tokio::test]
    async fn test_parse_ipv4_ptr() {
        let s = "1.2.3.4.in-addr.arpa.";
        let ip = ReverseLookupPlugin::parse_ptr(s).unwrap();
        assert_eq!(ip.to_string(), "4.3.2.1");
    }

    #[test]
    fn test_save_ips_after_and_lookup_cached() {
        use crate::dns::{Message, Question, RData, RecordClass, RecordType, ResourceRecord};
        use std::net::Ipv4Addr;

        // Build request with single question
        let mut req = Message::new();
        req.add_question(Question::new(
            "example.com".to_string(),
            RecordType::A,
            RecordClass::IN,
        ));

        // Build response with one A answer
        let mut resp = Message::new();
        resp.add_answer(ResourceRecord::new(
            "example.com".to_string(),
            RecordType::A,
            RecordClass::IN,
            300,
            RData::A(Ipv4Addr::new(1, 2, 3, 4)),
        ));

        let rl = ReverseLookupPlugin::new(ReverseLookupArgs::default());

        // Ensure no entry initially
        let ip = std::net::IpAddr::V4(Ipv4Addr::new(1, 2, 3, 4));
        assert!(rl.lookup_cached(&ip).is_none());

        // Invoke save hook
        rl.save_ips_after(&req, &resp);

        // Now the lookup should return the fqdn from the request
        let got = rl.lookup_cached(&ip).expect("expected cached name");
        assert_eq!(got, "example.com");
    }

    #[tokio::test]
    async fn test_reverse_lookup_answers_ptr() {
        use crate::dns::{Question, RecordClass, RecordType};
        use std::time::Duration;
        use std::time::Instant;

        // Create plugin and pre-fill its cache for IP 10.0.0.1 -> example.com
        let args = ReverseLookupArgs::default();
        let plugin = ReverseLookupPlugin::new(args);

        // insert entry for 10.0.0.1
        let key = "10.0.0.1".to_string();
        let entry: (String, Instant) = (
            "example.com".to_string(),
            Instant::now() + Duration::from_secs(60),
        );
        plugin.cache.insert(key.clone(), entry);

        // Build a PTR question for the corresponding reverse name
        let mut msg = Message::new();
        msg.add_question(Question::new(
            "1.0.0.10.in-addr.arpa".to_string(),
            RecordType::PTR,
            RecordClass::IN,
        ));

        let mut ctx = Context::new(msg);
        plugin.execute(&mut ctx).await.unwrap();

        let resp = ctx.response().unwrap();
        assert_eq!(resp.answer_count(), 1);
        assert_eq!(
            resp.answers()[0].rdata(),
            &RData::PTR("example.com".to_string())
        );
    }
}