From e07f9a5b9c662ab7299e314ee16a98eba80fe159 Mon Sep 17 00:00:00 2001
From: Jesse Hills <3060199+jesserockz@users.noreply.github.com>
Date: Mon, 19 Jun 2023 13:40:46 +1200
Subject: [PATCH] Convert wled light effect to use socket instead of WiFiUDP

---
 esphome/components/wled/__init__.py           |  3 +-
 esphome/components/wled/wled_light_effect.cpp | 70 +++++++++++--------
 esphome/components/wled/wled_light_effect.h   | 11 +--
 3 files changed, 45 insertions(+), 39 deletions(-)

diff --git a/esphome/components/wled/__init__.py b/esphome/components/wled/__init__.py
index 2795529203..18bc9c8e79 100644
--- a/esphome/components/wled/__init__.py
+++ b/esphome/components/wled/__init__.py
@@ -7,7 +7,8 @@ from esphome.const import CONF_NAME, CONF_PORT
 wled_ns = cg.esphome_ns.namespace("wled")
 WLEDLightEffect = wled_ns.class_("WLEDLightEffect", AddressableLightEffect)
 
-CONFIG_SCHEMA = cv.All(cv.Schema({}), cv.only_with_arduino)
+AUTO_LOAD = ["socket"]
+CONFIG_SCHEMA = cv.Schema({})
 
 
 @register_addressable_effect(
diff --git a/esphome/components/wled/wled_light_effect.cpp b/esphome/components/wled/wled_light_effect.cpp
index 8c68bca6e3..8b9457a421 100644
--- a/esphome/components/wled/wled_light_effect.cpp
+++ b/esphome/components/wled/wled_light_effect.cpp
@@ -1,17 +1,7 @@
-#ifdef USE_ARDUINO
-
 #include "wled_light_effect.h"
-#include "esphome/core/log.h"
+
 #include "esphome/core/helpers.h"
-
-#ifdef USE_ESP32
-#include <WiFi.h>
-#endif
-
-#ifdef USE_ESP8266
-#include <ESP8266WiFi.h>
-#include <WiFiUdp.h>
-#endif
+#include "esphome/core/log.h"
 
 namespace esphome {
 namespace wled {
@@ -35,9 +25,9 @@ void WLEDLightEffect::start() {
 void WLEDLightEffect::stop() {
   AddressableLightEffect::stop();
 
-  if (udp_) {
-    udp_->stop();
-    udp_.reset();
+  if (this->socket_) {
+    this->socket_->close();
+    this->socket_.reset();
   }
 }
 
@@ -50,27 +40,49 @@ void WLEDLightEffect::blank_all_leds_(light::AddressableLight &it) {
 
 void WLEDLightEffect::apply(light::AddressableLight &it, const Color &current_color) {
   // Init UDP lazily
-  if (!udp_) {
-    udp_ = make_unique<WiFiUDP>();
+  if (this->socket_ == nullptr) {
+    this->socket_ = socket::socket_ip(SOCK_DGRAM, IPPROTO_IP);
 
-    if (!udp_->begin(port_)) {
-      ESP_LOGW(TAG, "Cannot bind WLEDLightEffect to %d.", port_);
+    int enable = 1;
+    int err = this->socket_->setsockopt(SOL_SOCKET, SO_REUSEADDR, &enable, sizeof(int));
+    if (err != 0) {
+      ESP_LOGW(TAG, "Socket unable to set reuseaddr: errno %d", err);
+      // we can still continue
+    }
+    err = this->socket_->setblocking(false);
+    if (err != 0) {
+      ESP_LOGW(TAG, "Socket unable to set nonblocking mode: errno %d", err);
+      return;
+    }
+
+    struct sockaddr_storage server;
+
+    socklen_t sl = socket::set_sockaddr_any((struct sockaddr *) &server, sizeof(server), this->port_);
+    if (sl == 0) {
+      ESP_LOGW(TAG, "Socket unable to set sockaddr: errno %d", errno);
+      return;
+    }
+    server.ss_family = AF_INET;
+
+    err = this->socket_->bind((struct sockaddr *) &server, sizeof(server));
+    if (err != 0) {
+      ESP_LOGW(TAG, "Socket unable to bind: errno %d", errno);
       return;
     }
   }
 
   std::vector<uint8_t> payload;
-  while (uint16_t packet_size = udp_->parsePacket()) {
-    payload.resize(packet_size);
+  uint8_t buf[1460];
 
-    if (!udp_->read(&payload[0], payload.size())) {
-      continue;
-    }
+  ssize_t len = this->socket_->read(buf, sizeof(buf));
+  if (len == -1) {
+    return;
+  }
+  payload.resize(len);
+  memmove(&payload[0], buf, len);
 
-    if (!this->parse_frame_(it, &payload[0], payload.size())) {
-      ESP_LOGD(TAG, "Frame: Invalid (size=%zu, first=0x%02X).", payload.size(), payload[0]);
-      continue;
-    }
+  if (!this->parse_frame_(it, &payload[0], payload.size())) {
+    ESP_LOGD(TAG, "Frame: Invalid (size=%zu, first=0x%02X).", payload.size(), payload[0]);
   }
 
   // FIXME: Use roll-over safe arithmetic
@@ -248,5 +260,3 @@ bool WLEDLightEffect::parse_dnrgb_frame_(light::AddressableLight &it, const uint
 
 }  // namespace wled
 }  // namespace esphome
-
-#endif  // USE_ARDUINO
diff --git a/esphome/components/wled/wled_light_effect.h b/esphome/components/wled/wled_light_effect.h
index 8f239276d7..aec3c92c69 100644
--- a/esphome/components/wled/wled_light_effect.h
+++ b/esphome/components/wled/wled_light_effect.h
@@ -1,15 +1,12 @@
 #pragma once
 
-#ifdef USE_ARDUINO
-
-#include "esphome/core/component.h"
 #include "esphome/components/light/addressable_light_effect.h"
+#include "esphome/components/socket/socket.h"
+#include "esphome/core/component.h"
 
 #include <vector>
 #include <memory>
 
-class UDP;
-
 namespace esphome {
 namespace wled {
 
@@ -32,12 +29,10 @@ class WLEDLightEffect : public light::AddressableLightEffect {
   bool parse_dnrgb_frame_(light::AddressableLight &it, const uint8_t *payload, uint16_t size);
 
   uint16_t port_{0};
-  std::unique_ptr<UDP> udp_;
+  std::unique_ptr<socket::Socket> socket_;
   uint32_t blank_at_{0};
   uint32_t dropped_{0};
 };
 
 }  // namespace wled
 }  // namespace esphome
-
-#endif  // USE_ARDUINO