[online_image] Improve error handling (#8212)

esphome/components/online_image/image_decoder.cpp

@@ -25,6 +25,15 @@ void ImageDecoder::draw(int x, int y, int w, int h, const Color &color) {
   }
 }
 
+DownloadBuffer::DownloadBuffer(size_t size) : size_(size) {
+  this->buffer_ = this->allocator_.allocate(size);
+  this->reset();
+  if (!this->buffer_) {
+    ESP_LOGE(TAG, "Initial allocation of download buffer failed!");
+    this->size_ = 0;
+  }
+}
+
 uint8_t *DownloadBuffer::data(size_t offset) {
   if (offset > this->size_) {
     ESP_LOGE(TAG, "Tried to access beyond download buffer bounds!!!");
@@ -46,12 +55,13 @@ size_t DownloadBuffer::resize(size_t size) {
     return size;
   }
   this->allocator_.deallocate(this->buffer_, this->size_);
-  this->size_ = size;
   this->buffer_ = this->allocator_.allocate(size);
   this->reset();
   if (this->buffer_) {
+    this->size_ = size;
     return size;
   } else {
+    this->size_ = 0;
     return 0;
   }
 }

esphome/components/online_image/image_decoder.h

@@ -29,8 +29,12 @@ class ImageDecoder {
    * @brief Initialize the decoder.
    *
    * @param download_size The total number of bytes that need to be downloaded for the image.
+   * @return int          Returns 0 on success, a {@see DecodeError} value in case of an error.
    */
-  virtual void prepare(size_t download_size) { this->download_size_ = download_size; }
+  virtual int prepare(size_t download_size) {
+    this->download_size_ = download_size;
+    return 0;
+  }
 
   /**
    * @brief Decode a part of the image. It will try reading from the buffer.
@@ -83,10 +87,7 @@ class ImageDecoder {
 
 class DownloadBuffer {
  public:
-  DownloadBuffer(size_t size) : size_(size) {
-    this->buffer_ = this->allocator_.allocate(size);
-    this->reset();
-  }
+  DownloadBuffer(size_t size);
 
   virtual ~DownloadBuffer() { this->allocator_.deallocate(this->buffer_, this->size_); }
 

esphome/components/online_image/jpeg_image.cpp

@@ -41,13 +41,14 @@ static int draw_callback(JPEGDRAW *jpeg) {
   return 1;
 }
 
-void JpegDecoder::prepare(size_t download_size) {
+int JpegDecoder::prepare(size_t download_size) {
   ImageDecoder::prepare(download_size);
   auto size = this->image_->resize_download_buffer(download_size);
   if (size < download_size) {
-    ESP_LOGE(TAG, "Resize failed!");
-    // TODO: return an error code;
+    ESP_LOGE(TAG, "Download buffer resize failed!");
+    return DECODE_ERROR_OUT_OF_MEMORY;
   }
+  return 0;
 }
 
 int HOT JpegDecoder::decode(uint8_t *buffer, size_t size) {

esphome/components/online_image/jpeg_image.h

@@ -21,7 +21,7 @@ class JpegDecoder : public ImageDecoder {
   JpegDecoder(OnlineImage *image) : ImageDecoder(image) {}
   ~JpegDecoder() override {}
 
-  void prepare(size_t download_size) override;
+  int prepare(size_t download_size) override;
   int HOT decode(uint8_t *buffer, size_t size) override;
 
  protected:

esphome/components/online_image/online_image.cpp

@@ -178,7 +178,12 @@ void OnlineImage::update() {
     this->download_error_callback_.call();
     return;
   }
-  this->decoder_->prepare(total_size);
+  auto prepare_result = this->decoder_->prepare(total_size);
+  if (prepare_result < 0) {
+    this->end_connection_();
+    this->download_error_callback_.call();
+    return;
+  }
   ESP_LOGI(TAG, "Downloading image (Size: %d)", total_size);
   this->start_time_ = ::time(nullptr);
 }

esphome/components/online_image/png_image.cpp

@@ -40,11 +40,16 @@ static void draw_callback(pngle_t *pngle, uint32_t x, uint32_t y, uint32_t w, ui
   decoder->draw(x, y, w, h, color);
 }
 
-void PngDecoder::prepare(size_t download_size) {
+int PngDecoder::prepare(size_t download_size) {
   ImageDecoder::prepare(download_size);
+  if (!this->pngle_) {
+    ESP_LOGE(TAG, "PNG decoder engine not initialized!");
+    return DECODE_ERROR_OUT_OF_MEMORY;
+  }
   pngle_set_user_data(this->pngle_, this);
   pngle_set_init_callback(this->pngle_, init_callback);
   pngle_set_draw_callback(this->pngle_, draw_callback);
+  return 0;
 }
 
 int HOT PngDecoder::decode(uint8_t *buffer, size_t size) {

esphome/components/online_image/png_image.h

@@ -21,7 +21,7 @@ class PngDecoder : public ImageDecoder {
   PngDecoder(OnlineImage *image) : ImageDecoder(image), pngle_(pngle_new()) {}
   ~PngDecoder() override { pngle_destroy(this->pngle_); }
 
-  void prepare(size_t download_size) override;
+  int prepare(size_t download_size) override;
   int HOT decode(uint8_t *buffer, size_t size) override;
 
  protected: