From b1778c24c7d065a3d0114ef39ef9b754323a7274 Mon Sep 17 00:00:00 2001 From: Sylvain Zimmer Date: Thu, 23 Jul 2026 01:10:06 +0200 Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B(imports)=20harden=20imports=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/backend/core/models.py | 25 +++++++--- src/backend/core/services/importer/mbox.py | 57 +++++++++++++++------- src/backend/core/services/importer/pst.py | 15 +++++- 3 files changed, 70 insertions(+), 27 deletions(-) diff --git a/src/backend/core/models.py b/src/backend/core/models.py index cb796c7c3..609612bdd 100644 --- a/src/backend/core/models.py +++ b/src/backend/core/models.py @@ -2589,7 +2589,13 @@ def create_blob( # the content-addressed object lands first, then the row that # points at it commits. A crash in between leaves only an # idempotently-overwritable object for ``verify_blobs`` to - # reconcile. Any upload failure falls back to the PG tier. + # reconcile. Only the *upload* falls back to the PG tier — the + # object-storage row insert stays outside this handler so a DB + # error propagates and rolls the transaction back normally. + # (Swallowing it here and retrying ``self.create`` below would + # issue a second query inside an already-broken atomic block and + # raise ``TransactionManagementError``.) + uploaded = None try: staged = self.model( sha256=sha256_hash, @@ -2597,7 +2603,16 @@ def create_blob( encryption_key_id=encryption_key_id, compression=compression, ) - key_id, compression_used = service.upload_blob(staged) + uploaded = service.upload_blob(staged) + except Exception: # pylint: disable=broad-exception-caught + logger.warning( + "Direct-to-object-storage write failed; storing blob " + "in the PG tier instead (offload will retry later)", + exc_info=True, + ) + + if uploaded is not None: + key_id, compression_used = uploaded return self.create( sha256=sha256_hash, size=original_size, @@ -2609,12 +2624,6 @@ def create_blob( encryption_key_id=key_id, **kwargs, ) - except Exception: # pylint: disable=broad-exception-caught - logger.warning( - "Direct-to-object-storage write failed; storing blob " - "in the PG tier instead (offload will retry later)", - exc_info=True, - ) blob = self.create( sha256=sha256_hash, diff --git a/src/backend/core/services/importer/mbox.py b/src/backend/core/services/importer/mbox.py index f8914de8e..11d24475e 100644 --- a/src/backend/core/services/importer/mbox.py +++ b/src/backend/core/services/importer/mbox.py @@ -6,7 +6,6 @@ """ # pylint: disable=broad-exception-caught -import io from collections.abc import Callable from dataclasses import dataclass from datetime import datetime @@ -127,13 +126,11 @@ def index_mbox_messages( # Handle last message if message_start is not None: - # Get file end position - current_pos = file.tell() - file.seek(0, io.SEEK_END) - file_end = file.tell() - total_end = file_end - 1 - # Restore position for _extract_and_store_index - file.seek(current_pos) + # The scan loop only breaks at EOF, so ``buffer`` holds every byte from + # ``file_offset`` to the end of the file — the final offset follows from + # the tracked scan state without a seek. That keeps the documented + # optional-seek contract: a valid non-seekable stream stays indexable. + total_end = file_offset + len(buffer) - 1 if total_end >= message_start: _extract_and_store_index( file, @@ -147,32 +144,56 @@ def index_mbox_messages( return indices +# A message's Date can sit past a long Received/DKIM/ARC preamble, so header +# parsing must not stop at a fixed byte cut-off. Scanning still stops at the +# blank-line terminator (or this cap), so a malformed message with no terminator +# can't pull its whole body into memory just to look for a Date. +_MAX_HEADER_SCAN = 65536 + + def _extract_and_store_index( file, indices, msg_start, msg_end, buffer, buf_file_offset ): """Extract date from a message and add an index entry.""" - # Try to read first 2048 bytes of the message for header parsing - header_size = min(2048, msg_end - msg_start + 1) + header_bytes = _read_message_headers( + file, msg_start, msg_end, buffer, buf_file_offset + ) + date = extract_date_from_headers(header_bytes) + indices.append(MboxMessageIndex(start_byte=msg_start, end_byte=msg_end, date=date)) + - # Check if the header bytes are in our buffer +def _read_message_headers(file, msg_start, msg_end, buffer, buf_file_offset): + """Return the message's header block: bytes from ``msg_start`` up to and + including the first blank-line terminator, growing the read until the + terminator is found (or ``_MAX_HEADER_SCAN`` / end-of-message is reached). + + Reuses the caller's ``buffer`` when it already spans the requested range and + only falls back to a seek+read otherwise, restoring the file position. + """ + msg_len = msg_end - msg_start + 1 buf_start = buf_file_offset buf_end = buf_start + len(buffer) - 1 - if buf_start <= msg_start and msg_start + header_size - 1 <= buf_end: - offset_in_buf = msg_start - buf_start - header_bytes = buffer[offset_in_buf : offset_in_buf + header_size] - else: + def _slice(size): + if buf_start <= msg_start and msg_start + size - 1 <= buf_end: + offset_in_buf = msg_start - buf_start + return buffer[offset_in_buf : offset_in_buf + size] # Need to seek and read current_pos = file.tell() if hasattr(file, "tell") else None try: file.seek(msg_start) - header_bytes = file.read(header_size) + return file.read(size) finally: if current_pos is not None: file.seek(current_pos) - date = extract_date_from_headers(header_bytes) - indices.append(MboxMessageIndex(start_byte=msg_start, end_byte=msg_end, date=date)) + cap = min(msg_len, _MAX_HEADER_SCAN) + size = min(2048, cap) + while True: + chunk = _slice(size) + if b"\r\n\r\n" in chunk or b"\n\n" in chunk or size >= cap: + return chunk + size = min(size * 2, cap) def _mbox_plan( diff --git a/src/backend/core/services/importer/pst.py b/src/backend/core/services/importer/pst.py index 011688a6c..fd9195d7c 100644 --- a/src/backend/core/services/importer/pst.py +++ b/src/backend/core/services/importer/pst.py @@ -1192,7 +1192,20 @@ def reconstruct_eml( # of them) beyond the per-message size limit. A crafted PST could otherwise # hold a multi-GB attachment that OOMs the worker before ``deliver`` — which # would reject the oversized message anyway — ever gets to check the size. - attachment_budget = settings.MAX_INCOMING_EMAIL_SIZE + # + # The budget is tracked in *raw* attachment bytes, but the assembled EML + # base64-encodes every attachment (~4/3 inflation, plus line wrapping and + # per-part MIME headers) and also carries the body and boundaries. Deflate + # the raw budget by that overhead and reserve room for the body so raw data + # alone can't consume the whole limit — otherwise an attachment just under + # MAX_INCOMING_EMAIL_SIZE passes here yet builds an EML that ``deliver`` + # rejects for size (after ballooning memory past the intended cap). + _ENCODE_OVERHEAD = 1.4 # base64 (~1.37x) + line wrapping + per-part headers + _BODY_RESERVE = 64 * 1024 # text/html body + top-level MIME boundaries + attachment_budget = max( + 0, + int((settings.MAX_INCOMING_EMAIL_SIZE - _BODY_RESERVE) / _ENCODE_OVERHEAD), + ) try: num_attachments = message.number_of_attachments for i in range(num_attachments):