more thorough error handling

dnstp/src/message/question/mod.rs

@@ -137,7 +137,9 @@ pub fn questions_to_bytes(questions: &Vec<DNSQuestion>) -> Vec<u8>
 pub enum QuestionParseError {
     ShortLength(usize),
     QTypeParse(u16),
-    QClassParse(u16)
+    QClassParse(u16),
+    UTF8Parse,
+    URLDecode
 }
 
 pub fn questions_from_bytes(bytes: Vec<u8>, total_questions: u16) -> Result<(Vec<DNSQuestion>, Vec<u8>), QuestionParseError>
@@ -196,18 +198,34 @@ pub fn questions_from_bytes(bytes: Vec<u8>, total_questions: u16) -> Result<(Vec
                             match (two_byte_combine(qtype_1, qtype_2).try_into(),
                                    two_byte_combine(qclass_1, byte).try_into()) {
                                 (Ok(qtype), Ok(qclass)) => {
-                                    questions.push(DNSQuestion {
-                                        qname: decode(String::from_utf8(current_query.clone()).unwrap().as_str()).unwrap().to_string(),
-                                        qtype,
-                                        qclass
-                                    });
-
-                                    current_length = None;
-                                    remaining_length = byte;
-                                    current_query.clear();
-                                    current_qtype = (None, None);
-                                    current_qclass = (None, None);
-                                    trailers_reached = false;
+
+                                    match String::from_utf8(current_query.clone()) {
+                                        Ok(parsed_query) => {
+                                            match decode(parsed_query.as_str())
+                                            {
+                                                Ok(decoded_query) => {
+                                                    questions.push(DNSQuestion {
+                                                        qname: decoded_query.to_string(),
+                                                        qtype,
+                                                        qclass
+                                                    });
+
+                                                    current_length = None;
+                                                    remaining_length = byte;
+                                                    current_query.clear();
+                                                    current_qtype = (None, None);
+                                                    current_qclass = (None, None);
+                                                    trailers_reached = false;
+                                                }
+                                                Err(_) => {
+                                                    return Err(QuestionParseError::URLDecode);
+                                                }
+                                            }
+                                        }
+                                        Err(_) => {
+                                            return Err(QuestionParseError::UTF8Parse);
+                                        }
+                                    }
                                 }
                                 (Err(qtype_e), _) => {
                                     return Err(QuestionParseError::QTypeParse(qtype_e));

dnstp/src/processor/mod.rs

@@ -34,6 +34,12 @@ pub fn print_error(e: MessageParseError, peer: &SocketAddr)
                 QuestionParseError::QClassParse(ce) => {
                     error!("[{}] failed to parse questions of received message, qclass error: [{}]", peer, ce);
                 }
+                QuestionParseError::UTF8Parse => {
+                    error!("[{}] failed to parse questions of received message, failed to UTF-8 parse hostname", peer);
+                }
+                QuestionParseError::URLDecode => {
+                    error!("[{}] failed to parse questions of received message, failed to URL decode hostname", peer);
+                }
             }
         }
         MessageParseError::RecordParse(rp) => {

dnstp/src/processor/request/encryption.rs

@@ -1,10 +1,11 @@
 use std::net::Ipv4Addr;
 use p256::ecdh::EphemeralSecret;
 use crate::clients::Client;
-use crate::crypto::{asym_to_sym_key, fatten_public_key, get_random_asym_pair, get_shared_asym_secret, trim_public_key};
+use crate::crypto::{asym_to_sym_key, get_random_asym_pair, get_shared_asym_secret, trim_public_key};
 use crate::message::{ARdata, DNSMessage, QClass, QType, ResourceRecord};
 use crate::message::record::CnameRdata;
-use crate::string::{append_base_domain_to_key, encode_domain_name, strip_base_domain_from_key};
+use crate::string;
+use crate::string::{append_base_domain_to_key, encode_domain_name};
 
 /// Result of a client's handshake request including server key pair and prepared response
 pub struct KeySwapContext {
@@ -26,17 +27,8 @@ pub fn get_key_request_with_base_domain(base_domain: String) -> (EphemeralSecret
     (private, append_base_domain_to_key(trim_public_key(&public), &base_domain))
 }
 
-/// Extract the client's public key from the DNS message, turn the hostname back into the full fat public key with --- BEGIN KEY --- headers and trailers
-pub fn get_fattened_public_key(key_question: &String) -> (String, String)
-{
-    let public_key = key_question;
-    let (trimmed_public_key, base_domain) = strip_base_domain_from_key(public_key);
-
-    (fatten_public_key(&trimmed_public_key), base_domain)
-}
-
 #[derive(Ord, PartialOrd, Eq, PartialEq, Debug, Copy, Clone)]
-pub enum KeyDecodeError {
+pub enum DecodeKeyRequestError {
     QuestionCount(usize),
     FirstQuestionNotA(QType),
     SecondQuestionNotA(QType),
@@ -46,24 +38,24 @@ pub enum KeyDecodeError {
 /// Take a client's handshake request, process the crypto and prepare a response
 ///
 /// Includes generating a server key pair, using the public key in the response, deriving the shared secret.
-pub fn decode_key_request(message: &DNSMessage) -> Result<KeySwapContext, KeyDecodeError>
+pub fn decode_key_request(message: &DNSMessage) -> Result<KeySwapContext, DecodeKeyRequestError>
 {
     if message.questions.len() == 2 {
 
         if message.questions[0].qtype != QType::A
         {
-            return Err(KeyDecodeError::FirstQuestionNotA(message.questions[0].qtype));
+            return Err(DecodeKeyRequestError::FirstQuestionNotA(message.questions[0].qtype));
         }
 
         let key_question = &message.questions[1];
 
         if key_question.qtype != QType::A
         {
-            return Err(KeyDecodeError::SecondQuestionNotA(key_question.qtype));
+            return Err(DecodeKeyRequestError::SecondQuestionNotA(key_question.qtype));
         }
 
         // key is transmitted wihout --- BEGIN KEY -- header and trailer bits and with '.' instead of new lines
-        let (fattened_public_key, base_domain) = get_fattened_public_key(&key_question.qname);
+        let (fattened_public_key, base_domain) = string::get_fattened_public_key(&key_question.qname);
         // generate the servers public/private key pair
         let (server_private, server_public) = get_random_asym_pair();
 
@@ -116,12 +108,12 @@ pub fn decode_key_request(message: &DNSMessage) -> Result<KeySwapContext, KeyDec
                 });
             }
             Err(_) => {
-                return Err(KeyDecodeError::SharedSecretDerivation);
+                return Err(DecodeKeyRequestError::SharedSecretDerivation);
             }
         }
     }
     else
     {
-        return Err(KeyDecodeError::QuestionCount(message.questions.len()));
+        return Err(DecodeKeyRequestError::QuestionCount(message.questions.len()));
     }
 }

dnstp/src/processor/request/mod.rs

@@ -10,7 +10,7 @@ use crate::message::{DNSMessage, QType};
 use crate::net::{NetworkMessagePtr};
 use crate::message_parser::parse_message;
 use crate::processor::print_error;
-use crate::processor::request::encryption::{decode_key_request, KeyDecodeError};
+use crate::processor::request::encryption::{decode_key_request, DecodeKeyRequestError};
 use crate::{RequestError, send_message};
 
 pub mod encryption;
@@ -137,16 +137,16 @@ impl RequestProcesor {
             }
             Err(e) => {
                 match e {
-                    KeyDecodeError::QuestionCount(qc) => {
+                    DecodeKeyRequestError::QuestionCount(qc) => {
                         error!("[{}] failed to parse public key, wrong question count [{}]", peer, qc);
                     }
-                    KeyDecodeError::FirstQuestionNotA(qtype) => {
+                    DecodeKeyRequestError::FirstQuestionNotA(qtype) => {
                         error!("[{}] failed to parse public key, first question wasn't an A request [{}]", peer, qtype);
                     }
-                    KeyDecodeError::SecondQuestionNotA(qtype) => {
+                    DecodeKeyRequestError::SecondQuestionNotA(qtype) => {
                         error!("[{}] failed to parse public key, second question wasn't an A request [{}]", peer, qtype);
                     }
-                    KeyDecodeError::SharedSecretDerivation => {
+                    DecodeKeyRequestError::SharedSecretDerivation => {
                         error!("[{}] failed to parse public key, failed to derived shared secret", peer);
                     }
                 }

dnstp/src/processor/response/encryption.rs

@@ -2,10 +2,16 @@ use std::sync::{Arc, Mutex};
 use crate::client_crypto_context::ClientCryptoContext;
 use crate::crypto::{asym_to_sym_key, get_shared_asym_secret};
 use crate::message::DNSMessage;
-use crate::processor::request::encryption::get_fattened_public_key;
-use crate::string::decode_domain_name;
+use crate::string::get_fattened_public_key;
+use crate::string::{decode_domain_name, DomainDecodeError};
 
-pub fn decode_key_response(message: &DNSMessage, client_crypto_context: Arc<Mutex<ClientCryptoContext>>)
+#[derive(Ord, PartialOrd, Eq, PartialEq, Debug, Copy, Clone)]
+pub enum DecodeKeyResponseError {
+    DomainDecode(DomainDecodeError),
+    KeyDerivation
+}
+
+pub fn decode_key_response(message: &DNSMessage, client_crypto_context: Arc<Mutex<ClientCryptoContext>>) -> Result<(), DecodeKeyResponseError>
 {
     if message.answer_records.len() == 2 {
         // if message.questions[0].qtype != QType::A
@@ -20,19 +26,30 @@ pub fn decode_key_response(message: &DNSMessage, client_crypto_context: Arc<Mute
         //     return Err(KeyDecodeError::SecondQuestionNotA(key_answer.answer_type));
         // }
 
-        let data_string = decode_domain_name(key_answer.r_data.to_bytes());
-        // key is transmitted wihout --- BEGIN KEY -- header and trailer bits and with '.' instead of new lines
-        let (fattened_public_key, _) = get_fattened_public_key(&data_string);
+        match decode_domain_name(key_answer.r_data.to_bytes())
+        {
+            Ok(domain_name) => {
+                // key is transmitted wihout --- BEGIN KEY -- header and trailer bits and with '.' instead of new lines
+                let (fattened_public_key, _) = get_fattened_public_key(&domain_name);
 
-        let mut context = client_crypto_context.lock().unwrap();
+                let mut context = client_crypto_context.lock().unwrap();
 
-        match get_shared_asym_secret(&context.client_private, &fattened_public_key)
-        {
-            Ok(k) => {
-                context.server_public = Some(fattened_public_key);
-                context.shared_key = Some(asym_to_sym_key(&k));
+                match get_shared_asym_secret(&context.client_private, &fattened_public_key)
+                {
+                    Ok(k) => {
+                        context.server_public = Some(fattened_public_key);
+                        context.shared_key = Some(asym_to_sym_key(&k));
+                    }
+                    Err(_) => {
+                        return Err(DecodeKeyResponseError::KeyDerivation);
+                    }
+                }
+            }
+            Err(e) => {
+                return Err(DecodeKeyResponseError::DomainDecode(e));
             }
-            Err(_) => {}
         }
     }
+
+    Ok(())
 }

dnstp/src/processor/response/mod.rs

@@ -3,12 +3,13 @@ mod encryption;
 use std::sync::{Arc, mpsc, Mutex};
 use std::sync::mpsc::{Receiver, Sender};
 use std::thread;
-use log::info;
+use log::{error, info};
 use crate::client_crypto_context::ClientCryptoContext;
 use crate::net::raw_request::NetworkMessagePtr;
 use crate::message_parser::parse_message;
 use crate::processor::print_error;
-use crate::processor::response::encryption::decode_key_response;
+use crate::processor::response::encryption::{decode_key_response, DecodeKeyResponseError};
+use crate::string::DomainDecodeError;
 
 pub struct ResponseProcesor {
     message_channel: Option<Sender<NetworkMessagePtr>>,
@@ -40,15 +41,37 @@ impl ResponseProcesor {
                     Ok(r) => {
                         info!("received dns message: {:?}", r);
 
-                        decode_key_response(&r, crypto_context.clone());
+                        match decode_key_response(&r, crypto_context.clone())
+                        {
+                            Ok(_) => {
+                                info!("successfully decoded key response from server");
+                            }
+                            Err(e) => {
+                                match e {
+                                    DecodeKeyResponseError::DomainDecode(dd) => {
+                                        match dd {
+                                            DomainDecodeError::UTF8Parse => {
+                                                error!("failed to decode key response from server, failed to UTF-8 parse response");
+                                            }
+                                            DomainDecodeError::URLDecode => {
+                                                error!("failed to decode key response from server, failed to URL decode response");
+                                            }
+                                        }
+                                    }
+                                    DecodeKeyResponseError::KeyDerivation => {
+                                        error!("failed to decode key response from server, key derivation failed");
+                                    }
+                                }
+                            }
+                        }
                     }
                     Err(e) => {
                         print_error(e, &peer);
                     }
                 }
             }
 
-            info!("message processing thread finishing")
+            info!("message processing thread finishing");
         });
     }
 

dnstp/src/string/mod.rs

@@ -4,6 +4,7 @@
 mod tests;
 
 use urlencoding::{decode, encode};
+use crate::crypto::fatten_public_key;
 
 pub fn encode_domain_name(name: &String) -> Vec<u8>
 {
@@ -25,7 +26,13 @@ pub fn encode_domain_name(name: &String) -> Vec<u8>
     ret
 }
 
-pub fn decode_domain_name(name: Vec<u8>) -> String
+#[derive(Ord, PartialOrd, Eq, PartialEq, Debug, Copy, Clone)]
+pub enum DomainDecodeError {
+    UTF8Parse,
+    URLDecode
+}
+
+pub fn decode_domain_name(name: Vec<u8>) -> Result<String, DomainDecodeError>
 {
     let mut full_domain: String = String::new();
     let mut current_query: Vec<u8> = Vec::with_capacity(10);
@@ -42,18 +49,29 @@ pub fn decode_domain_name(name: Vec<u8>) -> String
             Some(_) => {
                 if remaining_length == 0 {
 
-                    let current_part = String::from_utf8(current_query.clone()).unwrap();
-                    let url_decoded = decode(current_part.as_str()).unwrap();
-
-                    full_domain.push_str(&url_decoded.to_string());
-
-                    if char != 0 {
-                        full_domain.push('.');
+                    match String::from_utf8(current_query.clone()) {
+                        Ok(parsed_query) => {
+                            match decode(parsed_query.as_str()) {
+                                Ok(decoded_query) => {
+                                    full_domain.push_str(&decoded_query.to_string());
+
+                                    if char != 0 {
+                                        full_domain.push('.');
+                                    }
+
+                                    current_query.clear();
+                                    current_length = Some(char);
+                                    remaining_length = char;
+                                }
+                                Err(_) => {
+                                    return Err(DomainDecodeError::URLDecode);
+                                }
+                            }
+                        }
+                        Err(_) => {
+                            return Err(DomainDecodeError::UTF8Parse);
+                        }
                     }
-
-                    current_query.clear();
-                    current_length = Some(char);
-                    remaining_length = char;
                 }
                 else {
                     current_query.push(char);
@@ -63,7 +81,7 @@ pub fn decode_domain_name(name: Vec<u8>) -> String
         }
     }
 
-    full_domain
+    Ok(full_domain)
 }
 
 pub fn strip_base_domain_from_key(public_key: &String) -> (String, String)
@@ -86,4 +104,13 @@ pub fn strip_base_domain_from_key(public_key: &String) -> (String, String)
 pub fn append_base_domain_to_key(trimmed_key: String, base_domain: &String) -> String
 {
     vec![trimmed_key, base_domain.to_string()].join(".")
-}
+}
+
+/// Extract the client's public key from the DNS message, turn the hostname back into the full fat public key with --- BEGIN KEY --- headers and trailers
+pub fn get_fattened_public_key(key_question: &String) -> (String, String)
+{
+    let public_key = key_question;
+    let (trimmed_public_key, base_domain) = strip_base_domain_from_key(public_key);
+
+    (fatten_public_key(&trimmed_public_key), base_domain)
+}