1 files changed, 21 insertions, 25 deletions

diff --git a/src/libslic3r/PNGRead.cpp b/src/libslic3r/PNGRead.cpp
index 5b5b9ffec..e66143b84 100644
--- a/src/libslic3r/PNGRead.cpp
+++ b/src/libslic3r/PNGRead.cpp
@@ -41,13 +41,6 @@ bool is_png(const ReadBuf &rb)
     return rb.sz >= PNG_SIG_BYTES && !png_sig_cmp(buf, 0, PNG_SIG_BYTES);
 }
 
-// A wrapper around ReadBuf to be read repeatedly like a stream. libpng needs
-// this form for its buffer read callback.
-struct ReadBufReader {
-    const ReadBuf &rdbuf; size_t pos;
-    ReadBufReader(const ReadBuf &rd): rdbuf{rd}, pos{0} {}
-};
-
 // Buffer read callback for libpng. It provides an allocated output buffer and
 // the amount of data it desires to read from the input.
 void png_read_callback(png_struct *png_ptr,
@@ -55,20 +48,21 @@ void png_read_callback(png_struct *png_ptr,
                        png_size_t  byteCountToRead)
 {
     // Retrieve our input buffer through the png_ptr
-    auto reader = static_cast<ReadBufReader *>(png_get_io_ptr(png_ptr));
+    auto reader = static_cast<IStream *>(png_get_io_ptr(png_ptr));
 
-    if (!reader || byteCountToRead > reader->rdbuf.sz - reader->pos) return;
+    if (!reader || !reader->is_ok()) return;
 
-    auto   buf = static_cast<const png_byte *>(reader->rdbuf.buf);
-    size_t pos = reader->pos;
-
-    std::copy(buf + pos, buf + (pos + byteCountToRead), outBytes);
-    reader->pos += byteCountToRead;
+    reader->read(static_cast<std::uint8_t *>(outBytes), byteCountToRead);
 }
 
-bool decode_png(const ReadBuf &rb, ImageGreyscale &img)
+bool decode_png(IStream &in_buf, ImageGreyscale &out_img)
 {
-    if (!is_png(rb)) return false;
+    static const constexpr int PNG_SIG_BYTES = 8;
+
+    std::vector<png_byte> sig(PNG_SIG_BYTES, 0);
+    in_buf.read(sig.data(), PNG_SIG_BYTES);
+    if (!png_check_sig(sig.data(), PNG_SIG_BYTES))
+        return false;
 
     PNGDescr dsc;
     dsc.png = png_create_read_struct(PNG_LIBPNG_VER_STRING, nullptr, nullptr,
@@ -77,26 +71,28 @@ bool decode_png(const ReadBuf &rb, ImageGreyscale &img)
     if(!dsc.png) return false;
 
     dsc.info = png_create_info_struct(dsc.png);
-    if(!dsc.info) return {};
+    if(!dsc.info) return false;
+
+    png_set_read_fn(dsc.png, static_cast<void *>(&in_buf), png_read_callback);
 
-    ReadBufReader reader {rb};
-    png_set_read_fn(dsc.png, static_cast<void *>(&reader), png_read_callback);
+    // Tell that we have already read the first bytes to check the signature
+    png_set_sig_bytes(dsc.png, PNG_SIG_BYTES);
 
     png_read_info(dsc.png, dsc.info);
 
-    img.cols = png_get_image_width(dsc.png, dsc.info);
-    img.rows = png_get_image_height(dsc.png, dsc.info);
+    out_img.cols = png_get_image_width(dsc.png, dsc.info);
+    out_img.rows = png_get_image_height(dsc.png, dsc.info);
     size_t color_type = png_get_color_type(dsc.png, dsc.info);
     size_t bit_depth  = png_get_bit_depth(dsc.png, dsc.info);
 
     if (color_type != PNG_COLOR_TYPE_GRAY || bit_depth != 8)
         return false;
 
-    img.buf.resize(img.rows * img.cols);
+    out_img.buf.resize(out_img.rows * out_img.cols);
 
-    auto readbuf = static_cast<png_bytep>(img.buf.data());
-    for (size_t r = 0; r < img.rows; ++r)
-        png_read_row(dsc.png, readbuf + r * img.cols, nullptr);
+    auto readbuf = static_cast<png_bytep>(out_img.buf.data());
+    for (size_t r = 0; r < out_img.rows; ++r)
+        png_read_row(dsc.png, readbuf + r * out_img.cols, nullptr);
 
     return true;
 }