fix(gmail-bulk-import): harden re-import against 8-bit headers + poison batches

Four audit findings (2026-07-14), each reproduced on crafted mboxes; the
Takeout corpus has proven 8-bit header bytes, so all are real re-import risks.

1. Whole-run crash on 8-bit Message-ID. compat32 .get() returns an
   email.header.Header (not str) for raw 8-bit bytes; the old
   Header.strip() raised AttributeError. _message_id ran BEFORE the
   per-message try, so one bad header killed the entire import.
   Fix: str() + sanitize_surrogates() before strip; and move
   _message_id/_parse_date/_parse_attachments INSIDE the per-message try —
   a broken message is now errors += 1, never run death.

2. Poison batch. pending.clear() ran only AFTER a successful insert, so a
   failed flush (DB down / bad row) left pending intact and every later
   message re-flushed the doomed batch; the final flush sat in try/finally
   with no except and propagated out, losing all stats. Fix: _flush always
   clears pending and counts a failed insert as db_insert_failed; the run
   always reaches import_complete.

3. Stats didn't reconcile with the DB. imported counts archive writes, not
   DB rows, so a partial-insert drift was invisible. Fix: separate
   db_inserted/db_insert_failed counters; main() exits non-zero on any
   error, DB drift, or a processed = imported + skipped + errors imbalance.

4. 8-bit Date → needless epoch_fallback. parsedate_to_datetime(Header)
   raised even when str(header) parses fine. Fix: str() before the epoch
   fallback.

Shared helper: _sanitize moved from gmail-header-backfill into
packages/kb-mail (kb_mail.text.sanitize_surrogates) and used by both jobs;
gmail-header-backfill now depends on kb-mail.

Tests: regression coverage for all four findings in gmail-bulk-import
(8-bit id/date, per-message guard, failed-insert non-poisoning, stats
balance) plus kb_mail.text unit tests. Full suites green:
kb-mail 27, gmail-bulk-import 33, gmail-header-backfill 43.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
oskar 2026-07-14 19:51:35 +02:00
parent 4746ebe0fb
commit 3755d7d6bb
7 changed files with 348 additions and 37 deletions

View file

@ -43,6 +43,7 @@ import structlog
from kb_mail.archive import save_eml from kb_mail.archive import save_eml
from kb_mail.envelope import Envelope from kb_mail.envelope import Envelope
from kb_mail.text import sanitize_surrogates as _sanitize
_log = structlog.get_logger(__name__) _log = structlog.get_logger(__name__)
@ -60,16 +61,28 @@ def _eml_ref(envelope_id: str, ts: datetime) -> str:
def _message_id(msg: mailbox.mboxMessage) -> str: def _message_id(msg: mailbox.mboxMessage) -> str:
"""Return a stable envelope id from Message-ID header or SHA-256 content hash.""" """Return a stable envelope id from Message-ID header or SHA-256 content hash.
mid = msg.get("Message-ID", "").strip().strip("<>")
compat32 .get() returns an email.header.Header (not str) when the raw value
holds 8-bit bytes str() + _sanitize keeps the id jsonb-safe and gives
.strip() a real string to work on. An unguarded .strip() on a Header raised
AttributeError and, because _message_id runs before the per-message try,
killed the whole run (the archive has proven 8-bit header bytes).
"""
mid = _sanitize(str(msg.get("Message-ID", ""))).strip().strip("<>")
if mid: if mid:
return mid return mid
return "sha256-" + hashlib.sha256(msg.as_bytes()).hexdigest()[:32] return "sha256-" + hashlib.sha256(msg.as_bytes()).hexdigest()[:32]
def _parse_date(msg: mailbox.mboxMessage) -> datetime: def _parse_date(msg: mailbox.mboxMessage) -> datetime:
"""Parse Date header to UTC-aware datetime. Returns epoch as last resort.""" """Parse Date header to UTC-aware datetime. Returns epoch as last resort.
date_str = msg.get("Date", "")
str() first: a compat32 8-bit Date surfaces as an email.header.Header, and
parsedate_to_datetime() raises on a Header ( needless epoch_fallback) even
when str(header) is perfectly parseable.
"""
date_str = str(msg.get("Date", ""))
if date_str: if date_str:
try: try:
dt = parsedate_to_datetime(date_str) dt = parsedate_to_datetime(date_str)
@ -152,8 +165,17 @@ async def run_import(
) -> dict[str, int]: ) -> dict[str, int]:
"""Import all messages from an mbox file. """Import all messages from an mbox file.
Returns stats: {processed, imported, skipped, errors, epoch_fallback, Returns stats: {processed, imported, skipped, errors, db_inserted,
msgs_with_attachments, total_attachments, total_attachment_bytes}. db_insert_failed, epoch_fallback, msgs_with_attachments, total_attachments,
total_attachment_bytes}.
Two invariants always hold (checked by main() for the exit code):
* processed == imported + skipped + errors
* with a DB connection: db_inserted + db_insert_failed == imported + skipped
`imported`/`skipped` count archive writes; `db_inserted`/`db_insert_failed`
count DB rows. They are tracked separately because a failed batch insert
leaves the archive ahead of the DB, and nothing else would catch the drift.
dry_run=True parses and counts everything without writing. dry_run=True parses and counts everything without writing.
limit=N stops after processing N messages (for sampling). limit=N stops after processing N messages (for sampling).
@ -163,6 +185,8 @@ async def run_import(
"imported": 0, "imported": 0,
"skipped": 0, "skipped": 0,
"errors": 0, "errors": 0,
"db_inserted": 0,
"db_insert_failed": 0,
"epoch_fallback": 0, "epoch_fallback": 0,
"msgs_with_attachments": 0, "msgs_with_attachments": 0,
"total_attachments": 0, "total_attachments": 0,
@ -176,8 +200,17 @@ async def run_import(
pending: list[Envelope] = [] pending: list[Envelope] = []
async def _flush() -> None: async def _flush() -> None:
# Always clear pending, even when the insert fails: a poisoned batch
# (DB down, one bad row) must never be re-flushed by every subsequent
# message. A failed insert leaves the archive ahead of the DB, counted
# as db_insert_failed so the drift is visible instead of silent.
if conn is not None and pending: if conn is not None and pending:
try:
await _insert_batch(conn, pending) await _insert_batch(conn, pending)
stats["db_inserted"] += len(pending)
except Exception:
_log.exception("db_insert_failed", count=len(pending))
stats["db_insert_failed"] += len(pending)
pending.clear() pending.clear()
try: try:
@ -188,6 +221,11 @@ async def run_import(
stats["processed"] += 1 stats["processed"] += 1
# The ENTIRE per-message body is guarded: a single malformed
# message (8-bit header, unwalkable MIME, bad save) counts as one
# error and never aborts the run.
envelope_id = "<unknown>"
try:
envelope_id = _message_id(msg) envelope_id = _message_id(msg)
ts = _parse_date(msg) ts = _parse_date(msg)
@ -205,7 +243,6 @@ async def run_import(
stats["imported"] += 1 stats["imported"] += 1
continue continue
try:
raw = msg.as_bytes(policy=email.policy.compat32) raw = msg.as_bytes(policy=email.policy.compat32)
try: try:
raw_ref = await save_eml(archive_root, envelope_id, "gmail", ts, raw) raw_ref = await save_eml(archive_root, envelope_id, "gmail", ts, raw)
@ -228,6 +265,9 @@ async def run_import(
if stats["processed"] % 500 == 0: if stats["processed"] % 500 == 0:
_log.info("progress", **stats) _log.info("progress", **stats)
# Final flush swallows its own insert error (see _flush) so the run
# always reaches import_complete with balanced stats instead of dying
# here with the counters lost.
await _flush() await _flush()
finally: finally:
@ -262,7 +302,17 @@ def main() -> None:
stats = asyncio.run( stats = asyncio.run(
run_import(args.mbox, args.archive, args.dsn or None, args.dry_run, args.limit) run_import(args.mbox, args.archive, args.dsn or None, args.dry_run, args.limit)
) )
sys.exit(1 if stats["errors"] > 0 else 0)
balanced = (
stats["processed"] == stats["imported"] + stats["skipped"] + stats["errors"]
)
if not balanced:
_log.error("stats_unbalanced", **stats)
# Non-zero exit on any message failure, any DB drift (archive ahead of DB),
# or a stats imbalance — the run must not report clean when it isn't.
failed = stats["errors"] > 0 or stats["db_insert_failed"] > 0 or not balanced
sys.exit(1 if failed else 0)
if __name__ == "__main__": if __name__ == "__main__":

View file

@ -2,6 +2,7 @@
from __future__ import annotations from __future__ import annotations
import hashlib import hashlib
import json
import mailbox import mailbox
from datetime import datetime, timezone from datetime import datetime, timezone
from email import encoders from email import encoders
@ -12,6 +13,7 @@ from pathlib import Path
import pytest import pytest
from gmail_bulk_import import importer
from gmail_bulk_import.importer import ( from gmail_bulk_import.importer import (
_message_id, _message_id,
_parse_attachments, _parse_attachments,
@ -281,3 +283,211 @@ class TestRunImport:
assert stats["imported"] == 1 assert stats["imported"] == 1
assert stats["total_attachments"] == 2 assert stats["total_attachments"] == 2
# ---------------------------------------------------------------------------
# Regression tests for the 2026-07-14 hardening audit (four fixes).
#
# The audit reproduced each of these on crafted mboxes; the corpus is proven to
# contain raw 8-bit header bytes, so every scenario below is a real re-import
# risk, not a hypothetical.
# ---------------------------------------------------------------------------
# Raw 8-bit byte (0xe9) inside Message-ID: compat32 .get() returns an
# email.header.Header, whose .strip() raised AttributeError. Because
# _message_id ran BEFORE the per-message try, one such header killed the run.
_EIGHTBIT_MID_EML = (
b"From: alice@example.com\r\n"
b"Message-ID: <bad\xe9id@example.com>\r\n"
b"Date: Tue, 10 Jun 2025 12:00:00 +0000\r\n"
b"Subject: hi\r\n\r\nbody"
)
# Raw 8-bit byte inside the Date header's trailing comment: .get() returns a
# Header, parsedate_to_datetime(Header) raises → needless epoch_fallback, even
# though str(header) parses to a real date (2009-05-19 10:27:09 +0200).
_EIGHTBIT_DATE_EML = (
b"From: alice@example.com\r\n"
b"Message-ID: <clean-date@example.com>\r\n"
b"Date: Tue, 19 May 2009 10:27:09 +0200 (Ho\xe9ra)\r\n"
b"Subject: hi\r\n\r\nbody"
)
def _make_mbox_raw(path: Path, raw_messages: list[bytes]) -> None:
mbox = mailbox.mbox(str(path), create=True)
for raw in raw_messages:
mbox.add(mailbox.mboxMessage(raw))
mbox.flush()
mbox.close()
class _FakeConn:
"""Stand-in asyncpg connection: records executemany rows or fails them."""
def __init__(self, fail: bool = False):
self.fail = fail
self.inserted: list = []
self.executemany_calls = 0
async def executemany(self, query, rows):
self.executemany_calls += 1
if self.fail:
raise RuntimeError("DB down")
self.inserted.extend(list(rows))
async def close(self):
pass
def _patch_connect(monkeypatch, conn: _FakeConn) -> None:
async def _fake_connect(dsn):
return conn
monkeypatch.setattr(importer.asyncpg, "connect", _fake_connect)
class TestEightBitMessageId:
def test_message_id_8bit_header_does_not_crash(self):
# Fix 1: str() + _sanitize gives .strip() a real string. Old code did
# Header.strip() → AttributeError.
msg = mailbox.mboxMessage(_EIGHTBIT_MID_EML)
mid = _message_id(msg)
assert isinstance(mid, str)
assert mid.endswith("id@example.com")
json.dumps(mid).encode("utf-8") # jsonb-safe: no lone surrogate
async def test_run_survives_8bit_message_id(self, tmp_path):
# Old code raised out of run_import here; the run must now complete and
# count the message normally.
mbox_path = tmp_path / "mail.mbox"
archive = tmp_path / "archive"
_make_mbox_raw(mbox_path, [_EIGHTBIT_MID_EML, _msg_with_id("<ok@e.com>").as_bytes()])
stats = await run_import(mbox_path, archive)
assert stats["processed"] == 2
assert stats["imported"] == 2
assert stats["errors"] == 0
class TestPerMessageGuard:
async def test_one_bad_message_counts_once_and_run_continues(self, tmp_path, monkeypatch):
# Fix 2: the parse helpers now live inside the per-message try, so a
# message that raises is one error, not the death of the run.
calls = {"n": 0}
real_message_id = importer._message_id
def flaky(msg):
calls["n"] += 1
if calls["n"] == 1:
raise RuntimeError("boom parsing first message")
return real_message_id(msg)
monkeypatch.setattr(importer, "_message_id", flaky)
mbox_path = tmp_path / "mail.mbox"
archive = tmp_path / "archive"
_make_mbox(mbox_path, [_msg_with_id("<a@e.com>"), _msg_with_id("<b@e.com>")])
stats = await run_import(mbox_path, archive)
assert stats["processed"] == 2
assert stats["errors"] == 1
assert stats["imported"] == 1
class TestFailedInsertDoesNotPoison:
async def test_final_flush_failure_does_not_kill_run(self, tmp_path, monkeypatch):
# Fix 3: the final flush used to be uncaught (try/finally, no except),
# so a failing insert lost all stats. Now it is swallowed and counted.
conn = _FakeConn(fail=True)
_patch_connect(monkeypatch, conn)
mbox_path = tmp_path / "mail.mbox"
archive = tmp_path / "archive"
_make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(3)])
stats = await run_import(mbox_path, archive, dsn="postgresql://fake")
# Run completes, archive is ahead of DB, and the drift is visible.
assert stats["imported"] == 3
assert stats["db_inserted"] == 0
assert stats["db_insert_failed"] == 3
assert stats["errors"] == 0
async def test_poisoned_batch_does_not_re_flush_forever(self, tmp_path, monkeypatch):
# Fix 3: pending is cleared even on insert failure, so a poisoned batch
# is not re-flushed by every subsequent message. With BATCH_SIZE=1 each
# message flushes exactly once; old code re-flushed a growing pending
# list and died in the uncaught final flush.
monkeypatch.setattr(importer, "BATCH_SIZE", 1)
conn = _FakeConn(fail=True)
_patch_connect(monkeypatch, conn)
mbox_path = tmp_path / "mail.mbox"
archive = tmp_path / "archive"
_make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(3)])
stats = await run_import(mbox_path, archive, dsn="postgresql://fake")
assert conn.executemany_calls == 3 # one per message, not a growing batch
assert stats["imported"] == 3
assert stats["db_insert_failed"] == 3
assert stats["errors"] == 0
class TestStatsBalance:
async def test_processed_balances_and_db_reconciles(self, tmp_path, monkeypatch):
# Fix 4: imported/skipped count archive writes; db_inserted counts DB
# rows. Both invariants must always hold.
conn = _FakeConn()
_patch_connect(monkeypatch, conn)
mbox_path = tmp_path / "mail.mbox"
archive = tmp_path / "archive"
_make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(3)])
stats1 = await run_import(mbox_path, archive, dsn="postgresql://fake")
# second run over the same archive: every message is skipped, but each
# still goes to the DB (ON CONFLICT DO NOTHING)
conn2 = _FakeConn()
_patch_connect(monkeypatch, conn2)
stats2 = await run_import(mbox_path, archive, dsn="postgresql://fake")
for s in (stats1, stats2):
assert s["processed"] == s["imported"] + s["skipped"] + s["errors"]
assert s["db_inserted"] + s["db_insert_failed"] == s["imported"] + s["skipped"]
assert stats1["imported"] == 3 and stats1["db_inserted"] == 3
assert stats2["skipped"] == 3 and stats2["db_inserted"] == 3
async def test_archive_only_mode_leaves_db_counters_zero(self, tmp_path):
mbox_path = tmp_path / "mail.mbox"
archive = tmp_path / "archive"
_make_mbox(mbox_path, [_msg_with_id(f"<m{i}@e.com>") for i in range(2)])
stats = await run_import(mbox_path, archive, dsn=None)
assert stats["imported"] == 2
assert stats["db_inserted"] == 0
assert stats["db_insert_failed"] == 0
assert stats["processed"] == stats["imported"] + stats["skipped"] + stats["errors"]
class TestParseDate8Bit:
def test_8bit_date_header_parses_via_str_not_epoch(self):
# Fix 4 (date): str(header) before epoch_fallback. Old code fed the
# Header straight to parsedate_to_datetime → AttributeError → epoch.
msg = mailbox.mboxMessage(_EIGHTBIT_DATE_EML)
ts = _parse_date(msg)
assert ts == datetime(2009, 5, 19, 8, 27, 9, tzinfo=timezone.utc)
async def test_8bit_date_not_counted_as_epoch_fallback(self, tmp_path):
mbox_path = tmp_path / "mail.mbox"
archive = tmp_path / "archive"
_make_mbox_raw(mbox_path, [_EIGHTBIT_DATE_EML])
stats = await run_import(mbox_path, archive)
assert stats["epoch_fallback"] == 0
assert stats["imported"] == 1

View file

@ -9,6 +9,7 @@ requires-python = ">=3.11"
dependencies = [ dependencies = [
"asyncpg>=0.29", "asyncpg>=0.29",
"structlog>=24.1", "structlog>=24.1",
"kb-mail",
] ]
[project.scripts] [project.scripts]

View file

@ -46,6 +46,8 @@ from typing import Optional
import asyncpg import asyncpg
import structlog import structlog
from kb_mail.text import sanitize_surrogates as _sanitize
_log = structlog.get_logger(__name__) _log = structlog.get_logger(__name__)
DEFAULT_ARCHIVE_ROOT = Path("/home/oskar/kb/mail/archive") DEFAULT_ARCHIVE_ROOT = Path("/home/oskar/kb/mail/archive")
@ -126,17 +128,6 @@ def parse_headers(raw: bytes) -> dict:
} }
def _sanitize(value: Optional[str]) -> Optional[str]:
"""Strip surrogates a compat32 bytes-parse can leave in header text.
Lone surrogates are not valid UTF-8 and postgres jsonb rejects them; each
undecodable byte degrades to '?' instead. Real Unicode passes through.
"""
if value is None:
return None
return value.encode("utf-8", errors="replace").decode("utf-8")
def parse_headers_fallback(raw: bytes) -> dict: def parse_headers_fallback(raw: bytes) -> dict:
"""Degraded parse for messages the typed (policy.default) path rejects. """Degraded parse for messages the typed (policy.default) path rejects.

View file

@ -1,5 +1,12 @@
from .archive import save_eml from .archive import save_eml
from .db import get_envelope, insert_envelope from .db import get_envelope, insert_envelope
from .envelope import Envelope from .envelope import Envelope
from .text import sanitize_surrogates
__all__ = ["Envelope", "insert_envelope", "get_envelope", "save_eml"] __all__ = [
"Envelope",
"insert_envelope",
"get_envelope",
"save_eml",
"sanitize_surrogates",
]

View file

@ -0,0 +1,20 @@
"""Text helpers shared across mail-parsing jobs."""
from __future__ import annotations
from typing import Optional
def sanitize_surrogates(value: Optional[str]) -> Optional[str]:
"""Strip surrogates a compat32 bytes-parse can leave in header text.
Lone surrogates are not valid UTF-8 and postgres jsonb rejects them; each
undecodable byte degrades to '?' instead. Real Unicode passes through.
Shared by gmail-header-backfill (header text) and gmail-bulk-import
(Message-ID envelope id): both parse Gmail Takeout .eml bytes under
compat32, where 8-bit header bytes surface as email.header.Header /
surrogate-laden strings that must not reach json.dumps or jsonb.
"""
if value is None:
return None
return value.encode("utf-8", errors="replace").decode("utf-8")

View file

@ -0,0 +1,32 @@
"""Unit tests for kb_mail.text.sanitize_surrogates."""
from __future__ import annotations
import json
from kb_mail.text import sanitize_surrogates
def test_none_stays_none():
assert sanitize_surrogates(None) is None
def test_plain_ascii_unchanged():
assert sanitize_surrogates("hello") == "hello"
def test_real_unicode_passes_through():
assert sanitize_surrogates("Kąpała") == "Kąpała"
def test_lone_surrogate_degraded_and_json_safe():
# A compat32 bytes-parse can leave a lone surrogate (undecodable byte);
# postgres jsonb and json.dumps().encode('utf-8') both reject it.
dirty = "Pr\udce9sent" # 0xe9 smuggled in as a surrogate escape
cleaned = sanitize_surrogates(dirty)
assert cleaned == "Pr?sent"
json.dumps(cleaned, ensure_ascii=False).encode("utf-8") # must not raise
def test_replacement_char_preserved():
# U+FFFD is already valid UTF-8 and must survive untouched.
assert sanitize_surrogates("bad<EFBFBD>id") == "bad<EFBFBD>id"