/*
 * Basic multicast DNS responder
 * 
 * Advertises the IP address, port, and characteristics of a service to other devices using multicast DNS on the same LAN,
 * so they can find devices with addresses dynamically allocated by DHCP. See avahi, Bonjour, etc
 * See RFC6762, RFC6763
 *
 * This sample code is in the public domain.
 *
 * by M J A Hamel 2016
 */


#include <espressif/esp_common.h>
#include <espressif/esp_wifi.h>

#include <string.h>
#include <stdio.h>

#include <FreeRTOS.h>
#include <task.h>

#include <lwip/err.h>
#include <lwip/sockets.h>
#include <lwip/sys.h>
#include <lwip/netdb.h>
#include <lwip/dns.h>
#include <lwip/udp.h>
#include <lwip/igmp.h>
#include <lwip/netif.h>

#include "mdnsresponder.h"

#define qDebugLog             // Log activity generally
#define qLogIncoming          // Log all arriving multicast packets
#define qLogAllTraffic        // Log and decode all mDNS packets

#define kMDNSStackSize            800

#define DNS_MULTICAST_ADDRESS   "224.0.0.251"   // RFC 6762
#define DNS_MDNS_PORT           5353            // RFC 6762

//-------------------------------------------------------------------

#ifdef PACK_STRUCT_USE_INCLUDES
#  include "arch/bpstruct.h"
#endif

PACK_STRUCT_BEGIN
/** DNS message header */
struct mdns_hdr {
    PACK_STRUCT_FIELD(u16_t id);
    PACK_STRUCT_FIELD(u8_t flags1);
    PACK_STRUCT_FIELD(u8_t flags2);
    PACK_STRUCT_FIELD(u16_t numquestions);
    PACK_STRUCT_FIELD(u16_t numanswers);
    PACK_STRUCT_FIELD(u16_t numauthrr);
    PACK_STRUCT_FIELD(u16_t numextrarr);
}PACK_STRUCT_STRUCT;
PACK_STRUCT_END

#define SIZEOF_DNS_HDR 12

PACK_STRUCT_BEGIN
/** MDNS query message structure */
struct mdns_query {
    /* MDNS query record starts with either a domain name or a pointer
     to a name already present somewhere in the packet. */
    PACK_STRUCT_FIELD(u16_t type);
    PACK_STRUCT_FIELD(u16_t class);
}PACK_STRUCT_STRUCT;
PACK_STRUCT_END

#define SIZEOF_DNS_QUERY 4

PACK_STRUCT_BEGIN
/** MDNS answer message structure */
struct mdns_answer {
    /* MDNS answer record starts with either a domain name or a pointer
     to a name already present somewhere in the packet. */
    PACK_STRUCT_FIELD(u16_t type);
    PACK_STRUCT_FIELD(u16_t class);
    PACK_STRUCT_FIELD(u32_t ttl);
    PACK_STRUCT_FIELD(u16_t len);
}PACK_STRUCT_STRUCT;
PACK_STRUCT_END
#define SIZEOF_DNS_ANSWER 10

PACK_STRUCT_BEGIN
struct mdns_rr_srv {
    /* RR SRV  */
    PACK_STRUCT_FIELD(u16_t prio);
    PACK_STRUCT_FIELD(u16_t weight);
    PACK_STRUCT_FIELD(u16_t port);
}PACK_STRUCT_STRUCT;
PACK_STRUCT_END
#define SIZEOF_DNS_RR_SRV 6


#ifdef PACK_STRUCT_USE_INCLUDES
#  include "arch/epstruct.h"
#endif
  
#define vTaskDelayMs(ms)    vTaskDelay((ms)/portTICK_PERIOD_MS)
#define UNUSED_ARG(x)       (void)x
#define kDummyDataSize      8           // arbitrary, dynamically resized
#define kMaxNameSize        64
#define kMaxQStr            128         // max incoming question key handled

typedef struct mdns_rsrc {
    struct mdns_rsrc*    rNext;
    u16_t     rType;
    u32_t    rTTL;
    u16_t    rKeySize;
    u16_t    rDataSize;
    char    rData[kDummyDataSize];      // Key, as C str with . seperators, followed by data in network-ready form
                                        // at rData[rKeySize]
} mdns_rsrc;

static struct udp_pcb* gMDNS_pcb = NULL;
static ip_addr_t       gMulticastAddr;      // == DNS_MULTICAST_ADDRESS
static mdns_rsrc*      gDictP = NULL;       // RR database, linked list

//---------------------- Debug/logging utilities -------------------------

#ifdef qDebugLog

    // DNS field TYPE used for "Resource Records", some additions
    #define DNS_RRTYPE_AAAA           28    /* IPv6 host address */
    #define DNS_RRTYPE_SRV            33    /* Service record */
    #define DNS_RRTYPE_OPT            41    /* EDNS0 OPT record */
    #define DNS_RRTYPE_NSEC           47    /* NSEC record */
    #define DNS_RRTYPE_TSIG           250   /* Transaction Signature */
    #define DNS_RRTYPE_ANY            255   /* Not a DNS type, but a DNS query type, meaning "all types"*/

    // DNS field CLASS used for "Resource Records" 
    #define DNS_RRCLASS_ANY           255  /* Any class (q) */

    #define DNS_FLAG1_RESP            0x80
    #define DNS_FLAG1_OPMASK          0x78
    #define DNS_FLAG1_AUTH            0x04
    #define DNS_FLAG1_TRUNC           0x02
    #define DNS_FLAG1_RD              0x01
    #define DNS_FLAG2_RA              0x80
    #define DNS_FLAG2_RESMASK         0x0F

    static char qstr[12];
    
    static char* mdns_qrtype(uint16_t typ)
    {
        switch(typ) {
            case DNS_RRTYPE_A    : return ("A");
            case DNS_RRTYPE_NS   : return ("NS");
            case DNS_RRTYPE_PTR  : return ("PTR");
            case DNS_RRTYPE_TXT  : return ("TXT ");
            case DNS_RRTYPE_AAAA : return ("AAAA");
            case DNS_RRTYPE_SRV  : return ("SRV ");
            case DNS_RRTYPE_NSEC : return ("NSEC ");
            case DNS_RRTYPE_ANY  : return ("ANY");
        }
        sprintf(qstr,"type %d",typ);
        return qstr;
    }
    
    #ifdef qLogAllTraffic
    
        static void mdns_printhex(u8_t* p, int n)
        {
            int i;
            for (i=0; i<n; i++) {
                printf("%02X ",*p++);
                if ((i % 32) == 31) printf("\n");
            }
            printf("\n");
        }

        static void mdns_print_pstr(u8_t* p)
        {
            int i, n;
            char* cp;
    
            n = *p++;
            cp = (char*)p;
            for (i=0; i<n; i++) putchar(*cp++);
        }

        static char cstr[16];
        
        static char* mdns_qclass(uint16_t cls)
        {
            switch(cls) {
                case DNS_RRCLASS_IN  : return ("In");
                case DNS_RRCLASS_ANY : return ("ANY");
            }
            sprintf(cstr,"class %d",cls);
            return cstr;
        }

        static u8_t* mdns_print_name(u8_t* p, struct mdns_hdr* hp)
        // Sequence of Pascal strings, terminated by zero-length string
        // Handles compression, returns ptr to next item
        {
            char* cp = (char*)p;
            int i, n;
    
            do {
                n = *cp++;
                if ((n & 0xC0) == 0xC0) {
                    n = (n & 0x3F) << 8;
                    n |= (u8_t)*cp++;
                    mdns_print_name( (u8_t*)hp + n ,hp );
                    n = 0;
                } else if (n & 0xC0) {
                    printf("<label $%X?>",n);
                    n = 0;
                } else {
                    for (i=0; i<n; i++)
                        putchar(*cp++);
                    if (n!=0) putchar('.');
                }
            } while (n>0);
            return (u8_t*)cp;
        }

    
        static u8_t* mdns_print_header(struct mdns_hdr* hdr)
        {
            if (hdr->flags1 & DNS_FLAG1_RESP) {
                printf("Response, ID $%X %s ", htons(hdr->id), (hdr->flags1 & DNS_FLAG1_AUTH) ? "Auth " : "Non-auth ");
                if (hdr->flags2 & DNS_FLAG2_RA) printf("RA ");
                if ((hdr->flags2 & DNS_FLAG2_RESMASK)==0) printf("noerr");
                                       else printf("err %d", hdr->flags2 & DNS_FLAG2_RESMASK);
            } else {
                printf("Query, ID $%X op %d", htons(hdr->id), (hdr->flags1>>4) & 0x7 );
            }
            if (hdr->flags1 & DNS_FLAG1_RD) printf("RD ");
            if (hdr->flags1 & DNS_FLAG1_TRUNC) printf("[TRUNC] ");
    
            printf(": %d questions", htons(hdr->numquestions) );
            if (hdr->numanswers != 0)
                printf(", %d answers",htons(hdr->numanswers));
            if (hdr->numauthrr != 0)
                printf(", %d auth RR",htons(hdr->numauthrr));
            if (hdr->numextrarr != 0)
                printf(", %d extra RR",htons(hdr->numextrarr));
            putchar('\n');
            return (u8_t*)hdr + SIZEOF_DNS_HDR;
        }

        static u8_t* mdns_print_query(u8_t* p)
        // Copy needed because it may be misaligned
        {
            struct mdns_query q;
            uint16_t c;
    
            memcpy(&q,p,SIZEOF_DNS_QUERY);
            c = htons(q.class);
            printf(" %s %s", mdns_qrtype(htons(q.type)), mdns_qclass(c & 0x7FFF) );
            if (c & 0x8000) printf(" unicast-req");
            printf("\n");
            return p + SIZEOF_DNS_QUERY;
        }

        static u8_t* mdns_print_answer(u8_t* p, struct mdns_hdr* hp)
        // Copy needed because it may be misaligned
        {
            struct mdns_answer ans;
            u16_t rrlen, atype, rrClass;;
    
            memcpy(&ans,p,SIZEOF_DNS_ANSWER);
            atype = htons(ans.type);
            rrlen = htons(ans.len);
            rrClass = htons(ans.class);
            printf(" %s %s TTL %d ", mdns_qrtype(atype), mdns_qclass(rrClass & 0x7FFF), htonl(ans.ttl));
            if (rrClass & 0x8000) 
                printf("cache-flush ");
            if (rrlen > 0) {
                u8_t* rp = p + SIZEOF_DNS_ANSWER;
                if (atype==DNS_RRTYPE_A && rrlen==4) {
                    printf("%d.%d.%d.%d\n",rp[0],rp[1],rp[2],rp[3]);
                } else if (atype==DNS_RRTYPE_PTR) {
                    mdns_print_name(rp, hp);
                    printf("\n");
                } else if (atype==DNS_RRTYPE_TXT) {
                    mdns_print_pstr(rp);
                    printf("\n");
                } else if (atype==DNS_RRTYPE_SRV && rrlen > SIZEOF_DNS_RR_SRV) {
                    struct mdns_rr_srv srvRR;
                    memcpy(&srvRR,rp,SIZEOF_DNS_RR_SRV);
                    printf("prio %d, weight %d, port %d, target ", srvRR.prio, srvRR.weight, srvRR.port);
                    mdns_print_name(rp + SIZEOF_DNS_RR_SRV, hp);
                    printf("\n");
                } else {
                    printf("%db:",rrlen);
                    mdns_printhex(rp,rrlen);
                }
            } else
                printf("\n");
            return p + SIZEOF_DNS_ANSWER + rrlen;
        }

        static int mdns_print_msg(u8_t* msgP, int msgLen)
        {
            int i;
            u8_t*    tp;
            u8_t*     limP = msgP + msgLen;
            struct mdns_hdr* hdr;
    
            hdr = (struct mdns_hdr*) msgP;
            tp = mdns_print_header(hdr);
            for (i=0; i<htons(hdr->numquestions); i++) {
                printf(" Q%d: ",i+1);
                tp = mdns_print_name(tp,hdr);
                tp = mdns_print_query(tp);
                if (tp > limP) return 0;
            }
    
            for (i=0; i<htons(hdr->numanswers); i++) {
                printf(" A%d: ",i+1);
                tp = mdns_print_name(tp,hdr);
                tp = mdns_print_answer(tp,hdr);
                if (tp > limP) return 0;
            }
    
            for (i=0; i<htons(hdr->numauthrr); i++) {
                printf(" AuRR%d: ",i+1);
                tp = mdns_print_name(tp,hdr);
                tp = mdns_print_answer(tp,hdr);
                if (tp > limP) return 0;
            }
    
            for (i=0; i<htons(hdr->numextrarr); i++) {
                printf(" ExRR%d: ",i+1);
                tp = mdns_print_name(tp,hdr);
                tp = mdns_print_answer(tp,hdr);
                if (tp > limP) return 0;
            }
            return 1;
        }
    #endif

#endif // qDebugLog

//---------------------------------------------------------------------------

static u8_t* mdns_labels2str(u8_t* hdrP, u8_t* p, char* qStr)
// Convert a DNS domain name label sequence into C string with . seperators
// Handles compression
{
    int i, n;
    
    do {
        n = *p++;
        if ((n & 0xC0) == 0xC0) {
            n = (n & 0x3F) << 8;
            n |= (u8_t)*p++;
            mdns_labels2str( hdrP, hdrP + n, qStr);
            return p;
        } else if (n & 0xC0) {
            printf(">>> mdns_labels2str,label $%X?",n);
            return p;
        } else {
            for (i=0; i<n; i++)
                *qStr++ = *p++;
            if (n==0) *qStr++ = 0;
                 else *qStr++ = '.';
        }
    } while (n>0);
    return p;
}

static int mdns_str2labels(const char* name, u8_t* lseq, int max)
// Encode a <string>.<string>.<string> as a sequence of labels, return length
{
    int i,n,sdx,idx = 0;
    int lc = 0;
    
    do {
        sdx = idx;
        while (name[idx] != '.' && name[idx] != 0) idx++;
        n = idx-sdx;
        *lseq++ = n;
        lc++;
        if (lc+n > max) {
            printf(">>> mdns_str2labels: oversize (%d)\n",lc+n);
            return 0;
        }
        for (i=0; i<n; i++) 
            *lseq++ = name[sdx+i];
        lc += n;
        if (name[idx]=='.')
            idx++;
    } while (n>0);
    return lc;
}

static u8_t* mdns_get_question(u8_t* hdrP, u8_t* qp, char* qStr, uint16_t* qClass, uint16_t* qType, u8_t* qUnicast)
// Unpack a DNS question RR at qp, return pointer to next RR
{
    struct mdns_query qr;
    uint16_t cls;
    
    qp = mdns_labels2str(hdrP, qp, qStr);
    memcpy(&qr,qp,SIZEOF_DNS_QUERY);
    *qType = htons(qr.type);
    cls = htons(qr.class);
    *qUnicast = cls>>15;
    *qClass = cls & 0x7FFF;
    return qp + SIZEOF_DNS_QUERY;
}

//---------------------------------------------------------------------------


static void mdns_add_response(const char* vKey, u16_t vType, u32_t ttl, const void* dataP, u16_t vDataSize)
// Add a record to the RR database list
{
    mdns_rsrc* rsrcP;
    int keyLen, recSize;
    
    keyLen = strlen(vKey) + 1;
    recSize = sizeof(mdns_rsrc) - kDummyDataSize + keyLen + vDataSize;
    rsrcP = (mdns_rsrc*)malloc(recSize);
    if (rsrcP==NULL)
        printf(">>> mdns_add_response: couldn't alloc %d\n",recSize);
    else {
        rsrcP->rType = vType;
        rsrcP->rTTL = ttl;
        rsrcP->rKeySize = keyLen;
        rsrcP->rDataSize = vDataSize;
        memcpy(rsrcP->rData, vKey, keyLen);
        memcpy(&rsrcP->rData[keyLen], dataP, vDataSize);
        rsrcP->rNext = gDictP;
        gDictP = rsrcP;
        #ifdef qDebugLog
            printf("mDNS added RR '%s' %s, %d bytes\n", vKey, mdns_qrtype(vType), vDataSize);
        #endif
    }
}

void mdns_add_PTR(const char* rKey, u32_t ttl, const char* nmStr)
{
    int nl;
    u8_t lBuff[kMaxNameSize];
    
    nl = mdns_str2labels(nmStr,lBuff,sizeof(lBuff));
    if (nl>0)
        mdns_add_response(rKey, DNS_RRTYPE_PTR, ttl, lBuff, nl);
}

void mdns_add_SRV(const char* rKey, u32_t ttl, u16_t rPort, const char* targName)
{
    typedef struct SrvRec  {
        struct mdns_rr_srv srvRR;
        u8_t lBuff[kMaxNameSize];
    } __attribute__((packed)) SrvRec;
    
    int     nl;
    SrvRec     temp;
    
    temp.srvRR.prio = 0;
    temp.srvRR.weight = 0;
    temp.srvRR.port = htons(rPort);
    nl = mdns_str2labels(targName,temp.lBuff,sizeof(temp.lBuff));
    if (nl>0)
        mdns_add_response(rKey, DNS_RRTYPE_SRV, ttl, &temp, SIZEOF_DNS_RR_SRV + nl);
}

void mdns_add_TXT(const char* rKey, u32_t ttl, const char* txStr)
// Single TXT str, can be concatenated
{
    char pstr[256];
    u16_t n = strlen(txStr);
    if (n > 255)
        printf(">>> mdns_add_TXT oversize (%d)\n",n);
    else {
        pstr[0] = n;
        memcpy(&pstr[1],txStr,n);
        mdns_add_response(rKey, DNS_RRTYPE_TXT, ttl, txStr, n+1);
    }
}

void mdns_add_A(const char* rKey, u32_t ttl, struct ip_addr addr)
{
    mdns_add_response(rKey, DNS_RRTYPE_A, ttl, &addr, sizeof(addr));
}

void mdns_add_facility( const char* instanceName,   // Friendly name, need not be unique
                        const char* serviceName,         // Must be _name
                        const char* addText,        // Must be <key>=<value>
                        mdns_flags  flags,                 // TCP or UDP
                        u16_t onPort,                 // port number
                        u32_t ttl                    // seconds
                      )
{
    char key[64];
    char fullName[128];
    char devName[96];
    struct ip_info ipInfo;
    
    #ifdef qDebugLog
        printf("\nmDNS advertising instance %s protocol %s text %s on port %d %s TTL %d secs\n",
                instanceName,serviceName,addText,onPort,(flags & mdns_UDP) ? "UDP" : "TCP", ttl);
    #endif
            
    snprintf(key,sizeof(key),"%s.%s.local.",serviceName,(flags & mdns_UDP) ? "_udp" :"_tcp");
    snprintf(fullName,sizeof(fullName),"%s.%s",instanceName,key);
    snprintf(devName,sizeof(devName),"%s.local.",instanceName);
    
    if (!sdk_wifi_get_ip_info(STATION_IF,&ipInfo))
        ipInfo.ip.addr = IPADDR_NONE;
    
    // Order has significance for extraRR feature
    mdns_add_TXT(fullName,ttl,addText);
    mdns_add_A(devName,ttl,ipInfo.ip);
    mdns_add_SRV(fullName,ttl,onPort,devName);
    mdns_add_PTR(key,ttl,fullName);
    
    // Optional, makes us browsable
    if (flags & mdns_Browsable)
        mdns_add_PTR("_services._dns-sd._udp.local.",ttl,key);
        
}

static void mdns_update_ipaddr(struct ip_info* ipInfo)
// IP address has been defined/changed: update any A records with the new IP
{
    mdns_rsrc* rp = gDictP;
    while (rp != NULL) {
        if (rp->rType==DNS_RRTYPE_A) {
            #ifdef qDebugLog
                printf("Updating A record for '%s' to %d.%d.%d.%d\n", rp->rData,
                    ip4_addr1(&ipInfo->ip), ip4_addr2(&ipInfo->ip), ip4_addr3(&ipInfo->ip), ip4_addr4(&ipInfo->ip));
            #endif
            memcpy(&rp->rData[rp->rKeySize], &ipInfo->ip, sizeof(ip_addr_t));
        }
        rp = rp->rNext;
    }
}

static mdns_rsrc* mdns_match(const char* qstr, u16_t qType)
{
    mdns_rsrc* rp = gDictP;
    while (rp != NULL) {
        if (rp->rType==qType || qType==DNS_RRTYPE_ANY) {
            if (strcasecmp(rp->rData,qstr)==0) {
                #ifdef qDebugLog
                    printf(" - matched '%s' %s\n",qstr,mdns_qrtype(rp->rType));
                #endif
                break;
            }
        }
        rp = rp->rNext;
    }
    return rp;
}

static int mdns_add_to_answer(mdns_rsrc* rsrcP, u8_t* resp, int respLen)
// Create answer RR and append to resp[respLen], return new length
{
    // Key is stored as C str, convert to labels
    respLen += mdns_str2labels(rsrcP->rData, &resp[respLen], DNS_MSG_SIZE-respLen);

    // Answer fields: may be misaligned, so build and memcpy
    struct mdns_answer ans;
    ans.type  = htons(rsrcP->rType);
    ans.class = htons(DNS_RRCLASS_IN);
    ans.ttl   = htonl(rsrcP->rTTL);
    ans.len   = htons(rsrcP->rDataSize);
    memcpy(&resp[respLen], &ans, SIZEOF_DNS_ANSWER);
    respLen += SIZEOF_DNS_ANSWER;

    // Data for this key
    memcpy(&resp[respLen], &rsrcP->rData[rsrcP->rKeySize], rsrcP->rDataSize);
    respLen += rsrcP->rDataSize;
    
    return respLen;
}

//---------------------------------------------------------------------------

static void mdns_send_mcast(u8_t* msgP, int nBytes)
// Send UDP to multicast address
{
    struct pbuf* p;
    err_t err;

    p = pbuf_alloc(PBUF_TRANSPORT, nBytes, PBUF_RAM);
    if (p) {
        memcpy(p->payload, msgP, nBytes);
        err = udp_sendto(gMDNS_pcb, p, &gMulticastAddr, DNS_MDNS_PORT);
        if (err==ERR_OK) {
            #ifdef qDebugLog
                printf(" - responded with %d bytes err %d\n",nBytes,err);
            #endif
        } else
            printf(">>> mdns_send failed %d\n",err);
        pbuf_free(p);
    } else
        printf(">>> mdns_send: alloc failed[%d]\n",nBytes);
}
    
static void mdns_reply(struct mdns_hdr* hdrP)
// Message has passed tests, may want to send an answer
{
    int i, nquestions, respLen;
    struct mdns_hdr* rHdr;
    mdns_rsrc* extra;
    u8_t* qBase = (u8_t*)hdrP;
    u8_t* qp;
    u8_t* mdns_response;
    
    mdns_response = malloc(DNS_MSG_SIZE);
    if (mdns_response==NULL) {
        printf(">>> mdns_reply could not alloc %d\n",DNS_MSG_SIZE);
        return;
    }
    
    // Build response header
    rHdr = (struct mdns_hdr*) mdns_response;
    rHdr->id = hdrP->id;
    rHdr->flags1 = DNS_FLAG1_RESP + DNS_FLAG1_AUTH;
    rHdr->flags2 = 0;
    rHdr->numquestions = 0;
    rHdr->numanswers = 0;
    rHdr->numauthrr = 0;
    rHdr->numextrarr = 0;
    respLen = SIZEOF_DNS_HDR;

    extra = NULL;
    qp = qBase + SIZEOF_DNS_HDR;
    nquestions = htons(hdrP->numquestions);
    
    for (i=0; i<nquestions; i++) {
        char  qStr[kMaxQStr];
        u16_t qClass, qType;
        u8_t  qUnicast;
        mdns_rsrc* rsrcP;
    
        qp = mdns_get_question(qBase, qp, qStr, &qClass, &qType, &qUnicast);
        if (qClass==DNS_RRCLASS_IN || qClass==DNS_RRCLASS_ANY) {
            rsrcP = mdns_match(qStr,qType);
            if (rsrcP) {
                respLen = mdns_add_to_answer(rsrcP, mdns_response, respLen);
                rHdr->numanswers = htons( htons(rHdr->numanswers) + 1 );
                // Extra RR logic: if SRV follows PTR, or A follows SRV, volunteer it in extraRR
                // Not required, but could do more here, see RFC6763 s12
                if (qType==DNS_RRTYPE_PTR) {
                    if (rsrcP->rNext && rsrcP->rNext->rType==DNS_RRTYPE_SRV) 
                        extra = rsrcP->rNext;
                } else if (qType==DNS_RRTYPE_SRV) {
                    if (rsrcP->rNext && rsrcP->rNext->rType==DNS_RRTYPE_A) 
                        extra = rsrcP->rNext;
                }        
            }
        }
    } // for nQuestions

    if (respLen > SIZEOF_DNS_HDR) {
        if (extra) {
            respLen = mdns_add_to_answer(extra, mdns_response, respLen);
            rHdr->numextrarr = htons( htons(rHdr->numextrarr) + 1 );
        }
        mdns_send_mcast(mdns_response,respLen);
    }
    free(mdns_response);
}

static void mdns_recv(void *arg, struct udp_pcb *pcb, struct pbuf *p, struct ip_addr *addr, u16_t port) 
// Callback from udp_recv
{
    UNUSED_ARG(pcb);
    UNUSED_ARG(port);
    
    u8_t* mdns_payload;
    int   plen;
    
    // Sanity checks on size
    plen = p->tot_len;
    if (plen > DNS_MSG_SIZE) {
        printf(">>> mdns_recv: pbuf too big\n");
    } else if (plen < (SIZEOF_DNS_HDR + SIZEOF_DNS_QUERY + 1 + SIZEOF_DNS_ANSWER + 1)) {
        printf(">>> mdns_recv: pbuf too small\n");
    } else {
        #ifdef qLogIncoming
            printf("\n\nmDNS got %d bytes from %d.%d.%d.%d\n",plen, ip4_addr1(addr), ip4_addr2(addr), ip4_addr3(addr), ip4_addr4(addr));
        #endif
        mdns_payload = malloc(plen);
        if (!mdns_payload)
            printf(">>> mdns_recv, could not alloc %d\n",plen);
        else {
            if (pbuf_copy_partial(p, mdns_payload, plen, 0) == plen) {
                struct mdns_hdr* hdrP = (struct mdns_hdr*) mdns_payload;
            
                #ifdef qLogAllTraffic
                    mdns_print_msg(mdns_payload, plen);
                #endif
                    
                if ( (hdrP->flags1 & (DNS_FLAG1_RESP + DNS_FLAG1_OPMASK + DNS_FLAG1_TRUNC) ) == 0 
                     && hdrP->numquestions > 0 )
                    mdns_reply(hdrP);
            }
            free(mdns_payload);
        }
    }
    pbuf_free(p);
}

static void mdns_start()
// If we are in station mode and have an IP address, start a multicast UDP receive
{
    struct ip_info ipInfo;
    err_t err;
    
    if (sdk_wifi_get_opmode() != STATION_MODE) {
        printf(">>> mDNS_start: wifi opmode not station\n");
        return;
    }
    
    if (!sdk_wifi_get_ip_info(STATION_IF,&ipInfo)) {
        printf(">>> mDNS_start: no IP addr\n");
        return;
    }

    mdns_update_ipaddr(&ipInfo);
    
    // Start IGMP on the netif for our interface: this isn't done for us
    struct netif* nfp = netif_list;
    while (nfp!=NULL) {        
        if ( ip_addr_cmp(&ipInfo.ip, &(nfp->ip_addr)) ) {
            if (!(nfp->flags & NETIF_FLAG_IGMP)) {
                nfp->flags |= NETIF_FLAG_IGMP;
                err = igmp_start(nfp);
                if (err != ERR_OK) {
                    printf(">>> mDNS_start: igmp_start on %c%c failed %d\n",nfp->name[0], nfp->name[1],err);
                    return;
                } 
            } 
        }
        nfp = nfp->next;
    }
                      
    gMDNS_pcb = udp_new();
    if (!gMDNS_pcb) {
        printf(">>> mDNS_start: udp_new failed\n");
        return;
    }
    
    if ((err=igmp_joingroup(&ipInfo.ip, &gMulticastAddr)) != ERR_OK) {
        printf(">>> mDNS_start: igmp_join failed %d\n",err);
        return;
    }
    
    if ((err=udp_bind(gMDNS_pcb, IP_ADDR_ANY, DNS_MDNS_PORT)) != ERR_OK) {
        printf(">>> mDNS_start: udp_bind failed %d\n",err);
        return;
    }
    
    udp_recv(gMDNS_pcb, mdns_recv, NULL);
}

static void mdns_close()
{
    udp_remove(gMDNS_pcb);
    gMDNS_pcb = NULL;
    #ifdef qDebugLog
        printf("Closing mDNS\n");
    #endif
}

static void mdns_task(void *pvParameters)
{
    uint8_t hasIP = 0;
    uint8_t status;
    UNUSED_ARG(pvParameters);

    ipaddr_aton(DNS_MULTICAST_ADDRESS, &gMulticastAddr);
    // Wait until we have joined AP and are assigned an IP
    while (1) {
        status = (sdk_wifi_station_get_connect_status() == STATION_GOT_IP);
        if (status != hasIP) {
            if (status) mdns_start();
                   else mdns_close();
            hasIP = status;
        }
        vTaskDelayMs(status ? 1000 : 100);
    }
}

void mdns_init()
{
    #if LWIP_IGMP
        xTaskCreate(mdns_task, "MDNS", kMDNSStackSize, NULL, 1, NULL);
    #else
        #error "LWIP_IGMP needs to be defined in lwipopts.h"
    #endif
}