From bd07b05d40631d8425f3ca7fb06a5b3f24b6dcdf Mon Sep 17 00:00:00 2001
From: s-hadinger <49731213+s-hadinger@users.noreply.github.com>
Date: Tue, 31 Jan 2023 22:53:23 +0100
Subject: [PATCH] Berry add ``mdns`` advanced features and query (#17842)

* Berry add ``mdns`` advanced features and query

* typo
---
 CHANGELOG.md                                  |   1 +
 .../berry_tasmota/src/be_mdns_module.cpp      | 195 +++++++++++++++++-
 2 files changed, 194 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 23ae173f8..db0e80491 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,7 @@ All notable changes to this project will be documented in this file.
 - ESP32 command ``EnergyCols 1..8`` to change number of GUI columns
 - ESP32 command ``EnergyDisplay 1..3`` to change GUI column presentation
 - support for SEN5X gas and air quality sensor by Tyeth Gundry (#17736)
+- Berry add ``mdns`` advanced features and query
 
 ### Breaking Changed
 - Berry energy_ctypes changed with new energy driver
diff --git a/lib/libesp32/berry_tasmota/src/be_mdns_module.cpp b/lib/libesp32/berry_tasmota/src/be_mdns_module.cpp
index 843090650..dad697ba0 100644
--- a/lib/libesp32/berry_tasmota/src/be_mdns_module.cpp
+++ b/lib/libesp32/berry_tasmota/src/be_mdns_module.cpp
@@ -12,6 +12,7 @@
 #ifdef USE_DISCOVERY
 #include "mdns.h"
 #include <string.h>
+#include "IPAddress.h"
 
 //
 // `mdsn.start([hostname:string]) -> nil`
@@ -52,7 +53,7 @@ static void m_mdns_set_hostname(struct bvm *vm, const char * hostname) {
 BE_FUNC_CTYPE_DECLARE(m_mdns_set_hostname, "", "@s")
 
 //
-// `mdns.add_service(service:string, proto:string, port:int, txt:map) -> nil`
+// `mdns.add_service(service:string, proto:string, port:int, txt:map [, instance:string, hostname:string]) -> nil`
 //
 // add a service declaration using the current hostname as instance name, and specify TXT fields as a `map`
 //
@@ -66,6 +67,14 @@ static int32_t m_mdns_add_service(struct bvm *vm) {
     const char* service_type = be_tostring(vm, 1);
     const char* proto = be_tostring(vm, 2);
     uint16_t port = be_toint(vm, 3);
+    const char * instance = nullptr;
+    const char * hostname = nullptr;
+    if (top >= 5 && be_isstring(vm, 5)) {
+      instance = be_tostring(vm, 5);
+    }
+    if (top >= 6 && be_isstring(vm, 6)) {
+      hostname = be_tostring(vm, 6);
+    }
     
     mdns_txt_item_t * txt_items = NULL;
     int32_t txt_num = 0;
@@ -98,7 +107,12 @@ static int32_t m_mdns_add_service(struct bvm *vm) {
         be_pop(vm, 1); /* pop iterator */
       }
     }
-    esp_err_t err = mdns_service_add(NULL, service_type, proto, port, txt_items, txt_num);
+    esp_err_t err;
+    if (hostname == nullptr) {
+      err = mdns_service_add(instance, service_type, proto, port, txt_items, txt_num);
+    } else {
+      err = mdns_service_add_for_host(instance, service_type, proto, hostname, port, txt_items, txt_num);
+    }
     // free all allocated memory
     if (txt_items != NULL) {
       for (uint32_t i = 0; i < txt_num; i++) {
@@ -115,12 +129,189 @@ static int32_t m_mdns_add_service(struct bvm *vm) {
   be_raise(vm, "value_error", "wrong or missing arguments");
 }
 
+// `mdns_service_subtype_add_for_host()` is only available in most recent esp-protocols version
+//
+// This alias makes sure that the compilation succeeds even if the function is not available
+//
+extern "C" esp_err_t __tasmota_mdns_service_subtype_add_for_host (const char *instance_name, const char *service_type, const char *proto, const char *hostname, const char *subtype)
+{ return ESP_OK; }
+extern "C" esp_err_t mdns_service_subtype_add_for_host(const char *instance_name, const char *service_type, const char *proto, const char *hostname, const char *subtype) __attribute__ ((weak, alias ("__tasmota_mdns_service_subtype_add_for_host")));
+
+//
+// `mdns.add_subtype(service:string, proto:string, instance:string, hostname:string, subtype:string) -> nil`
+//
+// add a subtype
+//
+static int32_t m_dns_add_subtype(struct bvm *vm) {
+  int32_t top = be_top(vm);
+  if (top >= 5 && be_isstring(vm, 1) && be_isstring(vm, 2) && be_isstring(vm, 3) && be_isstring(vm, 4) && be_isstring(vm, 5)) {
+    const char* service_type = be_tostring(vm, 1);
+    const char* proto = be_tostring(vm, 2);
+    const char* instance = be_tostring(vm, 3);
+    const char* hostname = be_tostring(vm, 4);
+    const char* subtype = be_tostring(vm, 5);
+
+    esp_err_t err;
+    // Waiting for support of `mdns_service_subtype_add_for_host()`
+    err = mdns_service_subtype_add_for_host(instance, service_type, proto, hostname, subtype);
+    if (err != ESP_OK) {
+      be_raisef(vm, "internal_error", "mdns add_subtype err=%i", err);
+    }
+    be_return_nil(vm);
+  }
+  be_raise(vm, "value_error", "wrong or missing arguments");
+}
+
+//
+// `mdns.add_hostname(hostname:string, ip:string [, ip:string]*) -> nil`
+//
+// add a delegate hostname
+//
+static int32_t m_dns_add_hostname(struct bvm *vm) {
+  int32_t top = be_top(vm);
+  if (top >= 2 && be_isstring(vm, 1) && be_isstring(vm, 2)) {
+    mdns_ip_addr_t * head = nullptr;     // head of the linked list of addresses
+    esp_err_t err = ESP_OK;             // return status
+
+    const char* hostname = be_tostring(vm, 1);
+    const char* ip_text = nullptr;
+    ip_addr_t a;
+    for (uint16_t i = 2; i <= top; i++) {
+      const char* ip_text = be_tostring(vm, i);
+      if (ip_text == nullptr || ip_text[0] == 0) { continue; }    // ignore empty string
+      IPAddress ip;
+      if (!ip.fromString(ip_text)) { err = -1; break; }
+
+      // allocate new slot
+      mdns_ip_addr_t *new_head = (mdns_ip_addr_t*) be_os_malloc(sizeof(mdns_ip_addr_t));
+      if (new_head == nullptr) { err = -2; break; }
+      new_head->next = head;
+      head = new_head;
+
+      // 
+      ip_addr_t *ip_addr = (ip_addr_t*) ip;
+      head->addr.type = ip_addr->type;
+      if (ip_addr->type == ESP_IPADDR_TYPE_V6) {
+        memcpy(head->addr.u_addr.ip6.addr, ip_addr->u_addr.ip6.addr, 16);
+      } else {
+        head->addr.u_addr.ip4.addr = ip_addr->u_addr.ip4.addr;
+      }
+    } while (0);
+
+    if (err == ESP_OK && head != nullptr) {
+      err = mdns_delegate_hostname_add(hostname, head);
+    }
+
+    // deallocate all memory
+    while (head != nullptr) {
+      mdns_ip_addr_t * next = head->next;
+      be_os_free(head);
+      head = next;
+    }
+
+    if (err == -1) {
+      be_raisef(vm, "value_error", "Invalid IP Address '%s'", ip_text);
+    } else if (err == -2) {
+      be_raise(vm, "memory_error", nullptr);
+    } else if (err != ESP_OK) {
+      be_raisef(vm, "value_error", "mdns_delegate_hostname_add err=%i", err);
+    }
+    be_return_nil(vm);
+
+  }
+  be_raise(vm, "value_error", "wrong or missing arguments");
+}
+
+//
+// `mdns.find_service(service:string, proto:string [timeout_ms:int(3000), max_responses:int(20)]) -> map`
+//
+static int32_t m_dns_find_service(struct bvm *vm) {
+  static const char * ip_protocol_str[] = {"v4", "v6", "max"};
+  int32_t top = be_top(vm);
+  if (top >= 2 && be_isstring(vm, 1) && be_isstring(vm, 2)) {
+    const char* service_name = be_tostring(vm, 1);
+    const char* proto = be_tostring(vm, 2);
+    int32_t timeout_ms = 3000;
+    if (top >= 3 && be_isint(vm, 3)) {
+      timeout_ms = be_toint(vm, 3);
+    }
+    int32_t max_responses = 20;
+    if (top >= 4 && be_isint(vm, 4)) {
+      max_responses = be_toint(vm, 4);
+    }
+
+    mdns_result_t * results = NULL;
+    esp_err_t err = mdns_query_ptr(service_name, proto, timeout_ms, max_responses,  &results);
+    if (err != ESP_OK) { be_raisef(vm, "value_error", "mdns_query_ptr err=%i", err); }
+    if (results == NULL) { be_return_nil(vm); }
+
+    mdns_result_t * r = results;
+    mdns_ip_addr_t * a = NULL;
+    be_newobject(vm, "list");
+    while (r) {
+      be_newobject(vm, "map");
+      be_map_insert_str(vm, "type", ip_protocol_str[r->ip_protocol]);
+      if (r->instance_name) { be_map_insert_str(vm, "instance", r->instance_name); }
+      if (r->hostname) { be_map_insert_str(vm, "hostname", r->hostname); }
+      // TXT
+      be_pushstring(vm, "txt");
+      be_newobject(vm, "map");
+      for (int32_t t=0; t < r->txt_count; t++){
+        be_map_insert_str(vm, r->txt[t].key, r->txt[t].value);
+      }
+      //
+      be_pop(vm, 1);
+      be_data_insert(vm, -3);
+      be_pop(vm, 2);
+      //
+
+      // IP addresses
+      be_pushstring(vm, "ip");
+      be_newobject(vm, "list");
+      //
+      for (a = r->addr; a != NULL; a = a->next) {
+        ip_addr_t ip_addr;
+        if (a->addr.type == IPADDR_TYPE_V6) {
+          ip_addr_copy_from_ip6(ip_addr, a->addr.u_addr.ip6);
+        } else if (a->addr.type == IPADDR_TYPE_V4) {
+          ip_addr_copy_from_ip4(ip_addr, a->addr.u_addr.ip4);
+        } else {
+          continue;
+        }
+        be_pushstring(vm, IPAddress(ip_addr).toString().c_str());
+        be_data_push(vm, -2);
+        be_pop(vm, 1);
+      }
+      //
+      be_pop(vm, 1);
+      be_data_insert(vm, -3);
+      be_pop(vm, 2);
+
+      be_pop(vm, 1);
+      be_data_push(vm, -2);
+      be_pop(vm, 1);
+
+      r = r->next;
+    }
+    be_pop(vm, 1);                  // now list is on top
+
+    mdns_query_results_free(results);
+    be_return(vm);
+  }
+  be_raise(vm, "value_error", "wrong or missing arguments");
+}
+
 /* @const_object_info_begin
 module mdns (scope: global) {
     start, ctype_func(m_mdns_start)
     stop, ctype_func(m_mdns_stop)
     set_hostname, ctype_func(m_mdns_set_hostname)
     add_service, func(m_mdns_add_service)
+    add_hostname, func(m_dns_add_hostname)
+    add_subtype, func(m_dns_add_subtype)
+
+    // querying
+    find_service, func(m_dns_find_service)
 }
 @const_object_info_end */
 #include "be_fixed_mdns.h"