Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 17 additions & 8 deletions src/backend/core/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -2589,15 +2589,30 @@ 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,
raw_content=encrypted_content,
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,
Expand All @@ -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,
Expand Down
57 changes: 39 additions & 18 deletions src/backend/core/services/importer/mbox.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,6 @@
"""

# pylint: disable=broad-exception-caught
import io
from collections.abc import Callable
from dataclasses import dataclass
from datetime import datetime
Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down
15 changes: 14 additions & 1 deletion src/backend/core/services/importer/pst.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading