From 759d9ec0688cd7b93948ef5e47b894115d3daa44 Mon Sep 17 00:00:00 2001
From: liuxueli <liuxueli@iie.ac.cn>
Date: Fri, 24 Dec 2021 12:44:09 +0300
Subject: [PATCH] =?UTF-8?q?TSG-9125:=20tsg=5Fmaster=E4=BB=85=E6=A0=B9?=
 =?UTF-8?q?=E6=8D=AE=E7=AB=AF=E5=8F=A3=E8=AF=86=E5=88=ABDNS=E5=8D=8F?=
 =?UTF-8?q?=E8=AE=AE=E8=BF=87=E4=BA=8E=E7=AE=80=E5=8D=95=EF=BC=8C=E5=A2=9E?=
 =?UTF-8?q?=E5=8A=A0=E5=AF=B9=E8=B4=9F=E8=BD=BD=E7=9A=84=E5=88=A4=E6=96=AD?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/tsg_entry.cpp | 65 +++++++++++++++++++++++++++++++----------------
 1 file changed, 43 insertions(+), 22 deletions(-)

diff --git a/src/tsg_entry.cpp b/src/tsg_entry.cpp
index 6136123..7705810 100644
--- a/src/tsg_entry.cpp
+++ b/src/tsg_entry.cpp
@@ -320,6 +320,46 @@ static int is_hited_allow(struct Maat_rule_t *result, int hit_cnt)
 	return 0;
 }
 
+static int is_dns_protocol(const struct streaminfo *a_stream)
+{
+	struct stream_tuple4_v4 *tpl4 = NULL;
+	struct stream_tuple4_v6 *tpl6 = NULL;
+
+	if(a_stream->pudpdetail==NULL || a_stream->pudpdetail->pdata==NULL || a_stream->pudpdetail->datalen<12)
+	{
+		return 0;
+	}
+	
+	switch(a_stream->addr.addrtype)
+	{
+		case ADDR_TYPE_IPV4:
+			tpl4=a_stream->addr.tuple4_v4;
+			if((ntohs(tpl4->source)!=53) && (ntohs(tpl4->dest)!=53))
+			{
+				return 0;
+			}
+			break;
+		case ADDR_TYPE_IPV6:
+			tpl6=a_stream->addr.tuple4_v6;
+			if((ntohs(tpl6->source)!=53) && (ntohs(tpl6->dest)!=53))
+			{
+				return 0;
+			}
+			break;
+		default:
+			return 0;
+			break;
+	}
+
+	struct _dns_hdr *dns_hdr=(struct _dns_hdr *)(a_stream->pudpdetail->pdata);
+	if(dns_hdr->qdcount==1 && dns_hdr->z==0)
+	{
+		return 1;
+	}
+
+	return 0;
+}
+
 int set_struct_project(const struct streaminfo *a_stream, int project_id, void *data)
 {
 	if(a_stream==NULL || project_id<0)
@@ -1289,29 +1329,10 @@ static int identify_application_protocol(const struct streaminfo *a_stream, stru
 		case	STREAM_TYPE_UDP:
 			if(g_tsg_para.proto_flag&(1<<PROTO_DNS))  //dns
 			{
-				struct stream_tuple4_v4 *tpl4 = NULL;
-				struct stream_tuple4_v6 *tpl6 = NULL;
-
-				switch(a_stream->addr.addrtype)
+				if(is_dns_protocol(a_stream))
 				{
-					case ADDR_TYPE_IPV4:
-						tpl4=a_stream->addr.tuple4_v4;
-						if((ntohs(tpl4->source)==53) || (ntohs(tpl4->dest)==53))
-						{
-							context->proto=PROTO_DNS;
-							return 1;
-						}
-						break;
-					case ADDR_TYPE_IPV6:
-						tpl6=a_stream->addr.tuple4_v6;
-						if((ntohs(tpl6->source)==53) || (ntohs(tpl6->dest)==53))
-						{
-							context->proto=PROTO_DNS;
-							return 1;
-						}
-						break;
-					default:
-						break;
+					context->proto=PROTO_DNS;
+					return 1;
 				}
 			}