From 0554b06b7e968901dd54c45ecb4ceb03fb4e7f5d Mon Sep 17 00:00:00 2001
From: Joakim Plate <elupus@ecce.se>
Date: Thu, 2 Feb 2023 05:18:58 +0100
Subject: [PATCH] Allow mdns services to be exposed by config (#4202)

---
 esphome/components/mdns/__init__.py        | 54 +++++++++++++++++++++-
 esphome/components/mdns/mdns_component.cpp |  2 +
 esphome/components/mdns/mdns_component.h   |  3 ++
 3 files changed, 58 insertions(+), 1 deletion(-)

diff --git a/esphome/components/mdns/__init__.py b/esphome/components/mdns/__init__.py
index 154e6b099f..e27786a98b 100644
--- a/esphome/components/mdns/__init__.py
+++ b/esphome/components/mdns/__init__.py
@@ -1,4 +1,10 @@
-from esphome.const import CONF_ID
+from esphome.const import (
+    CONF_ID,
+    CONF_PORT,
+    CONF_PROTOCOL,
+    CONF_SERVICES,
+    CONF_SERVICE,
+)
 import esphome.codegen as cg
 import esphome.config_validation as cv
 from esphome.core import CORE, coroutine_with_priority
@@ -8,6 +14,8 @@ DEPENDENCIES = ["network"]
 
 mdns_ns = cg.esphome_ns.namespace("mdns")
 MDNSComponent = mdns_ns.class_("MDNSComponent", cg.Component)
+MDNSTXTRecord = mdns_ns.struct("MDNSTXTRecord")
+MDNSService = mdns_ns.struct("MDNSService")
 
 
 def _remove_id_if_disabled(value):
@@ -17,18 +25,50 @@ def _remove_id_if_disabled(value):
     return value
 
 
+CONF_TXT = "txt"
+
+SERVICE_SCHEMA = cv.Schema(
+    {
+        cv.Required(CONF_SERVICE): cv.string,
+        cv.Required(CONF_PROTOCOL): cv.string,
+        cv.Optional(CONF_PORT, default=0): cv.Any(0, cv.port),
+        cv.Optional(CONF_TXT, default={}): {cv.string: cv.string},
+    }
+)
+
 CONF_DISABLED = "disabled"
 CONFIG_SCHEMA = cv.All(
     cv.Schema(
         {
             cv.GenerateID(): cv.declare_id(MDNSComponent),
             cv.Optional(CONF_DISABLED, default=False): cv.boolean,
+            cv.Optional(CONF_SERVICES, default=[]): cv.ensure_list(SERVICE_SCHEMA),
         }
     ),
     _remove_id_if_disabled,
 )
 
 
+def mdns_txt_record(key: str, value: str):
+    return cg.StructInitializer(
+        MDNSTXTRecord,
+        ("key", key),
+        ("value", value),
+    )
+
+
+def mdns_service(
+    service: str, proto: str, port: int, txt_records: list[dict[str, str]]
+):
+    return cg.StructInitializer(
+        MDNSService,
+        ("service_type", service),
+        ("proto", proto),
+        ("port", port),
+        ("txt_records", txt_records),
+    )
+
+
 @coroutine_with_priority(55.0)
 async def to_code(config):
     if CORE.using_arduino:
@@ -46,3 +86,15 @@ async def to_code(config):
 
     var = cg.new_Pvariable(config[CONF_ID])
     await cg.register_component(var, config)
+
+    for service in config[CONF_SERVICES]:
+        txt = [
+            mdns_txt_record(txt_key, txt_value)
+            for txt_key, txt_value in service[CONF_TXT].items()
+        ]
+
+        exp = mdns_service(
+            service[CONF_SERVICE], service[CONF_PROTOCOL], service[CONF_PORT], txt
+        )
+
+        cg.add(var.add_extra_service(exp))
diff --git a/esphome/components/mdns/mdns_component.cpp b/esphome/components/mdns/mdns_component.cpp
index 656e0382bf..cdb9aa8e74 100644
--- a/esphome/components/mdns/mdns_component.cpp
+++ b/esphome/components/mdns/mdns_component.cpp
@@ -90,6 +90,8 @@ void MDNSComponent::compile_records_() {
   }
 #endif
 
+  this->services_.insert(this->services_.end(), this->services_extra_.begin(), this->services_extra_.end());
+
   if (this->services_.empty()) {
     // Publish "http" service if not using native API
     // This is just to have *some* mDNS service so that .local resolution works
diff --git a/esphome/components/mdns/mdns_component.h b/esphome/components/mdns/mdns_component.h
index dbf7482f1f..5f9179b799 100644
--- a/esphome/components/mdns/mdns_component.h
+++ b/esphome/components/mdns/mdns_component.h
@@ -33,7 +33,10 @@ class MDNSComponent : public Component {
 #endif
   float get_setup_priority() const override { return setup_priority::AFTER_WIFI; }
 
+  void add_extra_service(MDNSService service) { services_extra_.push_back(std::move(service)); }
+
  protected:
+  std::vector<MDNSService> services_extra_{};
   std::vector<MDNSService> services_{};
   std::string hostname_;
   void compile_records_();