Merge pull request #12353 from alalek:imgcodecs_fix_webp

modules/imgcodecs/src/grfmt_webp.cpp

@@ -54,15 +54,21 @@
 
 #include "opencv2/imgproc.hpp"
 
-const size_t WEBP_HEADER_SIZE = 32;
+#include <opencv2/core/utils/configuration.private.hpp>
 
 namespace cv
 {
 
+// 64Mb limit to avoid memory DDOS
+static size_t param_maxFileSize = utils::getConfigurationParameterSizeT("OPENCV_IMGCODECS_WEBP_MAX_FILE_SIZE", 64*1024*1024);
+
+static const size_t WEBP_HEADER_SIZE = 32;
+
 WebPDecoder::WebPDecoder()
 {
     m_buf_supported = true;
     channels = 0;
+    fs_size = 0;
 }
 
 WebPDecoder::~WebPDecoder() {}
@@ -96,48 +102,29 @@ ImageDecoder WebPDecoder::newDecoder() const
 
 bool WebPDecoder::readHeader()
 {
+    uint8_t header[WEBP_HEADER_SIZE] = { 0 };
     if (m_buf.empty())
     {
-        FILE * wfile = NULL;
-
-        wfile = fopen(m_filename.c_str(), "rb");
-
-        if(wfile == NULL)
-        {
-            return false;
-        }
-
-        fseek(wfile, 0, SEEK_END);
-        long int wfile_size = ftell(wfile);
-        fseek(wfile, 0, SEEK_SET);
-
-        if(wfile_size > static_cast<long int>(INT_MAX))
-        {
-            fclose(wfile);
-            return false;
-        }
-
-        data.create(1, (int)wfile_size, CV_8U);
-
-        size_t data_size = fread(data.ptr(), 1, wfile_size, wfile);
-
-        if(wfile)
-        {
-            fclose(wfile);
-        }
-
-        if(static_cast<long int>(data_size) != wfile_size)
-        {
-            return false;
-        }
+        fs.open(m_filename.c_str(), std::ios::binary);
+        fs.seekg(0, std::ios::end);
+        fs_size = fs.tellg();
+        fs.seekg(0, std::ios::beg);
+        CV_Assert(fs && "File stream error");
+        CV_CheckGE(fs_size, WEBP_HEADER_SIZE, "File is too small");
+        CV_CheckLE(fs_size, param_maxFileSize, "File is too large. Increase OPENCV_IMGCODECS_WEBP_MAX_FILE_SIZE parameter if you want to process large files");
+
+        fs.read((char*)header, sizeof(header));
+        CV_Assert(fs && "Can't read WEBP_HEADER_SIZE bytes");
     }
     else
     {
+        CV_CheckGE(m_buf.total(), WEBP_HEADER_SIZE, "Buffer is too small");
+        memcpy(header, m_buf.ptr(), sizeof(header));
         data = m_buf;
     }
 
     WebPBitstreamFeatures features;
-    if(VP8_STATUS_OK == WebPGetFeatures(data.ptr(), WEBP_HEADER_SIZE, &features))
+    if (VP8_STATUS_OK == WebPGetFeatures(header, sizeof(header), &features))
     {
         m_width  = features.width;
         m_height = features.height;
@@ -161,41 +148,75 @@ bool WebPDecoder::readHeader()
 
 bool WebPDecoder::readData(Mat &img)
 {
-    if( m_width > 0 && m_height > 0 )
+    CV_CheckGE(m_width, 0, ""); CV_CheckGE(m_height, 0, "");
+
+    CV_CheckEQ(img.cols, m_width, "");
+    CV_CheckEQ(img.rows, m_height, "");
+
+    if (m_buf.empty())
     {
-        bool convert_grayscale = (img.type() == CV_8UC1); // IMREAD_GRAYSCALE requested
+        fs.seekg(0, std::ios::beg); CV_Assert(fs && "File stream error");
+        data.create(1, validateToInt(fs_size), CV_8UC1);
+        fs.read((char*)data.ptr(), fs_size);
+        CV_Assert(fs && "Can't read file data");
+        fs.close();
+    }
+    CV_Assert(data.type() == CV_8UC1); CV_Assert(data.rows == 1);
 
-        if (img.cols != m_width || img.rows != m_height || img.type() != m_type)
+    {
+        Mat read_img;
+        CV_CheckType(img.type(), img.type() == CV_8UC1 || img.type() == CV_8UC3 || img.type() == CV_8UC4, "");
+        if (img.type() != m_type)
+        {
+            read_img.create(m_height, m_width, m_type);
+        }
+        else
         {
-            img.create(m_height, m_width, m_type);
+            read_img = img;  // copy header
         }
 
-        uchar* out_data = img.ptr();
-        size_t out_data_size = img.cols * img.rows * img.elemSize();
+        uchar* out_data = read_img.ptr();
+        size_t out_data_size = read_img.dataend - out_data;
 
-        uchar *res_ptr = 0;
+        uchar *res_ptr = NULL;
         if (channels == 3)
         {
+            CV_CheckTypeEQ(read_img.type(), CV_8UC3, "");
             res_ptr = WebPDecodeBGRInto(data.ptr(), data.total(), out_data,
-                                        (int)out_data_size, (int)img.step);
+                                        (int)out_data_size, (int)read_img.step);
         }
         else if (channels == 4)
         {
+            CV_CheckTypeEQ(read_img.type(), CV_8UC4, "");
             res_ptr = WebPDecodeBGRAInto(data.ptr(), data.total(), out_data,
-                                         (int)out_data_size, (int)img.step);
+                                         (int)out_data_size, (int)read_img.step);
         }
 
-        if(res_ptr == out_data)
+        if (res_ptr != out_data)
+            return false;
+
+        if (read_img.data == img.data && img.type() == m_type)
         {
-            if (convert_grayscale)
-            {
-                cvtColor(img, img, COLOR_BGR2GRAY);
-            }
-            return true;
+            // nothing
+        }
+        else if (img.type() == CV_8UC1)
+        {
+            cvtColor(read_img, img, COLOR_BGR2GRAY);
+        }
+        else if (img.type() == CV_8UC3 && m_type == CV_8UC4)
+        {
+            cvtColor(read_img, img, COLOR_BGRA2BGR);
+        }
+        else if (img.type() == CV_8UC3 && m_type == CV_8UC4)
+        {
+            cvtColor(read_img, img, COLOR_BGRA2BGR);
+        }
+        else
+        {
+            CV_Error(Error::StsInternal, "");
         }
     }
-
-    return false;
+    return true;
 }
 
 WebPEncoder::WebPEncoder()
@@ -213,12 +234,9 @@ ImageEncoder WebPEncoder::newEncoder() const
 
 bool WebPEncoder::write(const Mat& img, const std::vector<int>& params)
 {
-    int channels = img.channels(), depth = img.depth();
-    int width = img.cols, height = img.rows;
+    CV_CheckDepthEQ(img.depth(), CV_8U, "WebP codec supports 8U images only");
 
-    const Mat *image = &img;
-    Mat temp;
-    size_t size = 0;
+    const int width = img.cols, height = img.rows;
 
     bool comp_lossless = true;
     float quality = 100.0f;
@@ -240,69 +258,64 @@ bool WebPEncoder::write(const Mat& img, const std::vector<int>& params)
         }
     }
 
-    uint8_t *out = NULL;
+    int channels = img.channels();
+    CV_Check(channels, channels == 1 || channels == 3 || channels == 4, "");
 
-    if(depth != CV_8U)
-    {
-        return false;
-    }
+    const Mat *image = &img;
+    Mat temp;
 
-    if(channels == 1)
+    if (channels == 1)
     {
         cvtColor(*image, temp, CV_GRAY2BGR);
         image = &temp;
         channels = 3;
     }
-    else if (channels == 2)
-    {
-        return false;
-    }
 
+    uint8_t *out = NULL;
+    size_t size = 0;
     if (comp_lossless)
     {
-        if(channels == 3)
+        if (channels == 3)
         {
             size = WebPEncodeLosslessBGR(image->ptr(), width, height, (int)image->step, &out);
         }
-        else if(channels == 4)
+        else if (channels == 4)
         {
             size = WebPEncodeLosslessBGRA(image->ptr(), width, height, (int)image->step, &out);
         }
     }
     else
     {
-        if(channels == 3)
+        if (channels == 3)
         {
             size = WebPEncodeBGR(image->ptr(), width, height, (int)image->step, quality, &out);
         }
-        else if(channels == 4)
+        else if (channels == 4)
         {
             size = WebPEncodeBGRA(image->ptr(), width, height, (int)image->step, quality, &out);
         }
     }
+#if WEBP_DECODER_ABI_VERSION >= 0x0206
+    Ptr<uint8_t> out_cleaner(out, WebPFree);
+#else
+    Ptr<uint8_t> out_cleaner(out, free);
+#endif
+
+    CV_Assert(size > 0);
 
-    if(size > 0)
+    if (m_buf)
     {
-        if(m_buf)
-        {
-            m_buf->resize(size);
-            memcpy(&(*m_buf)[0], out, size);
-        }
-        else
-        {
-            FILE *fd = fopen(m_filename.c_str(), "wb");
-            if(fd != NULL)
-            {
-                fwrite(out, size, sizeof(uint8_t), fd);
-                fclose(fd); fd = NULL;
-            }
-        }
+        m_buf->resize(size);
+        memcpy(&(*m_buf)[0], out, size);
     }
-
-    if (out != NULL)
+    else
     {
-        free(out);
-        out = NULL;
+        FILE *fd = fopen(m_filename.c_str(), "wb");
+        if (fd != NULL)
+        {
+            fwrite(out, size, sizeof(uint8_t), fd);
+            fclose(fd); fd = NULL;
+        }
     }
 
     return size > 0;

modules/imgcodecs/src/grfmt_webp.hpp

@@ -47,7 +47,7 @@
 
 #ifdef HAVE_WEBP
 
-
+#include <fstream>
 
 namespace cv
 {
@@ -61,7 +61,6 @@ public:
 
     bool readData( Mat& img ) CV_OVERRIDE;
     bool readHeader() CV_OVERRIDE;
-    void close();
 
     size_t signatureLength() const CV_OVERRIDE;
     bool checkSignature( const String& signature) const CV_OVERRIDE;
@@ -69,6 +68,8 @@ public:
     ImageDecoder newDecoder() const CV_OVERRIDE;
 
 protected:
+    std::ifstream fs;
+    size_t fs_size;
     Mat data;
     int channels;
 };

modules/imgcodecs/src/loadsave.cpp

@@ -709,11 +709,22 @@ static bool imwrite_( const String& filename, const std::vector<Mat>& img_vec,
 
     encoder->setDestination( filename );
     CV_Assert(params.size() <= CV_IO_MAX_IMAGE_PARAMS*2);
-    bool code;
-    if (!isMultiImg)
-        code = encoder->write( write_vec[0], params );
-    else
-        code = encoder->writemulti( write_vec, params ); //to be implemented
+    bool code = false;
+    try
+    {
+        if (!isMultiImg)
+            code = encoder->write( write_vec[0], params );
+        else
+            code = encoder->writemulti( write_vec, params ); //to be implemented
+    }
+    catch (const cv::Exception& e)
+    {
+        std::cerr << "imwrite_('" << filename << "'): can't write data: " << e.what() << std::endl << std::flush;
+    }
+    catch (...)
+    {
+        std::cerr << "imwrite_('" << filename << "'): can't write data: unknown exception" << std::endl << std::flush;
+    }
 
     //    CV_Assert( code );
     return code;

modules/imgcodecs/test/test_webp.cpp

@@ -96,12 +96,17 @@ TEST(Imgcodecs_WebP, encode_decode_with_alpha_webp)
     string output = cv::tempfile(".webp");
 
     EXPECT_NO_THROW(cv::imwrite(output, img));
-    cv::Mat img_webp = cv::imread(output);
+    cv::Mat img_webp = cv::imread(output, IMREAD_UNCHANGED);
+    cv::Mat img_webp_bgr = cv::imread(output); // IMREAD_COLOR by default
     EXPECT_EQ(0, remove(output.c_str()));
     EXPECT_FALSE(img_webp.empty());
     EXPECT_EQ(4,   img_webp.channels());
     EXPECT_EQ(512, img_webp.cols);
     EXPECT_EQ(512, img_webp.rows);
+    EXPECT_FALSE(img_webp_bgr.empty());
+    EXPECT_EQ(3,   img_webp_bgr.channels());
+    EXPECT_EQ(512, img_webp_bgr.cols);
+    EXPECT_EQ(512, img_webp_bgr.rows);
 }
 
 #endif // HAVE_WEBP