kapiti 0.0.3

use anyhow::Result;
use async_trait::async_trait;
use bytes::{BufMut, BytesMut};
use std::io::ErrorKind;

use crate::codec::encoder::ENCODER;
use crate::specs::enums_generated::OPTOptionCode;
use crate::specs::message::{IntEnum, Message, OPTOption, OPT};

static PADDED_REQUEST_BYTES: usize = 468;

#[async_trait]
pub trait DnsClient {
    /// Runs a query for the provided request, then returns the response.
    /// query_buffer may also be used as a scratch pad for handling the request.
    async fn query(
        &mut self,
        request: &Message,
        query_buffer: &mut BytesMut,
    ) -> Result<Option<Message>>;
}

fn is_timeout(kind: ErrorKind) -> bool {
    return match kind {
        ErrorKind::WouldBlock | ErrorKind::TimedOut => true,
        _ => false,
    };
}

/// Inserts an EDNS PADDING entry into the resulting query so that its size is a multiple of PADDED_REQUEST_BYTES.
/// This is specifically for encrypted client protocols where we want to reduce sniffing.
/// See slides: https://www.ietf.org/proceedings/105/slides/slides-105-pearg-encrypted-dns-privacy-a-traffic-analysis-perspective-01
/// Returns the resulting message length.
fn add_request_padding(
    request: &Message,
    max_size: u16,
    query_buffer: &mut BytesMut,
) -> Result<()> {
    // If the request already has padding, then just pass it through and call it good enough.
    // Shouldn't happen in practice since we're only serving UDP/TCP,
    // but nothing prevents a client from including this.
    if let Some(opt) = &request.opt {
        let padding_found = opt
            .option
            .iter()
            .any(|option| option.code == IntEnum::Enum(OPTOptionCode::PADDING));
        if padding_found {
            return Ok(());
        }
    }

    let initial_len = query_buffer.len();

    ENCODER.encode(request, Some(max_size), query_buffer)?;

    // Insert the resulting encoded size of the message into those leading two bytes that we'd reserved
    let message_len = query_buffer.len() - initial_len;

    // We want the padded result to be a multiple of PADDED_REQUEST_BYTES
    // size 32 => padded 468, size 500 => padded 936
    let desired_len = PADDED_REQUEST_BYTES * ((message_len / PADDED_REQUEST_BYTES) + 1);

    let mut request = request.clone();
    if let Some(opt) = &mut request.opt {
        // The OPT option we're adding has 4 bytes of overhead (2 code + 2 length),
        // so remove that from the padding data length.
        let padding_len = desired_len as i32 - message_len as i32 - initial_len as i32 - 4;
        if padding_len <= 0 {
            // No padding needed
            return Ok(());
        }
        opt.option.push(OPTOption {
            code: IntEnum::Enum(OPTOptionCode::PADDING),
            data: vec![0; padding_len as usize],
        });
    } else {
        // Create a plausible OPT entry for the request
        // In practice there should already be an OPT entry, but lets play it safe.

        // The OPT record we're assigning has 10 bytes of overhead (6 header + 2 code + 2 length),
        // so remove that from the padding data length.
        let padding_len = desired_len as i32 - message_len as i32 - initial_len as i32 - 10;
        if padding_len <= 0 {
            // No padding needed
            return Ok(());
        }

        request.opt = Some(OPT {
            option: vec![OPTOption {
                code: IntEnum::Enum(OPTOptionCode::PADDING),
                data: vec![0; padding_len as usize],
            }],
            udp_size: max_size,
            response_code: 0,
            version: 0,
            dnssec_ok: true,
        });
    }

    // Clear the buffer, then preserve any initial margin (e.g. for TCP/TLS)
    query_buffer.clear();
    query_buffer.put_bytes(0, initial_len);

    // Encode the newly padded copy
    ENCODER.encode(&request, Some(max_size), query_buffer)?;

    Ok(())
}

/// Removes any EDNS PADDING entry from the provided response.
/// This is specifically for encrypted client protocols where we want to reduce sniffing.
/// We want to remove the padding so that any UDP clients on our end are not seeing excessive sizes.
/// In theory any client could have PADDING, but it's best practice for encrypted transports.
fn remove_response_padding(response: &mut Message) {
    if let Some(opt) = &mut response.opt {
        opt.option = Vec::from_iter(
            opt.option
                .drain(..)
                .filter(|option| option.code != IntEnum::Enum(OPTOptionCode::PADDING)),
        );
    }
}

/// Parses configured upstream strings into DNS clients
pub mod upstream;

/// Client: DNS over HTTPS (DoH)
pub mod https;
/// Client: Host OS lookup
pub mod system;
/// Client: TCP (fallback for UDP)
pub mod tcp;
/// Client: DNS over TLS (DoT)
pub mod tls;
/// Client: UDP
pub mod udp;