Update high ResponseCode in EDNS Section if required

crates/proto/src/op/header.rs

@@ -285,6 +285,9 @@ impl Header {
     }
 
     /// The low response code (original response codes before EDNS extensions)
+    ///
+    /// Warning: This method can only used for basic ResponseCodes.
+    /// Some ResponseCode require the use of the EDNS record to represent the high bits.
     pub fn set_response_code(&mut self, response_code: ResponseCode) -> &mut Self {
         self.response_code = response_code;
         self

crates/proto/src/op/message.rs

@@ -126,6 +126,8 @@ impl Message {
 
     /// Returns a Message constructed with error details to return to a client
     ///
+    /// If the `response_code` requires an EDNS section, this function creates one if not already present.
+    ///
     /// # Arguments
     ///
     /// * `id` - message id should match the request message id
@@ -212,6 +214,9 @@ impl Message {
     /// see `Header::set_response_code`
     pub fn set_response_code(&mut self, response_code: ResponseCode) -> &mut Self {
         self.header.set_response_code(response_code);
+        if response_code.high() != 0 {
+            self.edns_mut().set_rcode_high(response_code.high());
+        }
         self
     }
 
@@ -508,11 +513,7 @@ impl Message {
 
     /// If edns is_none, this will create a new default Edns.
     pub fn edns_mut(&mut self) -> &mut Edns {
-        if self.edns.is_none() {
-            self.edns = Some(Edns::new());
-        }
-
-        self.edns.as_mut().unwrap()
+        self.edns.get_or_insert_with(Edns::new)
     }
 
     /// # Return value
@@ -589,7 +590,7 @@ impl Message {
                 query_count: self.queries.len(),
                 answer_count: self.answers.len(),
                 nameserver_count: self.name_servers.len(),
-                additional_count: self.additionals.len(),
+                additional_count: self.additionals.len() + if self.edns.is_some() { 1 } else { 0 },
             },
         );
         self
@@ -1054,6 +1055,21 @@ fn test_emit_and_read_records() {
     test_emit_and_read(message);
 }
 
+#[test]
+fn test_extended_response_code() {
+    let mut message = Message::new();
+    message.set_response_code(ResponseCode::BADCOOKIE);
+    message.update_counts();
+    test_emit_and_read(message);
+}
+
+#[test]
+fn test_extended_response_code_error() {
+    let mut message = Message::error_msg(0, OpCode::Query, ResponseCode::BADVERS);
+    message.update_counts();
+    test_emit_and_read(message);
+}
+
 #[cfg(test)]
 fn test_emit_and_read(message: Message) {
     let mut byte_vec: Vec<u8> = Vec::with_capacity(512);

crates/server/src/authority/message_response.rs

@@ -207,6 +207,8 @@ impl<'q> MessageResponseBuilder<'q> {
 
     /// Constructs a new error MessageResponse with associated settings
     ///
+    /// If the `response_code` requires an EDNS section, this function creates one if not already present.
+    ///
     /// # Arguments
     ///
     /// * `id` - request id to which this is a response
@@ -226,6 +228,11 @@ impl<'q> MessageResponseBuilder<'q> {
     > {
         let mut header = Header::response_from_request(request_header);
         header.set_response_code(response_code);
+        let mut edns = self.edns;
+        if response_code.high() != 0 {
+            edns.get_or_insert_with(Edns::new)
+                .set_rcode_high(response_code.high());
+        }
 
         MessageResponse {
             header,
@@ -235,7 +242,7 @@ impl<'q> MessageResponseBuilder<'q> {
             soa: Box::new(None.into_iter()),
             additionals: Box::new(None.into_iter()),
             sig0: self.sig0.unwrap_or_default(),
-            edns: self.edns,
+            edns,
         }
     }
 }
@@ -322,4 +329,26 @@ mod tests {
         assert_eq!(response.answer_count(), 0);
         assert!(response.name_server_count() > 1);
     }
+
+    #[test]
+    fn test_extended_response_codes() {
+        let mut buf = Vec::with_capacity(512);
+        {
+            let mut encoder = BinEncoder::new(&mut buf);
+            encoder.set_max_size(512);
+
+            let message = MessageResponseBuilder {
+                query: None,
+                sig0: None,
+                edns: None,
+            };
+            message
+                .error_msg(&Header::new(), ResponseCode::BADCOOKIE)
+                .destructive_emit(&mut encoder)
+                .expect("failed to encode");
+        }
+
+        let response = Message::from_vec(&buf).expect("failed to decode");
+        assert_eq!(response.response_code(), ResponseCode::BADCOOKIE);
+    }
 }