Preserve EDNS OPT record in a response if truncation occurs

crates/proto/src/op/edns.rs

@@ -123,6 +123,18 @@ impl Edns {
     pub fn set_option(&mut self, option: EdnsOption) {
         self.options.insert(option);
     }
+
+    pub(crate) fn len(&self) -> u16 {
+        1 //   Name::root
+        + 2 // RecordType::OPT
+        + 2 // DNSClass
+        + 4 // TTL (rcode_high, version and flags)
+        + self.options.options().iter().map(|o| {
+            2 //   Option code
+            + 2 // Len field
+            + o.1.len()
+        }).sum::<u16>()
+    }
 }
 
 impl<'a> From<&'a Record> for Edns {
@@ -208,6 +220,7 @@ impl BinEncodable for Edns {
         opt::emit(encoder, &self.options)?;
         let len = encoder.len_since_place(&place);
         assert!(len <= u16::max_value() as usize);
+        debug_assert_eq!(len, self.len() as usize);
 
         place.replace(encoder, len as u16)?;
         Ok(())

crates/proto/src/op/message.rs

@@ -16,7 +16,6 @@
 
 //! Basic protocol message for DNS
 
-use std::iter;
 use std::mem;
 use std::ops::Deref;
 use std::sync::Arc;
@@ -769,6 +768,23 @@ where
     N: EmitAndCount,
     D: EmitAndCount,
 {
+    let real_encoder_max_size = encoder.max_size();
+    if let Some(edns) = edns {
+        // From RFC 6891 section-7:
+        //   The minimal response MUST be the DNS header, question section, and an
+        //   OPT record. This MUST also occur when a truncated response (using
+        //   the DNS header's TC bit) is returned.
+        // Hence, we reserve some space for the EDNS OPT record
+        encoder.set_max_size(match real_encoder_max_size.checked_sub(edns.len()) {
+            Some(size) => size,
+            None => {
+                return Err(
+                    ProtoErrorKind::MaxBufferSizeExceeded(real_encoder_max_size as usize).into(),
+                )
+            }
+        });
+    }
+
     let include_sig0: bool = encoder.mode() != EncodeMode::Signing;
     let place = encoder.place::<Header>()?;
 
@@ -780,10 +796,10 @@ where
     let mut additional_count = count_was_truncated(additionals.emit(encoder))?;
 
     if let Some(edns) = edns {
+        encoder.set_max_size(real_encoder_max_size);
         // need to commit the error code
-        let count = count_was_truncated(encoder.emit_all(iter::once(&Record::from(edns))))?;
-        additional_count.0 += count.0;
-        additional_count.1 |= count.1;
+        Record::from(edns).emit(encoder)?;
+        additional_count.0 += 1;
     }
 
     // this is a little hacky, but if we are Verifying a signature, i.e. the original Message

crates/proto/src/serialize/binary/encoder.rs

@@ -23,6 +23,7 @@ use crate::op::Header;
 // this is private to make sure there is no accidental access to the inner buffer.
 mod private {
     use crate::error::{ProtoErrorKind, ProtoResult};
+    use std::convert::TryFrom;
 
     /// A wrapper for a buffer that guarantees writes never exceed a defined set of bytes
     pub struct MaximalBuf<'a> {
@@ -43,6 +44,11 @@ mod private {
             self.max_size = max as usize;
         }
 
+        /// Returns the enforced maximum size
+        pub fn max_size(&mut self) -> u16 {
+            u16::try_from(self.max_size).expect("max_size is known to fit into u16")
+        }
+
         /// returns an error if the maximum buffer size would be exceeded with the addition number of elements
         ///
         /// and reserves the additional space in the buffer
@@ -145,6 +151,11 @@ impl<'a> BinEncoder<'a> {
         self.buffer.set_max_size(max);
     }
 
+    /// Returns the enforced maximum size of the buffer
+    pub(crate) fn max_size(&mut self) -> u16 {
+        self.buffer.max_size()
+    }
+
     /// Returns a reference to the internal buffer
     pub fn into_bytes(self) -> &'a Vec<u8> {
         self.buffer.into_bytes()
@@ -400,8 +411,12 @@ impl<'a> BinEncoder<'a> {
         let len = T::size_of();
 
         // resize the buffer
-        self.buffer
-            .enforced_write(len, |buffer| buffer.resize(index + len, 0))?;
+        match (len + index).checked_sub(self.buffer.len()) {
+            None => {}
+            Some(resize_to) => self
+                .buffer
+                .enforced_write(resize_to, |buffer| buffer.resize(index + len, 0))?,
+        }
 
         // update the offset
         self.offset += len;

crates/server/src/authority/message_response.rs

@@ -223,11 +223,7 @@ mod tests {
             let mut encoder = BinEncoder::new(&mut buf);
             encoder.set_max_size(512);
 
-            let answer = Record::new()
-                .set_name(Name::from_str("www.example.com.").unwrap())
-                .set_rdata(RData::A(Ipv4Addr::new(93, 184, 216, 34)))
-                .set_dns_class(DNSClass::NONE)
-                .clone();
+            let answer = make_example_record();
 
             let message = MessageResponse {
                 header: Header::new(),
@@ -259,11 +255,7 @@ mod tests {
             let mut encoder = BinEncoder::new(&mut buf);
             encoder.set_max_size(512);
 
-            let answer = Record::new()
-                .set_name(Name::from_str("www.example.com.").unwrap())
-                .set_rdata(RData::A(Ipv4Addr::new(93, 184, 216, 34)))
-                .set_dns_class(DNSClass::NONE)
-                .clone();
+            let answer = make_example_record();
 
             let message = MessageResponse {
                 header: Header::new(),
@@ -286,4 +278,67 @@ mod tests {
         assert_eq!(response.answer_count(), 0);
         assert!(response.name_server_count() > 1);
     }
+
+    #[test]
+    fn test_edns_persists_in_truncated_message() {
+        let mut buf = Vec::with_capacity(512);
+        {
+            let mut encoder = BinEncoder::new(&mut buf);
+            encoder.set_max_size(512);
+
+            let answer = make_example_record();
+
+            let message = MessageResponse {
+                header: Header::new(),
+                queries: None,
+                answers: iter::repeat(&answer),
+                name_servers: iter::once(&answer),
+                soa: iter::once(&answer),
+                additionals: iter::once(&answer),
+                sig0: vec![],
+                edns: Some(Edns::new()),
+            };
+
+            message
+                .destructive_emit(&mut encoder)
+                .expect("failed to encode");
+        }
+
+        let response = Message::from_vec(&buf).expect("failed to decode");
+        assert!(response.header().truncated());
+        assert!(response.answer_count() > 1);
+        // should never have written the name server field...
+        assert_eq!(response.name_server_count(), 0);
+        assert!(response.edns().is_some());
+    }
+
+    #[test]
+    fn test_too_low_max_size() {
+        let mut buf = Vec::with_capacity(512);
+        let mut encoder = BinEncoder::new(&mut buf);
+        encoder.set_max_size(2);
+
+        let answer = make_example_record();
+
+        let message = MessageResponse {
+            header: Header::new(),
+            queries: None,
+            answers: iter::repeat(&answer),
+            name_servers: iter::once(&answer),
+            soa: iter::once(&answer),
+            additionals: iter::once(&answer),
+            sig0: vec![],
+            edns: Some(Edns::new()),
+        };
+
+        assert!(message.destructive_emit(&mut encoder).is_err())
+    }
+
+    fn make_example_record() -> Record {
+        Record::new()
+            .set_name(Name::from_str("www.example.com.").unwrap())
+            .set_rdata(RData::A(Ipv4Addr::new(93, 184, 216, 34)))
+            .set_dns_class(DNSClass::NONE)
+            .clone()
+    }
 }