Merge pull request #9257 from tacaswell/fix_png_segfault

anntzer · web-flow · commit 080d0f890547 · 2017-11-13T11:35:04.000-08:00
FIX: segfault on truncated png
diff --git a/lib/matplotlib/tests/test_png.py b/lib/matplotlib/tests/test_png.py
@@ -2,16 +2,15 @@
                         unicode_literals)
 
 import six
-
+from six import BytesIO
 import glob
 import os
-
 import numpy as np
+import pytest
 
 from matplotlib.testing.decorators import image_comparison
 from matplotlib import pyplot as plt
 import matplotlib.cm as cm
-
 import sys
 on_win = (sys.platform == 'win32')
 
@@ -46,3 +45,28 @@ def test_imread_png_uint16():
 
     assert (img.dtype == np.uint16)
     assert np.sum(img.flatten()) == 134184960
+
+
+def test_truncated_file(tmpdir):
+    d = tmpdir.mkdir('test')
+    fname = str(d.join('test.png'))
+    fname_t = str(d.join('test_truncated.png'))
+    plt.savefig(fname)
+    with open(fname, 'rb') as fin:
+        buf = fin.read()
+    with open(fname_t, 'wb') as fout:
+        fout.write(buf[:20])
+
+    with pytest.raises(Exception):
+        plt.imread(fname_t)
+
+
+def test_truncated_buffer():
+    b = BytesIO()
+    plt.savefig(b)
+    b.seek(0)
+    b2 = BytesIO(b.read(20))
+    b2.seek(0)
+
+    with pytest.raises(Exception):
+        plt.imread(b2)
diff --git a/src/_png.cpp b/src/_png.cpp
@@ -194,18 +194,18 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
 
     switch (channels) {
     case 1:
-	png_color_type = PNG_COLOR_TYPE_GRAY;
-	break;
+        png_color_type = PNG_COLOR_TYPE_GRAY;
+        break;
     case 3:
-	png_color_type = PNG_COLOR_TYPE_RGB;
-	break;
+        png_color_type = PNG_COLOR_TYPE_RGB;
+        break;
     case 4:
-	png_color_type = PNG_COLOR_TYPE_RGB_ALPHA;
-	break;
+        png_color_type = PNG_COLOR_TYPE_RGB_ALPHA;
+        break;
     default:
         PyErr_SetString(PyExc_ValueError,
-			"Buffer must be an NxMxD array with D in 1, 3, 4 "
-			"(grayscale, RGB, RGBA)");
+                        "Buffer must be an NxMxD array with D in 1, 3, 4 "
+                        "(grayscale, RGB, RGBA)");
         goto exit;
     }
 
@@ -349,20 +349,20 @@ static PyObject *Py_write_png(PyObject *self, PyObject *args, PyObject *kwds)
     sig_bit.alpha = 0;
     switch (png_color_type) {
     case PNG_COLOR_TYPE_GRAY:
-	sig_bit.gray = 8;
-	sig_bit.red = 0;
-	sig_bit.green = 0;
-	sig_bit.blue = 0;
-	break;
+        sig_bit.gray = 8;
+        sig_bit.red = 0;
+        sig_bit.green = 0;
+        sig_bit.blue = 0;
+        break;
     case PNG_COLOR_TYPE_RGB_ALPHA:
-	sig_bit.alpha = 8;
-	// fall through
+        sig_bit.alpha = 8;
+        // fall through
     case PNG_COLOR_TYPE_RGB:
-	sig_bit.gray = 0;
-	sig_bit.red = 8;
-	sig_bit.green = 8;
-	sig_bit.blue = 8;
-	break;
+        sig_bit.gray = 0;
+        sig_bit.red = 8;
+        sig_bit.green = 8;
+        sig_bit.blue = 8;
+        break;
     default:
         PyErr_SetString(PyExc_RuntimeError, "internal error, bad png_color_type");
         goto exit;
@@ -408,13 +408,21 @@ static void _read_png_data(PyObject *py_file_obj, png_bytep data, png_size_t len
     Py_ssize_t bufflen;
     if (read_method) {
         result = PyObject_CallFunction(read_method, (char *)"i", length);
-        if (PyBytes_AsStringAndSize(result, &buffer, &bufflen) == 0) {
-            if (bufflen == (Py_ssize_t)length) {
-                memcpy(data, buffer, length);
+        if (result) {
+            if (PyBytes_AsStringAndSize(result, &buffer, &bufflen) == 0) {
+                if (bufflen == (Py_ssize_t)length) {
+                    memcpy(data, buffer, length);
+                } else {
+                    PyErr_SetString(PyExc_IOError, "read past end of file");
+                }
             } else {
-                PyErr_SetString(PyExc_IOError, "read past end of file");
+                PyErr_SetString(PyExc_IOError, "failed to copy buffer");
             }
+        } else  {
+            PyErr_SetString(PyExc_IOError, "failed to read file");
         }
+
+
     }
     Py_XDECREF(read_method);
     Py_XDECREF(result);
@@ -424,6 +432,10 @@ static void read_png_data(png_structp png_ptr, png_bytep data, png_size_t length
 {
     PyObject *py_file_obj = (PyObject *)png_get_io_ptr(png_ptr);
     _read_png_data(py_file_obj, data, length);
+    if (PyErr_Occurred()) {
+        png_error(png_ptr, "failed to read file");
+    }
+
 }
 
 static PyObject *_read_png(PyObject *filein, bool float_result)
@@ -481,6 +493,9 @@ static PyObject *_read_png(PyObject *filein, bool float_result)
         }
         Py_XDECREF(read_method);
         _read_png_data(py_file, header, 8);
+        if (PyErr_Occurred()) {
+            goto exit;
+        }
     }
 
     if (png_sig_cmp(header, 0, 8)) {
@@ -503,7 +518,9 @@ static PyObject *_read_png(PyObject *filein, bool float_result)
     }
 
     if (setjmp(png_jmpbuf(png_ptr))) {
-        PyErr_SetString(PyExc_RuntimeError, "Error setting jump");
+        if (!PyErr_Occurred()) {
+            PyErr_SetString(PyExc_RuntimeError, "error setting jump");
+        }
         goto exit;
     }
 
