From 47894d689a051a6273d5a6112fe601e73b534cae Mon Sep 17 00:00:00 2001 From: Chris Guidry Date: Thu, 6 Aug 2026 10:19:34 -0400 Subject: [PATCH] Fail closed on every protected-snapshot failure, PBKDF2 always MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit With a key configured, a missing snapshot, a parse failure, or an apply failure now fails the task, not just a decryption failure. The key material is always stretched with PBKDF2, because length is not evidence of entropy, and an empty key is rejected. Addresses review feedback on #4772. 🤖 Generated with Claude Code Co-Authored-By: Claude Fable 5 --- fastmcp_slim/fastmcp/settings.py | 5 +- fastmcp_tasks/fastmcp_tasks/context.py | 36 ++++++++-- fastmcp_tasks/fastmcp_tasks/encryption.py | 32 +++++---- .../tasks/server/test_snapshot_encryption.py | 65 +++++++++++++++++++ 4 files changed, 116 insertions(+), 22 deletions(-) diff --git a/fastmcp_slim/fastmcp/settings.py b/fastmcp_slim/fastmcp/settings.py index 7533fa2c8..c62236e7f 100644 --- a/fastmcp_slim/fastmcp/settings.py +++ b/fastmcp_slim/fastmcp/settings.py @@ -279,8 +279,9 @@ class Settings(BaseSettings): sharing a task queue must set the same key; a worker that cannot decrypt a snapshot fails the task rather than running it as an anonymous caller. When unset, the snapshot is stored as - plaintext JSON. Any string works, because the Fernet key is - derived from it, but use at least 32 random characters. + plaintext JSON. The Fernet key is derived from this value with + PBKDF2, so any non-empty string works, but use at least 32 + random characters. """ ), ), diff --git a/fastmcp_tasks/fastmcp_tasks/context.py b/fastmcp_tasks/fastmcp_tasks/context.py index 90ce8c7ab..b7ba917c2 100644 --- a/fastmcp_tasks/fastmcp_tasks/context.py +++ b/fastmcp_tasks/fastmcp_tasks/context.py @@ -297,10 +297,11 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: still runs, and sync helpers return ``None`` as they would have before the snapshot was captured. - A snapshot that will not decrypt is the one exception. When the deployment - configured an encryption key, an unreadable snapshot means the caller cannot - be recovered, so the task fails rather than running under no identity at all - (#4747). + Configuring an encryption key changes that contract. The operator asked for + fail-closed protection, so any failure to retrieve, decrypt, parse, or apply + the snapshot, including a snapshot that is simply missing, escapes this + dependency and fails the task, rather than running the tool without the + submitting caller's identity (#4747). """ try: parts = parse_task_key(key) @@ -311,6 +312,11 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: from fastmcp.server.dependencies import get_server 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. + codec = snapshot_codec() + try: docket = get_server()._docket except RuntimeError: @@ -318,6 +324,11 @@ 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: + raise RuntimeError( + "No Docket backend is available to retrieve the protected " + "task snapshot, so the submitting caller cannot be recovered." + ) return task_scope = parts["task_scope"] @@ -328,8 +339,12 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot") ) if raw is None: - return - codec = snapshot_codec() + if codec is None: + 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) @@ -350,6 +365,15 @@ async def restore_task_snapshot(key: str = TaskKey()) -> None: ) raise except Exception: + if codec is not None: + _logger.error( + "Failed to restore the protected task snapshot for %s. The task " + "will fail rather than run without the submitting caller's " + "identity.", + key, + exc_info=True, + ) + raise _logger.warning("Failed to restore task snapshot for %s", key, exc_info=True) diff --git a/fastmcp_tasks/fastmcp_tasks/encryption.py b/fastmcp_tasks/fastmcp_tasks/encryption.py index 37f9d8d50..f1a3e5dbb 100644 --- a/fastmcp_tasks/fastmcp_tasks/encryption.py +++ b/fastmcp_tasks/fastmcp_tasks/encryption.py @@ -23,9 +23,10 @@ logger = get_logger(__name__) # over time, and each derives its own Fernet key from this shared material. _SNAPSHOT_KEY_SALT = "fastmcp-task-snapshot-key" -# Below this, the material is too weak to feed HKDF, which assumes its input is -# already high-entropy. Matches the OAuth proxy's threshold for the same reason. -_MINIMUM_HIGH_ENTROPY_LENGTH = 12 +# Below this, warn: the keyspace is small enough that the offline attacker this +# feature defends against can search it even through PBKDF2. Matches the OAuth +# proxy's threshold for its signing-key material. +_SHORT_KEY_WARNING_LENGTH = 12 class SnapshotDecryptionError(Exception): @@ -40,9 +41,10 @@ class SnapshotDecryptionError(Exception): class SnapshotCodec: """Encrypts and decrypts snapshot payloads with a key derived from material. - Any string works as ``material``. High-entropy material is stretched with - HKDF; a short passphrase goes through PBKDF2 instead, which is slower but - survives the weaker input. + 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 + comes from PBKDF2, never from HKDF. The stretch costs about a second, paid + once per process (see ``_codec_for``). """ def __init__(self, material: str) -> None: @@ -50,17 +52,19 @@ class SnapshotCodec: from fastmcp.server.auth.jwt_issuer import derive_jwt_key - if len(material) >= _MINIMUM_HIGH_ENTROPY_LENGTH: - key = derive_jwt_key( - high_entropy_material=material, salt=_SNAPSHOT_KEY_SALT + if not material: + raise ValueError( + "FASTMCP_ENCRYPTION_KEY must not be empty. Unset it to store " + "task snapshots as plaintext, or set at least 32 random " + "characters." ) - else: + if len(material) < _SHORT_KEY_WARNING_LENGTH: logger.warning( "The configured encryption key is shorter than %d characters; " "use at least 32 random characters.", - _MINIMUM_HIGH_ENTROPY_LENGTH, + _SHORT_KEY_WARNING_LENGTH, ) - key = derive_jwt_key(low_entropy_material=material, salt=_SNAPSHOT_KEY_SALT) + key = derive_jwt_key(low_entropy_material=material, salt=_SNAPSHOT_KEY_SALT) self._fernet = Fernet(key=key) @@ -90,8 +94,8 @@ class SnapshotCodec: def _codec_for(material: str) -> SnapshotCodec: """One codec per key, so the derivation cost is paid once per process. - PBKDF2 over a short passphrase takes about a second, and every task - submission and every restore needs a codec. + The PBKDF2 stretch takes about a second, and every task submission and + every restore needs a codec. """ return SnapshotCodec(material) diff --git a/tests/tasks/server/test_snapshot_encryption.py b/tests/tasks/server/test_snapshot_encryption.py index 49b83d496..2e1a25a82 100644 --- a/tests/tasks/server/test_snapshot_encryption.py +++ b/tests/tasks/server/test_snapshot_encryption.py @@ -13,6 +13,7 @@ from __future__ import annotations import json import logging from collections.abc import Iterator +from unittest.mock import patch import pytest from fastmcp_tasks.context import TaskContextSnapshot @@ -103,6 +104,11 @@ class TestSnapshotCodec: with pytest.raises(SnapshotDecryptionError): SnapshotCodec(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("") + def test_decode_accepts_bytes(self): """Redis hands back bytes on some backends.""" codec = SnapshotCodec(KEY) @@ -147,6 +153,15 @@ async def _write_stored_snapshot( await redis.set(key, payload) +async def _delete_stored_snapshot(mcp: FastMCP, task_scope: str, task_id: str) -> None: + """Remove a task's stored snapshot, as a TTL expiry would.""" + docket = mcp._docket + assert docket is not None + key = docket.key(f"{task_redis_prefix(task_scope)}:{task_id}:snapshot") + async with docket.redis() as redis: + await redis.delete(key) + + @pytest.fixture def echo_token_server() -> FastMCP: """A task server whose one tool reports the caller it restored.""" @@ -233,6 +248,56 @@ class TestEncryptedSnapshotRoundTrip: assert final.error is not None assert "FASTMCP_ENCRYPTION_KEY" in caplog.text + async def test_missing_snapshot_fails_the_task( + self, echo_token_server: FastMCP, encryption_key: str + ): + """Fail closed extends to a snapshot that is gone, not just unreadable. + + A missing snapshot is reachable in production through TTL expiry, and + it loses the caller just as completely as a wrong key does. + """ + token = make_access_token("client-a", "user-1") + + async with running_task_server(echo_token_server): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + await _delete_stored_snapshot( + echo_token_server, "client-a|user-1", created.task_id + ) + final = await wait_for_task( + echo_token_server, + created.task_id, + access_token=token, + target_states=frozenset({"failed"}), + ) + + assert final.status == "failed" + + async def test_unparseable_snapshot_fails_the_task( + self, echo_token_server: FastMCP, encryption_key: str + ): + """Fail closed extends past decryption: a parse failure also loses the + caller, so it must not degrade to an anonymous run.""" + token = make_access_token("client-a", "user-1") + + def boom(*_args, **_kwargs): + raise RuntimeError("simulated deserialization failure") + + async with running_task_server(echo_token_server): + with patch.object(TaskContextSnapshot, "from_json", boom): + created = await submit_task( + echo_token_server, "whoami", {}, access_token=token + ) + final = await wait_for_task( + echo_token_server, + created.task_id, + access_token=token, + target_states=frozenset({"failed"}), + ) + + assert final.status == "failed" + class TestUnencryptedByDefault: async def test_snapshot_stays_plaintext_without_a_key(