[api] Keep the protobuf write cursor in a register and emit sub-message tags inline (#20124)

tests/components/api/test_proto_sub_message.cpp

@@ -0,0 +1,88 @@
+#include <gtest/gtest.h>
+
+#include <cstdint>
+#include <vector>
+
+#include "esphome/components/api/api_buffer.h"
+#include "esphome/components/api/proto.h"
+
+namespace esphome::api::testing {
+
+// Sub-message whose body is raw bytes, so the body size is set directly by the test.
+struct BlobMessage {
+  const uint8_t *data;
+  uint32_t len;
+  static uint8_t *encode_msg(const void *self, uint8_t *pos PROTO_ENCODE_DEBUG_PARAM) {
+    const auto &msg = *static_cast<const BlobMessage *>(self);
+    // An empty vector's data() may be null, and memcpy needs a valid source even for zero bytes
+    if (msg.len == 0)
+      return pos;
+    return ProtoEncode::encode_raw(pos PROTO_ENCODE_DEBUG_ARG, msg.data, msg.len);
+  }
+  static uint32_t calc_size_msg(const void *self) { return static_cast<const BlobMessage *>(self)->len; }
+};
+
+static void append_varint(std::vector<uint8_t> &out, uint32_t value) {
+  while (value > 0x7F) {
+    out.push_back(static_cast<uint8_t>(value | 0x80));
+    value >>= 7;
+  }
+  out.push_back(static_cast<uint8_t>(value));
+}
+
+static std::vector<uint8_t> make_body(uint32_t len) {
+  std::vector<uint8_t> body(len);
+  for (uint32_t i = 0; i < len; i++)
+    body[i] = static_cast<uint8_t>(i * 7 + 1);
+  return body;
+}
+
+static std::vector<uint8_t> expected_field(uint32_t field_id, const std::vector<uint8_t> &body) {
+  std::vector<uint8_t> out;
+  append_varint(out, (field_id << 3) | 2);
+  append_varint(out, body.size());
+  out.insert(out.end(), body.begin(), body.end());
+  return out;
+}
+
+// Encodes into a buffer of exactly the expected size, so any overrun trips ASan or the debug bounds check.
+template<bool OPTIONAL> static void verify(uint32_t field_id, uint32_t body_len) {
+  std::vector<uint8_t> body = make_body(body_len);
+  std::vector<uint8_t> expected = expected_field(field_id, body);
+  if (OPTIONAL && body_len == 0)
+    expected.clear();
+  BlobMessage msg{body.data(), body_len};
+
+  APIBuffer buf;
+  ASSERT_TRUE(buf.resize(expected.empty() ? 1 : expected.size()));
+  uint8_t *pos = buf.data();
+#ifdef ESPHOME_DEBUG_API
+  uint8_t *proto_debug_end_ = buf.data() + buf.size();
+#endif
+  uint8_t *end;
+  if constexpr (OPTIONAL) {
+    end = ProtoEncode::encode_optional_sub_message(pos PROTO_ENCODE_DEBUG_ARG, field_id, msg);
+  } else {
+    end = ProtoEncode::encode_sub_message(pos PROTO_ENCODE_DEBUG_ARG, field_id, msg);
+  }
+
+  ASSERT_EQ(static_cast<size_t>(end - buf.data()), expected.size()) << "field " << field_id << " body " << body_len;
+  EXPECT_EQ(std::vector<uint8_t>(buf.data(), end), expected) << "field " << field_id << " body " << body_len;
+}
+
+TEST(ProtoSubMessage, OneByteTag) { verify<false>(4, 10); }
+TEST(ProtoSubMessage, TwoByteTag) { verify<false>(20, 10); }
+TEST(ProtoSubMessage, EmptyBody) { verify<false>(20, 0); }
+TEST(ProtoSubMessage, LongestOneByteLength) { verify<false>(20, 127); }
+// The length outgrows its reserved byte, so the body is moved forward
+TEST(ProtoSubMessage, TwoByteLength) {
+  verify<false>(20, 128);
+  verify<false>(25, 200);
+}
+TEST(ProtoSubMessage, ThreeByteLength) { verify<false>(4, 20000); }
+
+TEST(ProtoOptionalSubMessage, EmptyIsSkipped) { verify<true>(22, 0); }
+TEST(ProtoOptionalSubMessage, OneByteTag) { verify<true>(1, 10); }
+TEST(ProtoOptionalSubMessage, TwoByteTagAndLength) { verify<true>(22, 200); }
+
+}  // namespace esphome::api::testing