// Copyright 2015-2022 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://apache.org/licenses/LICENSE-2.0> or the MIT license <LICENSE-MIT or
// http://opensource.org/licenses/MIT>, at your option. This file may not be
// copied, modified, or distributed except according to those terms.

//! Integration tests for the recursor. These integration tests setup scenarios to verify that the recursor is able
//!  to recursively resolve various real world scenarios. As new scenarios are discovered, they should be added here.

use std::future::Future;
use std::net::*;
use std::pin::Pin;
use std::str::FromStr as _;
use std::sync::Arc;
use std::time::Instant;

use futures::stream::{self, Stream};
use futures::future;
use futures::StreamExt as _;
use futures::executor::block_on;

use hickory_client::op::{Message, MessageType, Query};
use hickory_client::rr::{rdata, Name, RData, Record, RecordType};
use hickory_integration::mock_client::*;
use hickory_proto::DnsHandle;
use hickory_proto::error::{ProtoError, ProtoErrorKind};
use hickory_proto::rr::LowerName;
use hickory_proto::serialize::txt::Parser;
use hickory_proto::xfer::{DnsRequest, DnsResponse};
use hickory_recursor::{ErrorKind, Recursor};
use hickory_resolver::config::*;
use hickory_resolver::error::ResolveErrorKind;
use hickory_resolver::name_server::ConnectionProvider;
use hickory_resolver::name_server::{NameServer, NameServerPool};
use hickory_server::authority::{Catalog, LookupError, LookupOptions, ZoneType};
use hickory_server::store::in_memory::InMemoryAuthority;


const NS_ROOT: IpAddr = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 1));
const NS_TLDS: IpAddr = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 2));
const NS_IANA: IpAddr = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 3));
const NS_ORG:  IpAddr = IpAddr::V4(Ipv4Addr::new(10, 0, 0, 4));


const ZONE_ROOT: &str = r#"
@                   IN SOA  a.root-servers.net. nstld.verisign-grs.com. (
                            20     ; SERIAL
                            7200   ; REFRESH
                            600    ; RETRY
                            3600000; EXPIRE
                            60)    ; MINIMUM

.                   NS      a.root-servers.net.
a.root-servers.net  A       10.0.0.1

com                 NS      a.gtld-servers.net.
a.gtld-servers.net  A       10.0.0.2

net                 NS      a.gtld-servers.net.

org                 NS      b0.org.afilias-nst.org.
b0.org.afilias-nst.org. A   10.0.0.4

us                  NS      b.cctld.us.
"#;

const ZONE_TLDS: &str = r#"
@                   IN SOA  av4.nstld.com. nstld.verisign-grs.com. (
                            20     ; SERIAL
                            7200   ; REFRESH
                            600    ; RETRY
                            3600000; EXPIRE
                            60)    ; MINIMUM

com.                TXT     my-text-record
example.com.        NS      a.iana-servers.net.
example.org.        NS      a.iana-servers.net.
iana-servers.net.   NS      a.iana-servers.net.
a.iana-servers.net. A       10.0.0.3
"#;


const ZONE_IANA: &str = r#"
@   IN SOA  a.iana-servers.net. noc.dns.icann.org. (
                            20     ; SERIAL
                            7200   ; REFRESH
                            600    ; RETRY
                            3600000; EXPIRE
                            60)    ; MINIMUM

a.iana-servers.net. A       10.0.0.3
"#;

const ZONE_EXAMPLE_COM: &str = r#"
@   IN SOA  a.iana-servers.net. noc.dns.icann.org. (
                            20     ; SERIAL
                            7200   ; REFRESH
                            600    ; RETRY
                            3600000; EXPIRE
                            60)    ; MINIMUM

example.com.        A       10.0.100.1
www.example.com.    A       10.0.100.1
cname.sub.example.com. CNAME www.example.com.
cname.example.com.     CNAME www.example.com.
example-org-cname.example.com. CNAME example.org.
inline-org-cname.example.com.  CNAME inline.org.
double-cname.sub.example.com.  CNAME inline-org-cname.example.org.
cycle-a.example.com.           CNAME cycle-b.org.
intractable-cname.example.com. CNAME nonexistent.org.
"#;

const ZONE_ORG: &str = r#"
@   IN SOA  a0.org.afilias-nst.info. hostmaster.donuts.email. 20 7200 600 3600000 60

example.org.        NS      a.iana-servers.net.
inline.org.         A       10.0.100.3
cycle-b.org.        CNAME   cycle-a.example.com.
"#;

const ZONE_EXAMPLE_ORG: &str = r#"
@   IN SOA  a.iana-servers.net. noc.dns.icann.org. (
                            20     ; SERIAL
                            7200   ; REFRESH
                            600    ; RETRY
                            3600000; EXPIRE
                            60)    ; MINIMUM

example.org.        A       10.0.100.2
inline-org-cname.example.org. CNAME inline.org.
"#;

type HardcodedNameServer = NameServer<HardcodedConnProvider>;

/// A ConnectionProvider which includes hard-coded DNS test data shared across tests.
#[derive(Clone, Default)]
struct HardcodedConnProvider;
impl ConnectionProvider for HardcodedConnProvider {
    type Conn = HardcodedDnsHandle;
    type FutureConn = Pin<Box<dyn Send + Future<Output = Result<Self::Conn, ProtoError>>>>;
    type RuntimeProvider = MockRuntimeProvider;

    fn new_connection(
        &self,
        config: &NameServerConfig,
        _options: &ResolverOpts,
    ) -> Self::FutureConn {
        println!("HardcodedConnProvider::new_connection");
        Box::pin(future::ok(HardcodedDnsHandle(config.socket_addr)))
    }
}

#[derive(Clone, Copy)]
struct HardcodedDnsHandle(SocketAddr);

impl DnsHandle for HardcodedDnsHandle {
    type Response = Pin<Box<dyn Stream<Item = Result<DnsResponse, ProtoError>> + Send>>;

    fn send<R: Into<DnsRequest>>(&self, request: R) -> Self::Response {
        let addr = self.0;
        let request = request.into();
        // let queries: Vec<_> = request.queries().iter().collect();
        println!("HardcodedDnsHandle::send request: {:?}", request.queries());
        let response_futures: Vec<_> = request.queries().iter().map(move |query| {
            let query = query.clone();
            stream::once(async move {
                println!("HardcodedDnsHandle: fielding query: {:?} {:?}", query.name(), query.query_type());
                let mut response = Message::new();
                response.add_query(query.clone());
                response.set_message_type(MessageType::Response);
                let query_lower = LowerName::new(query.name());
                let catalog = make_catalog_for(addr);
                let authority = catalog.find(&query_lower);
                if let Some(authority) = authority {
                    let lookup = authority.lookup(
                        &query_lower,
                        query.query_type(),
                        LookupOptions::default(),
                    ).await;
                    // println!("HardcodeDnsHandle: result {:?}", lookup);
                    match lookup {
                        Ok(mut lookup_obj) => {
                            response.add_answers(lookup_obj.iter().cloned());
                            if let Some(additionals_obj) = lookup_obj.take_additionals() {
                                response.add_additionals(additionals_obj.iter().cloned());
                            }
                        },
                        Err(LookupError::ResponseCode(code)) => {
                            response.set_response_code(code);
                        },
                        Err(LookupError::NameExists) => {
                            if !query.query_type().is_soa() {
                                println!("nx_domain and not soa, lookup: {:?} SOA", authority.origin());
                                let soa_lookup = authority.lookup(
                                    authority.origin(),
                                    RecordType::SOA,
                                    LookupOptions::default(),
                                ).await;
                                // println!("HardcodeDnsHandle: revised result {:?}", lookup);
                                if let Ok(r) = soa_lookup {
                                    response.add_name_servers(r.iter().cloned());
                                }
                            }
                        },
                        _ => (),
                    }
                }
                // response.set_header(*response.header().clone()
                //     .set_query_count(response.queries().len() as _)
                //     .set_answer_count(response.answers().len() as _)
                //     .set_name_server_count(response.name_servers().len() as _)
                //     .set_additional_count(response.additionals().len() as _)
                // );
                let resp = DnsResponse::from_message(response).unwrap();
                Ok(resp)
            })
        }).collect();

        Box::pin(stream::iter(response_futures).flatten())
    }
}


fn make_catalog_for(nameserver: SocketAddr) -> Catalog {
    println!("retrieving catalog for {:?}...", nameserver);
    let zones = match nameserver {
        s if s == SocketAddr::new(NS_ROOT, 53) => vec![
            (".", ZONE_ROOT),
        ],
        s if s == SocketAddr::new(NS_TLDS, 53) => vec![
            (".", ZONE_TLDS),
        ],
        s if s == SocketAddr::new(NS_IANA, 53) => vec![
            ("iana-servers.net.", ZONE_IANA),
            ("example.com.", ZONE_EXAMPLE_COM),
            ("example.org.", ZONE_EXAMPLE_ORG),
        ],
        s if s == SocketAddr::new(NS_ORG,  53) => vec![
            ("org.", ZONE_ORG),
        ],
        ns => panic!("unexpected nameserver {:?}", ns),
    };
    let mut catalog = Catalog::new();
    for (name, zone_text) in zones {
        let (origin, records) = Parser::new(zone_text, None, Some(Name::from_str(name).unwrap())).parse().unwrap();
        let authority = InMemoryAuthority::new(origin.clone(), records, ZoneType::Primary, false /* allow_axfr */).unwrap();
        catalog.upsert(
            LowerName::new(&origin),
            Box::new(Arc::new(authority)),
        )
    }
    catalog
}

pub fn logger(level: &str) {
    use tracing_subscriber::layer::SubscriberExt;
    use tracing_subscriber::util::SubscriberInitExt;
    // Setup tracing for logging based on input
    let subscriber = tracing_subscriber::EnvFilter::builder()
        .with_default_directive(tracing::metadata::LevelFilter::OFF.into())
        .parse(level)
        .expect("failed to configure tracing/logging");

    let formatter = tracing_subscriber::fmt::layer().compact();

    let _ = tracing_subscriber::registry()
        .with(formatter)
        .with(subscriber)
        .try_init();
}

fn mock_nameserver(addr: IpAddr) -> HardcodedNameServer {
    NameServer::new(
        NameServerConfig {
            socket_addr: SocketAddr::new(addr, 53),
            protocol: Protocol::Udp,
            tls_dns_name: None,
            trust_negative_responses: false,
            #[cfg(any(feature = "dns-over-rustls", feature = "dns-over-https-rustls"))]
            tls_config: None,
            bind_addr: None,
        },
        ResolverOpts::default(),
        HardcodedConnProvider,
    )
}

/// Tests a basic recursive resolution `a.recursive.test.` , `.` -> `test.` -> `recursive.test.` -> `a.recursive.test.`
///
/// There are three authorities needed for this test `.` which contains the `test` nameserver, `recursive.test` which is
///  target zone containing `a.recursive.test.`.
#[test]
fn test_basic_recursion() {
    // TBD
}

/// Query the root for its nameserver. This is a single DNS request, no recursion necessary.
#[test]
fn test_root_ns() {
    logger("DEBUG");
    let ns_query = Query::query(Name::from_str(".").unwrap(), RecordType::NS);
    let ns_record = ns_record(Name::from_str(".").unwrap(), Name::from_str("a.root-servers.net.").unwrap());

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(ns_query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &[ns_record]);
}

/// Query a top-level domain for its nameserver. `.` (NS) -> `com.` (NS).
/// This requires two DNS requests/iterations, but doesn't actually cross NS boundaries.
#[test]
fn test_tld_ns() {
    logger("DEBUG");

    let ns_query = Query::query(Name::from_str("com.").unwrap(), RecordType::NS);
    let com_ns_record = ns_record(Name::from_str("com.").unwrap(), Name::from_str("a.gtld-servers.net").unwrap());

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(ns_query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &[com_ns_record]);
}

/// Query a top-level domain for a non-NS record. `.` (NS) -> `com.` (NS) -> `com.` (TXT)
/// This requires two DNS requests/iterations, but doesn't actually cross NS boundaries.
#[test]
fn test_tld_txt() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("com.").unwrap(), RecordType::TXT);
    let expected_record = Record::from_rdata(
        Name::from_str("com.").unwrap(),
        86400,
        RData::TXT(rdata::TXT::new(vec!["my-text-record".to_string()])),
    );

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &[expected_record]);
}

#[test]
fn test_v4_domain() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("example.com.").unwrap(), RecordType::A);
    let expected_record = v4_record(
        Name::from_str("example.com.").unwrap(),
        Ipv4Addr::new(10, 0, 100, 1),
    );

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &[expected_record]);
}

/// Query a top-level domain for a record several layers deep.
/// `.` (NS) -> `com.` (NS) -> `example.com.` (NS) -> `www.example.com.` (NS, NXDomain), `www.example.com` (A)
#[test]
fn test_recurse() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("www.example.com.").unwrap(), RecordType::A);
    let expected_record = v4_record(
        Name::from_str("www.example.com.").unwrap(),
        Ipv4Addr::new(10, 0, 100, 1),
    );

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &[expected_record]);
}

/// If `example.com.` is a NS and `sub.example.com.` doesn't exist, then `cname.sub.example.com.`
/// triggers a different codepath than `cname.example.com.`. The recursor tries to query the
/// latter as a NS (and receives a CNAME); the former is queried directly as a CNAME.
#[test]
fn test_cname_below_nonexistent_parent() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("cname.sub.example.com.").unwrap(), RecordType::CNAME);
    let expected_record = cname_record(
        Name::from_str("cname.sub.example.com.").unwrap(),
        Name::from_str("www.example.com.").unwrap(),
    );

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &[expected_record]);
}

/// `.` (NS) -> `com.` (NS) -> `example.com.` (NS) -> `cname.example.com.`.
/// This results in a lookup for `cname.example.com. NS`, even though the record is actually CNAME.
#[test]
fn test_cname_under_ns() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("cname.example.com.").unwrap(), RecordType::CNAME);
    let expected_record = cname_record(
        Name::from_str("cname.example.com.").unwrap(),
        Name::from_str("www.example.com.").unwrap(),
    );

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &[expected_record]);
}

/// Query for an A record which is actually a CNAME record.
/// The server should answer with CNAME, and (since the target is in the same zone) an additional
/// A record.
#[test]
fn test_cname_queried_as_v4() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("cname.example.com.").unwrap(), RecordType::A);
    let expected_records = [
        cname_record(
            Name::from_str("cname.example.com.").unwrap(),
            Name::from_str("www.example.com.").unwrap(),
        ),
        v4_record(
            Name::from_str("www.example.com.").unwrap(),
            Ipv4Addr::new(10, 0, 100, 1),
        ),
    ];

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &expected_records);
}

/// Follow a CNAME from one zone to another, but where the target is still served by the same NS.
#[test]
fn test_cname_to_other_zone() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("example-org-cname.example.com.").unwrap(), RecordType::A);
    let expected_records = [
        cname_record(
            Name::from_str("example-org-cname.example.com.").unwrap(),
            Name::from_str("example.org.").unwrap(),
        ),
        v4_record(
            Name::from_str("example.org.").unwrap(),
            Ipv4Addr::new(10, 0, 100, 2),
        ),
    ];

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &expected_records);
}

/// Follow a CNAME from one zone to another zone, served by a different NS.
#[test]
fn test_cname_to_other_ns() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("inline-org-cname.example.com.").unwrap(), RecordType::A);
    let expected_records = [
        cname_record(
            Name::from_str("inline-org-cname.example.com.").unwrap(),
            Name::from_str("inline.org.").unwrap(),
        ),
        v4_record(
            Name::from_str("inline.org.").unwrap(),
            Ipv4Addr::new(10, 0, 100, 3),
        ),
    ];

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &expected_records);
}

/// Follow more than one CNAME to an existing record.
#[test]
fn test_double_cname() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("double-cname.sub.example.com.").unwrap(), RecordType::A);
    let expected_records = [
        cname_record(
            Name::from_str("double-cname.sub.example.com.").unwrap(),
            Name::from_str("inline-org-cname.example.org.").unwrap(),
        ),
        cname_record(
            Name::from_str("inline-org-cname.example.org.").unwrap(),
            Name::from_str("inline.org.").unwrap(),
        ),
        v4_record(
            Name::from_str("inline.org.").unwrap(),
            Ipv4Addr::new(10, 0, 100, 3),
        ),
    ];

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup = block_on(recursor.resolve(query, now)).unwrap();

    assert_eq!(&*lookup.records().to_vec(), &expected_records);
}

/// Test that the recursor does not get stuck when presented with an infinite cycle of CNAMEs.
#[test]
fn test_cname_cycle() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("cycle-a.example.com.").unwrap(), RecordType::A);
    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup_err = block_on(recursor.resolve(query, now)).unwrap_err();

    assert!(matches!(lookup_err.kind(), ErrorKind::Timeout));
}

/// Test that the recursor responds with an error when following CNAMEs to non-existent records,
/// rather than (for example) returning partial results.
#[test]
fn test_cname_to_nonexistent_record() {
    logger("DEBUG");

    let query = Query::query(Name::from_str("intractable-cname.example.com.").unwrap(), RecordType::A);

    let roots = NameServerPool::from_nameservers(
        Default::default(),
        vec![mock_nameserver(NS_ROOT)],
        vec![],
    );
    let recursor = Recursor::new_with_pool(roots, 1024, 1048576).unwrap();

    let now = Instant::now();
    let lookup_err = block_on(recursor.resolve(query, now)).unwrap_err();

    assert!(matches!(lookup_err.kind(),
        ErrorKind::Resolve(resolve_error) if matches!(resolve_error.kind(),
            ResolveErrorKind::Proto(proto_error) if matches!(proto_error.kind(),
                ProtoErrorKind::NoRecordsFound { .. }
            )
        )
    ));
}
