From 65053b7737ffe9641974df926d694efd801ba417 Mon Sep 17 00:00:00 2001 From: Lucas Alvares Gomes Date: Tue, 13 Oct 2015 17:10:37 +0100 Subject: [PATCH] Refactor the image download and checksum computation bits Prior to this patch downloading and computing the checksum of the image were done in different stages, after the download the file would need to be re-read and the checksum was computed. This patch is changing it by creating a ImageDownload class which computes the checksum at the same time the image is being downloaded, this saves time and also make the code more portable. Related-Bug: #1505685 Change-Id: I71f9f2bd9a62a6a6cc474d0ae519591cea6381d6 --- ironic_python_agent/extensions/standby.py | 116 ++++++++++-------- .../tests/unit/extensions/test_standby.py | 114 ++++++++--------- 2 files changed, 122 insertions(+), 108 deletions(-) diff --git a/ironic_python_agent/extensions/standby.py b/ironic_python_agent/extensions/standby.py index 2c8a9b739..afbd9f1b9 100644 --- a/ironic_python_agent/extensions/standby.py +++ b/ironic_python_agent/extensions/standby.py @@ -115,42 +115,75 @@ def _write_configdrive_to_partition(configdrive, device): totaltime)) -def _request_url(image_info, url): - no_proxy = image_info.get('no_proxy') - if no_proxy: - os.environ['no_proxy'] = no_proxy - proxies = image_info.get('proxies', {}) - resp = requests.get(url, stream=True, proxies=proxies) - if resp.status_code != 200: - msg = ('Received status code {0} from {1}, expected 200. Response ' - 'body: {2}').format(resp.status_code, url, resp.text) - raise errors.ImageDownloadError(image_info['id'], msg) - return resp +class ImageDownload(object): + """Helper class that opens a HTTP connection to download an image. + + This class opens a HTTP connection to download an image from a URL + and create an iterator so the image can be downloaded in chunks. The + MD5 hash of the image being downloaded is calculated on-the-fly. + """ + + def __init__(self, image_info, time_obj=None): + self._md5checksum = hashlib.md5() + self._time = time_obj or time.time() + self._request = None + + for url in image_info['urls']: + try: + LOG.info("Attempting to download image from {0}".format(url)) + self._request = self._download_file(image_info, url) + except errors.ImageDownloadError as e: + failtime = time.time() - self._time + log_msg = ('Image download failed. URL: {0}; time: {1} ' + 'seconds. Error: {2}') + LOG.warning(log_msg.format(url, failtime, e.details)) + continue + else: + break + else: + msg = 'Image download failed for all URLs.' + raise errors.ImageDownloadError(image_info['id'], msg) + + def _download_file(self, image_info, url): + no_proxy = image_info.get('no_proxy') + if no_proxy: + os.environ['no_proxy'] = no_proxy + proxies = image_info.get('proxies', {}) + resp = requests.get(url, stream=True, proxies=proxies) + if resp.status_code != 200: + msg = ('Received status code {0} from {1}, expected 200. Response ' + 'body: {2}').format(resp.status_code, url, resp.text) + raise errors.ImageDownloadError(image_info['id'], msg) + return resp + + def __iter__(self): + for chunk in self._request.iter_content(IMAGE_CHUNK_SIZE): + self._md5checksum.update(chunk) + yield chunk + + def md5sum(self): + return self._md5checksum.hexdigest() + + +def _verify_image(image_info, image_location, checksum): + LOG.debug('Verifying image at {0} against MD5 checksum ' + '{1}'.format(image_location, checksum)) + if checksum != image_info['checksum']: + LOG.error(errors.ImageChecksumError.details_str.format( + image_location, image_info['id'], + image_info['checksum'], checksum)) + raise errors.ImageChecksumError(image_location, image_info['id'], + image_info['checksum'], checksum) def _download_image(image_info): starttime = time.time() - resp = None - for url in image_info['urls']: - try: - LOG.info("Attempting to download image from {0}".format(url)) - resp = _request_url(image_info, url) - except errors.ImageDownloadError as e: - failtime = time.time() - starttime - log_msg = ('Image download failed. URL: {0}; time: {1} seconds. ' - 'Error: {2}') - LOG.warning(log_msg.format(url, failtime, e.details)) - continue - else: - break - if resp is None: - msg = 'Image download failed for all URLs.' - raise errors.ImageDownloadError(image_info['id'], msg) - image_location = _image_location(image_info) + image_download = ImageDownload(image_info, time_obj=starttime) + with open(image_location, 'wb') as f: try: - for chunk in resp.iter_content(IMAGE_CHUNK_SIZE): + for chunk in image_download: f.write(chunk) except Exception as e: msg = 'Unable to write image to {0}. Error: {1}'.format( @@ -160,30 +193,7 @@ def _download_image(image_info): totaltime = time.time() - starttime LOG.info("Image downloaded from {0} in {1} seconds".format(image_location, totaltime)) - - _verify_image(image_info, image_location) - - -def _verify_image(image_info, image_location): - checksum = image_info['checksum'] - log_msg = 'Verifying image at {0} against MD5 checksum {1}' - LOG.debug(log_msg.format(image_location, checksum)) - hash_ = hashlib.md5() - with open(image_location) as image: - while True: - data = image.read(IMAGE_CHUNK_SIZE) - if not data: - break - hash_.update(data) - hash_digest = hash_.hexdigest() - if hash_digest == checksum: - return True - - LOG.error(errors.ImageChecksumError.details_str.format( - image_location, image_info['id'], checksum, hash_digest)) - - raise errors.ImageChecksumError(image_location, image_info['id'], checksum, - hash_digest) + _verify_image(image_info, image_location, image_download.md5sum()) def _validate_image_info(ext, image_info=None, **kwargs): diff --git a/ironic_python_agent/tests/unit/extensions/test_standby.py b/ironic_python_agent/tests/unit/extensions/test_standby.py index 010d9f662..d924117ad 100644 --- a/ironic_python_agent/tests/unit/extensions/test_standby.py +++ b/ironic_python_agent/tests/unit/extensions/test_standby.py @@ -28,26 +28,27 @@ else: OPEN_FUNCTION_NAME = 'builtins.open' +def _build_fake_image_info(): + return { + 'id': 'fake_id', + 'urls': [ + 'http://example.org', + ], + 'checksum': 'abc123' + } + + class TestStandbyExtension(test_base.BaseTestCase): def setUp(self): super(TestStandbyExtension, self).setUp() self.agent_extension = standby.StandbyExtension() - def _build_fake_image_info(self): - return { - 'id': 'fake_id', - 'urls': [ - 'http://example.org', - ], - 'checksum': 'abc123' - } - def test_validate_image_info_success(self): - standby._validate_image_info(None, self._build_fake_image_info()) + standby._validate_image_info(None, _build_fake_image_info()) def test_validate_image_info_missing_field(self): for field in ['id', 'urls', 'checksum']: - invalid_info = self._build_fake_image_info() + invalid_info = _build_fake_image_info() del invalid_info[field] self.assertRaises(errors.InvalidCommandParamsError, @@ -55,7 +56,7 @@ class TestStandbyExtension(test_base.BaseTestCase): invalid_info) def test_validate_image_info_invalid_urls(self): - invalid_info = self._build_fake_image_info() + invalid_info = _build_fake_image_info() invalid_info['urls'] = 'this_is_not_a_list' self.assertRaises(errors.InvalidCommandParamsError, @@ -63,7 +64,7 @@ class TestStandbyExtension(test_base.BaseTestCase): invalid_info) def test_validate_image_info_empty_urls(self): - invalid_info = self._build_fake_image_info() + invalid_info = _build_fake_image_info() invalid_info['urls'] = [] self.assertRaises(errors.InvalidCommandParamsError, @@ -71,7 +72,7 @@ class TestStandbyExtension(test_base.BaseTestCase): invalid_info) def test_validate_image_info_invalid_checksum(self): - invalid_info = self._build_fake_image_info() + invalid_info = _build_fake_image_info() invalid_info['checksum'] = {'not': 'a string'} self.assertRaises(errors.InvalidCommandParamsError, @@ -79,7 +80,7 @@ class TestStandbyExtension(test_base.BaseTestCase): invalid_info) def test_validate_image_info_empty_checksum(self): - invalid_info = self._build_fake_image_info() + invalid_info = _build_fake_image_info() invalid_info['checksum'] = '' self.assertRaises(errors.InvalidCommandParamsError, @@ -92,14 +93,14 @@ class TestStandbyExtension(test_base.BaseTestCase): image_info={'foo': 'bar'}) def test_image_location(self): - image_info = self._build_fake_image_info() + image_info = _build_fake_image_info() location = standby._image_location(image_info) self.assertEqual(location, '/tmp/fake_id') @mock.patch(OPEN_FUNCTION_NAME, autospec=True) @mock.patch('ironic_python_agent.utils.execute', autospec=True) def test_write_image(self, execute_mock, open_mock): - image_info = self._build_fake_image_info() + image_info = _build_fake_image_info() device = '/dev/sda' location = standby._image_location(image_info) script = standby._path_to_script('shell/write_image.sh') @@ -202,7 +203,7 @@ class TestStandbyExtension(test_base.BaseTestCase): @mock.patch(OPEN_FUNCTION_NAME) @mock.patch('requests.get') def test_download_image(self, requests_mock, open_mock, md5_mock): - image_info = self._build_fake_image_info() + image_info = _build_fake_image_info() response = requests_mock.return_value response.status_code = 200 response.iter_content.return_value = ['some', 'content'] @@ -226,7 +227,7 @@ class TestStandbyExtension(test_base.BaseTestCase): @mock.patch.dict(os.environ, {}) def test_download_image_proxy( self, requests_mock, open_mock, md5_mock): - image_info = self._build_fake_image_info() + image_info = _build_fake_image_info() proxies = {'http': 'http://a.b.com', 'https': 'https://secure.a.b.com'} no_proxy = '.example.org,.b.com' @@ -252,57 +253,40 @@ class TestStandbyExtension(test_base.BaseTestCase): @mock.patch('requests.get', autospec=True) def test_download_image_bad_status(self, requests_mock): - image_info = self._build_fake_image_info() + image_info = _build_fake_image_info() response = requests_mock.return_value response.status_code = 404 self.assertRaises(errors.ImageDownloadError, standby._download_image, image_info) - @mock.patch('ironic_python_agent.extensions.standby._verify_image', - autospec=True) + @mock.patch('hashlib.md5', autospec=True) @mock.patch(OPEN_FUNCTION_NAME, autospec=True) @mock.patch('requests.get', autospec=True) def test_download_image_verify_fails(self, requests_mock, open_mock, - verify_mock): - image_info = self._build_fake_image_info() + md5_mock): + image_info = _build_fake_image_info() response = requests_mock.return_value response.status_code = 200 - verify_mock.side_effect = errors.ImageChecksumError( - 'foo', '/bar/foo', 'incorrect', 'correct') + hexdigest_mock = md5_mock.return_value.hexdigest + hexdigest_mock.return_value = 'invalid-checksum' self.assertRaises(errors.ImageChecksumError, standby._download_image, image_info) - @mock.patch(OPEN_FUNCTION_NAME) - @mock.patch('hashlib.md5') - def test_verify_image_success(self, md5_mock, open_mock): - image_info = self._build_fake_image_info() - hexdigest_mock = md5_mock.return_value.hexdigest - hexdigest_mock.return_value = image_info['checksum'] - file_mock = mock.Mock() - file_mock.read.return_value = None - open_mock.return_value.__enter__.return_value = file_mock + def test_verify_image_success(self): + image_info = _build_fake_image_info() image_location = '/foo/bar' + checksum = image_info['checksum'] + standby._verify_image(image_info, image_location, checksum) - verified = standby._verify_image(image_info, image_location) - self.assertTrue(verified) - self.assertEqual(1, md5_mock.call_count) - - @mock.patch(OPEN_FUNCTION_NAME) - @mock.patch('hashlib.md5') - def test_verify_image_failure(self, md5_mock, open_mock): - image_info = self._build_fake_image_info() - md5_mock.return_value.hexdigest.return_value = 'wrong hash' - file_mock = mock.Mock() - file_mock.read.return_value = None - open_mock.return_value.__enter__.return_value = file_mock + def test_verify_image_failure(self): + image_info = _build_fake_image_info() image_location = '/foo/bar' - + checksum = 'invalid-checksum' self.assertRaises(errors.ImageChecksumError, standby._verify_image, - image_info, image_location) - self.assertEqual(md5_mock.call_count, 1) + image_info, image_location, checksum) @mock.patch('ironic_python_agent.hardware.dispatch_to_managers', autospec=True) @@ -312,7 +296,7 @@ class TestStandbyExtension(test_base.BaseTestCase): autospec=True) def test_cache_image(self, download_mock, write_mock, dispatch_mock): - image_info = self._build_fake_image_info() + image_info = _build_fake_image_info() download_mock.return_value = None write_mock.return_value = None dispatch_mock.return_value = 'manager' @@ -337,7 +321,7 @@ class TestStandbyExtension(test_base.BaseTestCase): autospec=True) def test_cache_image_force(self, download_mock, write_mock, dispatch_mock): - image_info = self._build_fake_image_info() + image_info = _build_fake_image_info() self.agent_extension.cached_image_id = image_info['id'] download_mock.return_value = None write_mock.return_value = None @@ -365,7 +349,7 @@ class TestStandbyExtension(test_base.BaseTestCase): autospec=True) def test_cache_image_cached(self, download_mock, write_mock, dispatch_mock): - image_info = self._build_fake_image_info() + image_info = _build_fake_image_info() self.agent_extension.cached_image_id = image_info['id'] download_mock.return_value = None write_mock.return_value = None @@ -400,7 +384,7 @@ class TestStandbyExtension(test_base.BaseTestCase): write_mock, dispatch_mock, configdrive_copy_mock): - image_info = self._build_fake_image_info() + image_info = _build_fake_image_info() location_mock.return_value = '/tmp/configdrive' download_mock.return_value = None write_mock.return_value = None @@ -460,7 +444,7 @@ class TestStandbyExtension(test_base.BaseTestCase): write_mock, dispatch_mock, configdrive_copy_mock): - image_info = self._build_fake_image_info() + image_info = _build_fake_image_info() download_mock.return_value = None write_mock.return_value = None dispatch_mock.return_value = 'manager' @@ -530,3 +514,23 @@ class TestStandbyExtension(test_base.BaseTestCase): execute_mock.assert_called_once_with(*command, check_exit_code=[0]) self.assertEqual('FAILED', failed_result.command_status) + + +class TestImageDownload(test_base.BaseTestCase): + + @mock.patch('hashlib.md5', autospec=True) + @mock.patch('requests.get', autospec=True) + def test_download_image(self, requests_mock, md5_mock): + content = ['SpongeBob', 'SquarePants'] + response = requests_mock.return_value + response.status_code = 200 + response.iter_content.return_value = content + + image_info = _build_fake_image_info() + md5_mock.return_value.hexdigest.return_value = image_info['checksum'] + image_download = standby.ImageDownload(image_info) + + self.assertEqual(content, list(image_download)) + requests_mock.assert_called_once_with(image_info['urls'][0], + stream=True, proxies={}) + self.assertEqual(image_info['checksum'], image_download.md5sum())