1  // SPDX-FileCopyrightText: 2022 s60sc with changes by ladyada
  2  // SPDX-License-Identifier: GPL-3.0-or-later
  3  
  4   /*
  5     ESP32_AdBlocker acts as a DNS Sinkhole by returning 0.0.0.0 for any domain names in its blocked list, 
  6     else forwards to an external DNS server to resolve IP addresses. This prevents content being retrieved 
  7     from or sent to blocked domains. Searches generally take <200us.
  8  
  9     To use ESP32_AdBlocker, enter its IP address in place of the DNS server IPs in your router/devices.
 10     Currently it does have an IPv6 address and some devices use IPv6 by default, so disable IPv6 DNS on 
 11     your device / router to force it to use IPv4 DNS.
 12  
 13     Blocklist files can downloaded from hosting sites and should either be in HOSTS format 
 14     or Adblock format (only domain name entries processed)
 15  
 16     arduino-esp32 library DNSServer.cpp modified as custom AdBlockerDNSServer.cpp so that DNSServer::processNextRequest()
 17     calls checkBlocklist() in ESP32_AdBlocker to check if domain blocked, which returns the relevant IP. 
 18     Based on idea from https://github.com/rubfi/esphole
 19  
 20     Compile with partition scheme 'No OTA (2M APP/2M SPIFFS)'
 21  
 22     s60sc 2021
 23  */
 24  
 25  #include <WiFi.h>
 26  #include <string>
 27  #include <algorithm>
 28  #include <bits/stdc++.h>
 29  //#include <unordered_map>
 30  #include <WiFiClientSecure.h>
 31  #include <HTTPClient.h>
 32  //#include <ESPAsyncWebServer.h> // https://github.com/me-no-dev/ESPAsyncWebServer & https://github.com/me-no-dev/AsyncTCP
 33  #include <Preferences.h>
 34  
 35  #include <Adafruit_SPIFlash.h>
 36  #include <Adafruit_TinyUSB.h>
 37  #include <Adafruit_ST7789.h>
 38  #include <Fonts/FreeSans9pt7b.h>
 39  
 40  #include "AdBlockerDNSServer.h" // customised
 41  
 42  
 43  #include "config.h"
 44  
 45  //static AsyncWebServer webServer(80);
 46  static DNSServer dnsServer;
 47  
 48  //std::unordered_set <std::string> blockSet ; // holds hashed list of blocked domains in memory [too much memory overhead]
 49  std::vector<std::string> blockVec; // holds sorted list of blocked domains in memory
 50  static uint8_t* downloadBuff; // temporary store for ptocessing downloaded blocklist file
 51  static const char* BLOCKLISTFILE = {"/hosts.txt"};
 52  static uint32_t blockCnt, allowCnt = 0;
 53  
 54  Adafruit_ST7789 display = Adafruit_ST7789(TFT_CS, TFT_DC, TFT_RST);
 55  Adafruit_FlashTransport_ESP32 flashTransport;  // internal SPI flash access
 56  Adafruit_SPIFlash flash(&flashTransport); 
 57  FatFileSystem fatfs;  // file system object from SdFat
 58  
 59  char ssid[80] = DEFAULT_SSID;
 60  char password[80] = DEFAULT_PASSWORD;
 61  char hostname[80] = DEFAULT_HOSTNAME;
 62  char hostfile[255] = DEFAULT_HOSTFILE;
 63  
 64  
 65  bool init_filesystem(void);
 66  bool parseSecrets();
 67  
 68  static float inline inKB(size_t inVal) {
 69    return (float)(inVal / 1024.0);
 70  }
 71  
 72  void setup() {
 73  #if !defined(EXPOSE_FS_ON_MSD)  
 74    DBG_OUTPUT_PORT.begin(115200);
 75    while (!DBG_OUTPUT_PORT) delay(10);
 76    delay(1000);
 77  #endif
 78  
 79    if (!flash.begin()) {
 80      DBG_OUTPUT_PORT.println("failed to load flash");
 81      display.setTextColor(ST77XX_RED);
 82      display.println("Failed to load flash");
 83      while (1) yield();
 84    }  
 85  
 86    if (! init_filesystem()) {
 87      DBG_OUTPUT_PORT.printf("Aborting as no filesystem available");
 88      display.setTextColor(ST77XX_RED);
 89      display.println("Failed to load filesys");
 90      while (1) yield();
 91    }
 92    
 93    DBG_OUTPUT_PORT.printf("Sketch size %0.1fKB, PSRAM %s available\n", inKB(ESP.getSketchSize()), (psramFound()) ? "IS" : "NOT");
 94  
 95  
 96    // turn on TFT by default
 97    pinMode(TFT_I2C_POWER, OUTPUT);
 98    digitalWrite(TFT_I2C_POWER, HIGH);
 99    pinMode(TFT_BACKLITE, OUTPUT);
100    digitalWrite(TFT_BACKLITE, LOW);
101    delay(10);
102    display.init(135, 240);           // Init ST7789 240x135
103    display.setRotation(3);
104    display.fillScreen(ST77XX_BLACK);
105    digitalWrite(TFT_BACKLITE, HIGH);
106  
107    display.setTextColor(ST77XX_WHITE); 
108    display.setCursor(0, 0);
109    display.setTextSize(3);
110    display.print("ESPHole");
111  
112    display.setFont(&FreeSans9pt7b);
113    display.setTextSize(1);
114    display.setTextColor(ST77XX_WHITE); 
115    display.setCursor(0, 40);
116  
117  
118    parseSecrets();  
119    startWifi();
120    getNTP();
121  
122    bool haveBlocklist = fatfs.exists(BLOCKLISTFILE);
123    if (!haveBlocklist) {
124      DBG_OUTPUT_PORT.println("No blocklist stored, need to download ...");
125      display.println("Downloading blocklist");
126      if (downloadFile()) 
127        haveBlocklist = createBlocklist();
128    } else {
129      haveBlocklist = loadBlocklist();
130    }
131  
132    if (haveBlocklist) {
133      // DNS server startup
134      dnsServer.setErrorReplyCode(DNSReplyCode::ServerFailure);
135      if (dnsServer.start(DNS_PORT, "*", WiFi.localIP())) {
136        DBG_OUTPUT_PORT.printf("\nDNS Server started on %s:%d\n", WiFi.localIP().toString().c_str(), DNS_PORT);
137      } else {
138        puts("Aborting as DNS Server not running");
139        return;
140      }
141    } else {
142      puts("Aborting as no resources available");
143      return;
144    }
145    DBG_OUTPUT_PORT.printf("Free DRAM: %0.1fKB, Free PSRAM: %0.1fKB\n*** Ready...\n\n", inKB(ESP.getFreeHeap()), inKB(ESP.getFreePsram()));
146  }
147  
148  uint32_t timestamp = 0;
149  void loop() {
150    dnsServer.processNextRequest();
151    if (((timestamp + 2000) < millis()) || (timestamp > millis())){
152      Serial.print(".");
153      timestamp = millis();
154  #if !defined(DEBUG_ESP_DNS)
155      display.setFont(&FreeSans9pt7b);
156      display.setTextSize(2);
157      display.setTextColor(ST77XX_RED);
158      display.fillRect(0, 90, 240, 55, ST77XX_BLACK);
159      display.setCursor(0, 118);
160      display.print(blockCnt);
161      display.println(" Nahs!");
162  #endif
163  
164    }
165    //checkAlarm();
166  }
167  
168  static void startWifi() {
169    WiFi.mode(WIFI_AP_STA);
170    DBG_OUTPUT_PORT.printf("Connecting to SSID: %s\n", ssid);
171    //WiFi.config(ADBLOCKER, GATEWAY, SUBNET, RESOLVER);
172    WiFi.begin(ssid, password);
173    display.print("Conn'ing to ");
174    display.print(ssid);
175    display.print("...");
176    
177    while (WiFi.status() != WL_CONNECTED) {
178      delay(500);
179      DBG_OUTPUT_PORT.printf(".");
180    }
181    DBG_OUTPUT_PORT.println("");
182  
183    display.println("OK!");
184    DBG_OUTPUT_PORT.print("Connected! IP address: ");
185    DBG_OUTPUT_PORT.println(WiFi.localIP());
186    display.print("IP addr: ");
187    display.println(WiFi.localIP());
188  }
189  
190  IPAddress checkBlocklist(const char* domainName) {
191    // called from ESP32_DNSServer
192    IPAddress ip = IPAddress(0, 0, 0, 0); // IP to return for blocked domain
193    if (strlen(domainName)) {
194      uint64_t uselapsed = micros();
195      // search for domain in blocklist
196      //bool blocked = (blockSet.find(std::string(domainName)) == blockSet.end()) ? false : true;
197      bool blocked = binary_search(blockVec.begin(), blockVec.end(), std::string(domainName));
198      uint64_t checkTime = micros() - uselapsed;
199  
200      uint32_t mselapsed = millis();
201      // if not blocked, get IP address for domain from external DNS
202      if (!blocked)  WiFi.hostByName(domainName, ip);
203      uint32_t resolveTime = millis() - mselapsed;
204  
205      
206  #ifdef DEBUG_ESP_DNS
207      display.setFont(&FreeSans9pt7b);
208      display.setTextSize(1);
209      display.setTextColor(ST77XX_WHITE);
210      display.setTextWrap(true);
211      display.fillRect(0, 90, 240, 55, ST77XX_BLACK);
212      display.setCursor(0, 105);
213      display.println(domainName);
214  #endif
215      DBG_OUTPUT_PORT.printf("%s %s in %lluus", (blocked) ? "*Blocked*" : "Allowed", domainName, checkTime);
216      if (!blocked) {
217        DBG_OUTPUT_PORT.printf(", resolved to %s in %ums\n", ip.toString().c_str(), resolveTime);
218  #ifdef DEBUG_ESP_DNS
219        display.setTextColor(ST77XX_GREEN);
220        display.println(ip.toString().c_str());
221  #endif
222      }
223      else {
224        puts("");
225  #ifdef DEBUG_ESP_DNS
226        display.setTextColor(ST77XX_RED);
227        display.println("BLOCKED");
228  #endif
229      }
230  
231      if (blocked) blockCnt++;
232      else allowCnt++;
233    }
234    return ip;
235  }
236  
237  static bool loadBlocklist() {
238    // load blocklist file into memory from storage
239    size_t remaining = heap_caps_get_largest_free_block(MALLOC_CAP_8BIT);
240    DBG_OUTPUT_PORT.printf("Loading blocklist %s\n", hostfile);
241    uint32_t loadTime = millis();
242    if (psramFound()) heap_caps_malloc_extmem_enable(5); // force vector into psram
243    blockVec.reserve((psramFound()) ? MAX_DOMAINS : 1000);
244    if (psramFound()) heap_caps_malloc_extmem_enable(100000);
245    
246    File file = fatfs.open(BLOCKLISTFILE, FILE_READ);
247    if (file) {
248      DBG_OUTPUT_PORT.printf("File size %0.1fKB loading into %0.1fKB %s memory ...\n", inKB(file.size()),
249        inKB(remaining), (psramFound()) ? "PSRAM" : "DRAM");
250      char domainLine[MAX_LINELEN + 1];
251      int itemsLoaded = 0;
252      if (psramFound()) heap_caps_malloc_extmem_enable(5); // force vector into psram
253  
254      while (file.available()) {
255        size_t lineLen = file.readBytesUntil('\n', domainLine, MAX_LINELEN);
256        if (lineLen) {
257          domainLine[lineLen] = 0;
258          //blockSet.insert(std::string(domainLine));
259          blockVec.push_back(std::string(domainLine));
260          itemsLoaded++;
261          if (itemsLoaded%500 == 0) { // check memory too often triggers watchdog
262            remaining = heap_caps_get_largest_free_block(MALLOC_CAP_8BIT);
263            if (remaining < (MIN_MEMORY)) {
264              DBG_OUTPUT_PORT.printf("Blocklist truncated to avoid memory overflow, %u bytes remaining\n", remaining);
265              break;
266            }
267          }
268          if (itemsLoaded >= MAX_DOMAINS ) {
269            // max 64K vectors else esp32 crashes
270            DBG_OUTPUT_PORT.printf("Blocklist truncated as maximum number of domains loaded - %u\n", MAX_DOMAINS);
271            break;
272          }
273        }
274      }
275      file.close();
276      //for (int i=0; i < blockVec.size(); i++) DBG_OUTPUT_PORT.printf("%s\n", blockVec.at(i).c_str());
277  
278      if (psramFound()) heap_caps_malloc_extmem_enable(100000);
279      if (!itemsLoaded) {
280        puts("Aborting as empty blocklist read ..."); // SPIFFS issue
281        delay(100);
282        ESP.restart();
283      }
284      display.print("Checking ");
285      display.print(itemsLoaded);
286      display.println(" domains");
287        
288      DBG_OUTPUT_PORT.printf("Loaded %u domains from blocklist file in %.01f secs\n", itemsLoaded, (float)((millis() - loadTime) / 1000.0));
289      return true;
290    }
291    puts(" - not found");
292    return false;
293  }
294  
295  static bool sortBlocklist() {
296    // read downloaded blocklist from storage, sort in alpha order to allow binary search, rewrite to storage
297    free(downloadBuff);
298    delay(100);
299    if (loadBlocklist()) {
300      puts("Sorting blocklist alphabetically ...");
301      uint32_t sortTime = millis();
302      sort(blockVec.begin(), blockVec.end());
303      DBG_OUTPUT_PORT.printf("Sorted blocklist after %0.1f secs, saving to file ...\n", (float)((millis() - sortTime) / 1000.0));
304      sortTime = millis();
305      std::string previous = "";
306      int duplicate = 0;
307      
308      // rewrite file with sorted domains
309      fatfs.remove(BLOCKLISTFILE);
310      File file = fatfs.open(BLOCKLISTFILE, FILE_WRITE);
311      if (file) {
312        for (auto domain : blockVec) {
313          if (domain.compare(previous) != 0) {
314            // store unduplicated domain
315            file.write((uint8_t*)domain.c_str(), strlen(domain.c_str()));
316            file.write((uint8_t*)"\n", 1);
317            previous.assign(domain);
318          } else duplicate++;
319        }
320        file.close();
321        DBG_OUTPUT_PORT.printf("Saved into file removing %u duplicates after %0.1f secs, restarting ...\n", duplicate, (float)((millis() - sortTime) / 1000.0));
322        delay(100);
323        ESP.restart(); // quicker to restart than clear vector
324        return true;
325      } else puts ("Failed to write to blocklist file");
326    }
327    return false;
328  }
329  
330  static size_t inline extractDomains() {
331    // extract domain names from downloaded blocklist
332    size_t downloadBuffPtr = 0;
333    char *saveLine, *saveItem = NULL;
334    char* tokenLine = strtok_r((char*)downloadBuff, "\n", &saveLine);
335    char* tokenItem;
336  
337    // for each line
338    while (tokenLine != NULL) {
339      if (strncmp(tokenLine, "127.0.0.1", 9) == 0 || strncmp(tokenLine, "0.0.0.0", 7) == 0) {
340        // HOSTS file format matched, extract domain name
341        tokenItem = strtok_r(tokenLine, " \t", &saveItem); // skip over first token
342        if (tokenItem != NULL) tokenItem = strtok_r(NULL, " \t", &saveItem); // domain in second token
343      } else if (strncmp(tokenLine, "||", 2) == 0)
344        tokenItem = strtok_r(tokenLine, "|^", &saveItem); // Adblock format - domain in first token
345      else tokenItem = NULL; // no match
346      if (tokenItem != NULL) {
347        // write processed line back to buffer
348        size_t itemLen = strlen(tokenItem);
349        int wwwOffset = (strncmp(tokenItem, "www.", 4) == 0) ? 4 : 0;  // remove any leading "www."
350        memcpy(downloadBuff + downloadBuffPtr, tokenItem + wwwOffset, itemLen - wwwOffset);
351        downloadBuffPtr += itemLen;
352        memcpy(downloadBuff + downloadBuffPtr, (uint8_t*)"\n", 1);
353        downloadBuffPtr++;
354      }
355      tokenLine = strtok_r(NULL, "\n", &saveLine);
356    }
357    downloadBuff[downloadBuffPtr] = 0; // string terminator
358    return downloadBuffPtr;
359  }
360  
361  static bool createBlocklist() {
362    // after blocklist file downloaded, the content is parsed to extract the domain names, then written to storage
363    uint32_t createTime = millis();
364    size_t blocklistSize = extractDomains();
365              
366    // check storage space available, else abort
367    size_t storageAvailable = flashFreeSpace();
368    storageAvailable -= 1024*32; // leave overhead space
369    if (storageAvailable < blocklistSize) {
370      DBG_OUTPUT_PORT.printf("Aborting as insufficient storage %0.1fKB ...\n", inKB(storageAvailable));
371      delay(100);
372      ESP.restart();
373    }
374    DBG_OUTPUT_PORT.printf("Creating extracted unsorted blocklist of %0.1fKB\n", inKB(blocklistSize));
375  
376  
377    Serial.println("-----------------------------------------------");
378    Serial.println((const char *)downloadBuff);
379    Serial.println("-----------------------------------------------");
380  
381    File file = fatfs.open(BLOCKLISTFILE, FILE_WRITE);
382    if (file) {
383      // write buffer to file
384      size_t written = file.write(downloadBuff, blocklistSize);
385      if (!written) {
386        puts("Aborting as blocklist empty after writing ..."); // SPIFFS issue
387        delay(100);
388        ESP.restart();
389      }
390      file.close();
391      DBG_OUTPUT_PORT.printf("Blocklist file of %0.1fKB created in %.01f secs\n", inKB(written), (float)((millis() - createTime) / 1000.0));
392      return (sortBlocklist());
393    } else DBG_OUTPUT_PORT.printf("Failed to store blocklist %s\n", BLOCKLISTFILE);
394    free(downloadBuff);
395    return false;
396  }
397  
398  static bool downloadFile() {
399    size_t downloadBuffPtr = 0;
400    WiFiClientSecure *client = new WiFiClientSecure;
401    if (!client) {
402      puts("Failed to create secure client");
403      return false;
404    } else {
405      // scoping block
406      {
407        client->setInsecure(); // don't use a root cert
408        HTTPClient http;
409        // open connection to blocklist host
410        if (http.begin(*client, hostfile)) {
411          DBG_OUTPUT_PORT.printf("Downloading %s\n", hostfile);
412          uint32_t loadTime = millis();
413          // start connection and send HTTP header
414          int httpCode = http.GET();
415  
416          if (httpCode > 0) {
417            DBG_OUTPUT_PORT.printf("Response code: %d\n", httpCode);
418            if (httpCode == HTTP_CODE_OK || httpCode == HTTP_CODE_MOVED_PERMANENTLY) {
419              // get length of content (is -1 when Server sends no Content-Length header)
420              int len = http.getSize();
421              if (len > 0) {
422                DBG_OUTPUT_PORT.printf("File size: %0.1fKB", inKB(len));
423              }
424              else {
425                DBG_OUTPUT_PORT.printf("File size unknown");
426              }
427              size_t availableMem = heap_caps_get_largest_free_block(MALLOC_CAP_8BIT);
428              DBG_OUTPUT_PORT.printf(", with %0.1fKB memory available for download ...\n", inKB(availableMem));
429              if (len > 0 && len > availableMem) {
430                puts("Aborting as file too large");
431                delay(100);
432                ESP.restart();
433              }
434  
435              int chunk = 128; // amount to consume per read of stream
436              // allocate memory for download buffer
437              downloadBuff = (psramFound()) ? (uint8_t*)ps_malloc(availableMem) : (uint8_t*)malloc(availableMem);
438              WiFiClient * stream = http.getStreamPtr(); // stream data to client
439  
440              while (http.connected() && (len > 0 || len == -1)) {
441                size_t streamSize = stream->available();
442                if (streamSize) {
443                  // consume up to chunk bytes and write to memory
444                  int readc = stream->readBytes(downloadBuff + downloadBuffPtr, ((streamSize > chunk) ? chunk : streamSize));
445                  downloadBuffPtr += readc;
446                  if (len > 0) len -= readc;
447                  if ((downloadBuffPtr + chunk) >= availableMem) {
448                    puts("Aborting as file too large");
449                    delay(100);
450                    ESP.restart();
451                  }
452                }
453                delay(1);
454              }
455              downloadBuff[downloadBuffPtr] = 0; // string terminator
456              Serial.println("-----------------------------------------------");
457              Serial.println((const char *)downloadBuff);
458              Serial.println("-----------------------------------------------");
459              DBG_OUTPUT_PORT.printf("Download completed in %0.1f secs, stored %0.1fKB\n", ((millis() - loadTime) / 1000.0), inKB(downloadBuffPtr));
460            }
461          } else DBG_OUTPUT_PORT.printf("Connection failed with error: %s\n", http.errorToString(httpCode).c_str());
462        } else {
463          DBG_OUTPUT_PORT.printf("Unable to download %s\n", hostfile);
464          delay(100);
465          ESP.restart();
466        }
467        http.end();
468      }
469      delete client;
470    }
471    return (downloadBuffPtr) ? true : false;
472  }
473  
474  
475  static inline time_t getEpochSecs() {
476    struct timeval tv;
477    gettimeofday(&tv, NULL);
478    return tv.tv_sec;
479  }
480  
481  static void getNTP() {
482    // get current time from NTP server and apply to ESP32
483    const char* ntpServer = "0.adafruit.pool.ntp.org";
484    const long gmtOffset_sec = 0;  // offset from GMT
485    const int daylightOffset_sec = 3600; // daylight savings offset in secs
486    int i = 0;
487    do {
488      configTime(gmtOffset_sec, daylightOffset_sec, ntpServer);
489      delay(1000);
490    } while (getEpochSecs() < 1000 && i++ < 5); // try up to 5 times
491    // set timezone as required
492    setenv("TZ", TIMEZONE, 1);
493    if (getEpochSecs() > 1000) {
494      time_t currEpoch = getEpochSecs();
495      char timeFormat[20];
496      strftime(timeFormat, sizeof(timeFormat), "%d/%m/%Y %H:%M:%S", localtime(&currEpoch));
497      DBG_OUTPUT_PORT.printf("Got current time from NTP: %s\n", timeFormat);
498    }
499    else puts("Unable to sync with NTP");
500  }
501  
502  static void checkAlarm() {
503    // once per day at given time, load updated blocklist from host site
504    static time_t rolloverEpoch = setAlarm(UPDATE_HOUR);
505    if (getEpochSecs() >= rolloverEpoch) { 
506      puts("Scheduled restart to load updated blocklist ...");
507      delay(100);
508      ESP.restart();
509    }
510  }
511  
512  static time_t setAlarm(uint8_t alarmHour) {
513    // calculate future alarm datetime based on current datetime
514    struct tm* timeinfo;
515    time_t rawtime;
516    time(&rawtime);
517    timeinfo = localtime(&rawtime);
518    // set alarm date & time for next day at given hour
519    timeinfo->tm_mday += 1; 
520    timeinfo->tm_hour = alarmHour;
521    timeinfo->tm_min = 0;
522    // return future datetime as epoch seconds
523    return mktime(timeinfo);
524  }