fix(replays): Bound recording segment decompression at ingest

src/sentry/options/defaults.py

@@ -628,6 +628,14 @@
     default=False,
     flags=FLAG_AUTOMATOR_MODIFIABLE,
 )
+# Maximum number of bytes a single recording segment may decompress to at ingest. Segments
+# which inflate past this limit are dropped instead of stored.
+register(
+    "replay.consumer.max-segment-decompressed-size",
+    type=Int,
+    default=100 * 1024 * 1024,
+    flags=FLAG_AUTOMATOR_MODIFIABLE,
+)
 # Enable new database query caching.
 register(
     "replay.consumer.enable_new_query_caching_system",

src/sentry/replays/consumers/recording.py

@@ -1,6 +1,6 @@
 import logging
 import zlib
-from typing import cast
+from typing import NoReturn, cast
 
 import sentry_sdk
 from sentry_kafka_schemas.codecs import Codec, ValidationError
@@ -56,7 +56,9 @@ def process_message(message: bytes) -> ProcessedEvent | None:
 def parse_recording_event(message: bytes) -> Event:
     recording = parse_request_message(message)
     segment_id, payload = parse_headers(cast(bytes, recording["payload"]), recording["replay_id"])
-    compressed, decompressed = decompress_segment(payload)
+    compressed, decompressed = decompress_segment(
+        payload, max_size=options.get("replay.consumer.max-segment-decompressed-size")
+    )
 
     replay_event_json = recording.get("replay_event")
     if replay_event_json:
@@ -110,16 +112,46 @@ def parse_request_message(message: bytes) -> ReplayRecording:
 
 
 @traces.trace
-def decompress_segment(segment: bytes) -> tuple[bytes, bytes]:
+def decompress_segment(segment: bytes, max_size: int) -> tuple[bytes, bytes]:
+    """Return the compressed and decompressed forms of a recording segment.
+
+    Decompression is bounded by `max_size` bytes. Stored segments are inflated again on every
+    read, so a segment with an implausible compression ratio is a durable, repeatable cost on
+    the read path -- we drop it here rather than persist it.
+    """
     try:
-        return (segment, zlib.decompress(segment))
+        # Ask for one byte past the limit so an oversized segment is detectable without ever
+        # materializing more than `max_size + 1` bytes.
+        decompressor = zlib.decompressobj()
+        decompressed = decompressor.decompress(segment, max_size + 1)
+
+        # The incremental API returns partial output where `zlib.decompress` raised on a
+        # truncated stream. Normalize so malformed bodies take the branch below.
+        if not decompressor.eof and not decompressor.unconsumed_tail:
+            raise zlib.error("incomplete or truncated stream")
     except zlib.error:
         if segment and segment[0] == ord("["):
+            if len(segment) > max_size:
+                _drop_oversized_segment(len(segment), max_size)
             return (zlib.compress(segment), segment)
         else:
             logger.exception("Invalid recording body.")
             raise DropSilently()
 
+    if len(decompressed) > max_size:
+        _drop_oversized_segment(len(segment), max_size)
+
+    return (segment, decompressed)
+
+
+def _drop_oversized_segment(compressed_size: int, max_size: int) -> NoReturn:
+    metrics.incr("replays.consumer.recording.oversized_segment")
+    logger.warning(
+        "Dropped a recording segment which decompressed past the size limit.",
+        extra={"compressed_size": compressed_size, "max_size": max_size},
+    )
+    raise DropSilently()
+
 
 @traces.trace
 def parse_headers(recording: bytes, replay_id: str) -> tuple[int, bytes]:

tests/sentry/replays/integration/consumers/test_recording.py

@@ -9,10 +9,14 @@
 from sentry.replays.tasks import process_replay_recording
 from sentry.utils import json
 
+# Options read on this path are boolean toggles we want enabled, except for the segment size
+# limit which needs a real byte count.
+_OPTIONS: dict[str, object] = {"replay.consumer.max-segment-decompressed-size": 100 * 1024 * 1024}
+
 
 @mock.patch("sentry.options.get")
 def test_recording_task(options_get: MagicMock) -> None:
-    options_get.return_value = True
+    options_get.side_effect = lambda key, *args, **kwargs: _OPTIONS.get(key, True)
 
     headers = json.dumps({"segment_id": 42}).encode()
     recording_payload = headers + b"\n" + zlib.compress(json.dumps(MOCK_EVENTS).encode())

tests/sentry/replays/unit/consumers/test_recording.py

@@ -18,20 +18,23 @@
 from sentry.replays.usecases.ingest import ProcessedEvent
 from sentry.replays.usecases.ingest.event_parser import ParsedEventMeta
 from sentry.replays.usecases.pack import pack
+from sentry.testutils.helpers import override_options
 from sentry.testutils.pytest.fixtures import django_db_all
 from sentry.utils import json
 
+MAX_SIZE = 100 * 1024 * 1024
+
 
 def test_decompress_segment_success() -> None:
     """Test successful decompression of segment"""
     data = b"[hello, world!]"
     compressed_data = zlib.compress(data)
 
-    compressed, decompressed = decompress_segment(compressed_data)
+    compressed, decompressed = decompress_segment(compressed_data, MAX_SIZE)
     assert compressed == compressed_data
     assert decompressed == data
 
-    compressed, decompressed = decompress_segment(data)
+    compressed, decompressed = decompress_segment(data, MAX_SIZE)
     assert compressed == compressed_data
     assert decompressed == data
 
@@ -41,21 +44,97 @@ def test_decompress_segment_already_decompressed() -> None:
     data = b"[hello, world!]"
     compressed_data = zlib.compress(data)
 
-    compressed, decompressed = decompress_segment(data)
+    compressed, decompressed = decompress_segment(data, MAX_SIZE)
     assert compressed == compressed_data
     assert decompressed == data
 
 
 def test_decompress_segment_unexpected_start_character() -> None:
     """Test handling of invalid data that can't be decompressed"""
     with pytest.raises(DropSilently):
-        decompress_segment(b"hello, world!")
+        decompress_segment(b"hello, world!", MAX_SIZE)
 
 
 def test_decompress_segment_empty_data() -> None:
     """Test handling of empty data"""
     with pytest.raises(DropSilently):
-        decompress_segment(b"")
+        decompress_segment(b"", MAX_SIZE)
+
+
+def test_decompress_segment_truncated_stream() -> None:
+    """Test a truncated zlib stream is dropped rather than returning partial output"""
+    compressed_data = zlib.compress(b"[" + b"hello, world!" * 100 + b"]")
+
+    for cut in (1, 4, len(compressed_data) // 2):
+        with pytest.raises(DropSilently):
+            decompress_segment(compressed_data[:-cut], MAX_SIZE)
+
+
+def test_decompress_segment_invalid_checksum() -> None:
+    """Test a zlib stream with a corrupted Adler-32 checksum is dropped"""
+    compressed_data = bytearray(zlib.compress(b"[hello, world!]"))
+    compressed_data[-1] ^= 0xFF
+
+    with pytest.raises(DropSilently):
+        decompress_segment(bytes(compressed_data), MAX_SIZE)
+
+
+def test_decompress_segment_exceeds_max_size() -> None:
+    """Test a segment which inflates past the limit is dropped"""
+    payload = b"[" + b" " * (1024 * 1024) + b"]"
+    bomb = zlib.compress(payload, 9)
+    # A kilobyte of stored input, a megabyte of output on every read.
+    assert len(payload) // len(bomb) > 500
+
+    with pytest.raises(DropSilently):
+        decompress_segment(bomb, 1024)
+
+
+def test_decompress_segment_at_max_size() -> None:
+    """Test a segment which inflates to exactly the limit is accepted"""
+    data = b"[hello, world!]"
+    compressed_data = zlib.compress(data)
+
+    compressed, decompressed = decompress_segment(compressed_data, len(data))
+    assert compressed == compressed_data
+    assert decompressed == data
+
+    with pytest.raises(DropSilently):
+        decompress_segment(compressed_data, len(data) - 1)
+
+
+def test_decompress_segment_uncompressed_exceeds_max_size() -> None:
+    """Test an uncompressed segment larger than the limit is dropped"""
+    data = b"[" + b" " * 1024 + b"]"
+
+    with pytest.raises(DropSilently):
+        decompress_segment(data, 512)
+
+
+@django_db_all
+def test_process_message_oversized_segment() -> None:
+    """Test "process_message" function drops a segment which inflates past the limit."""
+    bomb = zlib.compress(b"[" + b" " * (1024 * 1024) + b"]", 9)
+    message = {
+        "type": "replay_recording_not_chunked",
+        "org_id": 3,
+        "project_id": 4,
+        "replay_id": "1",
+        "received": 2,
+        "retention_days": 30,
+        "payload": json.dumps({"segment_id": 42}).encode() + b"\n" + bomb,
+        "key_id": 1,
+        "replay_event": b"{}",
+        "replay_video": b"",
+        "version": 0,
+    }
+
+    kafka_message = make_kafka_message(message)
+    with override_options({"replay.consumer.max-segment-decompressed-size": 1024}):
+        assert process_message(kafka_message) is None
+
+    # The same message is accepted under the default limit.
+    assert process_message(kafka_message) is not None
 
 
 def test_parse_headers_success() -> None:
@@ -121,6 +200,7 @@ def test_parse_request_message_validation_error() -> None:
         parse_request_message(msgpack.packb(b"invalid"))
 
 
+@django_db_all
 def test_parse_recording_event_success() -> None:
     """Test successful parsing of recording event end-to-end"""
     # Create real compressed data
@@ -168,6 +248,7 @@ def test_parse_recording_event_success() -> None:
     assert result == expected
 
 
+@django_db_all
 def test_parse_recording_event_with_replay_event() -> None:
     """Test parsing recording event with replay_event JSON"""
     # Create real compressed data
@@ -232,6 +313,7 @@ def test_parse_recording_event_missing_payload() -> None:
         parse_recording_event(msgpack.packb(message))
 
 
+@django_db_all
 def test_parse_recording_event_invalid_compression() -> None:
     """Test that invalid compression in payload causes DropSilently"""
     message = {