Sample implementation of a header parser for DNS

dns_header.rs

use std::{convert::TryInto, time::Duration};

use nom::IResult;
use nom::{
    bits::{bits, streaming::take as take_bits},
    number::streaming::be_u16,
    sequence::tuple,
    Parser,
};
use nom_supreme::{error::ErrorTree, parser_ext::ParserExt};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QR {
    Query,
    Response,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum OpCode {
    Query,
    IQuery,
    Status,
    Notify,
    Update,
    DSO,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ResponseCode {
    NoError,
    FormatError,
    ServerFailure,
    NameError,
    NotImplemented,
    Refused,
    Other(u8),
}

#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub struct Flags {
    pub qr: QR,
    pub opcode: Option<OpCode>,
    pub authoritative: bool,
    pub truncated: bool,
    pub recursion_desired: bool,
    pub recursion_available: bool,
    pub response_code: ResponseCode,
}

fn parse_bool_bit(
    input: (&[u8], usize),
) -> IResult<(&[u8], usize), bool, ErrorTree<(&[u8], usize)>> {
    take_bits(1usize).map(|bit: u8| bit != 0).parse(input)
}

fn parse_flags(input: &[u8]) -> IResult<&[u8], Flags, ErrorTree<&[u8]>> {
    bits(tuple((
        parse_bool_bit
            .map(|bit| if bit { QR::Response } else { QR::Query })
            .context("QR bit"),
        take_bits(4usize)
            .map(|opcode: u8| match opcode {
                0 => Some(OpCode::Query),
                1 => Some(OpCode::IQuery),
                2 => Some(OpCode::Status),
                4 => Some(OpCode::Notify),
                5 => Some(OpCode::Update),
                6 => Some(OpCode::DSO),
                _ => None,
            })
            .context("opcode"),
        parse_bool_bit.context("authoritative bit"),
        parse_bool_bit.context("truncated bit"),
        parse_bool_bit.context("recursion desired bit"),
        parse_bool_bit.context("recursion available bit"),
        take_bits(3usize).map(|_: u8| ()).context("reserved bits"),
        take_bits(4usize)
            .map(|rcode: u8| match rcode {
                0 => ResponseCode::NoError,
                1 => ResponseCode::FormatError,
                2 => ResponseCode::ServerFailure,
                3 => ResponseCode::NameError,
                4 => ResponseCode::NotImplemented,
                5 => ResponseCode::Refused,
                code => ResponseCode::Other(code),
            })
            .context("response code"),
    )))
    .map(
        |(
            qr,
            opcode,
            authoritative,
            truncated,
            recursion_desired,
            recursion_available,
            (),
            response_code,
        )| Flags {
            qr,
            opcode,
            authoritative,
            truncated,
            recursion_desired,
            recursion_available,
            response_code,
        },
    )
    .parse(input)
}

pub struct Header {
    id: u16,

    flags: Flags,

    question_count: u16,
    answer_count: u16,
    ns_count: u16,
    ar_count: u16,
}

fn parse_header(input: &[u8]) -> IResult<&[u8], Header, ErrorTree<&[u8]>> {
    tuple((
        be_u16.context("id"),
        parse_flags.context("flags"),
        be_u16.context("question count"),
        be_u16.context("answer count"),
        be_u16.context("ns count"),
        be_u16.context("ar count"),
    ))
    .map(
        |(id, flags, question_count, answer_count, ns_count, ar_count)| Header {
            id,
            flags,
            question_count,
            answer_count,
            ns_count,
            ar_count,
        },
    )
    .parse(input)
}