Feature/Grdm-202403-Step3/2.1.3_remove-export-restriction-nextcloud-for-institution

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()

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()