fix: SQS instrumentation could cause InvalidOperationException (#2645) (#2646)

Author: tippmar-nr

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/SQSRequestHandler.cs

@@ -48,49 +48,50 @@ public static AfterWrappedMethodDelegate HandleSQSRequest(InstrumentedMethodCall
             switch (action)
             {
                 case MessageBrokerAction.Produce when requestType == "SendMessageRequest":
-                {
-                    if (request.MessageAttributes == null)
                     {
-                        agent.Logger.Debug("AwsSdkPipelineWrapper: requestContext.OriginalRequest.MessageAttributes is null, unable to insert distributed trace headers.");
-                    }
-                    else
-                    {
-                        SqsHelper.InsertDistributedTraceHeaders(transaction, request, dtHeaders.Count);
-                    }
+                        if (request.MessageAttributes == null)
+                        {
+                            agent.Logger.Debug("AwsSdkPipelineWrapper: requestContext.OriginalRequest.MessageAttributes is null, unable to insert distributed trace headers.");
+                        }
+                        else
+                        {
+                            SqsHelper.InsertDistributedTraceHeaders(transaction, request, dtHeaders.Count);
+                        }
 
-                    break;
-                }
+                        break;
+                    }
                 case MessageBrokerAction.Produce:
-                {
-                    if (requestType == "SendMessageBatchRequest")
                     {
-                        // loop through each message in the batch and insert distributed trace headers
-                        foreach (var message in request.Entries)
+                        if (requestType == "SendMessageBatchRequest")
                         {
-                            if (message.MessageAttributes == null)
+                            // loop through each message in the batch and insert distributed trace headers
+                            foreach (var message in request.Entries)
                             {
-                                agent.Logger.Debug("AwsSdkPipelineWrapper: requestContext.OriginalRequest.Entries.MessageAttributes is null, unable to insert distributed trace headers.");
-                            }
-                            else
-                            {
-                                SqsHelper.InsertDistributedTraceHeaders(transaction, message, dtHeaders.Count);
+                                if (message.MessageAttributes == null)
+                                {
+                                    agent.Logger.Debug("AwsSdkPipelineWrapper: requestContext.OriginalRequest.Entries.MessageAttributes is null, unable to insert distributed trace headers.");
+                                }
+                                else
+                                {
+                                    SqsHelper.InsertDistributedTraceHeaders(transaction, message, dtHeaders.Count);
+                                }
                             }
                         }
-                    }
 
-                    break;
-                }
+                        break;
+                    }
 
                 // modify the request to ask for DT headers in the response message attributes.
                 case MessageBrokerAction.Consume:
-                {
-                    if (request.MessageAttributeNames == null)
-                        request.MessageAttributeNames = new List<string>();
+                    {
+                        // create a new list or clone the existing one so we don't modify the original list
+                        request.MessageAttributeNames = request.MessageAttributeNames == null ? new List<string>() : new List<string>(request.MessageAttributeNames);
 
-                    foreach (var header in dtHeaders)
-                        request.MessageAttributeNames.Add(header);
-                    break;
-                }
+                        foreach (var header in dtHeaders)
+                            request.MessageAttributeNames.Add(header);
+
+                        break;
+                    }
             }
 
             return isAsync ?

tests/Agent/IntegrationTests/ContainerApplications/AwsSdkTestApp/AwsSdkExerciser/AwsSdkExerciser.cs

@@ -104,13 +104,19 @@ public async Task<IEnumerable<Message>> SQS_ReceiveMessageAsync(int maxMessagesT
                 throw new InvalidOperationException("Queue URL is not set. Call SQS_Initialize or SQS_SetQueueUrl first.");
             }
 
+            List<string> messageAttributeNames = ["All"];
+
             var response = await _amazonSqsClient.ReceiveMessageAsync(new ReceiveMessageRequest
             {
                 QueueUrl = _sqsQueueUrl,
                 MaxNumberOfMessages = maxMessagesToReceive,
-                MessageAttributeNames = ["All"]
+                MessageAttributeNames = messageAttributeNames
             });
 
+
+            if (messageAttributeNames.Count != 1)
+                throw new Exception("Expected messageAttributeNames to have a single element");
+
             foreach (var message in response.Messages)
             {
                 Console.WriteLine($"Message: {message.Body}");

tests/Agent/IntegrationTests/ContainerIntegrationTests/Fixtures/AwsSdkContainerTestFixtures.cs

@@ -41,6 +41,8 @@ public void ExerciseSQS_SendReceivePurge(string queueName)
     {
         var address = $"http://localhost:{Port}/awssdk";
 
+        // The exerciser will return a 500 error if the `RequestMessage.MessageAttributeNames` collection is modified by our instrumentation.
+        // See https://github.com/newrelic/newrelic-dotnet-agent/pull/2646 
         GetAndAssertStatusCode($"{address}/SQS_SendReceivePurge?queueName={queueName}", System.Net.HttpStatusCode.OK);
     }