From f362a487a5437c17405534948e34dad0ff3a923c Mon Sep 17 00:00:00 2001
From: Nick Peng <pymumu@gmail.com>
Date: Wed, 27 Jun 2018 23:55:04 +0800
Subject: [PATCH] TCP query Bugfix

---
 src/Makefile     |  9 +++++----
 src/dns_client.c | 12 ++++++++----
 src/dns_server.c | 12 ++++++------
 3 files changed, 19 insertions(+), 14 deletions(-)

diff --git a/src/Makefile b/src/Makefile
index f39cccc..ac1cee7 100644
--- a/src/Makefile
+++ b/src/Makefile
@@ -1,10 +1,11 @@
 
 BIN=smartdns 
-OBJS=smartdns.o fast_ping.o lib/bitops.o dns_client.o dns_server.o dns.o util.o tlog.o conf.o lib/rbtree.o lib/art.o
-CFLAGS=-g -O3 -Wall -Wstrict-prototypes -fno-omit-frame-pointer -Wstrict-aliasing 
+OBJS_LIB=lib/rbtree.o lib/art.o lib/bitops.o
+OBJS=smartdns.o fast_ping.o dns_client.o dns_server.o dns.o util.o tlog.o conf.o dns_cache.o $(OBJS_LIB)
+CFLAGS=-g -O0 -Wall -Wstrict-prototypes -fno-omit-frame-pointer -Wstrict-aliasing 
 CFLAGS +=-Iinclude
 CFLAGS += -DBASE_FILE_NAME=\"$(notdir $<)\"
-CXXFLAGS=-g -O3 -Wall -std=c++11 
+CXXFLAGS=-g -O0 -Wall -std=c++11 
 CXXFLAGS +=-Iinclude
 
 .PHONY: all
@@ -15,4 +16,4 @@ $(BIN) : $(OBJS)
 	$(CC) $(OBJS) -o $@ -lpthread 
 
 clean:
-	$(RM) $(OBJS) $(BIN)
\ No newline at end of file
+	$(RM) $(OBJS) $(BIN)
diff --git a/src/dns_client.c b/src/dns_client.c
index f233de7..ece404b 100644
--- a/src/dns_client.c
+++ b/src/dns_client.c
@@ -802,6 +802,7 @@ static int _dns_client_process_tcp(struct dns_server_info *server_info, struct e
 		/* tcp result format 
 	 	 * | len (short) | dns query result | 
 	 	 */
+		inpacket_data = server_info->recv_buff.data;
 		len = ntohs(*((unsigned short *)(inpacket_data)));
 		if (len <= 0 || len >= DNS_IN_PACKSIZE) {
 			/* data len is invalid */
@@ -820,7 +821,8 @@ static int _dns_client_process_tcp(struct dns_server_info *server_info, struct e
 		if (_dns_client_recv(inpacket_data, len, &server_info->addr, server_info->ai_addrlen) != 0) {
 			goto errout;
 		}
-		server_info->recv_buff.len -= (len + 2);
+		len += 2;
+		server_info->recv_buff.len -= len;
 
 		/* move to next result */
 		if (server_info->recv_buff.len > 0) {
@@ -836,8 +838,10 @@ errout:
 	pthread_mutex_lock(&client.server_list_lock);
 	server_info->recv_buff.len = 0;
 	server_info->send_buff.len = 0;
-	close(server_info->fd);
-	server_info->fd = 0;
+	if (server_info->fd > 0) {
+		close(server_info->fd);
+		server_info->fd = -1;
+	}
 	pthread_mutex_unlock(&client.server_list_lock);
 
 	return -1;
@@ -948,7 +952,7 @@ static int _dns_client_send_tcp(struct dns_server_info *server_info, void *packe
 	memcpy(inpacket + 2, packet, len);
 	len += 2;
 
-	send_len = send(server_info->fd, inpacket, len, 0);
+	send_len = send(server_info->fd, inpacket, len, MSG_NOSIGNAL);
 	if (send_len < 0) {
 		if (errno == EAGAIN) {
 			/* save data to buffer, and retry when EPOLLOUT is available */
diff --git a/src/dns_server.c b/src/dns_server.c
index 4053eb0..cee82e1 100644
--- a/src/dns_server.c
+++ b/src/dns_server.c
@@ -560,7 +560,7 @@ static int _dns_server_process_answer(struct dns_request *request, char *domain,
 				dns_get_SOA(rrs, name, 128, &ttl, &request->soa);
 				tlog(TLOG_INFO, "SOA: mname: %s, rname: %s, serial: %d, refresh: %d, retry: %d, expire: %d, minimum: %d", request->soa.mname,
 					 request->soa.rname, request->soa.serial, request->soa.refresh, request->soa.retry, request->soa.expire, request->soa.minimum);
-			}
+			} break;
 			default:
 				tlog(TLOG_INFO, "%s, qtype: %d", name, rrs->type);
 				break;
@@ -679,13 +679,13 @@ int _dns_server_art_iter_callback(void *data, const unsigned char *key, uint32_t
 
 static int _dns_server_art_domain_cmp(const art_leaf *n, const unsigned char *prefix, int prefix_len)
 {
-    // Fail if the prefix length is too short
-    if (n->key_len > (uint32_t)prefix_len) {
+	// Fail if the prefix length is too short
+	if (n->key_len > (uint32_t)prefix_len) {
 		return 1;
 	}
 
-    // Compare the keys
-    return memcmp(n->key, prefix, n->key_len);
+	// Compare the keys
+	return memcmp(n->key, prefix, n->key_len);
 }
 
 static struct dns_address *_dns_server_get_address_by_domain(char *domain, int qtype)
@@ -695,7 +695,7 @@ static struct dns_address *_dns_server_get_address_by_domain(char *domain, int q
 	char domain_key[DNS_MAX_CNAME_LEN];
 	char type = '4';
 
-	switch(qtype) {
+	switch (qtype) {
 	case DNS_T_A:
 		type = '4';
 		break;