From 6941a958e9a47be8cb2520c80e427f6972bff92d Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Thu, 6 Aug 2026 10:30:35 -0400 Subject: [PATCH] Split the snapshot codec into plaintext and encrypted variants MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit snapshot_codec() now always returns a codec, so save encodes unconditionally and the restore path reads codec.protected to pick between the fail-open and fail-closed contracts, instead of testing the codec for None. 🤖 Generated with Claude Code Co-Authored-By: Claude Fable 5 --- fastmcp_tasks/fastmcp_tasks/context.py | 19 +++---- fastmcp_tasks/fastmcp_tasks/encryption.py | 52 ++++++++++++++++--- .../tasks/server/test_snapshot_encryption.py | 33 +++++++----- 3 files changed, 73 insertions(+), 31 deletions(-) diff --git a/fastmcp_tasks/fastmcp_tasks/context.py b/fastmcp_tasks/fastmcp_tasks/context.py index b7ba917c2..b36180044 100644 --- a/fastmcp_tasks/fastmcp_tasks/context.py +++ b/fastmcp_tasks/fastmcp_tasks/context.py @@ -235,10 +235,7 @@ class TaskContextSnapshot: read it (#4747). """ key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot") - codec = snapshot_codec() - payload = self.to_json() - if codec is not None: - payload = codec.encode(payload) + payload = snapshot_codec().encode(self.to_json()) async with docket.redis() as redis: await redis.set(key, payload, ex=ttl_seconds) @@ -313,8 +310,8 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: from fastmcp_tasks.dependencies import _current_docket # Resolved before anything can fail: a misconfigured key (e.g. an empty - # string) raises here and fails the task, and the except blocks below read - # it to pick between the fail-open and fail-closed contracts. + # string) raises here and fails the task, and the branches below read + # `codec.protected` to pick between the fail-open and fail-closed contracts. codec = snapshot_codec() try: @@ -324,7 +321,7 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: if docket is None: docket = _current_docket.get() if docket is None: - if codec is not None: + if codec.protected: raise RuntimeError( "No Docket backend is available to retrieve the protected " "task snapshot, so the submitting caller cannot be recovered." @@ -339,15 +336,13 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot") ) if raw is None: - if codec is None: + if not codec.protected: return raise RuntimeError( "The task's context snapshot is missing (its TTL may have " "expired), so the submitting caller cannot be recovered." ) - if codec is not None: - raw = codec.decode(raw) - snapshot = TaskContextSnapshot.from_json(raw) + snapshot = TaskContextSnapshot.from_json(codec.decode(raw)) _remember_snapshot(task_id, snapshot) # Restore the ambient request context (auth token, headers) so core's # get_access_token()/get_http_headers() see the submitting caller inside @@ -365,7 +360,7 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: ) raise except Exception: - if codec is not None: + if codec.protected: _logger.error( "Failed to restore the protected task snapshot for %s. The task " "will fail rather than run without the submitting caller's " diff --git a/fastmcp_tasks/fastmcp_tasks/encryption.py b/fastmcp_tasks/fastmcp_tasks/encryption.py index f1a3e5dbb..ca2a52b2e 100644 --- a/fastmcp_tasks/fastmcp_tasks/encryption.py +++ b/fastmcp_tasks/fastmcp_tasks/encryption.py @@ -12,7 +12,9 @@ that restores a snapshot is rarely the one that captured it. from __future__ import annotations +from abc import ABC, abstractmethod from functools import lru_cache +from typing import ClassVar import fastmcp from fastmcp.utilities.logging import get_logger @@ -38,8 +40,39 @@ class SnapshotDecryptionError(Exception): """ -class SnapshotCodec: - """Encrypts and decrypts snapshot payloads with a key derived from material. +class SnapshotCodec(ABC): + """Transforms snapshot payloads on their way to and from the backend. + + ``protected`` tells the restore path which failure contract applies: a + protected snapshot that cannot be restored fails the task, an unprotected + one degrades to an anonymous run with a warning. + """ + + protected: ClassVar[bool] + + @abstractmethod + def encode(self, payload: str) -> str: + """Return the stored form of a serialized snapshot.""" + + @abstractmethod + def decode(self, stored: str | bytes) -> str: + """Return the serialized snapshot a stored value holds.""" + + +class PlaintextCodec(SnapshotCodec): + """Stores snapshots as-is; the contract when no encryption key is set.""" + + protected = False + + def encode(self, payload: str) -> str: + return payload + + def decode(self, stored: str | bytes) -> str: + return stored.decode() if isinstance(stored, bytes) else stored + + +class EncryptedCodec(SnapshotCodec): + """Encrypts snapshot payloads with a key derived from material. The material is a string from the environment, and nothing about a string proves it is random, so it is always treated as low-entropy: the Fernet key @@ -47,6 +80,8 @@ class SnapshotCodec: once per process (see ``_codec_for``). """ + protected = True + def __init__(self, material: str) -> None: from cryptography.fernet import Fernet @@ -90,21 +125,24 @@ class SnapshotCodec: ) from e +_PLAINTEXT_CODEC = PlaintextCodec() + + @lru_cache(maxsize=4) -def _codec_for(material: str) -> SnapshotCodec: +def _codec_for(material: str) -> EncryptedCodec: """One codec per key, so the derivation cost is paid once per process. The PBKDF2 stretch takes about a second, and every task submission and every restore needs a codec. """ - return SnapshotCodec(material) + return EncryptedCodec(material) -def snapshot_codec() -> SnapshotCodec | None: - """The codec for the configured key, or ``None`` when none is configured.""" +def snapshot_codec() -> SnapshotCodec: + """The codec for the configured key; the plaintext codec when none is set.""" key = fastmcp.settings.encryption_key if key is None: - return None + return _PLAINTEXT_CODEC return _codec_for(key.get_secret_value()) diff --git a/tests/tasks/server/test_snapshot_encryption.py b/tests/tasks/server/test_snapshot_encryption.py index 2e1a25a82..59d43caae 100644 --- a/tests/tasks/server/test_snapshot_encryption.py +++ b/tests/tasks/server/test_snapshot_encryption.py @@ -18,7 +18,8 @@ from unittest.mock import patch import pytest from fastmcp_tasks.context import TaskContextSnapshot from fastmcp_tasks.encryption import ( - SnapshotCodec, + EncryptedCodec, + PlaintextCodec, SnapshotDecryptionError, clear_codec_cache, snapshot_codec, @@ -84,41 +85,49 @@ def sensitive_snapshot() -> TaskContextSnapshot: class TestSnapshotCodec: def test_round_trips_a_payload(self): - codec = SnapshotCodec(KEY) + codec = EncryptedCodec(KEY) assert codec.decode(codec.encode('{"a": 1}')) == '{"a": 1}' def test_encoded_payload_hides_the_credentials( self, sensitive_snapshot: TaskContextSnapshot ): - encoded = SnapshotCodec(KEY).encode(sensitive_snapshot.to_json()) + encoded = EncryptedCodec(KEY).encode(sensitive_snapshot.to_json()) assert "token-client-a-user-1" not in encoded assert "authorization" not in encoded def test_decode_rejects_another_keys_payload(self): - encoded = SnapshotCodec(OTHER_KEY).encode('{"a": 1}') + encoded = EncryptedCodec(OTHER_KEY).encode('{"a": 1}') with pytest.raises(SnapshotDecryptionError): - SnapshotCodec(KEY).decode(encoded) + EncryptedCodec(KEY).decode(encoded) def test_decode_rejects_plaintext(self): """A snapshot written before the key was set must not be trusted.""" with pytest.raises(SnapshotDecryptionError): - SnapshotCodec(KEY).decode('{"access_token_json": null}') + EncryptedCodec(KEY).decode('{"access_token_json": null}') def test_empty_material_is_rejected(self): """An empty key would derive a universally reproducible Fernet key.""" with pytest.raises(ValueError, match="must not be empty"): - SnapshotCodec("") + EncryptedCodec("") def test_decode_accepts_bytes(self): """Redis hands back bytes on some backends.""" - codec = SnapshotCodec(KEY) + codec = EncryptedCodec(KEY) assert codec.decode(codec.encode('{"a": 1}').encode()) == '{"a": 1}' def test_same_key_reuses_one_codec(self, encryption_key: str): assert snapshot_codec() is snapshot_codec() - def test_no_codec_without_a_key(self, no_encryption_key: None): - assert snapshot_codec() is None + def test_plaintext_codec_without_a_key(self, no_encryption_key: None): + codec = snapshot_codec() + assert isinstance(codec, PlaintextCodec) + assert not codec.protected + + def test_plaintext_codec_is_a_pass_through(self): + codec = PlaintextCodec() + assert codec.encode('{"a": 1}') == '{"a": 1}' + assert codec.decode('{"a": 1}') == '{"a": 1}' + assert codec.decode(b'{"a": 1}') == '{"a": 1}' class TestSnapshotSerialization: @@ -227,7 +236,7 @@ class TestEncryptedSnapshotRoundTrip: cause has to come from the log. """ token = make_access_token("client-a", "user-1") - tampered = SnapshotCodec(OTHER_KEY).encode(TaskContextSnapshot().to_json()) + tampered = EncryptedCodec(OTHER_KEY).encode(TaskContextSnapshot().to_json()) with caplog.at_level(logging.ERROR, logger="fastmcp_tasks.context"): async with running_task_server(echo_token_server): @@ -348,7 +357,7 @@ class TestTaskStillResolvesAfterFailure: ): """A fail-closed task is still a well-formed `tasks/get` result.""" token = make_access_token("client-a", "user-1") - tampered = SnapshotCodec(OTHER_KEY).encode(TaskContextSnapshot().to_json()) + tampered = EncryptedCodec(OTHER_KEY).encode(TaskContextSnapshot().to_json()) async with running_task_server(echo_token_server): created = await submit_task(