feat(mdns): Support user defined allocators

Defines mem-alloc function optionally weak, so users can override them and implement their own allocation, or a static/bss memory for the mdns task stack.
david-cermak · Feb 18, 2025 · 88162d1 · 88162d1
1 parent e65f22a
commit 88162d1
Show file tree

Hide file tree

Showing 7 changed files with 109 additions and 62 deletions.
diff --git a/components/mdns/Kconfig b/components/mdns/Kconfig
@@ -61,27 +61,47 @@ menu "mDNS"
         default 0x0 if MDNS_TASK_AFFINITY_CPU0
         default 0x1 if MDNS_TASK_AFFINITY_CPU1
 
-    choice MDNS_TASK_MEMORY_ALLOC_FROM
-        prompt "Select mDNS task create on which type of memory"
-        default MDNS_TASK_CREATE_FROM_INTERNAL
-        config MDNS_TASK_CREATE_FROM_SPIRAM
-            bool "mDNS task creates on the SPIRAM"
-            depends on (SPIRAM_USE_CAPS_ALLOC || SPIRAM_USE_MALLOC)
-        config MDNS_TASK_CREATE_FROM_INTERNAL
-            bool "mDNS task creates on the internal RAM"
-    endchoice
+    menu "MDNS Memory Configuration"
 
-    choice MDNS_MEMORY_ALLOC_FROM
-        prompt "Select mDNS memory allocation type"
-        default MDNS_MEMORY_ALLOC_INTERNAL
+        choice MDNS_TASK_MEMORY_ALLOC_FROM
+            prompt "Select mDNS task create on which type of memory"
+            default MDNS_TASK_CREATE_FROM_INTERNAL
+            config MDNS_TASK_CREATE_FROM_SPIRAM
+                bool "mDNS task creates on the SPIRAM (READ HELP)"
+                depends on (SPIRAM_USE_CAPS_ALLOC || SPIRAM_USE_MALLOC)
+                help
+                    mDNS task creates on the SPIRAM.
+                    This option requires FreeRTOS component to allow creating
+                    tasks on the external memory.
+                    Please read the documentation about FREERTOS_TASK_CREATE_ALLOW_EXT_MEM
 
-        config MDNS_MEMORY_ALLOC_SPIRAM
-            bool "Allocate mDNS memory from SPIRAM"
-            depends on (SPIRAM_USE_CAPS_ALLOC || SPIRAM_USE_MALLOC)
+            config MDNS_TASK_CREATE_FROM_INTERNAL
+                bool "mDNS task creates on the internal RAM"
 
-        config MDNS_MEMORY_ALLOC_INTERNAL
-            bool "Allocate mDNS memory from internal RAM"
-    endchoice
+        endchoice
+
+        choice MDNS_MEMORY_ALLOC_FROM
+            prompt "Select mDNS memory allocation type"
+            default MDNS_MEMORY_ALLOC_INTERNAL
+
+            config MDNS_MEMORY_ALLOC_SPIRAM
+                bool "Allocate mDNS memory from SPIRAM"
+                depends on (SPIRAM_USE_CAPS_ALLOC || SPIRAM_USE_MALLOC)
+
+            config MDNS_MEMORY_ALLOC_INTERNAL
+                bool "Allocate mDNS memory from internal RAM"
+
+        endchoice
+
+        config MDNS_MEMORY_CUSTOM_IMPL
+            bool "Implement custom memory functions"
+            default n
+            help
+                Enable to implement custom memory functions for mDNS library.
+                This option is useful when the application wants to use custom
+                memory allocation functions for mDNS library.
+
+    endmenu # MDNS Memory Configuration
 
     config MDNS_SERVICE_ADD_TIMEOUT_MS
         int "mDNS adding service timeout (ms)"

diff --git a/components/mdns/mdns.c b/components/mdns/mdns.c
@@ -5466,12 +5466,10 @@ static esp_err_t _mdns_stop_timer(void)
 
 static esp_err_t _mdns_task_create_with_caps(void)
 {
-    ESP_LOGI(TAG, "mDNS task will be created from %s", MDNS_TASK_MEMORY_LOG);
     esp_err_t ret = ESP_OK;
     static StaticTask_t mdns_task_buffer;
 
-    // Allocate memory for the mDNS task's stack using the MDNS_TASK_MEMORY_CAPS
-    _mdns_stack_buffer = heap_caps_malloc(MDNS_SERVICE_STACK_DEPTH, MDNS_TASK_MEMORY_CAPS);
+    _mdns_stack_buffer = mdns_mem_task_malloc(MDNS_SERVICE_STACK_DEPTH);
     ESP_GOTO_ON_FALSE(_mdns_stack_buffer != NULL, ESP_FAIL, err, TAG, "failed to allocate memory for the mDNS task's stack");
 
     _mdns_service_task_handle = xTaskCreateStaticPinnedToCore(_mdns_service_task, "mdns", MDNS_SERVICE_STACK_DEPTH, NULL, MDNS_TASK_PRIORITY, _mdns_stack_buffer, &mdns_task_buffer, MDNS_TASK_AFFINITY);
@@ -5480,7 +5478,7 @@ static esp_err_t _mdns_task_create_with_caps(void)
     return ret;
 
 err:
-    heap_caps_free(_mdns_stack_buffer);
+    mdns_mem_task_free(_mdns_stack_buffer);
     return ret;
 }
 
@@ -5769,7 +5767,7 @@ void mdns_free(void)
     free_delegated_hostnames();
     _mdns_service_task_stop();
     // at this point, the service task is deleted, so we can destroy the stack size
-    heap_caps_free(_mdns_stack_buffer);
+    mdns_mem_task_free(_mdns_stack_buffer);
     for (i = 0; i < MDNS_MAX_INTERFACES; i++) {
         for (j = 0; j < MDNS_IP_PROTOCOL_MAX; j++) {
             mdns_pcb_deinit_local(i, j);

diff --git a/components/mdns/mdns_mem_caps.c b/components/mdns/mdns_mem_caps.c
@@ -4,31 +4,60 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 #include <string.h>
+#include "sdkconfig.h"
 #include "mdns_private.h"
 #include "mdns_mem_caps.h"
 #include "esp_heap_caps.h"
+#include "esp_log.h"
 
+#if CONFIG_MDNS_MEMORY_CUSTOM_IMPL
+#define ALLOW_WEAK __attribute__((weak))
+#else
+#define ALLOW_WEAK
+#endif
+
+#if CONFIG_MDNS_TASK_CREATE_FROM_SPIRAM
+#define MDNS_TASK_MEMORY_CAPS (MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT)
+#define MDNS_TASK_MEMORY_LOG "SPIRAM"
+#endif
+#if CONFIG_MDNS_TASK_CREATE_FROM_INTERNAL
+#define MDNS_TASK_MEMORY_CAPS (MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT)
+#define MDNS_TASK_MEMORY_LOG "internal RAM"
+#endif
+
+#if CONFIG_MDNS_MEMORY_ALLOC_SPIRAM
+#define MDNS_MEMORY_CAPS (MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT)
+#endif
+#if CONFIG_MDNS_MEMORY_ALLOC_INTERNAL
+#define MDNS_MEMORY_CAPS (MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT)
+#endif
+
+// Allocate memory from internal heap as default.
 #ifndef MDNS_MEMORY_CAPS
 #warning "No memory allocation method defined, using internal memory"
 #define MDNS_MEMORY_CAPS (MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT)
 #endif
+#ifndef MDNS_TASK_MEMORY_CAPS
+#define MDNS_TASK_MEMORY_CAPS (MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT)
+#define MDNS_TASK_MEMORY_LOG "internal RAM"
+#endif
 
-void *mdns_mem_malloc(size_t size)
+void ALLOW_WEAK *mdns_mem_malloc(size_t size)
 {
     return heap_caps_malloc(size, MDNS_MEMORY_CAPS);
 }
 
-void *mdns_mem_calloc(size_t num, size_t size)
+void ALLOW_WEAK *mdns_mem_calloc(size_t num, size_t size)
 {
     return heap_caps_calloc(num, size, MDNS_MEMORY_CAPS);
 }
 
-void mdns_mem_free(void *ptr)
+void ALLOW_WEAK mdns_mem_free(void *ptr)
 {
     heap_caps_free(ptr);
 }
 
-char *mdns_mem_strdup(const char *s)
+char ALLOW_WEAK *mdns_mem_strdup(const char *s)
 {
     if (!s) {
         return NULL;
@@ -41,7 +70,7 @@ char *mdns_mem_strdup(const char *s)
     return copy;
 }
 
-char *mdns_mem_strndup(const char *s, size_t n)
+char ALLOW_WEAK *mdns_mem_strndup(const char *s, size_t n)
 {
     if (!s) {
         return NULL;
@@ -54,3 +83,14 @@ char *mdns_mem_strndup(const char *s, size_t n)
     }
     return copy;
 }
+
+void ALLOW_WEAK *mdns_mem_task_malloc(size_t size)
+{
+    ESP_LOGI("mdns_mem", "mDNS task will be created from %s", MDNS_TASK_MEMORY_LOG);
+    return heap_caps_malloc(size, MDNS_TASK_MEMORY_CAPS);
+}
+
+void ALLOW_WEAK mdns_mem_task_free(void *ptr)
+{
+    heap_caps_free(ptr);
+}
diff --git a/components/mdns/private_include/mdns_mem_caps.h b/components/mdns/private_include/mdns_mem_caps.h
@@ -47,6 +47,19 @@ char *mdns_mem_strdup(const char *s);
  */
 char *mdns_mem_strndup(const char *s, size_t n);
 
+/**
+ * @brief Allocate memory for mDNS task.
+ * @param size Number of bytes to allocate.
+ * @return Pointer to allocated memory, or NULL on failure.
+ */
+void *mdns_mem_task_malloc(size_t size);
+
+/**
+ * @brief Free allocated memory for mDNS task.
+ * @param ptr Pointer to memory to free.
+ */
+void mdns_mem_task_free(void *ptr);
+
 #ifdef __cplusplus
 }
 #endif
diff --git a/components/mdns/private_include/mdns_private.h b/components/mdns/private_include/mdns_private.h
@@ -21,28 +21,6 @@
 #define _mdns_dbg_printf(...) printf(__VA_ARGS__)
 #endif
 
-#if CONFIG_MDNS_TASK_CREATE_FROM_SPIRAM
-#define MDNS_TASK_MEMORY_CAPS (MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT)
-#define MDNS_TASK_MEMORY_LOG "SPIRAM"
-#endif
-#if CONFIG_MDNS_TASK_CREATE_FROM_INTERNAL
-#define MDNS_TASK_MEMORY_CAPS (MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT)
-#define MDNS_TASK_MEMORY_LOG "internal RAM"
-#endif
-
-// Allocate memory from internal heap as default.
-#ifndef MDNS_TASK_MEMORY_CAPS
-#define MDNS_TASK_MEMORY_CAPS (MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT)
-#define MDNS_TASK_MEMORY_LOG "internal RAM"
-#endif
-
-#if CONFIG_MDNS_MEMORY_ALLOC_SPIRAM
-#define MDNS_MEMORY_CAPS (MALLOC_CAP_SPIRAM | MALLOC_CAP_8BIT)
-#endif
-#if CONFIG_MDNS_MEMORY_ALLOC_INTERNAL
-#define MDNS_MEMORY_CAPS (MALLOC_CAP_INTERNAL | MALLOC_CAP_8BIT)
-#endif
-
 /** Number of predefined interfaces */
 #ifndef CONFIG_MDNS_PREDEF_NETIF_STA
 #define CONFIG_MDNS_PREDEF_NETIF_STA 0

diff --git a/components/mdns/tests/test_afl_fuzz_host/esp32_mock.c b/components/mdns/tests/test_afl_fuzz_host/esp32_mock.c
@@ -126,16 +126,6 @@ uint32_t esp_log_timestamp(void)
     return 0;
 }
 
-void *heap_caps_malloc(size_t size, uint32_t caps)
-{
-    return malloc(size);
-}
-
-void heap_caps_free(void *ptr)
-{
-    free(ptr);
-}
-
 void *mdns_mem_malloc(size_t size)
 {
     return malloc(size);
@@ -160,3 +150,13 @@ char *mdns_mem_strndup(const char *s, size_t n)
 {
     return strndup(s, n);
 }
+
+void *mdns_mem_task_malloc(size_t size)
+{
+    return malloc(size);
+}
+
+void mdns_mem_task_free(void *ptr)
+{
+    free(ptr);
+}
diff --git a/components/mdns/tests/test_afl_fuzz_host/esp32_mock.h b/components/mdns/tests/test_afl_fuzz_host/esp32_mock.h
@@ -138,7 +138,5 @@ esp_err_t esp_event_handler_unregister(const char *event_base, int32_t event_id,
 TaskHandle_t xTaskGetCurrentTaskHandle(void);
 void xTaskNotifyGive(TaskHandle_t task);
 BaseType_t xTaskNotifyWait(uint32_t bits_entry_clear, uint32_t bits_exit_clear, uint32_t *value, TickType_t wait_time);
-void *heap_caps_malloc(size_t size, uint32_t caps);
-void heap_caps_free(void *ptr);
 
 #endif //_ESP32_COMPAT_H_