From c5877fb85776066d3dae7ae24c5332f705c2b9db Mon Sep 17 00:00:00 2001 From: Squeaky Date: Thu, 5 Sep 2024 18:56:07 +0200 Subject: [PATCH 1/4] zstd --- inbox/api/ns_api.py | 2 +- inbox/models/message.py | 4 +-- inbox/util/blockstore.py | 53 +++++++++++++++++++++++++++++++++++++--- requirements/prod.txt | 1 + 4 files changed, 54 insertions(+), 6 deletions(-) diff --git a/inbox/api/ns_api.py b/inbox/api/ns_api.py index 6eccba78c..f4e0cce14 100644 --- a/inbox/api/ns_api.py +++ b/inbox/api/ns_api.py @@ -619,7 +619,7 @@ def message_read_api(public_id): raise NotFoundError(f"Couldn't find message {public_id}") if request.headers.get("Accept", None) == "message/rfc822": - raw_message = blockstore.get_from_blockstore(message.data_sha256) + raw_message = blockstore.get_raw_mime(message.data_sha256) if raw_message is not None: return Response(raw_message, mimetype="message/rfc822") else: diff --git a/inbox/models/message.py b/inbox/models/message.py index 7cc150515..c1a41dae7 100644 --- a/inbox/models/message.py +++ b/inbox/models/message.py @@ -48,7 +48,7 @@ from inbox.security.blobstorage import decode_blob, encode_blob from inbox.sqlalchemy_ext.util import JSON, MAX_MYSQL_INTEGER, json_field_too_long from inbox.util.addr import HeaderTooBigException, parse_mimepart_address_header -from inbox.util.blockstore import save_to_blockstore +from inbox.util.blockstore import save_raw_mime from inbox.util.encoding import unicode_safe_truncate from inbox.util.html import HTMLParseError, plaintext2html, strip_tags from inbox.util.misc import get_internaldate, parse_references @@ -305,7 +305,7 @@ def create_from_synced( message.data_sha256 = sha256(body).hexdigest() # Persist the raw MIME message to disk/ S3 - save_to_blockstore(message.data_sha256, body) + save_raw_mime(message.data_sha256, body) # Persist the processed message to the database message.namespace_id = account.namespace.id diff --git a/inbox/util/blockstore.py b/inbox/util/blockstore.py index 079b40c8b..332edd24d 100644 --- a/inbox/util/blockstore.py +++ b/inbox/util/blockstore.py @@ -3,6 +3,8 @@ from hashlib import sha256 from typing import Optional +import zstandard + from inbox.config import config from inbox.logging import get_logger from inbox.util.stats import statsd_client @@ -16,6 +18,9 @@ from boto.s3.connection import S3Connection from boto.s3.key import Key +# https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#zstandard-frames +ZSTD_MAGIC_NUMBER_PREFIX = 0xFD2FB528.to_bytes(4, "little") + def _data_file_directory(h): return os.path.join( @@ -27,6 +32,29 @@ def _data_file_path(h): return os.path.join(_data_file_directory(h), h) +def save_raw_mime( + data_sha256: str, decompressed_raw_mime: bytes, *, compress: bool | None = None +) -> int: + if compress is None: + compress = config.get("COMPRESS_RAW_MIME", False) + + if compress: + assert not decompressed_raw_mime.startswith(ZSTD_MAGIC_NUMBER_PREFIX) + + compressed_raw_mime = zstandard.compress(decompressed_raw_mime) + + assert compressed_raw_mime.startswith(ZSTD_MAGIC_NUMBER_PREFIX) + + if len(compressed_raw_mime) > decompressed_raw_mime: + compressed_raw_mime = decompressed_raw_mime + else: + compressed_raw_mime = decompressed_raw_mime + + save_to_blockstore(data_sha256, compressed_raw_mime) + + return len(compressed_raw_mime) + + def save_to_blockstore(data_sha256: str, data: bytes) -> None: assert data is not None assert isinstance(data, bytes) @@ -86,7 +114,7 @@ def _save_to_s3_bucket(data_sha256: str, bucket_name: str, data: bytes) -> None: statsd_client.timing("s3_blockstore.save_latency", latency_millis) -def get_from_blockstore(data_sha256) -> Optional[bytes]: +def get_from_blockstore(data_sha256, *, check_sha=True) -> Optional[bytes]: if STORE_MSG_ON_S3: value = _get_from_s3(data_sha256) else: @@ -97,10 +125,29 @@ def get_from_blockstore(data_sha256) -> Optional[bytes]: log.warning("No data returned!") return value + if check_sha: + assert ( + data_sha256 == sha256(value).hexdigest() + ), "Returned data doesn't match stored hash!" + + return value + + +def get_raw_mime(data_sha256: str) -> Optional[bytes]: + compressed_raw_mime = get_from_blockstore(data_sha256, check_sha=False) + if compressed_raw_mime is None: + return None + + if compressed_raw_mime.startswith(ZSTD_MAGIC_NUMBER_PREFIX): + decompressed_raw_mime = zstandard.decompress(compressed_raw_mime) + else: + decompressed_raw_mime = compressed_raw_mime + assert ( - data_sha256 == sha256(value).hexdigest() + sha256(decompressed_raw_mime).hexdigest() == data_sha256 ), "Returned data doesn't match stored hash!" - return value + + return decompressed_raw_mime def _get_from_s3(data_sha256): diff --git a/requirements/prod.txt b/requirements/prod.txt index e2902b755..2b1454e9f 100644 --- a/requirements/prod.txt +++ b/requirements/prod.txt @@ -86,3 +86,4 @@ Werkzeug==3.0.3 zipp==3.6.0 zope.event==4.5.0 zope.interface==5.4.0 +zstandard==0.23.0 From b879061974627806126c188bd67d00e326e09fda Mon Sep 17 00:00:00 2001 From: Squeaky Date: Tue, 10 Sep 2024 11:41:52 +0200 Subject: [PATCH 2/4] Fix typing --- inbox/util/blockstore.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/inbox/util/blockstore.py b/inbox/util/blockstore.py index 332edd24d..85cabf456 100644 --- a/inbox/util/blockstore.py +++ b/inbox/util/blockstore.py @@ -33,7 +33,7 @@ def _data_file_path(h): def save_raw_mime( - data_sha256: str, decompressed_raw_mime: bytes, *, compress: bool | None = None + data_sha256: str, decompressed_raw_mime: bytes, *, compress: "bool | None" = None ) -> int: if compress is None: compress = config.get("COMPRESS_RAW_MIME", False) @@ -45,7 +45,7 @@ def save_raw_mime( assert compressed_raw_mime.startswith(ZSTD_MAGIC_NUMBER_PREFIX) - if len(compressed_raw_mime) > decompressed_raw_mime: + if len(compressed_raw_mime) > len(decompressed_raw_mime): compressed_raw_mime = decompressed_raw_mime else: compressed_raw_mime = decompressed_raw_mime @@ -133,7 +133,7 @@ def get_from_blockstore(data_sha256, *, check_sha=True) -> Optional[bytes]: return value -def get_raw_mime(data_sha256: str) -> Optional[bytes]: +def get_raw_mime(data_sha256: str) -> "bytes | None": compressed_raw_mime = get_from_blockstore(data_sha256, check_sha=False) if compressed_raw_mime is None: return None From 5d00c5ca4d2b7852df06b3cc9db8ac541d0009a1 Mon Sep 17 00:00:00 2001 From: Squeaky Date: Tue, 10 Sep 2024 12:32:53 +0200 Subject: [PATCH 3/4] Tests --- .gitattributes | 2 ++ tests/util/test_blockstore.py | 44 +++++++++++++++++++++++++ tests/util/tiny.eml | 60 +++++++++++++++++++++++++++++++++++ 3 files changed, 106 insertions(+) create mode 100644 tests/util/test_blockstore.py create mode 100644 tests/util/tiny.eml diff --git a/.gitattributes b/.gitattributes index c9bbacfd8..292c30f80 100644 --- a/.gitattributes +++ b/.gitattributes @@ -1,4 +1,6 @@ * text eol=lf +*.eml eol=crlf *.png binary *.jpg binary *.ics binary + diff --git a/tests/util/test_blockstore.py b/tests/util/test_blockstore.py new file mode 100644 index 000000000..239ff2b90 --- /dev/null +++ b/tests/util/test_blockstore.py @@ -0,0 +1,44 @@ +import hashlib +import pathlib + +import pytest + +from inbox.util import blockstore + + +@pytest.mark.usefixtures("blockstore_backend") +@pytest.mark.parametrize("blockstore_backend", ["disk", "s3"], indirect=True) +def test_save_to_and_get_from_blockstore(): + data = b"test data" + data_sha256 = hashlib.sha256(data).hexdigest() + blockstore.save_to_blockstore(data_sha256, data) + assert blockstore.get_from_blockstore(data_sha256) == data + + +@pytest.fixture +def tiny_email_data() -> bytes: + return (pathlib.Path(__file__).parent / "tiny.eml").read_bytes() + + +@pytest.mark.usefixtures("blockstore_backend") +@pytest.mark.parametrize("blockstore_backend", ["disk", "s3"], indirect=True) +def test_save_and_get_raw_mime_no_compression(tiny_email_data): + data_sha256 = hashlib.sha256(tiny_email_data).hexdigest() + stored_length = blockstore.save_raw_mime( + data_sha256, tiny_email_data, compress=False + ) + + assert stored_length == len(tiny_email_data) + assert blockstore.get_raw_mime(data_sha256) == tiny_email_data + + +@pytest.mark.usefixtures("blockstore_backend") +@pytest.mark.parametrize("blockstore_backend", ["disk", "s3"], indirect=True) +def test_save_and_get_raw_mime_with_compression(tiny_email_data): + data_sha256 = hashlib.sha256(tiny_email_data).hexdigest() + stored_length = blockstore.save_raw_mime( + data_sha256, tiny_email_data, compress=True + ) + + assert stored_length < len(tiny_email_data) + assert blockstore.get_raw_mime(data_sha256) == tiny_email_data diff --git a/tests/util/tiny.eml b/tests/util/tiny.eml new file mode 100644 index 000000000..ef57af496 --- /dev/null +++ b/tests/util/tiny.eml @@ -0,0 +1,60 @@ +Return-Path: +Received: from [192.168.0.12] ([47.61.146.222]) + by smtp.gmail.com with ESMTPSA id 5b1f17b1804b1-4279b1220f5sm6734315e9.14.2024.07.11.00.09.45 + (version=TLS1_3 cipher=TLS_AES_128_GCM_SHA256 bits=128/128); + Thu, 11 Jul 2024 00:09:45 -0700 (PDT) +Content-Type: multipart/mixed; boundary="------------PkQVzddmc0xE3KbbCmLoAoLd" +Message-ID: <085e74fe-aac3-4f61-a735-98bdb223f0ec@gmail.com> +Date: Thu, 11 Jul 2024 09:09:44 +0200 +MIME-Version: 1.0 +User-Agent: Mozilla Thunderbird +From: Example +Subject: Tiny Thunderbird +To: example@gmail.com +Content-Language: en-US + +This is a multi-part message in MIME format. +--------------PkQVzddmc0xE3KbbCmLoAoLd +Content-Type: multipart/alternative; + boundary="------------KIrHVhqjoMvCNpHdNK2emQ8B" +--------------KIrHVhqjoMvCNpHdNK2emQ8B +Content-Type: text/plain; charset=UTF-8; format=flowed +Content-Transfer-Encoding: 7bit +Text +Tiny png +--------------KIrHVhqjoMvCNpHdNK2emQ8B +Content-Type: multipart/related; + boundary="------------Uxc0eFahzcbfaq5vs0NlRK0z" +--------------Uxc0eFahzcbfaq5vs0NlRK0z +Content-Type: text/html; charset=UTF-8 +Content-Transfer-Encoding: 7bit + + + + + + + +

Text

+ Tiny png + + +--------------Uxc0eFahzcbfaq5vs0NlRK0z +Content-Type: image/png; name="tiny.png" +Content-Disposition: inline; filename="tiny.png" +Content-Id: +Content-Transfer-Encoding: base64 +iVBORw0KGgoAAAANSUhEUgAAAAEAAAABAQAAAAA3bvkkAAAACklEQVR4AWNgAAAAAgABc3UB +GAAAAABJRU5ErkJggg== +--------------Uxc0eFahzcbfaq5vs0NlRK0z-- + + +--------------KIrHVhqjoMvCNpHdNK2emQ8B-- +--------------PkQVzddmc0xE3KbbCmLoAoLd +Content-Type: image/gif; name="tiny.gif" +Content-Disposition: attachment; filename="tiny.gif" +Content-Transfer-Encoding: base64 +R0lGODlhAQABAIABAP///wAAACwAAAAAAQABAAACAkQBADs= +--------------PkQVzddmc0xE3KbbCmLoAoLd-- \ No newline at end of file From fb93daa12808013b75000b1511a9e921785e8696 Mon Sep 17 00:00:00 2001 From: Squeaky Date: Tue, 10 Sep 2024 12:57:30 +0200 Subject: [PATCH 4/4] Comments and docstrings --- inbox/util/blockstore.py | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/inbox/util/blockstore.py b/inbox/util/blockstore.py index 85cabf456..e71a736ab 100644 --- a/inbox/util/blockstore.py +++ b/inbox/util/blockstore.py @@ -19,6 +19,10 @@ from boto.s3.key import Key # https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#zstandard-frames +# > This value was selected to be less probable to find at the beginning of some random file. +# > It avoids trivial patterns (0x00, 0xFF, repeated bytes, increasing bytes, etc.), +# > contains byte values outside of ASCII range, and doesn't map into UTF8 space. +# > It reduces the chances that a text file represent this value by accident. ZSTD_MAGIC_NUMBER_PREFIX = 0xFD2FB528.to_bytes(4, "little") @@ -35,10 +39,28 @@ def _data_file_path(h): def save_raw_mime( data_sha256: str, decompressed_raw_mime: bytes, *, compress: "bool | None" = None ) -> int: + """ + Save the raw MIME data to the blockstore, optionally compressing it. + + Args: + data_sha256: The SHA256 hash of the *uncompressed* data. + decompressed_raw_mime: The raw MIME data. + compress: + Whether to compress the data before storing it. + If None, the value of `config["COMPRESS_RAW_MIME"]` is used + which defaults to False. + + Returns: + The length of the data in the datastore. + """ if compress is None: compress = config.get("COMPRESS_RAW_MIME", False) if compress: + # Raw MIME data will never start with the ZSTD magic number, + # because email messages always start with headers in 7-bit ASCII. + # ZSTD magic number contains bytes with the highest bit set to 1, + # so we can use it as a marker to check if the data is compressed. assert not decompressed_raw_mime.startswith(ZSTD_MAGIC_NUMBER_PREFIX) compressed_raw_mime = zstandard.compress(decompressed_raw_mime) @@ -46,6 +68,9 @@ def save_raw_mime( assert compressed_raw_mime.startswith(ZSTD_MAGIC_NUMBER_PREFIX) if len(compressed_raw_mime) > len(decompressed_raw_mime): + # This will not happen in practice, since even the most trivial email + # these days will have a lot of headers that can be compressed. + # But if it does, we should always store the smallest possible representation. compressed_raw_mime = decompressed_raw_mime else: compressed_raw_mime = decompressed_raw_mime @@ -134,10 +159,25 @@ def get_from_blockstore(data_sha256, *, check_sha=True) -> Optional[bytes]: def get_raw_mime(data_sha256: str) -> "bytes | None": + """ + Get the raw MIME data from the blockstore. + + The data may be compressed, so this function will decompress it if necessary. + + Args: + data_sha256: The SHA256 hash of the *uncompressed* data. + + Returns: + The raw MIME data, or None if it wasn't found. + """ compressed_raw_mime = get_from_blockstore(data_sha256, check_sha=False) if compressed_raw_mime is None: return None + # Raw MIME data will never start with the ZSTD magic number, + # because email messages always start with headers in 7-bit ASCII. + # ZSTD magic number contains bytes with the highest bit set to 1, + # so we can use it as a marker to check if the data is compressed. if compressed_raw_mime.startswith(ZSTD_MAGIC_NUMBER_PREFIX): decompressed_raw_mime = zstandard.decompress(compressed_raw_mime) else: