From f9446d687d0feff5fdbc451409c92cf119b98320 Mon Sep 17 00:00:00 2001
From: Dengfeng Liu <liudf0716@gmail.com>
Date: Wed, 12 Jun 2024 11:15:53 +0800
Subject: [PATCH] feat: using new method to parse dns's response

Signed-off-by: Dengfeng Liu <liudf0716@gmail.com>
---
 src/dns_forward.c | 37 +++++++++++++++++++++++++------------
 src/dns_forward.h | 17 +++++++++++++++++
 2 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/src/dns_forward.c b/src/dns_forward.c
index d29a820a4..c5ec5ea54 100644
--- a/src/dns_forward.c
+++ b/src/dns_forward.c
@@ -43,11 +43,23 @@ strrstr(const char *haystack, const char *needle)
 };
 
 static int 
-parse_dns_response_ex(unsigned char *response, int response_len) {
+process_dns_response(unsigned char *response, int response_len) {
     s_config *config = config_get_config();
-    unsigned char *ptr = response + 12; // Skip the DNS header
-    int qdcount = ntohs(*(unsigned short *)(response + 4));
-    int ancount = ntohs(*(unsigned short *)(response + 6));
+    
+
+    if (response_len <= sizeof(struct dns_header)) {
+        debug(LOG_WARNING, "Invalid DNS response");
+        return -1;
+    }
+    struct dns_header *header = (struct dns_header *)response;
+    if (header->qr != 1 || header->opcode != 0 || header->rcode != 0) {
+        debug(LOG_WARNING, "Invalid DNS response");
+        return -1;
+    }
+
+    int qdcount = ntohs(header->qdcount);
+    int ancount = ntohs(header->ancount);
+    unsigned char *ptr = response + sizeof(struct dns_header); // Skip the DNS header
 
     // Skip the question section
     for (int i = 0; i < qdcount; i++) {
@@ -72,24 +84,25 @@ parse_dns_response_ex(unsigned char *response, int response_len) {
         }
         domain[--len] = '\0'; // Null-terminate the domain string
         ptr++;
-
+        
         unsigned short type = ntohs(*(unsigned short *)ptr);
         ptr += 2;
         ptr += 8; // Skip class, TTL, and data length
         unsigned short data_len = ntohs(*(unsigned short *)ptr);
         ptr += 2;
-
-        if (type == 1 && data_len == 4) { // Type A
-            char ip[INET_ADDRSTRLEN];
-            inet_ntop(AF_INET, ptr, ip, sizeof(ip));
+        
+        debug(LOG_DEBUG, "DNS response type: %d, domain: %s", type, domain);
+        if (type == 1 && data_len == 4) { // Type A record
             t_domain_trusted *p = config->pan_domains_trusted;
             while (p) {
-                if (strstr(domain, p->domain)) {
+                if (strrstr(domain, p->domain)) {
+                    char ip[INET_ADDRSTRLEN] = {0};
+                    inet_ntop(AF_INET, ptr, ip, sizeof(ip));
                     t_ip_trusted *ip_entry = malloc(sizeof(t_ip_trusted));
                     strcpy(ip_entry->ip, ip);
                     ip_entry->next = p->ips_trusted;
                     p->ips_trusted = ip_entry;
-                    printf("Trusted domain: %s -> %s\n", domain, ip);
+                    debug(LOG_DEBUG, "Trusted domain: %s -> %s\n", domain, ip);
                     break;
                 }
                 p = p->next;
@@ -101,7 +114,7 @@ parse_dns_response_ex(unsigned char *response, int response_len) {
 }
 
 static void 
-process_dns_response(char *response, int response_len) 
+process_dns_response_ex(char *response, int response_len) 
 {
     ns_msg handle;
     s_config *config = config_get_config();
diff --git a/src/dns_forward.h b/src/dns_forward.h
index c7b9bec00..9c07a01b9 100644
--- a/src/dns_forward.h
+++ b/src/dns_forward.h
@@ -4,6 +4,23 @@
 #define DNS_FORWARD_PORT 15353
 #define LOCAL_DNS_PORT 53
 
+// define a structure to hold the DNS header
+struct dns_header {
+    unsigned short id;
+    unsigned char rd :1;
+    unsigned char tc :1;
+    unsigned char aa :1;
+    unsigned char opcode :4;
+    unsigned char qr :1;
+    unsigned char rcode :4;
+    unsigned char z :3;
+    unsigned char ra :1;
+    unsigned short qdcount;
+    unsigned short ancount;
+    unsigned short nscount;
+    unsigned short arcount;
+};
+
 void *dns_forward_thread(void *);
 
 #endif
\ No newline at end of file