From 9a0301c5c4a25ca027cbeee0bf2e83c97a3c9afe Mon Sep 17 00:00:00 2001
From: Marty Tippin <120425148+tippmar-nr@users.noreply.github.com>
Date: Tue, 25 Jun 2024 10:28:57 -0500
Subject: [PATCH] Cleanup, minor refactor

---
 .../AwsSdk/SqsHelper.cs                       | 30 ++++++++++++++-----
 .../Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs   | 12 +++++---
 2 files changed, 30 insertions(+), 12 deletions(-)
 rename src/Agent/NewRelic/Agent/Extensions/{Providers/Wrapper => NewRelic.Agent.Extensions}/AwsSdk/SqsHelper.cs (77%)

diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/SqsHelper.cs b/src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/AwsSdk/SqsHelper.cs
similarity index 77%
rename from src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/SqsHelper.cs
rename to src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/AwsSdk/SqsHelper.cs
index 28a2d7443..9b146ec7a 100644
--- a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/SqsHelper.cs
+++ b/src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/AwsSdk/SqsHelper.cs
@@ -3,12 +3,10 @@
 
 using System;
 using System.Collections.Generic;
-using System.Linq;
-using System.Net.NetworkInformation;
 using NewRelic.Agent.Api;
 using NewRelic.Agent.Extensions.Providers.Wrapper;
 
-namespace NewRelic.Providers.Wrapper.AwsSdk
+namespace NewRelic.Agent.Extensions.AwsSdk
 {
     public static class SqsHelper
     {
@@ -16,19 +14,35 @@ public static class SqsHelper
 
         private class SqsAttributes
         {
-            public string QueueName;
-            public string CloudId;
-            public string Region;
+            public string QueueName { get; }
+            public string CloudId { get; }
+            public string Region { get; }
 
             // https://sqs.us-east-2.amazonaws.com/123456789012/MyQueue
             public SqsAttributes(string url)
             {
+                if (string.IsNullOrEmpty(url))
+                {
+                    return;
+                }
+
                 var parts = url.Split('/');
+                if (parts.Length < 5)
+                {
+                    return;
+                }
+
+                CloudId = parts[3];
+                QueueName = parts[4];
+
                 var subdomain = parts[2].Split('.');
+                if (subdomain.Length < 2)
+                {
+                    return;
+                }
+
                 // subdomain[0] should always be "sqs"
                 Region = subdomain[1];
-                CloudId = parts[3];
-                QueueName = parts[4];
             }
         }
         public static ISegment GenerateSegment(ITransaction transaction, MethodCall methodCall, string url, MessageBrokerAction action)
diff --git a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs
index 02d9c920d..86e42c68e 100644
--- a/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs
+++ b/src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/AwsSdk/AwsSdkPipelineWrapper.cs
@@ -1,9 +1,8 @@
 // Copyright 2020 New Relic, Inc. All rights reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-using System;
-using System.Linq;
 using NewRelic.Agent.Api;
+using NewRelic.Agent.Extensions.AwsSdk;
 using NewRelic.Agent.Extensions.Providers.Wrapper;
 
 namespace NewRelic.Providers.Wrapper.AwsSdk
@@ -81,8 +80,13 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
             if (insertDistributedTraceHeaders)
             {
                 // This needs to happen at the end
-                //dynamic webRequest = requestContext.Request;
-                //SqsHelper.InsertDistributedTraceHeaders(transaction, webRequest);
+                if (requestContext.Request == null)
+                    agent.Logger.Debug("AwsSdkPipelineWrapper: requestContext.Request is null, unable to insert distributed trace headers.");
+                else
+                {
+                    dynamic webRequest = requestContext.Request;
+                    SqsHelper.InsertDistributedTraceHeaders(transaction, webRequest);
+                }
             }
 
             return Delegates.GetDelegateFor(segment);