Chapter 83: DNS Server and DNS-SD Implementation

#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "freertos/event_groups.h"
#include "esp_system.h"
#include "esp_wifi.h"
#include "esp_event.h"
#include "esp_log.h"
#include "nvs_flash.h"
#include "mdns.h"
#include "esp_netif.h"

// Wi-Fi Configuration - Replace with your network credentials
#define EXAMPLE_ESP_WIFI_SSID      "YOUR_WIFI_SSID"
#define EXAMPLE_ESP_WIFI_PASS      "YOUR_WIFI_PASSWORD"
#define EXAMPLE_ESP_MAXIMUM_RETRY  5

// mDNS Configuration
#define EXAMPLE_MDNS_HOSTNAME "my-esp32" // Hostname will be "my-esp32.local"
#define EXAMPLE_MDNS_INSTANCE "My ESP32 Device"

static EventGroupHandle_t s_wifi_event_group;
#define WIFI_CONNECTED_BIT BIT0
#define WIFI_FAIL_BIT      BIT1

static const char *TAG = "mdns_example";
static int s_retry_num = 0;

static void event_handler(void* arg, esp_event_base_t event_base,
                                int32_t event_id, void* event_data)
{
    if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_START) {
        esp_wifi_connect();
    } else if (event_base == WIFI_EVENT && event_id == WIFI_EVENT_STA_DISCONNECTED) {
        if (s_retry_num < EXAMPLE_ESP_MAXIMUM_RETRY) {
            esp_wifi_connect();
            s_retry_num++;
            ESP_LOGI(TAG, "retry to connect to the AP");
        } else {
            xEventGroupSetBits(s_wifi_event_group, WIFI_FAIL_BIT);
        }
        ESP_LOGI(TAG,"connect to the AP fail");
    } else if (event_base == IP_EVENT && event_id == IP_EVENT_STA_GOT_IP) {
        ip_event_got_ip_t* event = (ip_event_got_ip_t*) event_data;
        ESP_LOGI(TAG, "got ip:" IPSTR, IP2STR(&event->ip_info.ip));
        s_retry_num = 0;
        xEventGroupSetBits(s_wifi_event_group, WIFI_CONNECTED_BIT);
    }
}

void wifi_init_sta(void)
{
    s_wifi_event_group = xEventGroupCreate();

    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    esp_netif_create_default_wifi_sta();

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    esp_event_handler_instance_t instance_any_id;
    esp_event_handler_instance_t instance_got_ip;
    ESP_ERROR_CHECK(esp_event_handler_instance_register(WIFI_EVENT,
                                                        ESP_EVENT_ANY_ID,
                                                        &event_handler,
                                                        NULL,
                                                        &instance_any_id));
    ESP_ERROR_CHECK(esp_event_handler_instance_register(IP_EVENT,
                                                        IP_EVENT_STA_GOT_IP,
                                                        &event_handler,
                                                        NULL,
                                                        &instance_got_ip));

    wifi_config_t wifi_config = {
        .sta = {
            .ssid = EXAMPLE_ESP_WIFI_SSID,
            .password = EXAMPLE_ESP_WIFI_PASS,
        },
    };
    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_STA) );
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_STA, &wifi_config) );
    ESP_ERROR_CHECK(esp_wifi_start() );

    ESP_LOGI(TAG, "wifi_init_sta finished.");

    EventBits_t bits = xEventGroupWaitBits(s_wifi_event_group,
            WIFI_CONNECTED_BIT | WIFI_FAIL_BIT,
            pdFALSE,
            pdFALSE,
            portMAX_DELAY);

    if (bits & WIFI_CONNECTED_BIT) {
        ESP_LOGI(TAG, "connected to ap SSID:%s", EXAMPLE_ESP_WIFI_SSID);
    } else if (bits & WIFI_FAIL_BIT) {
        ESP_LOGI(TAG, "Failed to connect to SSID:%s", EXAMPLE_ESP_WIFI_SSID);
    } else {
        ESP_LOGE(TAG, "UNEXPECTED EVENT");
    }
}

void start_mdns_service(void)
{
    // Initialize mDNS service
    esp_err_t err = mdns_init();
    if (err) {
        ESP_LOGE(TAG, "mDNS Init failed: %d", err);
        return;
    }
    ESP_LOGI(TAG, "mDNS Initalized.");

    // Set hostname
    ESP_ERROR_CHECK(mdns_hostname_set(EXAMPLE_MDNS_HOSTNAME));
    ESP_LOGI(TAG, "mDNS hostname set to: [%s.local]", EXAMPLE_MDNS_HOSTNAME);

    // Set default instance name
    ESP_ERROR_CHECK(mdns_instance_name_set(EXAMPLE_MDNS_INSTANCE));
    ESP_LOGI(TAG, "mDNS instance name set to: [%s]", EXAMPLE_MDNS_INSTANCE);

    // Structure with TXT records - Can be empty if not advertising specific service attributes here
    // mdns_txt_item_t serviceTxtData[1] = {
    //     {"board", "ESP32"}
    // };

    // // Initialize service advertising (optional here, more relevant for DNS-SD)
    // ESP_ERROR_CHECK(mdns_service_add("ESP32-WebServer", "_http", "_tcp", 80, serviceTxtData, 1));
    // ESP_LOGI(TAG, "mDNS service _http._tcp added (port 80).");
    // ESP_LOGI(TAG, "You can now ping %s.local or browse for its services.", EXAMPLE_MDNS_HOSTNAME);
}

void app_main(void)
{
    // Initialize NVS
    esp_err_t ret = nvs_flash_init();
    if (ret == ESP_ERR_NVS_NO_FREE_PAGES || ret == ESP_ERR_NVS_NEW_VERSION_FOUND) {
      ESP_ERROR_CHECK(nvs_flash_erase());
      ret = nvs_flash_init();
    }
    ESP_ERROR_CHECK(ret);

    wifi_init_sta(); // Connect to Wi-Fi

    // Wait until Wi-Fi is connected before starting mDNS
    EventBits_t bits = xEventGroupWaitBits(s_wifi_event_group, WIFI_CONNECTED_BIT, pdFALSE, pdTRUE, portMAX_DELAY);
    if (bits & WIFI_CONNECTED_BIT) {
        start_mdns_service();
    } else {
        ESP_LOGE(TAG, "Wi-Fi connection failed. mDNS not started.");
    }
}

// ... (previous includes and Wi-Fi code) ...

// In start_mdns_service() function:
void start_mdns_service(void)
{
    // Initialize mDNS service
    esp_err_t err = mdns_init();
    if (err) {
        ESP_LOGE(TAG, "mDNS Init failed: %d", err);
        return;
    }
    ESP_LOGI(TAG, "mDNS Initalized.");

    // Set hostname
    ESP_ERROR_CHECK(mdns_hostname_set(EXAMPLE_MDNS_HOSTNAME));
    ESP_LOGI(TAG, "mDNS hostname set to: [%s.local]", EXAMPLE_MDNS_HOSTNAME);

    // Set default instance name
    ESP_ERROR_CHECK(mdns_instance_name_set(EXAMPLE_MDNS_INSTANCE));
    ESP_LOGI(TAG, "mDNS instance name set to: [%s]", EXAMPLE_MDNS_INSTANCE);

    // Define TXT records for the service
    mdns_txt_item_t serviceTxtData[3] = {
        {"board", "ESP32"},
        {"version", "1.0.0"},
        {"path", "/"} // Example: default path for the web server
    };

    // Add the HTTP service
    // Service name: "esp-web-server" (will appear as "esp-web-server._http._tcp.local")
    // Protocol: "_http"
    // Transport: "_tcp"
    // Port: 80
    ESP_ERROR_CHECK(mdns_service_add("esp-web-server", "_http", "_tcp", 80, serviceTxtData, 3));
    ESP_LOGI(TAG, "mDNS service _http._tcp added with instance name 'esp-web-server' on port 80.");
    ESP_LOGI(TAG, "Browse for HTTP services on your network or try accessing http://%s.local/", EXAMPLE_MDNS_HOSTNAME);

    // You can add more services, e.g., an FTP server
    // ESP_ERROR_CHECK(mdns_service_add("esp-ftp-server", "_ftp", "_tcp", 21, NULL, 0));
}

// ... (app_main and Wi-Fi init remain the same) ...

// Add these includes
#include "lwip/sockets.h" // For socket programming
#include "lwip/dns.h"     // For DNS header structures (may need to define them minimally if not fully exposed)

// AP Configuration
#define AP_ESP_WIFI_SSID      "ESP32_Config_AP"
#define AP_ESP_WIFI_PASS      "esp32config"
#define AP_ESP_WIFI_CHANNEL   1
#define AP_MAX_STA_CONN       4

// AP static IP (also used as DNS server IP for clients)
#define AP_IP_ADDR    "192.168.4.1"
#define AP_NETMASK_ADDR "255.255.255.0"
#define AP_GW_ADDR    "192.168.4.1"

#define DNS_PORT 53
#define DNS_MAX_LEN 512 // Max DNS packet length

// Simplified DNS Header Structure
typedef struct {
    uint16_t id;        // Identification
    uint16_t flags;     // Flags
    uint16_t qdcount;   // Question count
    uint16_t ancount;   // Answer count
    uint16_t nscount;   // Authority RR count
    uint16_t arcount;   // Additional RR count
} dns_header_t;

// Simplified DNS Question Structure (QNAME is variable)
// For this example, we don't parse QNAME deeply, just respond to any A query

static const char *TAG_AP_DNS = "ap_dns_server";
static esp_netif_t *ap_netif_g = NULL; // Global to store AP netif for IP

void dns_server_task(void *pvParameters) {
    uint8_t rx_buffer[DNS_MAX_LEN];
    struct sockaddr_in client_addr;
    socklen_t client_addr_len = sizeof(client_addr);
    int sock_fd;

    struct sockaddr_in server_addr;
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(DNS_PORT);
    server_addr.sin_addr.s_addr = htonl(INADDR_ANY); // Listen on AP interface IP

    sock_fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_IP);
    if (sock_fd < 0) {
        ESP_LOGE(TAG_AP_DNS, "Failed to create DNS socket");
        vTaskDelete(NULL);
        return;
    }

    if (bind(sock_fd, (struct sockaddr *)&server_addr, sizeof(server_addr)) < 0) {
        ESP_LOGE(TAG_AP_DNS, "Failed to bind DNS socket: %s", strerror(errno));
        close(sock_fd);
        vTaskDelete(NULL);
        return;
    }
    ESP_LOGI(TAG_AP_DNS, "DNS server listening on port %d", DNS_PORT);

    esp_netif_ip_info_t ip_info_ap;
    ESP_ERROR_CHECK(esp_netif_get_ip_info(ap_netif_g, &ip_info_ap)); // Get AP's current IP

    while (1) {
        int len = recvfrom(sock_fd, rx_buffer, sizeof(rx_buffer), 0, (struct sockaddr *)&client_addr, &client_addr_len);
        if (len < 0) {
            ESP_LOGE(TAG_AP_DNS, "DNS recvfrom failed: %s", strerror(errno));
            continue;
        }

        if (len < sizeof(dns_header_t)) {
            ESP_LOGW(TAG_AP_DNS, "Received too short DNS packet");
            continue;
        }
        
        dns_header_t *dns_hdr = (dns_header_t *)rx_buffer;

        // Basic check: is it a standard query? (QR=0, Opcode=0)
        // We only care about the question count being > 0
        if ((ntohs(dns_hdr->flags) & 0x8000) == 0 && dns_hdr->qdcount > 0) { // It's a query
            ESP_LOGI(TAG_AP_DNS, "Received DNS query (ID: 0x%X) from %s", ntohs(dns_hdr->id), inet_ntoa(client_addr.sin_addr));

            // Prepare response: copy query ID, set flags for response
            dns_hdr->flags = htons(0x8180); // Standard response, No error, Authoritative Answer
            dns_hdr->ancount = dns_hdr->qdcount; // Respond with one answer per question (simplified)
            // For simplicity, we assume only one question and it's an A record query.
            // A robust server would parse QTYPE.

            // Construct the answer part (directly after the question part)
            // Question part ends after QNAME (variable), QTYPE (2), QCLASS (2)
            // Find end of question section to append answer
            uint8_t *p_qname_start = rx_buffer + sizeof(dns_header_t);
            uint8_t *p_qname_end = p_qname_start;
            while (*p_qname_end != 0 && (p_qname_end < rx_buffer + len)) { // Find end of QNAME
                p_qname_end++;
            }
            if (*p_qname_end != 0) { // Malformed QNAME
                ESP_LOGW(TAG_AP_DNS, "Malformed QNAME in DNS query");
                continue;
            }
            p_qname_end++; // Skip the null terminator of QNAME

            uint8_t *p_answer_start = p_qname_end + 4; // Skip QTYPE (2 bytes) and QCLASS (2 bytes)
            
            // Check if there's enough space for our simple answer
            // Answer: Name (2, pointer), Type (2), Class (2), TTL (4), RDLENGTH (2), RDATA (4 for IPv4) = 16 bytes
            if ((p_answer_start + 16) > (rx_buffer + DNS_MAX_LEN)) {
                ESP_LOGW(TAG_AP_DNS, "Not enough buffer space for DNS answer");
                continue;
            }

            // Answer:
            // Name (pointer to the QNAME in the question section, 0xC00C assumes QNAME starts at offset 12)
            *p_answer_start++ = 0xC0;
            *p_answer_start++ = 0x0C; 
            // Type A (1)
            *p_answer_start++ = 0x00;
            *p_answer_start++ = 0x01;
            // Class IN (1)
            *p_answer_start++ = 0x00;
            *p_answer_start++ = 0x01;
            // TTL (e.g., 60 seconds)
            uint32_t ttl = htonl(60);
            memcpy(p_answer_start, &ttl, sizeof(ttl));
            p_answer_start += sizeof(ttl);
            // RDLENGTH (4 for IPv4)
            *p_answer_start++ = 0x00;
            *p_answer_start++ = 0x04;
            // RDATA (AP's IP address)
            memcpy(p_answer_start, &ip_info_ap.ip.addr, sizeof(ip_info_ap.ip.addr));
            p_answer_start += sizeof(ip_info_ap.ip.addr);

            int response_len = p_answer_start - rx_buffer;
            sendto(sock_fd, rx_buffer, response_len, 0, (struct sockaddr *)&client_addr, client_addr_len);
            ESP_LOGI(TAG_AP_DNS, "Sent DNS response to %s with IP " IPSTR, inet_ntoa(client_addr.sin_addr), IP2STR(&ip_info_ap.ip));
        }
    }
    close(sock_fd);
    vTaskDelete(NULL);
}

void wifi_init_softap_with_dns(void) {
    ESP_ERROR_CHECK(esp_netif_init());
    ESP_ERROR_CHECK(esp_event_loop_create_default());
    ap_netif_g = esp_netif_create_default_wifi_ap(); // Store ap_netif globally
    assert(ap_netif_g);

    // Configure static IP for AP
    esp_netif_ip_info_t ip_info;
    inet_pton(AF_INET, AP_IP_ADDR, &ip_info.ip);
    inet_pton(AF_INET, AP_GW_ADDR, &ip_info.gw);
    inet_pton(AF_INET, AP_NETMASK_ADDR, &ip_info.netmask);
    ESP_ERROR_CHECK(esp_netif_dhcps_stop(ap_netif_g)); // Stop DHCP server before setting IP
    ESP_ERROR_CHECK(esp_netif_set_ip_info(ap_netif_g, &ip_info));
    ESP_LOGI(TAG_AP_DNS, "AP Static IP configured: " IPSTR, IP2STR(&ip_info.ip));
    
    // Configure DHCP server to give AP's IP as DNS server
    esp_netif_dns_info_t dns_server_info;
    inet_pton(AF_INET, AP_IP_ADDR, &dns_server_info.ip.u_addr.ip4); // Set DNS server to AP's own IP
    dns_server_info.ip.type = ESP_IPADDR_TYPE_V4;
    dhcps_offer_t dhcps_dns_opt = OFFER_DNS; // Offer DNS server option
    ESP_ERROR_CHECK(esp_netif_dhcps_option(ap_netif_g, ESP_NETIF_OP_SET, ESP_NETIF_DOMAIN_NAME_SERVER, &dns_server_info, sizeof(dns_server_info)));
    ESP_ERROR_CHECK(esp_netif_dhcps_option(ap_netif_g, ESP_NETIF_OP_SET, ESP_NETIF_DHCP_OPTION_OFFER_DNS, &dhcps_dns_opt, sizeof(dhcps_dns_opt)));
    ESP_LOGI(TAG_AP_DNS, "DHCP server configured to offer AP IP as DNS.");

    ESP_ERROR_CHECK(esp_netif_dhcps_start(ap_netif_g)); // Start DHCP server
    ESP_LOGI(TAG_AP_DNS, "DHCP server started.");

    wifi_init_config_t cfg = WIFI_INIT_CONFIG_DEFAULT();
    ESP_ERROR_CHECK(esp_wifi_init(&cfg));

    // No specific Wi-Fi event handler needed for this simple DNS server example,
    // but you'd have one for AP client connections in a real app.

    wifi_config_t wifi_ap_config = {
        .ap = {
            .ssid = AP_ESP_WIFI_SSID,
            .ssid_len = strlen(AP_ESP_WIFI_SSID),
            .password = AP_ESP_WIFI_PASS,
            .max_connection = AP_MAX_STA_CONN,
            .authmode = WIFI_AUTH_WPA2_PSK,
            .channel = AP_ESP_WIFI_CHANNEL,
        },
    };
    if (strlen(AP_ESP_WIFI_PASS) == 0) {
        wifi_ap_config.ap.authmode = WIFI_AUTH_OPEN;
    }

    ESP_ERROR_CHECK(esp_wifi_set_mode(WIFI_MODE_AP));
    ESP_ERROR_CHECK(esp_wifi_set_config(WIFI_IF_AP, &wifi_ap_config));
    ESP_ERROR_CHECK(esp_wifi_start());

    ESP_LOGI(TAG_AP_DNS, "SoftAP started. SSID: %s Password: %s", AP_ESP_WIFI_SSID, AP_ESP_WIFI_PASS);

    // Start the DNS server task
    xTaskCreate(dns_server_task, "dns_server_task", 4096, NULL, 5, NULL);
}


// In app_main:
// void app_main(void) {
//     // NVS Init
//     // ...
//     wifi_init_softap_with_dns();
//     // Optionally, start mDNS as well if you want .local resolution on the AP network
//     // start_mdns_service(); // Ensure mDNS operates on ap_netif_g if called here
// }