Add SecretsManager and StepFunctions support.

Dockerfile

@@ -3,7 +3,9 @@
 #   The packages are installed in the `/autoinstrumentation` directory. This is required as when instrumenting the pod by CWOperator,
 #   one init container will be created to copy all the content in `/autoinstrumentation` directory to app's container. Then
 #   update the `PYTHONPATH` environment variable accordingly. Then in the second stage, copy the directory to `/autoinstrumentation`.
-FROM python:3.11 AS build
+
+# Stage 1: Install ADOT Python in the /operator-build folder
+FROM public.ecr.aws/docker/library/python:3.11 AS build
 
 WORKDIR /operator-build
 
@@ -18,11 +20,42 @@ RUN sed -i "/opentelemetry-exporter-otlp-proto-grpc/d" ./aws-opentelemetry-distr
 
 RUN mkdir workspace && pip install --target workspace ./aws-opentelemetry-distro
 
-FROM public.ecr.aws/amazonlinux/amazonlinux:minimal
+# Stage 2: Build the cp-utility binary
+FROM public.ecr.aws/docker/library/rust:1.75 as builder
+
+WORKDIR /usr/src/cp-utility
+COPY ./tools/cp-utility .
+
+## TARGETARCH is defined by buildx
+# https://docs.docker.com/engine/reference/builder/#automatic-platform-args-in-the-global-scope
+ARG TARGETARCH
+
+# Run validations and audit only on amd64 because it is faster and those two steps
+# are only used to validate the source code and don't require anything that is
+# architecture specific.
+
+# Validations
+# Validate formatting
+RUN if [ $TARGETARCH = "amd64" ]; then rustup component add rustfmt && cargo fmt --check ; fi
+
+# Audit dependencies
+RUN if [ $TARGETARCH = "amd64" ]; then cargo install cargo-audit && cargo audit ; fi
+
+
+# Cross-compile based on the target platform.
+RUN if [ $TARGETARCH = "amd64" ]; then export ARCH="x86_64" ; \
+    elif [ $TARGETARCH = "arm64" ]; then export ARCH="aarch64" ; \
+    else false; \
+    fi \
+    && rustup target add ${ARCH}-unknown-linux-musl \
+    && cargo test  --target ${ARCH}-unknown-linux-musl \
+    && cargo install --target ${ARCH}-unknown-linux-musl --path . --root .
+
+# Stage 3: Build the distribution image by copying the THIRD-PARTY-LICENSES, the custom built cp command from stage 2, and the installed ADOT Python from stage 1 to their respective destinations
+FROM scratch
 
 # Required to copy attribute files to distributed docker images
 ADD THIRD-PARTY-LICENSES ./THIRD-PARTY-LICENSES
 
+COPY --from=builder /usr/src/cp-utility/bin/cp-utility /bin/cp
 COPY --from=build /operator-build/workspace /autoinstrumentation
-
-RUN chmod -R go+r /autoinstrumentation

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_attribute_keys.py

@@ -16,3 +16,6 @@
 AWS_QUEUE_URL: str = "aws.sqs.queue_url"
 AWS_QUEUE_NAME: str = "aws.sqs.queue_name"
 AWS_STREAM_NAME: str = "aws.kinesis.stream_name"
+AWS_SECRET_ARN: str = "aws.secretsmanager.secret_arn"
+AWS_STATE_MACHINE_ARN: str = "aws.stepfunctions.state_machine_arn"
+AWS_ACTIVITY_ARN: str = "aws.stepfunctions.activity_arn"

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/_aws_metric_attribute_generator.py

@@ -6,6 +6,7 @@
 from urllib.parse import ParseResult, urlparse
 
 from amazon.opentelemetry.distro._aws_attribute_keys import (
+    AWS_ACTIVITY_ARN,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
     AWS_QUEUE_NAME,
@@ -14,7 +15,9 @@
     AWS_REMOTE_RESOURCE_IDENTIFIER,
     AWS_REMOTE_RESOURCE_TYPE,
     AWS_REMOTE_SERVICE,
+    AWS_SECRET_ARN,
     AWS_SPAN_KIND,
+    AWS_STATE_MACHINE_ARN,
     AWS_STREAM_NAME,
 )
 from amazon.opentelemetry.distro._aws_span_processing_util import (
@@ -78,6 +81,8 @@
 _NORMALIZED_KINESIS_SERVICE_NAME: str = "AWS::Kinesis"
 _NORMALIZED_S3_SERVICE_NAME: str = "AWS::S3"
 _NORMALIZED_SQS_SERVICE_NAME: str = "AWS::SQS"
+_NORMALIZED_SECRETSMANAGER_SERVICE_NAME: str = "AWS::SecretsManager"
+_NORMALIZED_STEPFUNCTIONS_SERVICE_NAME: str = "AWS::StepFunctions"
 _DB_CONNECTION_STRING_TYPE: str = "DB::Connection"
 
 # Special DEPENDENCY attribute value if GRAPHQL_OPERATION_TYPE attribute key is present.
@@ -290,7 +295,11 @@ def _normalize_remote_service_name(span: ReadableSpan, service_name: str) -> str
     resource format</a> as much as possible. Long term, we would like to normalize service name in the upstream.
     """
     if is_aws_sdk_span(span):
-        return "AWS::" + service_name
+        aws_sdk_service_mapping = {
+            "Secrets Manager": _NORMALIZED_SECRETSMANAGER_SERVICE_NAME,
+            "SFN": _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME,
+        }
+        return aws_sdk_service_mapping.get(service_name, "AWS::" + service_name)
     return service_name
 
 
@@ -372,6 +381,15 @@ def _set_remote_type_and_identifier(span: ReadableSpan, attributes: BoundedAttri
             remote_resource_identifier = _escape_delimiters(
                 SqsUrlParser.get_queue_name(span.attributes.get(AWS_QUEUE_URL))
             )
+        elif is_key_present(span, AWS_SECRET_ARN):
+            remote_resource_type = _NORMALIZED_SECRETSMANAGER_SERVICE_NAME + "::Secret"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_SECRET_ARN))
+        elif is_key_present(span, AWS_STATE_MACHINE_ARN):
+            remote_resource_type = _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME + "::StateMachine"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_STATE_MACHINE_ARN))
+        elif is_key_present(span, AWS_ACTIVITY_ARN):
+            remote_resource_type = _NORMALIZED_STEPFUNCTIONS_SERVICE_NAME + "::Activity"
+            remote_resource_identifier = _escape_delimiters(span.attributes.get(AWS_ACTIVITY_ARN))
     elif is_db_span(span):
         remote_resource_type = _DB_CONNECTION_STRING_TYPE
         remote_resource_identifier = _get_db_connection(span)

aws-opentelemetry-distro/src/amazon/opentelemetry/distro/patches/_botocore_patches.py

@@ -5,18 +5,21 @@
 
 from opentelemetry.instrumentation.botocore.extensions import _KNOWN_EXTENSIONS
 from opentelemetry.instrumentation.botocore.extensions.sqs import _SqsExtension
-from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension
+from opentelemetry.instrumentation.botocore.extensions.types import _AttributeMapT, _AwsSdkExtension, _BotoResultT
 from opentelemetry.semconv.trace import SpanAttributes
+from opentelemetry.trace.span import Span
 
 
 def _apply_botocore_instrumentation_patches() -> None:
     """Botocore instrumentation patches
 
-    Adds patches to provide additional support and Java parity for Kinesis, S3, and SQS.
+    Adds patches to provide additional support and Java parity for Kinesis, S3, SQS, SecretsManager and StepFunctions.
     """
     _apply_botocore_kinesis_patch()
     _apply_botocore_s3_patch()
     _apply_botocore_sqs_patch()
+    _apply_botocore_secretsmanager_patch()
+    _apply_botocore_stepfunctions_patch()
 
 
 def _apply_botocore_kinesis_patch() -> None:
@@ -65,6 +68,26 @@ def patch_extract_attributes(self, attributes: _AttributeMapT):
     _SqsExtension.extract_attributes = patch_extract_attributes
 
 
+def _apply_botocore_secretsmanager_patch() -> None:
+    """Botocore instrumentation patch for SecretsManager
+
+    This patch adds an extension to the upstream's list of known extension for SecretsManager.
+    Extensions allow for custom logic for adding service-specific information to spans,
+    such as attributes. Specifically, we are adding logic to add the AWS_SECRET_ARN attribute.
+    """
+    _KNOWN_EXTENSIONS["secretsmanager"] = _lazy_load(".", "_SecretsManagerExtension")
+
+
+def _apply_botocore_stepfunctions_patch() -> None:
+    """Botocore instrumentation patch for StepFunctions
+
+    This patch adds an extension to the upstream's list of known extension for StepFunctions.
+    Extensions allow for custom logic for adding service-specific information to spans,
+    such as attributes. Specifically, we are adding logic to add the AWS_STATE_MACHINE_ARN attribute.
+    """
+    _KNOWN_EXTENSIONS["stepfunctions"] = _lazy_load(".", "_StepFunctionsExtension")
+
+
 # The OpenTelemetry Authors code
 def _lazy_load(module, cls):
     """Clone of upstream opentelemetry.instrumentation.botocore.extensions.lazy_load
@@ -94,3 +117,48 @@ def extract_attributes(self, attributes: _AttributeMapT):
         stream_name = self._call_context.params.get("StreamName")
         if stream_name:
             attributes["aws.kinesis.stream_name"] = stream_name
+
+
+class _SecretsManagerExtension(_AwsSdkExtension):
+    def extract_attributes(self, attributes: _AttributeMapT):
+        """
+        SecretId can be secret name or secret arn, the function extracts attributes only if the SecretId parameter
+        is provided as arn which starts with 'arn:aws:secretsmanager:'.
+        """
+        secret_id = self._call_context.params.get("SecretId")
+        if secret_id and secret_id.startswith("arn:aws:secretsmanager:"):
+            attributes["aws.secretsmanager.secret_arn"] = secret_id
+
+    # pylint: disable=no-self-use
+    def on_success(self, span: Span, result: _BotoResultT):
+        secret_arn = result.get("ARN")
+        if secret_arn:
+            span.set_attribute(
+                "aws.secretsmanager.secret_arn",
+                secret_arn,
+            )
+
+
+class _StepFunctionsExtension(_AwsSdkExtension):
+    def extract_attributes(self, attributes: _AttributeMapT):
+        state_machine_arn = self._call_context.params.get("stateMachineArn")
+        if state_machine_arn:
+            attributes["aws.stepfunctions.state_machine_arn"] = state_machine_arn
+        activity_arn = self._call_context.params.get("activityArn")
+        if activity_arn:
+            attributes["aws.stepfunctions.activity_arn"] = activity_arn
+
+    # pylint: disable=no-self-use
+    def on_success(self, span: Span, result: _BotoResultT):
+        state_machine_arn = result.get("stateMachineArn")
+        if state_machine_arn:
+            span.set_attribute(
+                "aws.stepfunctions.state_machine_arn",
+                state_machine_arn,
+            )
+        activity_arn = result.get("activityArn")
+        if activity_arn:
+            span.set_attribute(
+                "aws.stepfunctions.activity_arn",
+                activity_arn,
+            )

aws-opentelemetry-distro/tests/amazon/opentelemetry/distro/test_aws_metric_attribute_generator.py

@@ -8,6 +8,7 @@
 from unittest.mock import MagicMock
 
 from amazon.opentelemetry.distro._aws_attribute_keys import (
+    AWS_ACTIVITY_ARN,
     AWS_CONSUMER_PARENT_SPAN_KIND,
     AWS_LOCAL_OPERATION,
     AWS_LOCAL_SERVICE,
@@ -17,7 +18,9 @@
     AWS_REMOTE_RESOURCE_IDENTIFIER,
     AWS_REMOTE_RESOURCE_TYPE,
     AWS_REMOTE_SERVICE,
+    AWS_SECRET_ARN,
     AWS_SPAN_KIND,
+    AWS_STATE_MACHINE_ARN,
     AWS_STREAM_NAME,
 )
 from amazon.opentelemetry.distro._aws_metric_attribute_generator import _AwsMetricAttributeGenerator
@@ -821,6 +824,8 @@ def test_normalize_remote_service_name_aws_sdk(self):
         self.validate_aws_sdk_service_normalization("Kinesis", "AWS::Kinesis")
         self.validate_aws_sdk_service_normalization("S3", "AWS::S3")
         self.validate_aws_sdk_service_normalization("SQS", "AWS::SQS")
+        self.validate_aws_sdk_service_normalization("Secrets Manager", "AWS::SecretsManager")
+        self.validate_aws_sdk_service_normalization("SFN", "AWS::StepFunctions")
 
     def validate_aws_sdk_service_normalization(self, service_name: str, expected_remote_service: str):
         self._mock_attribute([SpanAttributes.RPC_SYSTEM, SpanAttributes.RPC_SERVICE], ["aws-api", service_name])
@@ -977,6 +982,39 @@ def test_sdk_client_span_with_remote_resource_attributes(self):
         self._validate_remote_resource_attributes("AWS::DynamoDB::Table", "aws_table^^name")
         self._mock_attribute([SpanAttributes.AWS_DYNAMODB_TABLE_NAMES], [None])
 
+        # Validate behaviour of AWS_SECRET_ARN attribute, then remove it.
+        self._mock_attribute(
+            [AWS_SECRET_ARN], ["arn:aws:secretsmanager:us-east-1:123456789012:secret:secret_name-lERW9H"], keys, values
+        )
+        self._validate_remote_resource_attributes(
+            "AWS::SecretsManager::Secret", "arn:aws:secretsmanager:us-east-1:123456789012:secret:secret_name-lERW9H"
+        )
+        self._mock_attribute([AWS_SECRET_ARN], [None])
+
+        # Validate behaviour of AWS_STATE_MACHINE_ARN attribute, then remove it.
+        self._mock_attribute(
+            [AWS_STATE_MACHINE_ARN],
+            ["arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine"],
+            keys,
+            values,
+        )
+        self._validate_remote_resource_attributes(
+            "AWS::StepFunctions::StateMachine", "arn:aws:states:us-east-1:123456789012:stateMachine:test_state_machine"
+        )
+        self._mock_attribute([AWS_STATE_MACHINE_ARN], [None])
+
+        # Validate behaviour of AWS_ACTIVITY_ARN attribute, then remove it.
+        self._mock_attribute(
+            [AWS_ACTIVITY_ARN],
+            ["arn:aws:states:us-east-1:007003123456789012:activity:testActivity"],
+            keys,
+            values,
+        )
+        self._validate_remote_resource_attributes(
+            "AWS::StepFunctions::Activity", "arn:aws:states:us-east-1:007003123456789012:activity:testActivity"
+        )
+        self._mock_attribute([AWS_ACTIVITY_ARN], [None])
+
         self._mock_attribute([SpanAttributes.RPC_SYSTEM], [None])
 
     def test_client_db_span_with_remote_resource_attributes(self):