Mutate edns & remove edns options (#1363)

hickory-dns · Jan 21, 2021 · c3c4666 · c3c4666
1 parent 06e4579
commit c3c4666
Show file tree

Hide file tree

Showing 5 changed files with 24 additions and 6 deletions.
diff --git a/bin/tests/server_harness/mut_message_client.rs b/bin/tests/server_harness/mut_message_client.rs
@@ -36,7 +36,7 @@ impl<C: ClientHandle + Unpin> DnsHandle for MutMessageHandle<C> {
             edns.set_dnssec_ok(true);
 
             if let Some(supported_algs) = self.support_algorithms {
-                edns.set_option(EdnsOption::DAU(supported_algs));
+                edns.options_mut().insert(EdnsOption::DAU(supported_algs));
             }
         }
 

diff --git a/crates/proto/src/op/edns.rs b/crates/proto/src/op/edns.rs
@@ -88,6 +88,11 @@ impl Edns {
         &self.options
     }
 
+    /// Returns a mutable options portion of EDNS
+    pub fn options_mut(&mut self) -> &mut OPT {
+        &mut self.options
+    }
+
     /// Set the high order bits for the result code.
     pub fn set_rcode_high(&mut self, rcode_high: u8) {
         self.rcode_high = rcode_high
@@ -110,6 +115,7 @@ impl Edns {
     }
 
     /// Set the specified EDNS option
+    #[deprecated(note = "Please use options_mut().insert() to modify")]
     pub fn set_option(&mut self, option: EdnsOption) {
         self.options.insert(option);
     }
@@ -215,7 +221,8 @@ fn test_encode_decode() {
     edns.set_max_payload(0x8008);
     edns.set_version(0x40);
     edns.set_rcode_high(0x01);
-    edns.set_option(EdnsOption::DAU(SupportedAlgorithms::all()));
+    edns.options_mut()
+        .insert(EdnsOption::DAU(SupportedAlgorithms::all()));
 
     let record: Record = (&edns).into();
     let edns_decode: Edns = (&record).into();
@@ -225,4 +232,10 @@ fn test_encode_decode() {
     assert_eq!(edns.version(), edns_decode.version());
     assert_eq!(edns.rcode_high(), edns_decode.rcode_high());
     assert_eq!(edns.options(), edns_decode.options());
+
+    // re-insert and remove using mut
+    edns.options_mut()
+        .insert(EdnsOption::DAU(SupportedAlgorithms::all()));
+    edns.options_mut().remove(EdnsCode::DAU);
+    assert!(edns.option(EdnsCode::DAU).is_none());
 }
diff --git a/crates/proto/src/rr/rdata/opt.rs b/crates/proto/src/rr/rdata/opt.rs
@@ -195,6 +195,11 @@ impl OPT {
     pub fn insert(&mut self, option: EdnsOption) {
         self.options.insert((&option).into(), option);
     }
+
+    /// Remove an option, the key is derived from the `EdnsOption`
+    pub fn remove(&mut self, option: EdnsCode) {
+        self.options.remove(&option);
+    }
 }
 
 /// Read the RData from the given Decoder

diff --git a/crates/proto/src/xfer/dnssec_dns_handle.rs b/crates/proto/src/xfer/dnssec_dns_handle.rs
@@ -152,8 +152,8 @@ where
                 let dau = EdnsOption::DAU(algorithms);
                 let dhu = EdnsOption::DHU(algorithms);
 
-                edns.set_option(dau);
-                edns.set_option(dhu);
+                edns.options_mut().insert(dau);
+                edns.options_mut().insert(dhu);
             }
 
             request.set_authentic_data(true);

diff --git a/crates/server/src/authority/catalog.rs b/crates/server/src/authority/catalog.rs
@@ -59,8 +59,8 @@ fn send_response<R: ResponseHandler>(
         let dau = EdnsOption::DAU(algorithms);
         let dhu = EdnsOption::DHU(algorithms);
 
-        resp_edns.set_option(dau);
-        resp_edns.set_option(dhu);
+        resp_edns.options_mut().insert(dau);
+        resp_edns.options_mut().insert(dhu);
 
         response.set_edns(resp_edns);
     }