3 files changed, 114 insertions(+), 2 deletions(-)

M src/codec/mod.rs
A src/codec/optcookie.rs
M src/server.rs

@@ 19,4 19,7 @@ pub mod domain_name;
 
 pub mod message;
 
+/// Encodes and decodes OPT option entries
+pub mod optcookie;
+
 mod rdata;

@@ 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)
+}

@@ 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)? {