kapiti 0.0.3

use std::net::{SocketAddr, ToSocketAddrs};

use anyhow::{Context, Result};
use async_trait::async_trait;
use bytes::BytesMut;
use tracing::debug;

use crate::client::DnsClient;
use crate::codec::message::RequestInfo;
use crate::lookup;
use crate::specs::enums_generated;
use crate::specs::message::*;
use crate::specs::rdata;

pub struct Client {
    _priv: (),
}

/// DNS Client that uses the OS resolver.
impl Client {
    /// Constructs a new `Client` that uses the OS resolver.
    pub fn new() -> Self {
        Client { _priv: () }
    }
}

#[async_trait]
impl DnsClient for Client {
    async fn query(
        &mut self,
        request: &Message,
        _request_buffer: &mut BytesMut,
    ) -> Result<Option<Message>> {
        let (_request_question, request_info) = lookup::get_question(request)?
            .with_context(|| format!("Missing question in request: {}", request))?;

        let name_copy = request_info.name.clone();
        match smol::unblock(move || (name_copy, 0).to_socket_addrs()).await {
            Ok(socket_addrs) => {
                debug!(
                    "System resolver: {} => {:?}",
                    request_info.name, socket_addrs
                );
                // Find a socket addr matching the requested resource type, or None if none is found or if the requested type is not A/AAAA
                let socket_addr = match request_info.resource_type {
                    enums_generated::ResourceType::A => {
                        socket_addrs.filter(|addr| addr.is_ipv4()).next()
                    }
                    enums_generated::ResourceType::AAAA => {
                        socket_addrs.filter(|addr| addr.is_ipv6()).next()
                    }
                    _other => {
                        // Misc record type, for a domain that's got A and/or AAAA overrides: Record not found.
                        // It's a little ambiguous whether we should instead try going to another upstream if this happens.
                        // It might make sense to have a "not supported" resolver return value for situations like this.
                        // But for now, we do not support e.g. an MX record lookup using the system resolver.
                        None
                    }
                };

                Ok(Some(build_response(request_info, socket_addr)))
            }
            Err(e) => Err(e)
                .with_context(|| format!("Failed to query for hostname {}", request_info.name)),
        }
    }
}

fn build_response(mut request_info: RequestInfo, address: Option<SocketAddr>) -> Message {
    // Ensure domain ends with '.' before creating raw DNS message
    request_info.name.push('.');

    let answer_rdata = match address {
        Some(SocketAddr::V4(addr)) => {
            let octets = addr.ip().octets();
            Some(ResourceData::A(rdata::A {
                address1: octets[0],
                address2: octets[1],
                address3: octets[2],
                address4: octets[3],
            }))
        }
        Some(SocketAddr::V6(addr)) => {
            let segments = addr.ip().segments();
            Some(ResourceData::AAAA(rdata::AAAA {
                address1: segments[0],
                address2: segments[1],
                address3: segments[2],
                address4: segments[3],
                address5: segments[4],
                address6: segments[5],
                address7: segments[6],
                address8: segments[7],
            }))
        }
        None => None,
    };

    let mut answer = Vec::new();
    if let Some(rdata) = answer_rdata {
        answer.push(Resource {
            name: request_info.name.clone(),
            resource_type: IntEnum::Enum(request_info.resource_type),
            resource_class: IntEnum::Enum(enums_generated::ResourceClass::INTERNET),
            ttl: 600 as u32, // Let's just go with a 10 minute TTL for system responses
            rdata,
        });
    }

    let mut question = Vec::new();
    question.push(Question {
        name: request_info.name,
        resource_type: IntEnum::Enum(request_info.resource_type),
        resource_class: IntEnum::Enum(enums_generated::ResourceClass::INTERNET),
    });

    Message {
        header: Header {
            id: request_info.received_request_id,
            is_response: true,
            op_code: IntEnum::Enum(enums_generated::OpCode::QUERY),
            authoritative: false,
            truncated: false,
            // let's just assume recursion was desired, and that the client doesn't care if we are wrong
            recursion_desired: true,
            recursion_available: true,
            reserved_9: false,
            authentic_data: true,
            checking_disabled: false,
            response_code: IntEnum::Enum(enums_generated::ResponseCode::NOERROR),
        },
        opt: Some(OPT {
            option: Vec::new(),
            udp_size: 4096,
            response_code: 0,
            version: 0,
            dnssec_ok: true,
        }),
        question,
        answer,
        authority: Vec::new(),
        additional: Vec::new(),
    }
}