~nickbp/kapiti

025082b1e3b359a5343d5977bcdd37a3f57c3e0c — Nick Parker 1 year, 8 months ago 225b392 optcookie
Basic initial structure for OPT COOKIE support (#21)
3 files changed, 114 insertions(+), 2 deletions(-)

M src/codec/mod.rs
A src/codec/optcookie.rs
M src/server.rs
M src/codec/mod.rs => src/codec/mod.rs +3 -0
@@ 19,4 19,7 @@ pub mod domain_name;

pub mod message;

/// Encodes and decodes OPT option entries
pub mod optcookie;

mod rdata;

A src/codec/optcookie.rs => src/codec/optcookie.rs +106 -0
@@ 0,0 1,106 @@
#![deny(warnings, rust_2018_idioms)]

use std::convert::TryInto;
use std::io;
use std::net::IpAddr;

use sha2::{Digest, Sha256};

use crate::fbs::dns_enums_conv;
use crate::fbs::dns_enums_generated::OPTOptionCode;
use crate::fbs::dns_message_generated::OPT;

macro_rules! io_err {
    ($($arg:tt)*) => (Err(io::Error::new(io::ErrorKind::InvalidData, format!($($arg)+))))
}

// TODO server:
// - generator function that accepts client ip and client cookie and returns new response cookie
// - validator function that returns None (no cookie in request), a malformed request error (wrong number bytes), a cookie validation error (caller may optionally call generator to get correct server cookie), or OK (return request cookie in response as-is): Option<Result<bool, io::Error>>?


pub fn server_generate(client_addr: &IpAddr, client_cookie: &[u8; 8]) -> [u8; 32] {
    let sha256_builder = Sha256::default();
    hash_address(client_addr, &mut sha256_builder);
    sha256_builder.input("hello".as_bytes() /*TODO server secret*/);
    sha256_builder.input(client_cookie);
    // return sha256 as-is: server cookie max size is 32 bytes/256 bits
    sha256_builder.result().try_into()
        .expect("failed to convert sha256 slice to 32 byte array")
}

pub fn server_validate<'a>(opt: OPT<'a>, expect_server_cookie: &[u8]) -> Result<Option<bool>, io::Error> {
    if let Some(cookie) = get_cookie_data(opt)? {
        Ok(Some(false)) // TODO return true if cookie == expect_server_cookie
    } else {
        Ok(None)
    }
}

// TODO client:
// - generator function that accepts client ip(?) and server ip
// - updater function (accepts optional client cookie and optional server cookie: allow refreshing client cookie and adding/updating server cookie)
// - validator function that checks client cookie

// Generates a new client cookie for use with querying the server.
pub fn client_generate(client_addr: &IpAddr, server_addr: &IpAddr) -> [u8; 8] {
    let sha256_builder = Sha256::default();
    hash_address(client_addr, &mut sha256_builder);
    hash_address(server_addr, &mut sha256_builder);
    sha256_builder.input("hello".as_bytes() /*TODO client secret*/);
    sha256_builder.result()[0..8].try_into()
        .expect("failed to convert 8 byte slice to 8 byte array")
}

pub fn client_update(cookie: &mut [u8], client_cookie: &[u8; 8]) {
    cookie.copy
}

pub fn server_update(cookie: &mut [u8], server_cookie: &[u8]) {
}

pub fn client_validate<'a>(opt: OPT<'a>, expect_client_cookie: &[u8; 8]) -> Result<Option<bool>, io::Error> {
    if let Some(cookie) = get_cookie_data(opt)? {
        Ok(Some(false)) // TODO return true if first 8 bytes of cookie match expect_client_cookie
    } else {
        Ok(None)
    }
}

fn hash_address(addr: &IpAddr, sha256_builder: &mut Sha256) {
    match addr {
        IpAddr::V4(v4) => sha256_builder.input(v4.octets()),
        IpAddr::V6(v6) => sha256_builder.input(v6.octets()),
    }
}

fn get_cookie_data<'a>(opt: OPT<'a>) -> Result<Option<&'a [u8]>, io::Error> {
    let mut cookie_data = None;
    if let Some(options) = opt.option() {
        for i in 0..options.len() {
            let option = options.get(i);
            if Some(OPTOptionCode::OPTOPTION_COOKIE) == dns_enums_conv::optoptioncode_int(option.code() as usize) {
                if let Some(_existing_data) = cookie_data {
                    return io_err!("Message contains multiple COOKIE entries");
                }
                if let Some(data) = option.data() {
                    // Sanity check data length before returning it
                    match data.len() {
                        8 => {
                            // 8 byte client cookie only
                            cookie_data = Some(data);
                        },
                        16..=40 => {
                            // 8 byte client cookie + 8-32 byte server cookie
                            cookie_data = Some(data);
                        },
                        _ => {
                            return io_err!("OPT Cookie data should be 8 or 16-40 bytes, got {}", data.len());
                        },
                    }
                }
            }
        }
    }
    Ok(cookie_data)
}

M src/server.rs => src/server.rs +5 -2
@@ 9,7 9,7 @@ use log::{debug, log_enabled, trace, Level};
use scopeguard;

use crate::client::client;
use crate::codec::{decoder::DNSMessageDecoder, encoder::DNSMessageEncoder, message};
use crate::codec::{decoder::DNSMessageDecoder, encoder::DNSMessageEncoder, message, optcookie};
use crate::fbs::dns_enums_conv;
use crate::fbs::dns_enums_generated::{ResourceClass, ResourceType, ResponseCode};
use crate::fbs::dns_message_generated::{Message, Question};


@@ 43,7 43,7 @@ pub async fn handle_query<'a>(
) -> Result<(), io::Error> {
    let received_request_id: u16;
    let requested_udp_size: u16;

    let requested_cookie: Option<optcookie::COOKIE>;
    {
        // Reset request_fbb when exiting this block, success or error.
        let mut fbb = scopeguard::guard(&mut m.fbb, |fbb| {


@@ 62,10 62,13 @@ pub async fn handle_query<'a>(
        received_request_id = request.header().expect("missing request header").id();

        // For the response UDP size, lets just return whatever the client sent...
        let requested_cookie: Option<optcookie::COOKIE>;
        if let Some(opt) = request.opt() {
            requested_udp_size = opt.udp_size();
            requested_cookie = optcookie::read_opt_entry_cookie(opt)?;
        } else {
            requested_udp_size = 4096;
            requested_cookie = None;
        }

        if check_local_response(&request, received_request_id, packet_buffer, filter, requested_udp_size)? {