Compress RAW MIMEs with zstandard

.gitattributes

@@ -1,4 +1,6 @@
 * text eol=lf
+*.eml eol=crlf
 *.png binary
 *.jpg binary
 *.ics binary
+

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:

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

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,13 @@
 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
+# > 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")
+
 
 def _data_file_directory(h):
     return os.path.join(
@@ -27,6 +36,50 @@ 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:
+    """
+    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)
+
+        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
+
+    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 +139,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 +150,44 @@ 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) -> "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:
+        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):

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

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

tests/util/tiny.eml

@@ -0,0 +1,60 @@
+Return-Path: <example@gmail.com>
+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 <example@gmail.com>
+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
+
+<!DOCTYPE html>
+<html>
+  <head>
+    <meta http-equiv="content-type" content="text/html; charset=UTF-8">
+  </head>
+  <body>
+    <p>Text</p>
+    <img moz-do-not-send="false"
+      src="cid:part1.5cIgJBp4.8aHylIhm@gmail.com" alt="Tiny png"
+      class="" width="1" height="1">
+  </body>
+</html>
+--------------Uxc0eFahzcbfaq5vs0NlRK0z
+Content-Type: image/png; name="tiny.png"
+Content-Disposition: inline; filename="tiny.png"
+Content-Id: <part1.5cIgJBp4.8aHylIhm@gmail.com>
+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--