Repository: softwaremill/elasticmq (scala)

Issue:
"MD5 returned by SQS does not match the calculation on the original request" when using AWS JSON protocol and message attributes with custom data types
**Description**

When receiving a SQS message from ElasticMQ with an attribute with a custom data type ([docs](https://docs.aws.amazon.com/AWSSimpleQueueService/latest/SQSDeveloperGuide/sqs-message-metadata.html#message-attribute-data-types)) using AWS JSON protocol (AWS Java SDK v1 >= 1.12.584 / v2 >= 2.21.19), the following exception is thrown:

```
com.amazonaws.AmazonClientException: MD5 returned by SQS does not match the calculation on the original request. (MD5 calculated by the message attributes: "b6ba7933bdd4077bb2cc8686b0146cfa", MD5 checksum returned: "6735a2d7fe7f0e962a51db579c2d92cb")
	at com.amazonaws.services.sqs.MessageMD5ChecksumHandler.receiveMessageResultMd5Check(MessageMD5ChecksumHandler.java:146)
	at com.amazonaws.services.sqs.MessageMD5ChecksumHandler.afterResponse(MessageMD5ChecksumHandler.java:87)
	at com.amazonaws.handlers.RequestHandler2Adaptor.afterResponse(RequestHandler2Adaptor.java:49)
	at com.amazonaws.http.AmazonHttpClient$RequestExecutor.afterResponse(AmazonHttpClient.java:1070)
...
        at com.amazonaws.services.sqs.AmazonSQSClient.receiveMessage(AmazonSQSClient.java:2339)
```

An example message with a custom data type when the error occurs:


```
var request = new SendMessageRequest(
                getQueueUrl(endpoint),
                "msg"
        );
request.setMessageAttributes(Map.of(
        "timestamp", new MessageAttributeValue()
                .withStringValue("1")
                .withDataType("Number.java.lang.Long")
```


Minimal reproducible example (Java): https://github.com/mgocd/elasticmq-custom-attribute-data-type-issue
ElasticMQ version: 1.5.4

**Analysis**

This happens because [MessageAttributeSupport](https://github.com/softwaremill/elasticmq/blob/v1.5.4/rest/rest-sqs/src/main/scala/org/elasticmq/rest/sqs/MessageAttributesSupport.scala#L15) ignores custom data types when serializing to JSON (so `Number.java.lang.Long` becomes `Number`), while [SQSRestServerBuilder#md5AttributeDigest](https://github.com/softwaremill/elasticmq/blob/v1.5.4/rest/rest-sqs/src/main/scala/org/elasticmq/rest/sqs/SQSRestServerBuilder.scala#L396) calculates MD5 digest correctly using the custom data type (`Number.java.lang.Long`).

**Proposed solution**

`MessageAttributesSupport` should use the custom data type when serializing to JSON.

A pull request has been proposed to fix the issue.

Pull request description:
Closes #947

Additional info:
* Added custom data type to existing tests in `AmazonJavaSdkTestSuite` - note the tests will succeed also without changes in `MessageAttributesSupport` , because `aws-java-sdk-sqs` version used is 1.12.472 (pre-1.12.584 when JSON protocol was introduced)
* Added additional tests to `AmazonJavaSdkV2TestSuite` basing on `AmazonJavaSdkTestSuite` (some of these tests will not succeed without changes in `MessageAttributesSupport`)
* The new tests also are able to reproduce the issue described in https://github.com/softwaremill/elasticmq/issues/946, occurring when a binary message attribute is used - need to uncomment `"green" -> BinaryMessageAttribute("dog".getBytes("UTF-8"))` lines in `AmazonJavaSdkV2TestSuite`. This PR does not fix this issue.

Patch:
diff --git a/rest/rest-sqs/src/main/scala/org/elasticmq/rest/sqs/MessageAttributesSupport.scala b/rest/rest-sqs/src/main/scala/org/elasticmq/rest/sqs/MessageAttributesSupport.scala
index 619a128d..2e2f8915 100644
--- a/rest/rest-sqs/src/main/scala/org/elasticmq/rest/sqs/MessageAttributesSupport.scala
+++ b/rest/rest-sqs/src/main/scala/org/elasticmq/rest/sqs/MessageAttributesSupport.scala
@@ -13,12 +13,12 @@ trait MessageAttributesSupport {
   implicit val messageAttributeJsonFormat: RootJsonFormat[MessageAttribute] = new RootJsonFormat[MessageAttribute] {
 
     override def write(obj: MessageAttribute): JsValue = obj match {
-      case NumberMessageAttribute(value, _) =>
-        JsObject("DataType" -> JsString("Number"), "StringValue" -> JsString(value))
-      case StringMessageAttribute(value, _) =>
-        JsObject("DataType" -> JsString("String"), "StringValue" -> JsString(value))
+      case NumberMessageAttribute(value, customType) =>
+        JsObject("DataType" -> JsString("Number" + customTypeAsString(customType)), "StringValue" -> JsString(value))
+      case StringMessageAttribute(value, customType) =>
+        JsObject("DataType" -> JsString("String" + customTypeAsString(customType)), "StringValue" -> JsString(value))
       case msg: BinaryMessageAttribute =>
-        JsObject("DataType" -> JsString("Binary"), "BinaryValue" -> JsString(msg.asBase64))
+        JsObject("DataType" -> JsString("Binary" + customTypeAsString(msg.customType)), "BinaryValue" -> JsString(msg.asBase64))
     }
 
     override def read(json: JsValue): MessageAttribute = {
@@ -38,6 +38,8 @@ trait MessageAttributesSupport {
     }
 
     private def customType(appendix: String) = if (appendix.isEmpty) None else Some(appendix)
+
+    private def customTypeAsString(customType: Option[String]) = customType.fold("")(t => s".$t")
   }
 
 }


Test patch:
diff --git a/rest/rest-sqs-testing-amazon-java-sdk/src/test/scala/org/elasticmq/rest/sqs/AmazonJavaSdkTestSuite.scala b/rest/rest-sqs-testing-amazon-java-sdk/src/test/scala/org/elasticmq/rest/sqs/AmazonJavaSdkTestSuite.scala
index 30399d9a..7ea04e27 100644
--- a/rest/rest-sqs-testing-amazon-java-sdk/src/test/scala/org/elasticmq/rest/sqs/AmazonJavaSdkTestSuite.scala
+++ b/rest/rest-sqs-testing-amazon-java-sdk/src/test/scala/org/elasticmq/rest/sqs/AmazonJavaSdkTestSuite.scala
@@ -166,7 +166,8 @@ class AmazonJavaSdkTestSuite extends SqsClientServerCommunication with Matchers
         "red" -> StringMessageAttribute("fish"),
         "blue" -> StringMessageAttribute("cat"),
         "green" -> BinaryMessageAttribute("dog".getBytes("UTF-8")),
-        "yellow" -> NumberMessageAttribute("1234567890")
+        "yellow" -> NumberMessageAttribute("1234567890"),
+        "orange" -> NumberMessageAttribute("0987654321", Some("custom"))
       )
     )
   }
@@ -178,9 +179,10 @@ class AmazonJavaSdkTestSuite extends SqsClientServerCommunication with Matchers
         "red" -> StringMessageAttribute("fish"),
         "blue" -> StringMessageAttribute("cat"),
         "green" -> BinaryMessageAttribute("dog".getBytes("UTF-8")),
-        "yellow" -> NumberMessageAttribute("1234567890")
+        "yellow" -> NumberMessageAttribute("1234567890"),
+        "orange" -> NumberMessageAttribute("0987654321", Some("custom"))
       ),
-      List("red", "green")
+      List("red", "green", "orange")
     )
   }
 
@@ -191,7 +193,8 @@ class AmazonJavaSdkTestSuite extends SqsClientServerCommunication with Matchers
         "red" -> StringMessageAttribute("fish"),
         "blue" -> StringMessageAttribute("cat"),
         "green" -> BinaryMessageAttribute("dog".getBytes("UTF-8")),
-        "yellow" -> NumberMessageAttribute("1234567890")
+        "yellow" -> NumberMessageAttribute("1234567890"),
+        "orange" -> NumberMessageAttribute("0987654321", Some("custom"))
       ),
       List()
     )
diff --git a/rest/rest-sqs-testing-amazon-java-sdk/src/test/scala/org/elasticmq/rest/sqs/AmazonJavaSdkV2TestSuite.scala b/rest/rest-sqs-testing-amazon-java-sdk/src/test/scala/org/elasticmq/rest/sqs/AmazonJavaSdkV2TestSuite.scala
index 36c5a835..f78b83f1 100644
--- a/rest/rest-sqs-testing-amazon-java-sdk/src/test/scala/org/elasticmq/rest/sqs/AmazonJavaSdkV2TestSuite.scala
+++ b/rest/rest-sqs-testing-amazon-java-sdk/src/test/scala/org/elasticmq/rest/sqs/AmazonJavaSdkV2TestSuite.scala
@@ -1,6 +1,10 @@
 package org.elasticmq.rest.sqs
 
+import org.elasticmq.{BinaryMessageAttribute, MessageAttribute, NumberMessageAttribute, StringMessageAttribute}
 import org.scalatest.matchers.should.Matchers
+
+import scala.collection.JavaConverters._
+import software.amazon.awssdk.core.SdkBytes
 import software.amazon.awssdk.services.sqs.model._
 import software.amazon.awssdk.services.sqs.model.{GetQueueUrlRequest => AwsSdkGetQueueUrlRequest}
 
@@ -32,16 +36,159 @@ class AmazonJavaSdkV2TestSuite extends SqsClientServerWithSdkV2Communication wit
     thrown.awsErrorDetails().errorMessage() shouldBe "The specified queue does not exist."
   }
 
-  test("should send and receive message") {
+  test("should send and receive a simple message") {
+    doTestSendAndReceiveMessage("test msg 123")
+  }
+
+  test("should send and receive a simple message with message attributes") {
+    doTestSendAndReceiveMessageWithAttributes(
+      "Message 1",
+      Map(
+        "red" -> StringMessageAttribute("fish"),
+        "blue" -> StringMessageAttribute("cat"),
+        // affected by https://github.com/softwaremill/elasticmq/issues/946
+        // "green" -> BinaryMessageAttribute("dog".getBytes("UTF-8")),
+        "yellow" -> NumberMessageAttribute("1234567890"),
+        "orange" -> NumberMessageAttribute("0987654321", Some("custom"))
+      )
+    )
+  }
+
+  test("should send a simple message with message attributes and only receive requested attributes") {
+    doTestSendAndReceiveMessageWithAttributes(
+      "Message 1",
+      Map(
+        "red" -> StringMessageAttribute("fish"),
+        "blue" -> StringMessageAttribute("cat"),
+        // affected by https://github.com/softwaremill/elasticmq/issues/946
+        // "green" -> BinaryMessageAttribute("dog".getBytes("UTF-8")),
+        "yellow" -> NumberMessageAttribute("1234567890"),
+        "orange" -> NumberMessageAttribute("0987654321", Some("custom"))
+      ),
+      List("red", "green", "orange")
+    )
+  }
+
+  test("should send a simple message with message attributes and only receive no requested attributes by default") {
+    doTestSendAndReceiveMessageWithAttributes(
+      "Message 1",
+      Map(
+        "red" -> StringMessageAttribute("fish"),
+        "blue" -> StringMessageAttribute("cat"),
+        // affected by https://github.com/softwaremill/elasticmq/issues/946
+        // "green" -> BinaryMessageAttribute("dog".getBytes("UTF-8")),
+        "yellow" -> NumberMessageAttribute("1234567890"),
+        "orange" -> NumberMessageAttribute("0987654321", Some("custom"))
+      ),
+      List()
+    )
+  }
+
+  private def doTestSendAndReceiveMessage(content: String): Unit = {
+    doTestSendAndReceiveMessageWithAttributes(content, Map(), List())
+  }
+
+  private def doTestSendAndReceiveMessageWithAttributes(
+                                                         content: String,
+                                                         messageAttributes: Map[String, MessageAttribute]
+                                                       ): Unit = {
+    doTestSendAndReceiveMessageWithAttributes(content, messageAttributes, List("All"))
+  }
+
+  private def doTestSendAndReceiveMessageWithAttributes(
+                                                         content: String,
+                                                         messageAttributes: Map[String, MessageAttribute],
+                                                         requestedAttributes: List[String]
+                                                       ): Unit = {
+    // Given
     val queue = clientV2.createQueue(CreateQueueRequest.builder().queueName("testQueue1").build())
 
-    clientV2.sendMessage(SendMessageRequest.builder().queueUrl(queue.queueUrl()).messageBody("test msg 123").build())
+    // When
+    val attributes = messageAttributes.map {
+      case (k, v: StringMessageAttribute) =>
+        k -> MessageAttributeValue.builder().dataType(v.getDataType()).stringValue(v.stringValue).build()
+      case (k, v: NumberMessageAttribute) =>
+        k -> MessageAttributeValue.builder().dataType(v.getDataType()).stringValue(v.stringValue).build()
+      case (k, v: BinaryMessageAttribute) =>
+        k -> MessageAttributeValue
+          .builder()
+          .dataType(v.getDataType())
+          .binaryValue(SdkBytes.fromByteArray(v.binaryValue))
+          .build()
+    }
+
+    val sendMessageRequest = SendMessageRequest
+      .builder()
+      .queueUrl(queue.queueUrl())
+      .messageBody(content)
+      .messageAttributes(attributes.asJava)
+      .build()
 
-    val messages = clientV2.receiveMessage(ReceiveMessageRequest.builder().queueUrl(queue.queueUrl()).build())
+    clientV2.sendMessage(sendMessageRequest)
 
-    System.err.println(messages)
+    val message = receiveSingleMessageObject(queue.queueUrl(), requestedAttributes).orNull
+
+    // Then
+    message.body() shouldBe content
+    checkMessageAttributesMatchRequestedAttributes(messageAttributes, requestedAttributes, sendMessageRequest, message)
+  }
 
-    messages.messages().size() shouldBe 1
-    messages.messages().get(0).body() shouldBe "test msg 123"
+  private def receiveSingleMessageObject(queueUrl: String, requestedAttributes: List[String]): Option[Message] = {
+    clientV2
+      .receiveMessage(
+        ReceiveMessageRequest
+          .builder()
+          .queueUrl(queueUrl)
+          .messageAttributeNames(requestedAttributes.asJava)
+          .build()
+      )
+      .messages
+      .asScala
+      .headOption
+  }
+
+  private def checkMessageAttributesMatchRequestedAttributes(
+                                                              messageAttributes: Map[String, MessageAttribute],
+                                                              requestedAttributes: List[String],
+                                                              sendMessageRequest: SendMessageRequest,
+                                                              message: Message
+                                                            ) = {
+    val filteredSendMessageAttr =
+      filterBasedOnRequestedAttributes(requestedAttributes, sendMessageRequest.messageAttributes.asScala.toMap).asJava
+    val filteredMessageAttributes = filterBasedOnRequestedAttributes(requestedAttributes, messageAttributes)
+
+    message.messageAttributes should be(filteredSendMessageAttr)
+    message.messageAttributes.asScala.map { case (k, attr) =>
+      (
+        k,
+        if (attr.dataType.startsWith("String") && attr.stringValue != null) {
+          StringMessageAttribute(attr.stringValue).stringValue
+        } else if (attr.dataType.startsWith("Number") && attr.stringValue != null) {
+          NumberMessageAttribute(attr.stringValue).stringValue
+        } else {
+          BinaryMessageAttribute.fromByteBuffer(attr.binaryValue.asByteBuffer()).asBase64
+        }
+      )
+    } should be(filteredMessageAttributes.map { case (k, attr) =>
+      (
+        k,
+        attr match {
+          case s: StringMessageAttribute => s.stringValue
+          case n: NumberMessageAttribute => n.stringValue
+          case b: BinaryMessageAttribute => b.asBase64
+        }
+      )
+    })
+  }
+
+  private def filterBasedOnRequestedAttributes[T](
+                                                   requestedAttributes: List[String],
+                                                   messageAttributes: Map[String, T]
+                                                 ): Map[String, T] = {
+    if (requestedAttributes.contains("All")) {
+      messageAttributes
+    } else {
+      messageAttributes.filterKeys(k => requestedAttributes.contains(k)).toMap
+    }
   }
 }