Fail closed on every protected-snapshot failure, PBKDF2 always

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 <noreply@anthropic.com>

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.
                 """
             ),
         ),

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)
 
 

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)
 

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(