Fix #3680: MaskErrors does not mask errors for subscriptions

RELEASE.md

@@ -0,0 +1,25 @@
+Release type: patch
+
+Fixes an issue where schema extensions (like `MaskErrors`) were bypassed during WebSocket subscriptions. The extensions' `_process_result` hooks are now properly triggered for each yielded result in both `graphql-transport-ws` and `graphql-ws` protocols, ensuring errors are correctly formatted before being sent to the client.
+
+### Description
+Fixes an issue where schema extensions (such as `MaskErrors`) were being bypassed when streaming data over WebSockets.
+
+Previously, standard Queries and Mutations would pass their results through the extension pipeline, but Subscriptions would send raw `ExecutionResult` objects directly over the WebSocket. This caused internal/unmasked errors to leak to the client. This PR manually triggers `_process_result` on active extensions right before `send_next` and `send_data_message` dispatch the payload.
+
+### Migration guide
+No migration required.
+
+### Types of Changes
+- [ ] Core
+- [x] Bugfix
+- [ ] New feature
+- [ ] Enhancement/optimization
+- [ ] Documentation
+
+### Checklist
+- [x] My code follows the code style of this project.
+- [ ] My change requires a change to the documentation.
+- [x] I have read the CONTRIBUTING document.
+- [x] I have added tests to cover my changes.
+- [x] I have tested the changes and verified that they work and don't break anything (as well as I can manage).

strawberry/subscriptions/protocols/graphql_transport_ws/handlers.py

@@ -29,6 +29,7 @@
     PongMessage,
     SubscribeMessage,
 )
+from strawberry.subscriptions.utils import process_extensions
 from strawberry.types import ExecutionResult
 from strawberry.types.execution import PreExecutionError
 from strawberry.types.graphql import OperationType
@@ -394,6 +395,9 @@ async def send_initial_errors(self, errors: list[GraphQLError]) -> None:
         )
 
     async def send_next(self, execution_result: ExecutionResult) -> None:
+        extensions = getattr(self.handler.schema, "extensions", [])
+        process_extensions(execution_result, extensions)
+
         next_payload: NextMessagePayload = {"data": execution_result.data}
 
         if execution_result.errors:

strawberry/subscriptions/protocols/graphql_ws/handlers.py

@@ -24,6 +24,7 @@
     StartMessage,
     StopMessage,
 )
+from strawberry.subscriptions.utils import process_extensions
 from strawberry.types.execution import ExecutionResult, PreExecutionError
 from strawberry.types.unset import UnsetType
 
@@ -216,6 +217,9 @@ async def cleanup(self) -> None:
     async def send_data_message(
         self, execution_result: ExecutionResult, operation_id: str
     ) -> None:
+        extensions = getattr(self.schema, "extensions", [])
+        process_extensions(execution_result, extensions)
+
         data_message: DataMessage = {
             "type": "data",
             "id": operation_id,

strawberry/subscriptions/utils.py

@@ -0,0 +1,26 @@
+import inspect
+from typing import Any
+
+from strawberry.types import ExecutionResult
+
+
+def process_extensions(
+    execution_result: ExecutionResult, extensions: list[Any]
+) -> None:
+    """Run the execution result through active schema extensions."""
+    for ext in extensions:
+        if isinstance(ext, type):
+            # Inspect the constructor to see if it requires execution_context
+            sig = inspect.signature(ext.__init__)
+            if "execution_context" in sig.parameters:
+                extension_instance = ext(execution_context=None)
+            else:
+                extension_instance = ext()
+
+            # Explicitly set this ONLY for newly constructed instances
+            extension_instance.execution_context = None
+        else:
+            extension_instance = ext
+
+        if hasattr(extension_instance, "_process_result"):
+            extension_instance._process_result(execution_result)

tests/websockets/test_graphql_transport_ws.py

@@ -1216,3 +1216,116 @@ async def test_unexpected_client_disconnects_are_gracefully_handled(
 
         assert not process_errors.called
         assert Subscription.active_infinity_subscriptions == 0
+
+
+@patch.object(MyExtension, "_process_result", create=True)
+async def test_subscription_errors_trigger_extension_process_result(
+    mock: Mock, ws: WebSocketClient
+):
+    """Test that schema extensions are called to process results when a subscription yields an error."""
+    await ws.send_message(
+        {
+            "id": "sub1",
+            "type": "subscribe",
+            "payload": {
+                "query": 'subscription { exception(message: "TEST EXC") }',
+            },
+        }
+    )
+
+    next_message: NextMessage = await ws.receive_json()
+
+    assert next_message["type"] == "next"
+    assert next_message["id"] == "sub1"
+    assert "errors" in next_message["payload"]
+
+    # Error intercepted and extension called
+    mock.assert_called_once()
+
+
+async def test_subscription_error_masking_end_to_end(
+    http_client_class: type[HttpClient],
+):
+    """Test that the real MaskErrors extension successfully masks the payload text."""
+    import strawberry
+    from strawberry.extensions import MaskErrors
+    from tests.views.schema import Query, Subscription
+
+    # Create a custom schema with the real MaskErrors extension attached
+    custom_schema = strawberry.Schema(
+        query=Query,
+        subscription=Subscription,
+        extensions=[MaskErrors(error_message="Unexpected error.")],
+    )
+    test_client = http_client_class(custom_schema)
+
+    async with test_client.ws_connect(
+        "/graphql", protocols=[GRAPHQL_TRANSPORT_WS_PROTOCOL]
+    ) as ws:
+        await ws.send_message({"type": "connection_init"})
+        await ws.receive_json()
+
+        await ws.send_message(
+            {
+                "id": "sub1",
+                "type": "subscribe",
+                "payload": {
+                    "query": 'subscription { exception(message: "Super secret database error") }',
+                },
+            }
+        )
+
+        next_message: NextMessage = await ws.receive_json()
+
+        assert next_message["type"] == "next"
+        assert next_message["id"] == "sub1"
+        assert "errors" in next_message["payload"]
+
+        assert next_message["payload"]["errors"][0]["message"] == "Unexpected error."
+
+
+async def test_subscription_masking_with_class_extension(
+    http_client_class: type[HttpClient],
+):
+    """Test that passing an extension as a class (not instance) successfully masks errors."""
+    import strawberry
+    from strawberry.extensions import SchemaExtension
+    from tests.views.schema import Query, Subscription
+
+    class CustomClassExtension(SchemaExtension):
+        def __init__(self, execution_context=None):
+            self.execution_context = execution_context  # type: ignore
+
+        def _process_result(self, execution_result):
+            if execution_result.errors:
+                execution_result.errors[0].message = "Unexpected error."
+
+    # Create a custom schema passing the CLASS, not an instance
+    custom_schema = strawberry.Schema(
+        query=Query, subscription=Subscription, extensions=[CustomClassExtension]
+    )
+    test_client = http_client_class(custom_schema)
+
+    async with test_client.ws_connect(
+        "/graphql", protocols=[GRAPHQL_TRANSPORT_WS_PROTOCOL]
+    ) as ws:
+        await ws.send_message({"type": "connection_init"})
+        await ws.receive_json()
+
+        await ws.send_message(
+            {
+                "id": "sub_class",
+                "type": "subscribe",
+                "payload": {
+                    "query": 'subscription { exception(message: "Secret database error") }',
+                },
+            }
+        )
+
+        next_message: NextMessage = await ws.receive_json()
+
+        assert next_message["type"] == "next"
+        assert next_message["id"] == "sub_class"
+        assert "errors" in next_message["payload"]
+
+        assert next_message["payload"]["errors"][0]["message"] == "Unexpected error."

tests/websockets/test_graphql_ws.py

@@ -865,3 +865,128 @@ async def test_unexpected_client_disconnects_are_gracefully_handled(
 
         assert not process_errors.called
         assert Subscription.active_infinity_subscriptions == 0
+
+
+@mock.patch.object(MyExtension, "_process_result", create=True)
+async def test_subscription_errors_trigger_extension_process_result(
+    mock: mock.MagicMock, ws: WebSocketClient
+):
+    """Test that schema extensions are called to process results when a subscription yields an error."""
+    await ws.send_legacy_message(
+        {
+            "type": "start",
+            "id": "demo",
+            "payload": {
+                "query": 'subscription { exception(message: "TEST EXC") }',
+            },
+        }
+    )
+
+    data_message: DataMessage = await ws.receive_json()
+
+    assert data_message["type"] == "data"
+    assert data_message["id"] == "demo"
+    assert "errors" in data_message["payload"]
+
+    # Error intercepted and extension called
+    mock.assert_called_once()
+
+    await ws.send_legacy_message({"type": "stop", "id": "demo"})
+    complete_message = await ws.receive_json()
+    assert complete_message["type"] == "complete"
+
+
+async def test_subscription_error_masking_end_to_end(
+    http_client_class: type[HttpClient],
+):
+    """Test that the real MaskErrors extension successfully masks the legacy payload text."""
+    import strawberry
+    from strawberry.extensions import MaskErrors
+    from tests.views.schema import Query, Subscription
+
+    # Create a custom schema with the real MaskErrors extension attached
+    custom_schema = strawberry.Schema(
+        query=Query,
+        subscription=Subscription,
+        extensions=[MaskErrors(error_message="Unexpected error.")],
+    )
+    test_client = http_client_class(custom_schema)
+
+    async with test_client.ws_connect(
+        "/graphql", protocols=[GRAPHQL_WS_PROTOCOL]
+    ) as ws:
+        await ws.send_legacy_message({"type": "connection_init"})
+        await ws.receive_json()
+
+        await ws.send_legacy_message(
+            {
+                "type": "start",
+                "id": "demo",
+                "payload": {
+                    "query": 'subscription { exception(message: "Super secret database error") }',
+                },
+            }
+        )
+
+        data_message: DataMessage = await ws.receive_json()
+
+        assert data_message["type"] == "data"
+        assert data_message["id"] == "demo"
+        assert "errors" in data_message["payload"]
+
+        assert data_message["payload"]["errors"][0]["message"] == "Unexpected error."
+
+        # Clean up the socket
+        await ws.send_legacy_message({"type": "stop", "id": "demo"})
+        await ws.receive_json()
+
+
+async def test_subscription_masking_with_class_extension(
+    http_client_class: type[HttpClient],
+):
+    """Test that passing an extension as a class (not instance) successfully masks the legacy payload text."""
+    import strawberry
+    from strawberry.extensions import SchemaExtension
+    from tests.views.schema import Query, Subscription
+
+    class CustomClassExtension(SchemaExtension):
+        def __init__(self, execution_context=None):
+            self.execution_context = execution_context  # type: ignore
+
+        def _process_result(self, execution_result):
+            if execution_result.errors:
+                execution_result.errors[0].message = "Unexpected error."
+
+    # Create a custom schema passing the CLASS, not an instance
+    custom_schema = strawberry.Schema(
+        query=Query, subscription=Subscription, extensions=[CustomClassExtension]
+    )
+    test_client = http_client_class(custom_schema)
+
+    async with test_client.ws_connect(
+        "/graphql", protocols=[GRAPHQL_WS_PROTOCOL]
+    ) as ws:
+        await ws.send_legacy_message({"type": "connection_init"})
+        await ws.receive_json()
+
+        await ws.send_legacy_message(
+            {
+                "type": "start",
+                "id": "sub_class",
+                "payload": {
+                    "query": 'subscription { exception(message: "Secret database error") }',
+                },
+            }
+        )
+
+        data_message: DataMessage = await ws.receive_json()
+
+        assert data_message["type"] == "data"
+        assert data_message["id"] == "sub_class"
+        assert "errors" in data_message["payload"]
+
+        assert data_message["payload"]["errors"][0]["message"] == "Unexpected error."
+
+        # Clean up the socket
+        await ws.send_legacy_message({"type": "stop", "id": "sub_class"})
+        await ws.receive_json()