[runtime_image] Allow decoding into a caller-owned buffer (#17936)

esphome/components/runtime_image/runtime_image.cpp

@@ -248,9 +248,15 @@ void RuntimeImage::release() {
 
 void RuntimeImage::release_buffer_() {
   if (this->buffer_) {
-    ESP_LOGV(TAG, "Releasing buffer of size %zu", this->get_buffer_size_(this->buffer_width_, this->buffer_height_));
-    RAMAllocator<uint8_t> allocator;
-    allocator.deallocate(this->buffer_, this->get_buffer_size_(this->buffer_width_, this->buffer_height_));
+    if (this->external_buffer_) {
+      // The caller owns this memory and goes on using it after the image lets go of it.
+      ESP_LOGV(TAG, "Letting go of the external %dx%d buffer", this->buffer_width_, this->buffer_height_);
+      this->external_buffer_ = false;
+    } else {
+      ESP_LOGV(TAG, "Releasing buffer of size %zu", this->get_buffer_size(this->buffer_width_, this->buffer_height_));
+      RAMAllocator<uint8_t> allocator;
+      allocator.deallocate(this->buffer_, this->get_buffer_size(this->buffer_width_, this->buffer_height_));
+    }
     this->buffer_ = nullptr;
     this->data_start_ = nullptr;
     this->width_ = 0;
@@ -263,19 +269,46 @@ void RuntimeImage::release_buffer_() {
   }
 }
 
+bool RuntimeImage::set_external_buffer(uint8_t *buffer, int width, int height) {
+  this->release_buffer_();
+  if (buffer == nullptr || this->get_buffer_size(width, height) == 0) {
+    // Keep the released state rather than remembering a buffer that cannot be decoded into: an
+    // external buffer that is never handed back would otherwise block every later allocation.
+    ESP_LOGE(TAG, "Refusing an invalid external buffer for %dx%d", width, height);
+    return false;
+  }
+  this->buffer_ = buffer;
+  this->external_buffer_ = true;
+  this->buffer_width_ = width;
+  this->buffer_height_ = height;
+  return true;
+}
+
 size_t RuntimeImage::resize_buffer_(int width, int height) {
-  size_t new_size = this->get_buffer_size_(width, height);
+  size_t new_size = this->get_buffer_size(width, height);
+
+  // A buffer only ever exists with dimensions the image can decode at, so a match here means
+  // new_size is non-zero. Checking it before the invalid dimension case below lets the external
+  // buffer be let go of for every decode it cannot serve, not just for valid other dimensions.
+  if (this->buffer_ && this->buffer_width_ == width && this->buffer_height_ == height) {
+    // Buffer already allocated with correct size
+    return new_size;
+  }
+
+  if (this->external_buffer_) {
+    ESP_LOGE(TAG, "Image decoded to %dx%d, but the external buffer is %dx%d", width, height, this->buffer_width_,
+             this->buffer_height_);
+    // Let the buffer go rather than free memory that belongs to the caller. Dropping it also stops
+    // a decoder that ignores this failure from publishing a picture it never painted.
+    this->release_buffer_();
+    return 0;
+  }
 
   if (new_size == 0) {
     ESP_LOGE(TAG, "Refusing to allocate buffer for invalid image dimensions %dx%d", width, height);
     return 0;
   }
 
-  if (this->buffer_ && this->buffer_width_ == width && this->buffer_height_ == height) {
-    // Buffer already allocated with correct size
-    return new_size;
-  }
-
   // Release old buffer if dimensions changed
   if (this->buffer_) {
     this->release_buffer_();
@@ -300,7 +333,7 @@ size_t RuntimeImage::resize_buffer_(int width, int height) {
   return new_size;
 }
 
-size_t RuntimeImage::get_buffer_size_(int width, int height) const {
+size_t RuntimeImage::get_buffer_size(int width, int height) const {
   // Dimensions come from a remote image header; reject absurd values so the size math cannot overflow
   if (width <= 0 || height <= 0 || width > MAX_IMAGE_DIMENSION || height > MAX_IMAGE_DIMENSION) {
     return 0;

esphome/components/runtime_image/runtime_image.h

@@ -121,9 +121,48 @@ class RuntimeImage : public image::Image {
 
   /**
    * @brief Release the image buffer and free memory.
+   *
+   * An external buffer is let go of rather than freed.
    */
   void release();
 
+  /**
+   * @brief Decode into a buffer the caller owns, instead of one allocated here.
+   *
+   * The image never frees an external buffer and never resizes it: a decode that needs other
+   * dimensions fails as if the allocation had failed, and the buffer is let go of so a decoder
+   * that ignores that failure cannot publish a picture it did not paint. The caller keeps the
+   * buffer alive for as long as anything can draw the image, and calls release() (or hands over
+   * another buffer) before reusing it.
+   *
+   * Hand a buffer over before every decode. The image lets go of one whenever a decode fails and
+   * whenever release() is called, and it does not remember that it ever had one: a decode that
+   * starts without a buffer allocates its own, which is the runtime allocation this method exists
+   * to avoid.
+   *
+   * The buffer is decoded into as it is handed over, so the caller owns its initial contents.
+   * Zero it first if anything can draw the image before a decode has painted every pixel.
+   *
+   * Do not hand a buffer over while is_decoding() is true. A running decoder keeps scaling values
+   * for the buffer it started with.
+   *
+   * A null buffer or dimensions the image cannot decode at are refused, leaving the image with
+   * no buffer at all.
+   *
+   * @param buffer Memory for a picture of the given size, at least get_buffer_size() bytes.
+   * @param width Width of the buffer in pixels.
+   * @param height Height of the buffer in pixels.
+   * @return true if the image took the buffer, false if it was refused.
+   */
+  bool set_external_buffer(uint8_t *buffer, int width, int height);
+
+  /**
+   * @brief Get the buffer size in bytes needed for a picture of the given dimensions.
+   *
+   * Returns 0 for dimensions the image cannot decode at.
+   */
+  size_t get_buffer_size(int width, int height) const;
+
   /**
    * @brief Set whether to allow progressive display during decode.
    *
@@ -149,11 +188,6 @@ class RuntimeImage : public image::Image {
    */
   void release_buffer_();
 
-  /**
-   * @brief Get the buffer size in bytes for given dimensions.
-   */
-  size_t get_buffer_size_(int width, int height) const;
-
   /**
    * @brief Get the position in the buffer for a pixel.
    */
@@ -208,6 +242,8 @@ class RuntimeImage : public image::Image {
    * This is used to determine how to store 16 bit colors in the buffer.
    */
   bool is_big_endian_{false};
+  /** Whether buffer_ belongs to the caller, so it must not be freed or resized here. */
+  bool external_buffer_{false};
 };
 
 }  // namespace esphome::runtime_image