--amend

hickory-dns · Apr 4, 2022 · 184899c · 184899c
1 parent 71892c0
commit 184899c
Show file tree

Hide file tree

Showing 6 changed files with 53 additions and 49 deletions.
diff --git a/bin/tests/server_harness/mut_message_client.rs b/bin/tests/server_harness/mut_message_client.rs
@@ -35,7 +35,7 @@ impl<C: ClientHandle + Unpin> DnsHandle for MutMessageHandle<C> {
         #[cfg(feature = "dnssec")]
         {
             // mutable block
-            let edns = request.edns_mut();
+            let edns = request.edns_mut_or_insert();
             edns.set_dnssec_ok(true);
             edns.options_mut()
                 .insert(EdnsOption::DAU(self.lookup_options.supported_algorithms()));

diff --git a/crates/client/src/client/async_client.rs b/crates/client/src/client/async_client.rs
@@ -15,7 +15,6 @@ use futures_util::ready;
 use futures_util::stream::{Stream, StreamExt};
 use log::debug;
 use rand;
-use trust_dns_proto::op::Edns;
 
 use crate::client::Signer;
 use crate::error::*;
@@ -270,9 +269,9 @@ pub trait ClientHandle: 'static + Clone + DnsHandle<Error = ProtoError> + Send {
         // Extended dns
         if self.is_using_edns() {
             message
-                .set_edns(Edns::new())
-                .edns_mut()
-                .map(|edns| edns.set_max_payload(MAX_PAYLOAD_LEN).set_version(0));
+                .edns_mut_or_insert()
+                .set_max_payload(MAX_PAYLOAD_LEN)
+                .set_version(0);
         }
 
         // add the query

diff --git a/crates/client/src/op/update_message.rs b/crates/client/src/op/update_message.rs
@@ -9,8 +9,6 @@
 
 use std::fmt::Debug;
 
-use trust_dns_proto::op::Edns;
-
 use crate::client::async_client::MAX_PAYLOAD_LEN;
 use crate::op::{Message, MessageType, OpCode, Query};
 use crate::rr::rdata::{NULL, SOA};
@@ -188,9 +186,9 @@ pub fn create(rrset: RecordSet, zone_origin: Name, use_edns: bool) -> Message {
     // Extended dns
     if use_edns {
         message
-            .set_edns(Edns::new())
-            .edns_mut()
-            .map(|edns| edns.set_max_payload(MAX_PAYLOAD_LEN).set_version(0));
+            .edns_mut_or_insert()
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -259,9 +257,9 @@ pub fn append(rrset: RecordSet, zone_origin: Name, must_exist: bool, use_edns: b
     // Extended dns
     if use_edns {
         message
-            .set_edns(Edns::new())
-            .edns_mut()
-            .map(|edns| edns.set_max_payload(MAX_PAYLOAD_LEN).set_version(0));
+            .edns_mut_or_insert()
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -351,9 +349,9 @@ pub fn compare_and_swap(
     // Extended dns
     if use_edns {
         message
-            .set_edns(Edns::new())
-            .edns_mut()
-            .map(|edns| edns.set_max_payload(MAX_PAYLOAD_LEN).set_version(0));
+            .edns_mut_or_insert()
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -422,9 +420,9 @@ pub fn delete_by_rdata(mut rrset: RecordSet, zone_origin: Name, use_edns: bool)
     // Extended dns
     if use_edns {
         message
-            .set_edns(Edns::new())
-            .edns_mut()
-            .map(|edns| edns.set_max_payload(MAX_PAYLOAD_LEN).set_version(0));
+            .edns_mut_or_insert()
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -493,9 +491,9 @@ pub fn delete_rrset(mut record: Record, zone_origin: Name, use_edns: bool) -> Me
     // Extended dns
     if use_edns {
         message
-            .set_edns(Edns::new())
-            .edns_mut()
-            .map(|edns| edns.set_max_payload(MAX_PAYLOAD_LEN).set_version(0));
+            .edns_mut_or_insert()
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -561,9 +559,9 @@ pub fn delete_all(
     // Extended dns
     if use_edns {
         message
-            .set_edns(Edns::new())
-            .edns_mut()
-            .map(|edns| edns.set_max_payload(MAX_PAYLOAD_LEN).set_version(0));
+            .edns_mut_or_insert()
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -607,9 +605,9 @@ pub fn zone_transfer(zone_origin: Name, last_soa: Option<SOA>) -> Message {
     // Extended dns
     {
         message
-            .set_edns(Edns::new())
-            .edns_mut()
-            .map(|edns| edns.set_max_payload(MAX_PAYLOAD_LEN).set_version(0));
+            .edns_mut_or_insert()
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message

diff --git a/crates/proto/src/op/message.rs b/crates/proto/src/op/message.rs
@@ -524,6 +524,14 @@ impl Message {
         self.edns.as_mut()
     }
 
+    /// Optionally returns mutable reference to EDNS section
+    pub fn edns_mut_or_insert(&mut self) -> &mut Edns {
+        if self.edns.is_none() {
+            self.set_edns(Edns::new());
+        }
+        self.edns.as_mut().unwrap()
+    }
+
     /// # Return value
     ///
     /// the max payload value as it's defined in the EDNS section.

diff --git a/crates/proto/src/xfer/dns_handle.rs b/crates/proto/src/xfer/dns_handle.rs
@@ -12,9 +12,9 @@ use futures_util::stream::Stream;
 use log::debug;
 use rand;
 
+use crate::error::*;
 use crate::op::{Message, MessageType, OpCode, Query};
 use crate::xfer::{DnsRequest, DnsRequestOptions, DnsResponse, SerialMessage};
-use crate::{error::*, op::Edns};
 
 // TODO: this should be configurable
 // > An EDNS buffer size of 1232 bytes will avoid fragmentation on nearly all current networks.
@@ -85,9 +85,9 @@ fn build_message(query: Query, options: DnsRequestOptions) -> Message {
     // Extended dns
     if options.use_edns {
         message
-            .set_edns(Edns::new())
-            .edns_mut()
-            .map(|edns| edns.set_max_payload(MAX_PAYLOAD_LEN).set_version(0));
+            .edns_mut_or_insert()
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
     message
 }
diff --git a/crates/proto/src/xfer/dnssec_dns_handle.rs b/crates/proto/src/xfer/dnssec_dns_handle.rs
@@ -137,25 +137,24 @@ where
             // TODO: cache response of the server about understood algorithms
             #[cfg(feature = "dnssec")]
             {
-                if let Some(edns) = request.edns_mut() {
-                    edns.set_dnssec_ok(true);
-
-                    // send along the algorithms which are supported by this handle
-                    let mut algorithms = SupportedAlgorithms::new();
-                    #[cfg(feature = "ring")]
-                    {
-                        algorithms.set(Algorithm::ED25519);
-                    }
-                    algorithms.set(Algorithm::ECDSAP256SHA256);
-                    algorithms.set(Algorithm::ECDSAP384SHA384);
-                    algorithms.set(Algorithm::RSASHA256);
-
-                    let dau = EdnsOption::DAU(algorithms);
-                    let dhu = EdnsOption::DHU(algorithms);
+                let edns = request.edns_mut_or_insert();
+                edns.set_dnssec_ok(true);
 
-                    edns.options_mut().insert(dau);
-                    edns.options_mut().insert(dhu);
+                // send along the algorithms which are supported by this handle
+                let mut algorithms = SupportedAlgorithms::new();
+                #[cfg(feature = "ring")]
+                {
+                    algorithms.set(Algorithm::ED25519);
                 }
+                algorithms.set(Algorithm::ECDSAP256SHA256);
+                algorithms.set(Algorithm::ECDSAP384SHA384);
+                algorithms.set(Algorithm::RSASHA256);
+
+                let dau = EdnsOption::DAU(algorithms);
+                let dhu = EdnsOption::DHU(algorithms);
+
+                edns.options_mut().insert(dau);
+                edns.options_mut().insert(dhu);
             }
 
             request.set_authentic_data(true);