// SPDX-FileCopyrightText: 2022 Matteo Settenvini <matteo.settenvini@montecristosoftware.eu>
// SPDX-License-Identifier: GPL-3.0-or-later

use {
    crate::helpers::{set_if_valid, write_record_name_to_buf, write_vector_to_buf},
    crate::nss_bindings::{gaih_addrtuple, nss_status, HErrno},
    crate::policy_checker::POLICY_CHECKER,
    libc::{c_char, c_int, hostent, size_t, AF_INET, AF_INET6},
    nix::errno,
    nix::errno::Errno,
    std::ffi::CStr,
    std::mem::{align_of, discriminant, size_of},
    std::sync::Arc,
    trust_dns_proto::rr::record_type::RecordType,
    trust_dns_proto::rr::{RData, Record},
    trust_dns_proto::xfer::dns_request::DnsRequestOptions,
    trust_dns_resolver::TokioAsyncResolver,
    trust_dns_resolver::{lookup::Lookup, lookup_ip::LookupIp},
};

// See https://support.mozilla.org/en-US/kb/configuring-networks-disable-dns-over-https
const CANARY_HOSTNAME: &str = "use-application-dns.net";

#[derive(Debug)]
pub struct Args {
    pub name: *const c_char,
    pub family: c_int,
    pub buffer: *mut c_char,
    pub buflen: size_t,
    pub errnop: *mut Errno,
    pub h_errnop: *mut HErrno,
    pub ttlp: *mut i32,
    pub result: Result,
}

#[derive(Debug)]
pub enum Result {
    V3(*mut hostent),
    V4(*mut *mut gaih_addrtuple),
}

pub enum DnsResult {
    NxDomain,
    Found,
}

pub async unsafe fn with(args: &mut Args) -> nss_status {
    set_if_valid(args.errnop, errno::from_i32(0));
    set_if_valid(args.h_errnop, HErrno::Success);

    match POLICY_CHECKER.resolver(None).await {
        Ok(None) => {
            // no restrictions for user, the next NSS module will decide
            nss_status::NSS_STATUS_NOTFOUND
        }
        Ok(Some(resolver)) => resolve_hostname_through(resolver, args).await,
        Err(err) => {
            log::error!("{}", err);
            nss_status::NSS_STATUS_UNAVAIL
        }
    }
}

async unsafe fn resolve_hostname_through(
    resolver: Arc<TokioAsyncResolver>,
    args: &mut Args,
) -> nss_status {
    let name = match CStr::from_ptr(args.name).to_str() {
        Ok(name) => name,
        Err(_) => {
            set_if_valid(args.errnop, Errno::EINVAL);
            set_if_valid(args.h_errnop, HErrno::Internal);
            return nss_status::NSS_STATUS_TRYAGAIN;
        }
    };

    if name == CANARY_HOSTNAME {
        set_if_valid(args.h_errnop, HErrno::HostNotFound);
        return nss_status::NSS_STATUS_SUCCESS;
    }

    let lookup: std::result::Result<Lookup, _> = match args.family {
        libc::AF_UNSPEC => resolver.lookup_ip(name).await.map(LookupIp::into),
        libc::AF_INET => {
            resolver
                .lookup(name, RecordType::A, DnsRequestOptions::default())
                .await
        }
        libc::AF_INET6 => {
            resolver
                .lookup(name, RecordType::AAAA, DnsRequestOptions::default())
                .await
        }
        _ => return nss_status::NSS_STATUS_NOTFOUND,
    };

    match lookup {
        Ok(result) => prepare_response(args, result),
        Err(err) => {
            log::warn!("{}", err);
            nss_status::NSS_STATUS_UNAVAIL
        }
    }
}

unsafe fn prepare_response(
    args: &mut Args,
    lookup: trust_dns_resolver::lookup::Lookup,
) -> nss_status {
    if lookup.iter().peekable().peek().is_none() {
        set_if_valid(args.h_errnop, HErrno::HostNotFound);
        return nss_status::NSS_STATUS_SUCCESS;
    }

    let ttl = lookup
        .valid_until()
        .duration_since(std::time::Instant::now())
        .as_secs();
    set_if_valid(
        args.ttlp,
        if ttl < (i32::MAX as u64) {
            ttl as i32
        } else {
            i32::MAX
        },
    );

    let buf = std::slice::from_raw_parts_mut(args.buffer as *mut u8, args.buflen);
    let ret = match &mut args.result {
        Result::V3(hostent) => {
            if hostent.is_null() {
                set_if_valid(args.errnop, Errno::EINVAL);
                set_if_valid(args.h_errnop, HErrno::Internal);
                return nss_status::NSS_STATUS_TRYAGAIN;
            }

            match records_to_hostent(lookup.into(), &mut **hostent, buf) {
                Ok(DnsResult::Found) => nss_status::NSS_STATUS_SUCCESS,
                Ok(DnsResult::NxDomain) => {
                    set_if_valid(args.h_errnop, HErrno::HostNotFound);
                    nss_status::NSS_STATUS_SUCCESS
                }
                Err(err) => {
                    set_if_valid(
                        args.errnop,
                        err.raw_os_error()
                            .map(Errno::from_i32)
                            .unwrap_or(Errno::EAGAIN),
                    );
                    set_if_valid(args.h_errnop, HErrno::Internal);
                    nss_status::NSS_STATUS_TRYAGAIN
                }
            }
        }
        Result::V4(pat) => {
            if pat.is_null() {
                set_if_valid(args.errnop, Errno::EINVAL);
                set_if_valid(args.h_errnop, HErrno::Internal);
                return nss_status::NSS_STATUS_TRYAGAIN;
            }

            match records_to_gaih_addr(lookup.into(), buf) {
                Ok(addrs) => {
                    // DEBUG: eprintln!("{:?} => {:?}", addrs, *addrs);
                    **pat = addrs;
                    nss_status::NSS_STATUS_SUCCESS
                }
                Err(err) => {
                    set_if_valid(
                        args.errnop,
                        err.raw_os_error()
                            .map(Errno::from_i32)
                            .unwrap_or(Errno::EAGAIN),
                    );
                    set_if_valid(args.h_errnop, HErrno::Internal);
                    nss_status::NSS_STATUS_TRYAGAIN
                }
            }
        }
    };

    ret
}

// TODO: refactor to be less ugly
pub unsafe fn records_to_hostent(
    lookup: Lookup,
    host: &mut hostent,
    buf: &mut [u8],
) -> std::io::Result<DnsResult> {
    // In C struct hostent:
    //
    // - for the type of queries we perform, we can assume h_aliases
    //   is an empty array.
    // - hostent is limited to just one address type for all addresses
    //   in the list. We pick the type of first result, and only
    //   append addresses of the same type.

    let first_record = match lookup.record_iter().peekable().peek() {
        Some(record) => *record,
        None => return Ok(DnsResult::NxDomain),
    };

    let first_rdata = match first_record.data() {
        Some(rdata) => rdata,
        None => return Err(Errno::ENODATA.into()),
    };

    // First put the name in the buffer
    let (buf, name_dest) = write_record_name_to_buf(first_record, buf)?;

    // Then the empty aliases array (just a null pointer)
    let (mut buf, aliases_ptr) = write_vector_to_buf(buf, vec![std::ptr::null_mut()])?;

    // Then the address list
    let offset = buf.as_ptr().align_offset(align_of::<*mut c_char>());
    buf = &mut buf[offset..];

    let records = lookup
        .record_iter()
        .flat_map(Record::data)
        .filter(|r| discriminant(*r) == discriminant(first_rdata));

    let mut addresses = vec![];
    for record in records {
        let offset = buf.as_ptr().align_offset(align_of::<*mut c_char>());
        buf = &mut buf[offset..];
        addresses.push(buf.as_mut_ptr());

        let l = match record {
            RData::A(addr) => {
                let octets = addr.octets();
                let l = octets.len();
                if buf.len() < l {
                    return Err(Errno::ERANGE.into());
                }
                buf[..l].copy_from_slice(&octets);
                l
            }
            RData::AAAA(addr) => {
                let octets = addr.octets();
                let l = octets.len();
                if buf.len() < l {
                    return Err(Errno::ERANGE.into());
                }
                buf[..l].copy_from_slice(&octets);
                l
            }
            _ => unreachable!(),
        };

        buf = &mut buf[l..];
    }

    addresses.push(std::ptr::null_mut());
    let (_, addr_list) = write_vector_to_buf(buf, addresses)?;

    // Finally, populate the hostent structure
    host.h_name = name_dest;
    host.h_aliases = aliases_ptr;
    host.h_addr_list = addr_list as *mut *mut c_char;
    match first_rdata {
        RData::A(_) => {
            host.h_addrtype = AF_INET;
            host.h_length = 4;
        }
        RData::AAAA(_) => {
            host.h_addrtype = AF_INET6;
            host.h_length = 16;
        }
        _ => return Err(Errno::EBADMSG.into()),
    }

    Ok(DnsResult::Found)
}

// TODO: error handling codes chosen a bit sloppily
pub unsafe fn records_to_gaih_addr(
    lookup: Lookup,
    mut buf: &mut [u8],
) -> std::io::Result<*mut gaih_addrtuple> {
    const GAIH_ADDRTUPLE_SZ: usize = size_of::<gaih_addrtuple>();

    let mut ret = std::ptr::null_mut();
    let mut prev_link: *mut *mut gaih_addrtuple = std::ptr::null_mut();

    let mut name_dest;
    for record in lookup.record_iter() {
        // First add the name to the buffer
        (buf, name_dest) = write_record_name_to_buf(record, buf)?;

        // Then add the tuple with the address
        let offset = buf.as_ptr().align_offset(align_of::<gaih_addrtuple>());
        let l = GAIH_ADDRTUPLE_SZ;
        if buf.len() < offset + l {
            return Err(Errno::ERANGE.into());
        }

        let tuple = &mut *(buf.as_mut_ptr().add(offset) as *mut gaih_addrtuple);
        tuple.next = std::ptr::null_mut();
        tuple.name = name_dest;
        tuple.scopeid = 0; // how to set tuple.scopeid correctly??? Use a custom trust_dns_resolver::ConnectionProvider?
        set_if_valid(prev_link, &mut *tuple); // link from previous tuple to this tuple
        prev_link = &mut (*tuple).next;

        match record.data() {
            Some(RData::A(addr)) => {
                tuple.family = AF_INET;
                tuple.addr[0] = std::mem::transmute_copy(&addr.octets());
            }
            Some(RData::AAAA(addr)) => {
                tuple.family = AF_INET6;
                tuple.addr = std::mem::transmute_copy(&addr.octets());
            }
            Some(_) => return Err(Errno::EBADMSG.into()),
            None => return Err(Errno::ENODATA.into()),
        }

        if ret == std::ptr::null_mut() {
            ret = tuple;
        }

        buf = &mut buf[(offset + l)..];
    }

    Ok(ret)
}