kapiti 0.0.3

The Kapiti DNS Server
Documentation
use std::net::IpAddr;
use std::sync::{Arc, Mutex};

use anyhow::{bail, Result};
use bytes::BytesMut;
use lazy_static::lazy_static;
use tracing::{debug, level_enabled, trace, Level};

use crate::codec::{decoder::DNSMessageDecoder, encoder::ENCODER, message::RequestInfo};
use crate::filter::{filter::Filter, reader};
use crate::resolver::Resolver;
use crate::specs::enums_generated;
use crate::specs::message::{IntEnum, Message, Question, OPT};

lazy_static! {
    /// "File name" name use in filtered info responses about hardcoded targets
    static ref HARDCODED_SOURCE_NAME: String = "hardcoded".to_string();
}

/// Deserializes client requests and checks local filters before forwarding the request to a Resolver.
/// This is shared among external-facing clients, but is not used for internal host resolution.
/// This structure means that:
/// - external clients have both cache and filtering
/// - internal clients have cache but not filtering
pub struct Lookup {
    resolver: Resolver,
    filter: Arc<Mutex<Filter>>,
}

impl Lookup {
    pub fn new(resolver: Resolver, filter: Arc<Mutex<Filter>>) -> Lookup {
        Lookup { resolver, filter }
    }

    /// Receives and handles a single query provided by packet_buffer,
    /// and passes back the response again via packet_buffer.
    /// Responses may be processed locally by filters, or by remote upstream servers.
    pub async fn handle_query(self: &mut Lookup, packet_buffer: &mut BytesMut) -> Result<()> {
        // Decode the request message so that we can see what it's querying for
        if let Some(mut request) = DNSMessageDecoder::new().decode(packet_buffer)? {
            debug!("Incoming request: {}", request);

            // TODO implement OPTOptions server support here: COOKIE (#21), NSID (#35)
            //      (and figure out how to include them correctly in responses...)

            if let Some(opt) = &mut request.opt {
                // Per RFC6891, any OPT Options that we do not support should be removed.
                // We do this here to ensure that OPT Options are not included in:
                // - filter responses back to the client where they might be nonsensical (just below)
                // - upstream queries to other servers where they might change results or pollute the cache
                opt.option.clear();
            }

            if let Some((question, request_info)) = get_question(&request)? {
                if self.check_filter(packet_buffer, &request, question, &request_info)? {
                    // Filter hit: response written to packet_buffer
                    return Ok(());
                }

                // Reuse packet_buffer for resolver response.
                packet_buffer.clear();
                // Forward request to resolver, which will check cache and then upstream server(s).
                trace!(
                    "No filter entry found for {}, querying cache/upstreams",
                    request_info.name
                );
                self.resolver
                    .resolve_message(&request, &request_info, packet_buffer)
                    .await
            } else {
                bail!("Missing question in request");
            }
        } else {
            bail!("Failed to parse incomplete request");
        }
    }

    /// Checks internal filters for a result matching the request.
    /// If a match is found, a simulated DNS response is written to packet_buffer and Ok(true) is returned.
    fn check_filter(
        self: &mut Lookup,
        packet_buffer: &mut BytesMut,
        request: &Message,
        question: &Question,
        request_info: &RequestInfo,
    ) -> Result<bool> {
        let filter_result: Option<(String, reader::FilterEntry)>;
        match self.filter.lock() {
            Err(e) => bail!("Failed to lock query filter: {:?}", e),
            Ok(filter_locked) => {
                // Hold the lock as briefly as possible, these clones should be cheap
                filter_result =
                    filter_locked
                        .check(&request_info.name)
                        .map(|(file_info, filter_entry)| {
                            if let Some(f) = file_info {
                                (f.source_path.clone(), (*filter_entry).clone())
                            } else {
                                (HARDCODED_SOURCE_NAME.clone(), (*filter_entry).clone())
                            }
                        });
            }
        }
        if let Some((file_source_path, entry)) = filter_result {
            // If the filter rule says to allow the (sub)domain, then act as if no filter result was found
            // (e.g. allow 'good.foo.com' while blocking 'foo.com')
            if entry.dest_upstream {
                return Ok(false);
            }

            // Filter had a match (block or override), write filtered response to packet_buffer.
            packet_buffer.clear();
            write_filter_response(
                packet_buffer,
                &request_info,
                question,
                &request.opt,
                &file_source_path,
                &entry,
            )?;
            if level_enabled!(Level::TRACE) {
                if let Some(response) = DNSMessageDecoder::new().decode(packet_buffer)? {
                    trace!("Returning response from filter: {}", response);
                } else {
                    // Shouldn't happen for our own local data, implies parser bug
                    bail!(
                        "Failed to re-parse response from filter ({}b): {:02X?}",
                        packet_buffer.len(),
                        &packet_buffer[..]
                    );
                }
            }
            // Filter hit
            Ok(true)
        } else {
            // Filter miss
            Ok(false)
        }
    }
}

/// Writes the response to `packet_buffer` based on the destination info in the retrieved filter entry.
/// The response itself can either be NXDOMAIN for a blocked host, or NOERROR for an override.
/// The response also includes:
/// - A copy of the original request OPT, if any was provided
/// - An additional TXT record with filter info (name, linenum if any)
fn write_filter_response(
    packet_buffer: &mut BytesMut,
    request_info: &RequestInfo,
    question: &Question,
    opt: &Option<OPT>,
    filter_source: &String,
    entry: &reader::FilterEntry,
) -> Result<()> {
    // Filter debug info is included in the response as an additional TXT record.
    // Hardcoded filters don't have a line_num
    let filter_info = match entry.line_num {
        Some(line_num) => format!("filter={}:{}", filter_source, line_num),
        None => format!("filter={}", filter_source.to_string()),
    };
    if let (None, None) = (entry.dest_ipv4, entry.dest_ipv6) {
        // Return blocked domain
        debug!(
            "Got block entry for {} from {} line {:?}: dest=NXDOMAIN",
            request_info.name, filter_source, entry.line_num
        );
        ENCODER.encode_local_response(
            enums_generated::ResponseCode::NXDOMAIN,
            request_info.received_request_id,
            question,
            opt,
            &filter_info,
            None,
            Some(request_info.requested_udp_size),
            packet_buffer,
        )
    } else if request_info.resource_type == enums_generated::ResourceType::A {
        // Return configured IPv4/A override
        debug!(
            "Got override {:?} entry for {} from {} line {:?}: dest={:?}",
            request_info.resource_type,
            request_info.name,
            filter_source,
            entry.line_num,
            entry.dest_ipv4
        );
        ENCODER.encode_local_response(
            enums_generated::ResponseCode::NOERROR,
            request_info.received_request_id,
            question,
            opt,
            &filter_info,
            entry.dest_ipv4.map(|ip| IpAddr::V4(ip)),
            Some(request_info.requested_udp_size),
            packet_buffer,
        )
    } else if request_info.resource_type == enums_generated::ResourceType::AAAA {
        // Return configured IPv6/AAAA override
        debug!(
            "Got override {:?} entry for {} from {} line {:?}: dest={:?}",
            request_info.resource_type,
            request_info.name,
            filter_source,
            entry.line_num,
            entry.dest_ipv6
        );
        ENCODER.encode_local_response(
            enums_generated::ResponseCode::NOERROR,
            request_info.received_request_id,
            question,
            opt,
            &filter_info,
            entry.dest_ipv6.map(|ip| IpAddr::V6(ip)),
            Some(request_info.requested_udp_size),
            packet_buffer,
        )
    } else {
        // 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 upstream if this happens.
        // But if you had an upstream server with the right information, why would you be putting custom host entries locally?
        // Therefore we explicitly do NOT support misc record types like MX and SRV for hostnames that have an override entry.
        debug!(
            "Got override {:?} entry for {} from {} line {:?}: dest=NONE",
            request_info.resource_type, request_info.name, filter_source, entry.line_num
        );
        ENCODER.encode_local_response(
            enums_generated::ResponseCode::NOERROR,
            request_info.received_request_id,
            question,
            opt,
            &filter_info,
            None,
            Some(request_info.requested_udp_size),
            packet_buffer,
        )
    }
}

pub fn get_question<'a>(request: &'a Message) -> Result<Option<(&'a Question, RequestInfo)>> {
    let dnssec_ok = request.opt.as_ref().map_or(false, |opt| opt.dnssec_ok);

    for question in &request.question {
        if question.resource_class != IntEnum::Enum(enums_generated::ResourceClass::INTERNET) {
            continue;
        }

        if let IntEnum::Enum(resource_type) = question.resource_type {
            // Remove trailing '.': Filters do not include trailing '.'
            let mut name = question.name.to_string();
            if !name.is_empty() {
                name.pop();
            }

            let request_id = request.header.id;
            return Ok(Some((
                question,
                RequestInfo {
                    name,
                    resource_type,
                    dnssec_ok,
                    received_request_id: request_id,
                    // For the response UDP size, lets just return whatever the client sent...
                    requested_udp_size: request
                        .opt
                        .as_ref()
                        .map(|opt| opt.udp_size)
                        .unwrap_or(4096),
                },
            )));
        }
    }
    Ok(None)
}