#![deny(warnings)]
use std::net::IpAddr;
use anyhow::{Context, Result};
use bytes::BytesMut;
use packed_struct::prelude::*;
use crate::codec::{domain_name, message, rdata};
use crate::fbs::dns_enums_generated::{ResourceType, ResponseCode};
use crate::fbs::dns_message_generated::{Message, Question, OPT};
/// Encodes DNS `Message`s to bytes to be sent over the wire.
pub struct DNSMessageEncoder {}
/// TTL of NXDOMAIN responses due to a domain being blocked.
/// We don't want this to be too long, in case the domain is later unblocked.
static CACHED_NXDOMAIN_TTL: u32 = 300; // 5 minutes
impl DNSMessageEncoder {
/// Returns a `DNSMessageEncoder` for encoding DNS messages.
pub fn new() -> DNSMessageEncoder {
DNSMessageEncoder {}
}
/// Encodes a filter-based DNS response containing the provided overrides.
/// The response could be e.g. an NXDOMAIN value for a blocked domain, or a custom A record for an overridden domain.
pub fn encode_local_response<'a>(
&self,
response_code: ResponseCode,
orig_header_id: u16,
orig_question: &Question<'a>,
orig_opt: Option<OPT<'a>>,
ip: Option<IpAddr>,
udp_size_override: Option<u16>,
buf: &mut BytesMut,
) -> Result<()> {
// Cache of string offsets for use in name compression
let mut ptr_offsets = domain_name::LabelOffsets::new();
message::write_header_bits(
message::HeaderBits {
id: orig_header_id,
is_response: true,
op_code: Integer::from(0 /*QUERY*/),
authoritative: false,
truncated: false,
recursion_desired: true,
recursion_available: true,
reserved_9: false,
authentic_data: false,
checking_disabled: false,
response_code: Integer::from(response_code as u8),
question_count: 1,
answer_count: match ip {
Some(_) => 1, // A or AAA resource
None => 0,
},
authority_count: match ip {
Some(_) => 0,
None => 1, // SOA resource
},
additional_count: match orig_opt {
Some(_) => 1,
None => 0,
},
},
buf,
)?;
message::write_question(orig_question, buf, &mut ptr_offsets)?;
// The name should really be there - how else would we know that we wanted to filter for it?
let question_name = orig_question
.name()
.with_context(|| "missing resource.name")?;
// Write either an A/AAAA Answer resource or an SOA Authority resource, depending on ip
match ip {
Some(ip) => message::write_resource_fields(
&message::ResourceFields {
name: question_name,
resource_type: orig_question.resource_type(),
resource_class: orig_question.resource_class(),
ttl: 300,
rdata: message::RDataFields::IP(ip),
},
buf,
&mut ptr_offsets,
)?,
None => {
let rname = format!("kapiti.{}", question_name);
message::write_resource_fields(
&message::ResourceFields {
name: question_name,
resource_type: ResourceType::TYPE_SOA as u16,
resource_class: orig_question.resource_class(),
ttl: CACHED_NXDOMAIN_TTL,
rdata: match ip {
Some(ip) => message::RDataFields::IP(ip),
None => message::RDataFields::SOA(rdata::SOAFields {
mname: question_name,
rname: rname.as_str(),
serial: 42069,
refresh: 6 * 60 * 60,
retry: 60 * 60,
expire: 3 * 24 * 60 * 60,
minimum: 5 * 60,
}),
},
},
buf,
&mut ptr_offsets,
)?;
}
}
// Ensure OPT resource is written within the Additional section
if let Some(opt) = orig_opt {
message::write_opt(&opt, udp_size_override, buf)?;
}
Ok(())
}
/// Encodes a cached DNS response containing the provided overrides.
pub fn encode_cached_response<'a>(
&self,
message: &Message,
orig_header_id: u16,
ttl_subtract: u32,
udp_size_override: Option<u16>,
buf: &mut BytesMut,
) -> Result<()> {
// Cache of string offsets for use in name compression
let mut ptr_offsets = domain_name::LabelOffsets::new();
message::write_header_id(message, orig_header_id, buf)?;
if let Some(resources) = message.question() {
for i in 0..resources.len() {
message::write_question(&resources.get(i), buf, &mut ptr_offsets)?;
}
}
if let Some(resources) = message.answer() {
for i in 0..resources.len() {
message::write_resource_ttl_subtract(
&resources.get(i),
ttl_subtract,
buf,
&mut ptr_offsets,
)?;
}
}
if let Some(resources) = message.authority() {
for i in 0..resources.len() {
message::write_resource_ttl_subtract(
&resources.get(i),
ttl_subtract,
buf,
&mut ptr_offsets,
)?;
}
}
// Ensure OPT resource is written within the Additional section
if let Some(opt) = message.opt() {
message::write_opt(&opt, udp_size_override, buf)?;
}
if let Some(resources) = message.additional() {
for i in 0..resources.len() {
message::write_resource_ttl_subtract(
&resources.get(i),
ttl_subtract,
buf,
&mut ptr_offsets,
)?;
}
}
Ok(())
}
/// Encodes a response from the provided flatbuffer representation
pub fn encode(
&self,
message: &Message,
udp_size_override: Option<u16>,
buf: &mut BytesMut,
) -> Result<()> {
// Cache of string offsets for use in name compression
let mut ptr_offsets = domain_name::LabelOffsets::new();
message::write_header_id(
message,
message
.header()
.with_context(|| "required header is missing")?
.id(),
buf,
)?;
if let Some(resources) = message.question() {
for i in 0..resources.len() {
message::write_question(&resources.get(i), buf, &mut ptr_offsets)?;
}
}
if let Some(resources) = message.answer() {
for i in 0..resources.len() {
message::write_resource(&resources.get(i), buf, &mut ptr_offsets)?;
}
}
if let Some(resources) = message.authority() {
for i in 0..resources.len() {
message::write_resource(&resources.get(i), buf, &mut ptr_offsets)?;
}
}
// Ensure OPT resource is written within the Additional section
if let Some(opt) = message.opt() {
message::write_opt(&opt, udp_size_override, buf)?;
}
if let Some(resources) = message.additional() {
for i in 0..resources.len() {
message::write_resource(&resources.get(i), buf, &mut ptr_offsets)?;
}
}
Ok(())
}
}