esphome/esphome/components/api/api_frame_helper.cpp

#include "api_frame_helper.h"
#ifdef USE_API
#include "api_connection.h"  // For ClientInfo struct
#include "esphome/core/application.h"
#include "esphome/core/hal.h"
#include "esphome/core/helpers.h"
#include "esphome/core/log.h"
#include "proto.h"
#include <cstring>
#include <cinttypes>

namespace esphome::api {

static const char *const TAG = "api.frame_helper";

#define HELPER_LOG(msg, ...) ESP_LOGVV(TAG, "%s: " msg, this->client_info_->get_combined_info().c_str(), ##__VA_ARGS__)

#ifdef HELPER_LOG_PACKETS
#define LOG_PACKET_RECEIVED(buffer) ESP_LOGVV(TAG, "Received frame: %s", format_hex_pretty(buffer).c_str())
#define LOG_PACKET_SENDING(data, len) ESP_LOGVV(TAG, "Sending raw: %s", format_hex_pretty(data, len).c_str())
#else
#define LOG_PACKET_RECEIVED(buffer) ((void) 0)
#define LOG_PACKET_SENDING(data, len) ((void) 0)
#endif

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) {
    return "OK";
  } else if (err == APIError::WOULD_BLOCK) {
    return "WOULD_BLOCK";
  } else if (err == APIError::BAD_INDICATOR) {
    return "BAD_INDICATOR";
  } else if (err == APIError::BAD_DATA_PACKET) {
    return "BAD_DATA_PACKET";
  } else if (err == APIError::TCP_NODELAY_FAILED) {
    return "TCP_NODELAY_FAILED";
  } else if (err == APIError::TCP_NONBLOCKING_FAILED) {
    return "TCP_NONBLOCKING_FAILED";
  } else if (err == APIError::CLOSE_FAILED) {
    return "CLOSE_FAILED";
  } else if (err == APIError::SHUTDOWN_FAILED) {
    return "SHUTDOWN_FAILED";
  } else if (err == APIError::BAD_STATE) {
    return "BAD_STATE";
  } else if (err == APIError::BAD_ARG) {
    return "BAD_ARG";
  } else if (err == APIError::SOCKET_READ_FAILED) {
    return "SOCKET_READ_FAILED";
  } else if (err == APIError::SOCKET_WRITE_FAILED) {
    return "SOCKET_WRITE_FAILED";
  } else if (err == APIError::OUT_OF_MEMORY) {
    return "OUT_OF_MEMORY";
  } else if (err == APIError::CONNECTION_CLOSED) {
    return "CONNECTION_CLOSED";
  }
#ifdef USE_API_NOISE
  else if (err == APIError::BAD_HANDSHAKE_PACKET_LEN) {
    return "BAD_HANDSHAKE_PACKET_LEN";
  } else if (err == APIError::HANDSHAKESTATE_READ_FAILED) {
    return "HANDSHAKESTATE_READ_FAILED";
  } else if (err == APIError::HANDSHAKESTATE_WRITE_FAILED) {
    return "HANDSHAKESTATE_WRITE_FAILED";
  } else if (err == APIError::HANDSHAKESTATE_BAD_STATE) {
    return "HANDSHAKESTATE_BAD_STATE";
  } else if (err == APIError::CIPHERSTATE_DECRYPT_FAILED) {
    return "CIPHERSTATE_DECRYPT_FAILED";
  } else if (err == APIError::CIPHERSTATE_ENCRYPT_FAILED) {
    return "CIPHERSTATE_ENCRYPT_FAILED";
  } else if (err == APIError::HANDSHAKESTATE_SETUP_FAILED) {
    return "HANDSHAKESTATE_SETUP_FAILED";
  } else if (err == APIError::HANDSHAKESTATE_SPLIT_FAILED) {
    return "HANDSHAKESTATE_SPLIT_FAILED";
  } else if (err == APIError::BAD_HANDSHAKE_ERROR_BYTE) {
    return "BAD_HANDSHAKE_ERROR_BYTE";
  }
#endif
  return "UNKNOWN";
}

// Default implementation for loop - handles sending buffered data
APIError APIFrameHelper::loop() {
  if (!this->tx_buf_.empty()) {
    APIError err = try_send_tx_buf_();
    if (err != APIError::OK && err != APIError::WOULD_BLOCK) {
      return err;
    }
  }
  return APIError::OK;  // Convert WOULD_BLOCK to OK to avoid connection termination
}

// Common socket write error handling
APIError APIFrameHelper::handle_socket_write_error_() {
  if (errno == EWOULDBLOCK || errno == EAGAIN) {
    return APIError::WOULD_BLOCK;
  }
  HELPER_LOG("Socket write failed with errno %d", errno);
  this->state_ = State::FAILED;
  return APIError::SOCKET_WRITE_FAILED;
}

// Helper method to buffer data from IOVs
void APIFrameHelper::buffer_data_from_iov_(const struct iovec *iov, int iovcnt, uint16_t total_write_len,
                                           uint16_t offset) {
  SendBuffer buffer;
  buffer.size = total_write_len - offset;
  buffer.data = std::make_unique<uint8_t[]>(buffer.size);

  uint16_t to_skip = offset;
  uint16_t write_pos = 0;

  for (int i = 0; i < iovcnt; i++) {
    if (to_skip >= iov[i].iov_len) {
      // Skip this entire segment
      to_skip -= static_cast<uint16_t>(iov[i].iov_len);
    } else {
      // Include this segment (partially or fully)
      const uint8_t *src = reinterpret_cast<uint8_t *>(iov[i].iov_base) + to_skip;
      uint16_t len = static_cast<uint16_t>(iov[i].iov_len) - to_skip;
      std::memcpy(buffer.data.get() + write_pos, src, len);
      write_pos += len;
      to_skip = 0;
    }
  }
  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, uint16_t total_write_len) {
  // 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

  if (iovcnt == 0)
    return APIError::OK;  // Nothing to do, success

#ifdef HELPER_LOG_PACKETS
  for (int i = 0; i < iovcnt; i++) {
    LOG_PACKET_SENDING(reinterpret_cast<uint8_t *>(iov[i].iov_base), iov[i].iov_len);
  }
#endif

  // 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 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, 0);
      return APIError::OK;  // Success, data buffered
    }
  }

  // Try to send directly if no buffered data
  // Optimize for single iovec case (common for plaintext API)
  ssize_t sent =
      (iovcnt == 1) ? this->socket_->write(iov[0].iov_base, iov[0].iov_len) : this->socket_->writev(iov, iovcnt);

  if (sent == -1) {
    APIError err = this->handle_socket_write_error_();
    if (err == APIError::WOULD_BLOCK) {
      // Socket would block, buffer the data
      this->buffer_data_from_iov_(iov, iovcnt, total_write_len, 0);
      return APIError::OK;  // Success, data buffered
    }
    return err;  // Socket write failed
  } else if (static_cast<uint16_t>(sent) < total_write_len) {
    // Partially sent, buffer the remaining data
    this->buffer_data_from_iov_(iov, iovcnt, total_write_len, static_cast<uint16_t>(sent));
  }

  return APIError::OK;  // Success, all data sent or buffered
}

// 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) {
      return this->handle_socket_write_error_();
    } 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) {
    HELPER_LOG("Bad state for init %d", (int) state_);
    return APIError::BAD_STATE;
  }
  int err = this->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 = this->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;
  }
  return APIError::OK;
}

APIError APIFrameHelper::handle_socket_read_result_(ssize_t received) {
  if (received == -1) {
    if (errno == EWOULDBLOCK || errno == EAGAIN) {
      return APIError::WOULD_BLOCK;
    }
    state_ = State::FAILED;
    HELPER_LOG("Socket read failed with errno %d", errno);
    return APIError::SOCKET_READ_FAILED;
  } else if (received == 0) {
    state_ = State::FAILED;
    HELPER_LOG("Connection closed");
    return APIError::CONNECTION_CLOSED;
  }
  return APIError::OK;
}

}  // namespace esphome::api
#endif