[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:
J. Nick Koston
2026-10-07 13:21:05 -10:00
parent bbf1390fea
commit d7b5a8e9c8
8 changed files with 49 additions and 249 deletions
+2 -1
View File
@@ -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.)
+2 -27
View File
@@ -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
-6
View File
@@ -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
+2 -4
View File
@@ -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;
+2 -6
View File
@@ -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
+1 -1
View File
@@ -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);
+2 -7
View File
@@ -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)",
+38 -197
View File
@@ -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"