From 484c2310bbaed1357a19254f2b35c5b213ebe6f5 Mon Sep 17 00:00:00 2001 From: Yesudeep Mangalapilly Date: Sun, 1 Feb 2026 13:28:41 -0800 Subject: [PATCH] fix(py): implement SigV4 signing for AWS X-Ray OTLP exporter The AwsXRayOtlpExporter was not signing requests with AWS SigV4, which would cause authentication failures when sending traces to AWS X-Ray endpoints without a collector. This fix: - Adds SigV4SigningAdapter that intercepts HTTP requests and signs them with AWS SigV4 authentication using botocore - Creates _create_sigv4_session() helper that builds a requests Session with the signing adapter mounted for HTTPS - Passes the signing session to OTLPSpanExporter so all requests are automatically signed - Removes the unused _sign_request() method The implementation uses botocore's SigV4Auth to add Authorization, X-Amz-Date, and X-Amz-Security-Token headers to each request. Added tests: - SigV4SigningAdapter initialization and signing - Session creation with adapter mounting - Exporter session configuration verification Fixes: gemini-code-assist review on PR #4390 --- py/plugins/aws-bedrock/src/genkit/py.typed | 1 + .../genkit/plugins/aws/telemetry/tracing.py | 170 ++++++++++++++---- py/plugins/aws/tests/aws_telemetry_test.py | 127 +++++++++++++ py/samples/_common.sh | 32 ++++ py/samples/aws-bedrock-hello/run.sh | 5 +- py/samples/aws-hello/README.md | 6 +- py/samples/aws-hello/run.sh | 11 +- py/samples/aws-hello/src/main.py | 9 +- py/samples/google-genai-image/src/main.py | 2 +- 9 files changed, 314 insertions(+), 49 deletions(-) diff --git a/py/plugins/aws-bedrock/src/genkit/py.typed b/py/plugins/aws-bedrock/src/genkit/py.typed index e69de29bb2..8b13789179 100644 --- a/py/plugins/aws-bedrock/src/genkit/py.typed +++ b/py/plugins/aws-bedrock/src/genkit/py.typed @@ -0,0 +1 @@ + diff --git a/py/plugins/aws/src/genkit/plugins/aws/telemetry/tracing.py b/py/plugins/aws/src/genkit/plugins/aws/telemetry/tracing.py index 032b49b6d3..2060ada94a 100644 --- a/py/plugins/aws/src/genkit/plugins/aws/telemetry/tracing.py +++ b/py/plugins/aws/src/genkit/plugins/aws/telemetry/tracing.py @@ -224,8 +224,9 @@ import os import uuid from collections.abc import Callable, Mapping, MutableMapping, Sequence -from typing import Any +from typing import Any, cast +import requests import structlog from botocore.auth import SigV4Auth from botocore.awsrequest import AWSRequest @@ -238,6 +239,8 @@ from opentelemetry.sdk.trace import ReadableSpan, TracerProvider from opentelemetry.sdk.trace.export import SpanExporter, SpanExportResult from opentelemetry.sdk.trace.sampling import Sampler +from requests.adapters import HTTPAdapter +from urllib3.util.retry import Retry from genkit.core.environment import is_dev_environment from genkit.core.trace.adjusting_exporter import AdjustingTraceExporter, RedactedSpan @@ -446,12 +449,125 @@ def _inject_trace_context(self, event_dict: MutableMapping[str, Any]) -> Mutable return event_dict +class SigV4SigningAdapter(HTTPAdapter): + """HTTP adapter that signs requests with AWS SigV4 authentication. + + This adapter intercepts all HTTP requests and signs them with AWS SigV4 + before sending. This enables the OTLP exporter to authenticate with + AWS services like X-Ray that require SigV4. + + Example: + ```python + session = requests.Session() + adapter = SigV4SigningAdapter( + credentials=botocore_session.get_credentials(), + region='us-west-2', + service='xray', + ) + session.mount('https://', adapter) + ``` + """ + + def __init__( + self, + credentials: Any, # noqa: ANN401 - botocore credentials type + region: str, + service: str = 'xray', + **kwargs: Any, # noqa: ANN401 + ) -> None: + """Initialize the SigV4 signing adapter. + + Args: + credentials: Botocore credentials object. + region: AWS region for signing. + service: AWS service name for signing (default: 'xray'). + **kwargs: Additional arguments passed to HTTPAdapter. + """ + super().__init__(**kwargs) + self._credentials = credentials + self._region = region + self._service = service + + def send( # type: ignore[override] + self, + request: requests.PreparedRequest, + **kwargs: Any, # noqa: ANN401 + ) -> requests.Response: + """Send the request after signing it with SigV4. + + Args: + request: The prepared request to send. + **kwargs: Additional arguments passed to the parent send method. + + Returns: + The response from the server. + """ + if self._credentials is not None: + # Sign the request + aws_request = AWSRequest( + method=request.method or 'POST', + url=request.url or '', + headers=dict(request.headers) if request.headers else {}, + data=request.body or b'', + ) + SigV4Auth(self._credentials, self._service, self._region).add_auth(aws_request) + + # Update the original request with signed headers + if request.headers is None: + request.headers = cast(Any, {}) + request.headers.update(dict(aws_request.headers)) + + return super().send(request, **kwargs) + + +def _create_sigv4_session( + credentials: Any, # noqa: ANN401 - botocore credentials type + region: str, + service: str = 'xray', +) -> requests.Session: + """Create a requests Session that signs all requests with SigV4. + + Args: + credentials: Botocore credentials object. + region: AWS region for signing. + service: AWS service name for signing. + + Returns: + A configured requests Session with SigV4 signing. + """ + session = requests.Session() + + # Configure retry logic + retry = Retry( + total=3, + backoff_factor=0.5, + status_forcelist=[500, 502, 503, 504], + ) + + # Create signing adapter with retry + adapter = SigV4SigningAdapter( + credentials=credentials, + region=region, + service=service, + max_retries=retry, + ) + + # Mount the signing adapter for HTTPS requests + session.mount('https://', adapter) + + return session + + class AwsXRayOtlpExporter(SpanExporter): """OTLP/HTTP exporter with AWS SigV4 authentication for X-Ray. This exporter sends spans via OTLP/HTTP to the AWS X-Ray endpoint, signing each request with AWS SigV4 authentication using botocore. + The SigV4 signing is implemented via a custom requests Session adapter + that intercepts all outgoing HTTP requests and adds the required + Authorization, X-Amz-Date, and X-Amz-Security-Token headers. + Args: region: AWS region for the X-Ray endpoint. error_handler: Optional callback for export errors. @@ -485,47 +601,29 @@ def __init__( self._botocore_session = BotocoreSession() self._credentials = self._botocore_session.get_credentials() - # Create the underlying OTLP exporter - self._otlp_exporter = OTLPSpanExporter( - endpoint=self._endpoint, - ) - - def _sign_request(self, headers: dict[str, str], body: bytes) -> dict[str, str]: - """Sign the request with AWS SigV4. - - Args: - headers: Request headers. - body: Request body bytes. - - Returns: - Updated headers with SigV4 signature. - """ if self._credentials is None: - logger.warning('No AWS credentials found for SigV4 signing') - return headers - - # Create an AWS request for signing - aws_request = AWSRequest( - method='POST', - url=self._endpoint, - headers=headers, - data=body, - ) + logger.warning( + 'No AWS credentials found for SigV4 signing. X-Ray export will fail without valid credentials.' + ) - # Sign the request - SigV4Auth(self._credentials, 'xray', self._region).add_auth(aws_request) + # Create a session with SigV4 signing adapter + signing_session = _create_sigv4_session( + credentials=self._credentials, + region=region, + service='xray', + ) - # Return the signed headers - return dict(aws_request.headers) + # Create the underlying OTLP exporter with signing session + self._otlp_exporter = OTLPSpanExporter( + endpoint=self._endpoint, + session=signing_session, + ) def export(self, spans: Sequence[ReadableSpan]) -> SpanExportResult: - """Export spans to AWS X-Ray via OTLP/HTTP. + """Export spans to AWS X-Ray via OTLP/HTTP with SigV4 authentication. - Note: - The current implementation delegates to the underlying OTLP exporter - which may not include SigV4 headers. For production use with - collector-less export, consider using ADOT auto-instrumentation - or the ADOT collector. + All requests are automatically signed with AWS SigV4 using the + credentials from the botocore session. Args: spans: A sequence of OpenTelemetry ReadableSpan objects to export. diff --git a/py/plugins/aws/tests/aws_telemetry_test.py b/py/plugins/aws/tests/aws_telemetry_test.py index 474c720da4..49c35e21da 100644 --- a/py/plugins/aws/tests/aws_telemetry_test.py +++ b/py/plugins/aws/tests/aws_telemetry_test.py @@ -29,13 +29,16 @@ from unittest import mock import pytest +import requests from genkit.plugins.aws.telemetry.tracing import ( XRAY_OTLP_ENDPOINT_PATTERN, AwsAdjustingTraceExporter, AwsTelemetry, AwsXRayOtlpExporter, + SigV4SigningAdapter, TimeAdjustedSpan, + _create_sigv4_session, _resolve_region, add_aws_telemetry, ) @@ -116,6 +119,130 @@ def test_exporter_with_error_handler(self) -> None: ) assert exporter._error_handler is not None + def test_exporter_uses_sigv4_session(self) -> None: + """Exporter should use a session with SigV4 signing adapter mounted.""" + exporter = AwsXRayOtlpExporter(region='us-west-2') + # The OTLP exporter should have a session configured + assert exporter._otlp_exporter._session is not None + + +class TestSigV4SigningAdapter: + """Tests for the SigV4SigningAdapter class.""" + + def test_adapter_initialization(self) -> None: + """Adapter should initialize with credentials and region.""" + mock_credentials = mock.MagicMock() + adapter = SigV4SigningAdapter( + credentials=mock_credentials, + region='us-west-2', + service='xray', + ) + assert adapter._credentials == mock_credentials + assert adapter._region == 'us-west-2' + assert adapter._service == 'xray' + + def test_adapter_default_service(self) -> None: + """Adapter should default to xray service.""" + mock_credentials = mock.MagicMock() + adapter = SigV4SigningAdapter( + credentials=mock_credentials, + region='eu-west-1', + ) + assert adapter._service == 'xray' + + def test_adapter_signs_request(self) -> None: + """Adapter should add SigV4 headers to request.""" + # Create mock credentials + mock_credentials = mock.MagicMock() + mock_credentials.access_key = 'AKIAIOSFODNN7EXAMPLE' + mock_credentials.secret_key = 'wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY' + mock_credentials.token = None + + adapter = SigV4SigningAdapter( + credentials=mock_credentials, + region='us-west-2', + service='xray', + ) + + # Create a mock request + request = requests.PreparedRequest() + request.prepare( + method='POST', + url='https://xray.us-west-2.amazonaws.com/v1/traces', + headers={'Content-Type': 'application/x-protobuf'}, + data=b'test-payload', + ) + + # Mock the parent send method + with mock.patch.object( + adapter.__class__.__bases__[0], + 'send', + return_value=mock.MagicMock(status_code=200), + ): + adapter.send(request) + + # Verify that the request now has Authorization header + assert request.headers is not None + assert 'Authorization' in request.headers + assert 'AWS4-HMAC-SHA256' in request.headers['Authorization'] + + def test_adapter_handles_none_credentials(self) -> None: + """Adapter should handle None credentials gracefully.""" + adapter = SigV4SigningAdapter( + credentials=None, + region='us-west-2', + ) + + request = requests.PreparedRequest() + request.prepare( + method='POST', + url='https://xray.us-west-2.amazonaws.com/v1/traces', + headers={}, + data=b'test', + ) + + # Should not raise, just skip signing + with mock.patch.object( + adapter.__class__.__bases__[0], + 'send', + return_value=mock.MagicMock(status_code=200), + ): + adapter.send(request) + + # No Authorization header should be set + assert request.headers is not None + assert 'Authorization' not in request.headers + + +class TestCreateSigV4Session: + """Tests for the _create_sigv4_session function.""" + + def test_session_has_adapter_mounted(self) -> None: + """Session should have SigV4 adapter mounted for HTTPS.""" + mock_credentials = mock.MagicMock() + session = _create_sigv4_session( + credentials=mock_credentials, + region='us-west-2', + ) + + # Get the adapter for an HTTPS URL + adapter = session.get_adapter('https://xray.us-west-2.amazonaws.com') + assert isinstance(adapter, SigV4SigningAdapter) + + def test_session_adapter_has_correct_region(self) -> None: + """Session adapter should be configured with correct region.""" + mock_credentials = mock.MagicMock() + session = _create_sigv4_session( + credentials=mock_credentials, + region='eu-west-1', + service='xray', + ) + + adapter = session.get_adapter('https://xray.eu-west-1.amazonaws.com') + assert isinstance(adapter, SigV4SigningAdapter) + assert adapter._region == 'eu-west-1' + assert adapter._service == 'xray' + class TestAwsAdjustingTraceExporter: """Tests for the AwsAdjustingTraceExporter class.""" diff --git a/py/samples/_common.sh b/py/samples/_common.sh index 05b99965bc..b61b65a197 100644 --- a/py/samples/_common.sh +++ b/py/samples/_common.sh @@ -50,6 +50,38 @@ check_env_var() { local var_name="$1" local get_url="$2" + local current_val="${!var_name:-}" + + # Prompt if running interactively + # We check -t 0 (stdin is TTY) and also explicit check for /dev/tty availability + if [[ -t 0 ]] && [ -c /dev/tty ]; then + local display_val="${current_val}" + + # Simple masking for keys + if [[ "$var_name" == *"API_KEY"* || "$var_name" == *"SECRET"* ]]; then + if [[ -n "$current_val" ]]; then + display_val="******" + fi + fi + + echo -en "${BLUE}Enter ${var_name}${NC}" + if [[ -n "$display_val" ]]; then + echo -en " [${YELLOW}${display_val}${NC}]: " + else + echo -n ": " + fi + + local input_val + # Safely read from TTY + if read -r input_val < /dev/tty; then + if [[ -n "$input_val" ]]; then + export "$var_name"="$input_val" + fi + fi + # Only print newline if we actually prompted + echo "" + fi + if [[ -z "${!var_name:-}" ]]; then echo -e "${YELLOW}Warning: ${var_name} not set${NC}" if [[ -n "$get_url" ]]; then diff --git a/py/samples/aws-bedrock-hello/run.sh b/py/samples/aws-bedrock-hello/run.sh index 23c7880973..824fc5da98 100755 --- a/py/samples/aws-bedrock-hello/run.sh +++ b/py/samples/aws-bedrock-hello/run.sh @@ -46,7 +46,10 @@ esac print_banner "AWS Bedrock Hello World" "☁️" -check_env_var "AWS_REGION" "https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started.html" || true +# Set default region if not provided +export AWS_REGION="${AWS_REGION:-us-east-1}" + +check_env_var "AWS_REGION" "https://docs.aws.amazon.com/bedrock/latest/userguide/getting-started.html" install_deps diff --git a/py/samples/aws-hello/README.md b/py/samples/aws-hello/README.md index b0b6a52d6b..75ef202874 100644 --- a/py/samples/aws-hello/README.md +++ b/py/samples/aws-hello/README.md @@ -8,14 +8,14 @@ using the AWS plugin. 1. **AWS Credentials**: Configure via environment variables, credentials file, or IAM role 2. **AWS Region**: Set `AWS_REGION` environment variable 3. **IAM Permissions**: Attach `AWSXrayWriteOnlyPolicy` to your role/user -4. **Google AI API Key**: Set `GOOGLE_GENAI_API_KEY` (or use any other model provider) +4. **Gemini API Key**: Set `GEMINI_API_KEY` (or use any other model provider) ## Quick Start ```bash # Set required environment variables export AWS_REGION=us-west-2 -export GOOGLE_GENAI_API_KEY=your-api-key +export GEMINI_API_KEY=your-api-key # Run the demo ./run.sh @@ -49,7 +49,7 @@ appear as child spans in X-Ray. | `AWS_ACCESS_KEY_ID` | No* | AWS access key | | `AWS_SECRET_ACCESS_KEY` | No* | AWS secret key | | `AWS_PROFILE` | No* | AWS profile from credentials file | -| `GOOGLE_GENAI_API_KEY` | Yes | Google AI API key | +| `GEMINI_API_KEY` | Yes | Gemini API key | *At least one form of AWS credentials is required. diff --git a/py/samples/aws-hello/run.sh b/py/samples/aws-hello/run.sh index 6773570552..a57af79f45 100755 --- a/py/samples/aws-hello/run.sh +++ b/py/samples/aws-hello/run.sh @@ -10,7 +10,7 @@ # Prerequisites: # - AWS credentials configured (env vars, credentials file, or IAM role) # - AWS_REGION environment variable set -# - GOOGLE_GENAI_API_KEY for the model provider +# - GEMINI_API_KEY for the model provider # # Usage: # ./run.sh # Start the demo with Dev UI @@ -33,7 +33,7 @@ print_help() { echo " AWS_ACCESS_KEY_ID AWS access key ID (or use credentials file)" echo " AWS_SECRET_ACCESS_KEY AWS secret access key (or use credentials file)" echo " AWS_PROFILE AWS profile name from credentials file" - echo " GOOGLE_GENAI_API_KEY Required. Google AI API key for the model" + echo " GEMINI_API_KEY Required. Gemini API key for the model" echo "" echo "Setup Guides:" echo " AWS X-Ray: https://docs.aws.amazon.com/xray/" @@ -50,8 +50,11 @@ esac print_banner "AWS Telemetry Hello World" "📊" -check_env_var "AWS_REGION" "https://docs.aws.amazon.com/xray/" || true -check_env_var "GOOGLE_GENAI_API_KEY" "https://aistudio.google.com/app/apikey" +# Set default region if not provided +export AWS_REGION="${AWS_REGION:-us-west-2}" + +check_env_var "AWS_REGION" "https://docs.aws.amazon.com/xray/" +check_env_var "GEMINI_API_KEY" "https://aistudio.google.com/app/apikey" install_deps diff --git a/py/samples/aws-hello/src/main.py b/py/samples/aws-hello/src/main.py index d7f318ac08..7a05852ba8 100644 --- a/py/samples/aws-hello/src/main.py +++ b/py/samples/aws-hello/src/main.py @@ -63,13 +63,13 @@ Prerequisites: 1. AWS credentials configured (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY, or IAM role) 2. AWS_REGION environment variable set - 3. GOOGLE_GENAI_API_KEY environment variable set + 3. GEMINI_API_KEY environment variable set 4. IAM policy: AWSXrayWriteOnlyPolicy attached to your role/user Running the Demo: 1. Set environment variables: export AWS_REGION=us-west-2 - export GOOGLE_GENAI_API_KEY=your-api-key + export GEMINI_API_KEY=your-api-key 2. Run the sample: ./run.sh @@ -97,6 +97,7 @@ install_rich_traceback(show_locals=True, width=120, extra_lines=3) import asyncio +import os import structlog from pydantic import BaseModel, Field @@ -119,13 +120,13 @@ class HelloInput(BaseModel): # Enable AWS X-Ray telemetry # Traces will be exported to the AWS region specified in AWS_REGION -add_aws_telemetry() +add_aws_telemetry(region=os.environ.get('AWS_REGION')) # Configure the model provider # You can use any model provider - we use Google GenAI here for simplicity ai = Genkit( plugins=[GoogleAI()], - model='googleai/gemini-2.0-flash', + model='googleai/gemini-3-flash-preview', ) diff --git a/py/samples/google-genai-image/src/main.py b/py/samples/google-genai-image/src/main.py index 6868b56de9..42d6f435e0 100755 --- a/py/samples/google-genai-image/src/main.py +++ b/py/samples/google-genai-image/src/main.py @@ -290,7 +290,7 @@ async def photo_move_veo(_: object, context: ActionRunContext | None = None) -> 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==' ) - api_key = os.environ.get('GEMINI_API_KEY') or os.environ.get('GOOGLE_GENAI_API_KEY') + api_key = os.environ.get('GEMINI_API_KEY') if not api_key: raise ValueError('GEMINI_API_KEY not set')