Dns server cleanup (#5194)

Clean up the DNSServer class.

Removed member variables that are not required outside a
member call lifetime, and add destructor/checks.

Fixes #5179

Author: Makuna

Diff for: libraries/DNSServer/keywords.txt

@@ -33,6 +33,7 @@ stop	KEYWORD2
 DNS_QR_QUERY	LITERAL1		RESERVED_WORD_2
 DNS_QR_RESPONSE	LITERAL1		RESERVED_WORD_2
 DNS_OPCODE_QUERY	LITERAL1		RESERVED_WORD_2
+MAX_DNSNAME_LENGTH	LITERAL1		RESERVED_WORD_2
 NoError	LITERAL1		RESERVED_WORD_2
 FormError	LITERAL1		RESERVED_WORD_2
 ServerFailure	LITERAL1		RESERVED_WORD_2

Diff for: libraries/DNSServer/library.properties

@@ -1,5 +1,5 @@
 name=DNSServer
-version=1.1.0
+version=1.1.1
 author=Kristijan Novoselić
 maintainer=Kristijan Novoselić, <[email protected]>
 sentence=A simple DNS server for ESP8266.

Diff for: libraries/DNSServer/src/DNSServer.cpp

@@ -13,7 +13,7 @@ bool DNSServer::start(const uint16_t &port, const String &domainName,
                      const IPAddress &resolvedIP)
 {
   _port = port;
-  _buffer = NULL;
+  
   _domainName = domainName;
   _resolvedIP[0] = resolvedIP[0];
   _resolvedIP[1] = resolvedIP[1];
@@ -36,8 +36,6 @@ void DNSServer::setTTL(const uint32_t &ttl)
 void DNSServer::stop()
 {
   _udp.stop();
-  free(_buffer);
-  _buffer = NULL;
 }
 
 void DNSServer::downcaseAndRemoveWwwPrefix(String &domainName)
@@ -48,82 +46,106 @@ void DNSServer::downcaseAndRemoveWwwPrefix(String &domainName)
 
 void DNSServer::processNextRequest()
 {
-  _currentPacketSize = _udp.parsePacket();
-  if (_currentPacketSize)
+  size_t packetSize = _udp.parsePacket();
+
+  if (packetSize >= sizeof(DNSHeader))
   {
-    if (_buffer != NULL) free(_buffer);
-    _buffer = (unsigned char*)malloc(_currentPacketSize * sizeof(char));
-    if (_buffer == NULL) return;
-    _udp.read(_buffer, _currentPacketSize);
-    _dnsHeader = (DNSHeader*) _buffer;
-
-    if (_dnsHeader->QR == DNS_QR_QUERY &&
-        _dnsHeader->OPCode == DNS_OPCODE_QUERY &&
-        requestIncludesOnlyOneQuestion() &&
-        (_domainName == "*" || getDomainNameWithoutWwwPrefix() == _domainName)
+    uint8_t* buffer = reinterpret_cast<uint8_t*>(malloc(packetSize));
+    if (buffer == NULL) return;
+
+    _udp.read(buffer, packetSize);
+
+    DNSHeader* dnsHeader = reinterpret_cast<DNSHeader*>(buffer);
+
+    if (dnsHeader->QR == DNS_QR_QUERY &&
+        dnsHeader->OPCode == DNS_OPCODE_QUERY &&
+        requestIncludesOnlyOneQuestion(dnsHeader) &&
+        (_domainName == "*" || getDomainNameWithoutWwwPrefix(buffer, packetSize) == _domainName)
        )
     {
-      replyWithIP();
+      replyWithIP(buffer, packetSize);
     }
-    else if (_dnsHeader->QR == DNS_QR_QUERY)
+    else if (dnsHeader->QR == DNS_QR_QUERY)
     {
-      replyWithCustomCode();
+      replyWithCustomCode(buffer, packetSize);
     }
 
-    free(_buffer);
-    _buffer = NULL;
+    free(buffer);
   }
 }
 
-bool DNSServer::requestIncludesOnlyOneQuestion()
+bool DNSServer::requestIncludesOnlyOneQuestion(const DNSHeader* dnsHeader)
 {
-  return ntohs(_dnsHeader->QDCount) == 1 &&
-         _dnsHeader->ANCount == 0 &&
-         _dnsHeader->NSCount == 0 &&
-         _dnsHeader->ARCount == 0;
+  return ntohs(dnsHeader->QDCount) == 1 &&
+         dnsHeader->ANCount == 0 &&
+         dnsHeader->NSCount == 0 &&
+         dnsHeader->ARCount == 0;
 }
 
-String DNSServer::getDomainNameWithoutWwwPrefix()
+String DNSServer::getDomainNameWithoutWwwPrefix(const uint8_t* buffer, size_t packetSize)
 {
-  String parsedDomainName = "";
-  if (_buffer == NULL) return parsedDomainName;
-  unsigned char *start = _buffer + 12;
-  if (*start == 0)
-  {
-    return parsedDomainName;
-  }
-  int pos = 0;
-  while(true)
+  String parsedDomainName;
+  
+  const uint8_t* pos = buffer + sizeof(DNSHeader);
+  const uint8_t* end = buffer + packetSize;
+
+  // to minimize reallocations due to concats below
+  // we reserve enough space that a median or average domain
+  // name size cold be easily contained without a reallocation
+  // - max size would be 253, in 2013, average is 11 and max was 42
+  //
+  parsedDomainName.reserve(32); 
+
+  uint8_t labelLength = *pos;
+
+  while (true)
   {
-    unsigned char labelLength = *(start + pos);
-    for(int i = 0; i < labelLength; i++)
+    if (labelLength == 0)
+    {
+        // no more labels
+        downcaseAndRemoveWwwPrefix(parsedDomainName);
+        return parsedDomainName;
+    }
+
+    // append next label
+    for (int i = 0; i < labelLength && pos < end; i++)
     {
       pos++;
-      parsedDomainName += (char)*(start + pos);
+      parsedDomainName += static_cast<char>(*pos);
     }
-    pos++;
-    if (*(start + pos) == 0)
+
+    if (pos >= end)
     {
-      downcaseAndRemoveWwwPrefix(parsedDomainName);
-      return parsedDomainName;
+        // malformed packet, return an empty domain name
+        parsedDomainName = "";
+        return parsedDomainName;
     }
     else
     {
-      parsedDomainName += ".";
+        // next label
+        pos++;
+        labelLength = *pos;
+
+        // if there is another label, add delimiter
+        if (labelLength != 0)
+        {
+            parsedDomainName += ".";
+        }
     }
   }
 }
 
-void DNSServer::replyWithIP()
+void DNSServer::replyWithIP(uint8_t* buffer, size_t packetSize)
 {
-  if (_buffer == NULL) return;
-  _dnsHeader->QR = DNS_QR_RESPONSE;
-  _dnsHeader->ANCount = _dnsHeader->QDCount;
-  _dnsHeader->QDCount = _dnsHeader->QDCount; 
-  //_dnsHeader->RA = 1;  
+  DNSHeader* dnsHeader = reinterpret_cast<DNSHeader*>(buffer);
+
+  dnsHeader->QR = DNS_QR_RESPONSE;
+  dnsHeader->ANCount = dnsHeader->QDCount;
+  dnsHeader->QDCount = dnsHeader->QDCount; 
+  //dnsHeader->RA = 1;  
 
   _udp.beginPacket(_udp.remoteIP(), _udp.remotePort());
-  _udp.write(_buffer, _currentPacketSize);
+  _udp.write(buffer, packetSize);
 
   _udp.write((uint8_t)192); //  answer name is a pointer
   _udp.write((uint8_t)12);  // pointer to offset at 0x00c
@@ -142,22 +164,26 @@ void DNSServer::replyWithIP()
   _udp.write(_resolvedIP, sizeof(_resolvedIP));
   _udp.endPacket();
 
-
-
   #ifdef DEBUG_ESP_DNS
     DEBUG_ESP_PORT.printf("DNS responds: %s for %s\n",
-            IPAddress(_resolvedIP).toString().c_str(), getDomainNameWithoutWwwPrefix().c_str() );
+            IPAddress(_resolvedIP).toString().c_str(), getDomainNameWithoutWwwPrefix(buffer, packetSize).c_str() );
   #endif
 }
 
-void DNSServer::replyWithCustomCode()
+void DNSServer::replyWithCustomCode(uint8_t* buffer, size_t packetSize)
 {
-  if (_buffer == NULL) return;
-  _dnsHeader->QR = DNS_QR_RESPONSE;
-  _dnsHeader->RCode = (unsigned char)_errorReplyCode;
-  _dnsHeader->QDCount = 0;
+  if (packetSize < sizeof(DNSHeader))
+  {
+    return;
+  }
+
+  DNSHeader* dnsHeader = reinterpret_cast<DNSHeader*>(buffer);
+
+  dnsHeader->QR = DNS_QR_RESPONSE;
+  dnsHeader->RCode = (unsigned char)_errorReplyCode;
+  dnsHeader->QDCount = 0;
 
   _udp.beginPacket(_udp.remoteIP(), _udp.remotePort());
-  _udp.write(_buffer, sizeof(DNSHeader));
+  _udp.write(buffer, sizeof(DNSHeader));
   _udp.endPacket();
 }

Diff for: libraries/DNSServer/src/DNSServer.h

@@ -6,6 +6,8 @@
 #define DNS_QR_RESPONSE 1
 #define DNS_OPCODE_QUERY 0
 
+#define MAX_DNSNAME_LENGTH 253
+
 enum class DNSReplyCode
 {
   NoError = 0,
@@ -40,6 +42,9 @@ class DNSServer
 {
   public:
     DNSServer();
+    ~DNSServer() {
+        stop();
+    };
     void processNextRequest();
     void setErrorReplyCode(const DNSReplyCode &replyCode);
     void setTTL(const uint32_t &ttl);
@@ -56,16 +61,13 @@ class DNSServer
     uint16_t _port;
     String _domainName;
     unsigned char _resolvedIP[4];
-    int _currentPacketSize;
-    unsigned char* _buffer;
-    DNSHeader* _dnsHeader;
     uint32_t _ttl;
     DNSReplyCode _errorReplyCode;
 
     void downcaseAndRemoveWwwPrefix(String &domainName);
-    String getDomainNameWithoutWwwPrefix();
-    bool requestIncludesOnlyOneQuestion();
-    void replyWithIP();
-    void replyWithCustomCode();
+    String getDomainNameWithoutWwwPrefix(const uint8_t* buffer, size_t packetSize);
+    bool requestIncludesOnlyOneQuestion(const DNSHeader* dnsHeader);
+    void replyWithIP(uint8_t* buffer, size_t packetSize);
+    void replyWithCustomCode(uint8_t* buffer, size_t packetSize);
 };
 #endif