Merge branch 'develop' into resolve-state-block

Cargo.toml

@@ -16,6 +16,11 @@ tokio-stream = "0.1"
 thiserror = "1.0.20"
 futures-util = "0.3.28"
 async-trait = "0.1.77"
+sha2 = "0.10.2"
+hmac = "0.12.1"
+rust-crypto = "0.2"
+base64 = "0.22.1"
+
 lru = "0.12.3"
 
 [lib]

src/async_resolver.rs

@@ -19,10 +19,13 @@ use crate::message::rcode::Rcode;
 use crate::message::rdata::Rdata;
 use crate::message::resource_record::ResourceRecord;
 use crate::message::rrtype::Rrtype;
+use crate::tsig;
 use crate::message::{self, DnsMessage};
 use crate::resolver_cache::ResolverCache;
 use std::net::IpAddr;
+use std::time::SystemTime;
 use std::sync::{Arc, Mutex};
+use std::vec;
 
 /// Asynchronous resolver for DNS queries.
 ///
@@ -105,10 +108,13 @@ impl AsyncResolver {
         let response = self
             .inner_lookup(domain_name_struct, Rrtype::A, rclass.into())
             .await;
-
+        
         return self
             .check_error_from_msg(response)
             .and_then(|lookup_response| {
+                if lookup_response.to_dns_msg().get_header().get_tc() {
+                    self.config.set_protocol(ConnectionProtocol::TCP);
+                }
                 let rrs_iter = lookup_response.to_vec_of_rr().into_iter();
                 let ip_addresses: Result<Vec<IpAddr>, _> = rrs_iter
                     .map(|rr| AsyncResolver::from_rr_to_ip(rr))
@@ -168,6 +174,11 @@ impl AsyncResolver {
                 Rclass::from(rclass),
             )
             .await;
+
+
+        if self.config.get_tsig(){
+            self.verify_tsig(response.clone());
+        }
 
         return self.check_error_from_msg(response);
     }
@@ -316,13 +327,13 @@ impl AsyncResolver {
     /// answer section, it is always preferred.
     fn store_data_cache(&self, response: DnsMessage) {
         let truncated = response.get_header().get_tc();
-        let rcode = response.get_header().get_rcode();
         {
             let mut cache = self.cache.lock().unwrap();
             cache.timeout();
             if !truncated {
                 cache.add(response.clone());
             }
+
         }
     }
 
@@ -417,6 +428,44 @@ impl AsyncResolver {
             _ => Err(ClientError::ResponseError(rcode.into()))?,
         }
     }
+
+    /// Verifies tsig
+    pub fn verify_tsig(&self, 
+        response: Result<LookupResponse, ResolverError>)
+        -> Result<LookupResponse, ClientError> {
+        let lookup_response = match response {
+            Ok(val) => Ok(val),
+            Err(_) => Err(ClientError::TemporaryError("no DNS message found")),
+        };
+
+        let dns_response = lookup_response.unwrap().to_dns_msg();
+
+        let key_bytes = self.config.get_key();
+        let shared_key_name = self.config.get_key_name();
+        let time = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs();
+        let algorithm = String::from(self.config.get_algorithm());
+        let alg_list = vec![(algorithm, true)];
+        let mac = dns_response.get_mac();
+
+        let (_val, rcode) = tsig::process_tsig(
+            &dns_response,
+            &key_bytes,
+            shared_key_name.clone().unwrap(),
+            time,
+            alg_list,
+            mac,
+        );
+
+        match rcode {
+            Rcode::NOERROR => Ok(LookupResponse::new(dns_response)),
+            Rcode::FORMERR => Err(ClientError::FormatError("The name server was unable to interpret the query."))?,
+            Rcode::SERVFAIL => Err(ClientError::ServerFailure("The name server was unable to process this query due to a problem with the name server."))?,
+            Rcode::NXDOMAIN => Err(ClientError::NameError("The domain name referenced in the query does not exist."))?,
+            Rcode::NOTIMP => Err(ClientError::NotImplemented("The name server does not support the requested kind of query."))?,
+            Rcode::REFUSED => Err(ClientError::Refused("The name server refuses to perform the specified operation for policy reasons."))?,
+            _ => Err(ClientError::ResponseError(rcode.into()))?,
+        }
+    }
 }
 
 // Getters
@@ -2009,7 +2058,6 @@ mod async_resolver_test {
 
         resolver.save_negative_answers(dns_response.clone());
 
-        let rrtype_search = Rrtype::A;
         assert_eq!(dns_response.get_answer().len(), 0);
         assert_eq!(dns_response.get_additional().len(), 1);
         assert_eq!(

src/async_resolver/config.rs

@@ -1,8 +1,8 @@
 use crate::client::{udp_connection::ClientUDPConnection, tcp_connection::ClientTCPConnection,client_connection::ClientConnection };
 use crate::client::client_connection::ConnectionProtocol;
 use crate::message::DnsMessage;
+use crate::tsig::tsig_algorithm::TsigAlgorithm;
 use std::cmp::max;
-use std::option;
 use std::{net::{IpAddr,SocketAddr,Ipv4Addr}, time::Duration};
 
 use super::server_info::ServerInfo;
@@ -73,6 +73,14 @@ pub struct ResolverConfig {
     ends0_flags: u16,
     /// edns0 options for the resolver.
     ends0_options: Vec<u16>,
+    /// This is whether tsig is enabled or not.
+    tsig: bool,
+    /// This is the tsig keyname for the resolver.
+    key_name: Option<String>,
+    /// This is the tsig key for the resolver.
+    key: Vec<u8>,
+    /// algorithm for the tsig key
+    algorithm: TsigAlgorithm,
 }
 
 impl ResolverConfig {
@@ -113,6 +121,10 @@ impl ResolverConfig {
             ends0_version: 0,
             ends0_flags: 0,
             ends0_options: Vec::new(),
+            tsig: false,
+            key_name: None,
+            key: Vec::new(),
+            algorithm: TsigAlgorithm::HmacSha256,
         };
         resolver_config
     }
@@ -154,6 +166,10 @@ impl ResolverConfig {
             ends0_version: 0,
             ends0_flags: 0,
             ends0_options: Vec::new(),
+            tsig: false,
+            key_name: None,
+            key: Vec::new(),
+            algorithm: TsigAlgorithm::HmacSha256,
         };
         resolver_config
     }
@@ -245,6 +261,38 @@ impl ResolverConfig {
             message.add_edns0(Some(self.get_max_payload()), self.get_ends0_version(), self.get_ends0_flags(), Some(self.get_ends0_options()));
         }
     }
+
+    /// add tsig to the resolver
+    /// 
+    /// # Examples
+    /// ```
+    /// let mut resolver_config = ResolverConfig::default();
+    /// resolver_config.add_tsig("keyname".to_string(), b"key".to_vec(), Some(TsigAlgorithm::HmacSha256));
+    /// ```
+    pub fn add_tsig(&mut self, key_name: String, key: Vec<u8>, algorithm: Option<TsigAlgorithm>) {
+        self.tsig = true;
+        self.key_name = Some(key_name);
+        self.key = key;
+        if let Some(algorithm) = algorithm {
+            self.algorithm = algorithm;
+        }
+    }
+
+    /// add tsig from the resolver to a dns message
+    /// 
+    /// # Examples
+    /// ```
+    /// let mut resolver_config = ResolverConfig::default();
+    /// resolver_config.add_tsig("keyname".to_string(), b"key".to_vec(), Some(TsigAlgorithm::HmacSha256));
+    /// let message = Message::new();
+    /// resolver_config.add_tsig_to_message(&message, 300, vec![]);
+    /// ```
+    pub fn add_tsig_to_message(&self, message: &mut DnsMessage, fudge: Option<u16>, mac_request: Option<Vec<u8>>) {
+        if self.tsig {
+            message.add_tsig(self.key.clone(), self.algorithm.clone(), 
+            fudge.unwrap_or(300), self.key_name.clone(), mac_request.unwrap_or(Vec::new()));
+        }
+    }
 }
 
 ///Getters
@@ -317,6 +365,22 @@ impl ResolverConfig {
     pub fn get_ends0_options(&self) -> Vec<u16> {
         self.ends0_options.clone()
     }
+
+    pub fn get_tsig(&self) -> bool {
+        self.tsig
+    }
+
+    pub fn get_key_name(&self) -> Option<String> {
+        self.key_name.clone()
+    }
+
+    pub fn get_key(&self) -> Vec<u8> {
+        self.key.clone()
+    }
+
+    pub fn get_algorithm(&self) -> TsigAlgorithm {
+        self.algorithm.clone()
+    }
 }
 
 ///Setters
@@ -389,6 +453,22 @@ impl ResolverConfig{
     pub fn set_ends0_options(&mut self, ends0_options: Vec<u16>) {
         self.ends0_options = ends0_options;
     }
+
+    pub fn set_tsig(&mut self, tsig: bool) {
+        self.tsig = tsig;
+    }
+
+    pub fn set_key_name(&mut self, key_name: Option<String>) {
+        self.key_name = key_name;
+    }
+
+    pub fn set_key(&mut self, key: Vec<u8>) {
+        self.key = key;
+    }
+
+    pub fn set_algorithm(&mut self, algorithm: TsigAlgorithm) {
+        self.algorithm = algorithm;
+    }
 }
 
 

src/async_resolver/server_info.rs

@@ -12,6 +12,8 @@ pub struct ServerInfo {
     ip_addr: IpAddr,
     //The port of the server.
     port: u16,
+    //Tsig is enabled.
+    tsig: bool,
     //The key of the server.
     key: String,
     // The algorithm of the server.
@@ -29,6 +31,7 @@ impl ServerInfo {
         ServerInfo {
             ip_addr,
             port,
+            tsig: false,
             key,
             algorithm,
             udp_connection,
@@ -43,6 +46,7 @@ impl ServerInfo {
         ServerInfo {
             ip_addr,
             port,
+            tsig: false,
             key,
             algorithm,
             udp_connection,
@@ -59,13 +63,24 @@ impl ServerInfo {
         ServerInfo {
             ip_addr,
             port,
+            tsig: false,
             key,
             algorithm,
             udp_connection,
             tcp_connection,
         }
     }
 
+    /// Function to enable tsig.
+    pub fn enable_tsig(&mut self) {
+        self.tsig = true;
+    }
+
+    /// Function to disable tsig.
+    pub fn disable_tsig(&mut self) {
+        self.tsig = false;
+    }
+
     /// Implements get_ip_address
     /// Returns IpAddr.
     pub fn get_ip_addr(&self) -> IpAddr {
@@ -88,6 +103,11 @@ impl ServerInfo {
         self.port = port;
     }
 
+    /// Get the tsig of the server.
+    pub fn get_tsig(&self) -> bool {
+        self.tsig
+    }
+
     /// Get the key of the server.
     pub fn get_key(&self) -> &str {
         &self.key