diff --git a/clients/python/src/taskbroker_client/worker/client.py b/clients/python/src/taskbroker_client/worker/client.py index 717d0379..98aa1fa0 100644 --- a/clients/python/src/taskbroker_client/worker/client.py +++ b/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) diff --git a/clients/python/tests/worker/test_client.py b/clients/python/tests/worker/test_client.py index 33d9083a..3db69da9 100644 --- a/clients/python/tests/worker/test_client.py +++ b/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(