fix: SQS instrumentation could cause InvalidOperationException #2645

newrelic · Jul 25, 2024 · 6a59420 · 6a59420
1 parent 9efd346
commit 6a59420
Showing 1 changed file with 32 additions and 30 deletions.
diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/SQSRequestHandler.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/SQSRequestHandler.cs
@@ -4,6 +4,7 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Linq;
 using System.Threading.Tasks;
 using NewRelic.Agent.Api;
 using NewRelic.Agent.Extensions.AwsSdk;
@@ -48,49 +49,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 ?