Refactor API frame helpers to enable buffer reuse (#8825)

This commit is contained in:
J. Nick Koston 2025-05-18 17:05:20 -04:00 committed by GitHub
parent 574aabdede
commit b0c1e0e28c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 162 additions and 46 deletions

View File

@ -1645,7 +1645,7 @@ bool APIConnection::send_buffer(ProtoWriteBuffer buffer, uint32_t message_type)
return false;
}
APIError err = this->helper_->write_packet(message_type, buffer.get_buffer()->data(), buffer.get_buffer()->size());
APIError err = this->helper_->write_protobuf_packet(message_type, buffer);
if (err == APIError::WOULD_BLOCK)
return false;
if (err != APIError::OK) {

View File

@ -405,7 +405,14 @@ class APIConnection : public APIServerConnection {
ProtoWriteBuffer create_buffer(uint32_t reserve_size) override {
// FIXME: ensure no recursive writes can happen
this->proto_write_buffer_.clear();
this->proto_write_buffer_.reserve(reserve_size);
// Get header padding size - used for both reserve and insert
uint8_t header_padding = this->helper_->frame_header_padding();
// Reserve space for header padding + message + footer
// - Header padding: space for protocol headers (7 bytes for Noise, 6 for Plaintext)
// - Footer: space for MAC (16 bytes for Noise, 0 for Plaintext)
this->proto_write_buffer_.reserve(reserve_size + header_padding + this->helper_->frame_footer_size());
// Insert header padding bytes so message encoding starts at the correct position
this->proto_write_buffer_.insert(this->proto_write_buffer_.begin(), header_padding, 0);
return {&this->proto_write_buffer_};
}
bool try_to_clear_buffer(bool log_out_of_space);

View File

@ -493,9 +493,12 @@ void APINoiseFrameHelper::send_explicit_handshake_reject_(const std::string &rea
std::vector<uint8_t> data;
data.resize(reason.length() + 1);
data[0] = 0x01; // failure
for (size_t i = 0; i < reason.length(); i++) {
data[i + 1] = (uint8_t) reason[i];
// Copy error message in bulk
if (!reason.empty()) {
std::memcpy(data.data() + 1, reason.c_str(), reason.length());
}
// temporarily remove failed state
auto orig_state = state_;
state_ = State::EXPLICIT_REJECT;
@ -557,7 +560,7 @@ APIError APINoiseFrameHelper::read_packet(ReadPacketBuffer *buffer) {
return APIError::OK;
}
bool APINoiseFrameHelper::can_write_without_blocking() { return state_ == State::DATA && tx_buf_.empty(); }
APIError APINoiseFrameHelper::write_packet(uint16_t type, const uint8_t *payload, size_t payload_len) {
APIError APINoiseFrameHelper::write_protobuf_packet(uint16_t type, ProtoWriteBuffer buffer) {
int err;
APIError aerr;
aerr = state_action_();
@ -569,31 +572,36 @@ APIError APINoiseFrameHelper::write_packet(uint16_t type, const uint8_t *payload
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;
size_t frame_len = 3 + msg_len + noise_cipherstate_get_mac_length(send_cipher_);
auto tmpbuf = std::unique_ptr<uint8_t[]>{new (std::nothrow) uint8_t[frame_len]};
if (tmpbuf == nullptr) {
HELPER_LOG("Could not allocate for writing packet");
return APIError::OUT_OF_MEMORY;
}
tmpbuf[0] = 0x01; // indicator
// tmpbuf[1], tmpbuf[2] to be set later
// 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;
const uint8_t payload_offset = msg_offset + 4;
tmpbuf[msg_offset + 0] = (uint8_t) (type >> 8); // type
tmpbuf[msg_offset + 1] = (uint8_t) type;
tmpbuf[msg_offset + 2] = (uint8_t) (payload_len >> 8); // data_len
tmpbuf[msg_offset + 3] = (uint8_t) payload_len;
// copy data
std::copy(payload, payload + payload_len, &tmpbuf[payload_offset]);
// fill padding with zeros
std::fill(&tmpbuf[payload_offset + payload_len], &tmpbuf[frame_len], 0);
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);
noise_buffer_set_inout(mbuf, &tmpbuf[msg_offset], msg_len, frame_len - msg_offset);
// 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;
@ -602,11 +610,13 @@ APIError APINoiseFrameHelper::write_packet(uint16_t type, const uint8_t *payload
}
size_t total_len = 3 + mbuf.size;
tmpbuf[1] = (uint8_t) (mbuf.size >> 8);
tmpbuf[2] = (uint8_t) mbuf.size;
buf_start[1] = (uint8_t) (mbuf.size >> 8);
buf_start[2] = (uint8_t) mbuf.size;
struct iovec iov;
iov.iov_base = &tmpbuf[0];
// 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
@ -718,6 +728,8 @@ APIError APINoiseFrameHelper::check_handshake_finished_() {
return APIError::HANDSHAKESTATE_SPLIT_FAILED;
}
frame_footer_size_ = noise_cipherstate_get_mac_length(send_cipher_);
HELPER_LOG("Handshake complete!");
noise_handshakestate_free(handshake_);
handshake_ = nullptr;
@ -990,28 +1002,66 @@ APIError APIPlaintextFrameHelper::read_packet(ReadPacketBuffer *buffer) {
return APIError::OK;
}
bool APIPlaintextFrameHelper::can_write_without_blocking() { return state_ == State::DATA && tx_buf_.empty(); }
APIError APIPlaintextFrameHelper::write_packet(uint16_t type, const uint8_t *payload, size_t payload_len) {
APIError APIPlaintextFrameHelper::write_protobuf_packet(uint16_t type, ProtoWriteBuffer buffer) {
if (state_ != State::DATA) {
return APIError::BAD_STATE;
}
std::vector<uint8_t> header;
header.reserve(1 + api::ProtoSize::varint(static_cast<uint32_t>(payload_len)) +
api::ProtoSize::varint(static_cast<uint32_t>(type)));
header.push_back(0x00);
ProtoVarInt(payload_len).encode(header);
ProtoVarInt(type).encode(header);
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_;
struct iovec iov[2];
iov[0].iov_base = &header[0];
iov[0].iov_len = header.size();
if (payload_len == 0) {
return write_raw_(iov, 1);
// 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;
if (total_header_len > frame_header_padding_) {
// Header is too large to fit in the padding
return APIError::BAD_ARG;
}
iov[1].iov_base = const_cast<uint8_t *>(payload);
iov[1].iov_len = payload_len;
return write_raw_(iov, 2);
// 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);
}
APIError APIPlaintextFrameHelper::try_send_tx_buf_() {
// try send from tx_buf

View File

@ -16,6 +16,8 @@
namespace esphome {
namespace api {
class ProtoWriteBuffer;
struct ReadPacketBuffer {
std::vector<uint8_t> container;
uint16_t type;
@ -65,32 +67,46 @@ class APIFrameHelper {
virtual APIError loop() = 0;
virtual APIError read_packet(ReadPacketBuffer *buffer) = 0;
virtual bool can_write_without_blocking() = 0;
virtual APIError write_packet(uint16_t type, const uint8_t *data, size_t len) = 0;
virtual APIError write_protobuf_packet(uint16_t type, ProtoWriteBuffer buffer) = 0;
virtual std::string getpeername() = 0;
virtual int getpeername(struct sockaddr *addr, socklen_t *addrlen) = 0;
virtual APIError close() = 0;
virtual APIError shutdown(int how) = 0;
// Give this helper a name for logging
virtual void set_log_info(std::string info) = 0;
// Get the frame header padding required by this protocol
virtual uint8_t frame_header_padding() = 0;
// Get the frame footer size required by this protocol
virtual uint8_t frame_footer_size() = 0;
protected:
// Common implementation for writing raw data to socket
template<typename StateEnum>
APIError 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);
uint8_t frame_header_padding_{0};
uint8_t frame_footer_size_{0};
};
#ifdef USE_API_NOISE
class APINoiseFrameHelper : public APIFrameHelper {
public:
APINoiseFrameHelper(std::unique_ptr<socket::Socket> socket, std::shared_ptr<APINoiseContext> ctx)
: socket_(std::move(socket)), ctx_(std::move(std::move(ctx))) {}
: socket_(std::move(socket)), ctx_(std::move(ctx)) {
// Noise header structure:
// Pos 0: indicator (0x01)
// Pos 1-2: encrypted payload size (16-bit big-endian)
// Pos 3-6: encrypted type (16-bit) + data_len (16-bit)
// Pos 7+: actual payload data
frame_header_padding_ = 7;
}
~APINoiseFrameHelper() override;
APIError init() override;
APIError loop() override;
APIError read_packet(ReadPacketBuffer *buffer) override;
bool can_write_without_blocking() override;
APIError write_packet(uint16_t type, const uint8_t *payload, size_t len) override;
APIError write_protobuf_packet(uint16_t type, ProtoWriteBuffer buffer) override;
std::string getpeername() override { return this->socket_->getpeername(); }
int getpeername(struct sockaddr *addr, socklen_t *addrlen) override {
return this->socket_->getpeername(addr, addrlen);
@ -99,6 +115,10 @@ class APINoiseFrameHelper : public APIFrameHelper {
APIError shutdown(int how) override;
// Give this helper a name for logging
void set_log_info(std::string info) override { info_ = std::move(info); }
// Get the frame header padding required by this protocol
uint8_t frame_header_padding() override { return frame_header_padding_; }
// Get the frame footer size required by this protocol
uint8_t frame_footer_size() override { return frame_footer_size_; }
protected:
struct ParsedFrame {
@ -152,13 +172,20 @@ class APINoiseFrameHelper : public APIFrameHelper {
#ifdef USE_API_PLAINTEXT
class APIPlaintextFrameHelper : public APIFrameHelper {
public:
APIPlaintextFrameHelper(std::unique_ptr<socket::Socket> socket) : socket_(std::move(socket)) {}
APIPlaintextFrameHelper(std::unique_ptr<socket::Socket> socket) : socket_(std::move(socket)) {
// Plaintext header structure (worst case):
// Pos 0: indicator (0x00)
// Pos 1-3: payload size varint (up to 3 bytes)
// Pos 4-5: message type varint (up to 2 bytes)
// Pos 6+: actual payload data
frame_header_padding_ = 6;
}
~APIPlaintextFrameHelper() override = default;
APIError init() override;
APIError loop() override;
APIError read_packet(ReadPacketBuffer *buffer) override;
bool can_write_without_blocking() override;
APIError write_packet(uint16_t type, const uint8_t *payload, size_t len) override;
APIError write_protobuf_packet(uint16_t type, ProtoWriteBuffer buffer) override;
std::string getpeername() override { return this->socket_->getpeername(); }
int getpeername(struct sockaddr *addr, socklen_t *addrlen) override {
return this->socket_->getpeername(addr, addrlen);
@ -167,6 +194,10 @@ class APIPlaintextFrameHelper : public APIFrameHelper {
APIError shutdown(int how) override;
// Give this helper a name for logging
void set_log_info(std::string info) override { info_ = std::move(info); }
// Get the frame header padding required by this protocol
uint8_t frame_header_padding() override { return frame_header_padding_; }
// Get the frame footer size required by this protocol
uint8_t frame_footer_size() override { return frame_footer_size_; }
protected:
struct ParsedFrame {

View File

@ -83,6 +83,34 @@ class ProtoVarInt {
return static_cast<int64_t>(this->value_ >> 1);
}
}
/**
* Encode the varint value to a pre-allocated buffer without bounds checking.
*
* @param buffer The pre-allocated buffer to write the encoded varint to
* @param len The size of the buffer in bytes
*
* @note The caller is responsible for ensuring the buffer is large enough
* to hold the encoded value. Use ProtoSize::varint() to calculate
* the exact size needed before calling this method.
* @note No bounds checking is performed for performance reasons.
*/
void encode_to_buffer_unchecked(uint8_t *buffer, size_t len) {
uint64_t val = this->value_;
if (val <= 0x7F) {
buffer[0] = val;
return;
}
size_t i = 0;
while (val && i < len) {
uint8_t temp = val & 0x7F;
val >>= 7;
if (val) {
buffer[i++] = temp | 0x80;
} else {
buffer[i++] = temp;
}
}
}
void encode(std::vector<uint8_t> &out) {
uint64_t val = this->value_;
if (val <= 0x7F) {