From 86623e8c6689d6d149e902db93255953c0a763d1 Mon Sep 17 00:00:00 2001
From: Jade Mattsson <jmattsson@dius.com.au>
Date: Thu, 12 Dec 2024 17:47:06 +1100
Subject: [PATCH] Add mDNS module.

---
 components/modules/CMakeLists.txt    |   1 +
 components/modules/Kconfig           |   6 +
 components/modules/idf_component.yml |   3 +-
 components/modules/mdns.c            | 296 +++++++++++++++++++++++++++
 docs/modules/mdns.md                 | 204 ++++++++++++++++++
 mkdocs.yml                           |   1 +
 6 files changed, 510 insertions(+), 1 deletion(-)
 create mode 100644 components/modules/mdns.c
 create mode 100644 docs/modules/mdns.md

diff --git a/components/modules/CMakeLists.txt b/components/modules/CMakeLists.txt
index 0fab0691b..b1e543baf 100644
--- a/components/modules/CMakeLists.txt
+++ b/components/modules/CMakeLists.txt
@@ -25,6 +25,7 @@ set(module_srcs
   "i2c_hw_master.c"
   "i2c_hw_slave.c"
   "ledc.c"
+  "mdns.c"
   "mqtt.c"
   "net.c"
   "node.c"
diff --git a/components/modules/Kconfig b/components/modules/Kconfig
index c393ab6b9..7c15dd518 100644
--- a/components/modules/Kconfig
+++ b/components/modules/Kconfig
@@ -177,6 +177,12 @@ menu "NodeMCU modules"
         help
             Includes the LEDC module.
 
+    config NODEMCU_CMODULE_MDNS
+        bool "mDNS module"
+        default "n"
+        help
+            Includes the mDNS module.
+
     config NODEMCU_CMODULE_MQTT
         bool "MQTT module"
         default "n"
diff --git a/components/modules/idf_component.yml b/components/modules/idf_component.yml
index ae9c53017..a8458cf5e 100644
--- a/components/modules/idf_component.yml
+++ b/components/modules/idf_component.yml
@@ -4,4 +4,5 @@ dependencies:
   idf:
     version: ">=5.0.0"
   # Component dependencies
-  libsodium: "~1.0.20"
+  espressif/libsodium: "~1.0.20"
+  espressif/mdns: "^1.4.2"
diff --git a/components/modules/mdns.c b/components/modules/mdns.c
new file mode 100644
index 000000000..8405fdc55
--- /dev/null
+++ b/components/modules/mdns.c
@@ -0,0 +1,296 @@
+#include "module.h"
+#include "lauxlib.h"
+#include "ip_fmt.h"
+
+#include "esp_err.h"
+#include "mdns.h"
+
+// Table key names
+static const char *HOSTNAME = "hostname";
+static const char *INSTANCE_NAME = "instance_name";
+static const char *SERVICES = "services";
+static const char *SERVICE_TYPE = "service_type";
+static const char *PROTO = "protocol";
+static const char *PORT = "port";
+static const char *TXT = "txt";
+static const char *SUBTYPE = "subtype";
+static const char *QUERY_TYPE = "query_type";
+static const char *NAME = "name";
+static const char *TIMEOUT = "timeout";
+static const char *MAX_RESULTS = "max_results";
+static const char *ADDRESSES = "addresses";
+
+#define DEFAULT_TIMEOUT_MS 2000
+#define DEFAULT_MAX_RESULTS 10
+
+static bool started;
+
+static bool valid_query_type(int t)
+{
+  switch(t)
+  {
+    case MDNS_TYPE_A:
+    case MDNS_TYPE_PTR:
+    case MDNS_TYPE_TXT:
+    case MDNS_TYPE_AAAA:
+    case MDNS_TYPE_SRV:
+    //case MDNS_TYPE_OPT:
+    //case MDNS_TYPE_NSEC:
+    case MDNS_TYPE_ANY: return true;
+    default: return false;
+  }
+}
+
+
+static int lmdns_start(lua_State *L)
+{
+  luaL_checktable(L, 1);
+  lua_settop(L, 1);
+
+  if (started)
+    return luaL_error(L, "already started");
+
+  bool inited = false;
+  esp_err_t err = mdns_init();
+  if (err != ESP_OK)
+    goto mdns_err;
+  inited = true;
+
+  // Hostname
+  lua_getfield(L, 1, HOSTNAME);
+  const char *hostname = luaL_optstring(L, -1, NULL);
+  if (hostname)
+  {
+    err = mdns_hostname_set(hostname);
+    if (err != ESP_OK)
+      goto mdns_err;
+  }
+  lua_pop(L, 1);
+
+  // Instance name
+  lua_getfield(L, 1, INSTANCE_NAME);
+  const char *instname = luaL_optstring(L, -1, NULL);
+  if (instname)
+  {
+    err = mdns_instance_name_set(instname);
+    if (err != ESP_OK)
+      goto mdns_err;
+  }
+  lua_pop(L, 1);
+
+  // Services
+  lua_getfield(L, 1, SERVICES);
+  unsigned i = 1;
+  if (!lua_isnoneornil(L, 2)) // array of service entries
+  {
+    luaL_checktable(L, 2);
+    for (i = 1; true; ++i)
+    {
+      lua_rawgeti(L, 2, i);
+      if (!lua_istable(L, 3))
+        break;
+
+      lua_getfield(L, 3, SERVICE_TYPE);
+      const char *svctype = luaL_checkstring(L, -1);
+
+      lua_getfield(L, 3, PROTO);
+      const char *proto = luaL_checkstring(L, -1);
+
+      lua_getfield(L, 3, PORT);
+      int port = luaL_checkint(L, -1);
+
+      lua_getfield(L, 3, INSTANCE_NAME);
+      const char *instname2 = luaL_optstring(L, -1, NULL);
+
+      // Note: we add txt entries iteratively to avoid having to size and
+      // allocate a buffer to hold them all.
+      err = mdns_service_add(instname2, svctype, proto, port, NULL, 0);
+      if (err != ESP_OK)
+        goto mdns_err;
+
+      lua_pop(L, 4); // svctype, proto, port, instname2
+
+      lua_getfield(L, 3, TXT);
+      if (lua_istable(L, 4))
+      {
+        lua_pushnil(L); // 5 is now table key
+        while(lua_next(L, 4)) // replaces 5 with actual key
+        {
+          // copy key, value so we can safely tostring() them
+          lua_pushvalue(L, 5);
+          lua_pushvalue(L, 6);
+
+          const char *key = luaL_checkstring(L, -2);
+          const char *val = luaL_checkstring(L, -1);
+
+          err = mdns_service_txt_item_set_for_host(
+            instname2, svctype, proto, hostname, key, val);
+          if (err != ESP_OK)
+            goto mdns_err;
+
+          lua_pop(L, 3); // value, key, value
+        }
+      }
+      lua_pop(L, 1); // txt table
+
+      // Subtype
+      lua_getfield(L, 1, SUBTYPE);
+      const char *subtype = luaL_optstring(L, -1, NULL);
+      if (subtype)
+      {
+        err = mdns_service_subtype_add_for_host(
+          instname2, svctype, proto, hostname, subtype);
+        if (err != ESP_OK)
+          goto mdns_err;
+      }
+      lua_pop(L, 1); // subtype
+
+      lua_pop(L, 1); // services[i] table
+    }
+  }
+  lua_pop(L, 1); // services array
+
+  started = true;
+
+  // Return number of services we added
+  lua_pushinteger(L, i - 1);
+  return 1;
+
+mdns_err:
+  if (inited)
+  {
+    mdns_service_remove_all();
+    mdns_free();
+  }
+  return luaL_error(L, "mdns error: %s", esp_err_to_name(err));
+}
+
+
+static int lmdns_stop(lua_State *L)
+{
+  if (started)
+  {
+    mdns_service_remove_all();
+    started = false;
+  }
+  mdns_free();
+  return 0;
+}
+
+
+static int lmdns_query(lua_State *L)
+{
+  luaL_checktable(L, 1);
+  lua_settop(L, 1);
+
+  lua_getfield(L, 1, NAME);
+  const char *name = luaL_optstring(L, -1, NULL);
+
+  lua_getfield(L, 1, SERVICE_TYPE);
+  const char *svctype = luaL_optstring(L, -1, NULL);
+
+  lua_getfield(L, 1, PROTO);
+  const char *proto = luaL_optstring(L, -1, NULL);
+
+  lua_getfield(L, 1, QUERY_TYPE);
+  int qtype = luaL_checkint(L, -1);
+  if (!valid_query_type(qtype))
+    return luaL_error(L, "unknown mDNS query type");
+
+  lua_getfield(L, 1, TIMEOUT);
+  int timeout = luaL_optinteger(L, -1, DEFAULT_TIMEOUT_MS);
+
+  lua_getfield(L, 1, MAX_RESULTS);
+  int max_results = luaL_optinteger(L, -1, DEFAULT_MAX_RESULTS);
+
+  mdns_result_t *res = NULL;
+  esp_err_t err =
+    mdns_query(name, svctype, proto, qtype, timeout, max_results, &res);
+  if (err != ESP_OK)
+    return luaL_error(L, "mdns error: %s", esp_err_to_name(err));
+
+  lua_settop(L, 0);
+  lua_createtable(L, max_results, 0); // results array at idx 1
+
+  for (int n = 1; res; ++n, res = res->next)
+  {
+    // Reserve 5 slots, for SRV result host/port/instance/service_type/proto
+    lua_createtable(L, 0, 5); // result entry table at idx 2
+
+    if (res->instance_name)
+    {
+      lua_pushstring(L, res->instance_name);
+      lua_setfield(L, 2, INSTANCE_NAME);
+    }
+    if (res->service_type)
+    {
+      lua_pushstring(L, res->service_type);
+      lua_setfield(L, 2, SERVICE_TYPE);
+    }
+    if (res->proto)
+    {
+      lua_pushstring(L, res->proto);
+      lua_setfield(L, 2, PROTO);
+    }
+    if (res->hostname)
+    {
+      lua_pushstring(L, res->hostname);
+      lua_setfield(L, 2, HOSTNAME);
+    }
+    if (res->port)
+    {
+      lua_pushinteger(L, res->port);
+      lua_setfield(L, 2, PORT);
+    }
+    if (res->txt)
+    {
+      lua_createtable(L, 0, res->txt_count); // txt table at idx 3
+      for (int i = 0; i < res->txt_count; ++i)
+      {
+        lua_pushstring(L, res->txt[i].key);
+        if (res->txt[i].value)
+          lua_pushlstring(L, res->txt[i].value, res->txt_value_len[i]);
+        else
+          lua_pushliteral(L, "");
+        lua_settable(L, 3);
+      }
+      lua_setfield(L, 2, TXT);
+    }
+    if (res->addr)
+    {
+      lua_createtable(L, 1, 0); // address array table at idx 3
+      int i = 1;
+      for (mdns_ip_addr_t *a = res->addr; a; ++i, a = a->next)
+      {
+        char buf[IP_STR_SZ];
+        ipstr_esp(buf, &a->addr);
+        lua_pushstring(L, buf);
+        lua_rawseti(L, 3, i);
+      }
+      lua_setfield(L, 2, ADDRESSES);
+    }
+
+    lua_rawseti(L, 1, n); // insert into array of results
+  }
+
+  mdns_query_results_free(res);
+  return 1;
+}
+
+
+LROT_BEGIN(mdns, NULL, 0)
+  LROT_FUNCENTRY( start,     lmdns_start    )
+  LROT_FUNCENTRY( query,     lmdns_query    )
+  LROT_FUNCENTRY( stop,      lmdns_stop     )
+
+  LROT_NUMENTRY( TYPE_A,     MDNS_TYPE_A    )
+  LROT_NUMENTRY( TYPE_PTR,   MDNS_TYPE_PTR  )
+  LROT_NUMENTRY( TYPE_TXT,   MDNS_TYPE_TXT  )
+  LROT_NUMENTRY( TYPE_AAAA,  MDNS_TYPE_AAAA )
+  LROT_NUMENTRY( TYPE_SRV,   MDNS_TYPE_SRV  )
+  //LROT_NUMENTRY( TYPE_OPT,   MDNS_TYPE_OPT  )
+  //LROT_NUMENTRY( TYPE_NSEC,  MDNS_TYPE_NSEC )
+  LROT_NUMENTRY( TYPE_ANY,   MDNS_TYPE_ANY  )
+LROT_END(mdns, NULL, 0)
+
+NODEMCU_MODULE(MDNS, "mdns", mdns, NULL);
diff --git a/docs/modules/mdns.md b/docs/modules/mdns.md
new file mode 100644
index 000000000..1de5535f5
--- /dev/null
+++ b/docs/modules/mdns.md
@@ -0,0 +1,204 @@
+# mDNS Module
+| Since  | Origin / Contributor  | Maintainer  | Source  |
+| :----- | :-------------------- | :---------- | :------ |
+| 2024-12-11 | [Jade Mattsson](https://github.com/jmattsson) | [Jade Mattsson](https://github.com/jmattsson) | [mdns.c](../../components/modules/mdns.c)|
+
+This module provides access to the mDNS subsystem and allows both registering services that can be discovered and performing service discovery on the local network.
+
+Other names for mDNS include Bonjour and Avahi.
+
+Some aspects of the mDNS subsystem are compile-time configurable via Kconfig (`make menuconfig`). The defaults are likely sufficient for the vast majority of users.
+
+## mdns.start()
+
+Initialises the mDNS subsystem and registers any services for the device.
+
+#### Syntax
+`mdns.start(config)`
+
+#### Parameters
+- `config` Table containing the mDNS service configuration:
+  - `hostname` (Required if any services are to be registered) The hostname to use for mDNS.
+  - `instance_name` (Optional) The default service instance name. Defaults to the hostname if not set explicitly.
+  - `services` (Optional) An array of service entries to register with mDNS, with each entry being a table comprising these fields:
+    - `service_type` (Required) The service type to register, e.g. `"_http"`.
+    - `protocol` (Required) The protocol to register, one of `"_udp"` or `"_tcp"` typically.
+    - `port` (Required) The port number of the service, e.g. `80`.
+    - `subtype` (Optional) The service subtype, if applicable.
+    - `instance_name` (Optional) The instance name of the service. Defaults to the system-wide instance name if not set explicitly.
+    - `txt` (Optional) A table of key/value value pairs to add to the service's `TXT` entry.
+
+#### Returns
+The number of services registered with the mDNS subsystem.
+
+#### Examples
+Enabling mDNS discovery of this device for both HTTP and FTP services.
+
+```lua
+mdns.start({
+  hostname="esp32server",
+  instance_name="My cool ESP32",
+  services={
+    {
+      service_type="_http",
+      protocol="_tcp",
+      port=80,
+      txt={
+        path="/login.html"
+      }
+    },
+    {
+      service_type="_ftp",
+      protocol="_tcp",
+      port=21
+    }
+  }
+})
+```
+
+Starting mDNS without registering any services. Only useful for doing mDNS queries.
+
+```lua
+mdns.start({})
+```
+
+## mdns.query()
+
+Perform an mDNS query.
+
+#### Syntax
+`mdns.query(query)`
+
+#### Parameters
+- `query` Table with the query parameters. Most fields are optional depending on the query type.
+  - `query_type` (Required) The type of mDNS query to issue. One of:
+    - `mdns.TYPE_A` IPv4 address lookup query.
+    - `mdns.TYPE_AAAA` IPv6 address lookup query.
+    - `mdns.TYPE_PTR` PTR record query (find services).
+    - `mdns.TYPE_TXT` TXT record query.
+    - `mdns.TYPE_SRV` SRV record query (find hostname/port for service).
+    - `mdns.TYPE_ANY` Query all record types.
+  - `name` Name to query for.
+  - `service_type` The service type to query for.
+  - `protocol` The transport protocol of the service being queried for (e.g. `"_tcp"` or `"_udp"`.
+  - `timeout` Timeout in milliseconds to wait for responses. Default 2000.
+  - `max_results` Maximum number of responses to return. Default 10.
+
+#### Returns
+A Lua array with the results. Each result is a table. The fields in the table depend on the query type performed.
+
+```lua
+{ {
+    -- PTR results
+    instance_name=,
+    service_type=,
+    protocol=,
+    -- SRV results
+    hostname=,
+    port=,
+    -- TXT results
+    txt={
+      key1=,
+      key2=,
+      ...
+    },
+    -- A and AAAA results
+    addresses={ ip1str, ip2str, ...  }
+  },
+  ...
+}
+```
+
+## mdns.stop()
+
+Unregisters any services and shuts down the mDNS subsystems.
+
+#### Syntax
+`mdns.stop()`
+
+#### Parameters
+None
+
+#### Returns
+`nil`
+
+#### Examples
+Find SMB file shares on the network:
+
+```lua
+r=mdns.query({query_type=mdns.TYPE_PTR,service_type="_smb",protocol="_tcp"})
+dump(r)
+{
+  1={
+    service_type=_smb
+    protocol=_tcp
+    hostname=mynas
+    port=445
+    instance_name=mynas
+  }
+  2={
+    service_type=_smb
+    protocol=_tcp
+    hostname=desktop
+    port=445
+    instance_name=desktop
+  }
+}
+```
+
+Resolve IPv4 address:
+
+```lua
+r=mdns.query({query_type=mdns.TYPE_A,name="mynas"})
+dump(r)
+{
+  1={
+    addresses={
+      1=192.168.1.8
+    }
+    hostname=mynas
+  }
+}
+```
+
+Resolve IPv6 address:
+
+```lua
+r=mdns.query({query_type=mdns.TYPE_AAAA,name="Hue-Study"})
+dump(r)
+{
+  1={
+    addresses={
+      1=2001:44B8:221A:2132:217:88FF:FEB1:E364
+      2=2001:44B8:8C77:EB32:217:88FF:FEB1:E364
+      3=FE80::217:88FF:FEB1:E364
+    }
+    hostname=Hue-Study
+  }
+}
+```
+
+And in case someone wants the `dump` function, it's just:
+
+```lua
+function dump(x, ind, key)
+  local ind = ind or 0
+  local key = key or ""
+  local indent = string.rep("  ", ind)
+  local t = type(x)
+  local prefix=indent..key.."="
+  if x == nil then
+    print(prefix.."nil")
+  elseif (t == "table") then
+    print(prefix.."{")
+    for k,v in pairs(x) do
+      dump(v, ind + 1, k)
+    end
+    print(indent.."}")
+  elseif (t == "number" or t == "string" or t == "boolean") then
+    print(prefix..tostring(x))
+  else
+    print(prefix..t)
+  end
+end
+```
diff --git a/mkdocs.yml b/mkdocs.yml
index 01ea19c09..c62bf105e 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -59,6 +59,7 @@ nav:
       - 'i2c': 'modules/i2c.md'
       - 'i2s': 'modules/i2s.md'
       - 'ledc': 'modules/ledc.md'
+      - 'mdns': 'modules/mdns.md'
       - 'mqtt': 'modules/mqtt.md'
       - 'net': 'modules/net.md'
       - 'node': 'modules/node.md'