Image import "web-download" check downloaded size

If the downloaded data size is different from the expected one, the
task "web-download" in the image import process will now fail.

Change-Id: Ie260486d795a6f4af1632f6f3708abc92fb47a3a
Closes-Bug: #1895663

Author: yebinama

glance/async_/flows/_internal_plugins/web_download.py

@@ -41,6 +41,7 @@ def __init__(self, task_id, task_type, image_repo, image_id, uri):
         self.image_repo = image_repo
         self.image_id = image_id
         self.uri = uri
+        self._path = None
         super(_WebDownload, self).__init__(
             name='%s-WebDownload-%s' % (task_type, task_id))
 
@@ -117,8 +118,19 @@ def execute(self):
                           {"error": encodeutils.exception_to_unicode(e),
                            "task_id": self.task_id})
 
-        path = self.store.add(self.image_id, data, 0)[0]
-        return path
+        self._path, bytes_written = self.store.add(self.image_id, data, 0)[0:2]
+        try:
+            content_length = int(data.headers['content-length'])
+            if bytes_written != content_length:
+                msg = (_("Task %(task_id)s failed because downloaded data "
+                         "size %(data_size)i is different from expected %("
+                         "expected)i") %
+                       {"task_id": self.task_id, "data_size": bytes_written,
+                        "expected": content_length})
+                raise exception.ImportTaskError(msg)
+        except (KeyError, ValueError):
+            pass
+        return self._path
 
     def revert(self, result, **kwargs):
         if isinstance(result, failure.Failure):

glance/tests/unit/async_/flows/test_web_download.py

@@ -66,9 +66,50 @@ def test_web_download(self, mock_add):
             self.image_id, self.uri)
         with mock.patch.object(script_utils,
                                'get_image_data_iter') as mock_iter:
-            mock_iter.return_value = b"dddd"
-            web_download_task.execute()
-            mock_add.assert_called_once_with(self.image_id, b"dddd", 0)
+            mock_add.return_value = ["path", 4]
+            mock_iter.return_value.headers = {}
+            self.assertEqual(web_download_task.execute(), "path")
+            mock_add.assert_called_once_with(self.image_id,
+                                             mock_iter.return_value, 0)
+
+    @mock.patch.object(filesystem.Store, 'add')
+    def test_web_download_with_content_length(self, mock_add):
+        web_download_task = web_download._WebDownload(
+            self.task.task_id, self.task_type, self.task_repo,
+            self.image_id, self.uri)
+        with mock.patch.object(script_utils,
+                               'get_image_data_iter') as mock_iter:
+            mock_iter.return_value.headers = {'content-length': '4'}
+            mock_add.return_value = ["path", 4]
+            self.assertEqual(web_download_task.execute(), "path")
+            mock_add.assert_called_once_with(self.image_id,
+                                             mock_iter.return_value, 0)
+
+    @mock.patch.object(filesystem.Store, 'add')
+    def test_web_download_with_invalid_content_length(self, mock_add):
+        web_download_task = web_download._WebDownload(
+            self.task.task_id, self.task_type, self.task_repo,
+            self.image_id, self.uri)
+        with mock.patch.object(script_utils,
+                               'get_image_data_iter') as mock_iter:
+            mock_iter.return_value.headers = {'content-length': "not_valid"}
+            mock_add.return_value = ["path", 4]
+            self.assertEqual(web_download_task.execute(), "path")
+            mock_add.assert_called_once_with(self.image_id,
+                                             mock_iter.return_value, 0)
+
+    @mock.patch.object(filesystem.Store, 'add')
+    def test_web_download_fails_when_data_size_different(self, mock_add):
+        web_download_task = web_download._WebDownload(
+            self.task.task_id, self.task_type, self.task_repo,
+            self.image_id, self.uri)
+        with mock.patch.object(script_utils,
+                               'get_image_data_iter') as mock_iter:
+            mock_iter.return_value.headers = {'content-length': '4'}
+            mock_add.return_value = ["path", 3]
+            self.assertRaises(
+                glance.common.exception.ImportTaskError,
+                web_download_task.execute)
 
     def test_web_download_node_staging_uri_is_none(self):
         self.config(node_staging_uri=None)