1  // SPDX-FileCopyrightText: 2022 s60sc with changes by ladyada
  2  // SPDX-License-Identifier: GPL-3.0-or-later
  3  
  4  #include "AdBlockerDNSServer.h"
  5  #include <lwip/def.h>
  6  #include <Arduino.h>
  7  
  8  
  9  //// ESP32_AdBlocker
 10  // call back to check if domain blocked and receive IP address to be used
 11  IPAddress checkBlocklist(const char* domainName);
 12  //// ESP32_AdBlocker
 13  
 14  DNSServer::DNSServer()
 15  {
 16    _ttl = htonl(DNS_DEFAULT_TTL);
 17    _errorReplyCode = DNSReplyCode::NonExistentDomain;
 18    _dnsHeader    = (DNSHeader*) malloc( sizeof(DNSHeader) ) ;
 19    _dnsQuestion  = (DNSQuestion*) malloc( sizeof(DNSQuestion) ) ;     
 20    _buffer     = NULL;
 21    _currentPacketSize = 0;
 22    _port = 0;
 23  }
 24  
 25  bool DNSServer::start(const uint16_t &port, const String &domainName,
 26                       const IPAddress &resolvedIP)
 27  {
 28    _port = port;
 29    _buffer = NULL;
 30    _domainName = domainName;
 31    _resolvedIP[0] = resolvedIP[0];
 32    _resolvedIP[1] = resolvedIP[1];
 33    _resolvedIP[2] = resolvedIP[2];
 34    _resolvedIP[3] = resolvedIP[3];
 35    downcaseAndRemoveWwwPrefix(_domainName);
 36    return _udp.begin(_port) == 1;
 37  }
 38  
 39  void DNSServer::setErrorReplyCode(const DNSReplyCode &replyCode)
 40  {
 41    _errorReplyCode = replyCode;
 42  }
 43  
 44  void DNSServer::setTTL(const uint32_t &ttl)
 45  {
 46    _ttl = htonl(ttl);
 47  }
 48  
 49  void DNSServer::stop()
 50  {
 51    _udp.stop();
 52    free(_buffer);
 53    _buffer = NULL;
 54  }
 55  
 56  void DNSServer::downcaseAndRemoveWwwPrefix(String &domainName)
 57  {
 58    domainName.toLowerCase();
 59    domainName.replace("www.", "");
 60  }
 61  
 62  void DNSServer::processNextRequest()
 63  {
 64    _currentPacketSize = _udp.parsePacket();
 65    if (_currentPacketSize)
 66    {
 67      // Allocate buffer for the DNS query
 68      if (_buffer != NULL) 
 69        free(_buffer);
 70      _buffer = (unsigned char*)malloc(_currentPacketSize * sizeof(char));
 71      if (_buffer == NULL) 
 72        return;
 73  
 74      // Put the packet received in the buffer and get DNS header (beginning of message)
 75      // and the question
 76      _udp.read(_buffer, _currentPacketSize);
 77      memcpy( _dnsHeader, _buffer, DNS_HEADER_SIZE ) ; 
 78      if ( requestIncludesOnlyOneQuestion() )
 79      {
 80        // The QName has a variable length, maximum 255 bytes and is comprised of multiple labels.
 81        // Each label contains a byte to describe its length and the label itself. The list of 
 82        // labels terminates with a zero-valued byte. In "github.com", we have two labels "github" & "com"
 83        // Iterate through the labels and copy them as they come into a single buffer (for simplicity's sake)
 84        _dnsQuestion->QNameLength = 0 ;
 85        while ( _buffer[ DNS_HEADER_SIZE + _dnsQuestion->QNameLength ] != 0 )
 86        {
 87          memcpy( (void*) &_dnsQuestion->QName[_dnsQuestion->QNameLength], (void*) &_buffer[DNS_HEADER_SIZE + _dnsQuestion->QNameLength], _buffer[DNS_HEADER_SIZE + _dnsQuestion->QNameLength] + 1 ) ;
 88          _dnsQuestion->QNameLength += _buffer[DNS_HEADER_SIZE + _dnsQuestion->QNameLength] + 1 ; 
 89        }
 90        _dnsQuestion->QName[_dnsQuestion->QNameLength] = 0 ; 
 91        _dnsQuestion->QNameLength++ ;   
 92  
 93        // Copy the QType and QClass 
 94        memcpy( &_dnsQuestion->QType, (void*) &_buffer[DNS_HEADER_SIZE + _dnsQuestion->QNameLength], sizeof(_dnsQuestion->QType) ) ;
 95        memcpy( &_dnsQuestion->QClass, (void*) &_buffer[DNS_HEADER_SIZE + _dnsQuestion->QNameLength + sizeof(_dnsQuestion->QType)], sizeof(_dnsQuestion->QClass) ) ;
 96      }
 97      
 98  
 99      if (_dnsHeader->QR == DNS_QR_QUERY &&
100          _dnsHeader->OPCode == DNS_OPCODE_QUERY &&
101          requestIncludesOnlyOneQuestion() &&
102          (_domainName == "*" || getDomainNameWithoutWwwPrefix() == _domainName)
103         )
104      {
105        //// ESP32_AdBlocker
106        IPAddress IPtoUse = checkBlocklist(getDomainNameWithoutWwwPrefix().c_str());
107        for (int i=0; i<4; i++) _resolvedIP[i] = IPtoUse[i];
108        //// ESP32_AdBlocker
109        
110        replyWithIP();
111      }
112      else if (_dnsHeader->QR == DNS_QR_QUERY)
113      {
114        replyWithCustomCode();
115      }
116  
117      free(_buffer);
118      _buffer = NULL;
119    }
120  }
121  
122  bool DNSServer::requestIncludesOnlyOneQuestion()
123  {
124    return ntohs(_dnsHeader->QDCount) == 1 &&
125           _dnsHeader->ANCount == 0 &&
126           _dnsHeader->NSCount == 0 &&
127           _dnsHeader->ARCount == 0;
128  }
129  
130  
131  String DNSServer::getDomainNameWithoutWwwPrefix()
132  {
133    // Error checking : if the buffer containing the DNS request is a null pointer, return an empty domain
134    String parsedDomainName = "";
135    if (_buffer == NULL) 
136      return parsedDomainName;
137    
138    // Set the start of the domain just after the header (12 bytes). If equal to null character, return an empty domain
139    unsigned char *start = _buffer + DNS_OFFSET_DOMAIN_NAME;
140    if (*start == 0)
141    {
142      return parsedDomainName;
143    }
144  
145    int pos = 0;
146    while(true)
147    {
148      unsigned char labelLength = *(start + pos);
149      for(int i = 0; i < labelLength; i++)
150      {
151        pos++;
152        parsedDomainName += (char)*(start + pos);
153      }
154      pos++;
155      if (*(start + pos) == 0)
156      {
157        downcaseAndRemoveWwwPrefix(parsedDomainName);
158        return parsedDomainName;
159      }
160      else
161      {
162        parsedDomainName += ".";
163      }
164    }
165  }
166  
167  void DNSServer::replyWithIP()
168  {
169    if (_buffer == NULL) return;
170    
171    _udp.beginPacket(_udp.remoteIP(), _udp.remotePort());
172    
173    // Change the type of message to a response and set the number of answers equal to 
174    // the number of questions in the header
175    _dnsHeader->QR      = DNS_QR_RESPONSE;
176    _dnsHeader->ANCount = _dnsHeader->QDCount;
177    _udp.write( (unsigned char*) _dnsHeader, DNS_HEADER_SIZE ) ;
178  
179    // Write the question
180    _udp.write(_dnsQuestion->QName, _dnsQuestion->QNameLength) ;
181    _udp.write( (unsigned char*) &_dnsQuestion->QType, 2 ) ;
182    _udp.write( (unsigned char*) &_dnsQuestion->QClass, 2 ) ;
183  
184    // Write the answer 
185    // Use DNS name compression : instead of repeating the name in this RNAME occurence,
186    // set the two MSB of the byte corresponding normally to the length to 1. The following
187    // 14 bits must be used to specify the offset of the domain name in the message 
188    // (<255 here so the first byte has the 6 LSB at 0) 
189    _udp.write((uint8_t) 0xC0); 
190    _udp.write((uint8_t) DNS_OFFSET_DOMAIN_NAME);  
191  
192    // DNS type A : host address, DNS class IN for INternet, returning an IPv4 address 
193    uint16_t answerType = htons(DNS_TYPE_A), answerClass = htons(DNS_CLASS_IN), answerIPv4 = htons(DNS_RDLENGTH_IPV4)  ; 
194    _udp.write((unsigned char*) &answerType, 2 );
195    _udp.write((unsigned char*) &answerClass, 2 );
196    _udp.write((unsigned char*) &_ttl, 4);        // DNS Time To Live
197    _udp.write((unsigned char*) &answerIPv4, 2 );
198    _udp.write(_resolvedIP, sizeof(_resolvedIP)); // The IP address to return
199    _udp.endPacket();
200  
201    #ifdef DEBUG_ESP_DNS
202      DBG_OUTPUT_PORT.printf("DNS responds: %s for %s\n",
203              IPAddress(_resolvedIP).toString().c_str(), getDomainNameWithoutWwwPrefix().c_str() );
204    #endif  
205  }
206  
207  void DNSServer::replyWithCustomCode()
208  {
209    if (_buffer == NULL) return;
210    _dnsHeader->QR = DNS_QR_RESPONSE;
211    _dnsHeader->RCode = (unsigned char)_errorReplyCode;
212    _dnsHeader->QDCount = 0;
213  
214    _udp.beginPacket(_udp.remoteIP(), _udp.remotePort());
215    _udp.write(_buffer, sizeof(DNSHeader));
216    _udp.endPacket();
217  }