From aa80f21c608c567afd99924dd2533ac27223992a Mon Sep 17 00:00:00 2001
From: Litr0 <gonzalo.sobarzoa@gmail.com>
Date: Tue, 6 Aug 2024 16:40:06 -0400
Subject: [PATCH] add tsig rr form resolver to dns message

---
 src/async_resolver/config.rs | 16 ++++++++++++++++
 src/message.rs               |  4 ++--
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/src/async_resolver/config.rs b/src/async_resolver/config.rs
index 7fd739be..d4f20c21 100644
--- a/src/async_resolver/config.rs
+++ b/src/async_resolver/config.rs
@@ -3,6 +3,7 @@ use crate::client::client_connection::ConnectionProtocol;
 use crate::message::DnsMessage;
 use crate::tsig::tsig_algorithm::TsigAlgorithm;
 use std::cmp::max;
+use std::io::SeekFrom;
 use std::option;
 use std::{net::{IpAddr,SocketAddr,Ipv4Addr}, time::Duration};
 
@@ -278,6 +279,21 @@ impl ResolverConfig {
             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: u16, mac_request: Vec<u8>) {
+        if self.tsig {
+            message.add_tsig(self.key.clone(), self.algorithm.clone(), fudge, self.key_name.clone(), mac_request);
+        }
+    }
 }
 
 ///Getters
diff --git a/src/message.rs b/src/message.rs
index 85e91690..baf2874b 100644
--- a/src/message.rs
+++ b/src/message.rs
@@ -302,11 +302,11 @@ impl DnsMessage {
     /// dns_query_message.add_tsig(key, alg_name, fudge, key_name, mac_request);
     /// ```
     pub fn add_tsig(&mut self, key: Vec<u8>, alg_name: TsigAlgorithm, 
-        fudge: u16, key_name: String, mac_request: Vec<u8>) {
+        fudge: u16, key_name: Option<String>, mac_request: Vec<u8>) {
         let message = self;
         let time_signed = SystemTime::now().duration_since(SystemTime::UNIX_EPOCH).unwrap().as_secs();
         tsig::sign_tsig(message, &key, alg_name, 
-                                        fudge, time_signed, key_name, mac_request);
+                                        fudge, time_signed, key_name.unwrap_or("".to_string()), mac_request);
     }
 
     /// Gets the MAC from the TSIG RR.