diff --git a/tests/server/api/v1/test_metadata_mixin.py b/tests/server/api/v1/test_metadata_mixin.py index 9b7a83002..5e51a827c 100644 --- a/tests/server/api/v1/test_metadata_mixin.py +++ b/tests/server/api/v1/test_metadata_mixin.py @@ -13,6 +13,7 @@ mock_folder_children_provider_s3, auth, settings, credentials) from waterbutler.providers.s3 import S3Provider from waterbutler.providers.s3compat import S3CompatProvider +from tests.core.streams.fixtures import mock_response_stream_reader @pytest.fixture @@ -307,3 +308,54 @@ async def test_download_folder_as_zip_root(self, http_request, mock_stream): assert handler._headers['Content-Disposition'] == expected handler.write_stream.assert_called_once_with(mock_stream) + + @pytest.mark.asyncio + async def test_download_file_header_chunked_with_response_stream_reader(self, http_request, mock_response_stream_reader): + + handler = mock_handler(http_request) + mock_response_stream_reader.response.headers['Transfer-Encoding'] = 'chunked' + handler.provider.download = MockCoroutine(return_value=mock_response_stream_reader) + handler.path = WaterButlerPath('/test_file') + handler.bytes_downloaded = '100' + + await handler.download_file() + assert handler._headers['Content-Length'] == '100' + assert handler._headers['Content-Type'] == mock_response_stream_reader.content_type + disposition = 'attachment; filename="test stream"; filename*=UTF-8\'\'test%20stream' + assert handler._headers['Content-Disposition'] == disposition + + handler.write_stream.assert_awaited_once() + + @pytest.mark.asyncio + async def test_download_file_header_gzip_with_response_stream_reader(self, http_request, mock_response_stream_reader): + + handler = mock_handler(http_request) + mock_response_stream_reader.response.headers['Content-Encoding'] = 'gzip' + handler.provider.download = MockCoroutine(return_value=mock_response_stream_reader) + handler.path = WaterButlerPath('/test_file') + handler.bytes_downloaded = '100' + + await handler.download_file() + assert handler._headers['Content-Length'] == '100' + assert handler._headers['Content-Type'] == mock_response_stream_reader.content_type + disposition = 'attachment; filename="test stream"; filename*=UTF-8\'\'test%20stream' + assert handler._headers['Content-Disposition'] == disposition + + handler.write_stream.assert_awaited_once() + + @pytest.mark.asyncio + async def test_download_file_header_gzip_bytes_downloaded_is_none_with_response_stream_reader(self, http_request, mock_response_stream_reader): + + handler = mock_handler(http_request) + assert 'Content-Length' not in handler._headers + mock_response_stream_reader.response.headers['Content-Encoding'] = 'gzip' + handler.provider.download = MockCoroutine(return_value=mock_response_stream_reader) + handler.path = WaterButlerPath('/test_file') + handler.bytes_downloaded = None + + await handler.download_file() + assert handler._headers['Content-Type'] == mock_response_stream_reader.content_type + disposition = 'attachment; filename="test stream"; filename*=UTF-8\'\'test%20stream' + assert handler._headers['Content-Disposition'] == disposition + + handler.write_stream.assert_awaited_once() diff --git a/waterbutler/server/api/v1/provider/metadata.py b/waterbutler/server/api/v1/provider/metadata.py index 2bbb27feb..8c715f4d1 100644 --- a/waterbutler/server/api/v1/provider/metadata.py +++ b/waterbutler/server/api/v1/provider/metadata.py @@ -89,6 +89,11 @@ async def download_file(self): if isinstance(stream, str): return self.redirect(stream) + is_compression = is_chunked = False + if isinstance(stream, ResponseStreamReader): + is_compression = 'gzip' in stream.response.headers.get('Content-Encoding', '') + is_chunked = 'chunked' in stream.response.headers.get('Transfer-Encoding', '') + if getattr(stream, 'partial', None): # Use getattr here as not all stream may have a partial attribute # Plus it fixes tests @@ -99,7 +104,12 @@ async def download_file(self): self.set_header('Content-Type', stream.content_type) logger.debug('stream size is: {}'.format(stream.size)) - if stream.size is not None: + # - In chunked transfer encoding, the data stream is divided into a series of non-overlapping "chunks". + # The Content-Length header can be omitted in this case + # - Another case, If the Content-Encoding header is `gzip`, + # the Content-Length is equal to the length of the body after the Content-Encoding. + # We need adjust it after loading the entire body + if not (is_chunked or is_compression) and stream.size is not None: self.set_header('Content-Length', str(stream.size)) # Build `Content-Disposition` header from `displayName` override, @@ -112,13 +122,17 @@ async def download_file(self): self.set_header('Content-Disposition', make_disposition(name)) _, ext = os.path.splitext(name) - # If the file extention is in mime_types + # If the file extension is in mime_types # override the content type to fix issues with safari shoving in new file extensions if ext in mime_types: self.set_header('Content-Type', mime_types[ext]) await self.write_stream(stream) + # update response headers by the original size after decompressing + if (is_chunked or is_compression) and self.bytes_downloaded is not None: + self.set_header('Content-Length', str(self.bytes_downloaded)) + if getattr(stream, 'partial', False) and isinstance(stream, ResponseStreamReader): await stream.response.release()