Fixed #30565 -- Closed HttpResponse when wsgi.file_wrapper closes file-like object.

This commit is contained in:
Chris Jerdonek 2019-06-18 01:40:44 -07:00 committed by Carlton Gibson
parent 533311782f
commit cce47ff65a
3 changed files with 79 additions and 1 deletions

View File

@ -241,6 +241,9 @@ class HttpResponseBase:
# The WSGI server must call this method upon completion of the request. # The WSGI server must call this method upon completion of the request.
# See http://blog.dscpl.com.au/2012/10/obligations-for-calling-close-on.html # See http://blog.dscpl.com.au/2012/10/obligations-for-calling-close-on.html
# When wsgi.file_wrapper is used, the WSGI server instead calls close()
# on the file-like object. Django ensures this method is called in this
# case by replacing self.file_to_stream.close() with a wrapped version.
def close(self): def close(self):
for closable in self._closable_objects: for closable in self._closable_objects:
try: try:
@ -397,14 +400,39 @@ class FileResponse(StreamingHttpResponse):
self.filename = filename self.filename = filename
super().__init__(*args, **kwargs) super().__init__(*args, **kwargs)
def _wrap_file_to_stream_close(self, filelike):
"""
Wrap the file-like close() with a version that calls
FileResponse.close().
"""
closing = False
filelike_close = getattr(filelike, 'close', lambda: None)
def file_wrapper_close():
nonlocal closing
# Prevent an infinite loop since FileResponse.close() tries to
# close the objects in self._closable_objects.
if closing:
return
closing = True
try:
filelike_close()
finally:
self.close()
filelike.close = file_wrapper_close
def _set_streaming_content(self, value): def _set_streaming_content(self, value):
if not hasattr(value, 'read'): if not hasattr(value, 'read'):
self.file_to_stream = None self.file_to_stream = None
return super()._set_streaming_content(value) return super()._set_streaming_content(value)
self.file_to_stream = filelike = value self.file_to_stream = filelike = value
# Add to closable objects before wrapping close(), since the filelike
# might not have close().
if hasattr(filelike, 'close'): if hasattr(filelike, 'close'):
self._closable_objects.append(filelike) self._closable_objects.append(filelike)
self._wrap_file_to_stream_close(filelike)
value = iter(lambda: filelike.read(self.block_size), b'') value = iter(lambda: filelike.read(self.block_size), b'')
self.set_headers(filelike) self.set_headers(filelike)
super()._set_streaming_content(value) super()._set_streaming_content(value)

View File

@ -861,7 +861,8 @@ Methods
.. method:: HttpResponse.close() .. method:: HttpResponse.close()
This method is called at the end of the request directly by the WSGI This method is called at the end of the request directly by the WSGI
server. server, or when the WSGI server closes the file-like object, if
`wsgi.file_wrapper`_ is used for the request.
.. method:: HttpResponse.write(content) .. method:: HttpResponse.write(content)

View File

@ -80,3 +80,52 @@ class FileResponseTests(SimpleTestCase):
response['Content-Disposition'], response['Content-Disposition'],
"attachment; filename*=utf-8''%E7%A5%9D%E6%82%A8%E5%B9%B3%E5%AE%89.odt" "attachment; filename*=utf-8''%E7%A5%9D%E6%82%A8%E5%B9%B3%E5%AE%89.odt"
) )
def test_file_to_stream_closes_response(self):
# Closing file_to_stream calls FileResponse.close(), even when
# file-like object doesn't have a close() method.
class FileLike:
def read(self):
pass
class FileLikeWithClose(FileLike):
def __init__(self):
self.closed = False
def close(self):
self.closed = True
for filelike_cls in (FileLike, FileLikeWithClose):
with self.subTest(filelike_cls=filelike_cls.__name__):
filelike = filelike_cls()
response = FileResponse(filelike)
self.assertFalse(response.closed)
# Object with close() is added to the list of closable.
if hasattr(filelike, 'closed'):
self.assertEqual(response._closable_objects, [filelike])
else:
self.assertEqual(response._closable_objects, [])
file_to_stream = response.file_to_stream
file_to_stream.close()
if hasattr(filelike, 'closed'):
self.assertTrue(filelike.closed)
self.assertTrue(response.closed)
def test_file_to_stream_closes_response_on_error(self):
# Closing file_to_stream calls FileResponse.close(), even when
# closing file-like raises exceptions.
class FileLikeWithRaisingClose:
def read(self):
pass
def close(self):
raise RuntimeError()
filelike = FileLikeWithRaisingClose()
response = FileResponse(filelike)
self.assertFalse(response.closed)
self.assertEqual(response._closable_objects, [filelike])
file_to_stream = response.file_to_stream
with self.assertRaises(RuntimeError):
file_to_stream.close()
self.assertTrue(response.closed)