Handle TCP DNS replies and add is_reply flag

2 files changed, 14 insertions, 6 deletions

diff --git a/src/include/ndpi_typedefs.h b/src/include/ndpi_typedefs.h
index 3e71c2ab2..486213269 100644
--- a/src/include/ndpi_typedefs.h
+++ b/src/include/ndpi_typedefs.h
@@ -1172,7 +1172,7 @@ struct ndpi_flow_struct {
   union {
     /* the only fields useful for nDPI and ntopng */
     struct {
-      u_int8_t num_queries, num_answers, reply_code;
+      u_int8_t num_queries, num_answers, reply_code, is_query;
       u_int16_t query_type, query_class, rsp_type;
       ndpi_ip_addr_t rsp_addr; /* The first address in a DNS response packet */
     } dns;
diff --git a/src/lib/protocols/dns.c b/src/lib/protocols/dns.c
index 7c3ead514..91598f47c 100644
--- a/src/lib/protocols/dns.c
+++ b/src/lib/protocols/dns.c
@@ -199,8 +199,13 @@ static int search_dns_again(struct ndpi_detection_module_struct *ndpi_struct, st
   /* possibly dissect the DNS reply */
   ndpi_search_dns(ndpi_struct, flow);
 
-  /* stop extra processing */
-  return(0);
+  if(flow->protos.dns.num_answers > 0) {
+    /* stop extra processing */
+    return(0);
+  }
+
+  /* Possibly more processing */
+  return(1);
 }
 
 /* *********************************************** */
@@ -280,16 +285,19 @@ void ndpi_search_dns(struct ndpi_detection_module_struct *ndpi_struct, struct nd
       else
 	ret.master_protocol = NDPI_PROTOCOL_DNS;
     }
+
+    /* Report if this is a DNS query or reply */
+    flow->protos.dns.is_query = is_query;
     
-    if(is_query && (ndpi_struct->dns_dont_dissect_response == 0) && (flow->num_processed_pkts == 1)) {
+    if(is_query && (ndpi_struct->dns_dont_dissect_response == 0) && (flow->check_extra_packets == 0)) {
       /* In this case we say that the protocol has been detected just to let apps carry on with their activities */
       ndpi_set_detected_protocol(ndpi_struct, flow, ret.app_protocol, ret.master_protocol);
 
       /* This is necessary to inform the core to call this dissector again */
       flow->check_extra_packets = 1;
 
-      /* Dissect at most 1 more packets, hopefully the DNS response */
-      flow->max_extra_packets_to_check = 1;
+      /* Don't use just 1 as in TCP DNS more packets could be returned (e.g. ACK). */
+      flow->max_extra_packets_to_check = 5;
       flow->extra_packets_func = search_dns_again;
       return; /* The response will set the verdict */
     }