dns: split header and body parsing

As part of extra header validation, split out DNS body parsing to
avoid the overhead of parsing the header twice.

rust/src/dns/dns.rs

@@ -383,18 +383,23 @@ impl DNSState {
         tx.tx_data.set_event(event as u8);
     }
 
-    fn validate_header(&self, input: &[u8]) -> bool {
-        parser::dns_parse_header(input)
-            .map(|(_, header)| probe_header_validity(header, input.len()).0)
-            .unwrap_or(false)
+    fn validate_header<'a>(&self, input: &'a [u8]) -> Option<(&'a [u8], DNSHeader)> {
+        if let Ok((body, header)) = parser::dns_parse_header(input) {
+            if probe_header_validity(&header, input.len()).0 {
+                return Some((body, header));
+            }
+        }
+        None
     }
 
     fn parse_request(&mut self, input: &[u8], is_tcp: bool) -> bool {
-        if !self.validate_header(input) {
+        let (body, header) = if let Some((body, header)) = self.validate_header(input) {
+            (body, header)
+        } else {
             return !is_tcp;
-        }
+        };
 
-        match parser::dns_parse_request(input) {
+        match parser::dns_parse_request_body(body, input, header) {
             Ok((_, request)) => {
                 if request.header.flags & 0x8000 != 0 {
                     SCLogDebug!("DNS message is not a request");
@@ -460,11 +465,13 @@ impl DNSState {
     }
 
     pub fn parse_response(&mut self, input: &[u8], is_tcp: bool) -> bool {
-        if !self.validate_header(input) {
+        let (body, header) = if let Some((body, header)) = self.validate_header(input) {
+            (body, header)
+        } else {
             return !is_tcp;
-        }
+        };
 
-        match parser::dns_parse_response(input) {
+        match parser::dns_parse_response_body(body, input, header) {
             Ok((_, response)) => {
                 SCLogDebug!("Response header flags: {}", response.header.flags);
 
@@ -653,7 +660,7 @@ impl DNSState {
 
 const DNS_HEADER_SIZE: usize = 12;
 
-fn probe_header_validity(header: DNSHeader, rlen: usize) -> (bool, bool, bool) {
+fn probe_header_validity(header: &DNSHeader, rlen: usize) -> (bool, bool, bool) {
     let min_msg_size = 2
         * (header.additional_rr as usize
             + header.answer_rr as usize
@@ -686,7 +693,7 @@ fn probe(input: &[u8], dlen: usize) -> (bool, bool, bool) {
     // parse a complete message, so perform header validation only.
     if input.len() < dlen {
         if let Ok((_, header)) = parser::dns_parse_header(input) {
-            return probe_header_validity(header, dlen);
+            return probe_header_validity(&header, dlen);
         } else {
             return (false, false, false);
         }
@@ -694,11 +701,11 @@ fn probe(input: &[u8], dlen: usize) -> (bool, bool, bool) {
 
     match parser::dns_parse_request(input) {
         Ok((_, request)) => {
-            return probe_header_validity(request.header, dlen);
+            return probe_header_validity(&request.header, dlen);
         }
         Err(Err::Incomplete(_)) => match parser::dns_parse_header(input) {
             Ok((_, header)) => {
-                return probe_header_validity(header, dlen);
+                return probe_header_validity(&header, dlen);
             }
             Err(Err::Incomplete(_)) => (false, false, true),
             Err(_) => (false, false, false),

rust/src/dns/parser.rs

@@ -49,7 +49,7 @@ pub fn dns_parse_header(i: &[u8]) -> IResult<&[u8], DNSHeader> {
 ///
 /// Parameters:
 ///   start: the start of the name
-///   message: the complete message that start is a part of
+///   message: the complete message that start is a part of with the DNS header
 pub fn dns_parse_name<'b>(start: &'b [u8], message: &'b [u8]) -> IResult<&'b [u8], Vec<u8>> {
     let mut pos = start;
     let mut pivot = start;
@@ -192,12 +192,18 @@ fn dns_parse_answer<'a>(
 }
 
 /// Parse a DNS response.
-pub fn dns_parse_response(slice: &[u8]) -> IResult<&[u8], DNSResponse> {
-    let i = slice;
+pub fn dns_parse_response(message: &[u8]) -> IResult<&[u8], DNSResponse> {
+    let i = message;
     let (i, header) = dns_parse_header(i)?;
-    let (i, queries) = count(|b| dns_parse_query(b, slice), header.questions as usize)(i)?;
-    let (i, answers) = dns_parse_answer(i, slice, header.answer_rr as usize)?;
-    let (i, authorities) = dns_parse_answer(i, slice, header.authority_rr as usize)?;
+    dns_parse_response_body(i, message, header)
+}
+
+pub fn dns_parse_response_body<'a>(
+    i: &'a [u8], message: &'a [u8], header: DNSHeader,
+) -> IResult<&'a [u8], DNSResponse> {
+    let (i, queries) = count(|b| dns_parse_query(b, message), header.questions as usize)(i)?;
+    let (i, answers) = dns_parse_answer(i, message, header.answer_rr as usize)?;
+    let (i, authorities) = dns_parse_answer(i, message, header.authority_rr as usize)?;
     Ok((
         i,
         DNSResponse {
@@ -348,7 +354,14 @@ pub fn dns_parse_rdata<'a>(
 pub fn dns_parse_request(input: &[u8]) -> IResult<&[u8], DNSRequest> {
     let i = input;
     let (i, header) = dns_parse_header(i)?;
-    let (i, queries) = count(|b| dns_parse_query(b, input), header.questions as usize)(i)?;
+    dns_parse_request_body(i, input, header)
+}
+
+pub fn dns_parse_request_body<'a>(
+    input: &'a [u8], message: &'a [u8], header: DNSHeader,
+) -> IResult<&'a [u8], DNSRequest> {
+    let i = input;
+    let (i, queries) = count(|b| dns_parse_query(b, message), header.questions as usize)(i)?;
     Ok((i, DNSRequest { header, queries }))
 }