Chore: consolidate pickle hmac signing (#13899)

src/paperless/celery.py

@@ -1,12 +1,12 @@
-import hmac
 import os
-import pickle
-from hashlib import sha256
 
 from celery import Celery
 from celery.signals import worker_process_init
 from kombu.serialization import register
 
+from paperless.signed_pickle import signed_pickle_dumps
+from paperless.signed_pickle import signed_pickle_loads
+
 # Set the default Django settings module for the 'celery' program.
 os.environ.setdefault("DJANGO_SETTINGS_MODULE", "paperless.settings")
 
@@ -18,34 +18,6 @@ os.environ.setdefault("DJANGO_SETTINGS_MODULE", "paperless.settings")
 # on the worker side using Django's SECRET_KEY.
 # ---------------------------------------------------------------------------
 
-HMAC_SIZE = 32  # SHA-256 digest length
-
-
-def _get_signing_key() -> bytes:
-    from django.conf import settings
-
-    return settings.SECRET_KEY.encode()
-
-
-def signed_pickle_dumps(obj: object) -> bytes:
-    data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
-    signature = hmac.new(_get_signing_key(), data, sha256).digest()
-    return signature + data
-
-
-def signed_pickle_loads(payload: bytes) -> object:
-    if len(payload) < HMAC_SIZE:
-        msg = "Signed-pickle payload too short"
-        raise ValueError(msg)
-    signature = payload[:HMAC_SIZE]
-    data = payload[HMAC_SIZE:]
-    expected = hmac.new(_get_signing_key(), data, sha256).digest()
-    if not hmac.compare_digest(signature, expected):
-        msg = "Signed-pickle HMAC verification failed — message may have been tampered with"
-        raise ValueError(msg)
-    return pickle.loads(data)
-
-
 register(
     "signed-pickle",
     signed_pickle_dumps,

src/paperless/signed_pickle.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+import hmac
+import pickle
+from hashlib import sha256
+from typing import Any
+
+from django.conf import settings
+
+HMAC_SIZE = sha256().digest_size
+
+
+class SignedPickleError(ValueError):
+    """Raised when a signed pickle payload cannot be authenticated."""
+
+
+def _get_signing_key() -> bytes:
+    return settings.SECRET_KEY.encode()
+
+
+def signed_pickle_dumps(obj: object) -> bytes:
+    data = pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL)
+    signature = hmac.new(_get_signing_key(), data, sha256).digest()
+    return signature + data
+
+
+def signed_pickle_loads(payload: bytes) -> Any:
+    if len(payload) <= HMAC_SIZE:
+        msg = "Signed-pickle payload too short"
+        raise SignedPickleError(msg)
+
+    signature = payload[:HMAC_SIZE]
+    data = payload[HMAC_SIZE:]
+    expected = hmac.new(_get_signing_key(), data, sha256).digest()
+    if not hmac.compare_digest(signature, expected):
+        msg = "Signed-pickle HMAC verification failed; payload may have been tampered with"
+        raise SignedPickleError(msg)
+
+    return pickle.loads(data)

src/paperless/tests/test_celery.py

@@ -6,9 +6,9 @@ from pathlib import Path
 import pytest
 from django.test import override_settings
 
-from paperless.celery import HMAC_SIZE
 from paperless.celery import signed_pickle_dumps
 from paperless.celery import signed_pickle_loads
+from paperless.signed_pickle import HMAC_SIZE
 
 
 class TestSignedPickleSerializer: