Repository: localstack/localstack

Issue:
Enable X-Ray propagation when publishing SNS to SQS
<!-- Love localstack? Please consider supporting our collective:
👉  https://opencollective.com/localstack/donate -->

# Type of request: This is a ...

[ ] bug report
[x] feature request

# Detailed description

As per AWS documentation (https://docs.aws.amazon.com/xray/latest/devguide/xray-services-sns.html) it is possible to propagate trace context when publishing messages from SNS to SQS. Value of HTTP header "X-Amzn-Trace-Id" set upon publish to SNS topic is converted into "AWSTraceHeader" system attribute of a SQS message.

## Expected behavior

If HTTP header "X-Amzn-Trace-Id" is set upon publish to SNS topic, resulting message gets system attribute "AWSTraceHeader".

## Actual behavior

Header value is not propagated


# Steps to reproduce

Groovy / Spock test case used (note - X-Ray header is added by the OpenTelemetry instrumentation):

```
  @Shared
  int sqsPort
  @Shared
  LocalStackContainer localstack
  @Shared
  AmazonSQSAsyncClient sqsClient

  @Shared
  int snsPort
  @Shared
  LocalStackContainer sns
  @Shared
  AmazonSNSAsyncClient snsClient


  def setupSpec() {

    localstack = new LocalStackContainer(DockerImageName.parse("localstack/localstack:latest"))
      .withServices(LocalStackContainer.Service.SQS, LocalStackContainer.Service.SNS)
      .withEnv("DEBUG", "1")
      .withEnv("SQS_PROVIDER", "elasticmq")
    localstack.start()

    sqsPort = localstack.getMappedPort(4566)

    sqsClient = AmazonSQSAsyncClient.asyncBuilder()
      .withEndpointConfiguration(localstack.getEndpointConfiguration(LocalStackContainer.Service.SQS))
      .withCredentials(localstack.getDefaultCredentialsProvider())
      .build()

    snsPort = localstack.getMappedPort(4566)
    snsClient = AmazonSNSAsyncClient.asyncBuilder()
      .withEndpointConfiguration(localstack.getEndpointConfiguration(LocalStackContainer.Service.SNS))
      .withCredentials(localstack.getDefaultCredentialsProvider())
      .build()
  }

  def createQueue(String queueName) {
    CreateQueueResult cqr = sqsClient.createQueue(queueName)
    sqsClient.setQueueAttributes(cqr.getQueueUrl(), Collections.singletonMap("Policy", policy(queueName)))
    return cqr.getQueueUrl()
  }

  def getQueueArn(String queueUrl) {
    return sqsClient.getQueueAttributes(
      new GetQueueAttributesRequest(queueUrl)
        .withAttributeNames("QueueArn")).getAttributes()
      .get("QueueArn");
  }

  def createAndSubscribeTopic(String topicName, String queueArn) {
    CreateTopicResult ctr = snsClient.createTopic(topicName)
    snsClient.subscribe(ctr.getTopicArn(), "sqs", queueArn)
    return ctr.getTopicArn()
  }

  def "simple SNS producer - SQS consumer services"() {
    setup:
    String queueName = "snsToSqsTestQueue"
    String topicName = "snsToSqsTestTopic"

    String queueUrl = createQueue(queueName)
    String queueArn = getQueueArn(queueUrl)
    String topicArn = createAndSubscribeTopic(topicName, queueArn)

    when:
    snsClient.publish(topicArn, "Hello There")
    ReceiveMessageRequest rmr = new ReceiveMessageRequest(queueUrl).withMessageAttributeNames("test")
    sqsClient.receiveMessage(rmr)

```



Discussion:

Resolving patch:
diff --git a/localstack/constants.py b/localstack/constants.py
--- a/localstack/constants.py
+++ b/localstack/constants.py
@@ -90,7 +90,7 @@
  'mapper-murmur3', 'mapper-size', 'analysis-phonetic', 'analysis-smartcn', 'analysis-stempel', 'analysis-ukrainian']
 # Default ES modules to exclude (save apprx 66MB in the final image)
 ELASTICSEARCH_DELETE_MODULES = ['ingest-geoip']
-ELASTICMQ_JAR_URL = 'https://s3-eu-west-1.amazonaws.com/softwaremill-public/elasticmq-server-0.15.7.jar'
+ELASTICMQ_JAR_URL = 'https://s3-eu-west-1.amazonaws.com/softwaremill-public/elasticmq-server-1.1.0.jar'
 STS_JAR_URL = 'https://repo1.maven.org/maven2/com/amazonaws/aws-java-sdk-sts/1.11.14/aws-java-sdk-sts-1.11.14.jar'
 STEPFUNCTIONS_ZIP_URL = 'https://s3.amazonaws.com/stepfunctionslocal/StepFunctionsLocal.zip'
 KMS_URL_PATTERN = 'https://s3-eu-west-2.amazonaws.com/local-kms/localstack/v3/local-kms.<arch>.bin'
diff --git a/localstack/services/sns/sns_listener.py b/localstack/services/sns/sns_listener.py
--- a/localstack/services/sns/sns_listener.py
+++ b/localstack/services/sns/sns_listener.py
@@ -127,7 +127,7 @@ def forward_request(self, method, path, data, headers):
                     if topic_arn not in SNS_SUBSCRIPTIONS:
                         return make_error(code=404, code_string='NotFound', message='Topic does not exist')
 
-                message_id = publish_message(topic_arn, req_data)
+                message_id = publish_message(topic_arn, req_data, headers)
 
                 # return response here because we do not want the request to be forwarded to SNS backend
                 return make_response(req_action, message_id=message_id)
@@ -269,7 +269,8 @@ def unsubscribe_sqs_queue(queue_url):
                 subscriptions.remove(subscriber)
 
 
-def message_to_subscribers(message_id, message, topic_arn, req_data, subscription_arn=None, skip_checks=False):
+def message_to_subscribers(message_id, message, topic_arn, req_data, headers, subscription_arn=None, skip_checks=False):
+
     subscriptions = SNS_SUBSCRIPTIONS.get(topic_arn, [])
     for subscriber in list(subscriptions):
         if subscription_arn not in [None, subscriber['SubscriptionArn']]:
@@ -313,6 +314,7 @@ def message_to_subscribers(message_id, message, topic_arn, req_data, subscriptio
                     QueueUrl=queue_url,
                     MessageBody=create_sns_message_body(subscriber, req_data, message_id),
                     MessageAttributes=create_sqs_message_attributes(subscriber, message_attributes),
+                    MessageSystemAttributes=create_sqs_system_attributes(headers),
                     **kwargs
                 )
             except Exception as exc:
@@ -394,7 +396,7 @@ def message_to_subscribers(message_id, message, topic_arn, req_data, subscriptio
             LOG.warning('Unexpected protocol "%s" for SNS subscription' % subscriber['Protocol'])
 
 
-def publish_message(topic_arn, req_data, subscription_arn=None, skip_checks=False):
+def publish_message(topic_arn, req_data, headers, subscription_arn=None, skip_checks=False):
     message = req_data['Message'][0]
     message_id = str(uuid.uuid4())
 
@@ -405,7 +407,8 @@ def publish_message(topic_arn, req_data, subscription_arn=None, skip_checks=Fals
 
     LOG.debug('Publishing message to TopicArn: %s | Message: %s' % (topic_arn, message))
     start_thread(
-        lambda _: message_to_subscribers(message_id, message, topic_arn, req_data, subscription_arn, skip_checks))
+        lambda _: message_to_subscribers(
+            message_id, message, topic_arn, req_data, headers, subscription_arn, skip_checks))
     return message_id
 
 
@@ -461,7 +464,7 @@ def do_subscribe(topic_arn, endpoint, protocol, subscription_arn, attributes, fi
                         'To confirm the subscription, visit the SubscribeURL included in this message.'],
             'SubscribeURL': ['%s/?Action=ConfirmSubscription&TopicArn=%s&Token=%s' % (external_url, topic_arn, token)]
         }
-        publish_message(topic_arn, confirmation, subscription_arn, skip_checks=True)
+        publish_message(topic_arn, confirmation, {}, subscription_arn, skip_checks=True)
 
 
 def do_unsubscribe(subscription_arn):
@@ -613,6 +616,17 @@ def create_sqs_message_attributes(subscriber, attributes):
     return message_attributes
 
 
+def create_sqs_system_attributes(headers):
+
+    system_attributes = {}
+    if 'X-Amzn-Trace-Id' in headers:
+        system_attributes['AWSTraceHeader'] = {
+            'DataType': 'String',
+            'StringValue': str(headers['X-Amzn-Trace-Id'])
+        }
+    return system_attributes
+
+
 def get_message_attributes(req_data):
     attributes = {}
     x = 1


Test patch:
diff --git a/tests/integration/test_sns.py b/tests/integration/test_sns.py
--- a/tests/integration/test_sns.py
+++ b/tests/integration/test_sns.py
@@ -20,6 +20,7 @@
 from localstack.services.sns import sns_listener
 from .lambdas import lambda_integration
 from .test_lambda import TEST_LAMBDA_PYTHON, LAMBDA_RUNTIME_PYTHON36, TEST_LAMBDA_LIBS
+from localstack.services.install import SQS_BACKEND_IMPL
 
 TEST_TOPIC_NAME = 'TestTopic_snsTest'
 TEST_QUEUE_NAME = 'TestQueue_snsTest'
@@ -786,6 +787,40 @@ def test_publish_sqs_from_sns(self):
         )
         self.assertEqual(len(response['Messages']), 1)
 
+    def add_xray_header(self, request, **kwargs):
+        request.headers['X-Amzn-Trace-Id'] = \
+            'Root=1-3152b799-8954dae64eda91bc9a23a7e8;Parent=7fa8c0f79203be72;Sampled=1'
+
+    def test_publish_sqs_from_sns_with_xray_propagation(self):
+        if SQS_BACKEND_IMPL != 'elasticmq':
+            return
+
+        self.sns_client.meta.events.register('before-send.sns.Publish', self.add_xray_header)
+
+        topic = self.sns_client.create_topic(Name='test_topic4')
+        topic_arn = topic['TopicArn']
+        test_queue = self.sqs_client.create_queue(QueueName='test_queue4')
+
+        queue_url = test_queue['QueueUrl']
+        self.sns_client.subscribe(TopicArn=topic_arn, Protocol='sqs', Endpoint=queue_url)
+        self.sns_client.publish(TargetArn=topic_arn, Message='X-Ray propagation test msg')
+
+        response = self.sqs_client.receive_message(
+            QueueUrl=queue_url,
+            AttributeNames=['SentTimestamp', 'AWSTraceHeader'],
+            MaxNumberOfMessages=1,
+            MessageAttributeNames=['All'],
+            VisibilityTimeout=2,
+            WaitTimeSeconds=2,
+        )
+
+        self.assertEqual(len(response['Messages']), 1)
+        message = response['Messages'][0]
+        self.assertTrue('Attributes' in message)
+        self.assertTrue('AWSTraceHeader' in message['Attributes'])
+        self.assertEqual(message['Attributes']['AWSTraceHeader'],
+                         'Root=1-3152b799-8954dae64eda91bc9a23a7e8;Parent=7fa8c0f79203be72;Sampled=1')
+
     def test_create_topic_after_delete_with_new_tags(self):
         topic_name = 'test-%s' % short_uid()
         topic = self.sns_client.create_topic(Name=topic_name, Tags=[{'Key': 'Name', 'Value': 'pqr'}])