diff --git a/src/documents/classifier.py b/src/documents/classifier.py index 1cdef4419..bac810cb8 100644 --- a/src/documents/classifier.py +++ b/src/documents/classifier.py @@ -13,6 +13,9 @@ if TYPE_CHECKING: from collections.abc import Callable from collections.abc import Iterator from datetime import datetime + from types import TracebackType + from typing import BinaryIO + from typing import Self from numpy import ndarray @@ -68,6 +71,49 @@ RE_DIGIT = re.compile(r"\d") RE_WORD = re.compile(r"\b[\w]+\b") # words that may contain digits +class _SignedFileWriter: + """ + Atomically writes a file made of an HMAC signature followed by the data, + signing the data as it streams to disk rather than holding it in memory. + + The signature is only known once everything is written, so its space is + reserved at the start of the file and filled in on exit. The target is only + replaced once the file is complete; on error the partial file is removed. + """ + + def __init__(self, target: Path, mac: hmac.HMAC) -> None: + self._target = target + self._temp = target.with_name(f"{target.name}.part") + self._mac = mac + self._file: BinaryIO + + def __enter__(self) -> Self: + self._file = self._temp.open("wb") + self._file.write(bytes(self._mac.digest_size)) + return self + + def write(self, data: bytes | memoryview) -> int: + self._mac.update(data) + return self._file.write(data) + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + try: + with self._file: + if exc_type is None: + self._file.seek(0) + self._file.write(self._mac.digest()) + if exc_type is None: + self._temp.rename(self._target) + finally: + # A no-op after a successful rename, otherwise removes the partial file + self._temp.unlink(missing_ok=True) + + class IncompatibleClassifierVersionError(Exception): def __init__(self, message: str, *args: object) -> None: self.message: str = message @@ -159,13 +205,15 @@ class DocumentClassifier: pickle.dumps(self.data_vectorizer), ).hexdigest() + @staticmethod + def _new_hmac() -> hmac.HMAC: + return hmac.new(settings.SECRET_KEY.encode(), digestmod=sha256) + @staticmethod def _compute_hmac(data: bytes | memoryview) -> bytes: - return hmac.new( - settings.SECRET_KEY.encode(), - data, - sha256, - ).digest() + mac = DocumentClassifier._new_hmac() + mac.update(data) + return mac.digest() def load(self) -> None: from sklearn.exceptions import InconsistentVersionWarning @@ -226,29 +274,24 @@ class DocumentClassifier: raise IncompatibleClassifierVersionError("sklearn version update") def save(self) -> None: - target_file: Path = settings.MODEL_FILE - target_file_temp: Path = target_file.with_suffix(".pickle.part") - - data = pickle.dumps( - ( - self.FORMAT_VERSION, - self.last_doc_change_time, - self.last_auto_type_hash, - self.data_vectorizer, - self.tags_binarizer, - self.tags_classifier, - self.correspondent_classifier, - self.document_type_classifier, - self.storage_path_classifier, - ), - ) - - signature = self._compute_hmac(data) - - with target_file_temp.open("wb") as f: - f.write(signature + data) - - target_file_temp.rename(target_file) + # Stream to disk instead of building the payload in memory. Protocol 5+ + # pickles numpy arrays without copying them (the default is 4 before 3.14). + with _SignedFileWriter(settings.MODEL_FILE, self._new_hmac()) as f: + pickle.dump( + ( + self.FORMAT_VERSION, + self.last_doc_change_time, + self.last_auto_type_hash, + self.data_vectorizer, + self.tags_binarizer, + self.tags_classifier, + self.correspondent_classifier, + self.document_type_classifier, + self.storage_path_classifier, + ), + f, + protocol=pickle.HIGHEST_PROTOCOL, + ) def train( self, diff --git a/src/documents/tests/test_classifier.py b/src/documents/tests/test_classifier.py index 9066c4d41..99d7196f3 100644 --- a/src/documents/tests/test_classifier.py +++ b/src/documents/tests/test_classifier.py @@ -1,5 +1,8 @@ +import pickle import re import warnings +from datetime import UTC +from datetime import datetime from pathlib import Path from unittest import mock @@ -10,6 +13,7 @@ from django.db import connection from django.test import TestCase from django.test import override_settings from django.test.utils import CaptureQueriesContext +from pytest_django.fixtures import Settings from pytest_mock import MockerFixture from documents.classifier import ClassifierModelCorruptError @@ -914,6 +918,85 @@ class TestClassifier(DirectoriesMixin, TestCase): load_classifier(raise_exception=True) +class TestClassifierSave: + @pytest.fixture + def model_file(self, tmp_path: Path, settings: Settings) -> Path: + settings.MODEL_FILE = tmp_path / "classifier.pickle" + return settings.MODEL_FILE + + @pytest.fixture + def classifier(self) -> DocumentClassifier: + classifier = DocumentClassifier() + classifier.last_doc_change_time = datetime(2026, 1, 2, 3, 4, 5, tzinfo=UTC) + classifier.last_auto_type_hash = b"\x01" * 32 + return classifier + + def test_save_writes_signed_pickle( + self, + model_file: Path, + classifier: DocumentClassifier, + ) -> None: + """ + GIVEN: + - A classifier with training state + WHEN: + - The classifier is saved + THEN: + - The file is the HMAC of the pickled state followed by that pickle + - The pickle uses the highest protocol + - The saved state loads back into a new classifier + - No temporary file is left behind + """ + classifier.save() + + raw = model_file.read_bytes() + signature = raw[: DocumentClassifier.HMAC_SIZE] + data = raw[DocumentClassifier.HMAC_SIZE :] + assert signature == DocumentClassifier._compute_hmac(data) + # A pickle opens with the PROTO opcode followed by the protocol number + assert data[:2] == bytes([pickle.PROTO[0], pickle.HIGHEST_PROTOCOL]) + assert pickle.loads(data)[:3] == ( + DocumentClassifier.FORMAT_VERSION, + classifier.last_doc_change_time, + classifier.last_auto_type_hash, + ) + + loaded = DocumentClassifier() + loaded.load() + assert loaded.last_doc_change_time == classifier.last_doc_change_time + assert loaded.last_auto_type_hash == classifier.last_auto_type_hash + + assert not model_file.with_name(f"{model_file.name}.part").exists() + + def test_save_failure_removes_partial_file( + self, + model_file: Path, + classifier: DocumentClassifier, + mocker: MockerFixture, + ) -> None: + """ + GIVEN: + - An existing classifier model file + WHEN: + - Saving a new classifier fails part way through writing + THEN: + - The error is raised + - The partially written temporary file is removed + - The existing model file is left untouched + """ + model_file.write_bytes(b"existing model") + mocker.patch( + "documents.classifier.pickle.dump", + side_effect=RuntimeError("disk full"), + ) + + with pytest.raises(RuntimeError, match="disk full"): + classifier.save() + + assert not model_file.with_name(f"{model_file.name}.part").exists() + assert model_file.read_bytes() == b"existing model" + + class _StubProbaClassifier: """ A fake scikit-learn classifier exposing just enough of the API for