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.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 = "<unknown>"
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__":

View file

@ -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: <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 = [
"asyncpg>=0.29",
"structlog>=24.1",
"kb-mail",
]
[project.scripts]

View file

@ -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.

View file

@ -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",
]

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"