From 14aa27f5e2df02dbe251f72ef0869e87405c15a0 Mon Sep 17 00:00:00 2001
From: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
Date: Wed, 18 Oct 2023 14:26:47 +1300
Subject: [PATCH] esp32_improv advertise capabilities and state in ble service
 data (#5553)

---
 .../components/esp32_ble/ble_advertising.cpp  | 21 +++++++++++++------
 .../components/esp32_ble/ble_advertising.h    |  1 +
 .../esp32_improv/esp32_improv_component.cpp   | 19 +++++++++++++++++
 3 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/esphome/components/esp32_ble/ble_advertising.cpp b/esphome/components/esp32_ble/ble_advertising.cpp
index 072bb38c07..59d2398829 100644
--- a/esphome/components/esp32_ble/ble_advertising.cpp
+++ b/esphome/components/esp32_ble/ble_advertising.cpp
@@ -2,9 +2,9 @@
 
 #ifdef USE_ESP32
 
-#include "ble_uuid.h"
-#include <cstring>
 #include <cstdio>
+#include <cstring>
+#include "ble_uuid.h"
 #include "esphome/core/log.h"
 
 namespace esphome {
@@ -16,8 +16,8 @@ BLEAdvertising::BLEAdvertising() {
   this->advertising_data_.set_scan_rsp = false;
   this->advertising_data_.include_name = true;
   this->advertising_data_.include_txpower = true;
-  this->advertising_data_.min_interval = 0x20;
-  this->advertising_data_.max_interval = 0x40;
+  this->advertising_data_.min_interval = 0;
+  this->advertising_data_.max_interval = 0;
   this->advertising_data_.appearance = 0x00;
   this->advertising_data_.manufacturer_len = 0;
   this->advertising_data_.p_manufacturer_data = nullptr;
@@ -42,6 +42,17 @@ void BLEAdvertising::remove_service_uuid(ESPBTUUID uuid) {
                                  this->advertising_uuids_.end());
 }
 
+void BLEAdvertising::set_service_data(const std::vector<uint8_t> &data) {
+  delete[] this->advertising_data_.p_service_data;
+  this->advertising_data_.p_service_data = nullptr;
+  this->advertising_data_.service_data_len = data.size();
+  if (!data.empty()) {
+    // NOLINTNEXTLINE(cppcoreguidelines-owning-memory)
+    this->advertising_data_.p_service_data = new uint8_t[data.size()];
+    memcpy(this->advertising_data_.p_service_data, data.data(), data.size());
+  }
+}
+
 void BLEAdvertising::set_manufacturer_data(const std::vector<uint8_t> &data) {
   delete[] this->advertising_data_.p_manufacturer_data;
   this->advertising_data_.p_manufacturer_data = nullptr;
@@ -85,8 +96,6 @@ void BLEAdvertising::start() {
     this->scan_response_data_.set_scan_rsp = true;
     this->scan_response_data_.include_name = true;
     this->scan_response_data_.include_txpower = true;
-    this->scan_response_data_.min_interval = 0;
-    this->scan_response_data_.max_interval = 0;
     this->scan_response_data_.manufacturer_len = 0;
     this->scan_response_data_.appearance = 0;
     this->scan_response_data_.flag = 0;
diff --git a/esphome/components/esp32_ble/ble_advertising.h b/esphome/components/esp32_ble/ble_advertising.h
index 9e4e2b7701..16a7dd1d8e 100644
--- a/esphome/components/esp32_ble/ble_advertising.h
+++ b/esphome/components/esp32_ble/ble_advertising.h
@@ -21,6 +21,7 @@ class BLEAdvertising {
   void set_scan_response(bool scan_response) { this->scan_response_ = scan_response; }
   void set_min_preferred_interval(uint16_t interval) { this->advertising_data_.min_interval = interval; }
   void set_manufacturer_data(const std::vector<uint8_t> &data);
+  void set_service_data(const std::vector<uint8_t> &data);
 
   void start();
   void stop();
diff --git a/esphome/components/esp32_improv/esp32_improv_component.cpp b/esphome/components/esp32_improv/esp32_improv_component.cpp
index 8a901a79e5..19340c3dd8 100644
--- a/esphome/components/esp32_improv/esp32_improv_component.cpp
+++ b/esphome/components/esp32_improv/esp32_improv_component.cpp
@@ -189,6 +189,25 @@ void ESP32ImprovComponent::set_state_(improv::State state) {
     if (state != improv::STATE_STOPPED)
       this->status_->notify();
   }
+  std::vector<uint8_t> service_data(8, 0);
+  service_data[0] = 0x77;  // PR
+  service_data[1] = 0x46;  // IM
+  service_data[2] = static_cast<uint8_t>(state);
+
+  uint8_t capabilities = 0x00;
+#ifdef USE_OUTPUT
+  if (this->status_indicator_ != nullptr)
+    capabilities |= improv::CAPABILITY_IDENTIFY;
+#endif
+
+  service_data[3] = capabilities;
+  service_data[4] = 0x00;  // Reserved
+  service_data[5] = 0x00;  // Reserved
+  service_data[6] = 0x00;  // Reserved
+  service_data[7] = 0x00;  // Reserved
+
+  esp32_ble::global_ble->get_advertising()->set_service_data(service_data);
+  esp32_ble::global_ble->get_advertising()->start();
 }
 
 void ESP32ImprovComponent::set_error_(improv::Error error) {