From 85c5577955131eb63a12dff03422a74f5eb756ec Mon Sep 17 00:00:00 2001
From: cpq <valenok@gmail.com>
Date: Fri, 11 Dec 2020 14:02:52 +0000
Subject: [PATCH] mg_dns_parse_name() recursion protection

---
 mongoose.c       | 3 ++-
 src/dns.c        | 7 ++++---
 test/unit_test.c | 3 +++
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/mongoose.c b/mongoose.c
index a6a61893..54c8cb08 100644
--- a/mongoose.c
+++ b/mongoose.c
@@ -165,11 +165,12 @@ void mg_resolve_cancel(struct mg_mgr *mgr, struct mg_connection *c) {
 static size_t mg_dns_parse_name(const uint8_t *s, const uint8_t *e, size_t off,
                                 char *to, size_t tolen, int depth) {
   size_t i = 0, j = 0;
+  // if (depth > 5) return 0;
   while (&s[off + i + 1] < e && s[off + i] > 0) {
     size_t n = s[off + i];
     if (n & 0xc0) {
       size_t ptr = (((n & 0x3f) << 8) | s[off + i + 1]) - 12;  // 12 is hdr len
-      if (&s[ptr + 1] < e && (s[ptr] & 0xc0) == 0 && depth < 5) {
+      if (&s[ptr + 1] < e && (s[ptr] & 0xc0) == 0) {
         j = mg_dns_parse_name(s, e, ptr, to, tolen, depth + 1);
       }
       i++;
diff --git a/src/dns.c b/src/dns.c
index 61825f36..63cf0bac 100644
--- a/src/dns.c
+++ b/src/dns.c
@@ -39,11 +39,12 @@ void mg_resolve_cancel(struct mg_mgr *mgr, struct mg_connection *c) {
 static size_t mg_dns_parse_name(const uint8_t *s, const uint8_t *e, size_t off,
                                 char *to, size_t tolen, int depth) {
   size_t i = 0, j = 0;
+  if (depth > 5) return 0;
   while (&s[off + i + 1] < e && s[off + i] > 0) {
     size_t n = s[off + i];
     if (n & 0xc0) {
       size_t ptr = (((n & 0x3f) << 8) | s[off + i + 1]) - 12;  // 12 is hdr len
-      if (&s[ptr + 1] < e && (s[ptr] & 0xc0) == 0 && depth < 5) {
+      if (&s[ptr + 1] < e && (s[ptr] & 0xc0) == 0) {
         j = mg_dns_parse_name(s, e, ptr, to, tolen, depth + 1);
       }
       i++;
@@ -67,9 +68,9 @@ static size_t mg_dns_parse_name(const uint8_t *s, const uint8_t *e, size_t off,
 int mg_dns_parse(const uint8_t *buf, size_t len, struct mg_dns_message *dm) {
   struct mg_dns_header *h = (struct mg_dns_header *) buf;
   const uint8_t *s = buf + sizeof(*h), *e = &buf[len];
-  size_t i, j, n, ok = 0;
+  size_t i, j = 0, n, ok = 0;
   if (len < sizeof(*h)) return ok;
-  for (i = j = 0; i < mg_ntohs(h->num_questions); i++) {
+  for (i = 0; i < mg_ntohs(h->num_questions); i++) {
     j += mg_dns_parse_name(s, e, j, dm->name, sizeof(dm->name), 0) + 5;
     // LOG(LL_INFO, ("QUE [%s]", name));
   }
diff --git a/test/unit_test.c b/test/unit_test.c
index f933fad7..804d6916 100644
--- a/test/unit_test.c
+++ b/test/unit_test.c
@@ -864,7 +864,10 @@ static void test_str(void) {
 
 static void test_dns(void) {
   struct mg_dns_message dm;
+  char *data = mg_file_read("data.txt");
   ASSERT(mg_dns_parse(NULL, 0, &dm) == 0);
+  ASSERT(mg_dns_parse((uint8_t *) data, strlen(data), &dm) == 0);
+  free(data);
 }
 
 static void test_util(void) {