From 11569bbe03e45902bea17c26ad5d4b4a65ec0020 Mon Sep 17 00:00:00 2001
From: Evan Cameron <cameron.evan@gmail.com>
Date: Wed, 20 Jan 2021 13:56:41 -0500
Subject: [PATCH] proto: add mut and remove for Edns

---
 crates/proto/src/op/edns.rs      | 20 ++++++++++++++++++++
 crates/proto/src/rr/rdata/opt.rs |  5 +++++
 2 files changed, 25 insertions(+)

diff --git a/crates/proto/src/op/edns.rs b/crates/proto/src/op/edns.rs
index bae7e9f6a7..4016a40f83 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
@@ -113,6 +118,11 @@ impl Edns {
     pub fn set_option(&mut self, option: EdnsOption) {
         self.options.insert(option);
     }
+
+    /// Remove the option specified by the `EdnsCode`
+    pub fn remove_option(&mut self, option: EdnsCode) {
+        self.options.remove(option);
+    }
 }
 
 impl<'a> From<&'a Record> for Edns {
@@ -225,4 +235,14 @@ 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());
+
+    // remove option
+    edns.remove_option(EdnsCode::DAU);
+    assert!(edns.option(EdnsCode::DAU).is_none());
+
+    // re-insert and remove r 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 f08c51ff57..90cfa6cc9f 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