Alfis/src/blockchain/filter.rs

use std::collections::HashMap;
use std::net::{IpAddr, SocketAddr};
use std::sync::{Arc, Mutex};
use std::time::Instant;

#[allow(unused_imports)]
use log::{debug, error, info, trace, warn};
use rand::seq::SliceRandom;

use crate::blockchain::transaction::DomainData;
use crate::dns::filter::DnsFilter;
use crate::dns::protocol::{DnsPacket, DnsQuestion, DnsRecord, QueryType, ResultCode, TransientTtl};
use crate::Context;
use crate::dns::client::{DnsClient, DnsNetworkClient};

const NAME_SERVER: &str = "ns.alfis.name";
const SERVER_ADMIN: &str = "admin.alfis.name";

/// Unbound-style RTT band width in milliseconds.
/// Servers within min_rtt + BAND are considered equally good.
const RTT_BAND_MS: f64 = 100.0;
/// EWMA smoothing factor: 87.5% history, 12.5% new measurement.
const EWMA_WEIGHT: f64 = 7.0 / 8.0;
/// Penalty RTT assigned on timeout/failure (ms).
const TIMEOUT_PENALTY_MS: f64 = 5000.0;
/// Stats older than this are expired so the server gets re-probed.
const STATS_EXPIRE_SECS: u64 = 900;

struct NsStats {
    rtt: f64,
    last_update: Instant,
}

pub struct BlockchainFilter {
    context: Arc<Mutex<Context>>,
    ns_stats: Arc<Mutex<HashMap<IpAddr, NsStats>>>,
}

impl BlockchainFilter {
    pub fn new(context: Arc<Mutex<Context>>) -> Self {
        BlockchainFilter {
            context,
            ns_stats: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    fn add_soa_record(zone: String, serial: u32, packet: &mut DnsPacket) {
        packet.authorities.push(DnsRecord::SOA {
            domain: zone,
            m_name: String::from(NAME_SERVER),
            r_name: String::from(SERVER_ADMIN),
            serial,
            refresh: 3600,
            retry: 300,
            expire: 604800,
            minimum: 60,
            ttl: TransientTtl(60)
        });
    }

    fn get_zone_response(&self, zone: &str, serial: u32, packet: &mut DnsPacket) -> bool {
        let have_zone = self.context.lock().unwrap().chain.is_available_zone(zone);
        if have_zone {
            BlockchainFilter::add_soa_record(zone.to_owned(), serial, packet);
        }
        have_zone
    }

    fn lookup_from_ns(qname: &str, qtype: QueryType, servers: &[IpAddr], ns_stats: &Arc<Mutex<HashMap<IpAddr, NsStats>>>) -> Option<DnsPacket> {
        let mut dns_client = DnsNetworkClient::new();
        dns_client.run().unwrap();
        let timeout = std::time::Duration::from_secs(2);

        // Build ordered server list using RTT banding
        let ordered = Self::select_servers(servers, ns_stats);

        for server in &ordered {
            let addr = SocketAddr::new(*server, 53);
            let start = Instant::now();
            match dns_client.send_udp_query(qname, qtype, addr, false, timeout) {
                Ok(res) => {
                    let elapsed = start.elapsed().as_secs_f64() * 1000.0;
                    Self::update_ns_stats(ns_stats, *server, elapsed);
                    dns_client.stop();
                    return Some(res);
                }
                Err(_) => {
                    Self::update_ns_stats(ns_stats, *server, TIMEOUT_PENALTY_MS);
                }
            }
        }
        dns_client.stop();
        None
    }

    /// Select servers using Unbound-style RTT banding.
    /// Servers with no stats or expired stats are treated as preferred (to be probed).
    /// Among known servers, those within min_rtt + RTT_BAND_MS are preferred.
    /// Each group is shuffled, then preferred servers come first.
    fn select_servers(servers: &[IpAddr], ns_stats: &Arc<Mutex<HashMap<IpAddr, NsStats>>>) -> Vec<IpAddr> {
        let now = Instant::now();
        let stats = ns_stats.lock().unwrap();

        // Separate into known (with valid stats) and unknown
        let mut known: Vec<(IpAddr, f64)> = Vec::new();
        let mut unknown: Vec<IpAddr> = Vec::new();
        for &ip in servers {
            match stats.get(&ip) {
                Some(s) if now.duration_since(s.last_update).as_secs() < STATS_EXPIRE_SECS => {
                    known.push((ip, s.rtt));
                }
                _ => {
                    unknown.push(ip);
                }
            }
        }
        drop(stats);

        let mut rng = rand::thread_rng();

        if known.is_empty() {
            // No stats yet — shuffle all and probe
            unknown.shuffle(&mut rng);
            return unknown;
        }

        let min_rtt = known.iter().map(|(_, rtt)| *rtt).fold(f64::INFINITY, f64::min);
        let band_threshold = min_rtt + RTT_BAND_MS;

        let mut preferred: Vec<IpAddr> = Vec::new();
        let mut fallback: Vec<IpAddr> = Vec::new();
        for (ip, rtt) in &known {
            if *rtt <= band_threshold {
                preferred.push(*ip);
            } else {
                fallback.push(*ip);
            }
        }

        // Unknown servers join the preferred group to get probed
        preferred.extend(unknown);
        preferred.shuffle(&mut rng);
        fallback.shuffle(&mut rng);
        preferred.extend(fallback);
        preferred
    }

    fn update_ns_stats(ns_stats: &Arc<Mutex<HashMap<IpAddr, NsStats>>>, ip: IpAddr, rtt_ms: f64) {
        let mut stats = ns_stats.lock().unwrap();
        let entry = stats.entry(ip).or_insert(NsStats {
            rtt: rtt_ms,
            last_update: Instant::now(),
        });
        entry.rtt = entry.rtt * EWMA_WEIGHT + rtt_ms * (1.0 - EWMA_WEIGHT);
        entry.last_update = Instant::now();
    }

    fn create_packet(&self, qname: &str, qtype: QueryType, zone: String, answers: Vec<DnsRecord>, ns_records: Vec<DnsRecord>, glue_records: Vec<DnsRecord>) -> Option<DnsPacket> {
        if !answers.is_empty() {
            // Create DnsPacket with answers
            let mut packet = DnsPacket::new();
            packet.header.authoritative_answer = true;
            packet.questions.push(DnsQuestion::new(String::from(qname), qtype));
            for answer in answers {
                packet.answers.push(answer);
            }
            // Add NS records to authority section
            for ns_record in ns_records {
                packet.authorities.push(ns_record);
            }
            // Add GLUE records to additional section (resources)
            for glue_record in glue_records {
                packet.resources.push(glue_record);
            }
            //trace!("Returning packet: {:?}", &packet);
            Some(packet)
        } else {
            // Create DnsPacket without answers
            let mut packet = DnsPacket::new();
            packet.header.authoritative_answer = true;
            packet.header.rescode = ResultCode::NXDOMAIN;
            packet.questions.push(DnsQuestion::new(String::from(qname), qtype));
            let serial = self.context.lock().unwrap().chain.get_soa_serial();
            BlockchainFilter::add_soa_record(zone, serial, &mut packet);
            //trace!("Returning packet: {:?}", &packet);
            Some(packet)
        }
    }

    fn resolve_by_ns(qname: &str, qtype: QueryType, top_domain: &String, data: &DomainData, recursive: bool, ns_stats: &Arc<Mutex<HashMap<IpAddr, NsStats>>>) -> (bool, Option<DnsPacket>) {
        // First we search for NS records, collecting nameserver domains
        let mut hosts = Vec::new();
        for record in data.records.iter() {
            if record.get_querytype() == QueryType::NS {
                match &record {
                    DnsRecord::NS { domain, host, .. } if domain == "@" => {
                        hosts.push(host.to_owned());
                    }
                    _ => ()
                }
            }
        }

        if hosts.is_empty() {
            return (false, None);
        }

        // If non-recursive, return a referral response with NS and GLUE records
        if !recursive {
            trace!("Non-recursive query for delegated domain {}, returning referral", qname);
            let ns_records = BlockchainFilter::get_ns_records(data, top_domain);
            let glue_records = BlockchainFilter::get_glue_records(data, top_domain, &hosts);

            let mut packet = DnsPacket::new();
            packet.header.authoritative_answer = false;  // Not authoritative for the answer, but for the zone
            packet.questions.push(DnsQuestion::new(String::from(qname), qtype));
            // Add NS records to authority section
            for ns_record in ns_records {
                packet.authorities.push(ns_record);
            }
            // Add GLUE records to additional section (resources)
            for glue_record in glue_records {
                packet.resources.push(glue_record);
            }
            return (true, Some(packet));
        }

        // For recursive queries, search for glue records to query external servers
        let mut servers = Vec::new();
        for record in data.records.iter() {
            match &record {
                DnsRecord::A { domain, addr, .. } => {
                    let domain = format!("{}.{}", &domain, &top_domain);
                    for host in &hosts {
                        if &domain == host {
                            servers.push(IpAddr::from(addr.clone()));
                        }
                    }
                }
                DnsRecord::AAAA { domain, addr, .. } => {
                    let domain = format!("{}.{}", &domain, &top_domain);
                    for host in &hosts {
                        if &domain == host {
                            servers.push(IpAddr::from(addr.clone()));
                        }
                    }
                }
                _ => ()
            }
        }

        if !servers.is_empty() {
            trace!("Found NS servers for domain {}: {:?}", &qname, &servers);
            let answer = BlockchainFilter::lookup_from_ns(qname, qtype, &servers, ns_stats);
            if let Some(packet) = &answer {
                trace!("Resolved {:?} from NS: {:?}", (qname, qtype), &packet.answers);
            }
            return (true, answer);
        }

        (false, None)
    }

    /// Extract NS records from domain data and return them
    fn get_ns_records(data: &DomainData, top_domain: &str) -> Vec<DnsRecord> {
        data.records.iter()
            .filter_map(|record| {
                if let DnsRecord::NS { domain, host, ttl } = record {
                    if domain == "@" {
                        return Some(DnsRecord::NS {
                            domain: String::from(top_domain),
                            host: host.clone(),
                            ttl: *ttl
                        });
                    }
                }
                None
            })
            .collect()
    }

    /// Extract GLUE records (A/AAAA records for NS hosts within the same domain)
    fn get_glue_records(data: &DomainData, top_domain: &str, ns_hosts: &[String]) -> Vec<DnsRecord> {
        let mut glue_records = Vec::new();

        for record in data.records.iter() {
            match record {
                DnsRecord::A { domain, addr, ttl } => {
                    let full_domain = if domain == "@" {
                        String::from(top_domain)
                    } else {
                        format!("{}.{}", domain, top_domain)
                    };

                    if ns_hosts.iter().any(|ns| ns == &full_domain) {
                        glue_records.push(DnsRecord::A {
                            domain: full_domain,
                            addr: addr.clone(),
                            ttl: *ttl
                        });
                    }
                }
                DnsRecord::AAAA { domain, addr, ttl } => {
                    let full_domain = if domain == "@" {
                        String::from(top_domain)
                    } else {
                        format!("{}.{}", domain, top_domain)
                    };

                    if ns_hosts.iter().any(|ns| ns == &full_domain) {
                        glue_records.push(DnsRecord::AAAA {
                            domain: full_domain,
                            addr: addr.clone(),
                            ttl: *ttl
                        });
                    }
                }
                _ => {}
            }
        }

        glue_records
    }
}

impl DnsFilter for BlockchainFilter {
    fn lookup(&self, qname: &str, qtype: QueryType, recursive: bool) -> Option<DnsPacket> {
        // Lowercase for case-insensitive lookup (blockchain stores domains as lowercase)
        let qname_lower = qname.to_lowercase();
        let top_domain;
        let subdomain;
        let parts: Vec<&str> = qname_lower.rsplitn(3, '.').collect();
        match parts.len() {
            1 => {
                let mut packet = DnsPacket::new();
                let serial = self.context.lock().unwrap().chain.get_soa_serial();
                if self.get_zone_response(parts[0], serial, &mut packet) {
                    return Some(packet);
                }
                return None;
            }
            2 => {
                top_domain = format!("{}.{}", parts[1], parts[0]);
                subdomain = String::new();
            }
            _ => {
                top_domain = format!("{}.{}", parts[1], parts[0]);
                subdomain = String::from(parts[2]);
            }
        }
        //trace!("Searching record type '{:?}', name '{}' for domain '{}'", &qtype, &subdomain, &search);

        let data = self.context.lock().unwrap().chain.get_domain_info(&top_domain);
        let zone = parts[0].to_owned();
        match data {
            None => {
                if self.context.lock().unwrap().chain.is_available_zone(&zone) {
                    trace!("Not found data for domain {}", &top_domain);
                    // Create DnsPacket
                    let mut packet = DnsPacket::new();
                    packet.questions.push(DnsQuestion::new(String::from(qname), qtype));
                    packet.header.rescode = ResultCode::NXDOMAIN;
                    packet.header.authoritative_answer = true;
                    let serial = self.context.lock().unwrap().chain.get_soa_serial();
                    BlockchainFilter::add_soa_record(zone, serial, &mut packet);
                    //trace!("Returning packet: {:?}", &packet);
                    return Some(packet);
                }
            }
            Some(data) => {
                trace!("Found data for domain {}", &top_domain);
                let mut data: DomainData = match serde_json::from_str(&data) {
                    Err(_) => {
                        return None;
                    }
                    Ok(data) => data
                };

                // Check if this domain has NS records and needs to resolve all records through them
                // But skip this if we're querying for NS records themselves - return them directly
                if qtype != QueryType::NS {
                    let (has_ns, result) = Self::resolve_by_ns(qname, qtype, &top_domain, &data, recursive, &self.ns_stats);
                    if has_ns {
                        return result;
                    }
                }

                let mut answers: Vec<DnsRecord> = Vec::new();
                let mut cname: Option<DnsRecord> = None;
                for mut record in data.records.iter_mut() {
                    if record.get_querytype() == qtype || record.get_querytype() == QueryType::CNAME {
                        match &mut record {
                            DnsRecord::A { domain, .. }
                            | DnsRecord::AAAA { domain, .. }
                            | DnsRecord::NS { domain, .. }
                            | DnsRecord::CNAME { domain, .. }
                            | DnsRecord::SRV { domain, .. }
                            | DnsRecord::TLSA { domain, .. }
                            | DnsRecord::MX { domain, .. }
                            | DnsRecord::UNKNOWN { domain, .. }
                            | DnsRecord::SOA { domain, .. }
                            | DnsRecord::TXT { domain, .. } if (domain == "@" && subdomain.is_empty()) || domain == &subdomain => {
                                *domain = String::from(qname);
                            }
                            _ => ()
                        }

                        match record.get_domain() {
                            None => {}
                            Some(domain) => {
                                if domain == qname || domain == subdomain {
                                    if record.get_querytype() == QueryType::CNAME {
                                        cname = Some(record.clone());
                                    } else {
                                        answers.push(record.clone());
                                    }
                                }
                            }
                        }
                    }
                }
                if answers.is_empty() && cname.is_some() {
                    answers.push(cname.unwrap());
                }
                let mut domain_exists = !answers.is_empty() || subdomain.is_empty();
                if answers.is_empty() {
                    // If there are no records found we search for *.domain.tld record
                    for mut record in data.records.iter_mut() {
                        let record_domain = record.get_domain().unwrap_or(String::new());
                        if record.get_querytype() == qtype && record_domain == "*" {
                            match &mut record {
                                DnsRecord::A { domain, .. }
                                | DnsRecord::AAAA { domain, .. }
                                | DnsRecord::NS { domain, .. }
                                | DnsRecord::CNAME { domain, .. }
                                | DnsRecord::SRV { domain, .. }
                                | DnsRecord::TLSA { domain, .. }
                                | DnsRecord::MX { domain, .. }
                                | DnsRecord::UNKNOWN { domain, .. }
                                | DnsRecord::SOA { domain, .. }
                                | DnsRecord::TXT { domain, .. } => {
                                    *domain = String::from(qname);
                                }
                                _ => ()
                            }
                            answers.push(record.clone());
                        }
                        if !domain_exists && (record_domain == subdomain || record_domain == "*") {
                            domain_exists = true;
                        }
                    }
                }

                // Extract NS records and GLUE records for the response
                let ns_records = BlockchainFilter::get_ns_records(&data, &top_domain);
                let ns_hosts: Vec<String> = ns_records.iter()
                    .filter_map(|record| {
                        if let DnsRecord::NS { host, .. } = record {
                            Some(host.clone())
                        } else {
                            None
                        }
                    })
                    .collect();
                let glue_records = BlockchainFilter::get_glue_records(&data, &top_domain, &ns_hosts);

                if let Some(mut packet) = self.create_packet(qname, qtype, zone, answers, ns_records, glue_records) {
                    if domain_exists && packet.answers.is_empty() {
                        packet.header.rescode = ResultCode::NOERROR;
                    }
                    return Some(packet);
                }
            }
        }

        None
    }
}

#[cfg(test)]
mod tests {
    // TODO write tests for this filter
}