Rework MDNS parsing to support longer domains and use progmem

Author: sandeepmistry

examples/MDNS_WiFiWebServer/MDNS_WiFiWebServer.ino

@@ -94,7 +94,7 @@ void setup() {
 void loop() {
   // Call the update() function on the MDNS responder every loop iteration to
   // make sure it can detect and respond to name requests.
-  mdns.update();
+  mdns.poll();
 
   // listen for incoming clients
   WiFiClient client = server.available();

keywords.txt

@@ -37,6 +37,7 @@ RSSI	KEYWORD2
 encryptionType	KEYWORD2
 getResult	KEYWORD2
 getSocket	KEYWORD2
+poll	KEYWORD2
 WiFiClient	KEYWORD2
 WiFiServer	KEYWORD2
 WiFiSSLClient	KEYWORD2

src/WiFiMdns.cpp

@@ -21,177 +21,189 @@
 // License along with this library; if not, write to the Free Software
 // Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 
+#include <avr/pgmspace.h>
+
 #include "Arduino.h"
 #include "WiFiMdns.h"
 
 // Important RFC's for reference:
 // - DNS request and response: http://www.ietf.org/rfc/rfc1035.txt
 // - Multicast DNS: http://www.ietf.org/rfc/rfc6762.txt
 
-#define READ_BUFFER_SIZE 20
 #define HEADER_SIZE 12
-#define QDCOUNT_OFFSET 4
-#define A_RECORD_SIZE 14
-#define NSEC_RECORD_SIZE 20
 #define TTL_OFFSET 4
 #define IP_OFFSET 10
 
+const uint8_t expectedRequestHeader[HEADER_SIZE] PROGMEM = {
+  0x00, 0x00,
+  0x00, 0x00,
+  0x00, 0x01,
+  0x00, 0x00,
+  0x00, 0x00,
+  0x00, 0x00
+};
 
-MDNSResponder::MDNSResponder()
-  : _expected(NULL)
-  , _expectedLen(0)
-  , _index(0)
-  , _response(NULL)
-  , _responseLen(0)
-  , _mdnsSocket()
-{ }
+const uint8_t responseHeader[] PROGMEM = {
+  0x00, 0x00,   // ID = 0
+  0x84, 0x00,   // Flags = response + authoritative answer
+  0x00, 0x00,   // Question count = 0
+  0x00, 0x01,   // Answer count = 1
+  0x00, 0x00,   // Name server records = 0
+  0x00, 0x01    // Additional records = 1
+};
 
-MDNSResponder::~MDNSResponder() {
-  if (_expected != NULL) {
-    free(_expected);
-  }
-  if (_response != NULL) {
-    free(_response);
-  }
+// Generate positive response for IPV4 address
+const uint8_t aRecord[] PROGMEM = {
+  0x00, 0x01,                // Type = 1, A record/IPV4 address
+  0x80, 0x01,                // Class = Internet, with cache flush bit
+  0x00, 0x00, 0x00, 0x00,    // TTL in seconds, to be filled in later
+  0x00, 0x04,                // Length of record
+  0x00, 0x00, 0x00, 0x00     // IP address, to be filled in later
+};
+
+// Generate negative response for IPV6 address (CC3000 doesn't support IPV6)
+const uint8_t nsecRecord[] PROGMEM = { 
+  0xC0, 0x0C,                // Name offset
+  0x00, 0x2F,                // Type = 47, NSEC (overloaded by MDNS)
+  0x80, 0x01,                // Class = Internet, with cache flush bit
+  0x00, 0x00, 0x00, 0x00,    // TTL in seconds, to be filled in later
+  0x00, 0x08,                // Length of record
+  0xC0, 0x0C,                // Next domain = offset to FQDN
+  0x00,                      // Block number = 0
+  0x04,                      // Length of bitmap = 4 bytes
+  0x40, 0x00, 0x00, 0x00     // Bitmap value = Only first bit (A record/IPV4) is set
+};
+
+const uint8_t domain[] PROGMEM = { 'l', 'o', 'c', 'a', 'l' };
+
+MDNSResponder::MDNSResponder() :
+  expectedRequestLength(0)
+{
 }
 
-bool MDNSResponder::begin(const char* domain, uint32_t ttlSeconds)
+MDNSResponder::~MDNSResponder()
 {
-  // Construct DNS request/response fully qualified domain name of form:
-  // <domain length>, <domain characters>, 5, "local"
-  size_t n = strlen(domain);
-  if (n > 255) {
-    // Can only handle domains that are 255 chars in length.
-    return false;
-  }
-  _expectedLen = 12 + n;
-  if (_expected != NULL) {
-    free(_expected);
-  }
-  _expected = (uint8_t*) malloc(_expectedLen);
-  if (_expected == NULL) {
-    return false;
-  }
-  _expected[0] = (uint8_t)n;
-  // Copy in domain characters as lowercase
-  for (unsigned int i = 0; i < n; ++i) {
-    _expected[1+i] = tolower(domain[i]);
-  }
-  // Values for:
-  //  - 5 (length)
-  //  - "local"
-  //  - 0x00 (end of domain)
-  //  - 0x00 0x01 (A record query)
-  //  - 0x00 0x01 (Class IN)
-  uint8_t local[] = { 0x05, 0x6C, 0x6F, 0x63, 0x61, 0x6C, 0x00, 0x00, 0x01, 0x00, 0x01 };
-  memcpy(&_expected[1+n], local, 11);
-
-  // Construct DNS query response
-  // TODO: Move these to flash or just construct in code.
-  uint8_t respHeader[] = { 0x00, 0x00,   // ID = 0
-                           0x84, 0x00,   // Flags = response + authoritative answer
-                           0x00, 0x00,   // Question count = 0
-                           0x00, 0x01,   // Answer count = 1
-                           0x00, 0x00,   // Name server records = 0
-                           0x00, 0x01    // Additional records = 1
-  };
-  // Generate positive response for IPV4 address
-  uint8_t aRecord[] = { 0x00, 0x01,                // Type = 1, A record/IPV4 address
-                        0x80, 0x01,                // Class = Internet, with cache flush bit
-                        0x00, 0x00, 0x00, 0x00,    // TTL in seconds, to be filled in later
-                        0x00, 0x04,                // Length of record
-                        0x00, 0x00, 0x00, 0x00     // IP address, to be filled in later
-  };
-  // Generate negative response for IPV6 address (CC3000 doesn't support IPV6)
-  uint8_t nsecRecord[] = {  0xC0, 0x0C,                // Name offset
-                            0x00, 0x2F,                // Type = 47, NSEC (overloaded by MDNS)
-                            0x80, 0x01,                // Class = Internet, with cache flush bit
-                            0x00, 0x00, 0x00, 0x00,    // TTL in seconds, to be filled in later
-                            0x00, 0x08,                // Length of record
-                            0xC0, 0x0C,                // Next domain = offset to FQDN
-                            0x00,                      // Block number = 0
-                            0x04,                      // Length of bitmap = 4 bytes
-                            0x40, 0x00, 0x00, 0x00     // Bitmap value = Only first bit (A record/IPV4) is set
-  };
-  // Allocate memory for response.
-  int queryFQDNLen = _expectedLen - 4;
-  _responseLen = HEADER_SIZE + queryFQDNLen + A_RECORD_SIZE + NSEC_RECORD_SIZE;
-  if (_response != NULL) {
-    free(_response);
-  }
-  _response = (uint8_t*) malloc(_responseLen);
-  if (_response == NULL) {
+}
+
+bool MDNSResponder::begin(const char* _name, uint32_t _ttlSeconds)
+{
+  int nameLength = strlen(_name);
+
+  if (nameLength > 255) {
+    // Can only handle domains that are upto 255 chars in length.
+    expectedRequestLength = 0;
     return false;
   }
-  // Copy data into response.
-  memcpy(_response, respHeader, HEADER_SIZE);
-  memcpy(_response + HEADER_SIZE, _expected, queryFQDNLen);
-  uint8_t* records = _response + HEADER_SIZE + queryFQDNLen;
-  memcpy(records, aRecord, A_RECORD_SIZE);
-  memcpy(records + A_RECORD_SIZE, nsecRecord, NSEC_RECORD_SIZE);
-  // Add TTL to records.
-  uint8_t ttl[4] = { (uint8_t)(ttlSeconds >> 24), (uint8_t)(ttlSeconds >> 16), (uint8_t)(ttlSeconds >> 8), (uint8_t)ttlSeconds };
-  memcpy(records + TTL_OFFSET, ttl, 4);
-  memcpy(records + A_RECORD_SIZE + 2 + TTL_OFFSET, ttl, 4);
-  // Add IP address to response
-  uint32_t ipAddress = WiFi.localIP();
-  records[IP_OFFSET]     = (uint8_t) ipAddress;
-  records[IP_OFFSET + 1] = (uint8_t)(ipAddress >> 8);
-  records[IP_OFFSET + 2] = (uint8_t)(ipAddress >> 16);
-  records[IP_OFFSET + 3] = (uint8_t)(ipAddress >> 24);
+
+  name = _name;
+  ttlSeconds = _ttlSeconds;
+
+  name.toLowerCase();
+  expectedRequestLength = HEADER_SIZE + 1 + nameLength + 1 + sizeof(domain) + 5;
 
   // Open the MDNS UDP listening socket on port 5353 with multicast address
   // 224.0.0.251 (0xE00000FB)
-  if (!_mdnsSocket.beginMulti(IPAddress(224, 0, 0, 251), 5353)) {
+  if (!udpSocket.beginMulti(IPAddress(224, 0, 0, 251), 5353)) {
     return false;
   }
 
   return true;
 }
 
-void MDNSResponder::update() {
-  // Check if there's data to read from the UDP socket.
-  int available = _mdnsSocket.parsePacket();
-  // Stop processing if no data is available.
-  if (available <= 0) {
-    return;
-  }
-  // Otherwise there is data to read so grab it all and parse the data.
-  uint8_t buffer[READ_BUFFER_SIZE];
-  int n = _mdnsSocket.read((unsigned char*)&buffer, sizeof(buffer));
-  if (n < 1) {
-    // Error getting data.
-    return;
+void MDNSResponder::poll()
+{
+  if (parseRequest()) {
+    replyToRequest();
   }
-  // Look for domain name in request and respond with canned response if found.
-  for (int i = 0; i < n; ++i) {
-    uint8_t ch = tolower(buffer[i]);
-    // Check character matches expected.
-    if (ch == _expected[_index])
-    {
-      _index++;
-      // Check if domain name was found and send a response.
-      if (_index >= _expectedLen) {
-        // Send response to multicast address.
-        _broadcastResponse();
-        _index = 0;
+}
+
+bool MDNSResponder::parseRequest()
+{
+  if (udpSocket.parsePacket()) {
+    // check if parsed packet matches expected request length
+    if (udpSocket.available() != expectedRequestLength) {
+      // it does not, read the full packet in and drop data
+      while(udpSocket.available()) {
+        udpSocket.read();
       }
+
+      return false;
     }
-    else if (ch == _expected[0]) {
-      // Found a character that doesn't match, but does match the start of the domain.
-      _index = 1;
-    }
-    else {
-      // Found a character that doesn't match the expected character or start of domain.
-      _index = 0;
+
+    // read packet
+    uint8_t request[expectedRequestLength];
+    udpSocket.read(request, expectedRequestLength);
+
+    // parse request
+    uint8_t requestNameLength   = request[HEADER_SIZE];
+    uint8_t* requestName        = &request[HEADER_SIZE + 1];
+    uint8_t requestDomainLength = request[HEADER_SIZE + 1 + requestNameLength];
+    uint8_t* requestDomain      = &request[HEADER_SIZE + 1 + requestNameLength + 1];
+    uint16_t requestQtype;
+    uint16_t requestQclass;
+
+    memcpy(&requestQtype, &request[expectedRequestLength - 4], sizeof(requestQtype));
+    memcpy(&requestQclass, &request[expectedRequestLength - 2], sizeof(requestQclass));
+
+    requestQtype = _ntohs(requestQtype);
+    requestQclass = _ntohs(requestQclass);
+
+    // compare request
+    if (memcmp_P(request, expectedRequestHeader, HEADER_SIZE) == 0 &&            // request header match
+        requestNameLength == name.length() &&                                    // name length match
+        strncasecmp(name.c_str(), (char*)requestName, requestNameLength) == 0 && // name match
+        requestDomainLength == sizeof(domain) &&                                 // domain length match
+        memcmp_P(requestDomain, domain, requestDomainLength) == 0 &&             // suffix match
+        requestQtype == 0x0001 &&                                                // request QType match
+        (requestQclass == 0x0001 || requestQclass == 0x8001) ) {                 // request QClass match
+
+      return true;
     }
   }
+
+  return false;
 }
 
-void MDNSResponder::_broadcastResponse() {
-  // Send a MDNS name query response for this device.
-  // Use the MDNS multicast address 224.0.0.251 port 5353.
-  _mdnsSocket.beginPacket(IPAddress(224, 0, 0, 251), 5353);
-  _mdnsSocket.write(_response, _responseLen);
-  _mdnsSocket.endPacket();
+void MDNSResponder::replyToRequest()
+{
+  int nameLength = name.length();
+  int domainLength = sizeof(domain);
+  uint32_t ipAddress = WiFi.localIP();
+  uint32_t ttl = _htonl(ttlSeconds);
+
+  int responseSize = sizeof(responseHeader) + 1 + nameLength + 1 + domainLength + 1 + sizeof(aRecord) + sizeof(nsecRecord);
+  uint8_t response[responseSize];
+  uint8_t* r = response;
+
+  // copy header
+  memcpy_P(r, responseHeader, sizeof(responseHeader));
+  r += sizeof(responseHeader);
+  
+  // copy name
+  *r = nameLength;
+  memcpy(r + 1, name.c_str(), nameLength);
+  r += (1 + nameLength);
+
+  // copy domain
+  *r = domainLength;
+  memcpy_P(r + 1, domain, domainLength);
+  r += (1 + domainLength);
+
+  // terminator
+  *r = 0x00;
+  r++;
+
+  // copy A record
+  memcpy_P(r, aRecord, sizeof(aRecord));
+  memcpy(r + TTL_OFFSET, &ttl, sizeof(ttl));            // replace TTL value
+  memcpy(r + IP_OFFSET, &ipAddress, sizeof(ipAddress)); // replace IP address value
+  r += sizeof(aRecord);
+
+  // copy NSEC record
+  memcpy_P(r, nsecRecord, sizeof(nsecRecord));
+  r += sizeof(nsecRecord);
+
+  udpSocket.beginPacket(IPAddress(224, 0, 0, 251), 5353);
+  udpSocket.write(response, responseSize);
+  udpSocket.endPacket();
 }

src/WiFiMdns.h

@@ -31,23 +31,21 @@ class MDNSResponder {
 public:
   MDNSResponder();
   ~MDNSResponder();
-  bool begin(const char* domain, uint32_t ttlSeconds = 3600);
-  void update();
+  bool begin(const char* _name, uint32_t _ttlSeconds = 3600);
+  void poll();
 
 private:
-  void _broadcastResponse();
-
-  // Expected query values
-  static uint8_t _queryHeader[];
-  uint8_t* _expected;
-  int _expectedLen;
-  // Current parsing state
-  int _index;
-  // Response data
-  uint8_t* _response;
-  int _responseLen;
+  bool parseRequest();
+  void replyToRequest();
+
+private:
+  String name;
+  uint32_t ttlSeconds;
+
+  int expectedRequestLength;
+
   // UDP socket for receiving/sending MDNS data.
-  WiFiUDP _mdnsSocket;
+  WiFiUDP udpSocket;
 };
 
 #endif