feat: set Content-MD5 header to avoid silent data corruption

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,14 +183,20 @@ 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:
                 log.info("create_artifact {}: {}".format(path, resp.status))
                 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

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="<Error><Code>BadDigest</Code></Error>")
+
+    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