Enable tracing S3 id pairs and fix SNS naming issue. (#168)

* Enable tracing S3 id pairs and fix SNS naming issue.

* Changed AWSSDK.Core version

* Changed parameter names and updated comments.

* Set parameters as default

sdk/src/Core/AWSXRayRecorder.Core.csproj

@@ -38,7 +38,7 @@
   </ItemGroup>
 
   <ItemGroup>
-    <PackageReference Include="AWSSDK.Core" Version="[3.3.25.1,4.0)" />
+    <PackageReference Include="AWSSDK.Core" Version="3.3.25.1" />
   </ItemGroup>
 
   <ItemGroup Condition="'$(TargetFramework)' == 'netstandard2.0'">

sdk/src/Handlers/AwsSdk/Internal/AWSXRaySDKUtils.cs

@@ -26,6 +26,13 @@ internal class AWSXRaySDKUtils
     {
         private static readonly String XRayServiceName = "XRay";
         private static readonly ISet<String> WhitelistedOperations = new HashSet<String> { "GetSamplingRules", "GetSamplingTargets" };
+
+        // Collection to uniform service names across X-Ray SDKs.
+        private static readonly Dictionary<string, string> FormattedServiceNames = new Dictionary<string, string>()
+        {
+            { "SimpleNotificationService" , "SNS" },
+        };
+
         internal static bool IsBlacklistedOperation(String serviceName, string operation)
         {
             if (string.Equals(serviceName, XRayServiceName) && WhitelistedOperations.Contains(operation))
@@ -34,5 +41,15 @@ internal static bool IsBlacklistedOperation(String serviceName, string operation
             }
             return false;
         }
+
+        internal static string FormatServiceName(string serviceName)
+        {
+            if (FormattedServiceNames.TryGetValue(serviceName, out string formattedName))
+            {
+                return formattedName;
+            }
+
+            return serviceName;
+        }
     }
 }

sdk/src/Handlers/AwsSdk/Internal/XRayPipelineHandler.cs

@@ -44,9 +44,6 @@ public class XRayPipelineHandler : PipelineHandler
         private const string DefaultAwsWhitelistManifestResourceName = "Amazon.XRay.Recorder.Handlers.AwsSdk.DefaultAWSWhitelist.json";
         private static readonly Logger _logger = Runtime.Internal.Util.Logger.GetLogger(typeof(AWSXRayRecorder));
         private AWSXRayRecorder _recorder;
-        private const String S3RequestIdHeaderKey = "x-amz-request-id";
-        private const String S3ExtendedRequestIdHeaderKey = "x-amz-id-2";
-        private const String ExtendedRquestIdSegmentKey = "id_2";
 
         /// <summary>
         /// Gets AWS service manifest of operation parameter whitelist.
@@ -319,7 +316,8 @@ private void ProcessBeginRequest(IExecutionContext executionContext)
                 _recorder.TraceContext.HandleEntityMissing(_recorder, e, "Cannot get entity while processing AWS SDK request");
             }
 
-            _recorder.BeginSubsegment(RemoveAmazonPrefixFromServiceName(request.ServiceName));
+            var serviceName = RemoveAmazonPrefixFromServiceName(request.ServiceName);
+            _recorder.BeginSubsegment(AWSXRaySDKUtils.FormatServiceName(serviceName));
             _recorder.SetNamespace("aws");
 
             entity = entity == null ? null : _recorder.GetEntity();
@@ -380,20 +378,27 @@ private void ProcessEndRequest(IExecutionContext executionContext)
                 // s3 doesn't follow request header id convention
                 else
                 {
-                    if (requestContext.Request.Headers.TryGetValue(S3RequestIdHeaderKey, out requestId))
+                    if (requestContext.Request.Headers.TryGetValue("x-amz-request-id", out requestId))
                     {
                         subsegment.Aws["request_id"] = requestId;
                     }
 
-                    if (requestContext.Request.Headers.TryGetValue(S3ExtendedRequestIdHeaderKey, out requestId))
+                    if (requestContext.Request.Headers.TryGetValue("x-amz-id-2", out requestId))
                     {
-                        subsegment.Aws[ExtendedRquestIdSegmentKey] = requestId;
+                        subsegment.Aws["id_2"] = requestId;
                     }
                 }
             }
             else
             {
                 subsegment.Aws["request_id"] = responseContext.Response.ResponseMetadata.RequestId;
+
+                // try getting x-amz-id-2 if dealing with s3 request
+                if (responseContext.Response.ResponseMetadata.Metadata.TryGetValue("x-amz-id-2", out string extendedRequestId))
+                {
+                    subsegment.Aws["id_2"] = extendedRequestId;
+                }
+
                 AddResponseSpecificInformation(serviceName, operation, responseContext.Response, subsegment.Aws);
             }
 
@@ -451,6 +456,13 @@ private void ProcessException(AmazonServiceException ex, Entity subsegment)
             _recorder.AddHttpInformation("response", responseAttributes);
 
             subsegment.Aws["request_id"] = ex.RequestId;
+
+            // AmazonId2 property in AmazonS3Exception corresponds to the x-amz-id-2 Http header
+            var property = ex.GetType().GetProperty("AmazonId2");
+            if (property != null)
+            {
+                subsegment.Aws["id_2"] = (string)property.GetValue(ex, null);
+            }
         }
 
         private void AddRequestSpecificInformation(string serviceName, string operation, AmazonWebServiceRequest request, IDictionary<string, object> aws)

sdk/test/UnitTests/AWSSDKHandlerTests.cs

@@ -26,6 +26,7 @@
 using Amazon.Lambda.Model;
 using Amazon.Runtime;
 using Amazon.S3;
+using Amazon.SimpleNotificationService;
 using Amazon.XRay.Recorder.Core;
 using Amazon.XRay.Recorder.Core.Internal.Utils;
 using Amazon.XRay.Recorder.Handlers.AwsSdk;
@@ -80,7 +81,7 @@ public void TestContextMissingStrategyForAWSSDKHandler()
             AWSXRayRecorder.Instance.ContextMissingStrategy = Core.Strategies.ContextMissingStrategy.LOG_ERROR;
             AWSSDKHandler.RegisterXRayForAllServices();
             var dynamo = new AmazonDynamoDBClient(new AnonymousAWSCredentials(), RegionEndpoint.USEast1);
-            CustomResponses.SetResponse(dynamo, null, null, true);
+            CustomResponses.SetResponse(dynamo);
 
             AWSXRayRecorder.Instance.BeginSegment("test dynamo", TraceId);
             var segment = AWSXRayRecorder.Instance.TraceContext.GetEntity();
@@ -101,7 +102,7 @@ public void TestS3SubsegmentNameIsCorrectForAWSSDKHandler()
         {
             AWSSDKHandler.RegisterXRay<IAmazonS3>();
             var s3 = new AmazonS3Client(new AnonymousAWSCredentials(), RegionEndpoint.USEast1);
-            CustomResponses.SetResponse(s3, null, null, true);
+            CustomResponses.SetResponse(s3, null, "TestAmazonId");
 
             _recorder.BeginSegment("test s3", TraceId);
 #if NET45
@@ -115,6 +116,7 @@ public void TestS3SubsegmentNameIsCorrectForAWSSDKHandler()
             Assert.IsTrue(segment.Subsegments[0].Aws.ContainsKey("version_id"));
             Assert.AreEqual(segment.Subsegments[0].Aws["bucket_name"], "testBucket");
             Assert.AreEqual(segment.Subsegments[0].Aws["operation"], "GetObject");
+            Assert.AreEqual(segment.Subsegments[0].Aws["id_2"], "TestAmazonId");
         }
 
         [TestMethod]
@@ -126,7 +128,7 @@ public void TestDynamoDbClient()
             using (var client = new AmazonDynamoDBClient(new AnonymousAWSCredentials(), RegionEndpoint.USEast1))
             {
                 string requestId = @"fakerequ-esti-dfak-ereq-uestidfakere";
-                CustomResponses.SetResponse(client, null, requestId, true);
+                CustomResponses.SetResponse(client, requestId);
 
                 _recorder.BeginSegment("test", TraceId);
 #if NET45
@@ -184,7 +186,7 @@ public void TestRequestResponseParameterAndDescriptorForAWSSDKHandler()
         {
             using (var client = new AmazonDynamoDBClient(new AnonymousAWSCredentials(), RegionEndpoint.USEast1))
             {
-                CustomResponses.SetResponse(client, null, null, true);
+                CustomResponses.SetResponse(client);
                 _recorder.BeginSegment("test", TraceId);
                 var segment = AWSXRayRecorder.Instance.TraceContext.GetEntity();
 
@@ -255,7 +257,7 @@ await client.GetItemAsync(
         public void TestDynamoSubsegmentNameIsCorrectForAWSSDKHandler()
         {
             var dynamo = new AmazonDynamoDBClient(new AnonymousAWSCredentials(), RegionEndpoint.USEast1);
-            CustomResponses.SetResponse(dynamo, null, null, true);
+            CustomResponses.SetResponse(dynamo);
             _recorder.BeginSegment("test dynamo", TraceId);
 #if NET45
             dynamo.ListTables();
@@ -272,7 +274,7 @@ public void TestDynamoSubsegmentNameIsCorrectForAWSSDKHandler()
         public void TestManifestFileNoLambda() //At this point, current manifest file doen't contain Lambda service.
         {
             var lambda = new AmazonLambdaClient(new AnonymousAWSCredentials(), RegionEndpoint.USEast1);
-            CustomResponses.SetResponse(lambda, null, null, true);
+            CustomResponses.SetResponse(lambda);
             _recorder.BeginSegment("lambda", TraceId);
 #if NET45
             lambda.Invoke(new InvokeRequest
@@ -296,7 +298,7 @@ public void TestLambdaInvokeSubsegmentContainsFunctionNameForAWSSDKHandler()
             String temp_path = $"JSONs{Path.DirectorySeparatorChar}AWSRequestInfoWithLambda.json"; //registering manifest file with Lambda
             AWSSDKHandler.RegisterXRayManifest(temp_path);
             var lambda = new AmazonLambdaClient(new AnonymousAWSCredentials(), RegionEndpoint.USEast1);
-            CustomResponses.SetResponse(lambda, null, null, true);
+            CustomResponses.SetResponse(lambda);
             _recorder.BeginSegment("lambda", TraceId);
 #if NET45
             lambda.Invoke(new InvokeRequest
@@ -325,7 +327,7 @@ public void TestRegisterXRayManifestWithStreamLambdaForAWSSDKHandler()
                 AWSSDKHandler.RegisterXRayManifest(stream);
             }
             var lambda = new AmazonLambdaClient(new AnonymousAWSCredentials(), RegionEndpoint.USEast1);
-            CustomResponses.SetResponse(lambda, null, null, true);
+            CustomResponses.SetResponse(lambda);
             _recorder.BeginSegment("lambda", TraceId);
 #if NET45
             lambda.Invoke(new InvokeRequest
@@ -343,5 +345,32 @@ public void TestRegisterXRayManifestWithStreamLambdaForAWSSDKHandler()
 
             Assert.AreEqual("Invoke", segment.Subsegments[0].Aws["operation"]);
         }
+
+        [TestMethod]
+        public void TestSNSSubsegment()
+        {
+            AWSSDKHandler.RegisterXRay<IAmazonSimpleNotificationService>();
+            var sns = new AmazonSimpleNotificationServiceClient(new AnonymousAWSCredentials(), RegionEndpoint.USEast1);
+            CustomResponses.SetResponse(sns);
+
+            _recorder.BeginSegment("test sns", TraceId);
+
+            try
+            {
+                sns.ListTopicsAsync().Wait();
+            }
+            catch
+            {
+                // will throw exception for using anonymous AWS credentials, but will not affect generating traces
+            }
+
+            var segment = _recorder.TraceContext.GetEntity();
+            _recorder.EndSegment();
+            Assert.AreEqual(1, segment.Subsegments.Count);
+            Assert.AreEqual("SNS", segment.Subsegments[0].Name); // Name should be SNS instead of SimpleNotificationService
+            Assert.AreEqual("ListTopics", segment.Subsegments[0].Aws["operation"]);
+            Assert.AreEqual("us-east-1", segment.Subsegments[0].Aws["region"]);
+            Assert.AreEqual("aws", segment.Subsegments[0].Namespace);
+        }
     }
 }

sdk/test/UnitTests/AWSXRayRecorder.UnitTests.csproj

@@ -40,6 +40,7 @@
     <PackageReference Include="AWSSDK.DynamoDBv2" Version="3.5.1.1" />
     <PackageReference Include="AWSSDK.Lambda" Version="3.5.0.25" />
     <PackageReference Include="AWSSDK.S3" Version="3.5.3.1" />
+    <PackageReference Include="AWSSDK.SimpleNotificationService" Version="3.5.0.26" />
     <PackageReference Include="Moq" Version="4.7.137" />
   </ItemGroup>
   <ItemGroup Condition="'$(TargetFramework)' == 'netcoreapp2.0'">

sdk/test/UnitTests/Tools/CustomResponses.cs

@@ -31,9 +31,9 @@ public static class CustomResponses
     {
 #if NET45
         public static void SetResponse(
-            AmazonServiceClient client, string content, string requestId, bool isOK)
+            AmazonServiceClient client, string requestId = null, string s3ExtendedRequestId = null, string content = null, bool isOK = true)
         {
-            var response = Create(content, requestId, isOK);
+            var response = Create(content, requestId, s3ExtendedRequestId, isOK);
             SetResponse(client, response);
         }
 
@@ -54,7 +54,7 @@ public static void SetResponse(
         }
 
         private static Func<MockHttpRequest, HttpWebResponse> Create(
-            string content, string requestId, bool isOK)
+            string content, string requestId, string s3ExtendedRequestId, bool isOK)
         {
             var status = isOK ? HttpStatusCode.OK : HttpStatusCode.NotFound;
 
@@ -65,6 +65,13 @@ private static Func<MockHttpRequest, HttpWebResponse> Create(
                 {
                     headers.Add(HeaderKeys.RequestIdHeader, requestId);
                 }
+
+                // https://aws.amazon.com/premiumsupport/knowledge-center/s3-request-id-values/
+                if (!string.IsNullOrEmpty(s3ExtendedRequestId))
+                {
+                    headers.Add(HeaderKeys.XAmzId2Header, s3ExtendedRequestId);
+                }
+
                 var response = MockWebResponse.Create(status, headers, content);
 
                 if (isOK)
@@ -77,9 +84,9 @@ private static Func<MockHttpRequest, HttpWebResponse> Create(
         }
 #else
 
-        public static void SetResponse(AmazonServiceClient client, string content, string requestId, bool isOK)
+        public static void SetResponse(AmazonServiceClient client, string requestId = null, string s3ExtendedRequestId = null, string content = null, bool isOK = true)
         {
-            var response = Create(content, requestId, isOK);
+            var response = Create(content, requestId, s3ExtendedRequestId, isOK);
             SetResponse(client, response);
         }
 
@@ -98,7 +105,7 @@ public static void SetResponse(AmazonServiceClient client, Func<MockHttpRequest,
         }
 
         private static Func<MockHttpRequest, HttpResponseMessage> Create(
-            string content, string requestId, bool isOK)
+            string content, string requestId, string s3ExtendedRequestId, bool isOK)
         {
             var status = isOK ? HttpStatusCode.OK : HttpStatusCode.NotFound;
 
@@ -109,6 +116,13 @@ private static Func<MockHttpRequest, HttpResponseMessage> Create(
                 {
                     headers.Add(HeaderKeys.RequestIdHeader, requestId);
                 }
+
+                // https://aws.amazon.com/premiumsupport/knowledge-center/s3-request-id-values/
+                if (!string.IsNullOrEmpty(s3ExtendedRequestId))
+                {
+                    headers.Add(HeaderKeys.XAmzId2Header, s3ExtendedRequestId);
+                }
+
                 var response = MockWebResponse.Create(status, headers, content);
 
                 if (isOK)