diff --git a/src/harness_sdk/instrumentation/__init__.py b/src/harness_sdk/instrumentation/__init__.py index 93af348..147ddb4 100644 --- a/src/harness_sdk/instrumentation/__init__.py +++ b/src/harness_sdk/instrumentation/__init__.py @@ -25,11 +25,16 @@ class BaseInstrumentorWrapper: def __init__(self): '''constructor''' logger.debug('Entering BaseInstrumentorWrapper constructor.') - self._process_request_headers = Config()._instance.config.data_capture.http_headers.request.value - self._process_response_headers = Config()._instance.config.data_capture.http_headers.response.value - self._process_request_body = Config()._instance.config.data_capture.http_body.request.value - self._process_response_body = Config()._instance.config.data_capture.http_body.response.value - self._max_body_size = Config()._instance.config.data_capture.body_max_size_bytes.value + data_capture = Config()._instance.config.data_capture + self._process_request_headers = data_capture.http_headers.request.value + self._process_response_headers = data_capture.http_headers.response.value + self._process_request_body = data_capture.http_body.request.value + self._process_response_body = data_capture.http_body.response.value + self._process_rpc_request_metadata = data_capture.rpc_metadata.request.value + self._process_rpc_response_metadata = data_capture.rpc_metadata.response.value + self._process_rpc_request_body = data_capture.rpc_body.request.value + self._process_rpc_response_body = data_capture.rpc_body.response.value + self._max_body_size = data_capture.body_max_size_bytes.value proto_allowed_content_types = Config()._instance.config.data_capture.allowed_content_types self._allowed_content_types = [item.value for item in proto_allowed_content_types] @@ -151,11 +156,11 @@ def generic_rpc_request_handler(self, logger.debug('Span is Recording!') lowercased_headers = self.lowercase_headers(request_headers) - # Log rpc metatdata if requested - if self._process_request_headers: + # Log rpc metadata if requested + if self._process_rpc_request_metadata: self.add_headers_to_span(self.RPC_REQUEST_METADATA_PREFIX, span, lowercased_headers) # Log rpc body if requested - if self._process_response_body: + if self._process_rpc_request_body: request_body_str = str(request_body) request_body_str = self.grab_first_n_bytes(request_body_str) span.set_attribute(self.RPC_REQUEST_BODY_PREFIX, @@ -182,12 +187,12 @@ def generic_rpc_response_handler(self, logger.debug('Span is Recording!') lowercased_headers = self.lowercase_headers(response_headers) - # Log rpc metadata if requested? - if self._process_response_headers: + # Log rpc metadata if requested + if self._process_rpc_response_metadata: logger.debug('Dumping Response Headers:') self.add_headers_to_span(self.RPC_RESPONSE_METADATA_PREFIX, span, lowercased_headers) # Log rpc body if requested - if self._process_response_body: + if self._process_rpc_response_body: response_body_str = str(response_body) logger.debug('Processing response body') response_body_str = self.grab_first_n_bytes(response_body_str) diff --git a/test/instrumentation/grpc/rpc_body_disabled_config.yaml b/test/instrumentation/grpc/rpc_body_disabled_config.yaml new file mode 100644 index 0000000..df27049 --- /dev/null +++ b/test/instrumentation/grpc/rpc_body_disabled_config.yaml @@ -0,0 +1,7 @@ +data_capture: + rpc_body: + request: false + response: false + rpc_metadata: + request: false + response: false diff --git a/test/instrumentation/grpc/test_grpc_rpc_body_config.py b/test/instrumentation/grpc/test_grpc_rpc_body_config.py new file mode 100644 index 0000000..4ccffaf --- /dev/null +++ b/test/instrumentation/grpc/test_grpc_rpc_body_config.py @@ -0,0 +1,89 @@ +"""Verify rpc_body / rpc_metadata config gates gRPC span body capture.""" +import json +import os +from concurrent import futures + +import grpc +import pytest +from opentelemetry.trace import SpanKind + +from harness_sdk.agent import Agent +from harness_sdk.config.config import Config +from harness_sdk.instrumentation.instrumentation_definitions import _uninstrument_all +from test import configure_inmemory_span_exporter +from test.instrumentation.grpc import helloworld_pb2, helloworld_pb2_grpc + + +@pytest.fixture +def agent_with_rpc_body_disabled(): + config_dir = os.path.dirname(__file__) + os.environ["HA_CONFIG_FILE"] = os.path.join(config_dir, "rpc_body_disabled_config.yaml") + os.environ["HA_ENABLE_CONSOLE_SPAN_EXPORTER"] = "true" + os.environ["HARNESS_ENABLE_API"] = "true" + for key in ( + "HARNESS_ENABLE_AI_LITELLM", + "HARNESS_ENABLE_AI_OPENAI", + "HARNESS_ENABLE_AI_ANTHROPIC", + "HARNESS_ENABLE_AI_GOOGLE_GENAI", + "HARNESS_ENABLE_AI_MCP", + ): + os.environ.pop(key, None) + _uninstrument_all() + Config._instance = None + Agent._instance = None + agent = Agent() + agent._init.init_trace_provider() # pylint: disable=protected-access + yield agent + _uninstrument_all() + Config._instance = None + Agent._instance = None + os.environ.pop("HA_CONFIG_FILE", None) + + +@pytest.fixture +def exporter(agent_with_rpc_body_disabled): + exporter = configure_inmemory_span_exporter(agent_with_rpc_body_disabled) + yield exporter + exporter.clear() + + +def _server_span(spans): + for span in spans: + attrs = span.attributes or {} + if attrs.get("rpc.system") == "grpc" and span.kind == SpanKind.SERVER: + return json.loads(span.to_json()) + raise AssertionError("No gRPC server span found") + + +def test_grpc_respects_rpc_body_disabled(agent_with_rpc_body_disabled, exporter): + agent_with_rpc_body_disabled.instrument() + + class Greeter(helloworld_pb2_grpc.GreeterServicer): + def SayHello(self, request, context): # pylint: disable=unused-argument + return helloworld_pb2.HelloReply(message=f"Hello, {request.name}!") + + executor = futures.ThreadPoolExecutor(max_workers=10) + server = grpc.server(executor) + try: + helloworld_pb2_grpc.add_GreeterServicer_to_server(Greeter(), server) + port = server.add_insecure_port("[::]:0") + server.start() + + with grpc.insecure_channel(f"127.0.0.1:{port}") as channel: + stub = helloworld_pb2_grpc.GreeterStub(channel) + response = stub.SayHello(helloworld_pb2.HelloRequest(name="world")) + assert response.message == "Hello, world!" + + span_object = _server_span(exporter.get_finished_spans()) + attrs = span_object["attributes"] + + assert attrs["rpc.method"] == "SayHello" + assert attrs["rpc.grpc.status_code"] == 0 + assert "rpc.request.body" not in attrs + assert "rpc.response.body" not in attrs + assert not any(key.startswith("rpc.request.metadata.") for key in attrs) + assert not any(key.startswith("rpc.response.metadata.") for key in attrs) + exporter.clear() + finally: + server.stop(grace=0) + executor.shutdown(wait=False, cancel_futures=True)