Merge pull request #1039 from kdmurray91/fix_freeimage_segfault

Fix freeimage segfault with a broken/truncated JPEG
This commit is contained in:
Stefan van der Walt
2014-06-30 15:22:29 +02:00
6 changed files with 69 additions and 13 deletions
Binary file not shown.

After

Width:  |  Height:  |  Size: 400 B

+49 -13
View File
@@ -34,15 +34,29 @@ def _generate_candidate_libs():
return lib_dirs, lib_paths
if sys.platform == 'win32':
LOADER = ctypes.windll
FUNCTYPE = ctypes.WINFUNCTYPE
else:
LOADER = ctypes.cdll
FUNCTYPE = ctypes.CFUNCTYPE
def handle_errors():
global FT_ERROR_STR
if FT_ERROR_STR:
tmp = FT_ERROR_STR
FT_ERROR_STR = None
raise RuntimeError(tmp)
FT_ERROR_STR = None
# This MUST happen in module scope, or the function pointer is garbage
# collected, leading to a segfault when error_handler is called.
@FUNCTYPE(None, ctypes.c_int, ctypes.c_char_p)
def c_error_handler(fif, message):
global FT_ERROR_STR
FT_ERROR_STR = 'FreeImage error: %s' % message
def load_freeimage():
if sys.platform == 'win32':
loader = ctypes.windll
functype = ctypes.WINFUNCTYPE
else:
loader = ctypes.cdll
functype = ctypes.CFUNCTYPE
freeimage = None
errors = []
# First try a few bare library names that ctypes might be able to find
@@ -54,7 +68,7 @@ def load_freeimage():
lib_paths = bare_libs + lib_paths
for lib in lib_paths:
try:
freeimage = loader.LoadLibrary(lib)
freeimage = LOADER.LoadLibrary(lib)
break
except Exception:
if lib not in bare_libs:
@@ -81,11 +95,7 @@ def load_freeimage():
'\n'.join(lib_dirs))
# FreeImage found
@functype(None, ctypes.c_int, ctypes.c_char_p)
def error_handler(fif, message):
raise RuntimeError('FreeImage error: %s' % message)
freeimage.FreeImage_SetOutputMessage(error_handler)
freeimage.FreeImage_SetOutputMessage(c_error_handler)
return freeimage
_FI = load_freeimage()
@@ -189,13 +199,17 @@ class FI_TYPES(object):
@classmethod
def get_type_and_shape(cls, bitmap):
w = _FI.FreeImage_GetWidth(bitmap)
handle_errors()
h = _FI.FreeImage_GetHeight(bitmap)
handle_errors()
fi_type = _FI.FreeImage_GetImageType(bitmap)
handle_errors()
if not fi_type:
raise ValueError('Unknown image pixel type')
dtype = cls.dtypes[fi_type]
if fi_type == cls.FIT_BITMAP:
bpp = _FI.FreeImage_GetBPP(bitmap)
handle_errors()
if bpp == 8:
extra_dims = []
elif bpp == 24:
@@ -377,9 +391,11 @@ class METADATA_DATATYPE(object):
def _process_bitmap(filename, flags, process_func):
filename = asbytes(filename)
ftype = _FI.FreeImage_GetFileType(filename, 0)
handle_errors()
if ftype == -1:
raise ValueError('Cannot determine type of file %s' % filename)
bitmap = _FI.FreeImage_Load(ftype, filename, flags)
handle_errors()
bitmap = ctypes.c_void_p(bitmap)
if not bitmap:
raise ValueError('Could not load file %s' % filename)
@@ -387,6 +403,7 @@ def _process_bitmap(filename, flags, process_func):
return process_func(bitmap)
finally:
_FI.FreeImage_Unload(bitmap)
handle_errors()
def read(filename, flags=0):
@@ -414,6 +431,7 @@ def read_metadata(filename):
def _process_multipage(filename, flags, process_func):
filename = asbytes(filename)
ftype = _FI.FreeImage_GetFileType(filename, 0)
handle_errors()
if ftype == -1:
raise ValueError('Cannot determine type of file %s' % filename)
create_new = False
@@ -422,14 +440,17 @@ def _process_multipage(filename, flags, process_func):
multibitmap = _FI.FreeImage_OpenMultiBitmap(ftype, filename, create_new,
read_only, keep_cache_in_memory,
flags)
handle_errors()
multibitmap = ctypes.c_void_p(multibitmap)
if not multibitmap:
raise ValueError('Could not open %s as multi-page image.' % filename)
try:
pages = _FI.FreeImage_GetPageCount(multibitmap)
handle_errors()
out = []
for i in range(pages):
bitmap = _FI.FreeImage_LockPage(multibitmap, i)
handle_errors()
bitmap = ctypes.c_void_p(bitmap)
if not bitmap:
raise ValueError('Could not open %s as a multi-page image.'
@@ -438,9 +459,11 @@ def _process_multipage(filename, flags, process_func):
out.append(process_func(bitmap))
finally:
_FI.FreeImage_UnlockPage(multibitmap, bitmap, False)
handle_errors()
return out
finally:
_FI.FreeImage_CloseMultiBitmap(multibitmap, 0)
handle_errors()
def read_multipage(filename, flags=0):
@@ -469,6 +492,7 @@ def _wrap_bitmap_bits_in_array(bitmap, shape, dtype):
"""
pitch = _FI.FreeImage_GetPitch(bitmap)
handle_errors()
height = shape[-1]
byte_size = height * pitch
itemsize = dtype.itemsize
@@ -478,6 +502,7 @@ def _wrap_bitmap_bits_in_array(bitmap, shape, dtype):
else:
strides = (itemsize, pitch)
bits = _FI.FreeImage_GetBits(bitmap)
handle_errors()
array = numpy.ndarray(shape, dtype=dtype,
buffer=(ctypes.c_char * byte_size).from_address(bits),
strides=strides)
@@ -502,6 +527,7 @@ def _array_from_bitmap(bitmap):
g = n(array[1])
r = n(array[2])
if shape[0] == 3:
handle_errors()
return numpy.dstack((r, g, b))
elif shape[0] == 4:
a = n(array[3])
@@ -523,6 +549,7 @@ def _read_metadata(bitmap):
for model_name, number in models:
mdhandle = _FI.FreeImage_FindFirstMetadata(number, bitmap,
ctypes.byref(tag))
handle_errors()
mdhandle = ctypes.c_void_p(mdhandle)
if mdhandle:
more = True
@@ -530,8 +557,10 @@ def _read_metadata(bitmap):
tag_name = asstr(_FI.FreeImage_GetTagKey(tag))
tag_type = _FI.FreeImage_GetTagType(tag)
byte_size = _FI.FreeImage_GetTagLength(tag)
handle_errors()
char_ptr = ctypes.c_char * byte_size
tag_str = char_ptr.from_address(_FI.FreeImage_GetTagValue(tag))
handle_errors()
if tag_type == METADATA_DATATYPE.FIDT_ASCII:
tag_val = asstr(tag_str.value)
else:
@@ -541,7 +570,9 @@ def _read_metadata(bitmap):
tag_val = tag_val[0]
metadata[(model_name, tag_name)] = tag_val
more = _FI.FreeImage_FindNextMetadata(mdhandle, ctypes.byref(tag))
handle_errors()
_FI.FreeImage_FindCloseMetadata(mdhandle)
handle_errors()
return metadata
@@ -556,6 +587,7 @@ def write(array, filename, flags=0):
array = numpy.asarray(array)
filename = asbytes(filename)
ftype = _FI.FreeImage_GetFIFFromFilename(filename)
handle_errors()
if ftype == -1:
raise ValueError('Cannot determine type for %s' % filename)
bitmap, fi_type = _array_to_bitmap(array)
@@ -563,16 +595,20 @@ def write(array, filename, flags=0):
if fi_type == FI_TYPES.FIT_BITMAP:
can_write = _FI.FreeImage_FIFSupportsExportBPP(ftype,
_FI.FreeImage_GetBPP(bitmap))
handle_errors()
else:
can_write = _FI.FreeImage_FIFSupportsExportType(ftype, fi_type)
handle_errors()
if not can_write:
raise TypeError('Cannot save image of this format '
'to this file type')
res = _FI.FreeImage_Save(ftype, bitmap, filename, flags)
handle_errors()
if not res:
raise RuntimeError('Could not save image properly.')
finally:
_FI.FreeImage_Unload(bitmap)
handle_errors()
def write_multipage(arrays, filename, flags=0):
+5
View File
@@ -37,6 +37,11 @@ def test_imread():
assert img.shape == (370, 371, 3)
assert all(img[274, 135] == [0, 130, 253])
@skipif(not FI_available)
def test_imread_truncated_jpg():
assert_raises((RuntimeError, ValueError),
sio.imread,
os.path.join(si.data_dir, 'truncated.jpg'))
@skipif(not FI_available)
def test_imread_uint16():
+5
View File
@@ -37,6 +37,11 @@ def test_imread_palette():
img = imread(os.path.join(data_dir, 'palette_color.png'))
assert img.ndim == 3
@skipif(not imread_available)
def test_imread_truncated_jpg():
assert_raises((RuntimeError, ValueError),
sio.imread,
os.path.join(si.data_dir, 'truncated.jpg'))
@skipif(not imread_available)
def test_bilevel():
+4
View File
@@ -95,6 +95,10 @@ def test_repr_png():
assert np.all(original_img == round_trip)
@skipif(not PIL_available)
def test_imread_truncated_jpg():
assert_raises((IOError, ValueError), imread,
os.path.join(data_dir, 'truncated.jpg'))
# Big endian images not correctly loaded for PIL < 1.1.7
# Renable test when PIL 1.1.7 is more common.
+6
View File
@@ -51,6 +51,12 @@ def test_bilevel():
np.testing.assert_array_equal(img, expected)
@skipif(not sitk_available)
def test_imread_truncated_jpg():
assert_raises((RuntimeError, ValueError),
sio.imread,
os.path.join(si.data_dir, 'truncated.jpg'))
@skipif(not sitk_available)
def test_imread_uint16():
expected = np.load(os.path.join(data_dir, 'chessboard_GRAY_U8.npy'))