deprecate edns methods, add extensions

bin/tests/server_harness/mut_message_client.rs

@@ -1,7 +1,7 @@
-use trust_dns_client::client::*;
 use trust_dns_client::proto::xfer::{DnsHandle, DnsRequest};
 #[cfg(feature = "dnssec")]
 use trust_dns_client::rr::rdata::opt::EdnsOption;
+use trust_dns_client::{client::*, op::Edns};
 use trust_dns_server::authority::LookupOptions;
 
 #[derive(Clone)]
@@ -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.extensions_mut().get_or_insert_with(Edns::new);
             edns.set_dnssec_ok(true);
             edns.options_mut()
                 .insert(EdnsOption::DAU(self.lookup_options.supported_algorithms()));

crates/client/src/client/async_client.rs

@@ -270,9 +270,10 @@ 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));
+                .extensions_mut()
+                .get_or_insert_with(Edns::new)
+                .set_max_payload(MAX_PAYLOAD_LEN)
+                .set_version(0);
         }
 
         // add the query

crates/client/src/op/update_message.rs

@@ -188,9 +188,10 @@ 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));
+            .extensions_mut()
+            .get_or_insert_with(Edns::new)
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -259,9 +260,10 @@ 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));
+            .extensions_mut()
+            .get_or_insert_with(Edns::new)
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -351,9 +353,10 @@ 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));
+            .extensions_mut()
+            .get_or_insert_with(Edns::new)
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -422,9 +425,10 @@ 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));
+            .extensions_mut()
+            .get_or_insert(Edns::new())
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -493,9 +497,10 @@ 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));
+            .extensions_mut()
+            .get_or_insert_with(Edns::new)
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -561,9 +566,10 @@ 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));
+            .extensions_mut()
+            .get_or_insert_with(Edns::new)
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message
@@ -607,9 +613,10 @@ 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));
+            .extensions_mut()
+            .get_or_insert_with(Edns::new)
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
 
     message

crates/proto/src/op/message.rs

@@ -155,6 +155,12 @@ impl Message {
         truncated
     }
 
+    /// Sets the `Header` with provided
+    pub fn set_header(&mut self, header: Header) -> &mut Self {
+        self.header = header;
+        self
+    }
+
     /// see `Header::set_id`
     pub fn set_id(&mut self, id: u16) -> &mut Self {
         self.header.set_id(id);
@@ -515,13 +521,30 @@ impl Message {
     /// # Return value
     ///
     /// Optionally returns a reference to EDNS section
+    #[deprecated(note = "Please use `extensions()`")]
     pub fn edns(&self) -> Option<&Edns> {
         self.edns.as_ref()
     }
 
     /// Optionally returns mutable reference to EDNS section
-    pub fn edns_mut(&mut self) -> Option<&mut Edns> {
-        self.edns.as_mut()
+    #[deprecated(
+        note = "Please use `extensions_mut()`. You can chain `.get_or_insert_with(Edns::new)` to recover original behavior of adding Edns if not present"
+    )]
+    pub fn edns_mut(&mut self) -> &mut Edns {
+        if self.edns.is_none() {
+            self.set_edns(Edns::new());
+        }
+        self.edns.as_mut().unwrap()
+    }
+
+    /// Returns reference of Edns section
+    pub fn extensions(&self) -> &Option<Edns> {
+        &self.edns
+    }
+
+    /// Returns mutable reference of Edns section
+    pub fn extensions_mut(&mut self) -> &mut Option<Edns> {
+        &mut self.edns
     }
 
     /// # Return value
@@ -1104,23 +1127,23 @@ fn test_emit_and_read(message: Message) {
 #[rustfmt::skip]
 fn test_legit_message() {
     let buf: Vec<u8> = vec![
-  0x10,0x00,0x81,0x80, // id = 4096, response, op=query, recursion_desired, recursion_available, no_error
-  0x00,0x01,0x00,0x01, // 1 query, 1 answer,
-  0x00,0x00,0x00,0x00, // 0 namesservers, 0 additional record
-
-  0x03,b'w',b'w',b'w', // query --- www.example.com
-  0x07,b'e',b'x',b'a', //
-  b'm',b'p',b'l',b'e', //
-  0x03,b'c',b'o',b'm', //
-  0x00,                // 0 = endname
-  0x00,0x01,0x00,0x01, // ReordType = A, Class = IN
-
-  0xC0,0x0C,           // name pointer to www.example.com
-  0x00,0x01,0x00,0x01, // RecordType = A, Class = IN
-  0x00,0x00,0x00,0x02, // TTL = 2 seconds
-  0x00,0x04,           // record length = 4 (ipv4 address)
-  0x5D,0xB8,0xD8,0x22, // address = 93.184.216.34
-  ];
+    0x10,0x00,0x81,0x80, // id = 4096, response, op=query, recursion_desired, recursion_available, no_error
+    0x00,0x01,0x00,0x01, // 1 query, 1 answer,
+    0x00,0x00,0x00,0x00, // 0 namesservers, 0 additional record
+
+    0x03,b'w',b'w',b'w', // query --- www.example.com
+    0x07,b'e',b'x',b'a', //
+    b'm',b'p',b'l',b'e', //
+    0x03,b'c',b'o',b'm', //
+    0x00,                // 0 = endname
+    0x00,0x01,0x00,0x01, // ReordType = A, Class = IN
+
+    0xC0,0x0C,           // name pointer to www.example.com
+    0x00,0x01,0x00,0x01, // RecordType = A, Class = IN
+    0x00,0x00,0x00,0x02, // TTL = 2 seconds
+    0x00,0x04,           // record length = 4 (ipv4 address)
+    0x5D,0xB8,0xD8,0x22, // address = 93.184.216.34
+    ];
 
     let mut decoder = BinDecoder::new(&buf);
     let message = Message::read(&mut decoder).unwrap();

crates/proto/src/xfer/dns_handle.rs

@@ -85,9 +85,10 @@ 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));
+            .extensions_mut()
+            .get_or_insert_with(Edns::new)
+            .set_max_payload(MAX_PAYLOAD_LEN)
+            .set_version(0);
     }
     message
 }

crates/proto/src/xfer/dnssec_dns_handle.rs

@@ -19,7 +19,6 @@ use futures_util::stream;
 use futures_util::stream::{Stream, TryStreamExt};
 use log::{debug, trace};
 
-use crate::error::*;
 use crate::op::{OpCode, Query};
 use crate::rr::dnssec::rdata::{DNSSECRData, DNSKEY, SIG};
 #[cfg(feature = "dnssec")]
@@ -29,6 +28,7 @@ use crate::rr::rdata::opt::EdnsOption;
 use crate::rr::{DNSClass, Name, RData, Record, RecordType};
 use crate::xfer::dns_handle::DnsHandle;
 use crate::xfer::{DnsRequest, DnsRequestOptions, DnsResponse, FirstAnswer};
+use crate::{error::*, op::Edns};
 
 #[derive(Debug)]
 struct Rrset {
@@ -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.extensions_mut().get_or_insert_with(Edns::new);
+                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);

crates/resolver/src/name_server/name_server.rs

@@ -144,7 +144,7 @@ impl<C: DnsHandle<Error = ResolveError>, P: ConnectionProvider<Conn = C>> NameSe
                     ResolveError::from_response(response, self.config.trust_nx_responses)?;
 
                 // TODO: consider making message::take_edns...
-                let remote_edns = response.edns().cloned();
+                let remote_edns = response.extensions().clone();
 
                 // take the remote edns options and store them
                 self.state.establish(remote_edns);

tests/integration-tests/tests/client_future_tests.rs

@@ -219,7 +219,8 @@ fn test_query_edns(client: &mut AsyncClient) -> impl Future<Output = ()> {
     .set_op_code(OpCode::Query)
     .set_recursion_desired(true)
     .set_edns(edns)
-    .edns_mut()
+    .extensions_mut()
+    .as_mut()
     .map(|edns| edns.set_max_payload(1232).set_version(0));
 
     client
@@ -238,9 +239,14 @@ fn test_query_edns(client: &mut AsyncClient) -> impl Future<Output = ()> {
             assert_eq!(record.name(), &name);
             assert_eq!(record.rr_type(), RecordType::A);
             assert_eq!(record.dns_class(), DNSClass::IN);
-            assert!(response.edns().is_some());
+            assert!(response.extensions().is_some());
             assert_eq!(
-                response.edns().unwrap().option(EdnsCode::Subnet).unwrap(),
+                response
+                    .extensions()
+                    .as_ref()
+                    .unwrap()
+                    .option(EdnsCode::Subnet)
+                    .unwrap(),
                 &EdnsOption::Unknown(EdnsCode::Subnet.into(), vec![0, 1, 16, 0, 1, 2])
             );
             if let RData::A(ref address) = *record.data().unwrap() {

tests/integration-tests/tests/client_tests.rs

@@ -156,7 +156,8 @@ where
     .set_op_code(OpCode::Query)
     .set_recursion_desired(true)
     .set_edns(edns)
-    .edns_mut()
+    .extensions_mut()
+    .as_mut()
     .map(|edns| edns.set_max_payload(1232).set_version(0));
 
     let response = client.send(msg).remove(0).expect("Query failed");
@@ -173,9 +174,14 @@ where
     assert_eq!(record.name(), &name);
     assert_eq!(record.rr_type(), RecordType::A);
     assert_eq!(record.dns_class(), DNSClass::IN);
-    assert!(response.edns().is_some());
+    assert!(response.extensions().is_some());
     assert_eq!(
-        response.edns().unwrap().option(EdnsCode::Subnet).unwrap(),
+        response
+            .extensions()
+            .as_ref()
+            .unwrap()
+            .option(EdnsCode::Subnet)
+            .unwrap(),
         &EdnsOption::Unknown(EdnsCode::Subnet.into(), vec![0, 1, 16, 0, 1, 2])
     );
 
@@ -227,7 +233,11 @@ where
         .expect("Query failed");
 
     println!("response records: {:?}", response);
-    assert!(response.edns().expect("edns not here").dnssec_ok());
+    assert!(response
+        .extensions()
+        .as_ref()
+        .expect("edns not here")
+        .dnssec_ok());
 
     let record = &response.answers()[0];
     assert_eq!(record.name(), &name);

tests/integration-tests/tests/dnssec_client_handle_tests.rs

@@ -50,7 +50,11 @@ where
         .expect("query failed");
 
     println!("response records: {:?}", response);
-    assert!(response.edns().expect("edns not here").dnssec_ok());
+    assert!(response
+        .extensions()
+        .as_ref()
+        .expect("edns not here")
+        .dnssec_ok());
 
     assert!(!response.answers().is_empty());
     let record = &response.answers()[0];