Merge pull request micropython#8262 from thzinc/add-mdns-txt-record-support

Add mDNS TXT record support

Author: tannewt

locale/circuitpython.pot

@@ -1018,6 +1018,15 @@ msgstr ""
 msgid "Failed to acquire mutex, err 0x%04x"
 msgstr ""
 
+#: ports/raspberrypi/common-hal/mdns/Server.c
+msgid "Failed to add service TXT record"
+msgstr ""
+
+#: shared-bindings/mdns/Server.c
+msgid ""
+"Failed to add service TXT record; non-string or bytes found in txt_records"
+msgstr ""
+
 #: shared-module/rgbmatrix/RGBMatrix.c
 msgid "Failed to allocate %q buffer"
 msgstr ""

ports/espressif/common-hal/mdns/Server.c

@@ -208,10 +208,25 @@ mp_obj_t common_hal_mdns_server_find(mdns_server_obj_t *self, const char *servic
     return MP_OBJ_FROM_PTR(tuple);
 }
 
-void common_hal_mdns_server_advertise_service(mdns_server_obj_t *self, const char *service_type, const char *protocol, mp_int_t port) {
+void common_hal_mdns_server_advertise_service(mdns_server_obj_t *self, const char *service_type, const char *protocol, mp_int_t port, const char *txt_records[], size_t num_txt_records) {
     if (mdns_service_exists(service_type, protocol, NULL)) {
         mdns_service_port_set(service_type, protocol, port);
     } else {
+        // TODO: Add support for TXT record
+        /* NOTE: The `mdns_txt_item_t *txt` argument of mdns_service_add uses a struct
+         * that splits out the TXT record into keys and values, though it seems little
+         * is done with those fields aside from concatenating them with an optional
+         * equals sign and calculating the total length of the concatenated string.
+         *
+         * There should be little issue with the underlying implementation to populate
+         * the mdns_txt_item_t struct with only a key containing exactly the desired TXT
+         * record. As long as the underlying implementation calculates the length of the
+         * key + NULL value correctly, it should work.
+         *
+         * Ref: RFC 6763, section 6.1:
+         * > The format of each constituent string within the DNS TXT record is a single
+         * > length byte, followed by 0-255 bytes of text data.
+         */
         mdns_service_add(NULL, service_type, protocol, port, NULL, 0);
     }
 }

ports/raspberrypi/common-hal/mdns/Server.c

@@ -30,6 +30,7 @@
 #include "py/runtime.h"
 #include "shared/runtime/interrupt_char.h"
 #include "shared-bindings/mdns/RemoteService.h"
+#include "shared-bindings/mdns/Server.h"
 #include "shared-bindings/wifi/__init__.h"
 #include "supervisor/shared/tick.h"
 
@@ -295,7 +296,27 @@ mp_obj_t common_hal_mdns_server_find(mdns_server_obj_t *self, const char *servic
     return MP_OBJ_FROM_PTR(tuple);
 }
 
-void common_hal_mdns_server_advertise_service(mdns_server_obj_t *self, const char *service_type, const char *protocol, mp_int_t port) {
+STATIC void srv_txt_cb(struct mdns_service *service, void *ptr) {
+    mdns_server_obj_t *self = ptr;
+    err_t res;
+    for (size_t i = 0; i < self->num_txt_records; i++) {
+        res = mdns_resp_add_service_txtitem(service, self->txt_records[i], strlen(self->txt_records[i]));
+        if (res != ERR_OK) {
+            mp_raise_RuntimeError(MP_ERROR_TEXT("Failed to add service TXT record"));
+            return;
+        }
+    }
+}
+
+STATIC void assign_txt_records(mdns_server_obj_t *self, const char *txt_records[], size_t num_txt_records) {
+    size_t allowed_num_txt_records = MDNS_MAX_TXT_RECORDS < num_txt_records ? MDNS_MAX_TXT_RECORDS : num_txt_records;
+    self->num_txt_records = allowed_num_txt_records;
+    for (size_t i = 0; i < allowed_num_txt_records; i++) {
+        self->txt_records[i] = txt_records[i];
+    }
+}
+
+void common_hal_mdns_server_advertise_service(mdns_server_obj_t *self, const char *service_type, const char *protocol, mp_int_t port, const char *txt_records[], size_t num_txt_records) {
     enum mdns_sd_proto proto = DNSSD_PROTO_UDP;
     if (strcmp(protocol, "_tcp") == 0) {
         proto = DNSSD_PROTO_TCP;
@@ -313,7 +334,9 @@ void common_hal_mdns_server_advertise_service(mdns_server_obj_t *self, const cha
     if (existing_slot < MDNS_MAX_SERVICES) {
         mdns_resp_del_service(NETIF_STA, existing_slot);
     }
-    int8_t slot = mdns_resp_add_service(NETIF_STA, self->instance_name, service_type, proto, port, NULL, NULL);
+
+    assign_txt_records(self, txt_records, num_txt_records);
+    int8_t slot = mdns_resp_add_service(NETIF_STA, self->instance_name, service_type, proto, port, srv_txt_cb, self);
     if (slot < 0) {
         mp_raise_RuntimeError(MP_ERROR_TEXT("Out of MDNS service slots"));
         return;

ports/raspberrypi/common-hal/mdns/Server.h

@@ -30,12 +30,16 @@
 
 #include "lwip/apps/mdns_opts.h"
 
+#define MDNS_MAX_TXT_RECORDS 32
+
 typedef struct {
     mp_obj_base_t base;
     const char *hostname;
     const char *instance_name;
     char default_hostname[sizeof("cpy-XXXXXX")];
     const char *service_type[MDNS_MAX_SERVICES];
+    size_t num_txt_records;
+    const char *txt_records[MDNS_MAX_TXT_RECORDS];
     // Track if this object owns access to the underlying MDNS service.
     bool inited;
 } mdns_server_obj_t;

shared-bindings/mdns/Server.c

@@ -27,7 +27,9 @@
 
 #include <string.h>
 
+#include "py/obj.h"
 #include "py/objproperty.h"
+#include "py/objstr.h"
 #include "py/runtime.h"
 #include "shared-bindings/mdns/__init__.h"
 #include "shared-bindings/mdns/Server.h"
@@ -173,20 +175,26 @@ STATIC MP_DEFINE_CONST_FUN_OBJ_KW(mdns_server_find_obj, 1, _mdns_server_find);
 //|
 //|         If web workflow is active, the port it uses can't also be used to advertise a service.
 //|
+//|         **Limitations**: Publishing up to 32 TXT records is only supported on the RP2040 Pico W board at
+//|         this time.
+//|
 //|         :param str service_type: The service type such as "_http"
 //|         :param str protocol: The service protocol such as "_tcp"
-//|         :param int port: The port used by the service"""
+//|         :param int port: The port used by the service
+//|         :param Sequence[str] txt_records: An optional sequence of strings to serve as TXT records along with the service
+//|         """
 //|         ...
 //|
 STATIC mp_obj_t mdns_server_advertise_service(mp_uint_t n_args, const mp_obj_t *pos_args, mp_map_t *kw_args) {
     mdns_server_obj_t *self = MP_OBJ_TO_PTR(pos_args[0]);
     check_for_deinit(self);
 
-    enum { ARG_service_type, ARG_protocol, ARG_port };
+    enum { ARG_service_type, ARG_protocol, ARG_port, ARG_txt_records };
     static const mp_arg_t allowed_args[] = {
         { MP_QSTR_service_type, MP_ARG_KW_ONLY | MP_ARG_REQUIRED | MP_ARG_OBJ },
         { MP_QSTR_protocol, MP_ARG_KW_ONLY | MP_ARG_REQUIRED | MP_ARG_OBJ },
         { MP_QSTR_port, MP_ARG_KW_ONLY | MP_ARG_REQUIRED | MP_ARG_INT },
+        { MP_QSTR_txt_records, MP_ARG_KW_ONLY | MP_ARG_OBJ, {.u_obj = mp_const_none} },
     };
 
     mp_arg_val_t args[MP_ARRAY_SIZE(allowed_args)];
@@ -195,7 +203,21 @@ STATIC mp_obj_t mdns_server_advertise_service(mp_uint_t n_args, const mp_obj_t *
     const char *service_type = mp_obj_str_get_str(args[ARG_service_type].u_obj);
     const char *protocol = mp_obj_str_get_str(args[ARG_protocol].u_obj);
 
-    common_hal_mdns_server_advertise_service(self, service_type, protocol, args[ARG_port].u_int);
+    const mp_obj_t txt_records = args[ARG_txt_records].u_obj;
+    const size_t num_txt_records = txt_records == mp_const_none
+        ? 0
+        : (size_t)MP_OBJ_SMALL_INT_VALUE(mp_obj_len(txt_records));
+
+    const char *txt_records_array[num_txt_records];
+    for (size_t i = 0; i < num_txt_records; i++) {
+        mp_obj_t txt_record = mp_obj_subscr(txt_records, MP_OBJ_NEW_SMALL_INT(i), MP_OBJ_SENTINEL);
+        if (!mp_obj_is_str_or_bytes(txt_record)) {
+            mp_raise_ValueError(MP_ERROR_TEXT("Failed to add service TXT record; non-string or bytes found in txt_records"));
+        }
+        txt_records_array[i] = mp_obj_str_get_str(txt_record);
+    }
+
+    common_hal_mdns_server_advertise_service(self, service_type, protocol, args[ARG_port].u_int, txt_records_array, num_txt_records);
     return mp_const_none;
 }
 STATIC MP_DEFINE_CONST_FUN_OBJ_KW(mdns_server_advertise_service_obj, 1, mdns_server_advertise_service);

shared-bindings/mdns/Server.h

@@ -42,7 +42,18 @@ void common_hal_mdns_server_set_hostname(mdns_server_obj_t *self, const char *ho
 const char *common_hal_mdns_server_get_instance_name(mdns_server_obj_t *self);
 void common_hal_mdns_server_set_instance_name(mdns_server_obj_t *self, const char *instance_name);
 mp_obj_t common_hal_mdns_server_find(mdns_server_obj_t *self, const char *service_type, const char *protocol, mp_float_t timeout);
-void common_hal_mdns_server_advertise_service(mdns_server_obj_t *self, const char *service_type, const char *protocol, mp_int_t port);
+
+/**
+ * @brief Advertises service
+ *
+ * @param self
+ * @param service_type A string indicating the DNS-SD type of service being advertised (e.g., _http)
+ * @param protocol A string indicating the DNS-SD protocol of the service (e.g., _tcp or _udp)
+ * @param port The TCP or UDP port number of the service
+ * @param txt_records An array of strings representing TXT records to publish along with the service
+ * @param num_txt_records Number of records expected in txt_records
+ */
+void common_hal_mdns_server_advertise_service(mdns_server_obj_t *self, const char *service_type, const char *protocol, mp_int_t port, const char *txt_records[], size_t num_txt_records);
 
 // For internal use.
 void mdns_server_construct(mdns_server_obj_t *self, bool workflow);

supervisor/shared/web_workflow/web_workflow.c

@@ -355,7 +355,7 @@ bool supervisor_start_web_workflow(bool reload) {
             }
         }
         if (!common_hal_mdns_server_deinited(&mdns)) {
-            common_hal_mdns_server_advertise_service(&mdns, "_circuitpython", "_tcp", web_api_port);
+            common_hal_mdns_server_advertise_service(&mdns, "_circuitpython", "_tcp", web_api_port, NULL, 0);
         }
         #endif