diff --git a/jobs/gmail-bulk-import/src/gmail_bulk_import/importer.py b/jobs/gmail-bulk-import/src/gmail_bulk_import/importer.py index 65735f7..e936b21 100644 --- a/jobs/gmail-bulk-import/src/gmail_bulk_import/importer.py +++ b/jobs/gmail-bulk-import/src/gmail_bulk_import/importer.py @@ -43,6 +43,7 @@ import structlog from kb_mail.archive import save_eml from kb_mail.envelope import Envelope +from kb_mail.text import sanitize_surrogates as _sanitize _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: - """Return a stable envelope id from Message-ID header or SHA-256 content hash.""" - mid = msg.get("Message-ID", "").strip().strip("<>") + """Return a stable envelope id from Message-ID header or SHA-256 content hash. + + 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: return mid return "sha256-" + hashlib.sha256(msg.as_bytes()).hexdigest()[:32] def _parse_date(msg: mailbox.mboxMessage) -> datetime: - """Parse Date header to UTC-aware datetime. Returns epoch as last resort.""" - date_str = msg.get("Date", "") + """Parse Date header to UTC-aware datetime. Returns epoch as last resort. + + 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: try: dt = parsedate_to_datetime(date_str) @@ -152,8 +165,17 @@ async def run_import( ) -> dict[str, int]: """Import all messages from an mbox file. - Returns stats: {processed, imported, skipped, errors, epoch_fallback, - msgs_with_attachments, total_attachments, total_attachment_bytes}. + Returns stats: {processed, imported, skipped, errors, db_inserted, + 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. limit=N stops after processing N messages (for sampling). @@ -163,6 +185,8 @@ async def run_import( "imported": 0, "skipped": 0, "errors": 0, + "db_inserted": 0, + "db_insert_failed": 0, "epoch_fallback": 0, "msgs_with_attachments": 0, "total_attachments": 0, @@ -176,8 +200,17 @@ async def run_import( pending: list[Envelope] = [] 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: - await _insert_batch(conn, pending) + try: + 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() try: @@ -188,24 +221,28 @@ async def run_import( stats["processed"] += 1 - envelope_id = _message_id(msg) - ts = _parse_date(msg) - - if ts == _EPOCH: - stats["epoch_fallback"] += 1 - _log.warning("epoch_fallback", id=envelope_id) - - attachments = _parse_attachments(msg) - stats["total_attachments"] += len(attachments) - stats["total_attachment_bytes"] += sum(a["size"] for a in attachments) - if attachments: - stats["msgs_with_attachments"] += 1 - - if dry_run: - stats["imported"] += 1 - continue - + # 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 = "" try: + envelope_id = _message_id(msg) + ts = _parse_date(msg) + + if ts == _EPOCH: + stats["epoch_fallback"] += 1 + _log.warning("epoch_fallback", id=envelope_id) + + attachments = _parse_attachments(msg) + stats["total_attachments"] += len(attachments) + stats["total_attachment_bytes"] += sum(a["size"] for a in attachments) + if attachments: + stats["msgs_with_attachments"] += 1 + + if dry_run: + stats["imported"] += 1 + continue + raw = msg.as_bytes(policy=email.policy.compat32) try: 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: _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() finally: @@ -262,7 +302,17 @@ def main() -> None: stats = asyncio.run( 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__": diff --git a/jobs/gmail-bulk-import/tests/test_importer.py b/jobs/gmail-bulk-import/tests/test_importer.py index 5a0de47..bc6c46f 100644 --- a/jobs/gmail-bulk-import/tests/test_importer.py +++ b/jobs/gmail-bulk-import/tests/test_importer.py @@ -2,6 +2,7 @@ from __future__ import annotations import hashlib +import json import mailbox from datetime import datetime, timezone from email import encoders @@ -12,6 +13,7 @@ from pathlib import Path import pytest +from gmail_bulk_import import importer from gmail_bulk_import.importer import ( _message_id, _parse_attachments, @@ -281,3 +283,211 @@ class TestRunImport: assert stats["imported"] == 1 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: \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: \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("").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(""), _msg_with_id("")]) + + 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"") 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"") 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"") 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"") 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 diff --git a/jobs/gmail-header-backfill/pyproject.toml b/jobs/gmail-header-backfill/pyproject.toml index 25549e1..161dc7c 100644 --- a/jobs/gmail-header-backfill/pyproject.toml +++ b/jobs/gmail-header-backfill/pyproject.toml @@ -9,6 +9,7 @@ requires-python = ">=3.11" dependencies = [ "asyncpg>=0.29", "structlog>=24.1", + "kb-mail", ] [project.scripts] diff --git a/jobs/gmail-header-backfill/src/gmail_header_backfill/backfill.py b/jobs/gmail-header-backfill/src/gmail_header_backfill/backfill.py index e7ce14b..1ea6b19 100644 --- a/jobs/gmail-header-backfill/src/gmail_header_backfill/backfill.py +++ b/jobs/gmail-header-backfill/src/gmail_header_backfill/backfill.py @@ -46,6 +46,8 @@ from typing import Optional import asyncpg import structlog +from kb_mail.text import sanitize_surrogates as _sanitize + _log = structlog.get_logger(__name__) 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: """Degraded parse for messages the typed (policy.default) path rejects. diff --git a/packages/kb-mail/src/kb_mail/__init__.py b/packages/kb-mail/src/kb_mail/__init__.py index fc5bd07..131ecb2 100644 --- a/packages/kb-mail/src/kb_mail/__init__.py +++ b/packages/kb-mail/src/kb_mail/__init__.py @@ -1,5 +1,12 @@ from .archive import save_eml from .db import get_envelope, insert_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", +] diff --git a/packages/kb-mail/src/kb_mail/text.py b/packages/kb-mail/src/kb_mail/text.py new file mode 100644 index 0000000..acfd4f4 --- /dev/null +++ b/packages/kb-mail/src/kb_mail/text.py @@ -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") diff --git a/packages/kb-mail/tests/test_text.py b/packages/kb-mail/tests/test_text.py new file mode 100644 index 0000000..f37d62a --- /dev/null +++ b/packages/kb-mail/tests/test_text.py @@ -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�id") == "bad�id"