diff --git a/lib/matplotlib/tests/test_png.py b/lib/matplotlib/tests/test_png.py index 9047eed846c7..9e5dfae06dae 100644 --- 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 index 06e4b87543f2..5a6a46c37332 100644 --- 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; }