mirror of
https://github.com/esphome/esphome.git
synced 2026-10-18 12:12:48 +00:00
[store_yaml] Use the shared zstd, progmem and raw api client helpers
GetYamlResponse marks its data (progmem), so every platform streams the blob straight from flash and the ESP8266 bounce buffer is gone. The zstd import and the integration test's plaintext client use the shared helpers.
This commit is contained in:
@@ -3099,7 +3099,8 @@ message GetYamlResponse {
|
||||
option (source) = SOURCE_SERVER;
|
||||
option (ifdef) = "USE_STORE_YAML";
|
||||
|
||||
bytes data = 1 [(force) = true];
|
||||
// The blob is in flash, and the device sends it from there
|
||||
bytes data = 1 [(force) = true, (pointer_to_buffer) = true, (progmem) = true];
|
||||
bool done = 2 [(force) = true];
|
||||
// Sent on the first chunk only — the client is expected to cache it.
|
||||
// (firmware bandwidth/flash is expensive; the client has gigabytes.)
|
||||
|
||||
@@ -1222,26 +1222,12 @@ void APIConnection::on_camera_image_request(const CameraImageRequest &msg) {
|
||||
#endif
|
||||
|
||||
#ifdef USE_STORE_YAML
|
||||
#ifdef USE_ESP8266
|
||||
// On ESP8266 the blob lives in instruction flash and can't be read directly, so
|
||||
// each chunk is bounced through a heap buffer via progmem_memcpy. The buffer
|
||||
// exists only while a transfer is in flight; retrieval is rare, so no RAM is
|
||||
// held for the firmware's lifetime. Every other platform sends straight from
|
||||
// the blob, zero-copy, in MTU-sized chunks.
|
||||
static constexpr size_t STORE_YAML_CHUNK_SIZE = 512;
|
||||
#endif
|
||||
|
||||
void APIConnection::on_get_yaml_request() {
|
||||
// A re-request while a transfer is in flight is ignored; see the
|
||||
// GetYamlRequest comment in api.proto. A client that wants to restart
|
||||
// must reconnect.
|
||||
if (this->store_yaml_pos_ != std::numeric_limits<size_t>::max())
|
||||
return;
|
||||
#ifdef USE_ESP8266
|
||||
// OOM aborts by design (NEW_OOM_ABORT): a heap that cannot supply this
|
||||
// buffer is already failing, and the post-reset retry gets a clean heap.
|
||||
this->store_yaml_chunk_buf_ = std::make_unique<uint8_t[]>(STORE_YAML_CHUNK_SIZE);
|
||||
#endif
|
||||
// All responses go through the loop-driven retry below, so a full TX
|
||||
// buffer at request time can't strand the client without a terminal frame.
|
||||
this->store_yaml_pos_ = 0;
|
||||
@@ -1260,11 +1246,7 @@ void APIConnection::try_send_store_yaml_() {
|
||||
// Codegen always embeds a non-empty blob, so total > 0 here.
|
||||
const size_t total = comp->get_size();
|
||||
|
||||
#ifdef USE_ESP8266
|
||||
const size_t chunk_size = STORE_YAML_CHUNK_SIZE;
|
||||
#else
|
||||
const size_t chunk_size = MAX_BATCH_PACKET_SIZE;
|
||||
#endif
|
||||
|
||||
// Camera-style streaming: advance the position only after a successful send,
|
||||
// so a WOULD_BLOCK simply retries the same chunk on the next loop iteration.
|
||||
@@ -1276,12 +1258,8 @@ void APIConnection::try_send_store_yaml_() {
|
||||
const size_t to_send = std::min(remaining, chunk_size);
|
||||
|
||||
GetYamlResponse resp;
|
||||
#ifdef USE_ESP8266
|
||||
progmem_memcpy(this->store_yaml_chunk_buf_.get(), comp->get_data() + this->store_yaml_pos_, to_send);
|
||||
resp.set_data(this->store_yaml_chunk_buf_.get(), to_send);
|
||||
#else
|
||||
resp.set_data(comp->get_data() + this->store_yaml_pos_, to_send);
|
||||
#endif
|
||||
resp.data = comp->get_data() + this->store_yaml_pos_;
|
||||
resp.data_len = to_send;
|
||||
if (this->store_yaml_pos_ == 0) {
|
||||
resp.total_size = static_cast<uint32_t>(total);
|
||||
resp.encoding = StringRef(store_yaml::ENCODING);
|
||||
@@ -1298,9 +1276,6 @@ void APIConnection::try_send_store_yaml_() {
|
||||
|
||||
// Final response (with done=true) sent successfully.
|
||||
this->store_yaml_pos_ = std::numeric_limits<size_t>::max();
|
||||
#ifdef USE_ESP8266
|
||||
this->store_yaml_chunk_buf_.reset();
|
||||
#endif
|
||||
}
|
||||
#endif
|
||||
|
||||
|
||||
@@ -445,12 +445,6 @@ class APIConnection final : public APIServerConnectionBase {
|
||||
void try_send_store_yaml_();
|
||||
// Streaming offset into the PROGMEM blob; max() means "not streaming".
|
||||
size_t store_yaml_pos_{std::numeric_limits<size_t>::max()};
|
||||
#ifdef USE_ESP8266
|
||||
// Bounce buffer for progmem_memcpy, alive only while a transfer is in
|
||||
// flight; retrieval is rare, so the RAM is not held for the firmware's
|
||||
// lifetime. Freed on the terminal frame or with the connection.
|
||||
std::unique_ptr<uint8_t[]> store_yaml_chunk_buf_;
|
||||
#endif
|
||||
#endif
|
||||
|
||||
#ifdef USE_API_HOMEASSISTANT_STATES
|
||||
|
||||
@@ -4251,9 +4251,7 @@ uint32_t BluetoothSetConnectionParamsResponse::calc_size_msg(const void *self) {
|
||||
#ifdef USE_STORE_YAML
|
||||
uint8_t *GetYamlResponse::encode_msg(const void *self, uint8_t *__restrict__ pos PROTO_ENCODE_DEBUG_PARAM) {
|
||||
const auto &msg = *static_cast<const GetYamlResponse *>(self);
|
||||
pos = ProtoEncode::write_raw_byte(pos PROTO_ENCODE_DEBUG_ARG, 10);
|
||||
pos = ProtoEncode::encode_varint_raw(pos PROTO_ENCODE_DEBUG_ARG, msg.data_len_);
|
||||
pos = ProtoEncode::encode_raw(pos PROTO_ENCODE_DEBUG_ARG, msg.data_ptr_, msg.data_len_);
|
||||
pos = ProtoEncode::encode_progmem_bytes_force(pos PROTO_ENCODE_DEBUG_ARG, 1, msg.data, msg.data_len);
|
||||
pos = ProtoEncode::write_raw_byte(pos PROTO_ENCODE_DEBUG_ARG, 16);
|
||||
pos = ProtoEncode::write_raw_byte(pos PROTO_ENCODE_DEBUG_ARG, msg.done ? 0x01 : 0x00);
|
||||
pos = ProtoEncode::encode_uint32(pos PROTO_ENCODE_DEBUG_ARG, 3, msg.total_size);
|
||||
@@ -4263,7 +4261,7 @@ uint8_t *GetYamlResponse::encode_msg(const void *self, uint8_t *__restrict__ pos
|
||||
uint32_t GetYamlResponse::calc_size_msg(const void *self) {
|
||||
const auto &msg = *static_cast<const GetYamlResponse *>(self);
|
||||
uint32_t size = 0;
|
||||
size += ProtoSize::calc_length_force(1, msg.data_len_);
|
||||
size += ProtoSize::calc_length_force(1, msg.data_len);
|
||||
size += ProtoSize::calc_bool_force(1);
|
||||
size += ProtoSize::calc_uint32(1, msg.total_size);
|
||||
size += !msg.encoding.empty() ? 2 + msg.encoding.size() : 0;
|
||||
|
||||
@@ -4166,12 +4166,8 @@ class GetYamlResponse final : public ProtoMessage {
|
||||
#ifdef HAS_PROTO_MESSAGE_DUMP
|
||||
const LogString *message_name() const override { return LOG_STR("get_yaml_response"); }
|
||||
#endif
|
||||
const uint8_t *data_ptr_{nullptr};
|
||||
size_t data_len_{0};
|
||||
void set_data(const uint8_t *data, size_t len) {
|
||||
this->data_ptr_ = data;
|
||||
this->data_len_ = len;
|
||||
}
|
||||
const uint8_t *data{nullptr};
|
||||
uint16_t data_len{0};
|
||||
bool done{false};
|
||||
uint32_t total_size{0};
|
||||
StringRef encoding{nullptr, 0}; // null until set, encode only
|
||||
|
||||
@@ -2957,7 +2957,7 @@ const char *BluetoothSetConnectionParamsResponse::dump_to(DumpBuffer &out) const
|
||||
#ifdef USE_STORE_YAML
|
||||
const char *GetYamlResponse::dump_to(DumpBuffer &out) const {
|
||||
MessageDumpHelper helper(out, ESPHOME_PSTR("GetYamlResponse"));
|
||||
dump_bytes_field(out, ESPHOME_PSTR("data"), this->data_ptr_, this->data_len_);
|
||||
dump_progmem_bytes_field(out, ESPHOME_PSTR("data"), this->data, this->data_len);
|
||||
dump_field(out, ESPHOME_PSTR("done"), this->done);
|
||||
dump_field(out, ESPHOME_PSTR("total_size"), this->total_size);
|
||||
dump_field(out, ESPHOME_PSTR("encoding"), this->encoding);
|
||||
|
||||
@@ -15,14 +15,9 @@ import esphome.config_validation as cv
|
||||
from esphome.const import CONF_API, CONF_ID, CONF_KEY, CONF_RAW_DATA_ID
|
||||
from esphome.core import CORE, EsphomeError, HexInt, Lambda
|
||||
import esphome.final_validate as fv
|
||||
from esphome.helpers import ensure_unique_string
|
||||
from esphome.helpers import ensure_unique_string, zstd_module
|
||||
from esphome.types import ConfigType
|
||||
|
||||
try:
|
||||
from compression import zstd # Python 3.14+ stdlib
|
||||
except ImportError:
|
||||
from backports import zstd # pinned in requirements.txt for Python < 3.14
|
||||
|
||||
_LOGGER = logging.getLogger(__name__)
|
||||
|
||||
CODEOWNERS = ["@bdraco"]
|
||||
@@ -528,7 +523,7 @@ async def to_code(config: ConfigType) -> None:
|
||||
if (note := _uncaptured_note(discovered.unresolved, remote_packages)) is not None:
|
||||
files.append(note)
|
||||
envelope = _pack_envelope(files)
|
||||
compressed = zstd.compress(envelope, level=ZSTD_LEVEL)
|
||||
compressed = zstd_module().compress(envelope, level=ZSTD_LEVEL)
|
||||
|
||||
_LOGGER.info(
|
||||
"store_yaml: embedding %d file(s) as %d bytes (%d uncompressed, %.1f%% ratio)",
|
||||
|
||||
@@ -1,159 +1,28 @@
|
||||
"""End-to-end test for the `store_yaml` recovery flow over the native API.
|
||||
|
||||
Talks plaintext API to a host build directly via asyncio sockets rather than
|
||||
through aioesphomeapi: the released aioesphomeapi shipped with this PR does
|
||||
not yet know about `GetYamlRequest` / `GetYamlResponse`, so the high-level
|
||||
client would silently drop the streamed bytes as "unknown message type".
|
||||
|
||||
The raw client implements just enough of the plaintext framing
|
||||
(``0x00 | varint(size) | varint(msg_type) | payload``, see
|
||||
``api_frame_helper_plaintext.cpp``) to send the empty `GetYamlRequest`
|
||||
(message type 159) and accumulate every `GetYamlResponse` (message type 160)
|
||||
until ``done=true``.
|
||||
Uses the raw plaintext client: the released aioesphomeapi does not know
|
||||
`GetYamlRequest` / `GetYamlResponse` yet, so the high-level client would drop
|
||||
the streamed bytes as an unknown message type.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
|
||||
from aioesphomeapi import api_pb2
|
||||
import pytest
|
||||
|
||||
# The component resolves the stdlib-vs-backport zstd import once; reuse it.
|
||||
from esphome.components.store_yaml import unpack_envelope, zstd
|
||||
from esphome.components.store_yaml import unpack_envelope
|
||||
from esphome.helpers import zstd_module
|
||||
from esphome.yaml_util import find_secret_references
|
||||
|
||||
from .raw_api_client import MESSAGE_TYPE_OF, RawApiClient, decode_fields
|
||||
from .types import RunCompiledFunction
|
||||
|
||||
# Message IDs from esphome/components/api/api.proto.
|
||||
HELLO_REQUEST = 1
|
||||
HELLO_RESPONSE = 2
|
||||
DEVICE_CAPABILITIES_REQUEST = 149
|
||||
DEVICE_CAPABILITIES_RESPONSE = 150
|
||||
# Not in the released aioesphomeapi yet; see esphome/components/api/api.proto.
|
||||
GET_YAML_REQUEST = 159
|
||||
GET_YAML_RESPONSE = 160
|
||||
|
||||
|
||||
def _encode_varint(value: int) -> bytes:
|
||||
"""Encode an unsigned integer as a protobuf varint."""
|
||||
out = bytearray()
|
||||
while True:
|
||||
byte = value & 0x7F
|
||||
value >>= 7
|
||||
if value:
|
||||
out.append(byte | 0x80)
|
||||
else:
|
||||
out.append(byte)
|
||||
return bytes(out)
|
||||
|
||||
|
||||
def _read_varint(buf: bytes, pos: int) -> tuple[int, int]:
|
||||
result = 0
|
||||
shift = 0
|
||||
while True:
|
||||
b = buf[pos]
|
||||
pos += 1
|
||||
result |= (b & 0x7F) << shift
|
||||
if not (b & 0x80):
|
||||
return result, pos
|
||||
shift += 7
|
||||
|
||||
|
||||
def _parse_get_yaml_response(payload: bytes) -> tuple[bytes, bool, int, str]:
|
||||
"""Hand-rolled parser for `GetYamlResponse`.
|
||||
|
||||
Returns ``(data, done, total_size, encoding)``.
|
||||
"""
|
||||
data = b""
|
||||
done = False
|
||||
total_size = 0
|
||||
encoding = ""
|
||||
pos = 0
|
||||
while pos < len(payload):
|
||||
tag, pos = _read_varint(payload, pos)
|
||||
field_number = tag >> 3
|
||||
wire_type = tag & 0x07
|
||||
if wire_type == 0: # varint
|
||||
value, pos = _read_varint(payload, pos)
|
||||
if field_number == 2:
|
||||
done = bool(value)
|
||||
elif field_number == 3:
|
||||
total_size = value
|
||||
elif wire_type == 2: # length-delimited
|
||||
length, pos = _read_varint(payload, pos)
|
||||
chunk = payload[pos : pos + length]
|
||||
pos += length
|
||||
if field_number == 1:
|
||||
data = chunk
|
||||
elif field_number == 4:
|
||||
encoding = chunk.decode("utf-8")
|
||||
else:
|
||||
raise AssertionError(f"unexpected wire type {wire_type}")
|
||||
return data, done, total_size, encoding
|
||||
|
||||
|
||||
def _parse_store_yaml_supported(payload: bytes) -> bool:
|
||||
"""Read `store_yaml.supported` (field 6, field 1) from `DeviceCapabilitiesResponse`."""
|
||||
pos = 0
|
||||
while pos < len(payload):
|
||||
tag, pos = _read_varint(payload, pos)
|
||||
wire_type = tag & 0x07
|
||||
if wire_type == 0:
|
||||
_, pos = _read_varint(payload, pos)
|
||||
continue
|
||||
assert wire_type == 2, f"unexpected wire type {wire_type}"
|
||||
length, pos = _read_varint(payload, pos)
|
||||
chunk = payload[pos : pos + length]
|
||||
pos += length
|
||||
if tag >> 3 != 6:
|
||||
continue
|
||||
sub_pos = 0
|
||||
while sub_pos < len(chunk):
|
||||
sub_tag, sub_pos = _read_varint(chunk, sub_pos)
|
||||
value, sub_pos = _read_varint(chunk, sub_pos)
|
||||
if sub_tag == 0x08:
|
||||
return bool(value)
|
||||
return False
|
||||
|
||||
|
||||
async def _read_varint_from(reader: asyncio.StreamReader) -> int:
|
||||
"""Read a protobuf varint byte-by-byte from a stream."""
|
||||
result = 0
|
||||
shift = 0
|
||||
while True:
|
||||
byte = (await reader.readexactly(1))[0]
|
||||
result |= (byte & 0x7F) << shift
|
||||
if not (byte & 0x80):
|
||||
return result
|
||||
shift += 7
|
||||
|
||||
|
||||
class _PlaintextClient:
|
||||
"""Just-enough plaintext API client for one short streaming exchange."""
|
||||
|
||||
def __init__(
|
||||
self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter
|
||||
) -> None:
|
||||
self._reader = reader
|
||||
self._writer = writer
|
||||
|
||||
async def send(self, msg_type: int, payload: bytes = b"") -> None:
|
||||
# Frame: 0x00 | varint(payload_size) | varint(message_id) | payload
|
||||
frame = (
|
||||
b"\x00" + _encode_varint(len(payload)) + _encode_varint(msg_type) + payload
|
||||
)
|
||||
self._writer.write(frame)
|
||||
await self._writer.drain()
|
||||
|
||||
async def recv(self) -> tuple[int, bytes]:
|
||||
# Read preamble byte (must be 0x00 for plaintext).
|
||||
preamble = await self._reader.readexactly(1)
|
||||
assert preamble == b"\x00", f"unexpected preamble {preamble!r}"
|
||||
|
||||
payload_size = await _read_varint_from(self._reader)
|
||||
msg_type = await _read_varint_from(self._reader)
|
||||
payload = await self._reader.readexactly(payload_size) if payload_size else b""
|
||||
return msg_type, payload
|
||||
# DeviceCapabilitiesResponse.store_yaml and StoreYamlCapabilities.supported
|
||||
CAPABILITIES_STORE_YAML_FIELD = 6
|
||||
STORE_YAML_SUPPORTED_FIELD = 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@@ -162,62 +31,34 @@ async def test_store_yaml_recovery(
|
||||
run_compiled: RunCompiledFunction,
|
||||
unused_tcp_port: int,
|
||||
) -> None:
|
||||
"""Compile a host build with `store_yaml`, ask it to stream the YAML back,
|
||||
decompress, and verify the recovered file tree matches the source fixture."""
|
||||
async with run_compiled(yaml_config):
|
||||
# Open a raw TCP connection to the API server.
|
||||
reader, writer = await asyncio.wait_for(
|
||||
asyncio.open_connection("127.0.0.1", unused_tcp_port),
|
||||
timeout=10.0,
|
||||
"""Stream the embedded YAML back from a host build and check the recovered files."""
|
||||
async with run_compiled(yaml_config), RawApiClient(unused_tcp_port) as client:
|
||||
await client.connect("store_yaml integration test")
|
||||
|
||||
# Clients learn the device can answer get_yaml from its capabilities.
|
||||
await client.send_message(api_pb2.DeviceCapabilitiesRequest())
|
||||
capabilities = decode_fields(
|
||||
await client.read_frame(MESSAGE_TYPE_OF[api_pb2.DeviceCapabilitiesResponse])
|
||||
)
|
||||
(store_yaml,) = capabilities[CAPABILITIES_STORE_YAML_FIELD]
|
||||
assert decode_fields(store_yaml).get(STORE_YAML_SUPPORTED_FIELD) == [1], (
|
||||
"expected DeviceCapabilitiesResponse to report store_yaml.supported"
|
||||
)
|
||||
client = _PlaintextClient(reader, writer)
|
||||
try:
|
||||
# HelloRequest: client_info (field 1, length-delimited string).
|
||||
# Password auth (the old ConnectRequest/Response exchange at message
|
||||
# IDs 3/4) was removed in 2026.1.0, so a successful HelloResponse is
|
||||
# all the handshake we need before issuing application requests.
|
||||
client_info = b"store_yaml integration test"
|
||||
api_version = b"\x10\x01\x18\x0e" # api_version_major=1, minor=14
|
||||
hello_payload = (
|
||||
b"\x0a" + _encode_varint(len(client_info)) + client_info + api_version
|
||||
)
|
||||
await client.send(HELLO_REQUEST, hello_payload)
|
||||
msg_type, _ = await asyncio.wait_for(client.recv(), timeout=5.0)
|
||||
assert msg_type == HELLO_RESPONSE, f"expected HelloResponse, got {msg_type}"
|
||||
|
||||
# Clients learn the device can answer get_yaml from its capabilities.
|
||||
await client.send(DEVICE_CAPABILITIES_REQUEST, b"")
|
||||
while True:
|
||||
msg_type, payload = await asyncio.wait_for(client.recv(), timeout=5.0)
|
||||
if msg_type == DEVICE_CAPABILITIES_RESPONSE:
|
||||
break
|
||||
assert _parse_store_yaml_supported(payload), (
|
||||
"expected DeviceCapabilitiesResponse to report store_yaml.supported"
|
||||
)
|
||||
|
||||
# The actual request under test.
|
||||
await client.send(GET_YAML_REQUEST, b"")
|
||||
|
||||
chunks: list[bytes] = []
|
||||
advertised_total: int | None = None
|
||||
advertised_encoding: str | None = None
|
||||
done = False
|
||||
while not done:
|
||||
msg_type, payload = await asyncio.wait_for(client.recv(), timeout=5.0)
|
||||
if msg_type != GET_YAML_RESPONSE:
|
||||
# Tolerate intervening server messages (e.g. pings).
|
||||
continue
|
||||
chunk, done, total_size, encoding = _parse_get_yaml_response(payload)
|
||||
if encoding:
|
||||
advertised_encoding = encoding
|
||||
if total_size and advertised_total is None:
|
||||
advertised_total = total_size
|
||||
if chunk:
|
||||
chunks.append(chunk)
|
||||
finally:
|
||||
writer.close()
|
||||
with contextlib.suppress(ConnectionError, OSError):
|
||||
await writer.wait_closed()
|
||||
await client.send_raw(GET_YAML_REQUEST, b"")
|
||||
chunks: list[bytes] = []
|
||||
advertised_total: int | None = None
|
||||
advertised_encoding: str | None = None
|
||||
done = False
|
||||
while not done:
|
||||
fields = decode_fields(await client.read_frame(GET_YAML_RESPONSE))
|
||||
if data := fields.get(1):
|
||||
chunks.append(data[0])
|
||||
done = fields.get(2) == [1]
|
||||
if total := fields.get(3):
|
||||
advertised_total = total[0]
|
||||
if encoding := fields.get(4):
|
||||
advertised_encoding = encoding[0].decode("utf-8")
|
||||
|
||||
compressed = b"".join(chunks)
|
||||
assert advertised_encoding == "zstd", (
|
||||
@@ -227,7 +68,7 @@ async def test_store_yaml_recovery(
|
||||
f"server advertised {advertised_total} bytes but we received {len(compressed)}"
|
||||
)
|
||||
|
||||
envelope = zstd.decompress(compressed)
|
||||
envelope = zstd_module().decompress(compressed)
|
||||
files = unpack_envelope(envelope)
|
||||
|
||||
assert files, "envelope should contain at least one file"
|
||||
|
||||
Reference in New Issue
Block a user