diff --git a/src/scriptworker/artifacts.py b/src/scriptworker/artifacts.py index 0960eb37..a7a269d9 100644 --- a/src/scriptworker/artifacts.py +++ b/src/scriptworker/artifacts.py @@ -6,6 +6,7 @@ """ import asyncio +import base64 import fnmatch import gzip import logging @@ -21,7 +22,15 @@ from scriptworker.client import validate_artifact_url from scriptworker.exceptions import DownloadError, ScriptWorkerRetryException, ScriptWorkerTaskException from scriptworker.task import get_decision_task_id, get_run_id, get_task_id -from scriptworker.utils import add_enumerable_item_to_dict, download_file, get_loggable_url, raise_future_exceptions, retry_async, semaphore_wrapper +from scriptworker.utils import ( + add_enumerable_item_to_dict, + download_file, + get_hash, + get_loggable_url, + raise_future_exceptions, + retry_async, + semaphore_wrapper, +) log = logging.getLogger(__name__) @@ -163,6 +172,8 @@ async def create_artifact(context, path, target_path, content_type, content_enco payload = {"storageType": storage_type, "expires": expires or get_expiration_arrow(context).isoformat(), "contentType": content_type} args = [get_task_id(context.claim_task), get_run_id(context.claim_task), target_path, payload] + content_md5 = await asyncio.to_thread(get_content_md5, path) + tc_response = await context.temp_queue.createArtifact(*args) skip_auto_headers = [aiohttp.hdrs.CONTENT_TYPE] loggable_url = get_loggable_url(tc_response["putUrl"]) @@ -172,7 +183,7 @@ async def create_artifact(context, path, target_path, content_type, content_enco async with context.session.put( tc_response["putUrl"], data=fh, - headers=_craft_artifact_put_headers(content_type, content_encoding), + headers=_craft_artifact_put_headers(content_type, content_md5, content_encoding), skip_auto_headers=skip_auto_headers, compress=False, ) as resp: @@ -180,6 +191,12 @@ async def create_artifact(context, path, target_path, content_type, content_enco response_text = await resp.text() log.info(response_text) if resp.status not in (200, 204): + if content_md5 is not None and "BadDigest" in (response_text or ""): + log.error( + "{} was corrupted in transit: the storage backend rejected the body as not matching Content-MD5 {}. Retrying.".format( + target_path, content_md5 + ) + ) raise ScriptWorkerRetryException("Bad status {}".format(resp.status)) @@ -206,9 +223,22 @@ async def create_link_artifact(context, target_path, link_to, content_type, expi await context.temp_queue.createArtifact(*args) -def _craft_artifact_put_headers(content_type, encoding=None): +def get_content_md5(path): + """Get the base64-encoded md5 digest of a file, formatted for ``Content-MD5``. + + Args: + path (str): the path to the file to digest. + + Returns: + str: the base64-encoded md5 digest. + + """ + return base64.b64encode(bytes.fromhex(get_hash(path, hash_alg="md5"))).decode("ascii") + + +def _craft_artifact_put_headers(content_type, content_md5, encoding=None): log.debug("{} {}".format(content_type, encoding)) - headers = {aiohttp.hdrs.CONTENT_TYPE: content_type} + headers = {aiohttp.hdrs.CONTENT_TYPE: content_type, aiohttp.hdrs.CONTENT_MD5: content_md5} if encoding is not None: headers[aiohttp.hdrs.CONTENT_ENCODING] = encoding diff --git a/tests/test_artifacts.py b/tests/test_artifacts.py index 6621ab1e..4380bbee 100644 --- a/tests/test_artifacts.py +++ b/tests/test_artifacts.py @@ -1,10 +1,13 @@ import asyncio +import base64 import gzip +import hashlib import itertools import json import os import tempfile +import aiohttp import arrow import mock import pytest @@ -18,6 +21,7 @@ download_artifacts, get_and_check_single_upstream_artifact_full_path, get_artifact_url, + get_content_md5, get_expiration_arrow, get_optional_artifacts_per_task_id, get_single_upstream_artifact_full_path, @@ -27,7 +31,7 @@ ) from scriptworker.exceptions import ScriptWorkerRetryException, ScriptWorkerTaskException -from . import touch +from . import FakeResponse, touch @pytest.fixture(scope="function") @@ -162,6 +166,81 @@ async def test_create_artifact_retry(context, fake_session_500, successful_queue await create_artifact(context, path, "public/env/one.log", content_type="text/plain", content_encoding=None, expires=expires) +def _write(path, contents=b"hello world"): + with open(path, "wb") as fh: + fh.write(contents) + return contents + + +async def _capture_put(context, fake_session, successful_queue, path, target_path="public/env/one.txt", content_type="text/plain", content_encoding=None): + captured = {} + original = fake_session._request + + async def capture(method, url, *args, **kwargs): + captured.update(kwargs) + return await original(method, url, *args, **kwargs) + + fake_session._request = capture + context.session = fake_session + context.temp_queue = successful_queue + await create_artifact(context, path, target_path, content_type=content_type, content_encoding=content_encoding, expires=arrow.utcnow().isoformat()) + return captured + + +@pytest.mark.asyncio +async def test_create_artifact_sends_content_md5(context, fake_session, successful_queue): + path = os.path.join(context.config["artifact_dir"], "one.txt") + contents = _write(path) + + captured = await _capture_put(context, fake_session, successful_queue, path) + + expected = base64.b64encode(hashlib.md5(contents).digest()).decode("ascii") + assert captured["headers"][aiohttp.hdrs.CONTENT_MD5] == expected + + +@pytest.mark.asyncio +async def test_create_artifact_content_md5_covers_compressed_bytes(context, fake_session, successful_queue): + """The digest has to cover what goes on the wire, which for a gzipped artifact is the compressed file.""" + path = os.path.join(context.config["artifact_dir"], "one.log") + original_contents = _write(path, b"12:00:00 Foo bar") + content_type, content_encoding = compress_artifact_if_supported(path) + assert content_encoding == "gzip" + + captured = await _capture_put( + context, fake_session, successful_queue, path, target_path="public/logs/one.log", content_type=content_type, content_encoding=content_encoding + ) + + with open(path, "rb") as fh: + on_disk = fh.read() + assert on_disk != original_contents + assert captured["headers"][aiohttp.hdrs.CONTENT_MD5] == base64.b64encode(hashlib.md5(on_disk).digest()).decode("ascii") + assert captured["headers"][aiohttp.hdrs.CONTENT_ENCODING] == "gzip" + + +@pytest.mark.asyncio +async def test_create_artifact_bad_digest_retries(context, fake_session, successful_queue, caplog): + path = os.path.join(context.config["artifact_dir"], "one.txt") + _write(path) + + async def bad_digest(method, url, *args, **kwargs): + return FakeResponse(method, url, status=400, payload="BadDigest") + + fake_session._request = bad_digest + context.session = fake_session + context.temp_queue = successful_queue + + with pytest.raises(ScriptWorkerRetryException): + await create_artifact(context, path, "public/env/one.txt", content_type="text/plain", content_encoding=None, expires=arrow.utcnow().isoformat()) + + assert "was corrupted in transit" in caplog.text + + +def test_get_content_md5(tmpdir): + path = os.path.join(tmpdir, "one.txt") + contents = _write(path, b"some artifact contents") + assert get_content_md5(path) == base64.b64encode(hashlib.md5(contents).digest()).decode("ascii") + + @pytest.mark.asyncio async def test_create_link_artifact(context, successful_queue): expires = arrow.utcnow().isoformat() @@ -191,9 +270,13 @@ async def test_create_link_artifact(context, successful_queue): def test_craft_artifact_put_headers(): - assert _craft_artifact_put_headers("text/plain") == {"Content-Type": "text/plain"} - assert _craft_artifact_put_headers("text/plain", encoding=None) == {"Content-Type": "text/plain"} - assert _craft_artifact_put_headers("text/plain", "gzip") == {"Content-Type": "text/plain", "Content-Encoding": "gzip"} + assert _craft_artifact_put_headers("text/plain", "deadbeef==") == {"Content-Type": "text/plain", "Content-MD5": "deadbeef=="} + assert _craft_artifact_put_headers("text/plain", "deadbeef==", encoding=None) == {"Content-Type": "text/plain", "Content-MD5": "deadbeef=="} + assert _craft_artifact_put_headers("text/plain", "deadbeef==", "gzip") == { + "Content-Type": "text/plain", + "Content-Encoding": "gzip", + "Content-MD5": "deadbeef==", + } # get_artifact_url {{{1