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.

(cherry picked from commit d720ead470)

rust/src/dns/dns.rs

@@ -29,7 +29,7 @@ use crate::core::STREAM_TOSERVER;
 use crate::core::{self, AppProto, ALPROTO_UNKNOWN, IPPROTO_UDP, IPPROTO_TCP};
 use crate::dns::parser;
 
-use nom::IResult;
+use nom::{Err, IResult};
 use nom::number::streaming::be_u16;
 
 /// DNS record types.
@@ -500,18 +500,23 @@ impl DNSState {
                                                         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");
@@ -562,11 +567,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);
@@ -727,7 +734,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
@@ -756,7 +763,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);
         }
@@ -764,17 +771,15 @@ fn probe(input: &[u8], dlen: usize) -> (bool, bool, bool) {
 
     match parser::dns_parse_request(input) {
         Ok((_, request)) => {
-            return probe_header_validity(request.header, dlen);
-        },
-        Err(nom::Err::Incomplete(_)) => {
-            match parser::dns_parse_header(input) {
-                Ok((_, header)) => {
-                    return probe_header_validity(header, dlen);
-                }
-                Err(nom::Err::Incomplete(_)) => (false, false, true),
-                Err(_) => (false, false, false),
-            }
+            return probe_header_validity(&request.header, dlen);
         }
+        Err(Err::Incomplete(_)) => match parser::dns_parse_header(input) {
+            Ok((_, header)) => {
+                return probe_header_validity(&header, dlen);
+            }
+            Err(Err::Incomplete(_)) => (false, false, true),
+            Err(_) => (false, false, false),
+        },
         Err(_) => (false, false, false),
     }
 }

rust/src/dns/parser.rs

@@ -20,6 +20,7 @@
 use nom::IResult;
 use nom::combinator::rest;
 use nom::error::ErrorKind;
+use nom::multi::count;
 use nom::number::streaming::{be_u8, be_u16, be_u32};
 use nom;
 use crate::dns::dns::*;
@@ -50,10 +51,8 @@ named!(pub dns_parse_header<DNSHeader>,
 ///
 /// Parameters:
 ///   start: the start of the name
-///   message: the complete message that start is a part of
-pub fn dns_parse_name<'a, 'b>(start: &'b [u8],
-                              message: &'b [u8])
-                              -> IResult<&'b [u8], Vec<u8>> {
+///   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;
     let mut name: Vec<u8> = Vec::with_capacity(32);
@@ -197,26 +196,27 @@ fn dns_parse_answer<'a>(slice: &'a [u8], message: &'a [u8], count: usize)
 
 
 /// Parse a DNS response.
-pub fn dns_parse_response<'a>(slice: &'a [u8])
-                              -> IResult<&[u8], DNSResponse> {
-    do_parse!(
-        slice,
-        header: dns_parse_header
-            >> queries: count!(
-                call!(dns_parse_query, slice), header.questions as usize)
-            >> answers: call!(
-                dns_parse_answer, slice, header.answer_rr as usize)
-            >> authorities: call!(
-                dns_parse_answer, slice, header.authority_rr as usize)
-            >> (
-                DNSResponse{
-                    header: header,
-                    queries: queries,
-                    answers: answers,
-                    authorities: authorities,
-                }
-            )
-    )
+pub fn dns_parse_response(message: &[u8]) -> IResult<&[u8], DNSResponse> {
+    let i = message;
+    let (i, header) = dns_parse_header(i)?;
+    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 {
+            header,
+            queries,
+            answers,
+            authorities,
+        },
+    ))
 }
 
 /// Parse a single DNS query.
@@ -344,19 +344,18 @@ pub fn dns_parse_rdata<'a>(input: &'a [u8], message: &'a [u8], rrtype: u16)
 }
 
 /// Parse a DNS request.
-pub fn dns_parse_request<'a>(input: &'a [u8]) -> IResult<&[u8], DNSRequest> {
-    do_parse!(
-        input,
-        header: dns_parse_header >>
-        queries: count!(call!(dns_parse_query, input),
-                        header.questions as usize) >>
-            (
-                DNSRequest{
-                    header: header,
-                    queries: queries,
-                }
-            )
-    )
+pub fn dns_parse_request(input: &[u8]) -> IResult<&[u8], DNSRequest> {
+    let i = input;
+    let (i, header) = dns_parse_header(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 }))
 }
 
 #[cfg(test)]