Support custom SQS queue attributes (localstack#1520)

whummer · web-flow · commit 5eab8f64c1b8 · 2019-09-01T15:42:52.000+02:00
diff --git a/localstack/services/sqs/sqs_listener.py b/localstack/services/sqs/sqs_listener.py
@@ -15,7 +15,6 @@
 
 XMLNS_SQS = 'http://queue.amazonaws.com/doc/2012-11-05/'
 
-
 SUCCESSFUL_SEND_MESSAGE_XML_TEMPLATE = """
     <?xml version="1.0"?>
     <SendMessageResponse xmlns="%s">
@@ -30,32 +29,29 @@
     </SendMessageResponse>
 """.strip() % XMLNS_SQS
 
+# list of valid attribute names, and names not supported by the backend (elasticmq)
+VALID_ATTRIBUTE_NAMES = ['DelaySeconds', 'MaximumMessageSize', 'MessageRetentionPeriod',
+    'Policy', 'ReceiveMessageWaitTimeSeconds', 'RedrivePolicy', 'VisibilityTimeout']
+UNSUPPORTED_ATTRIBUTE_NAMES = ['MaximumMessageSize', 'MessageRetentionPeriod', 'Policy', 'RedrivePolicy']
+
+# maps queue URLs to attributes set via the API
+QUEUE_ATTRIBUTES = {}
+
 
 class ProxyListenerSQS(ProxyListener):
 
     def forward_request(self, method, path, data, headers):
         req_data = self.parse_request_data(method, path, data)
 
         if req_data:
-            if req_data.get('Action', [None])[0] == 'SendMessage':
-                queue_url = req_data.get('QueueUrl', [path.partition('?')[0]])[0]
-                queue_name = queue_url[queue_url.rindex('/') + 1:]
-                message_body = req_data.get('MessageBody', [None])[0]
-                message_attributes = self.format_message_attributes(req_data)
-                region_name = extract_region_from_auth_header(headers)
-
-                process_result = lambda_api.process_sqs_message(message_body,
-                    message_attributes, queue_name, region_name=region_name)
-                if process_result:
-                    # If a Lambda was listening, do not add the message to the queue
-                    new_response = Response()
-                    new_response._content = SUCCESSFUL_SEND_MESSAGE_XML_TEMPLATE.format(
-                        message_attr_hash=md5(data),
-                        message_body_hash=md5(message_body),
-                        message_id=str(uuid.uuid4())
-                    )
-                    new_response.status_code = 200
+            action = req_data.get('Action', [None])[0]
+            if action == 'SendMessage':
+                new_response = self._send_message(path, data, req_data, headers)
+                if new_response:
                     return new_response
+            elif action == 'SetQueueAttributes':
+                self._set_queue_attributes(req_data)
+
             if 'QueueName' in req_data:
                 encoded_data = urlencode(req_data, doseq=True) if method == 'POST' else ''
                 modified_url = None
@@ -76,6 +72,73 @@ def parse_request_data(self, method, path, data):
             return urlparse.parse_qs(parsed_path.query)
         return {}
 
+    def return_response(self, method, path, data, headers, response, request_handler):
+        if method == 'OPTIONS' and path == '/':
+            # Allow CORS preflight requests to succeed.
+            return 200
+
+        if method == 'POST' and path == '/':
+            region_name = extract_region_from_auth_header(headers)
+            req_data = urlparse.parse_qs(to_str(data))
+            action = req_data.get('Action', [None])[0]
+            content_str = content_str_original = to_str(response.content)
+
+            self._fire_event(req_data, response)
+
+            # patch the response and add missing attributes
+            if action == 'GetQueueAttributes':
+                content_str = self._add_queue_attributes(req_data, content_str)
+
+            # patch the response and return the correct endpoint URLs / ARNs
+            if action in ('CreateQueue', 'GetQueueUrl', 'ListQueues', 'GetQueueAttributes'):
+                if config.USE_SSL and '<QueueUrl>http://' in content_str:
+                    # return https://... if we're supposed to use SSL
+                    content_str = re.sub(r'<QueueUrl>\s*http://', r'<QueueUrl>https://', content_str)
+                # expose external hostname:port
+                external_port = SQS_PORT_EXTERNAL or get_external_port(headers, request_handler)
+                content_str = re.sub(r'<QueueUrl>\s*([a-z]+)://[^<]*:([0-9]+)/([^<]*)\s*</QueueUrl>',
+                    r'<QueueUrl>\1://%s:%s/\3</QueueUrl>' % (HOSTNAME_EXTERNAL, external_port), content_str)
+                # fix queue ARN
+                content_str = re.sub(r'<([a-zA-Z0-9]+)>\s*arn:aws:sqs:elasticmq:([^<]+)</([a-zA-Z0-9]+)>',
+                    r'<\1>arn:aws:sqs:%s:\2</\3>' % (region_name), content_str)
+
+            if content_str_original != content_str:
+                # if changes have been made, return patched response
+                new_response = Response()
+                new_response.status_code = response.status_code
+                new_response.headers = response.headers
+                new_response._content = content_str
+                new_response.headers['content-length'] = len(new_response._content)
+                return new_response
+
+            # Since the following 2 API calls are not implemented in ElasticMQ, we're mocking them
+            # and letting them to return an empty response
+            if action == 'TagQueue':
+                new_response = Response()
+                new_response.status_code = 200
+                new_response._content = ("""
+                    <?xml version="1.0"?>
+                    <TagQueueResponse>
+                        <ResponseMetadata>
+                            <RequestId>{}</RequestId>
+                        </ResponseMetadata>
+                    </TagQueueResponse>
+                """).strip().format(uuid.uuid4())
+                return new_response
+            elif action == 'ListQueueTags':
+                new_response = Response()
+                new_response.status_code = 200
+                new_response._content = ("""
+                    <?xml version="1.0"?>
+                    <ListQueueTagsResponse xmlns="{}">
+                        <ListQueueTagsResult/>
+                        <ResponseMetadata>
+                            <RequestId>{}</RequestId>
+                        </ResponseMetadata>
+                    </ListQueueTagsResponse>
+                """).strip().format(XMLNS_SQS, uuid.uuid4())
+                return new_response
+
     # Format of the message Name attribute is MessageAttribute.<int id>.<field>
     # Format of the Value attributes is MessageAttribute.<int id>.Value.DataType
     # and MessageAttribute.<int id>.Value.<Type>Value
@@ -115,10 +178,10 @@ def parse_request_data(self, method, path, data):
     #      dataType: 'String'
     #    }
     # }
-
     def format_message_attributes(self, data):
+        prefix = 'MessageAttribute'
         names = []
-        for (k, name) in [(k, data[k]) for k in data if k.startswith('MessageAttribute') and k.endswith('.Name')]:
+        for (k, name) in [(k, data[k]) for k in data if k.startswith(prefix) and k.endswith('.Name')]:
             attr_name = name[0]
             k_id = k.split('.')[1]
             names.append((attr_name, k_id))
@@ -128,7 +191,7 @@ def format_message_attributes(self, data):
             msg_attrs[key_name] = {}
             # Find vals for each key_id
             attrs = [(k, data[k]) for k in data
-                if k.startswith('MessageAttribute.{}.'.format(key_id)) and not k.endswith('.Name')]
+                if k.startswith('{}.{}.'.format(prefix, key_id)) and not k.endswith('.Name')]
             for (attr_k, attr_v) in attrs:
                 attr_name = attr_k.split('.')[3]
                 msg_attrs[key_name][attr_name[0].lower() + attr_name[1:]] = attr_v[0]
@@ -141,78 +204,78 @@ def format_message_attributes(self, data):
 
         return msg_attrs
 
-    def return_response(self, method, path, data, headers, response, request_handler):
-        if method == 'OPTIONS' and path == '/':
-            # Allow CORS preflight requests to succeed.
-            return 200
+    # Format attributes as dict. Example input:
+    #  {
+    #    'Attribute.1.Name': ['Policy'],
+    #    'Attribute.1.Value': ['...']
+    #  }
+    def _format_attributes(self, req_data):
+        result = {}
+        for i in range(1, 500):
+            key1 = 'Attribute.%s.Name' % i
+            key2 = 'Attribute.%s.Value' % i
+            if key1 not in req_data:
+                break
+            key_name = req_data[key1][0]
+            key_value = req_data[key2][0]
+            result[key_name] = key_value
+        return result
 
-        if method == 'POST' and path == '/':
-            region_name = extract_region_from_auth_header(headers)
-            req_data = urlparse.parse_qs(to_str(data))
-            action = req_data.get('Action', [None])[0]
-            event_type = None
-            queue_url = None
-            if action == 'CreateQueue':
-                event_type = event_publisher.EVENT_SQS_CREATE_QUEUE
-                response_data = xmltodict.parse(response.content)
-                if 'CreateQueueResponse' in response_data:
-                    queue_url = response_data['CreateQueueResponse']['CreateQueueResult']['QueueUrl']
-            elif action == 'DeleteQueue':
-                event_type = event_publisher.EVENT_SQS_DELETE_QUEUE
-                queue_url = req_data.get('QueueUrl', [None])[0]
-
-            if event_type and queue_url:
-                event_publisher.fire_event(event_type, payload={'u': event_publisher.get_hash(queue_url)})
+    def _send_message(self, path, data, req_data, headers):
+        queue_url = req_data.get('QueueUrl', [path.partition('?')[0]])[0]
+        queue_name = queue_url[queue_url.rindex('/') + 1:]
+        message_body = req_data.get('MessageBody', [None])[0]
+        message_attributes = self.format_message_attributes(req_data)
+        region_name = extract_region_from_auth_header(headers)
 
-            # patch the response and return the correct endpoint URLs / ARNs
-            if action in ('CreateQueue', 'GetQueueUrl', 'ListQueues', 'GetQueueAttributes'):
-                content_str = content_str_original = to_str(response.content)
-                new_response = Response()
-                new_response.status_code = response.status_code
-                new_response.headers = response.headers
-                if config.USE_SSL and '<QueueUrl>http://' in content_str:
-                    # return https://... if we're supposed to use SSL
-                    content_str = re.sub(r'<QueueUrl>\s*http://', r'<QueueUrl>https://', content_str)
-                # expose external hostname:port
-                external_port = SQS_PORT_EXTERNAL or get_external_port(headers, request_handler)
-                content_str = re.sub(r'<QueueUrl>\s*([a-z]+)://[^<]*:([0-9]+)/([^<]*)\s*</QueueUrl>',
-                    r'<QueueUrl>\1://%s:%s/\3</QueueUrl>' % (HOSTNAME_EXTERNAL, external_port), content_str)
-                # fix queue ARN
-                content_str = re.sub(r'<([a-zA-Z0-9]+)>\s*arn:aws:sqs:elasticmq:([^<]+)</([a-zA-Z0-9]+)>',
-                    r'<\1>arn:aws:sqs:%s:\2</\3>' % (region_name), content_str)
-                new_response._content = content_str
-                if content_str_original != new_response._content:
-                    # if changes have been made, return patched response
-                    new_response.headers['content-length'] = len(new_response._content)
-                    return new_response
+        process_result = lambda_api.process_sqs_message(message_body,
+            message_attributes, queue_name, region_name=region_name)
+        if process_result:
+            # If a Lambda was listening, do not add the message to the queue
+            new_response = Response()
+            new_response._content = SUCCESSFUL_SEND_MESSAGE_XML_TEMPLATE.format(
+                message_attr_hash=md5(data),
+                message_body_hash=md5(message_body),
+                message_id=str(uuid.uuid4())
+            )
+            new_response.status_code = 200
+            return new_response
 
-            # Since the following 2 API calls are not implemented in ElasticMQ, we're mocking them
-            # and letting them to return an empty response
-            if action == 'TagQueue':
-                new_response = Response()
-                new_response.status_code = 200
-                new_response._content = ("""
-                    <?xml version="1.0"?>
-                    <TagQueueResponse>
-                        <ResponseMetadata>
-                            <RequestId>{}</RequestId>
-                        </ResponseMetadata>
-                    </TagQueueResponse>
-                """).strip().format(uuid.uuid4())
-                return new_response
-            elif action == 'ListQueueTags':
-                new_response = Response()
-                new_response.status_code = 200
-                new_response._content = ("""
-                    <?xml version="1.0"?>
-                    <ListQueueTagsResponse xmlns="{}">
-                        <ListQueueTagsResult/>
-                        <ResponseMetadata>
-                            <RequestId>{}</RequestId>
-                        </ResponseMetadata>
-                    </ListQueueTagsResponse>
-                """).strip().format(XMLNS_SQS, uuid.uuid4())
-                return new_response
+    def _set_queue_attributes(self, req_data):
+        queue_url = req_data['QueueUrl'][0]
+        attrs = self._format_attributes(req_data)
+        # select only the attributes in UNSUPPORTED_ATTRIBUTE_NAMES
+        attrs = dict([(k, v) for k, v in attrs.items() if k in UNSUPPORTED_ATTRIBUTE_NAMES])
+        QUEUE_ATTRIBUTES[queue_url] = QUEUE_ATTRIBUTES.get(queue_url) or {}
+        QUEUE_ATTRIBUTES[queue_url].update(attrs)
+
+    def _add_queue_attributes(self, req_data, content_str):
+        flags = re.MULTILINE | re.DOTALL
+        queue_url = req_data['QueueUrl'][0]
+        regex = r'(.*<GetQueueAttributesResult>)(.*)(</GetQueueAttributesResult>.*)'
+        attrs = re.sub(regex, r'\2', content_str, flags=flags)
+        for key, value in QUEUE_ATTRIBUTES.get(queue_url, {}).items():
+            if not re.match(r'<Name>\s*%s\s*</Name>' % key, attrs, flags=flags):
+                attrs += '<Attribute><Name>%s</Name><Value>%s</Value></Attribute>' % (key, value)
+        content_str = (re.sub(regex, r'\1', content_str, flags=flags) +
+            attrs + re.sub(regex, r'\3', content_str, flags=flags))
+        return content_str
+
+    def _fire_event(self, req_data, response):
+        action = req_data.get('Action', [None])[0]
+        event_type = None
+        queue_url = None
+        if action == 'CreateQueue':
+            event_type = event_publisher.EVENT_SQS_CREATE_QUEUE
+            response_data = xmltodict.parse(response.content)
+            if 'CreateQueueResponse' in response_data:
+                queue_url = response_data['CreateQueueResponse']['CreateQueueResult']['QueueUrl']
+        elif action == 'DeleteQueue':
+            event_type = event_publisher.EVENT_SQS_DELETE_QUEUE
+            queue_url = req_data.get('QueueUrl', [None])[0]
+
+        if event_type and queue_url:
+            event_publisher.fire_event(event_type, payload={'u': event_publisher.get_hash(queue_url)})
 
 
 # extract the external port used by the client to make the request
diff --git a/requirements.txt b/requirements.txt
@@ -28,7 +28,7 @@ jsonpath-rw==1.4.0
 localstack-ext[full]>=0.10.2
 localstack-ext>=0.10.2         #basic-lib
 localstack-client>=0.10        #basic-lib
-moto-ext==1.3.14.dev0
+moto-ext>=1.3.14.1
 nose>=1.3.7
 nose-timer>=0.7.5
 psutil==5.4.8
diff --git a/tests/integration/test_sqs.py b/tests/integration/test_sqs.py
@@ -1,27 +1,64 @@
 import unittest
 from localstack.utils.aws import aws_stack
+from localstack.utils.common import short_uid
 
 
 TEST_QUEUE_NAME = 'TestQueue'
 
+TEST_POLICY = """
+{
+  "Version":"2012-10-17",
+  "Statement":[
+    {
+      "Effect": "Allow",
+      "Principal": { "AWS": "*" },
+      "Action": "sqs:SendMessage",
+      "Resource": "'$sqs_queue_arn'",
+      "Condition":{
+        "ArnEquals":{
+        "aws:SourceArn":"'$sns_topic_arn'"
+        }
+      }
+    }
+  ]
+}
+ """
+
 
 class SQSTest(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        cls.client = aws_stack.connect_to_service('sqs')
+
     def test_list_queue_tags(self):
         # Since this API call is not implemented in ElasticMQ, we're mocking it
         # and letting it return an empty response
-        sqs_client = aws_stack.connect_to_service('sqs')
-        queue_info = sqs_client.create_queue(QueueName=TEST_QUEUE_NAME)
+        queue_info = self.client.create_queue(QueueName=TEST_QUEUE_NAME)
         queue_url = queue_info['QueueUrl']
-        res = sqs_client.list_queue_tags(QueueUrl=queue_url)
+        res = self.client.list_queue_tags(QueueUrl=queue_url)
 
         # Apparently, if there are no tags, then `Tags` should NOT appear in the response.
         assert 'Tags' not in res
 
     def test_create_fifo_queue(self):
         fifo_queue = 'my-queue.fifo'
-        sqs_client = aws_stack.connect_to_service('sqs')
-        queue_info = sqs_client.create_queue(QueueName=fifo_queue, Attributes={'FifoQueue': 'true'})
+        queue_info = self.client.create_queue(QueueName=fifo_queue, Attributes={'FifoQueue': 'true'})
         queue_url = queue_info['QueueUrl']
 
         # it should preserve .fifo in the queue name
         self.assertIn(fifo_queue, queue_url)
+
+    def test_set_queue_policy(self):
+        fifo_queue = 'queue-%s' % short_uid()
+        queue_info = self.client.create_queue(QueueName=fifo_queue)
+        queue_url = queue_info['QueueUrl']
+
+        attributes = {
+            'Policy': TEST_POLICY
+        }
+        self.client.set_queue_attributes(QueueUrl=queue_url, Attributes=attributes)
+
+        attrs = self.client.get_queue_attributes(QueueUrl=queue_url, AttributeNames=['All'])['Attributes']
+        self.assertIn('sqs:SendMessage', attrs['Policy'])
+        attrs = self.client.get_queue_attributes(QueueUrl=queue_url, AttributeNames=['Policy'])['Attributes']
+        self.assertIn('sqs:SendMessage', attrs['Policy'])
