From 10c5a19503bec70af4825bfc466dc921c7d3ebf5 Mon Sep 17 00:00:00 2001
From: "J. Nick Koston" <nick@home-assistant.io>
Date: Fri, 26 Sep 2025 22:01:44 -0500
Subject: [PATCH] optimize

---
 .../components/esphome/ota/ota_esphome.cpp    | 69 ++++++++-----------
 esphome/components/esphome/ota/ota_esphome.h  |  3 +
 2 files changed, 32 insertions(+), 40 deletions(-)

diff --git a/esphome/components/esphome/ota/ota_esphome.cpp b/esphome/components/esphome/ota/ota_esphome.cpp
index 8edcba99c9..e231cf336e 100644
--- a/esphome/components/esphome/ota/ota_esphome.cpp
+++ b/esphome/components/esphome/ota/ota_esphome.cpp
@@ -156,33 +156,9 @@ void ESPHomeOTAComponent::handle_handshake_() {
   while (true) {
     switch (this->ota_state_) {
       case OTAState::MAGIC_READ: {
-        // Try to read remaining magic bytes
-        if (this->handshake_buf_pos_ < 5) {
-          // Read as many bytes as available
-          uint8_t bytes_to_read = 5 - this->handshake_buf_pos_;
-          ssize_t read = this->client_->read(this->handshake_buf_ + this->handshake_buf_pos_, bytes_to_read);
-
-          if (read == -1 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
-            return;  // No data yet, try again next loop
-          }
-
-          if (read <= 0) {
-            // Error or connection closed
-            if (read == -1) {
-              this->log_socket_error_(LOG_STR("reading magic bytes"));
-            } else {
-              this->log_remote_closed_(LOG_STR("handshake"));
-            }
-            this->cleanup_connection_();
-            return;
-          }
-
-          this->handshake_buf_pos_ += read;
-        }
-
-        // Check if we have all 5 magic bytes
-        if (this->handshake_buf_pos_ != 5) {
-          break;
+        // Try to read remaining magic bytes (5 total)
+        if (!this->try_read_(5, LOG_STR("reading magic bytes"), LOG_STR("handshake"))) {
+          return;
         }
 
         // Validate magic bytes
@@ -237,19 +213,7 @@ void ESPHomeOTAComponent::handle_handshake_() {
 
       case OTAState::FEATURE_READ: {
         // Read features - 1 byte
-        ssize_t read = this->client_->read(this->handshake_buf_, 1);
-
-        if (read == -1 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
-          return;  // No data yet, try again next loop
-        }
-
-        if (read <= 0) {
-          if (read == -1) {
-            this->log_socket_error_(LOG_STR("reading features"));
-          } else {
-            this->log_remote_closed_(LOG_STR("feature read"));
-          }
-          this->cleanup_connection_();
+        if (!this->try_read_(1, LOG_STR("reading features"), LOG_STR("feature read"))) {
           return;
         }
 
@@ -594,6 +558,31 @@ void ESPHomeOTAComponent::log_remote_closed_(const LogString *during) {
   ESP_LOGW(TAG, "Remote closed during %s", LOG_STR_ARG(during));
 }
 
+bool ESPHomeOTAComponent::try_read_(size_t to_read, const LogString *error_desc, const LogString *close_desc) {
+  // Read bytes into handshake buffer, starting at handshake_buf_pos_
+  size_t bytes_to_read = to_read - this->handshake_buf_pos_;
+  ssize_t read = this->client_->read(this->handshake_buf_ + this->handshake_buf_pos_, bytes_to_read);
+
+  if (read == -1 && (errno == EAGAIN || errno == EWOULDBLOCK)) {
+    return false;  // No data yet, try again next loop
+  }
+
+  if (read <= 0) {
+    // Error or connection closed
+    if (read == -1) {
+      this->log_socket_error_(error_desc);
+    } else {
+      this->log_remote_closed_(close_desc);
+    }
+    this->cleanup_connection_();
+    return false;
+  }
+
+  this->handshake_buf_pos_ += read;
+  // Return true only if we have all the requested bytes
+  return this->handshake_buf_pos_ >= to_read;
+}
+
 void ESPHomeOTAComponent::cleanup_connection_() {
   this->client_->close();
   this->client_ = nullptr;
diff --git a/esphome/components/esphome/ota/ota_esphome.h b/esphome/components/esphome/ota/ota_esphome.h
index f50444c6ce..c73fe7e732 100644
--- a/esphome/components/esphome/ota/ota_esphome.h
+++ b/esphome/components/esphome/ota/ota_esphome.h
@@ -46,6 +46,9 @@ class ESPHomeOTAComponent : public ota::OTAComponent {
 #endif  // USE_OTA_PASSWORD
   bool readall_(uint8_t *buf, size_t len);
   bool writeall_(const uint8_t *buf, size_t len);
+
+  bool try_read_(size_t to_read, const LogString *error_desc, const LogString *close_desc);
+
   void log_socket_error_(const LogString *msg);
   void log_read_error_(const LogString *what);
   void log_start_(const LogString *phase);