From c3c466641654a365d6bda1595a7fa303df69f0b0 Mon Sep 17 00:00:00 2001
From: Evan Cameron <cameron.evan@gmail.com>
Date: Thu, 21 Jan 2021 08:22:49 -0500
Subject: [PATCH] Mutate edns & remove edns options (#1363)

---
 bin/tests/server_harness/mut_message_client.rs |  2 +-
 crates/proto/src/op/edns.rs                    | 15 ++++++++++++++-
 crates/proto/src/rr/rdata/opt.rs               |  5 +++++
 crates/proto/src/xfer/dnssec_dns_handle.rs     |  4 ++--
 crates/server/src/authority/catalog.rs         |  4 ++--
 5 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/bin/tests/server_harness/mut_message_client.rs b/bin/tests/server_harness/mut_message_client.rs
index abe54a87f0..6545e425bc 100644
--- 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
index 0b9503cf3e..5ec3b1a611 100644
--- 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
index 75c60c0c12..153e5f2bdf 100644
--- 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
index 6a2c3333db..2eb0c9404a 100644
--- 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
index 3ea5b8cc36..5048ef0b08 100644
--- 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);
     }