esphome 2025.5.1__py3-none-any.whl → 2025.6.0b1__py3-none-any.whl

esphome/components/api/api_frame_helper.cpp

@@ -1,26 +1,19 @@
 #include "api_frame_helper.h"
 #ifdef USE_API
-#include "esphome/core/log.h"
+#include "esphome/core/application.h"
 #include "esphome/core/hal.h"
 #include "esphome/core/helpers.h"
-#include "esphome/core/application.h"
+#include "esphome/core/log.h"
 #include "proto.h"
 #include "api_pb2_size.h"
 #include <cstring>
+#include <cinttypes>
 
 namespace esphome {
 namespace api {
 
 static const char *const TAG = "api.socket";
 
-/// Is the given return value (from write syscalls) a wouldblock error?
-bool is_would_block(ssize_t ret) {
-  if (ret == -1) {
-    return errno == EWOULDBLOCK || errno == EAGAIN;
-  }
-  return ret == 0;
-}
-
 const char *api_error_to_str(APIError err) {
   // not using switch to ensure compiler doesn't try to build a big table out of it
   if (err == APIError::OK) {
@@ -73,92 +66,154 @@ const char *api_error_to_str(APIError err) {
   return "UNKNOWN";
 }
 
-// Common implementation for writing raw data to socket
-template<typename StateEnum>
-APIError APIFrameHelper::write_raw_(const struct iovec *iov, int iovcnt, socket::Socket *socket,
-                                    std::vector<uint8_t> &tx_buf, const std::string &info, StateEnum &state,
-                                    StateEnum failed_state) {
-  // This method writes data to socket or buffers it
+// Helper method to buffer data from IOVs
+void APIFrameHelper::buffer_data_from_iov_(const struct iovec *iov, int iovcnt, uint16_t total_write_len) {
+  SendBuffer buffer;
+  buffer.data.reserve(total_write_len);
+  for (int i = 0; i < iovcnt; i++) {
+    const uint8_t *data = reinterpret_cast<uint8_t *>(iov[i].iov_base);
+    buffer.data.insert(buffer.data.end(), data, data + iov[i].iov_len);
+  }
+  this->tx_buf_.push_back(std::move(buffer));
+}
+
+// This method writes data to socket or buffers it
+APIError APIFrameHelper::write_raw_(const struct iovec *iov, int iovcnt) {
   // Returns APIError::OK if successful (or would block, but data has been buffered)
-  // Returns APIError::SOCKET_WRITE_FAILED if socket write failed, and sets state to failed_state
+  // Returns APIError::SOCKET_WRITE_FAILED if socket write failed, and sets state to FAILED
 
   if (iovcnt == 0)
     return APIError::OK;  // Nothing to do, success
 
-  size_t total_write_len = 0;
+  uint16_t total_write_len = 0;
   for (int i = 0; i < iovcnt; i++) {
 #ifdef HELPER_LOG_PACKETS
     ESP_LOGVV(TAG, "Sending raw: %s",
               format_hex_pretty(reinterpret_cast<uint8_t *>(iov[i].iov_base), iov[i].iov_len).c_str());
 #endif
-    total_write_len += iov[i].iov_len;
+    total_write_len += static_cast<uint16_t>(iov[i].iov_len);
   }
 
-  if (!tx_buf.empty()) {
-    // try to empty tx_buf first
-    while (!tx_buf.empty()) {
-      ssize_t sent = socket->write(tx_buf.data(), tx_buf.size());
-      if (is_would_block(sent)) {
-        break;
-      } else if (sent == -1) {
-        ESP_LOGVV(TAG, "%s: Socket write failed with errno %d", info.c_str(), errno);
-        state = failed_state;
-        return APIError::SOCKET_WRITE_FAILED;  // Socket write failed
-      }
-      // TODO: inefficient if multiple packets in txbuf
-      // replace with deque of buffers
-      tx_buf.erase(tx_buf.begin(), tx_buf.begin() + sent);
+  // Try to send any existing buffered data first if there is any
+  if (!this->tx_buf_.empty()) {
+    APIError send_result = try_send_tx_buf_();
+    // If real error occurred (not just WOULD_BLOCK), return it
+    if (send_result != APIError::OK && send_result != APIError::WOULD_BLOCK) {
+      return send_result;
     }
-  }
 
-  if (!tx_buf.empty()) {
-    // tx buf not empty, can't write now because then stream would be inconsistent
-    // Reserve space upfront to avoid multiple reallocations
-    tx_buf.reserve(tx_buf.size() + total_write_len);
-    for (int i = 0; i < iovcnt; i++) {
-      tx_buf.insert(tx_buf.end(), reinterpret_cast<uint8_t *>(iov[i].iov_base),
-                    reinterpret_cast<uint8_t *>(iov[i].iov_base) + iov[i].iov_len);
+    // If there is still data in the buffer, we can't send, buffer
+    // the new data and return
+    if (!this->tx_buf_.empty()) {
+      this->buffer_data_from_iov_(iov, iovcnt, total_write_len);
+      return APIError::OK;  // Success, data buffered
     }
-    return APIError::OK;  // Success, data buffered
   }
 
-  ssize_t sent = socket->writev(iov, iovcnt);
-  if (is_would_block(sent)) {
-    // operation would block, add buffer to tx_buf
-    // Reserve space upfront to avoid multiple reallocations
-    tx_buf.reserve(tx_buf.size() + total_write_len);
-    for (int i = 0; i < iovcnt; i++) {
-      tx_buf.insert(tx_buf.end(), reinterpret_cast<uint8_t *>(iov[i].iov_base),
-                    reinterpret_cast<uint8_t *>(iov[i].iov_base) + iov[i].iov_len);
+  // Try to send directly if no buffered data
+  ssize_t sent = this->socket_->writev(iov, iovcnt);
+
+  if (sent == -1) {
+    if (errno == EWOULDBLOCK || errno == EAGAIN) {
+      // Socket would block, buffer the data
+      this->buffer_data_from_iov_(iov, iovcnt, total_write_len);
+      return APIError::OK;  // Success, data buffered
     }
-    return APIError::OK;  // Success, data buffered
-  } else if (sent == -1) {
-    // an error occurred
-    ESP_LOGVV(TAG, "%s: Socket write failed with errno %d", info.c_str(), errno);
-    state = failed_state;
+    // Socket error
+    ESP_LOGVV(TAG, "%s: Socket write failed with errno %d", this->info_.c_str(), errno);
+    this->state_ = State::FAILED;
     return APIError::SOCKET_WRITE_FAILED;  // Socket write failed
-  } else if ((size_t) sent != total_write_len) {
-    // partially sent, add end to tx_buf
-    size_t remaining = total_write_len - sent;
-    // Reserve space upfront to avoid multiple reallocations
-    tx_buf.reserve(tx_buf.size() + remaining);
+  } else if (static_cast<uint16_t>(sent) < total_write_len) {
+    // Partially sent, buffer the remaining data
+    SendBuffer buffer;
+    uint16_t to_consume = static_cast<uint16_t>(sent);
+    uint16_t remaining = total_write_len - static_cast<uint16_t>(sent);
+
+    buffer.data.reserve(remaining);
 
-    size_t to_consume = sent;
     for (int i = 0; i < iovcnt; i++) {
       if (to_consume >= iov[i].iov_len) {
-        to_consume -= iov[i].iov_len;
+        // This segment was fully sent
+        to_consume -= static_cast<uint16_t>(iov[i].iov_len);
       } else {
-        tx_buf.insert(tx_buf.end(), reinterpret_cast<uint8_t *>(iov[i].iov_base) + to_consume,
-                      reinterpret_cast<uint8_t *>(iov[i].iov_base) + iov[i].iov_len);
+        // This segment was partially sent or not sent at all
+        const uint8_t *data = reinterpret_cast<uint8_t *>(iov[i].iov_base) + to_consume;
+        uint16_t len = static_cast<uint16_t>(iov[i].iov_len) - to_consume;
+        buffer.data.insert(buffer.data.end(), data, data + len);
         to_consume = 0;
       }
     }
-    return APIError::OK;  // Success, data buffered
+
+    this->tx_buf_.push_back(std::move(buffer));
   }
-  return APIError::OK;  // Success, all data sent
+
+  return APIError::OK;  // Success, all data sent or buffered
 }
 
-#define HELPER_LOG(msg, ...) ESP_LOGVV(TAG, "%s: " msg, info_.c_str(), ##__VA_ARGS__)
+// Common implementation for trying to send buffered data
+// IMPORTANT: Caller MUST ensure tx_buf_ is not empty before calling this method
+APIError APIFrameHelper::try_send_tx_buf_() {
+  // Try to send from tx_buf - we assume it's not empty as it's the caller's responsibility to check
+  bool tx_buf_empty = false;
+  while (!tx_buf_empty) {
+    // Get the first buffer in the queue
+    SendBuffer &front_buffer = this->tx_buf_.front();
+
+    // Try to send the remaining data in this buffer
+    ssize_t sent = this->socket_->write(front_buffer.current_data(), front_buffer.remaining());
+
+    if (sent == -1) {
+      if (errno != EWOULDBLOCK && errno != EAGAIN) {
+        // Real socket error (not just would block)
+        ESP_LOGVV(TAG, "%s: Socket write failed with errno %d", this->info_.c_str(), errno);
+        this->state_ = State::FAILED;
+        return APIError::SOCKET_WRITE_FAILED;  // Socket write failed
+      }
+      // Socket would block, we'll try again later
+      return APIError::WOULD_BLOCK;
+    } else if (sent == 0) {
+      // Nothing sent but not an error
+      return APIError::WOULD_BLOCK;
+    } else if (static_cast<uint16_t>(sent) < front_buffer.remaining()) {
+      // Partially sent, update offset
+      // Cast to ensure no overflow issues with uint16_t
+      front_buffer.offset += static_cast<uint16_t>(sent);
+      return APIError::WOULD_BLOCK;  // Stop processing more buffers if we couldn't send a complete buffer
+    } else {
+      // Buffer completely sent, remove it from the queue
+      this->tx_buf_.pop_front();
+      // Update empty status for the loop condition
+      tx_buf_empty = this->tx_buf_.empty();
+      // Continue loop to try sending the next buffer
+    }
+  }
+
+  return APIError::OK;  // All buffers sent successfully
+}
+
+APIError APIFrameHelper::init_common_() {
+  if (state_ != State::INITIALIZE || this->socket_ == nullptr) {
+    ESP_LOGVV(TAG, "%s: Bad state for init %d", this->info_.c_str(), (int) state_);
+    return APIError::BAD_STATE;
+  }
+  int err = this->socket_->setblocking(false);
+  if (err != 0) {
+    state_ = State::FAILED;
+    ESP_LOGVV(TAG, "%s: Setting nonblocking failed with errno %d", this->info_.c_str(), errno);
+    return APIError::TCP_NONBLOCKING_FAILED;
+  }
+
+  int enable = 1;
+  err = this->socket_->setsockopt(IPPROTO_TCP, TCP_NODELAY, &enable, sizeof(int));
+  if (err != 0) {
+    state_ = State::FAILED;
+    ESP_LOGVV(TAG, "%s: Setting nodelay failed with errno %d", this->info_.c_str(), errno);
+    return APIError::TCP_NODELAY_FAILED;
+  }
+  return APIError::OK;
+}
+
+#define HELPER_LOG(msg, ...) ESP_LOGVV(TAG, "%s: " msg, this->info_.c_str(), ##__VA_ARGS__)
 // uncomment to log raw packets
 //#define HELPER_LOG_PACKETS
 
@@ -206,23 +261,9 @@ std::string noise_err_to_str(int err) {
 
 /// Initialize the frame helper, returns OK if successful.
 APIError APINoiseFrameHelper::init() {
-  if (state_ != State::INITIALIZE || socket_ == nullptr) {
-    HELPER_LOG("Bad state for init %d", (int) state_);
-    return APIError::BAD_STATE;
-  }
-  int err = socket_->setblocking(false);
-  if (err != 0) {
-    state_ = State::FAILED;
-    HELPER_LOG("Setting nonblocking failed with errno %d", errno);
-    return APIError::TCP_NONBLOCKING_FAILED;
-  }
-
-  int enable = 1;
-  err = socket_->setsockopt(IPPROTO_TCP, TCP_NODELAY, &enable, sizeof(int));
-  if (err != 0) {
-    state_ = State::FAILED;
-    HELPER_LOG("Setting nodelay failed with errno %d", errno);
-    return APIError::TCP_NODELAY_FAILED;
+  APIError err = init_common_();
+  if (err != APIError::OK) {
+    return err;
   }
 
   // init prologue
@@ -234,17 +275,16 @@ APIError APINoiseFrameHelper::init() {
 /// Run through handshake messages (if in that phase)
 APIError APINoiseFrameHelper::loop() {
   APIError err = state_action_();
-  if (err == APIError::WOULD_BLOCK)
-    return APIError::OK;
-  if (err != APIError::OK)
+  if (err != APIError::OK && err != APIError::WOULD_BLOCK) {
     return err;
-  if (!tx_buf_.empty()) {
+  }
+  if (!this->tx_buf_.empty()) {
     err = try_send_tx_buf_();
-    if (err != APIError::OK) {
+    if (err != APIError::OK && err != APIError::WOULD_BLOCK) {
       return err;
     }
   }
-  return APIError::OK;
+  return APIError::OK;  // Convert WOULD_BLOCK to OK to avoid connection termination
 }
 
 /** Read a packet into the rx_buf_. If successful, stores frame data in the frame parameter
@@ -270,8 +310,8 @@ APIError APINoiseFrameHelper::try_read_frame_(ParsedFrame *frame) {
   // read header
   if (rx_header_buf_len_ < 3) {
     // no header information yet
-    size_t to_read = 3 - rx_header_buf_len_;
-    ssize_t received = socket_->read(&rx_header_buf_[rx_header_buf_len_], to_read);
+    uint8_t to_read = 3 - rx_header_buf_len_;
+    ssize_t received = this->socket_->read(&rx_header_buf_[rx_header_buf_len_], to_read);
     if (received == -1) {
       if (errno == EWOULDBLOCK || errno == EAGAIN) {
         return APIError::WOULD_BLOCK;
@@ -284,8 +324,8 @@ APIError APINoiseFrameHelper::try_read_frame_(ParsedFrame *frame) {
       HELPER_LOG("Connection closed");
       return APIError::CONNECTION_CLOSED;
     }
-    rx_header_buf_len_ += received;
-    if ((size_t) received != to_read) {
+    rx_header_buf_len_ += static_cast<uint8_t>(received);
+    if (static_cast<uint8_t>(received) != to_read) {
       // not a full read
       return APIError::WOULD_BLOCK;
     }
@@ -317,8 +357,8 @@ APIError APINoiseFrameHelper::try_read_frame_(ParsedFrame *frame) {
 
   if (rx_buf_len_ < msg_size) {
     // more data to read
-    size_t to_read = msg_size - rx_buf_len_;
-    ssize_t received = socket_->read(&rx_buf_[rx_buf_len_], to_read);
+    uint16_t to_read = msg_size - rx_buf_len_;
+    ssize_t received = this->socket_->read(&rx_buf_[rx_buf_len_], to_read);
     if (received == -1) {
       if (errno == EWOULDBLOCK || errno == EAGAIN) {
         return APIError::WOULD_BLOCK;
@@ -331,8 +371,8 @@ APIError APINoiseFrameHelper::try_read_frame_(ParsedFrame *frame) {
       HELPER_LOG("Connection closed");
       return APIError::CONNECTION_CLOSED;
     }
-    rx_buf_len_ += received;
-    if ((size_t) received != to_read) {
+    rx_buf_len_ += static_cast<uint16_t>(received);
+    if (static_cast<uint16_t>(received) != to_read) {
       // not all read
       return APIError::WOULD_BLOCK;
     }
@@ -381,6 +421,8 @@ APIError APINoiseFrameHelper::state_action_() {
     if (aerr != APIError::OK)
       return aerr;
     // ignore contents, may be used in future for flags
+    // Reserve space for: existing prologue + 2 size bytes + frame data
+    prologue_.reserve(prologue_.size() + 2 + frame.msg.size());
     prologue_.push_back((uint8_t) (frame.msg.size() >> 8));
     prologue_.push_back((uint8_t) frame.msg.size());
     prologue_.insert(prologue_.end(), frame.msg.begin(), frame.msg.end());
@@ -389,16 +431,20 @@ APIError APINoiseFrameHelper::state_action_() {
   }
   if (state_ == State::SERVER_HELLO) {
     // send server hello
+    const std::string &name = App.get_name();
+    const std::string &mac = get_mac_address();
+
     std::vector<uint8_t> msg;
+    // Reserve space for: 1 byte proto + name + null + mac + null
+    msg.reserve(1 + name.size() + 1 + mac.size() + 1);
+
     // chosen proto
     msg.push_back(0x01);
 
     // node name, terminated by null byte
-    const std::string &name = App.get_name();
     const uint8_t *name_ptr = reinterpret_cast<const uint8_t *>(name.c_str());
     msg.insert(msg.end(), name_ptr, name_ptr + name.size() + 1);
     // node mac, terminated by null byte
-    const std::string &mac = get_mac_address();
     const uint8_t *mac_ptr = reinterpret_cast<const uint8_t *>(mac.c_str());
     msg.insert(msg.end(), mac_ptr, mac_ptr + mac.size() + 1);
 
@@ -505,7 +551,6 @@ void APINoiseFrameHelper::send_explicit_handshake_reject_(const std::string &rea
   write_frame_(data.data(), data.size());
   state_ = orig_state;
 }
-
 APIError APINoiseFrameHelper::read_packet(ReadPacketBuffer *buffer) {
   int err;
   APIError aerr;
@@ -533,7 +578,7 @@ APIError APINoiseFrameHelper::read_packet(ReadPacketBuffer *buffer) {
     return APIError::CIPHERSTATE_DECRYPT_FAILED;
   }
 
-  size_t msg_size = mbuf.size;
+  uint16_t msg_size = mbuf.size;
   uint8_t *msg_data = frame.msg.data();
   if (msg_size < 4) {
     state_ = State::FAILED;
@@ -559,11 +604,22 @@ APIError APINoiseFrameHelper::read_packet(ReadPacketBuffer *buffer) {
   buffer->type = type;
   return APIError::OK;
 }
-bool APINoiseFrameHelper::can_write_without_blocking() { return state_ == State::DATA && tx_buf_.empty(); }
 APIError APINoiseFrameHelper::write_protobuf_packet(uint16_t type, ProtoWriteBuffer buffer) {
-  int err;
-  APIError aerr;
-  aerr = state_action_();
+  std::vector<uint8_t> *raw_buffer = buffer.get_buffer();
+  uint16_t payload_len = static_cast<uint16_t>(raw_buffer->size() - frame_header_padding_);
+
+  // Resize to include MAC space (required for Noise encryption)
+  raw_buffer->resize(raw_buffer->size() + frame_footer_size_);
+
+  // Use write_protobuf_packets with a single packet
+  std::vector<PacketInfo> packets;
+  packets.emplace_back(type, 0, payload_len);
+
+  return write_protobuf_packets(buffer, packets);
+}
+
+APIError APINoiseFrameHelper::write_protobuf_packets(ProtoWriteBuffer buffer, const std::vector<PacketInfo> &packets) {
+  APIError aerr = state_action_();
   if (aerr != APIError::OK) {
     return aerr;
   }
@@ -572,77 +628,67 @@ APIError APINoiseFrameHelper::write_protobuf_packet(uint16_t type, ProtoWriteBuf
     return APIError::WOULD_BLOCK;
   }
 
-  std::vector<uint8_t> *raw_buffer = buffer.get_buffer();
-  // Message data starts after padding
-  size_t payload_len = raw_buffer->size() - frame_header_padding_;
-  size_t padding = 0;
-  size_t msg_len = 4 + payload_len + padding;
-
-  // We need to resize to include MAC space, but we already reserved it in create_buffer
-  raw_buffer->resize(raw_buffer->size() + frame_footer_size_);
-
-  // Write the noise header in the padded area
-  // Buffer layout:
-  // [0]    - 0x01 indicator byte
-  // [1-2]  - Size of encrypted payload (filled after encryption)
-  // [3-4]  - Message type (encrypted)
-  // [5-6]  - Payload length (encrypted)
-  // [7...] - Actual payload data (encrypted)
-  uint8_t *buf_start = raw_buffer->data();
-  buf_start[0] = 0x01;  // indicator
-  // buf_start[1], buf_start[2] to be set later after encryption
-  const uint8_t msg_offset = 3;
-  buf_start[msg_offset + 0] = (uint8_t) (type >> 8);         // type high byte
-  buf_start[msg_offset + 1] = (uint8_t) type;                // type low byte
-  buf_start[msg_offset + 2] = (uint8_t) (payload_len >> 8);  // data_len high byte
-  buf_start[msg_offset + 3] = (uint8_t) payload_len;         // data_len low byte
-  // payload data is already in the buffer starting at position 7
-
-  NoiseBuffer mbuf;
-  noise_buffer_init(mbuf);
-  // The capacity parameter should be msg_len + frame_footer_size_ (MAC length) to allow space for encryption
-  noise_buffer_set_inout(mbuf, buf_start + msg_offset, msg_len, msg_len + frame_footer_size_);
-  err = noise_cipherstate_encrypt(send_cipher_, &mbuf);
-  if (err != 0) {
-    state_ = State::FAILED;
-    HELPER_LOG("noise_cipherstate_encrypt failed: %s", noise_err_to_str(err).c_str());
-    return APIError::CIPHERSTATE_ENCRYPT_FAILED;
+  if (packets.empty()) {
+    return APIError::OK;
   }
 
-  size_t total_len = 3 + mbuf.size;
-  buf_start[1] = (uint8_t) (mbuf.size >> 8);
-  buf_start[2] = (uint8_t) mbuf.size;
-
-  struct iovec iov;
-  // Point iov_base to the beginning of the buffer (no unused padding in Noise)
-  // We send the entire frame: indicator + size + encrypted(type + data_len + payload + MAC)
-  iov.iov_base = buf_start;
-  iov.iov_len = total_len;
-
-  // write raw to not have two packets sent if NAGLE disabled
-  return write_raw_(&iov, 1);
-}
-APIError APINoiseFrameHelper::try_send_tx_buf_() {
-  // try send from tx_buf
-  while (state_ != State::CLOSED && !tx_buf_.empty()) {
-    ssize_t sent = socket_->write(tx_buf_.data(), tx_buf_.size());
-    if (sent == -1) {
-      if (errno == EWOULDBLOCK || errno == EAGAIN)
-        break;
+  std::vector<uint8_t> *raw_buffer = buffer.get_buffer();
+  this->reusable_iovs_.clear();
+  this->reusable_iovs_.reserve(packets.size());
+
+  // We need to encrypt each packet in place
+  for (const auto &packet : packets) {
+    uint16_t type = packet.message_type;
+    uint16_t offset = packet.offset;
+    uint16_t payload_len = packet.payload_size;
+    uint16_t msg_len = 4 + payload_len;  // type(2) + data_len(2) + payload
+
+    // The buffer already has padding at offset
+    uint8_t *buf_start = raw_buffer->data() + offset;
+
+    // Write noise header
+    buf_start[0] = 0x01;  // indicator
+    // buf_start[1], buf_start[2] to be set after encryption
+
+    // Write message header (to be encrypted)
+    const uint8_t msg_offset = 3;
+    buf_start[msg_offset + 0] = (uint8_t) (type >> 8);         // type high byte
+    buf_start[msg_offset + 1] = (uint8_t) type;                // type low byte
+    buf_start[msg_offset + 2] = (uint8_t) (payload_len >> 8);  // data_len high byte
+    buf_start[msg_offset + 3] = (uint8_t) payload_len;         // data_len low byte
+    // payload data is already in the buffer starting at offset + 7
+
+    // Make sure we have space for MAC
+    // The buffer should already have been sized appropriately
+
+    // Encrypt the message in place
+    NoiseBuffer mbuf;
+    noise_buffer_init(mbuf);
+    noise_buffer_set_inout(mbuf, buf_start + msg_offset, msg_len, msg_len + frame_footer_size_);
+
+    int err = noise_cipherstate_encrypt(send_cipher_, &mbuf);
+    if (err != 0) {
       state_ = State::FAILED;
-      HELPER_LOG("Socket write failed with errno %d", errno);
-      return APIError::SOCKET_WRITE_FAILED;
-    } else if (sent == 0) {
-      break;
+      HELPER_LOG("noise_cipherstate_encrypt failed: %s", noise_err_to_str(err).c_str());
+      return APIError::CIPHERSTATE_ENCRYPT_FAILED;
     }
-    // TODO: inefficient if multiple packets in txbuf
-    // replace with deque of buffers
-    tx_buf_.erase(tx_buf_.begin(), tx_buf_.begin() + sent);
+
+    // Fill in the encrypted size
+    buf_start[1] = (uint8_t) (mbuf.size >> 8);
+    buf_start[2] = (uint8_t) mbuf.size;
+
+    // Add iovec for this encrypted packet
+    struct iovec iov;
+    iov.iov_base = buf_start;
+    iov.iov_len = 3 + mbuf.size;  // indicator + size + encrypted data
+    this->reusable_iovs_.push_back(iov);
   }
 
-  return APIError::OK;
+  // Send all encrypted packets in one writev call
+  return this->write_raw_(this->reusable_iovs_.data(), this->reusable_iovs_.size());
 }
-APIError APINoiseFrameHelper::write_frame_(const uint8_t *data, size_t len) {
+
+APIError APINoiseFrameHelper::write_frame_(const uint8_t *data, uint16_t len) {
   uint8_t header[3];
   header[0] = 0x01;  // indicator
   header[1] = (uint8_t) (len >> 8);
@@ -652,12 +698,12 @@ APIError APINoiseFrameHelper::write_frame_(const uint8_t *data, size_t len) {
   iov[0].iov_base = header;
   iov[0].iov_len = 3;
   if (len == 0) {
-    return write_raw_(iov, 1);
+    return this->write_raw_(iov, 1);
   }
   iov[1].iov_base = const_cast<uint8_t *>(data);
   iov[1].iov_len = len;
 
-  return write_raw_(iov, 2);
+  return this->write_raw_(iov, 2);
 }
 
 /** Initiate the data structures for the handshake.
@@ -752,58 +798,25 @@ APINoiseFrameHelper::~APINoiseFrameHelper() {
   }
 }
 
-APIError APINoiseFrameHelper::close() {
-  state_ = State::CLOSED;
-  int err = socket_->close();
-  if (err == -1)
-    return APIError::CLOSE_FAILED;
-  return APIError::OK;
-}
-APIError APINoiseFrameHelper::shutdown(int how) {
-  int err = socket_->shutdown(how);
-  if (err == -1)
-    return APIError::SHUTDOWN_FAILED;
-  if (how == SHUT_RDWR) {
-    state_ = State::CLOSED;
-  }
-  return APIError::OK;
-}
 extern "C" {
 // declare how noise generates random bytes (here with a good HWRNG based on the RF system)
 void noise_rand_bytes(void *output, size_t len) {
   if (!esphome::random_bytes(reinterpret_cast<uint8_t *>(output), len)) {
-    ESP_LOGE(TAG, "Failed to acquire random bytes, rebooting!");
+    ESP_LOGE(TAG, "Acquiring random bytes failed; rebooting");
     arch_restart();
   }
 }
 }
 
-// Explicit template instantiation for Noise
-template APIError APIFrameHelper::write_raw_<APINoiseFrameHelper::State>(
-    const struct iovec *iov, int iovcnt, socket::Socket *socket, std::vector<uint8_t> &tx_buf_, const std::string &info,
-    APINoiseFrameHelper::State &state, APINoiseFrameHelper::State failed_state);
 #endif  // USE_API_NOISE
 
 #ifdef USE_API_PLAINTEXT
 
 /// Initialize the frame helper, returns OK if successful.
 APIError APIPlaintextFrameHelper::init() {
-  if (state_ != State::INITIALIZE || socket_ == nullptr) {
-    HELPER_LOG("Bad state for init %d", (int) state_);
-    return APIError::BAD_STATE;
-  }
-  int err = socket_->setblocking(false);
-  if (err != 0) {
-    state_ = State::FAILED;
-    HELPER_LOG("Setting nonblocking failed with errno %d", errno);
-    return APIError::TCP_NONBLOCKING_FAILED;
-  }
-  int enable = 1;
-  err = socket_->setsockopt(IPPROTO_TCP, TCP_NODELAY, &enable, sizeof(int));
-  if (err != 0) {
-    state_ = State::FAILED;
-    HELPER_LOG("Setting nodelay failed with errno %d", errno);
-    return APIError::TCP_NODELAY_FAILED;
+  APIError err = init_common_();
+  if (err != APIError::OK) {
+    return err;
   }
 
   state_ = State::DATA;
@@ -814,14 +827,13 @@ APIError APIPlaintextFrameHelper::loop() {
   if (state_ != State::DATA) {
     return APIError::BAD_STATE;
   }
-  // try send pending TX data
-  if (!tx_buf_.empty()) {
+  if (!this->tx_buf_.empty()) {
     APIError err = try_send_tx_buf_();
-    if (err != APIError::OK) {
+    if (err != APIError::OK && err != APIError::WOULD_BLOCK) {
       return err;
     }
   }
-  return APIError::OK;
+  return APIError::OK;  // Convert WOULD_BLOCK to OK to avoid connection termination
 }
 
 /** Read a packet into the rx_buf_. If successful, stores frame data in the frame parameter
@@ -841,12 +853,15 @@ APIError APIPlaintextFrameHelper::try_read_frame_(ParsedFrame *frame) {
 
   // read header
   while (!rx_header_parsed_) {
-    uint8_t data;
-    // Reading one byte at a time is fastest in practice for ESP32 when
-    // there is no data on the wire (which is the common case).
-    // This results in faster failure detection compared to
-    // attempting to read multiple bytes at once.
-    ssize_t received = socket_->read(&data, 1);
+    // Now that we know when the socket is ready, we can read up to 3 bytes
+    // into the rx_header_buf_ before we have to switch back to reading
+    // one byte at a time to ensure we don't read past the message and
+    // into the next one.
+
+    // Read directly into rx_header_buf_ at the current position
+    // Try to get to at least 3 bytes total (indicator + 2 varint bytes), then read one byte at a time
+    ssize_t received =
+        this->socket_->read(&rx_header_buf_[rx_header_buf_pos_], rx_header_buf_pos_ < 3 ? 3 - rx_header_buf_pos_ : 1);
     if (received == -1) {
       if (errno == EWOULDBLOCK || errno == EAGAIN) {
         return APIError::WOULD_BLOCK;
@@ -860,64 +875,74 @@ APIError APIPlaintextFrameHelper::try_read_frame_(ParsedFrame *frame) {
       return APIError::CONNECTION_CLOSED;
     }
 
-    // Successfully read a byte
-
-    // Process byte according to current buffer position
-    if (rx_header_buf_pos_ == 0) {  // Case 1: First byte (indicator byte)
-      if (data != 0x00) {
+    // If this was the first read, validate the indicator byte
+    if (rx_header_buf_pos_ == 0 && received > 0) {
+      if (rx_header_buf_[0] != 0x00) {
         state_ = State::FAILED;
-        HELPER_LOG("Bad indicator byte %u", data);
+        HELPER_LOG("Bad indicator byte %u", rx_header_buf_[0]);
         return APIError::BAD_INDICATOR;
       }
-      // We don't store the indicator byte, just increment position
-      rx_header_buf_pos_ = 1;  // Set to 1 directly
-      continue;                // Need more bytes before we can parse
     }
 
-    // Check buffer overflow before storing
-    if (rx_header_buf_pos_ == 5) {  // Case 2: Buffer would overflow (5 bytes is max allowed)
+    rx_header_buf_pos_ += received;
+
+    // Check for buffer overflow
+    if (rx_header_buf_pos_ >= sizeof(rx_header_buf_)) {
       state_ = State::FAILED;
       HELPER_LOG("Header buffer overflow");
       return APIError::BAD_DATA_PACKET;
     }
 
-    // Store byte in buffer (adjust index to account for skipped indicator byte)
-    rx_header_buf_[rx_header_buf_pos_ - 1] = data;
-
-    // Increment position after storing
-    rx_header_buf_pos_++;
-
-    // Case 3: If we only have one varint byte, we need more
-    if (rx_header_buf_pos_ == 2) {  // Have read indicator + 1 byte
-      continue;                     // Need more bytes before we can parse
+    // Need at least 3 bytes total (indicator + 2 varint bytes) before trying to parse
+    if (rx_header_buf_pos_ < 3) {
+      continue;
     }
 
     // At this point, we have at least 3 bytes total:
-    //   - Validated indicator byte (0x00) but not stored
+    //   - Validated indicator byte (0x00) stored at position 0
     //   - At least 2 bytes in the buffer for the varints
     // Buffer layout:
-    //   First 1-3 bytes: Message size varint (variable length)
-    //     - 2 bytes would only allow up to 16383, which is less than noise's 65535
+    //   [0]: indicator byte (0x00)
+    //   [1-3]: Message size varint (variable length)
+    //     - 2 bytes would only allow up to 16383, which is less than noise's UINT16_MAX (65535)
     //     - 3 bytes allows up to 2097151, ensuring we support at least as much as noise
-    //   Remaining 1-2 bytes: Message type varint (variable length)
+    //   [2-5]: Message type varint (variable length)
     // We now attempt to parse both varints. If either is incomplete,
     // we'll continue reading more bytes.
 
+    // Skip indicator byte at position 0
+    uint8_t varint_pos = 1;
     uint32_t consumed = 0;
-    auto msg_size_varint = ProtoVarInt::parse(&rx_header_buf_[0], rx_header_buf_pos_ - 1, &consumed);
+
+    auto msg_size_varint = ProtoVarInt::parse(&rx_header_buf_[varint_pos], rx_header_buf_pos_ - varint_pos, &consumed);
     if (!msg_size_varint.has_value()) {
       // not enough data there yet
       continue;
     }
 
-    rx_header_parsed_len_ = msg_size_varint->as_uint32();
+    if (msg_size_varint->as_uint32() > std::numeric_limits<uint16_t>::max()) {
+      state_ = State::FAILED;
+      HELPER_LOG("Bad packet: message size %" PRIu32 " exceeds maximum %u", msg_size_varint->as_uint32(),
+                 std::numeric_limits<uint16_t>::max());
+      return APIError::BAD_DATA_PACKET;
+    }
+    rx_header_parsed_len_ = msg_size_varint->as_uint16();
+
+    // Move to next varint position
+    varint_pos += consumed;
 
-    auto msg_type_varint = ProtoVarInt::parse(&rx_header_buf_[consumed], rx_header_buf_pos_ - 1 - consumed, &consumed);
+    auto msg_type_varint = ProtoVarInt::parse(&rx_header_buf_[varint_pos], rx_header_buf_pos_ - varint_pos, &consumed);
     if (!msg_type_varint.has_value()) {
       // not enough data there yet
       continue;
     }
-    rx_header_parsed_type_ = msg_type_varint->as_uint32();
+    if (msg_type_varint->as_uint32() > std::numeric_limits<uint16_t>::max()) {
+      state_ = State::FAILED;
+      HELPER_LOG("Bad packet: message type %" PRIu32 " exceeds maximum %u", msg_type_varint->as_uint32(),
+                 std::numeric_limits<uint16_t>::max());
+      return APIError::BAD_DATA_PACKET;
+    }
+    rx_header_parsed_type_ = msg_type_varint->as_uint16();
     rx_header_parsed_ = true;
   }
   // header reading done
@@ -929,8 +954,8 @@ APIError APIPlaintextFrameHelper::try_read_frame_(ParsedFrame *frame) {
 
   if (rx_buf_len_ < rx_header_parsed_len_) {
     // more data to read
-    size_t to_read = rx_header_parsed_len_ - rx_buf_len_;
-    ssize_t received = socket_->read(&rx_buf_[rx_buf_len_], to_read);
+    uint16_t to_read = rx_header_parsed_len_ - rx_buf_len_;
+    ssize_t received = this->socket_->read(&rx_buf_[rx_buf_len_], to_read);
     if (received == -1) {
       if (errno == EWOULDBLOCK || errno == EAGAIN) {
         return APIError::WOULD_BLOCK;
@@ -943,8 +968,8 @@ APIError APIPlaintextFrameHelper::try_read_frame_(ParsedFrame *frame) {
       HELPER_LOG("Connection closed");
       return APIError::CONNECTION_CLOSED;
     }
-    rx_buf_len_ += received;
-    if ((size_t) received != to_read) {
+    rx_buf_len_ += static_cast<uint16_t>(received);
+    if (static_cast<uint16_t>(received) != to_read) {
       // not all read
       return APIError::WOULD_BLOCK;
     }
@@ -962,7 +987,6 @@ APIError APIPlaintextFrameHelper::try_read_frame_(ParsedFrame *frame) {
   rx_header_parsed_ = false;
   return APIError::OK;
 }
-
 APIError APIPlaintextFrameHelper::read_packet(ReadPacketBuffer *buffer) {
   APIError aerr;
 
@@ -990,7 +1014,7 @@ APIError APIPlaintextFrameHelper::read_packet(ReadPacketBuffer *buffer) {
                          "Bad indicator byte";
       iov[0].iov_base = (void *) msg;
       iov[0].iov_len = 19;
-      write_raw_(iov, 1);
+      this->write_raw_(iov, 1);
     }
     return aerr;
   }
@@ -1001,108 +1025,89 @@ APIError APIPlaintextFrameHelper::read_packet(ReadPacketBuffer *buffer) {
   buffer->type = rx_header_parsed_type_;
   return APIError::OK;
 }
-bool APIPlaintextFrameHelper::can_write_without_blocking() { return state_ == State::DATA && tx_buf_.empty(); }
 APIError APIPlaintextFrameHelper::write_protobuf_packet(uint16_t type, ProtoWriteBuffer buffer) {
-  if (state_ != State::DATA) {
-    return APIError::BAD_STATE;
-  }
-
   std::vector<uint8_t> *raw_buffer = buffer.get_buffer();
-  // Message data starts after padding (frame_header_padding_ = 6)
-  size_t payload_len = raw_buffer->size() - frame_header_padding_;
-
-  // Calculate varint sizes for header components
-  size_t size_varint_len = api::ProtoSize::varint(static_cast<uint32_t>(payload_len));
-  size_t type_varint_len = api::ProtoSize::varint(static_cast<uint32_t>(type));
-  size_t total_header_len = 1 + size_varint_len + type_varint_len;
+  uint16_t payload_len = static_cast<uint16_t>(raw_buffer->size() - frame_header_padding_);
 
-  if (total_header_len > frame_header_padding_) {
-    // Header is too large to fit in the padding
-    return APIError::BAD_ARG;
-  }
+  // Use write_protobuf_packets with a single packet
+  std::vector<PacketInfo> packets;
+  packets.emplace_back(type, 0, payload_len);
 
-  // Calculate where to start writing the header
-  // The header starts at the latest possible position to minimize unused padding
-  //
-  // Example 1 (small values): total_header_len = 3, header_offset = 6 - 3 = 3
-  // [0-2]  - Unused padding
-  // [3]    - 0x00 indicator byte
-  // [4]    - Payload size varint (1 byte, for sizes 0-127)
-  // [5]    - Message type varint (1 byte, for types 0-127)
-  // [6...] - Actual payload data
-  //
-  // Example 2 (medium values): total_header_len = 4, header_offset = 6 - 4 = 2
-  // [0-1]  - Unused padding
-  // [2]    - 0x00 indicator byte
-  // [3-4]  - Payload size varint (2 bytes, for sizes 128-16383)
-  // [5]    - Message type varint (1 byte, for types 0-127)
-  // [6...] - Actual payload data
-  //
-  // Example 3 (large values): total_header_len = 6, header_offset = 6 - 6 = 0
-  // [0]    - 0x00 indicator byte
-  // [1-3]  - Payload size varint (3 bytes, for sizes 16384-2097151)
-  // [4-5]  - Message type varint (2 bytes, for types 128-32767)
-  // [6...] - Actual payload data
-  uint8_t *buf_start = raw_buffer->data();
-  size_t header_offset = frame_header_padding_ - total_header_len;
-
-  // Write the plaintext header
-  buf_start[header_offset] = 0x00;  // indicator
-
-  // Encode size varint directly into buffer
-  ProtoVarInt(payload_len).encode_to_buffer_unchecked(buf_start + header_offset + 1, size_varint_len);
-
-  // Encode type varint directly into buffer
-  ProtoVarInt(type).encode_to_buffer_unchecked(buf_start + header_offset + 1 + size_varint_len, type_varint_len);
-
-  struct iovec iov;
-  // Point iov_base to the beginning of our header (skip unused padding)
-  // This ensures we only send the actual header and payload, not the empty padding bytes
-  iov.iov_base = buf_start + header_offset;
-  iov.iov_len = total_header_len + payload_len;
-
-  return write_raw_(&iov, 1);
+  return write_protobuf_packets(buffer, packets);
 }
-APIError APIPlaintextFrameHelper::try_send_tx_buf_() {
-  // try send from tx_buf
-  while (state_ != State::CLOSED && !tx_buf_.empty()) {
-    ssize_t sent = socket_->write(tx_buf_.data(), tx_buf_.size());
-    if (is_would_block(sent)) {
-      break;
-    } else if (sent == -1) {
-      state_ = State::FAILED;
-      HELPER_LOG("Socket write failed with errno %d", errno);
-      return APIError::SOCKET_WRITE_FAILED;
-    }
-    // TODO: inefficient if multiple packets in txbuf
-    // replace with deque of buffers
-    tx_buf_.erase(tx_buf_.begin(), tx_buf_.begin() + sent);
+
+APIError APIPlaintextFrameHelper::write_protobuf_packets(ProtoWriteBuffer buffer,
+                                                         const std::vector<PacketInfo> &packets) {
+  if (state_ != State::DATA) {
+    return APIError::BAD_STATE;
   }
 
-  return APIError::OK;
-}
+  if (packets.empty()) {
+    return APIError::OK;
+  }
 
-APIError APIPlaintextFrameHelper::close() {
-  state_ = State::CLOSED;
-  int err = socket_->close();
-  if (err == -1)
-    return APIError::CLOSE_FAILED;
-  return APIError::OK;
-}
-APIError APIPlaintextFrameHelper::shutdown(int how) {
-  int err = socket_->shutdown(how);
-  if (err == -1)
-    return APIError::SHUTDOWN_FAILED;
-  if (how == SHUT_RDWR) {
-    state_ = State::CLOSED;
+  std::vector<uint8_t> *raw_buffer = buffer.get_buffer();
+  this->reusable_iovs_.clear();
+  this->reusable_iovs_.reserve(packets.size());
+
+  for (const auto &packet : packets) {
+    uint16_t type = packet.message_type;
+    uint16_t offset = packet.offset;
+    uint16_t payload_len = packet.payload_size;
+
+    // Calculate varint sizes for header layout
+    uint8_t size_varint_len = api::ProtoSize::varint(static_cast<uint32_t>(payload_len));
+    uint8_t type_varint_len = api::ProtoSize::varint(static_cast<uint32_t>(type));
+    uint8_t total_header_len = 1 + size_varint_len + type_varint_len;
+
+    // Calculate where to start writing the header
+    // The header starts at the latest possible position to minimize unused padding
+    //
+    // Example 1 (small values): total_header_len = 3, header_offset = 6 - 3 = 3
+    // [0-2]  - Unused padding
+    // [3]    - 0x00 indicator byte
+    // [4]    - Payload size varint (1 byte, for sizes 0-127)
+    // [5]    - Message type varint (1 byte, for types 0-127)
+    // [6...] - Actual payload data
+    //
+    // Example 2 (medium values): total_header_len = 4, header_offset = 6 - 4 = 2
+    // [0-1]  - Unused padding
+    // [2]    - 0x00 indicator byte
+    // [3-4]  - Payload size varint (2 bytes, for sizes 128-16383)
+    // [5]    - Message type varint (1 byte, for types 0-127)
+    // [6...] - Actual payload data
+    //
+    // Example 3 (large values): total_header_len = 6, header_offset = 6 - 6 = 0
+    // [0]    - 0x00 indicator byte
+    // [1-3]  - Payload size varint (3 bytes, for sizes 16384-2097151)
+    // [4-5]  - Message type varint (2 bytes, for types 128-32767)
+    // [6...] - Actual payload data
+    //
+    // The message starts at offset + frame_header_padding_
+    // So we write the header starting at offset + frame_header_padding_ - total_header_len
+    uint8_t *buf_start = raw_buffer->data() + offset;
+    uint32_t header_offset = frame_header_padding_ - total_header_len;
+
+    // Write the plaintext header
+    buf_start[header_offset] = 0x00;  // indicator
+
+    // Encode size varint directly into buffer
+    ProtoVarInt(payload_len).encode_to_buffer_unchecked(buf_start + header_offset + 1, size_varint_len);
+
+    // Encode type varint directly into buffer
+    ProtoVarInt(type).encode_to_buffer_unchecked(buf_start + header_offset + 1 + size_varint_len, type_varint_len);
+
+    // Add iovec for this packet (header + payload)
+    struct iovec iov;
+    iov.iov_base = buf_start + header_offset;
+    iov.iov_len = total_header_len + payload_len;
+    this->reusable_iovs_.push_back(iov);
   }
-  return APIError::OK;
+
+  // Send all packets in one writev call
+  return write_raw_(this->reusable_iovs_.data(), this->reusable_iovs_.size());
 }
 
-// Explicit template instantiation for Plaintext
-template APIError APIFrameHelper::write_raw_<APIPlaintextFrameHelper::State>(
-    const struct iovec *iov, int iovcnt, socket::Socket *socket, std::vector<uint8_t> &tx_buf_, const std::string &info,
-    APIPlaintextFrameHelper::State &state, APIPlaintextFrameHelper::State failed_state);
 #endif  // USE_API_PLAINTEXT
 
 }  // namespace api