fix(client): Reject push requests with a missing or invalid signature

clients/python/src/taskbroker_client/worker/client.py

@@ -34,9 +34,9 @@
 
 logger = logging.getLogger(__name__)
 
-# gRPC runs the unary request deserializer and the servicer on the same thread for a given call
-# If HMAC verification fails inside a deserializer wrapper, raising turns into INTERNAL, so we set this flag and abort UNAUTHENTICATED from the servicer
-_RPC_SIGNATURE_AUTH_TLS = threading.local()
+# Raising in the deserializer would return INTERNAL, so a request that fails the HMAC check
+# is replaced with this marker and the servicer aborts with UNAUTHENTICATED.
+_AUTH_FAILED = object()
 
 MAX_ACTIVATION_SIZE = 1024 * 1024 * 10
 """Max payload size we will process."""
@@ -217,18 +217,15 @@ def intercept_service(
         original = handler.unary_unary
 
         def request_deserializer(serialized_request: bytes) -> Any:
-            _RPC_SIGNATURE_AUTH_TLS.failed = False
-
             if not verify_rpc_request_signature_hmac(
                 self._secrets, method, serialized_request, signature
             ):
-                _RPC_SIGNATURE_AUTH_TLS.failed = True
-                return inner_deserializer(b"")
+                return _AUTH_FAILED
 
             return inner_deserializer(serialized_request)
 
         def unary_unary(request: Any, context: grpc.ServicerContext) -> Any:
-            if getattr(_RPC_SIGNATURE_AUTH_TLS, "failed", False):
+            if request is _AUTH_FAILED:
                 context.abort(grpc.StatusCode.UNAUTHENTICATED, "Authentication failed")
 
             return original(request, context)

clients/python/tests/worker/test_client.py

@@ -6,6 +6,7 @@
 import time
 from collections import defaultdict
 from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor
 from pathlib import Path
 from typing import Any
 from unittest.mock import Mock, patch
@@ -20,10 +21,16 @@
     GetTaskRequest,
     GetTaskResponse,
     PushTaskRequest,
+    PushTaskResponse,
     SetTaskStatusRequest,
     SetTaskStatusResponse,
     TaskActivation,
 )
+from sentry_protos.taskbroker.v1.taskbroker_pb2_grpc import (
+    WorkerServiceServicer,
+    WorkerServiceStub,
+    add_WorkerServiceServicer_to_server,
+)
 
 from taskbroker_client.constants import DEFAULT_WORKER_HEALTH_CHECK_SEC_PER_TOUCH
 from taskbroker_client.metrics import NoOpMetricsBackend
@@ -394,6 +401,53 @@ def test_request_signature_server_interceptor_skips_grpc_health_check() -> None:
     assert interceptor.intercept_service(lambda _: handler, handler_call_details) is handler
 
 
+@pytest.mark.parametrize(
+    ("signature", "expected"),
+    [
+        (None, grpc.StatusCode.UNAUTHENTICATED),
+        (lambda body: _push_task_hmac(b"wrong", body), grpc.StatusCode.UNAUTHENTICATED),
+        (lambda body: _push_task_hmac(b"secret", body + b"x"), grpc.StatusCode.UNAUTHENTICATED),
+        (lambda body: _push_task_hmac(b"secret", body), grpc.StatusCode.OK),
+    ],
+    ids=["missing", "wrong_secret", "wrong_body", "valid"],
+)
+def test_request_signature_server_interceptor_on_grpc_server(
+    signature: Callable[[bytes], str] | None, expected: grpc.StatusCode
+) -> None:
+    received: list[PushTaskRequest] = []
+
+    class Worker(WorkerServiceServicer):
+        def PushTask(
+            self, request: PushTaskRequest, context: grpc.ServicerContext
+        ) -> PushTaskResponse:
+            received.append(request)
+            return PushTaskResponse()
+
+    server = grpc.server(
+        ThreadPoolExecutor(max_workers=2),
+        interceptors=[RequestSignatureServerInterceptor(["secret"])],
+    )
+    add_WorkerServiceServicer_to_server(Worker(), server)
+    port = server.add_insecure_port("127.0.0.1:0")
+    server.start()
+    try:
+        with grpc.insecure_channel(f"127.0.0.1:{port}") as channel:
+            request = PushTaskRequest(task=TaskActivation(id="abc123", taskname="do_thing"))
+            metadata: tuple[tuple[str, str | bytes], ...] = ()
+            if signature is not None:
+                metadata = (("sentry-signature", signature(request.SerializeToString())),)
+            try:
+                WorkerServiceStub(channel).PushTask(request, metadata=metadata, timeout=5)
+                code = grpc.StatusCode.OK
+            except grpc.RpcError as e:
+                code = e.code()
+    finally:
+        server.stop(None)
+
+    assert code == expected
+    assert received == ([request] if expected == grpc.StatusCode.OK else [])
+
+
 def test_get_task_with_namespace() -> None:
     channel = MockChannel()
     channel.add_response(