dns: improve probing parser

Checks opcode is valid
Checks additional_rr do not exceed message length
Better logic for incomplete cases

rust/src/dns/dns.rs

@@ -659,23 +659,53 @@ impl DNSState {
     }
 }
 
+const DNS_HEADER_SIZE: usize = 12;
+
+fn probe_header_validity(header: DNSHeader, rlen: usize) -> (bool, bool, bool) {
+    let opcode = ((header.flags >> 11) & 0xf) as u8;
+    if opcode >= 7 {
+        //unassigned opcode
+        return (false, false, false);
+    }
+    if 2 * (header.additional_rr as usize
+        + header.answer_rr as usize
+        + header.authority_rr as usize
+        + header.questions as usize)
+        + DNS_HEADER_SIZE
+        > rlen
+    {
+        //not enough data for such a DNS record
+        return (false, false, false);
+    }
+    let is_request = header.flags & 0x8000 == 0;
+    return (true, is_request, false);
+}
+
 /// Probe input to see if it looks like DNS.
-fn probe(input: &[u8]) -> (bool, bool) {
-    match parser::dns_parse_request(input) {
+fn probe(input: &[u8], dlen: usize) -> (bool, bool, bool) {
+    let i2 = if input.len() <= dlen { input } else { &input[..dlen] };
+    match parser::dns_parse_request(i2) {
         Ok((_, request)) => {
-            let is_request = request.header.flags & 0x8000 == 0;
-            return (true, is_request);
+            return probe_header_validity(request.header, dlen);
         },
-        Err(_) => (false, false),
+        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),
+            }
+        }
+        Err(_) => (false, false, false),
     }
 }
 
 /// Probe TCP input to see if it looks like DNS.
 pub fn probe_tcp(input: &[u8]) -> (bool, bool, bool) {
-    match be_u16(input) as IResult<&[u8],_> {
-        Ok((rem, _)) => {
-            let r = probe(rem);
-            return (r.0, r.1, false);
+    match be_u16(input) as IResult<&[u8],u16> {
+        Ok((rem, dlen)) => {
+            return probe(rem, dlen as usize);
         },
         Err(nom::Err::Incomplete(_)) => {
             return (false, false, true);
@@ -961,7 +991,7 @@ pub extern "C" fn rs_dns_probe(
     let slice: &[u8] = unsafe {
         std::slice::from_raw_parts(input as *mut u8, len as usize)
     };
-    let (is_dns, is_request) = probe(slice);
+    let (is_dns, is_request, _) = probe(slice, slice.len());
     if is_dns {
         let dir = if is_request {
             core::STREAM_TOSERVER