1 files changed, 270 insertions, 0 deletions

diff --git a/main/dns_server.c b/main/dns_server.c
new file mode 100644
index 0000000..f7977c6
--- /dev/null
+++ b/main/dns_server.c
@@ -0,0 +1,270 @@
+#include "dns_server.h"
+#include "esp_log.h"
+#include "freertos/FreeRTOS.h"
+#include "freertos/task.h"
+#include "lwip/sockets.h"
+#include "lwip/netdb.h"
+#include <string.h>
+#include <sys/param.h>
+
+#define MAX_AUTH_IPS 10
+#define MAX_PENDING 50
+#define DNS_BUF_SIZE 512
+#define DNS_PORT 53
+#define DNS_TASK_STACK 4096
+#define DNS_TASK_PRIO 5
+#define DNS_FORWARD_TIMEOUT_MS 2000
+#define NXDOMAIN_TTL 30
+
+static const char *TAG = "dns_server";
+
+#pragma pack(push, 1)
+typedef struct {
+    uint16_t id;
+    uint16_t flags;
+    uint16_t qdcount;
+    uint16_t ancount;
+    uint16_t nscount;
+    uint16_t arcount;
+} dns_header_t;
+#pragma pack(pop)
+
+#pragma pack(push, 1)
+typedef struct {
+    uint16_t name;
+    uint16_t type;
+    uint16_t class;
+    uint32_t ttl;
+    uint16_t len;
+    uint32_t addr;
+} dns_answer_t;
+#pragma pack(pop)
+
+typedef struct {
+    uint32_t ip;
+} auth_entry_t;
+
+static auth_entry_t s_auth_list[MAX_AUTH_IPS];
+static int s_auth_count = 0;
+static TaskHandle_t s_dns_task = NULL;
+static volatile bool s_dns_running = false;
+static esp_ip4_addr_t s_ap_ip;
+static esp_ip4_addr_t s_upstream_dns;
+
+static bool is_authenticated(uint32_t ip)
+{
+    for (int i = 0; i < s_auth_count; i++) {
+        if (s_auth_list[i].ip == ip) return true;
+    }
+    return false;
+}
+
+static void parse_dns_name(const uint8_t *buf, int buf_len, int offset, char *out, int out_len)
+{
+    int pos = offset;
+    int out_pos = 0;
+    int jumped = 0;
+    int jump_pos = 0;
+    while (pos < buf_len && out_pos < out_len - 1) {
+        uint8_t len = buf[pos];
+        if (len == 0) break;
+        if ((len & 0xC0) == 0xC0) {
+            if (!jumped) jump_pos = pos + 2;
+            pos = ((len & 0x3F) << 8) | buf[pos + 1];
+            jumped = 1;
+            continue;
+        }
+        if (out_pos > 0 && out_pos < out_len - 1) out[out_pos++] = '.';
+        pos++;
+        for (int i = 0; i < len && pos < buf_len && out_pos < out_len - 1; i++) {
+            out[out_pos++] = buf[pos++];
+        }
+    }
+    out[out_pos] = '\0';
+}
+
+static int build_nxdomain(uint8_t *response, int req_len)
+{
+    memcpy(response, response, req_len);
+    dns_header_t *hdr = (dns_header_t *)response;
+    hdr->flags = htons(0x8403);
+    hdr->ancount = 0;
+    hdr->nscount = 0;
+    hdr->arcount = 0;
+    return req_len;
+}
+
+static int build_redirect_response(uint8_t *response, int req_len)
+{
+    memcpy(response, response, req_len);
+    dns_header_t *hdr = (dns_header_t *)response;
+    hdr->flags = htons(0x8180);
+    hdr->ancount = htons(1);
+    hdr->nscount = 0;
+    hdr->arcount = 0;
+    int resp_len = req_len;
+    dns_answer_t ans;
+    ans.name = htons(0xC00C);
+    ans.type = htons(1);
+    ans.class = htons(1);
+    ans.ttl = htonl(NXDOMAIN_TTL);
+    ans.len = htons(4);
+    ans.addr = s_ap_ip.addr;
+    memcpy(response + resp_len, &ans, sizeof(ans));
+    resp_len += sizeof(ans);
+    return resp_len;
+}
+
+static int forward_dns(const uint8_t *req, int req_len, uint8_t *resp, int resp_buf_len,
+                       const struct sockaddr_in *client_addr, uint16_t txn_id)
+{
+    int upstream_sock = socket(AF_INET, SOCK_DGRAM, 0);
+    if (upstream_sock < 0) return -1;
+
+    struct timeval tv = { .tv_sec = DNS_FORWARD_TIMEOUT_MS / 1000, .tv_usec = (DNS_FORWARD_TIMEOUT_MS % 1000) * 1000 };
+    setsockopt(upstream_sock, SOL_SOCKET, SO_RCVTIMEO, &tv, sizeof(tv));
+
+    struct sockaddr_in upstream_addr = {
+        .sin_family = AF_INET,
+        .sin_port = htons(DNS_PORT),
+        .sin_addr.s_addr = s_upstream_dns.addr,
+    };
+
+    sendto(upstream_sock, req, req_len, 0, (struct sockaddr *)&upstream_addr, sizeof(upstream_addr));
+
+    int n = recvfrom(upstream_sock, resp, resp_buf_len, 0, NULL, NULL);
+    close(upstream_sock);
+
+    if (n > 0) {
+        if (n >= sizeof(dns_header_t)) {
+            dns_header_t *hdr = (dns_header_t *)resp;
+            hdr->id = htons(txn_id);
+        }
+    }
+    return n;
+}
+
+static void dns_server_task(void *arg)
+{
+    int sock = socket(AF_INET, SOCK_DGRAM, 0);
+    if (sock < 0) {
+        ESP_LOGE(TAG, "Failed to create DNS socket");
+        s_dns_running = false;
+        vTaskDelete(NULL);
+        return;
+    }
+
+    struct sockaddr_in bind_addr = {
+        .sin_family = AF_INET,
+        .sin_port = htons(DNS_PORT),
+        .sin_addr.s_addr = INADDR_ANY,
+    };
+    if (bind(sock, (struct sockaddr *)&bind_addr, sizeof(bind_addr)) < 0) {
+        ESP_LOGE(TAG, "Failed to bind DNS socket");
+        close(sock);
+        s_dns_running = false;
+        vTaskDelete(NULL);
+        return;
+    }
+
+    ESP_LOGI(TAG, "DNS server started on port %d, AP IP=" IPSTR ", upstream DNS=" IPSTR,
+             DNS_PORT, IP2STR(&s_ap_ip), IP2STR(&s_upstream_dns));
+
+    uint8_t rx_buf[DNS_BUF_SIZE];
+    uint8_t tx_buf[DNS_BUF_SIZE + sizeof(dns_answer_t)];
+
+    while (s_dns_running) {
+        struct sockaddr_in client_addr;
+        socklen_t client_len = sizeof(client_addr);
+        int n = recvfrom(sock, rx_buf, sizeof(rx_buf), 0,
+                         (struct sockaddr *)&client_addr, &client_len);
+        if (n < (int)sizeof(dns_header_t)) continue;
+
+        uint32_t client_ip = client_addr.sin_addr.s_addr;
+        dns_header_t *hdr = (dns_header_t *)rx_buf;
+        uint16_t txn_id = ntohs(hdr->id);
+        bool is_query = (ntohs(hdr->flags) & 0x8000) == 0;
+        uint16_t qdcount = ntohs(hdr->qdcount);
+
+        if (!is_query || qdcount == 0) continue;
+
+        int q_offset = sizeof(dns_header_t);
+        while (q_offset < n && rx_buf[q_offset] != 0) {
+            q_offset += rx_buf[q_offset] + 1;
+        }
+        if (q_offset + 5 > n) continue;
+        uint16_t qtype = (rx_buf[q_offset + 1] << 8) | rx_buf[q_offset + 2];
+        int req_len = q_offset + 5;
+
+        if (is_authenticated(client_ip)) {
+            int resp_len = forward_dns(rx_buf, req_len, tx_buf, sizeof(tx_buf), &client_addr, txn_id);
+            if (resp_len > 0) {
+                sendto(sock, tx_buf, resp_len, 0, (struct sockaddr *)&client_addr, client_len);
+            }
+        } else {
+            if (qtype == 1) {
+                int resp_len = build_redirect_response(rx_buf, req_len);
+                memcpy(tx_buf, rx_buf, resp_len);
+                dns_header_t *resp_hdr = (dns_header_t *)tx_buf;
+                resp_hdr->id = htons(txn_id);
+                sendto(sock, tx_buf, resp_len, 0, (struct sockaddr *)&client_addr, client_len);
+            } else if (qtype == 28) {
+                int resp_len = build_nxdomain(rx_buf, req_len);
+                memcpy(tx_buf, rx_buf, resp_len);
+                dns_header_t *resp_hdr = (dns_header_t *)tx_buf;
+                resp_hdr->id = htons(txn_id);
+                sendto(sock, tx_buf, resp_len, 0, (struct sockaddr *)&client_addr, client_len);
+            } else {
+                int resp_len = forward_dns(rx_buf, req_len, tx_buf, sizeof(tx_buf), &client_addr, txn_id);
+                if (resp_len > 0) {
+                    sendto(sock, tx_buf, resp_len, 0, (struct sockaddr *)&client_addr, client_len);
+                }
+            }
+        }
+    }
+
+    close(sock);
+    ESP_LOGI(TAG, "DNS server stopped");
+    vTaskDelete(NULL);
+}
+
+esp_err_t dns_server_start(esp_ip4_addr_t ap_ip, esp_ip4_addr_t upstream_dns)
+{
+    if (s_dns_running) return ESP_OK;
+    s_ap_ip = ap_ip;
+    s_upstream_dns = upstream_dns;
+    s_dns_running = true;
+    xTaskCreate(dns_server_task, "dns_server", DNS_TASK_STACK, NULL, DNS_TASK_PRIO, &s_dns_task);
+    return ESP_OK;
+}
+
+void dns_server_stop(void)
+{
+    s_dns_running = false;
+    vTaskDelay(pdMS_TO_TICKS(200));
+    s_dns_task = NULL;
+}
+
+void dns_server_set_client_authenticated(uint32_t client_ip, bool authenticated)
+{
+    if (authenticated) {
+        if (is_authenticated(client_ip)) return;
+        if (s_auth_count < MAX_AUTH_IPS) {
+            s_auth_list[s_auth_count].ip = client_ip;
+            s_auth_count++;
+        }
+    } else {
+        for (int i = 0; i < s_auth_count; i++) {
+            if (s_auth_list[i].ip == client_ip) {
+                s_auth_list[i] = s_auth_list[s_auth_count - 1];
+                s_auth_count--;
+                return;
+            }
+        }
+    }
+}
+
+bool dns_server_is_running(void)
+{
+    return s_dns_running;
+}